#%% 0. Imports
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


#%% 1. Open Data
data = pd.read_excel("/Users/mikas/Library/CloudStorage/GoogleDrive-mikeamlima@gmail.com/O meu disco/Aulas/Meteorologia/PLs/US_Standard_Simplified_0-120km_0.5km.xlsx")

altitude = data["altitude_km"].values
pressure = data["pressure_Pa"].values / 100
density = data["density_kg_m3"].values

temperature = data["temperature_K"].values
N2 = data["N2_molefrac"].values * pressure
O  = data["O_molefrac"].values * pressure
O2 = data["O2_molefrac"].values * pressure
O3 = data["O3_ppmv"].values * pressure
Ar = data["Ar_molefrac"].values * pressure


#%% 2. Exercise a)
#%%% 2.1. Plot in function of altitude
#### Initialize figure
fig = plt.figure(figsize=(12,12))
ax = fig.add_subplot(1,1,1)
ax.grid(ls='--', alpha=0.6)

#### Plot variables
ax.plot(temperature, altitude, color='red', lw=2, label='Temperature [K]')

#### Adjust plot
ax.set_title('Composição vertical média da atmosfera', pad=10, fontsize=20)
ax.set_ylabel('Altitude [km]', fontsize=18); ax.set_ylim(0,120)
ax.set_xlabel('Temperatura [K]', fontsize=18); ax.set_xlim(180,290)
ax.tick_params(axis='both', which='major', labelsize=16)

#%%% 2.2. Plot in function of pressure without log scale
#### Initialize figure
fig = plt.figure(figsize=(12,12))
ax = fig.add_subplot(1,1,1)
ax.grid(ls='--', alpha=0.6)

#### Plot variables
ax.plot(temperature, pressure, color='red', lw=2, label='Temperature [K]')

#### Adjust plot
ax.set_title('Composição vertical média da atmosfera', pad=10, fontsize=20)
ax.set_ylabel('Pressure [hPa]', fontsize=18); ax.invert_yaxis(); ax.set_ylim(np.max(pressure),np.min(pressure))
ax.set_xlabel('Temperatura [K]', fontsize=18); ax.set_xlim(180,290)
ax.tick_params(axis='both', which='major', labelsize=16)

#%%% 2.3. Plot in function of pressure with log scale
#### Initialize figure
fig = plt.figure(figsize=(12,12))
ax = fig.add_subplot(1,1,1)
ax.grid(ls='--', alpha=0.6)

#### Plot variables
ax.plot(temperature, pressure, color='red', lw=2, label='Temperature [K]')

#### Adjust plot
ax.set_title('Composição vertical média da atmosfera', pad=10, fontsize=20)
ax.set_ylabel('Pressure [hPa]', fontsize=18); ax.set_yscale("log"); ax.invert_yaxis(); ax.set_ylim(np.max(pressure),np.min(pressure))
ax.set_xlabel('Temperatura [K]', fontsize=18); ax.set_xlim(180,290)
ax.tick_params(axis='both', which='major', labelsize=16)


#%% 3. Exercise b)
#%%% 3.1. Plot in function of altitude
#### Initialize figure
fig = plt.figure(figsize=(12,12))
ax = fig.add_subplot(1,1,1)
ax.grid(ls='--', alpha=0.6)

#### Plot variables
ax.plot(temperature, altitude, color='red', lw=2, label='Temperature [K]')

#### Adjust plot
ax.set_title('Composição vertical média da atmosfera', pad=10, fontsize=20)
ax.set_ylabel('Altitude [km]', fontsize=18); ax.set_ylim(0,120)
ax.set_xlabel('Temperatura [K]', fontsize=18); ax.set_xlim(180,290)
ax.tick_params(axis='both', which='major', labelsize=16)

#### Mark the layers
ax.text(182,4,'Troposfera', fontsize=24)
ax.hlines(11,180,300, ls='--', color='tab:green', lw=3)
ax.text(265,14,'Tropopausa', fontsize=24, color='tab:green')

ax.text(182,27,'Estratosfera', fontsize=24)
ax.hlines(47,180,300, ls='--', color='tab:green', lw=3)
ax.text(182,50,'Estratopausa', fontsize=24, color='tab:green')

ax.text(266,63,'Mesosfera', fontsize=24)
ax.hlines(85,180,300, ls='--', color='tab:green', lw=3)
ax.text(265,88,'Mesopausa', fontsize=24, color='tab:green')

ax.text(265,105,'Termosfera', fontsize=24)


#%% 4. Exercise c)
#%%% 4.1. Plot in function of altitude
#### Initialize figure
fig = plt.figure(figsize=(10,14))
ax = fig.add_subplot(1,1,1)
ax.grid(ls='--', alpha=0.6)

# main (bottom) axis
ax.plot(temperature, altitude, color='red', lw=2, label='Temperature [K]')
ax.set_ylabel('Altitude [km]', fontsize=18)
ax.set_ylim(0,120)
ax.set_xlabel('Temperatura [K]', fontsize=18, color='red')
ax.set_xlim(180,290)
ax.tick_params(axis='both', which='major', labelsize=16)
ax.tick_params(axis='x', which='major', labelsize=16, colors='red')

# position box to overlay new axes
pos = ax.get_position()

# ---------------------------
# first top axis: partial pressures (green)
# ---------------------------
ax_add = fig.add_axes(pos, sharey=ax, frameon=False)   # independent x-axis, shared y-axis
ax_add.plot(O2, altitude, lw=2, label='$O_2$', ls='-', color='green')
ax_add.plot(N2, altitude, lw=2, label='$N_2$', ls='--', color='green')
ax_add.plot(Ar, altitude, lw=2, label='$Ar$', ls=':', color='green')
ax_add.plot(O, altitude, lw=2, label='$O$', ls='-.', color='green')
ax_add.set_xscale('log')

ax_add.xaxis.set_ticks_position('top')
ax_add.xaxis.set_label_position('top')
ax_add.spines['top'].set_position(('axes', 1.0))   # at the normal top
ax_add.patch.set_visible(False)                    # make axis background transparent
ax_add.set_xlabel('Pressão parcial [hPa]', fontsize=18, color='green')
ax_add.tick_params(axis='x', which='major', labelsize=16, colors='green')
ax_add.legend(fontsize=18, loc='upper right')

# ---------------------------
# second top axis: O3 (blue), stacked above previous top axis
# ---------------------------
ax_add2 = fig.add_axes(pos, sharey=ax, frameon=True)  # independent x-axis as well
ax_add2.plot(O3, altitude, color='tab:blue', lw=2, label='$O_3$')

ax_add2.xaxis.set_ticks_position('top')
ax_add2.xaxis.set_label_position('top')
ax_add2.spines['top'].set_position(('axes', 1.07))     # push it above ax_add
ax_add2.patch.set_visible(False)
ax_add2.set_xlabel('Concentração [ppmv]', fontsize=18, color='tab:blue')
ax_add2.tick_params(axis='x', which='major', labelsize=16, colors='tab:blue')

# stack order so main axis is behind the top axes
ax_add.get_yaxis().set_visible(False)
ax_add2.get_yaxis().set_visible(False)
ax.set_zorder(1)
ax_add.set_zorder(2)
ax_add2.set_zorder(3)