# -*- coding: utf-8 -*-

# Water retention related implementation.
# Based on Notebook implemenation created by Marko Petrovic

# Numerical data processing
import pandas as pd
import numpy as np
from scipy import stats

# Time series operations
import glob
from datetime import datetime

class Calibration():
    """Water retention and impact estimate related parameters setting."""
    # Conversion coefficient meter to millimeter
    m_to_mm = 1

    # Extinction coefficient =0.7 for trees and 0.3 for shrubs (Wang et al. 2008)
    kappa = 0.7

    # Shade factor, is the percentage of sky covered by foliage and branches within
    # the perimeter of individual tree crowns, can vary by species from about
    # 60% to 95% when trees are in-leaf (McPherson, 1984).
    # The value below is set according to Glasgow mean
    avg_shade_factor = 0.85

    # Specific leaf storage of water (sl =0.0002 m).
    leaf_storage = 0.0002 * m_to_mm

    # Leaf-on leaf_off transition days (Wang_et_al_2008).
    leaf_transition_days = 28

    # Specific impervious cover storage of water (=0.0015 m).
    maximum_impervious_cover_storage = 0.0015 * m_to_mm

    # Specific pervious cover storage of water (=0.001 m).
    maximum_pervious_cover_storage = 0.001 * m_to_mm

    def __init__(
        self,
        leaf_transition_days = 28,
        leaf_storage = 0.0002,
        pervios_storage_max = 0.0010,
        impervios_storage_max = 0.0015
        ):
        """The constructor method.

        Args:
            leaf_transition_days: (:obj:`int`): Leaf-on leaf_off transition days
            leaf_storage: (:obj:`float`): Specific leaf storage of water.
            pervios_storage_max: (:obj:`float`): Specific pervious cover storage of water
            impervios_storage_max: (:obj:`float`): Specific impervious cover storage of water
        Returns:
            None
        Note:

        TODO:
            * pass optional settings as key word parameters.
        """
        self.leaf_storage = leaf_storage * Calibration.m_to_mm
        self.leaf_transition_days = leaf_transition_days
        self.maximum_impervious_cover_storage = impervios_storage_max * Calibration.m_to_mm
        self.maximum_pervious_cover_storage = pervios_storage_max * Calibration.m_to_mm

    def set_surface_storage_rates(
        self, leaf_storage=None,
        pervios_storage_max=None,
        impervios_storage_max=None):
        """Setting specific water storage rates. 

        Args:
            leaf_storage: (:obj:`float`): Specific leaf storage of water.
            pervios_storage_max: (:obj:`float`): Specific pervious cover storage of water
            impervios_storage_max: (:obj:`float`): Specific impervious cover storage of water
        Returns:
            None
        Note:

        Todo:

        """
        if leaf_storage:
            self.leaf_storage = leaf_storage * Calibration.m_to_mm
        if impervios_storage_max:
            self.maximum_impervious_cover_storage = impervios_storage_max * Calibration.m_to_mm
        if pervios_storage_max:
            self.maximum_pervious_cover_storage = pervios_storage_max * Calibration.m_to_mm
    

def compute_leaf_area_index(
    dbh,
    tree_height,
    crown_height,
    crown_width,
    crown_missing = 0,
    shade_factor = Calibration.avg_shade_factor):
    """The function given allometrics of a tree computes its leaf, bark and plant area indices. 

    Args:
        dbh: (:obj:`float`): the diameter in cm of the trunk usually measured at 1.3m from the ground.
        tree_height: (:obj:`float`): The tree height in meters.
        crown_height: (:obj:`float`): The vertical length of tree crown in meters.
        crown_width: (:obj:`float`): The horizontal length (diameter) of tree crown in meters.
        crown_missing: (:obj:`float`): The percentage loss of the crown.
        shade_factor: (:obj:`float`): The percentage of sky covered by foliage and branches.

    Returns:
        (:obj:`tuple`): the tuple returns the tree indices (LAI,BAI,PAI)
    Note:
        The beta multipliers and the main equation is based on Nowak (1996).

    TODO:
        Parametrize beta multipliers.
    """
    loss = crown_missing
    th = tree_height
    cw = crown_width
    ch = crown_height
    sf = shade_factor
    beta_0 = -4.3309
    beta_1 = 0.2942
    beta_2 = 0.7312
    beta_3 = 5.7217
    beta_4 = 0.0148

    def compute_under_canopy_area(crown_width):
        return pow((crown_width/2),2) * np.pi
    
    def compute_bark_area(dbh, tree_height, crown_height):
        # * 0.01 converts DBH(cm) into meter.
        return np.pi * (dbh * 0.01) * (tree_height - crown_height)

    # Outer surface area estimate below is based on Gacka-Grzesikiewicz (1980).
    under_canopy = compute_under_canopy_area(cw)
    crown_surface = np.pi  * crown_width * (crown_height + crown_width)/2
    bark_area = compute_bark_area(dbh,th,ch)
    leaf_area = (1-loss) * np.exp(beta_0 + beta_1 * th + beta_2 * cw + beta_3 * sf - beta_4 * crown_surface)
    leaf_area_index = leaf_area / under_canopy
    bark_area_index = bark_area / under_canopy
    plant_area_index = leaf_area_index + bark_area_index
    return (leaf_area_index, bark_area_index, plant_area_index)


