# -*- coding: UTF-8 -*-
# Python
"""
04-07-2023
@author: jeremy auclair

Test file
"""

import xarray as xr
# from dask.distributed import Client
import os
import numpy as np
import rasterio as rio
from typing import List, Tuple, Union
import netCDF4 as nc
from tqdm import tqdm
from parameters.params_samir_class import samir_parameters
from config.config import config
from time import time
# import webbrowser  # to open dask dashboard


def rasterize_samir_parameters(csv_param_file: str, empty_dataset: xr.Dataset, land_cover_raster: str, chunk_size: dict) -> Tuple[xr.Dataset, dict]:
    """
    Creates a raster `xarray` dataset from the csv parameter file, the land cover raster and an empty dataset
    that contains the right structure (emptied ndvi dataset for example). For each parameter, the function loops
    on land cover classes to fill the raster.

    # Arguments
    1. csv_param_file: `str`
        path to csv paramter file
    2. empty_dataset: `xr.Dataset`
        empty dataset that contains the right structure (emptied ndvi dataset for example).
    3. land_cover_raster: `str`
        path to land cover netcdf raster
    4. chunk_size: `dict`
        chunk_size for dask computation

    # Returns
    1. parameter_dataset: `xr.Dataset`
        the dataset containing all the rasterized Parameters
    2. scale_factor: `dict`
        dictionnary containing the scale factors for each parameter
    """

    # Load samir params into an object
    table_param = samir_parameters(csv_param_file)

    # Set general variables
    class_count = table_param.table.shape[1] - 2  # remove dtype and default columns

    # Open land cover raster
    land_cover = xr.open_dataarray(land_cover_raster, chunks=chunk_size)

    # Create dataset
    parameter_dataset = empty_dataset.copy(deep=True)

    # Create dictionnary containing the scale factors
    scale_factor = {}

    # Loop on samir parameters and create
    for parameter in table_param.table.index[1:]:

        # Create new variable and set attributes
        parameter_dataset[parameter] = land_cover.copy(deep=True).astype('f4')
        parameter_dataset[parameter].attrs['name'] = parameter
        parameter_dataset[parameter].attrs['description'] = 'cf SAMIR Doc for detail'
        parameter_dataset[parameter].attrs['scale factor'] = str(
            table_param.table.loc[table_param.table.index == parameter]['scale_factor'].values[0])

        # Assigne value in dictionnary
        scale_factor[parameter] = 1/int(table_param.table.loc[table_param.table.index == parameter]['scale_factor'].values[0])

        # Loop on classes to set parameter values for each class
        for class_val, class_name in zip(range(1, class_count + 1), table_param.table.columns[2:]):

            # Parameter values are multiplied by the scale factor in order to store all values as int16 types
            # These values are then rounded to make sure there isn't any decimal point issues when casting the values to int16
            # TODO vr: formule trop longue, trouver un moyen de rendre lisible
            parameter_dataset[parameter].values = np.where(parameter_dataset[parameter].values == class_val, round(table_param.table.loc[table_param.table.index == parameter][class_name].values[0] * table_param.table.loc[table_param.table.index == parameter]['scale_factor'].values[0]), parameter_dataset[parameter].values).astype('f4')

    # Return dataset converted to 'int16' data type to reduce memory usage
    # and scale_factor dictionnary for later conversion
    return parameter_dataset, scale_factor


def setup_time_loop(calculation_variables_t1: List[str], calculation_variables_t2: List[str], empty_dataset: xr.Dataset) -> Tuple[xr.Dataset, xr.Dataset]:
    """
    Creates two temporary `xarray Datasets` that will be used in the SAMIR time loop.
    `variables_t1` corresponds to the variables for the previous day and `variables_t2`
    corresponds to the variables for the current day. After each loop, `variables_t1`
    takes the value of `variables_t2` for the corresponding variables.

    # Arguments
    1. calculation_variables_t1: `List[str]`
        list of strings containing the variable names
        for the previous day dataset
    2. calculation_variables_t2: `List[str]`
        list of strings containing the variable names
        for the current day dataset
    3. empty_dataset: `xr.Dataset`
        empty dataset that contains the right structure

    # Returns
    1. variables_t1: `xr.Dataset`
        output dataset for previous day
    2. variables_t2: `xr.Dataset`
        output dataset for current day
    """

    # Create new dataset
    variables_t1 = empty_dataset.copy(deep=True)

    # Create empty DataArray for each variable
    for variable in calculation_variables_t1:

        # Assign new empty DataArray
        variables_t1[variable] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype='float32'))
        variables_t1[variable].attrs['name'] = variable  # set name in attributes

    # Create new dataset
    variables_t2 = empty_dataset.copy(deep=True)

    # Create empty DataArray for each variable
    for variable in calculation_variables_t2:

        # Assign new empty DataArray
        variables_t2[variable] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype='float32'))
        variables_t2[variable].attrs['name'] = variable  # set name in attributes

    return variables_t1, variables_t2


