from forecast import get_forecast_data
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))
    return round(Ya, 2)


def get_yield_forecast(latitude: float, longitude: float, scenario: str = "pessimist", shading_coef: float = 0.):
    monthly_forecast = get_forecast_data(latitude, longitude, scenario=scenario, shading_coef=shading_coef)

    Kc, Ky = get_cultural_coefficients()

    Yx = get_maximum_theoretical_yield()

    soil_properties = get_soil_properties(latitude, longitude)

    ETo = monthly_forecast["et0"]

    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, ETx, ETa, Ky)

    return projected_yield