def pai_seasons(x, leaf_on_start, leaf_off_start, leaf_transition_days):
    """The method updates Plant Area Index (PAI) with respect to leaf on-off seasons.

    Args:
        A data frame "x" with following variables:
        Date_time: (:obj:`time`): date and time
        BAI: Bark Area Index
        LAI: Leaf Area Index
        conifers: it taks value "true" if conifers and "false" otherwise 
        
        leaf_on_start: a day in the year when the leaf on season starts. In the case of Glasgow it is April 14 (day 105).
                        See also other sources: https://weatherspark.com/y/147740/Average-Weather-at-Glasgow-Airport-United-Kingdom-Year-Round
        leaf_off_start: a day in the year when the leaf off season starts. In the case of Glasgow it is November 2 (day 307).
        leaf_transition_days: the number of days that the leaf on-off transition last.
        

    Returns:
        An array of values of Plant Area Indexes (PAI) over time.

    Note:
        The function is built upon the following paper:
        Wang, Jun, Theodore A. Endreny, and David J. Nowak. “Mechanistic Simulation of Tree 
        Effects in an Urban Water Balance Model 1.” JAWRA Journal of the American Water Resources 
        Association 44, no. 1 (February 2008): 75–85. https://doi.org/10.1111/j.1752-1688.2007.00139.x   

    Todo:
    """
    x = x.assign(PAI = np.where(x['Conifers'], x['BAI']+x['LAI'],
                                np.where(x.Date_time.dt.day_of_year < leaf_on_start, x['BAI'],
                                    np.where(((x.Date_time.dt.day_of_year >= leaf_on_start) & (x.Date_time.dt.day_of_year < leaf_off_start)),
                                    x['LAI']/(1+np.exp(-0.37*(x.Date_time.dt.day_of_year-(leaf_on_start+leaf_transition_days/2)))) + x['BAI'],
                                    x['LAI']/(1+np.exp(-0.37*((leaf_off_start+leaf_transition_days/2)-x.Date_time.dt.day_of_year))) + x['BAI'])))
    )
    return x['PAI']

def lmbd(temperature):
    """The method calculates latent heat of vaporization.

    Args:
        Temperature in [C]     

    Returns:
        latent heat of vaporization in [MJ/Kg]

    Note:
        The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015.  

    Todo:
    """
    return 2.501-0.002361*temperature

def e_s(temperature):
    """The method calculates saturated vapor pressure.

    Args:
        Temperature in [C]   

    Returns:
        Saturated vapor pressure in [kPa]

    Note:
            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    return 0.6108*np.exp(17.27*temperature/(237.3+temperature))

def e(dew_point_temperature):
    """The method calculates vapor pressure.

    Args:
        Dew point temperature in [C]   

    Returns:
        Vapor pressure in [kPa]

    Note:
            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    return 0.6108*np.exp(17.27*dew_point_temperature/(237.3+dew_point_temperature))

def DELTA(temperature):
    """The method calculates the slope of vapor pressure temperature curve.

    Args:
        Temperature in [C]   

    Returns:
        Slope of vapor pressure temperature curve in [kPa/C]

    Note:
        It includes in calculation saturated vapor pressure in [kPa] which is calculated with function e_s()!

        The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    return 4098*e_s(temperature)/(237.3+temperature)**2

def rho_a(temperature, surface_pressure):
    """The method calculates the density of air.

    Args:
        Temperature in [C]
        Surface preasure in [kPa]  

    Returns:
        The density of air in [kg/m^3]

    Note:
            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    return 3.486 * surface_pressure/(275+temperature)