def prepare_outputs(empty_dataset: xr.Dataset, additional_outputs: List[str] = None) -> xr.Dataset:
    """
    Creates the `xarray Dataset` containing the outputs of the SAMIR model that will be saved.
    Additional variables can be saved by adding their names to the `additional_outputs` list.

    # Arguments
    1. empty_dataset: `xr.Dataset`
        empty dataset that contains the right structure
    2. additional_outputs: `List[str]`
        list of additional variable names to be saved

    # Returns
    1. model_outputs: `xr.Dataset`
        model outputs to be saved
    """

    # Evaporation and Transpiraion
    model_outputs = empty_dataset.copy(deep=True)

    model_outputs['E'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype='int16'))
    model_outputs['E'].attrs['units'] = 'mm'
    model_outputs['E'].attrs['standard_name'] = 'Evaporation'
    model_outputs['E'].attrs['description'] = 'Accumulated daily evaporation in milimeters'
    model_outputs['E'].attrs['scale factor'] = '1000'

    model_outputs['Tr'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d]
                           for d in list(empty_dataset.dims)), dtype='int16'))
    model_outputs['Tr'].attrs['units'] = 'mm'
    model_outputs['Tr'].attrs['standard_name'] = 'Transpiration'
    model_outputs['Tr'].attrs['description'] = 'Accumulated daily plant transpiration in milimeters'
    model_outputs['Tr'].attrs['scale factor'] = '1000'

    # Soil Water Content
    model_outputs['SWCe'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype='int16'))
    model_outputs['SWCe'].attrs['units'] = 'mm'
    model_outputs['SWCe'].attrs['standard_name'] = 'Soil Water Content of the evaporative zone'
    model_outputs['SWCe'].attrs['description'] = 'Soil water content of the evaporative zone in milimeters'
    model_outputs['SWCe'].attrs['scale factor'] = '1000'

    model_outputs['SWCr'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d]
                             for d in list(empty_dataset.dims)), dtype='int16'))
    model_outputs['SWCr'].attrs['units'] = 'mm'
    model_outputs['SWCr'].attrs['standard_name'] = 'Soil Water Content of the root zone'
    model_outputs['SWCr'].attrs['description'] = 'Soil water content of the root zone in milimeters'
    model_outputs['SWCr'].attrs['scale factor'] = '1000'

    # Irrigation
    model_outputs['Irr'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype='int16'))
    model_outputs['Irr'].attrs['units'] = 'mm'
    model_outputs['Irr'].attrs['standard_name'] = 'Irrigation'
    model_outputs['Irr'].attrs['description'] = 'Simulated daily irrigation in milimeters'
    model_outputs['Irr'].attrs['scale factor'] = '1000'

    # Deep Percolation
    model_outputs['DP'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype='int16'))
    model_outputs['DP'].attrs['units'] = 'mm'
    model_outputs['DP'].attrs['standard_name'] = 'Deep Percolation'
    model_outputs['DP'].attrs['description'] = 'Simulated daily Deep Percolation in milimeters'
    model_outputs['DP'].attrs['scale factor'] = '1000'

    if additional_outputs:
        for var in additional_outputs:
            model_outputs[var] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype='int16'))

    return model_outputs


def xr_maximum(ds: xr.DataArray, value: Union[xr.DataArray, float, int]) -> xr.DataArray:
    """
    Equivalent of `numpy.maximum(ds, value)` for xarray DataArrays

    # Arguments
    1. ds: `xr.DataArray`
        datarray to compare
    2. value: `Union[xr.DataArray, float, int]`
        value (scalar or dataarray) to compare

    # Returns
    1. output: `xr.DataArray`
        resulting dataarray with maximum value element-wise
    """
    return xr.where(ds <= value, value, ds, keep_attrs=True)


def xr_minimum(ds: xr.DataArray, value: Union[xr.DataArray, float, int]) -> xr.DataArray:
    """
    Equivalent of `numpy.minimum(ds, value)` for xarray DataArrays

    # Arguments
    1. ds: `xr.DataArray`
        datarray to compare
    2. value: `Union[xr.DataArray, float, int]`
        value (scalar or dataarray) to compare

    # Returns
    1. output: `xr.DataArray`
        resulting dataarray with minimum value element-wise
    """
    return xr.where(ds >= value, value, ds, keep_attrs=True)


def calculate_diff_re(TAW: np.ndarray, Dr: np.ndarray, Zr: np.ndarray, RUE: np.ndarray, De: np.ndarray, Wfc: np.ndarray, Ze: np.ndarray, DiffE: np.ndarray) -> np.ndarray:
    """
    Calculates the diffusion between the top soil layer and the root layer.

    # Arguments
    1. TAW: `np.ndarray`
        water capacity of root layer
    2. Dr: `np.ndarray`
        depletion of root layer
    3. Zr: `np.ndarray`
        height of root layer
    4. RUE: `np.ndarray`
        total available surface water = (Wfc-Wwp)*Ze
    5. De: `np.ndarray`
        depletion of the evaporative layer
    6. Wfc: `np.ndarray`
        field capacity
    7. Ze: `np.ndarray`
        height of evaporative layer (paramter)
    8. DiffE: `np.ndarray`
        diffusion coefficient between evaporative
        and root layers (unitless, parameter)

    # Returns
    1. diff_re: `np.ndarray`
        the diffusion between the top soil layer and
        the root layer
    """

    # Temporary variables to make calculation easier to read
    tmp1 = ((TAW - Dr) / Zr - (RUE - De) / Ze) / (Wfc * DiffE)
    tmp2 = ((TAW * Ze) - (RUE - De - Dr) * Zr) / (Zr + Ze) - Dr

    # Calculate diffusion according to SAMIR equation
    diff_re = np.where(tmp1 < 0, np.maximum(tmp1, tmp2), np.minimum(tmp1, tmp2))

    # Return zero values where the 'DiffE' parameter is equal to 0
    return np.where(DiffE == 0, 0, diff_re)


