import numpy as np


## INPUT variables

# maximum temperature in Celsius
# minimum temperature in Celsius
# mean temperature in Celsius
# minimum relative humidity in percentage
# maximum relative humidity in percentage
# wind speed at 2 meters in m/s
# solar radiation in MJ/m^2/day
# doy is day of the year (1-365)

# Constants
Gsc = 0.0820  # Solar constant (MJ/m^2/min)
sigma = 4.903e-9  # Stefan-Boltzmann constant (MJ/K^4/m^2/day)

# Functions to calculate various components needed for ET0 calculation
def calculate_extraterrestrial_radiation(latitude, doy):
    phi = np.radians(latitude)
    dr = 1 + 0.033 * np.cos(2 * np.pi / 365 * doy)
    delta = 0.409 * np.sin(2 * np.pi / 365 * doy - 1.39)
    omega = np.arccos(-np.tan(phi) * np.tan(delta))
    Ra = (24 * 60 / np.pi) * Gsc * dr * (omega * np.sin(phi) * np.sin(delta) + np.cos(phi) * np.cos(delta) * np.sin(omega))
    return Ra

def calculate_clear_sky_radiation(Ra, elevation):
    return (0.75 + 2e-5 * elevation) * Ra

def calculate_net_solar_radiation(Rs):
    return (1 - 0.23) * Rs

def calculate_net_outgoing_longwave_radiation(T_max, T_min, Rs, Rso, ea):
    T_max_K = T_max + 273.15  # Convert to Kelvin
    T_min_K = T_min + 273.15  # Convert to Kelvin
    return sigma * ((T_max_K ** 4 + T_min_K ** 4) / 2) * (0.34 - 0.14 * np.sqrt(ea)) * (1.35 * Rs / Rso - 0.35)

def calculate_slope_of_saturation_vapor_pressure_curve(T_mean):
    return 4098 * (0.6108 * np.exp(17.27 * T_mean / (T_mean + 237.3))) / (T_mean + 237.3) ** 2

def calculate_psychrometric_constant(elevation):
    P = 101.3 * ((293 - 0.0065 * elevation) / 293) ** 5.26
    return 0.665e-3 * P

def calculate_saturation_vapor_pressure(T):
    return 0.6108 * np.exp(17.27 * T / (T + 237.3))

def calculate_mean_saturation_vapor_pressure(T_max, T_min):
    es_Tmax = calculate_saturation_vapor_pressure(T_max)
    es_Tmin = calculate_saturation_vapor_pressure(T_min)
    return (es_Tmax + es_Tmin) / 2

def calculate_actual_vapor_pressure(T_min, T_max, RH_min, RH_max):
    es_min = calculate_saturation_vapor_pressure(T_min)
    es_max = calculate_saturation_vapor_pressure(T_max)
    ea = (es_min * RH_max / 100 + es_max * RH_min / 100) / 2
    return ea

def calculate_reference_evapotranspiration(T_max, T_min, T_mean, RH_min, RH_max, u2, Rs, elevation, latitude, doy):
    # Calculate radiation components
    Ra = calculate_extraterrestrial_radiation(latitude, doy)
    Rso = calculate_clear_sky_radiation(Ra, elevation)
    Rns = calculate_net_solar_radiation(Rs)

    # Calculate vapor pressure components
    es = calculate_mean_saturation_vapor_pressure(T_max, T_min)
    ea = calculate_actual_vapor_pressure(T_min, T_max, RH_min, RH_max)

    # Calculate net outgoing longwave radiation
    Rnl = calculate_net_outgoing_longwave_radiation(T_max, T_min, Rs, Rso, ea)

    # Calculate net radiation
    Rn = Rns - Rnl

    # Calculate slope of saturation vapor pressure curve
    Delta = calculate_slope_of_saturation_vapor_pressure_curve(T_mean)

    # Calculate psychrometric constant
    gamma = calculate_psychrometric_constant(elevation)

    # Soil heat flux density (G) is 0 for daily timestep
    G = 0

    # Calculate reference evapotranspiration (ET0)
    ET0 = (0.408 * Delta * (Rn - G) + gamma * (900 / (T_mean + 273)) * u2 * (es - ea)) / (Delta + gamma * (1 + 0.34 * u2))

    return ET0



# # Define Kc values based on FAO guidelines
def calculate_kc_fao(Kcini, Kcmid, Kcend, doy):
    """Calculate the Kc value based on the day of the year (doy)"""

    # Define stage durations (these are approximations, adjust based on your specific vineyard phenology)
    start_day = 90  # Bud break (end of March)
    mid_day = start_day + 60  # Initial stage ends
    full_cover_day = mid_day + 40  # Crop development stage ends
    harvest_day = full_cover_day + 70  # Midseason stage ends
    end_season_day = harvest_day + 40  # Late season ends

    if doy < mid_day:
        return Kcini
    elif doy < full_cover_day:
        return Kcini + (Kcmid - Kcini) * (doy - mid_day) / (full_cover_day - mid_day)
    elif doy < harvest_day:
        return Kcmid
    elif doy < end_season_day:
        return Kcmid + (Kcend - Kcmid) * (doy - harvest_day) / (end_season_day - harvest_day)
    else:
        return Kcend
    
# Generic arable crops
def kc_generic(NDVI):
    return 1.457 * NDVI - 0.1725

# Campos et al. 2010 (irrigated vineyards, La Mancha region)
def kc_campos(NDVI):
    return 1.44 * NDVI - 0.10

# Er-Raki et al., 2013 (table grapes in northwest Mexico)
def kc_erraki(NDVI):
    return 0.1808 * np.exp(1.3138 * NDVI)

def kc_tasumi(NDVI):
    return 1.18 * NDVI + 0.04



def calculate_etc(T_max, T_min, T_mean, RH_min, RH_max, u2, Rs, elevation, latitude, doy, Kcini=0.30, Kcmid=0.70, Kcend=0.45):
    et0 = calculate_reference_evapotranspiration(T_max, T_min, T_mean, RH_min, RH_max, u2, Rs, elevation, latitude, doy)
    # ndvi = get_ndvi(NIR, RED) # NDVI = (NIR - RED) / (NIR + RED)
    kc = calculate_kc_fao(Kcini, Kcmid, Kcend, doy)
    return et0*kc