from cmt.components import Meteorology, SnowDegreeDay
met, sno = Meteorology(), SnowDegreeDay()

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


%cd input

met.initialize('meteorology-2.cfg')
sno.initialize('snow_degree_day-1.cfg')

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


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

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


count = 1
while sno.get_current_time() < sno.get_end_time():
    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')
    P_snow = met.get_value('atmosphere_water__snowfall_leq-volume_flux')
    T_surf = met.get_value('land_surface__temperature')
    rho_H2O = met.get_value('water-liquid__mass-per-volume_density')
    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('water-liquid__mass-per-volume_density', rho_H2O)

    sno.update(sno.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())
    
    count += 1

print time

print snow_depth

print air_temp

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

Get variables from SnowDegreeDay; set into Meteorology
Advance Meteorology
Get variables from Meteorology; set into SnowDegreeDay
Advance SnowDegreeDay
End of explanation
"""


met.finalize(), sno.finalize()

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


%matplotlib inline
from matplotlib import pyplot as plt

plt.plot(time[1:], snow_depth[1:])
plt.title('Snow depth versus time')
plt.xlabel('Time [s]')
plt.ylabel('Snow depth [m]')

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