def calculate_diff_dr(TAW: np.ndarray, TDW: np.ndarray, Dr: np.ndarray, Zr: np.ndarray, Dd: np.ndarray, Wfc: np.ndarray, Zsoil: np.ndarray, DiffR: np.ndarray) -> np.ndarray:
    """
    Calculates the diffusion between the root layer and the deep layer.

    # Arguments
    1. TAW: `np.ndarray`
        water capacity of root layer
    2. TDW: `np.ndarray`
        water capacity of deep layer
    3. Dr: `np.ndarray`
        depletion of root layer
    4. Zr: `np.ndarray`
        height of root layer
    5. Dd: `np.ndarray`
        depletion of deep layer
    6. Wfc: `np.ndarray`
        field capacity
    7. Zsoil: `np.ndarray`
        total height of soil (paramter)
    8. DiffR: `np.ndarray`
        Diffusion coefficient between root
        and deep layers (unitless, parameter)

    # Returns
    1. Diff_dr: `np.ndarray`
        the diffusion between the root layer and the
        deep layer
    """

    # Temporary variables to make calculation easier to read
    tmp1 = (((TDW - Dd) / (Zsoil - Zr) - (TAW - Dr) / Zr) / Wfc) * DiffR
    tmp2 = (TDW * Zr - (TAW - Dr - Dd) * (Zsoil - Zr)) / Zsoil - Dd

    # Calculate diffusion according to SAMIR equation
    Diff_dr = np.where(tmp1 < 0, np.maximum(tmp1, tmp2), np.minimum(tmp1, tmp2))

    # Return zero values where the 'DiffR' parameter is equal to 0
    return np.where(DiffR == 0, 0, Diff_dr)


def calculate_W(TEW: np.ndarray, Dei: np.ndarray, Dep: np.ndarray, fewi: np.ndarray, fewp: np.ndarray) -> np.ndarray:
    """
    Calculate W, the weighting factor to split the energy available
    for evaporation depending on the difference in water availability
    in the two evaporation components, ensuring that the larger and
    the wetter, the more the evaporation occurs from that component

    # Arguments
    1. TEW: `np.ndarray`
        water capacity of evaporative layer
    2. Dei: `np.ndarray`
        depletion of the evaporative layer
        (irrigation part)
    3. Dep: `np.ndarray`
        depletion of the evaporative layer
        (precipitation part)
    4. fewi: `np.ndarray`
        soil fraction which is wetted by irrigation
        and exposed to evaporation
    5. fewp: `np.ndarray`
        soil fraction which is wetted by precipitation
        and exposed to evaporation

    # Returns
    1. W: `np.ndarray`
        weighting factor W
    """

    # Temporary variables to make calculation easier to read
    tmp = fewi * (TEW - Dei)

    # Calculate the weighting factor to split the energy available for evaporation
    W = 1 / (1 + (fewp * (TEW - Dep) / tmp))

    # Return W
    return np.where(tmp > 0, W, 0)


def calculate_Kr(TEW: np.ndarray, De: np.ndarray, REW: np.ndarray) -> np.ndarray:
    """
    calculates of the reduction coefficient for evaporation dependent
    on the amount of water in the soil using the FAO-56 method.

    # Arguments
    1. TEW: `np.ndarray`
        water capacity of evaporative layer
    2. De: `np.ndarray`
        depletion of evaporative layer
    3. REW: `np.ndarray`
        fixed readily evaporable water

    # Returns
    1. Kr: `np.ndarray`
        Kr coefficient
    """

    # Formula for calculating Kr
    Kr = (TEW - De) / (TEW - REW)

    # Return Kr
    return np.maximum(0, np.minimum(Kr, 1))


def update_Dr_from_root(TAW: np.ndarray, Zr: np.ndarray, TAW0: np.ndarray, TDW0: np.ndarray, Dr0: np.ndarray, Dd0: np.ndarray, Zr0: np.ndarray) -> np.ndarray:
    """
    Return the updated depletion for the root layer.

    # Arguments
    1. TAW: `np.ndarray`
        water capacity of root layer for current day
    2. Zr: `np.ndarray`
        root layer height for current day
    3. TAW0: `np.ndarray`
        water capacity of root layer for previous day
    4. TDW0: `np.ndarray`
        water capacity of deep layer for previous day
    5. Dr0: `np.ndarray`
        depletion of the root layer for previous day
    6. Dd0: `np.ndarray`
        depletion of the deep laye for previous day
    7. Zr0: `np.ndarray`
        root layer height for previous day

    # Returns
    1. output: `np.ndarray`
        updated depletion for the root layer
    """
    # Temporary variables to make calculation easier to read
    # tmp1 = np.minimum(Dr0 + Dd0 * (TAW - TAW0) / TDW0, TAW)
    # tmp2 = np.maximum(Dr0 + Dd0 * (TAW - TAW0) / TDW0, 0)

    # TODO vr: Updated version from xls
    tmp1 = np.maximum(Dr0 + Dd0 * (TAW - TAW0)/TDW0, 0)
    tmp2 = np.maximum(Dr0 + Dr0 * (TAW - TAW0)/TAW0, 0)

    # Return updated Dr
    return np.where(Zr > Zr0, tmp1, tmp2)


def update_Dd_from_root(TAW: np.ndarray, TDW: np.ndarray, Zr: np.ndarray, TAW0: np.ndarray, TDW0: np.ndarray, Dr0: np.ndarray, Dd0: np.ndarray, Zr0: np.ndarray) -> np.ndarray:
    """
    Return the updated depletion for the deep layer

    # Arguments
    1. TAW: `np.ndarray`
        water capacity of root layer for current day
    2. TDW: `np.ndarray`
        water capacity of deep layer for current day
    3. TAW0: `np.ndarray`
        water capacity of root layer for previous day
    5. TDW0: `np.ndarray`
        water capacity of deep layer for previous day
    6. Dd0: `np.ndarray`
        depletion of the deep laye for previous day
    7. Zr0: `np.ndarray`
        root layer height for previous day

    # Returns
    1. output: `np.ndarray`
        updated depletion for the deep layer
    """

    # Temporary variables to make calculation easier to read
    # tmp1 = np.maximum(Dd0 - Dd0 * (TAW - TAW0) / TDW0, 0)
    # tmp2 = np.minimum(Dd0 - Dd0 * (TAW - TAW0) / TDW0, TDW)

    # TODO vr: Updated version from xls
    tmp1 = np.maximum(Dd0 + Dd0 * (TAW - TAW0)/TDW0, 0)
    tmp2 = np.maximum(Dd0 + Dr0 * (TAW - TAW0)/TAW0, 0)

    # Return updated Dd
    return np.where(Zr > Zr0, tmp1, tmp2)