def rho_w(temperature):
    """The method calculates the density of water.

    Args:
        Temperature in [C]

    Returns:
        The density of water in [kg/m^3]

    Note:
            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    return 999.88 + 0.018*temperature - 0.0051*temperature**2

def D(temperature, dew_point_temperature):
    """The method calculates vapor pressure deficit.

    Args:
        Temperature in [C]
        Dew point temperature in [C]  

    Returns:
        Vapor pressure deficit in [kPa]

    Note:
            It includes in calculation vapor pressure as well as saturated vapor pressure, both in [kPa], which 
            are calculated with functions e() e_s().

            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    return np.maximum(e_s(temperature)-e(dew_point_temperature),0)

def U_t(wind_speed, wind_estimate_height):
    """The method estimates wind speed at the tree top.

    Args:
        Measured wind speed [m/s] 
        Wind estimate height (tree height) [m]

 
    Returns:
        Wind speed at the tree top in [m/s]

    Note:
            The method is using two other parameters:
                Wind measurement height Z_u (usually 10m) and roughness height for water d_w which are set as constants.

            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    return wind_speed*np.log(wind_estimate_height/d_w)/np.log(Z_u/d_w)

def r_a(wind_speed, wind_estimate_height, roughness_height):
    """The method calculates aerodynamic resistance.

    Args:
        Measured wind speed [m/s] 
        Wind estimate height (tree height) [m]
        Roughness_height [m]
 
    Returns:
        Aerodynamic resistance in [m/s]

    Note:
            It includes in calculation wind speed at the tree top U_t.
            If roughness height is negative, a new equation is applied to calculate aerodynamic resistance for
            the evapotranspiration from the soil.

            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    if roughness_height >= 0:
        return 4.72*np.log(wind_estimate_height/(Z_ov*roughness_height))/(1+0.53*U_t(wind_speed, wind_estimate_height))
    else:
        return 208/U_t(wind_speed, wind_estimate_height)

def gamma(temperature,surface_pressure):
    """The method calculates psychrometric constant.

    Args:
        Temperature [C] 
        Surface pressure [kPa]
 
    Returns:
        Psychrometric constant in [kPA/C]

    Note:
            It includes in calculation latent heat of vaporization lmbd().

            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    #return 10**(-3)*c_p*surface_pressure*10/(lmbd(temperature)*0.622)# pressure in mbar
    return 10**(-3)*c_p*surface_pressure/(lmbd(temperature)*0.622)# pressure in kPa

def r_s(pai):
    """The method calculates stomatal resistance.

    Args:
        Plant area index (PAI) 
 
    Returns:
        Stomatal resistance in [s/m]

    Note:
            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    return 200/pai

def C_leaf(temperature):
    """The method calculates water vapor concentration at the evaporating surfaces within the leaf.

    Args:
        Temperature [C] 

    Returns:
        Water vapor concentration at the evaporating surfaces within the leaf in [g/m^3]

    Note:
            Temerature in celsius [C] is converted to kelvins [K] such that T(K) = T(C)+273.15.
            The formula also includes saturated vapor pressure e_s() in [kPa].  

            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    return 2165*e_s(temperature)/(temperature+273.15)

def C_air(dew_point_temperature, temperature):
    """The method calculates water vapor concentration in the air.

    Args:
        Temperature [C] 
        Dew point temperature [C]

    Returns:
        Water vapor concentration in the air in [g/m^3]

    Note:
            Temerature in celsius [C] is converted to kelvins [K] such that T(K) = T(C)+273.15.
            The formula also includes vapor pressure e() in [kPa].  

            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    return 2165*e(dew_point_temperature)/(temperature+273.15)

