# -*- coding: utf-8 -*-
"""
Code to run the ETLook model.
"""
import os
import numpy as np
import pywapor.et_look_dev as ETLook_dev
import pywapor.et_look_v2 as ETLook_v2
import pywapor.general as g
import pywapor.general.processing_functions as PF
import pywapor.general.pre_defaults as defaults
import xarray as xr
import pandas as pd
from pywapor.general.compositer import calculate_ds
import copy

def main(input_data, et_look_version = "v2", export_vars = "default", export_to_tif = False):
    """Runs the ETLook model over the provided input data.

    Parameters
    ----------
    input_data : str | xr.Dataset
        Dataset generated by `pywapor.pre_et_look`.
    et_look_version : "v2" | "dev", optional
        Which version of the ETLook model to use, by default "v2".
    export_vars : "default" | "all" | list, optional
        Specify which variables to save inside the output file. `"Default"` stores `int_mm`,
        `t_24_mm`, `e_24_mm`, `et_24_mm`, `et_ref_24_mm`, `se_root`, `biomass_prod`,
        `epoch_ends` and `epoch_starts`. `"all"` stores all calculated variables. Use a
        list to specify a custom output set, by default "default".
    export_to_tif : bool, optional
        Return the variables selected with `export_vars` as netCDF (False) 
        or geoTIFF (True), by default False.

    Returns
    -------
    xr.Dataset | dict
        Returns a dataset or a dictionary with variable names as keys and
        lists with paths to geoTIFF files as values, depending on `export_to_tif`.
    """

    # Version
    if et_look_version == "v2":
        ETLook = ETLook_v2
        print("--> Running ETLook_v2")
    elif et_look_version == "dev":
        ETLook = ETLook_dev
        print("--> Running ETLook_dev")

    # Allow skipping of et_look-functions if not all of its required inputs are
    # available.
    g.lazifier.decorate_submods(ETLook_v2, g.lazifier.etlook_decorator)
    g.lazifier.decorate_submods(ETLook_dev, g.lazifier.etlook_decorator)

    # Inputs
    if isinstance(input_data, str):
        ds = xr.open_dataset(input_data)
    else:
        ds = input_data
        input_data = ds.encoding["source"]

    # Add constants to ds
    # ds = ds.assign(defaults.constants_defaults())

    ds = g.variables.initiate_ds(ds)

    if et_look_version == "dev":
        ds["nd_min"] = 0.1
        ds["tenacity"] = 1.0

    # Constants TODO: move to pre_et_look
    doy_epoch_start = [int(pd.Timestamp(x).strftime("%j")) for x in ds["epoch_starts"].values]
    doy_epoch_end = [int(pd.Timestamp(x).strftime("%j")) for x in ds["epoch_ends"].values]
    doy = [int((x+y)/2) for x, y in zip(doy_epoch_start, doy_epoch_end)]
    ds["doy"] = xr.DataArray(doy, coords = ds["epoch_starts"].coords)

    # ds["sc"] = ETLook.solar_radiation.seasonal_correction(ds["doy"])
    ds["decl"] = ETLook.solar_radiation.declination(ds["doy"])
    ds["iesd"] = ETLook.solar_radiation.inverse_earth_sun_distance(ds["doy"])

    ######################## MODEL ETLOOK ####################################

    # **effective_leaf_area_index*********************************************
    ds["vc"] = ETLook.leaf.vegetation_cover(ds["ndvi"], nd_min = ds["nd_min"], nd_max = ds["nd_max"], vc_pow = ds["vc_pow"])
    ds["lai"] = ETLook.leaf.leaf_area_index(ds["vc"], vc_min = ds["vc_min"], vc_max = ds["vc_max"], lai_pow = ds["lai_pow"])
    ds["lai_eff"] = ETLook.leaf.effective_leaf_area_index(ds["lai"])
    
    # *******TRANSPIRATION COMPONENT******************************************
    
    # **soil fraction*********************************************************
    ds["sf_soil"] = ETLook.radiation.soil_fraction(ds["lai"])

    # **atmospheric canopy resistance***********************************************
    ds["lat_rad"] = ETLook.solar_radiation.latitude_rad(ds["lat_deg"])
    ds["ws"] = ETLook.solar_radiation.sunset_hour_angle(ds["lat_rad"], ds["decl"])

    ds["ra_24_toa_flat"] = ETLook.solar_radiation.daily_solar_radiation_toa_flat(ds["decl"], ds["iesd"], ds["lat_rad"], ds["ws"])
    ds["trans_24"] = ETLook.solar_radiation.transmissivity(ds["ra_24"], ds["ra_24_toa_flat"])

    ds["stress_rad"] = ETLook.stress.stress_radiation(ds["ra_24"])
    ds["p_air_0_24_mbar"] = ETLook.meteo.air_pressure_kpa2mbar(ds["p_air_0_24"])
    ds["p_air_24"] = ETLook.meteo.air_pressure_daily(ds["z"], ds["p_air_0_24_mbar"])
    ds["vp_24"] = ETLook.meteo.vapour_pressure_from_specific_humidity_daily(ds["qv_24"], ds["p_air_24"])

    ds["svp_24"] = ETLook.meteo.saturated_vapour_pressure(ds["t_air_24"])
    
    if et_look_version == "dev":
        ds["svp_24"] = ETLook.meteo.saturated_vapour_pressure_average(
                    ETLook.meteo.saturated_vapour_pressure_maximum(ds["t_air_max_24"]),
                    ETLook.meteo.saturated_vapour_pressure_minimum(ds["t_air_min_24"]))

    ds["vpd_24"] = ETLook.meteo.vapour_pressure_deficit_daily(ds["svp_24"], ds["vp_24"])
    ds["stress_vpd"] = ETLook.stress.stress_vpd(ds["vpd_24"], ds["vpd_slope"])
    ds["stress_temp"] = ETLook.stress.stress_temperature(ds["t_air_24"], t_opt = ds["t_opt"], t_min = ds["t_min"], t_max = ds["t_max"])

    if et_look_version == "dev":
        if isinstance(ds["land_mask"], xr.DataArray):
            ds["rs_min"] = xr.where(ds["land_mask"] == 3, 400, 100)
        else:
            ds["rs_min"] = np.where(ds["land_mask"] == 3, 400, 100)

    ds["r_canopy_0"] = ETLook.resistance.atmospheric_canopy_resistance(ds["lai_eff"], ds["stress_rad"], ds["stress_vpd"], ds["stress_temp"], rs_min = ds["rs_min"], rcan_max = ds["rcan_max"])

    # **net radiation canopy******************************************************
    ds["t_air_k_24"] = ETLook.meteo.air_temperature_kelvin_daily(ds["t_air_24"])
    ds["l_net"] = ETLook.radiation.longwave_radiation_fao(ds["t_air_k_24"], ds["vp_24"], ds["trans_24"], vp_slope = ds["vp_slope"], vp_offset = ds["vp_offset"], lw_slope = ds["lw_slope"], lw_offset = ds["lw_offset"])
    ds["int_mm"] = ETLook.evapotranspiration.interception_mm(ds["p_24"], ds["vc"], ds["lai"], int_max = ds["int_max"])
    ds["lh_24"] = ETLook.meteo.latent_heat_daily(ds["t_air_24"])
    ds["int_wm2"] = ETLook.radiation.interception_wm2(ds["int_mm"], ds["lh_24"])
    ds["rn_24"] = ETLook.radiation.net_radiation(ds["r0"], ds["ra_24"], ds["l_net"], ds["int_wm2"])
    ds["rn_24_canopy"] = ETLook.radiation.net_radiation_canopy(ds["rn_24"], ds["sf_soil"])

    # find water region using ndvi
    if et_look_version == "dev":
        if ds.pixel_size < 250:
            if isinstance(ds["land_mask"], xr.DataArray):
                ds["land_mask"] = xr.where(ds["land_mask"] == 2, 1, ds["land_mask"])
                ds["land_mask"] = xr.where(ds["ndvi"] < 0, 2, ds["land_mask"])
            else:
                ds["land_mask"] = np.where(ds["land_mask"] == 2, 1, ds["land_mask"])
                ds["land_mask"] = np.where(ds["ndvi"] < 0, 2, ds["land_mask"])

    ds["stress_moist"] = ETLook.stress.stress_moisture(ds["se_root"], tenacity = ds["tenacity"])
    ds["r_canopy"] = ETLook.resistance.canopy_resistance(ds["r_canopy_0"], ds["stress_moist"], rcan_max = ds["rcan_max"])

    # **initial canopy aerodynamic resistance***********************************************************

    ds["z_obst"] = ETLook.roughness.obstacle_height(ds["ndvi"], ds["z_obst_max"], ndvi_obs_min = ds["ndvi_obs_min"], ndvi_obs_max = ds["ndvi_obs_max"], obs_fr = ds["obs_fr"])
    ds["z0m"] = ETLook.roughness.roughness_length(ds["lai"], ds["z_oro"], ds["z_obst"], ds["z_obst_max"], ds["land_mask"])
    ds["u_24"] = ETLook.meteo.wind_speed(ds["u2m_24"], ds["v2m_24"])
    ds["ra_canopy_init"] = ETLook.neutral.initial_canopy_aerodynamic_resistance(ds["u_24"], ds["z0m"], z_obs = ds["z_obs"])

    # **windspeed blending height daily***********************************************************
    ds["u_b_24"] = ETLook.meteo.wind_speed_blending_height_daily(ds["u_24"], z_obs = ds["z_obs"], z_b = ds["z_b"])

    # **ETLook.neutral.initial_daily_transpiration***********************************************************
    ds["ad_dry_24"] = ETLook.meteo.dry_air_density_daily(ds["p_air_24"], ds["vp_24"], ds["t_air_k_24"])
    ds["ad_moist_24"] = ETLook.meteo.moist_air_density_daily(ds["vp_24"], ds["t_air_k_24"])
    ds["ad_24"] = ETLook.meteo.air_density_daily(ds["ad_dry_24"], ds["ad_moist_24"])
    ds["psy_24"] = ETLook.meteo.psychrometric_constant_daily(ds["p_air_24"], ds["lh_24"])
    ds["ssvp_24"] = ETLook.meteo.slope_saturated_vapour_pressure_daily(ds["t_air_24"])
    ds["t_24_init"] = ETLook.neutral.initial_daily_transpiration(ds["rn_24_canopy"], ds["ssvp_24"], ds["ad_24"], ds["vpd_24"], ds["psy_24"], ds["r_canopy"], ds["ra_canopy_init"])

    # **ETLook.unstable.initial_sensible_heat_flux_canopy_daily***********************************************************
    ds["h_canopy_24_init"] = ETLook.unstable.initial_sensible_heat_flux_canopy_daily(ds["rn_24_canopy"], ds["t_24_init"])

    # **ETLook.unstable.initial_friction_velocity_daily***********************************************************
    ds["disp"] = ETLook.roughness.displacement_height(ds["lai"], ds["z_obst"], land_mask = ds["land_mask"], c1 = ds["c1"])
    ds["u_star_24_init"] = ETLook.unstable.initial_friction_velocity_daily(ds["u_b_24"], ds["z0m"], ds["disp"], z_b = ds["z_b"])

    # **ETLook.unstable.transpiration***********************************************************
    ds["t_24"] = ETLook.unstable.transpiration(ds["rn_24_canopy"], ds["ssvp_24"], ds["ad_24"], ds["vpd_24"], ds["psy_24"], ds["r_canopy"], ds["h_canopy_24_init"], ds["t_air_k_24"], ds["u_star_24_init"], ds["z0m"], ds["disp"], ds["u_b_24"], z_obs = ds["z_obs"], z_b = ds["z_b"], iter_h = ds["iter_h"])
    ds["t_24_mm"] = ETLook.unstable.transpiration_mm(ds["t_24"], ds["lh_24"])

    if et_look_version == "dev":
        ds["r_canopy_unstressed"] = ETLook.resistance.unstressed_canopy_resistance(ds["r_canopy"], ds["stress_moist"])
        ds["tpot_24"] = ETLook.unstable.transpiration(ds["rn_24_canopy"], ds["ssvp_24"], ds["ad_24"], ds["vpd_24"], ds["psy_24"], ds["r_canopy_unstressed"], ds["h_canopy_24_init"], ds["t_air_k_24"], ds["u_star_24_init"], ds["z0m"], ds["disp"], ds["u_b_24"], z_obs = ds["z_obs"], z_b = ds["z_b"], iter_h = ds["iter_h"])
        ds["tpot_24_mm"] = ETLook.unstable.transpiration_mm(ds["tpot_24"], ds["lh_24"])

    #*******EVAPORATION COMPONENT****************************************************************

    # **ETLook.radiation.net_radiation_soil***********************************************************
    ds["rn_24_soil"] = ETLook.radiation.net_radiation_soil(ds["rn_24"], ds["sf_soil"])

    # **ETLook.resistance.soil_resistance***********************************************************
    ds["r_soil"] = ETLook.resistance.soil_resistance(ds["se_root"], land_mask = ds["land_mask"], r_soil_pow = ds["r_soil_pow"], r_soil_min = ds["r_soil_min"])

    # **ETLook.resistance.soil_resistance***********************************************************
    ds["ra_soil_init"] = ETLook.neutral.initial_soil_aerodynamic_resistance(ds["u_24"], z_obs = ds["z_obs"])

    # **ETLook.unstable.initial_friction_velocity_soil_daily***********************************************************
    ds["u_star_24_soil_init"] = ETLook.unstable.initial_friction_velocity_soil_daily(ds["u_b_24"], ds["disp"], z_b = ds["z_b"])

    # **ETLook.unstable.initial_sensible_heat_flux_soil_daily***********************************************************
    ds["stc"] = ETLook.radiation.soil_thermal_conductivity(ds["se_root"])
    ds["vhc"] = ETLook.radiation.volumetric_heat_capacity(ds["se_root"], porosity = ds["porosity"])    
    
    if et_look_version == "dev":
        ds["stc"] = ETLook.radiation.soil_thermal_conductivity(se_top = ds["se_top"])
        ds["vhc"] = ETLook.radiation.volumetric_heat_capacity(se_top = ds["se_top"], porosity = ds["porosity"])

    ds["dd"] = ETLook.radiation.damping_depth(ds["stc"], ds["vhc"])
    ds["g0_bs"] = ETLook.radiation.bare_soil_heat_flux(ds["doy"], ds["dd"], ds["stc"], ds["t_amp_year"], ds["lat_rad"])
    ds["g0_24"] = ETLook.radiation.soil_heat_flux(ds["g0_bs"], ds["sf_soil"], ds["land_mask"], ds["rn_24_soil"], ds["trans_24"], ds["ra_24"], ds["l_net"], ds["rn_slope"], ds["rn_offset"])
    ds["e_24_init"] = ETLook.neutral.initial_daily_evaporation(ds["rn_24_soil"], ds["g0_24"], ds["ssvp_24"], ds["ad_24"], ds["vpd_24"], ds["psy_24"], ds["r_soil"], ds["ra_soil_init"])
    ds["h_soil_24_init"] = ETLook.unstable.initial_sensible_heat_flux_soil_daily(ds["rn_24_soil"], ds["e_24_init"], ds["g0_24"])

    # **ETLook.unstable.evaporation***********************************************************
    ds["e_24"] = ETLook.unstable.evaporation(ds["rn_24_soil"], ds["g0_24"], ds["ssvp_24"], ds["ad_24"], ds["vpd_24"], ds["psy_24"], ds["r_soil"], ds["h_soil_24_init"], ds["t_air_k_24"], ds["u_star_24_soil_init"], ds["disp"], ds["u_b_24"], z_b = ds["z_b"], z_obs = ds["z_obs"], iter_h = ds["iter_h"])
    ds["e_24_mm"] = ETLook.unstable.evaporation_mm(ds["e_24"], ds["lh_24"])
    ds["et_24_mm"] = ETLook.evapotranspiration.et_actual_mm(ds["e_24_mm"], ds["t_24_mm"])
    
    if et_look_version == "dev":
        ds["e_24_mm"] = np.clip(ds["e_24_mm"], 0, np.inf)
        ds["et_24_mm"] = np.clip(ds["et_24_mm"], 0, np.inf)

    # **ETLook.unstable.evaporation***********************************************************
    ds["rn_24_grass"] = ETLook.radiation.net_radiation_grass(ds["ra_24"], ds["l_net"], r0_grass = ds["r0_grass"])
    ds["et_ref_24"] = ETLook.evapotranspiration.et_reference(ds["rn_24_grass"], ds["ad_24"], ds["psy_24"], ds["vpd_24"], ds["ssvp_24"], ds["u_24"])
    ds["et_ref_24_mm"] = ETLook.evapotranspiration.et_reference_mm(ds["et_ref_24"], ds["lh_24"])

    if et_look_version == "dev":
        ds["lue"] = ETLook_dev.biomass.lue(ds["lue_max"], ds["stress_temp"], ds["stress_moist"], ds["eps_a"])
        ds["fpar"] = ETLook_dev.leaf.fpar(ds["vc"], ds["ndvi"])
        ds["apar"] = ETLook_dev.leaf.apar(ds["ra_24"], ds["fpar"])       
        ds["biomass_prod"] = ETLook_dev.biomass.biomass(ds["apar"], ds["lue"])         
    else:
        ds["t_air_k_min"] = ETLook.meteo.air_temperature_kelvin_daily(ds["t_air_min_24"])
        ds["t_air_k_max"] = ETLook.meteo.air_temperature_kelvin_daily(ds["t_air_max_24"])

        ds["t_air_k_12"] = ETLook.meteo.mean_temperature_kelvin_daytime(ds["t_air_k_min"], ds["t_air_k_max"])

        ds["t_dep"] = ETLook.biomass.temperature_dependency(ds["t_air_k_12"], dh_ap=ds["dh_ap"], d_s=ds["d_s"], dh_dp=ds["dh_dp"])
        ds["k_m"] = ETLook.biomass.affinity_constant_co2(ds["t_air_k_12"])
        ds["k_0"] = ETLook.biomass.inhibition_constant_o2(ds["t_air_k_12"])
        ds["tau_co2_o2"] = ETLook.biomass.co2_o2_specificity_ratio(ds["t_air_k_12"])

        ds["year"] = ds.epoch_starts.dt.year

        ds["co2_act"] = ETLook.biomass.co2_level_annual(ds["year"])
        ds["a_d"] = ETLook.biomass.autotrophic_respiration(ds["t_air_k_24"], ar_slo=ds["ar_slo"], ar_int=ds["ar_int"])

        ds["apar"] = ETLook.biomass.par(ds["ra_24"])
        ds["f_par"] = ETLook.biomass.fpar(ds["ndvi"], fpar_slope=ds["fpar_slope"], fpar_offset=ds["fpar_offset"])

        ds["co2_fert"] = ETLook.biomass.co2_fertilisation(ds["tau_co2_o2"], ds["k_m"], ds["k_0"], ds["co2_act"], o2=ds["o2"], co2_ref=ds["co2_ref"])
        ds["npp_max"] = ETLook.biomass.net_primary_production_max(ds["t_dep"], ds["co2_fert"], ds["a_d"], ds["apar"], gcgdm=ds["gcgdm"])
        ds["npp"] = ETLook.biomass.net_primary_production(ds["npp_max"], ds["f_par"], ds["stress_moist"], phot_eff=ds["phot_eff"])


    ds = ds.drop_vars([x for x in ds.variables if ds[x].dtype == object])

    fp, fn = os.path.split(input_data)

    ds = g.variables.fill_attrs(ds)
    g.network.create_network(ds, os.path.join(fp, f"et_look_{et_look_version}.html"))

    if export_vars == "all":
        ...
    elif export_vars == "default":
        keep_vars = ['int_mm',
                    't_24_mm',
                    'e_24_mm',
                    'et_24_mm',
                    'et_ref_24_mm',
                    'se_root',
                    'npp',
                    'epoch_ends',
                    'epoch_starts']
        ds = PF.ds_remove_except(ds, keep_vars)
    elif isinstance(export_vars, list):
        keep_vars = copy.copy(export_vars)
        keep_vars = np.unique(keep_vars + ['epoch_ends', 'epoch_starts']).tolist()
        ds = PF.ds_remove_except(ds, keep_vars)
    else:
        raise ValueError

    ds = ds.transpose("epoch", "lat", "lon") # set dimension order the same for all vars.
    
    fn = fn.replace("_input", "_output")
    ds, fh = calculate_ds(ds, os.path.join(fp, fn), "--> Saving outputs.")

    if export_to_tif:
        files = PF.export_ds_to_tif(ds, keep_vars, None)
        ds.close()
        os.remove(fh)
        return files
    else:
        return ds

if __name__ == "__main__":

    # project_folder = r"/Volumes/Data/FAO/WaPOR_vs_pyWaPOR/pyWAPOR_v1"
    # startdate = date = "2021-07-01"

    # level = "level_1"
    et_look_version = "v2"

    input_data = r"/Users/hmcoerver/Downloads/et_look_input.nc"
    input_data = xr.open_dataset(input_data)
    # input_data = input_data.drop_vars(["ndvi"])

    input_data = input_data.assign(defaults.constants_defaults())

    ds = main(input_data,
                et_look_version=et_look_version, 
                export_vars="default")

    # input_data["year"] = input_data.epoch_starts.dt.year