def format_duration(timedelta: float) -> None:
    """
    Print formatted timedelta in human readable format
    (days, hours, minutes, seconds, microseconds, milliseconds, nanoseconds).

    # Arguments
    timedelta: `float`
        time value in seconds to format

    # Returns
    `None`
    """

    if timedelta < 0.9e-6:
        print(round(timedelta*1e9, 1), 'ns')
    elif timedelta < 0.9e-3:
        print(round(timedelta*1e6, 1), 'µs')
    elif timedelta < 0.9:
        print(round(timedelta*1e3, 1), 'ms')
    elif timedelta < 60:
        print(round(timedelta, 1), 's')
    elif timedelta < 3.6e3:
        print(round(timedelta//60), 'm', round(timedelta % 60, 1),  's')
    elif timedelta < 24*3.6e3:
        print(round((timedelta/3.6e3)//1), 'h', round((timedelta/3.6e3) % 1*60//1), 'm', round((timedelta/3.6e3) % 1*60 % 1*60, 1), 's')
    elif timedelta < 48*3.6e3:
        print(round((timedelta/(24*3.6e3))//1), 'day,', round(((timedelta/(24*3.6e3)) % 1*24)//1), 'h,', round((timedelta/(24*3.6e3)) % 1*24 % 1*60//1), 'm,',  round((timedelta/(24*3.6e3)) % 1*24 % 1*60 % 1*60, 1), 's')
    else:
        print(round((timedelta/(24*3.6e3))//1), 'days,', round(((timedelta/(24*3.6e3)) % 1*24)//1), 'h,', round((timedelta/(24*3.6e3)) % 1*24 % 1*60//1), 'm,',  round((timedelta/(24*3.6e3)) % 1*24 % 1*60 % 1*60, 1), 's')

    return None


# @profile  # type: ignore
def run_samir(json_config_file: str, csv_param_file: str, ndvi_cube_path: str, Rain_path: str, ET0_path: str, soil_params_path: str, land_cover_path: str, chunk_size: dict, save_path: str, additional_outputs: List[str] = None, additional_outputs_scale: List[float] = None, max_ram_GB: int = 2) -> None:

    # Test inputs
    if len(additional_outputs) != len(additional_outputs_scale):
        print('\nadditional outpus name and scale list length do not match\n')
        return None

    # Turn off numpy warings
    np.seterr(divide='ignore', invalid='ignore')

    # ============ General parameters ============#
    # TODO vr: config_params jamais utilisé...?  ça viendra
    config_params = config(json_config_file)
    
    # calculation_variables_t2 contains the list of variables used for current day calculations
    calculation_variables_t2 = ['Diff_rei', 'Diff_rep', 'Diff_dr', 'Dd', 'De', 'Dei', 'Dep', 'DP', 'Dr', 'FCov', 'Irrig', 'Kcb', 'Kei', 'Kep', 'Kri', 'Krp', 'Ks', 'Kti', 'Ktp', 'RUE', 'SWCe', 'SWCr', 'TAW', 'TDW', 'TEW', 'Tei', 'Tep', 'W', 'Zr', 'fewi', 'fewp', 'p_cor']
    # calculation_variables_t1 contains the list of variables of the previous day necessary for current day calculations
    calculation_variables_t1 = ['Dr', 'Dd', 'TAW', 'TDW', 'Zr']

    # ============ Manage inputs ============#
    # NDVI (to have a correct empty dataset structure)
    ndvi_cube = xr.open_mfdataset(ndvi_cube_path, chunks=chunk_size, parallel=True)

    # SAMIR Parameters
    param_dataset, scale_factor = rasterize_samir_parameters(
        csv_param_file, ndvi_cube.drop_vars(['NDVI', 'time']), land_cover_path, chunk_size=chunk_size)

    # SAMIR Variables
    variables_t1, variables_t2 = setup_time_loop(
        calculation_variables_t1, calculation_variables_t2, ndvi_cube.drop_vars(['NDVI', 'time']))

    # ============ Prepare outputs ============#
    model_outputs = prepare_outputs(ndvi_cube.drop_vars(['NDVI']), additional_outputs=additional_outputs)

    # Create encoding dictionnary
    for variable in list(model_outputs.keys()):
        # Write encoding dict
        encoding_dict = {}
        encod = {}
        encod['dtype'] = 'i2'
        encoding_dict[variable] = encod

    # Save empty output
    model_outputs.to_netcdf(save_path, encoding=encoding_dict)
    model_outputs.close()

    # ============ Prepare time iterations ============#
    dates = ndvi_cube.time.values
    ndvi_cube.close()

    # ============ Create aliases for better readability ============#

    # Variables for current day
    # var = da.from_array(dataarray, chunks = (5, 5))
    Diff_rei = variables_t2.Diff_rei.to_numpy()
    Diff_rep = variables_t2.Diff_rep.to_numpy()
    Diff_dr = variables_t2.Diff_dr.to_numpy()
    Dd = variables_t2.Dd.to_numpy()
    De = variables_t2.De.to_numpy()
    Dei = variables_t2.Dei.to_numpy()
    Dep = variables_t2.Dep.to_numpy()
    DP = variables_t2.DP.to_numpy()
    Dr = variables_t2.Dr.to_numpy()
    FCov = variables_t2.FCov.to_numpy()
    Irrig = variables_t2.Irrig.to_numpy()
    Kcb = variables_t2.Kcb.to_numpy()
    Kei = variables_t2.Kei.to_numpy()
    Kep = variables_t2.Kep.to_numpy()
    Kri = variables_t2.Kri.to_numpy()
    Krp = variables_t2.Krp.to_numpy()
    Ks = variables_t2.Ks.to_numpy()
    Kti = variables_t2.Kti.to_numpy()
    Ktp = variables_t2.Ktp.to_numpy()
    RUE = variables_t2.RUE.to_numpy()
    SWCe = variables_t2.SWCe.to_numpy()
    SWCr = variables_t2.SWCr.to_numpy()
    TAW = variables_t2.TAW.to_numpy()
    TDW = variables_t2.TDW.to_numpy()
    TEW = variables_t2.TEW.to_numpy()
    Tei = variables_t2.Tei.to_numpy()
    Tep = variables_t2.Tep.to_numpy()
    Zr = variables_t2.Zr.to_numpy()
    W = variables_t2.W.to_numpy()
    fewi = variables_t2.fewi.to_numpy()
    fewp = variables_t2.fewp.to_numpy()
    p_cor = variables_t2.p_cor.to_numpy()

    # Variables for previous day
    TAW0 = variables_t1.TAW.to_numpy()
    TDW0 = variables_t1.TDW.to_numpy()
    Dr0 = variables_t1.Dr.to_numpy()
    Dd0 = variables_t1.Dd.to_numpy()
    Zr0 = variables_t1.Zr.to_numpy()

    # Parameters
    # Parameters have an underscore (_) behind their name for recognition
    DiffE_ = param_dataset.DiffE.to_numpy()
    DiffR_ = param_dataset.DiffR.to_numpy()
    FW_ = param_dataset.FW.to_numpy()
    Fc_stop_ = param_dataset.Fc_stop.to_numpy()
    FmaxFC_ = param_dataset.FmaxFC.to_numpy()
    Foffset_ = param_dataset.Foffset.to_numpy()
    Fslope_ = param_dataset.Fslope.to_numpy()
    Init_RU_ = param_dataset.Init_RU.to_numpy()
    Irrig_auto_ = param_dataset.Irrig_auto.to_numpy()
    # Kcmax_ = param_dataset.Kcmax.to_numpy() # Inutile?
    KmaxKcb_ = param_dataset.KmaxKcb.to_numpy()
    Koffset_ = param_dataset.Koffset.to_numpy()
    Kslope_ = param_dataset.Kslope.to_numpy()
    Lame_max_ = param_dataset.Lame_max.to_numpy()
    REW_ = param_dataset.REW.to_numpy()
    Ze_ = param_dataset.Ze.to_numpy()
    Zsoil_ = param_dataset.Zsoil.to_numpy()
    maxZr_ = param_dataset.maxZr.to_numpy()
    minZr_ = param_dataset.minZr.to_numpy()
    p_ = param_dataset.p.to_numpy()

    # scale factors
    # Scale factors have the following name scheme : s_ + parameter_name
    s_FW = scale_factor['FW']
    s_Fc_stop = scale_factor['Fc_stop']
    s_FmaxFC = scale_factor['FmaxFC']
    s_Foffset = scale_factor['Foffset']
    s_Fslope = scale_factor['Fslope']
    s_Init_RU = scale_factor['Init_RU']
    # s_Irrig_auto = scale_factor['Irrig_auto']
    # s_Kcmax = scale_factor['Kcmax']  # TODO vr: variable jamais utilisée?
    s_KmaxKcb = scale_factor['KmaxKcb']
    s_Koffset = scale_factor['Koffset']
    s_Kslope = scale_factor['Kslope']
    s_Lame_max = scale_factor['Lame_max']
    s_REW = scale_factor['REW']
    s_Ze = scale_factor['Ze']
    s_DiffE = scale_factor['DiffE']
    s_DiffR = scale_factor['DiffR']
    s_Zsoil = scale_factor['Zsoil']
    s_maxZr = scale_factor['maxZr']
    s_minZr = scale_factor['minZr']
    s_p = scale_factor['p']

    # input data
    with nc.Dataset(ndvi_cube_path, mode='r') as ds:
        # Dimensions of ndvi dataset : (time, x, y)
        NDVI = ds.variables['NDVI'][0, :, :] / 255
    with nc.Dataset(soil_params_path, mode='r') as ds:
        Wfc = ds.variables['Wfc'][:, :]
        Wwp = ds.variables['Wwp'][:, :]
        print('soil Wfc and Wwp:', Wfc[0, 0], Wwp[0, 0])
    with rio.open(Rain_path, mode='r') as ds:
        Rain = ds.read(1) / 1000
    with rio.open(ET0_path, mode='r') as ds:
        ET0 = ds.read(1) / 1000

    # Create progress bar
    progress_bar = tqdm(total=len(dates), desc='Running model', unit=' days')

    # ============ First day initialization ============#
    # Fraction cover
    FCov = s_Fslope * Fslope_ * NDVI + s_Foffset * Foffset_
    FCov = np.minimum(np.maximum(FCov, 0), s_Fc_stop * Fc_stop_)

    # Root depth upate
    # TODO vr: il semblerait que dans le xls, Fcmax soit le Fc max de la simulation, alors que dans python, c'est un parametre
    Zr = np.maximum(s_minZr * minZr_ + (FCov / (s_FmaxFC * FmaxFC_)) * (s_maxZr * maxZr_ - s_minZr * minZr_), s_Ze * Ze_ + 0.001)

    # Water capacities
    TEW = (Wfc - Wwp/2) * s_Ze * Ze_
    RUE = (Wfc - Wwp) * s_Ze * Ze_
    TAW = (Wfc - Wwp) * Zr
    TDW = (Wfc - Wwp) * (s_Zsoil * Zsoil_ - Zr)  # Zd = Zsoil - Zr

    # Depletions
    Dei = RUE * (1 - s_Init_RU * Init_RU_)
    Dep = RUE * (1 - s_Init_RU * Init_RU_)
    Dr = TAW * (1 - s_Init_RU * Init_RU_)
    Dd = TDW * (1 - s_Init_RU * Init_RU_)

    # p_cor
    p_cor = s_p * p_

    # Irrigation  TODO : find correct method for irrigation
    Irrig = np.minimum(np.maximum(Dr - Rain, 0), s_Lame_max * Lame_max_) * Irrig_auto_
    Irrig = np.where(Dr > TAW * p_cor, Irrig, 0)

    # Kcb
    Kcb = np.minimum(np.maximum(s_Kslope * Kslope_ * NDVI + s_Koffset * Koffset_, 0), s_KmaxKcb * KmaxKcb_)

    # Update depletions with Rainfall and/or irrigation

    # DP
    DP = - np.minimum(Dd + np.minimum(Dr - Rain - Irrig, 0), 0)

    # De
    Dei = np.minimum(np.maximum(Dei - Rain - Irrig / (s_FW * FW_), 0), TEW)
    Dep = np.minimum(np.maximum(Dep - Rain, 0), TEW)

    fewi = np.minimum(1 - FCov, (s_FW * FW_))
    fewp = 1 - FCov - fewi

    De = np.nansum((Dei * fewi, Dep * fewp)) / np.nansum([fewi, fewp])
    De = np.where(np.isfinite(De), De, Dei * (s_FW * FW_) + Dep * (1 - (s_FW * FW_)))

    # Dr
    Dr = np.minimum(np.maximum(Dr - Rain - Irrig, 0), TAW)

    # Dd
    Dd = np.minimum(np.maximum(Dd + np.minimum(Dr - Rain - Irrig, 0), 0), TDW)

    # Diffusion coefficients
    Diff_rei = calculate_diff_re(TAW, Dr, Zr, RUE, Dei, Wfc, s_Ze*Ze_, s_DiffE*DiffE_)
    Diff_rep = calculate_diff_re(TAW, Dr, Zr, RUE, Dep, Wfc, s_Ze*Ze_, s_DiffE*DiffE_)
    Diff_dr = calculate_diff_dr(TAW, TDW, Dr, Zr, Dd, Wfc, s_Zsoil*Zsoil_, s_DiffR*DiffR_)

    # Weighing factor W
    W = calculate_W(TEW, Dei, Dep, fewi, fewp)

    # Soil water content of evaporative layer
    SWCe = 1 - De/TEW
    # Soil water content of root layer
    SWCr = 1 - Dr/TAW

    # Water Stress coefficient
    Ks = np.minimum((TAW - Dr) / (TAW * (1 - p_cor)), 1)

    # Reduction coefficient for evaporation
    Kri = calculate_Kr(TEW, Dei, s_REW*REW_)
    Krp = calculate_Kr(TEW, Dep, s_REW*REW_)
    Kei = np.minimum(W * Kri * (s_KmaxKcb*KmaxKcb_ - Kcb), fewi * s_KmaxKcb*KmaxKcb_)
    Kep = np.minimum((1 - W) * Krp * (s_KmaxKcb*KmaxKcb_ - Kcb), fewp * s_KmaxKcb*KmaxKcb_)

    # Prepare coefficients for evapotranspiration
    Kti = np.minimum(((s_Ze*Ze_ / Zr)**0.6) * (1 - Dei / TEW) / np.maximum(1 - Dr / TAW, 0.001), 1)
    Ktp = np.minimum(((s_Ze*Ze_ / Zr)**0.6) * (1 - Dep / TEW) / np.maximum(1 - Dr / TAW, 0.001), 1)
    Tei = Kti * Ks * Kcb * ET0
    Tep = Ktp * Ks * Kcb * ET0

    # Update depletions
    Dei = np.where(fewi > 0, np.minimum(np.maximum(Dei + ET0 * Kei / fewi + Tei - Diff_rei, 0), TEW),
                   np.minimum(np.maximum(Dei + Tei - Diff_rei, 0), TEW))
    Dep = np.where(fewp > 0, np.minimum(np.maximum(Dep + ET0 * Kep / fewp + Tep - Diff_rep, 0), TEW),
                   np.minimum(np.maximum(Dep + Tep - Diff_rep, 0), TEW))

    De = np.nansum((Dei * fewi, Dep * fewp)) / np.nansum([fewi, fewp])
    De = np.where(np.isfinite(De), De, Dei * (s_FW * FW_) + Dep * (1 - (s_FW * FW_)))

    # Evaporation
    E = np.maximum((Kei + Kep) * ET0, 0)

    # Transpiration
    Tr = Kcb * Ks * ET0

    # Potential evapotranspiration and evaporative fraction ??

    # Update depletions (root and deep zones) at the end of the day
    Dr = np.minimum(np.maximum(Dr + E + Tr - Diff_dr, 0), TAW)
    Dd = np.minimum(np.maximum(Dd + Diff_dr, 0), TDW)

    # Write outputs
    with nc.Dataset(save_path, mode='r+') as outputs:
        # Dimensions of output dataset : (x, y, time)
        # Deep percolation
        outputs.variables['DP'][:, :, 0] = np.round(DP * 1000)
        # Soil water content of the evaporative layer
        outputs.variables['SWCe'][:, :, 0] = np.round(SWCe * 1000)
        # Soil water content of the root layer
        outputs.variables['SWCr'][:, :, 0] = np.round(SWCr * 1000)
        # Evaporation
        outputs.variables['E'][:, :, 0] = np.round(E * 1000)
        # Transpiration
        outputs.variables['Tr'][:, :, 0] = np.round(Tr * 1000)
        # Irrigation
        outputs.variables['Irr'][:, :, 0] = np.round(Irrig * 1000)
        # Additionnal outputs
        for var, scale in zip(additional_outputs, additional_outputs_scale):
            outputs.variables[var][:, :, 0] = np.round(eval(var) * scale)

    # Update previous day values
    TAW0 = TAW
    TDW0 = TDW
    Dr0 = Dr
    Dd0 = Dd
    Zr0 = Zr

    # Update progress bar
    progress_bar.update()

    # %% ============ Time loop ============#
    for i in range(1, len(dates)):

        # Reset input aliases
        # input data
        with nc.Dataset(ndvi_cube_path, mode='r') as ds:
            # Dimensions of ndvi dataset : (time, x, y)
            NDVI = ds.variables['NDVI'][i, :, :] / 255
        with rio.open(Rain_path, mode='r') as ds:
            Rain = ds.read(i+1) / 1000
        with rio.open(ET0_path, mode='r') as ds:
            ET0 = ds.read(i+1) / 1000
            # ET0_previous = ds.read(i) / 1000  # TODO vr: variable jamais utilisée!

        # Update variables
        # Fraction cover
        FCov = s_Fslope * Fslope_ * NDVI + s_Foffset * Foffset_
        FCov = np.minimum(np.maximum(FCov, 0), s_Fc_stop * Fc_stop_)

        # Root depth upate
        Zr = np.maximum(s_minZr * minZr_ + (FCov / (s_FmaxFC * FmaxFC_))
                        * s_maxZr * (maxZr_ - minZr_), s_Ze * Ze_ + 0.001)

        # Water capacities
        TAW = (Wfc - Wwp) * Zr
        TDW = (Wfc - Wwp) * (s_Zsoil * Zsoil_ - Zr)  # Zd = Zsoil - Zr

        # Update depletions from root increase
        Dr = update_Dr_from_root(TAW, Zr, TAW0, TDW0, Dr0, Dd0, Zr0)
        Dd = update_Dd_from_root(TAW, TDW, Zr, TAW0, TDW0, Dr0, Dd0, Zr0)

        # Update param p
        # TODO: Calcul p_cor différent entre la doc ou l'excel et le code samir parcelle
        # p_cor = np.minimum(np.maximum(s_p * p_ + 0.04 * (5 - (E + Tr)), 0.1), 0.8)
        # Calcul p_cor différent entre la doc ou l'excel et le code samir parcelle
        p_cor = s_p * p_ + 0.04 * (5 - (E + Tr))

        # Irrigation   TODO : find correct method for irrigation
        Irrig = np.minimum(np.maximum(Dr - Rain, 0), s_Lame_max * Lame_max_) * Irrig_auto_
        Irrig = np.where(Dr > TAW * p_cor, Irrig, 0)

        # Kcb
        Kcb = np.minimum(np.maximum(s_Kslope * Kslope_ * NDVI + s_Koffset * Koffset_, 0), s_KmaxKcb * KmaxKcb_)

        # DP (Deep percolation)
        DP = - np.minimum(Dd + np.minimum(Dr - Rain - Irrig, 0), 0)

        # Update depletions with Rainfall and/or irrigation

        # De
        Dei = np.minimum(np.maximum(Dei - Rain - Irrig / (s_FW * FW_), 0), TEW)
        Dep = np.minimum(np.maximum(Dep - Rain, 0), TEW)

        fewi = np.minimum(1 - FCov, (s_FW * FW_))
        fewp = 1 - FCov - fewi

        De = np.nansum((Dei * fewi, Dep * fewp)) / np.nansum([fewi, fewp])
        De = np.where(np.isfinite(De), De, Dei * (s_FW * FW_) + Dep * (1 - (s_FW * FW_)))

        # Update depletions from rain and irrigation
        Dr = np.minimum(np.maximum(Dr - Rain - Irrig, 0), TAW)
        Dd = np.minimum(np.maximum(Dd + np.minimum(Dr - Rain - Irrig, 0), 0), TDW)

        # Diffusion coefficients
        Diff_rei = calculate_diff_re(TAW, Dr, Zr, RUE, Dei, Wfc, s_Ze*Ze_, s_DiffE*DiffE_)
        Diff_rep = calculate_diff_re(TAW, Dr, Zr, RUE, Dep, Wfc, s_Ze*Ze_, s_DiffE*DiffE_)
        Diff_dr = calculate_diff_dr(TAW, TDW, Dr, Zr, Dd, Wfc, s_Zsoil*Zsoil_, s_DiffR*DiffR_)

        # Weighing factor W
        W = calculate_W(TEW, Dei, Dep, fewi, fewp)

        # Soil water content of evaporative layer
        SWCe = 1 - De/TEW
        # Soil water content of root layer
        SWCr = 1 - Dr/TAW

        # Water Stress coefficient
        Ks = np.minimum((TAW - Dr) / (TAW * (1 - p_cor)), 1)

        # Reduction coefficient for evaporation
        Kri = calculate_Kr(TEW, Dei, REW_*s_REW)
        Krp = calculate_Kr(TEW, Dep, REW_*s_REW)
        Kei = np.minimum(W * Kri * (s_KmaxKcb * KmaxKcb_ - Kcb), fewi * s_KmaxKcb * KmaxKcb_)
        Kep = np.minimum((1 - W) * Krp * (s_KmaxKcb * KmaxKcb_ - Kcb), fewp * s_KmaxKcb * KmaxKcb_)

        # Prepare coefficients for evapotranspiration
        Kti = np.minimum(((s_Ze * Ze_ / Zr)**0.6) * (1 - Dei / TEW) / np.maximum(1 - Dr / TAW, 0.001), 1)
        Ktp = np.minimum(((s_Ze * Ze_ / Zr)**0.6) * (1 - Dep / TEW) / np.maximum(1 - Dr / TAW, 0.001), 1)
        Tei = Kti * Ks * Kcb * ET0
        Tep = Ktp * Ks * Kcb * ET0

        # Update depletions
        Dei = np.where(fewi > 0, np.minimum(np.maximum(Dei + ET0 * Kei / fewi + Tei - Diff_rei, 0), TEW),
                       np.minimum(np.maximum(Dei + Tei - Diff_rei, 0), TEW))
        Dep = np.where(fewp > 0, np.minimum(np.maximum(Dep + ET0 * Kep / fewp + Tep - Diff_rep, 0), TEW),
                       np.minimum(np.maximum(Dep + Tep - Diff_rep, 0), TEW))

        De = np.nansum((Dei * fewi, Dep * fewp)) / np.nansum([fewi, fewp])
        De = np.where(np.isfinite(De), De, Dei * (s_FW * FW_) + Dep * (1 - (s_FW * FW_)))

        # Evaporation
        E = np.maximum((Kei + Kep) * ET0, 0)

        # Transpiration
        Tr = Kcb * Ks * ET0

        # Potential evapotranspiration and evaporative fraction ??

        # Update depletions (root and deep zones) at the end of the day
        Dr = np.minimum(np.maximum(Dr + E + Tr - Diff_dr, 0), TAW)
        Dd = np.minimum(np.maximum(Dd + Diff_dr, 0), TDW)

        # Write outputs
        with nc.Dataset(save_path, mode='r+') as outputs:
            # Dimensions of output dataset : (x, y, time)
            # Deep percolation
            outputs.variables['DP'][:, :, i] = np.round(DP * 1000)
            # Soil water content of the evaporative layer
            outputs.variables['SWCe'][:, :, i] = np.round(SWCe * 1000)
            # Soil water content of the root layer
            outputs.variables['SWCr'][:, :, i] = np.round(SWCr * 1000)
            # Evaporation
            outputs.variables['E'][:, :, i] = np.round(E * 1000)
            # Transpiration
            outputs.variables['Tr'][:, :, i] = np.round(Tr * 1000)
            # Irrigation
            outputs.variables['Irr'][:, :, i] = np.round(Irrig * 1000)
            # Additionnal outputs
            for var, scale in zip(additional_outputs, additional_outputs_scale):
                outputs.variables[var][:, :, i] = np.round(eval(var) * scale)

        # Update previous day values
        TAW0 = TAW
        TDW0 = TDW
        Dr0 = Dr
        Dd0 = Dd
        Zr0 = Zr

        # Update progress bar
        progress_bar.update()

    # Close progress bar
    progress_bar.close()

    return None


# %% MAIN
if __name__ == '__main__':

    data_path = '/mnt/e/DATA/DEV_inputs_test'
    # data_path = './DEV_inputs_test'

    size = 10

    NDVI_path = data_path + os.sep + 'NDVI_' + str(size) + '.nc'
    Rain_path = data_path + os.sep + 'Rain_' + str(size) + '.tif'
    ET0_path = data_path + os.sep + 'ET0_' + str(size) + '.tif'
    land_cover_path = data_path + os.sep + 'land_cover_' + str(size) + '.nc'
    # json_config_file = '/home/auclairj/GIT/modspa-pixel/config/config_modspa.json'
    json_config_file = './config/config_modspa.json'
    # param_file = '/home/auclairj/GIT/modspa-pixel/parameters/csv_files/params_samir_test.csv'
    param_file = './parameters/csv_files/params_samir_test.csv'
    soil_path = data_path + os.sep + 'soil_' + str(size) + '.nc'
    save_path = data_path + os.sep + 'outputs_' + str(size) + '.nc'
    if os.path.exists(save_path):
        os.remove(save_path)

    # Validation sets
    xls_NDVI_path = data_path + os.sep + 'xls_NDVI_10.nc'
    xls_Rain_path = data_path + os.sep + 'xls_Rain_10.tif'
    val_ET0_path = data_path + os.sep + 'xls_ET0_10.tif'

    output_save_path = data_path + os.sep + 'pix_outputs_10.nc'

    additional_outputs = ['Zr', 'Dei', 'Dep', 'Dr', 'Dd', 'Kei', 'Kep', 'Ks', 'W', 'Kcb', 'Kri', 'Krp', 'NDVI', 'fewi', 'fewp', 'TDW', 'TAW', 'FCov', 'Tei', 'Tep', 'Diff_rei', 'Diff_rep', 'Diff_dr', 'Rain', 'p_cor']
    additional_outputs_scale = [100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100]

    chunk_size = {'x': 250, 'y': 250, 'time': -1}

    t = time()

    # client = Client()
    # webbrowser.open('http://127.0.0.1:8787/status', new=2, autoraise=True)

    # run_samir(json_config_file, param_file, ndvi_path, Rain_path, ET0_path, soil_path, land_cover_path, chunk_size, save_path)
    run_samir(json_config_file, param_file, xls_NDVI_path, xls_Rain_path, val_ET0_path, soil_path, land_cover_path,
              chunk_size, output_save_path, additional_outputs=additional_outputs, additional_outputs_scale=additional_outputs_scale)

    print('')
    print('Writting Output:', output_save_path)
    format_duration(time() - t)

    # client.close()