from cmt.components import EvapEnergyBalance, Meteorology, SnowEnergyBalance
evp, met, sno = EvapEnergyBalance(), Meteorology(), SnowEnergyBalance()

"""
Explanation: EvapEnergyBalance-Meteorology-SnowEnergyBalance coupling
Goal: Try to successfully run a coupled EvapEnergyBalance-Meteorology-SnowEnergyBalance simulation, with EvapEnergyBalance as the driver.
Import the Babel-wrapped EvapEnergyBalance, Meteorology and SnowEnergyBalance components and create instances:
End of explanation
"""


%cd input

evp.initialize('evap_energy_balance-1.cfg')
met.initialize('meteorology-2.cfg')
sno.initialize('snow_energy_balance-1.cfg')

"""
Explanation: Initialize the components with cfg files that, for simplicity, use the same time step and run duration:
End of explanation
"""


time = [met.get_current_time()]
snow_depth = [sno.get_value('snowpack__depth').max()]
air_temp = [met.get_value('atmosphere_bottom_air__temperature').max()]
evap_flux = [evp.get_value('land_surface_water__evaporation_volume_flux').max()]

"""
Explanation: Store initial values of time, snow depth, and air temperature:
End of explanation
"""


count = 1
while evp.get_current_time() < evp.get_end_time():  
    rho_air = met.get_value('atmosphere_bottom_air__mass-per-volume_density')
    Cp_air = met.get_value('atmosphere_bottom_air__mass-specific_isobaric_heat_capacity')
    T_air = met.get_value('atmosphere_bottom_air__temperature')
    P_snow = met.get_value('atmosphere_water__snowfall_leq-volume_flux')
    T_surf = met.get_value('land_surface__temperature')
    Q_sum = met.get_value('land_surface_net-total-energy__energy_flux')
    rho_H2O = met.get_value('water-liquid__mass-per-volume_density')
    sno.set_value('atmosphere_bottom_air__mass-per-volume_density', rho_air)
    sno.set_value('atmosphere_bottom_air__mass-specific_isobaric_heat_capacity', Cp_air)
    sno.set_value('atmosphere_bottom_air__temperature', T_air)
    sno.set_value('atmosphere_water__snowfall_leq-volume_flux', P_snow)
    sno.set_value('land_surface__temperature', T_surf)
    sno.set_value('land_surface_net-total-energy__energy_flux', Q_sum)
    sno.set_value('water-liquid__mass-per-volume_density', rho_H2O)
    
    sno.update(sno.get_time_step()*count)

    rho_snow = sno.get_value('snowpack__z_mean_of_mass-per-volume_density')
    h_snow = sno.get_value('snowpack__depth')
    h_swe = sno.get_value('snowpack__liquid-equivalent_depth')
    SM = sno.get_value('snowpack__melt_volume_flux')
    met.set_value('snowpack__z_mean_of_mass-per-volume_density', rho_snow)
    met.set_value('snowpack__depth', h_snow)
    met.set_value('snowpack__liquid-equivalent_depth', h_swe)
    met.set_value('snowpack__melt_volume_flux', SM)

    met.update(met.get_time_step()*count)
    
    T_air = met.get_value('atmosphere_bottom_air__temperature')
    Qe = met.get_value('atmosphere_bottom_air_land_net-latent-heat__energy_flux')
    Q_sum = met.get_value('land_surface_net-total-energy__energy_flux')
    T_surf = met.get_value('land_surface__temperature')
    h_snow = sno.get_value('snowpack__depth')
    evp.set_value('atmosphere_bottom_air__temperature', T_air)
    evp.set_value('atmosphere_bottom_air_land_net-latent-heat__energy_flux', Qe)
    evp.set_value('land_surface_net-total-energy__energy_flux', Q_sum)
    evp.set_value('land_surface__temperature', T_surf)
    evp.set_value('snowpack__depth', h_snow)
    
    evp.update(evp.get_time_step()*count)
    
    time.append(met.get_current_time())
    snow_depth.append(sno.get_value('snowpack__depth').max())
    air_temp.append(met.get_value('atmosphere_bottom_air__temperature').max())
    evap_flux.append(evp.get_value('land_surface_water__evaporation_volume_flux').max())
    
    count += 1

print time

print snow_depth

print air_temp

print evap_flux

"""
Explanation: Run the coupled models to completion. In each time step, perform the following actions:

Get variables from Meteorology; set into SnowEnergyBalance
Advance SnowEnergyBalance
Get variables from SnowEnergyBalance; set into Meteorology
Advance Meteorology
Get variables from Meteorology and SnowEnergyBalance; set into EvapEnergyBalance
Advance EvapEnergyBalance
End of explanation
"""


evp.finalize(), met.finalize(), sno.finalize()

"""
Explanation: Finalize the components:
End of explanation
"""


%matplotlib inline
from matplotlib import pyplot as plt

fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(12, 5))

snow_depth_plot = axes[0].plot(time[1:], snow_depth[1:], 'b')
axes[0].set_title('Snow depth versus time')
axes[0].set_xlabel('Time [s]')
axes[0].set_ylabel('Snow depth [m]')

evap_flux_plot = axes[1].plot(time[1:], evap_flux[1:], 'r')
axes[1].set_title('Evaporative flux versus time')
axes[1].set_xlabel('Time [s]')
axes[1].set_ylabel('Evaporative flux [m s-1]')

"""
Explanation: Plot snow depth versus time.
End of explanation
"""