def potential_evaporation(temperature, dew_point, solar_radiation, sea_level_pressure, pai, wind_speed, wind_estimate_height, roughness_height):
    """The method calculate potential evaporation.

    Args: a data frame
    
    Returns: the updated ata frame 
        
    Note:

    Todo:
    """
    potential_evaporation = M_TO_MM*(1/(lmbd(temperature)*rho_w(temperature)))*(DELTA(temperature)*solar_radiation+(rho_a(temperature, sea_level_pressure)*c_p*D(temperature, dew_point))/r_a(wind_speed, wind_estimate_height, roughness_height))/(DELTA(temperature)+gamma(temperature, sea_level_pressure)*(1+(r_s(pai)/r_a(wind_speed, wind_estimate_height, roughness_height))))

    return potential_evaporation

def potential_evapotranspiration(df):
    """The method calculate potential evapotranspiration.

    Args: a data frame
    
    Returns: the updated ata frame 
        
    Note:

    Todo:
    """
    df = df.assign(
        Potential_evaporation_v = lambda x: potential_evaporation(x.Temperature, x.Dew_Point, x.Solar_Radiation, x.Sea_Level_Pressure, x.PAI, x.Wind_Speed, x.height, roughness_height_trees),
        Potential_evaporation_g = lambda x: potential_evaporation(x.Temperature, x.Dew_Point, x.Solar_Radiation, x.Sea_Level_Pressure, 1    , x.Wind_Speed, 1       , roughness_height_bare_soil),
        PET                     = lambda x: potential_evaporation(x.Temperature, x.Dew_Point, x.Solar_Radiation, x.Sea_Level_Pressure, x.PAI, x.Wind_Speed, x.height, -1),
        Transpiration = lambda x: 10**(-6)*(3600/x.PAI)*np.maximum((C_leaf(x.Temperature)-C_air(x.Dew_Point, x.Temperature)),0)/(r_s(x.PAI)+r_a(x.Wind_Speed, x.height, roughness_height_trees))
        )
    df = df.assign(
        TF_average_ratio = lambda x: np.where(((x.PET>x.Transpiration) & (x.PET>0)), x.Transpiration/x.PET, np.nan)
        )

    df['TF_average_ratio'] = df.groupby(['AgentID','Step'])['TF_average_ratio'].transform('mean')
    df = df.assign(
        Transpiration = lambda x: np.where(((x.PAI<(x.LAI+x.BAI)) | (x.Transpiration > x.PET)), x.TF_average_ratio*x.PET, x.Transpiration)
        )

    df = df.drop(columns='TF_average_ratio')
    return df

