import pandas as pd

from forecast import get_forecast_data
from retrieve_coefs_max_yield import get_coefs_Kc_Ky_and_max_yield
from utils.soil_utils import get_soil_properties


def calculate_ETx(Kc, ETo):
    """
    Calculate the maximum evapotranspiration (ETx) using the crop coefficient (Kc) and reference evapotranspiration (ETo).

    Parameters:
    Kc (float): Crop coefficient
    ETo (float): Reference evapotranspiration (mm)

    Returns:
    float: Maximum evapotranspiration (ETx) in mm
    """
    ETx = Kc * ETo
    return ETx

def calculate_ETa(ETx, soil_moisture, field_capacity, wilting_point):
    """
    Calculate the actual evapotranspiration (ETa) using the maximum evapotranspiration (ETx), soil moisture, field capacity, and wilting point.

    Parameters:
    ETx (float): Maximum evapotranspiration (mm)
    soil_moisture (float): Current soil moisture content (%)
    field_capacity (float): Field capacity of the soil (%)
    wilting_point (float): Wilting point of the soil (%)

    Returns:
    float: Actual evapotranspiration (ETa) in mm
    """
    if soil_moisture > field_capacity:
        ETa = ETx
    elif soil_moisture < wilting_point:
        ETa = 0
    else:
        ETa = ETx * ((soil_moisture - wilting_point) / (field_capacity - wilting_point))
    
    return ETa


def calculate_yield_projection(Yx, ETx, ETa, Ky):
    """
    Calculate the agricultural yield projection using the FAO water production function.

    Parameters:
    Yx (float): Maximum yield (quintal/ha)
    ETx (float): Maximum evapotranspiration (mm)
    ETa (float): Actual evapotranspiration (mm)
    Ky (float): Yield response factor

    Returns:
    float: Projected yield (quintal/ha)
    """

    Ya = Yx * (1 - Ky * (1 - ETa / ETx))
    Ya.loc[ETx == 0] = 0

    return round(Ya, 2)


def add_cultural_coefs(monthly_forecast: pd.DataFrame, cultural_coefs: pd.DataFrame) -> pd.DataFrame:
    monthly_forecast["Kc"] = 0
    monthly_forecast["Ky"] = 0
    for month in range(1, 13):
        Kc = cultural_coefs["Kc"][cultural_coefs.Mois == month].iloc[0]
        Ky = cultural_coefs["Ky"][cultural_coefs.Mois == month].iloc[0]
        monthly_forecast.loc[(monthly_forecast.month==month).to_numpy(), "Kc"] = Kc
        monthly_forecast.loc[(monthly_forecast.month==month).to_numpy(), "Ky"] = Ky
    return monthly_forecast


def compute_yield_forecast(
        latitude: float,
        longitude: float,
        culture: str = "Colza d'hiver",
        region: str = "Bourgogne-Franche-Comté",
        scenario: str = "pessimist",
        shading_coef: float = 0.,
):
    monthly_forecast = get_forecast_data(latitude, longitude, scenario=scenario, shading_coef=shading_coef)

    cultural_coefs, max_yield = get_coefs_Kc_Ky_and_max_yield(culture, region)
    monthly_forecast = add_cultural_coefs(monthly_forecast, cultural_coefs)
    Kc = monthly_forecast["Kc"]
    Ky = monthly_forecast["Ky"]

    soil_properties = get_soil_properties(latitude, longitude)

    ETo = monthly_forecast["Evaporation (including sublimation and transpiration) (kg m-2 s-1)"]

    ETx = calculate_ETx(Kc, ETo)

    ETa = calculate_ETa(
        ETx,
        soil_properties["soil_moisture"],
        soil_properties["field_capacity"],
        soil_properties["wilting_point"],
    )

    projected_yield = calculate_yield_projection(
        Yx=max_yield,
        ETx=ETx,
        ETa=ETa,
        Ky=Ky)
    monthly_forecast["Estimated yield (quintal/ha)"] = projected_yield

    return monthly_forecast


def get_annual_yield(monthly_forecast: pd.DataFrame) -> pd.Series:
    yield_forecast = pd.Series(
        index=monthly_forecast["time"],
        data=monthly_forecast["Estimated yield (quintal/ha)"].to_numpy(),
    )
    yield_forecast = yield_forecast.resample("1YE").mean()
    return yield_forecast


if __name__ == '__main__':
    monthly_forecast = compute_yield_forecast(
        latitude=47,
        longitude=5,
        culture="Colza d'hiver",
        scenario="pessimist",
        shading_coef=0.,
    )
    print(monthly_forecast.head())

    yield_forecast = get_annual_yield(monthly_forecast)
    print(yield_forecast)