def ped1(Date_time, Precipitation, Potential_evaporation, Maximum_storage, evp):
    """The method computes precipitation-evaporation dynamics for one type at a time, which can be trees, pervious cover, impervious cover...

    Args:
        Date_time [time identifier]
        Precipitation [amount of water hitting the surface area of the type]
        Potential_evaporation [of the type for the given weather conditions]
        Maximum_storage [parameter: maximum storage of the type]
        evp [evaporation coefficient for the type]

    Returns:
        Precipitation storage of the type for each hour.

    Note:

            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    df=pd.DataFrame({'Date_time':Date_time, 'Precipitation':Precipitation, 
    'Potential_evaporation':Potential_evaporation})
    df["Storage"] = 0
    df['Potential_evaporation_lag1'] = df.Potential_evaporation.shift(1).interpolate(limit_direction ='backward')
    x = 0
    def func2(row):
        # non local variable ==> will use pre_value from the ped function
        nonlocal x
        new_value =  new_value =  np.maximum(0,(np.minimum(Maximum_storage,x) + row['Precipitation']-
        np.minimum(np.minimum(Maximum_storage,x),((np.minimum(Maximum_storage,x)/
        Maximum_storage)**(evp)*row['Potential_evaporation_lag1']))))
        x = new_value
        return new_value
    # This line might throw a SettingWithCopyWarning warning
    df.loc[0:,'Storage'] = df.loc[0:,:].apply(func2, axis=1)
    return df["Storage"]

def ped3(df, evp_v, evp_g):
    """The method computes precipitation-evaporation dynamics for each tree as well as the impervious and pervious covers.

    Args:
        A data frame including the following variables:
        Evaporation coefficient for trees: evp_v
        Evaporation coefficient for ground: evp_g

    Returns:
        A data frame with new variables: Vegetation_storage, Impervious_cover_storage_v, Pervious_cover_storage_v for each hour.

    Note:

            The function is built upon the following book:
            Maidment, David R and others. “Handbook of Hydrology, McGraw-Hil.” Inc., New York, NY, 1992.
            
            And a paper:
                Hirabayashi, Satoshi. “I-Tree Eco United States County-Based Hydrologic Estimates.” 
                Washington, DC: US Department of Agriculture, Forest Service, 2015. 

    Todo:
    """
    df['Storage'] = 0
    df['Potential_evaporation_v_lag1'] = df.Potential_evaporation_v.shift(1).interpolate(limit_direction ='backward')
    df['Potential_evaporation_g_lag1'] = df.Potential_evaporation_g.shift(1).interpolate(limit_direction ='backward')
    x = 0
    y = 0
    z = 0
    def func2(row):
        # non local variable ==> will use pre_value from the ped function
        nonlocal x
        nonlocal y
        nonlocal z
        new_value_x =  np.maximum(0,(np.minimum(row['Maximum_vegetation_storage'],x) + row['Canopy_interception']-
        np.minimum(np.minimum(row['Maximum_vegetation_storage'],x),((np.minimum(row['Maximum_vegetation_storage'],x)/
        row['Maximum_vegetation_storage'])**(evp_v)*row['Potential_evaporation_v_lag1']))))

        Precipitation_on_the_ground_with_vegetation = np.maximum(0,(new_value_x-row['Maximum_vegetation_storage']))+row['Through_canopy_precipitation']
        
        new_value_y = np.maximum(0,(np.minimum(MAXIMUM_IMPERVIOUS_COVER_STORAGE,y) + Precipitation_on_the_ground_with_vegetation-
        np.minimum(np.minimum(MAXIMUM_IMPERVIOUS_COVER_STORAGE,y),((np.minimum(MAXIMUM_IMPERVIOUS_COVER_STORAGE,y)/
        MAXIMUM_IMPERVIOUS_COVER_STORAGE)**(evp_g)*row['Potential_evaporation_g_lag1']))))
        
        new_value_z = np.maximum(0,(np.minimum(MAXIMUM_PERVIOUS_COVER_STORAGE,z) + Precipitation_on_the_ground_with_vegetation-
        np.minimum(np.minimum(MAXIMUM_PERVIOUS_COVER_STORAGE,z),((np.minimum(MAXIMUM_PERVIOUS_COVER_STORAGE,z)/
        MAXIMUM_PERVIOUS_COVER_STORAGE)**(evp_g)*row['Potential_evaporation_g_lag1']))))
        x = new_value_x
        y = new_value_y
        z = new_value_z
        return [new_value_x, new_value_y, new_value_z]
    # This line might throw a SettingWithCopyWarning warning
    df.loc[0:,'Storage'] = df.loc[0:,:].apply(func2, axis=1)
    df['Vegetation_storage'] = pd.DataFrame(df['Storage'].tolist())[0]
    df['Impervious_cover_storage_v'] = pd.DataFrame(df['Storage'].tolist())[1]
    df['Pervious_cover_storage_v'] = pd.DataFrame(df['Storage'].tolist())[2]
    df = df.drop(columns='Storage')
    return df

def ecosystem_services(i):
    """The method parallelize the computation of water retention benefits.

    Args: agent index "i"
    
    Returns: water retention benefits for each agent [data frame] 
        
    Note:

    Todo:
    """
    output = df_scenario[tree_population.AgentID == AGENTS[i]]
    output = output[['Step', 'AgentID', 'height','BAI', 'LA','LAI', 'PAI', 'Conifers', 'Under_canopy_area', 'Total_under_canopy_area','Scenario', 'Precipitation_scale', 
    'SAMPLE_AREA', 'IMPERVIOUS_COVER_SHARE', 'PERVIOUS_COVER_SHARE', 'POPULATION_AREA', 'POPULATION_SAMPLE_TREE_RATIO']]
    output = pd.merge(output, weather_forcast, on=['Step'], how='left').sort_values(['Step', 'AgentID'])
    # Calculate plant area index (PAI) in leaf-on and leaf-off seasons
    output['PAI'] = pai_seasons(output[['Date_time', 'BAI', 'LAI', 'Conifers']],
                                        Leaf_on_transition_day_start, Leaf_off_transition_day_start, 
                                        LEAF_TRANSITION_DAYS)

    
    # Calculate potential evapotranspiration over the vegetation and ground areas
    output = potential_evapotranspiration(output)

    output = output.assign(
        Canopy_cover_fraction = lambda x: 1-np.exp(-KAPPA*x.PAI),
        Maximum_vegetation_storage = lambda x: SL*x.PAI,
        Through_canopy_precipitation = lambda x: x.Precipitation*(1-x.Canopy_cover_fraction),
        Canopy_interception = lambda x: x.Precipitation-x.Through_canopy_precipitation
    )

    output = ped3(output,
        evp_v = 2/3, 
        evp_g = 1
    )
    output = output.assign(
        Canopy_drip = lambda x: np.maximum(0,(x.Vegetation_storage-x.Maximum_vegetation_storage)),
        Evaporation_from_vegetation = lambda x: np.maximum(0,(x.Canopy_interception-x.Canopy_drip)),
        Precipitation_on_the_ground_with_vegetation = lambda x: x.Canopy_drip + x.Through_canopy_precipitation,
        Vegetation_storage = lambda x: np.minimum(x.Vegetation_storage, x.Maximum_vegetation_storage),
        Run_off_v = lambda x: np.maximum(0,(x.Impervious_cover_storage_v-MAXIMUM_IMPERVIOUS_COVER_STORAGE)),
        Evaporation_from_impervious_cover_v = lambda x: np.maximum(0,(x.Precipitation_on_the_ground_with_vegetation-x.Run_off_v)),
        Impervious_cover_storage_v = lambda x: np.minimum(x.Impervious_cover_storage_v, MAXIMUM_IMPERVIOUS_COVER_STORAGE),
        Infiltration_v = lambda x: np.maximum(0,(x.Pervious_cover_storage_v-MAXIMUM_PERVIOUS_COVER_STORAGE)),
        Evaporation_from_pervious_cover_v = lambda x: np.maximum(0,(x.Precipitation_on_the_ground_with_vegetation-x.Infiltration_v)),
        Pervious_cover_storage_v = lambda x: np.minimum(x.Pervious_cover_storage_v, MAXIMUM_PERVIOUS_COVER_STORAGE)
    )

    # Annual aggregation
    output = output.groupby(['Scenario','Precipitation_scale','AgentID','Step']).agg(
        Leaf_area = pd.NamedAgg(column='LA', aggfunc=np.mean),
        Under_canopy_area = pd.NamedAgg(column='Under_canopy_area', aggfunc=np.mean),
        SAMPLE_AREA = pd.NamedAgg(column='SAMPLE_AREA', aggfunc=np.mean),
        IMPERVIOUS_COVER_SHARE = pd.NamedAgg(column='IMPERVIOUS_COVER_SHARE', aggfunc=np.mean),
        PERVIOUS_COVER_SHARE = pd.NamedAgg(column='PERVIOUS_COVER_SHARE', aggfunc=np.mean),
        POPULATION_AREA = pd.NamedAgg(column='POPULATION_AREA', aggfunc=np.mean),
        POPULATION_SAMPLE_TREE_RATIO = pd.NamedAgg(column='POPULATION_SAMPLE_TREE_RATIO', aggfunc=np.mean),
        Annual_precipitation = pd.NamedAgg(column='Precipitation', aggfunc=np.sum),
        Annual_canopy_interception_loss = pd.NamedAgg(column='Evaporation_from_vegetation', aggfunc=np.sum),
        Annual_run_off_v = pd.NamedAgg(column='Run_off_v', aggfunc=np.sum),
        Annual_evaporation_from_impervious_cover_v = pd.NamedAgg(column='Evaporation_from_impervious_cover_v', aggfunc=np.sum),
        Annual_infiltration_v=pd.NamedAgg(column="Infiltration_v", aggfunc=np.sum),
        Annual_evaporation_from_pervious_cover_v=pd.NamedAgg(column="Evaporation_from_pervious_cover_v", aggfunc=np.sum),
        Annual_transpiration=pd.NamedAgg(column="Transpiration", aggfunc=np.sum),
        Total_under_canopy_area = pd.NamedAgg(column='Total_under_canopy_area', aggfunc=np.mean),
    ).reset_index()
    output['Annual_stormwater_retention'] = output.Annual_transpiration + output.Annual_canopy_interception_loss + output.Annual_evaporation_from_impervious_cover_v*output.IMPERVIOUS_COVER_SHARE + output.Annual_evaporation_from_pervious_cover_v*output.PERVIOUS_COVER_SHARE

    return output
