# -*- coding: UTF-8 -*-
# Python
"""
Created on Wed 2023.

@author: auclairj
Adapated from notebook: rivallandv
"""
import os
import xarray as xr
import matplotlib.pyplot as plt
from matplotlib import rcParams
import matplotlib.dates as mdates

# Settings for plots
plt.style.use('default')
rcParams['mathtext.fontset'] = 'stix'
rcParams['font.family'] = 'STIXGeneral'
rcParams.update({'font.size': 18})
# Date format
date_plot_format = mdates.WeekdayLocator(interval=6)
date_format = mdates.DateFormatter('%Y-%m')


def plot_time_series(ds1: xr.Dataset, var: str, ds2: xr.Dataset = None, label1: str = 'Dataset1', label2: str = 'Dataset2',
                     scale_factor1: float = 1000, scale_factor2: float = 1000, unit: str = 'mm', title: str = 'variable comparison') -> None:
    """
    Plot times series of a uniform modspa dataset.
    Select first pixel (upper left corner) and plot
    its value over time.

    ## Arguments
    1. ds1: `xr.Dataset`
        first dataset to plot
    2. var: `str`
        name of variable to plot
    3. ds2: `xr.Dataset` `default = None`
        second dataset to plot, optional
    4. label1: `str` `default = 'Dataset1'`
        label for first dataset
    5. label2: `str` `default = 'Dataset2'`
        label for second dataset, optional
    6. scale_factor1: `float` `default = 1000`
        scale factor for first dataset to
        divide the time series in order to
        plot the correct variable values
    7. scale_factor2: `float` `default = 1000`
        scale factor for second dataset to
        divide the time series in order to
        plot the correct variable values
    8. unit: `str` `default = 'mm'`
        unit of value
    9. title: `str` `default = 'variable comparison'`
        title of plot

    ## Returns
    `None`
    """
    plt.figure(figsize=(14, 7))
    plt.grid(color='silver', linestyle='--', axis='both', linewidth=1)
    plt.gca().xaxis.set_major_formatter(date_format)
    plt.gca().xaxis.set_major_locator(date_plot_format)
    (ds1.isel(x=0, y=0)[var] / scale_factor1).plot(label=label1, color='b', marker='o', markersize=2, alpha=0.8)
    if ds2:
        (ds2.isel(x=0, y=0)[var] / scale_factor2).plot(label=label2, color='r', marker='o', markersize=2, alpha=0.8)
    plt.title(var + ' ' + title)
    plt.ylabel(var + ' [' + unit + ']')
    plt.legend()

    return None


# %% Path and Size

# IN/OUT working path:
data_path = './DEV_inputs_test'

# Test size
size = 10

# Reference meteo forcing
Rain_path = data_path + os.sep + 'xls_Rain_' + str(size) + '.tif'
ET0_path = data_path + os.sep + 'xls_ET0_' + str(size) + '.tif'

# Modspa-pixel output
pix_outputs_path = data_path + os.sep + 'pix_outputs_' + str(size) + '.nc'

# Excel output (validation or reeference)
# File generated by: 1_format_ref_samir_xls.py
xls_outputs_path = data_path + os.sep + 'xls_outputs_' + str(size) + '.nc'

# %% Open/Read REFERENCE and MODSPA datasets
pix_outputs = xr.open_dataset(pix_outputs_path)  # modspa pixel
xls_outputs = xr.open_dataset(xls_outputs_path)  # reference xls

# %% Forçage meteo: inutile car dans les outputs
# Rain = xr.open_dataset(Rain_path).rename({'band_data': 'Rain', 'band': 'time'})
# Rain['time'] = pix_outputs.time.values
# ET0 = xr.open_dataset(ET0_path).rename({'band_data': 'ET0', 'band': 'time'})
# ET0['time'] = pix_outputs.time.values


# %% Compute differences
pix_variables = list(pix_outputs.keys())  # list of output var from samir-pixel
xls_variables = list(xls_outputs.keys())  # list of output var from samir-xls

# recherche de l'intersection => variables communes de part et d'autre (xls & pix)
intersect_variables = set(pix_variables).intersection(set(xls_variables))

diff = pix_outputs.drop_vars(intersect_variables).copy(deep=True)
for var in intersect_variables:
    # NOTE vr: Pourquoi pix*10 et pas xls?
    if var in ['E', 'Tr', 'SWCe', 'SWCr', 'DP']:
        diff[var] = (pix_outputs[var].copy(deep=True)*10 - xls_outputs[var].copy(deep=True)) / 1000
    else:
        diff[var] = (pix_outputs[var].copy(deep=True) - xls_outputs[var].copy(deep=True)) / 100

        # Get values
differences_sum = {}
differences_mean = {}
for var in intersect_variables:
    differences_sum[var] = round(float(diff[var].sum(dim='time').mean().values), 3)
    differences_mean[var] = round(float(diff[var].mean(dim='time').mean().values), 3)

print('Difference on sum :', differences_sum)
print('Difference on mean :', differences_mean)


plt.figure()
plt.bar(range(len(differences_mean)), list(differences_mean.values()), align='center')
plt.xticks(range(len(differences_mean)), list(differences_mean.keys()), rotation='vertical')
plt.rc('xtick', labelsize=4)
plt.title('Difference on mean')
plt.show()

# %% Figures Specifiques
# Dataset pixel, Tr, E, SWCe, SWCr, Irr, DP : scale factor: 1000

plot_time_series(ds1=xls_outputs, ds2=pix_outputs, var='E', label1='Excel values',
                 label2='Pixel values', scale_factor1=100, scale_factor2=1000, unit='mm')

plot_time_series(ds1=xls_outputs, ds2=pix_outputs, var='Tr', label1='Excel values',
                 label2='Pixel values', scale_factor1=100, scale_factor2=1000, unit='mm')

plot_time_series(ds1=xls_outputs, ds2=pix_outputs, var='SWCe', label1='Excel values',
                 label2='Pixel values', scale_factor1=100, scale_factor2=1000, unit='mm')

plot_time_series(ds1=xls_outputs, ds2=pix_outputs, var='SWCr', label1='Excel values',
                 label2='Pixel values', scale_factor1=100, scale_factor2=1000, unit='mm')

plot_time_series(ds1=xls_outputs, ds2=pix_outputs, var='DP', label1='Excel values',
                 label2='Pixel values', scale_factor1=100, scale_factor2=1000, unit='mm')

# %% All Figures

plot_variables = intersect_variables
plot_variables.remove('E')
plot_variables.remove('Tr')
plot_variables.remove('SWCe')
plot_variables.remove('SWCr')
plot_variables.remove('DP')

for var in plot_variables:
    plot_time_series(ds1=xls_outputs, ds2=pix_outputs, var=var, label1='Excel values',
                     label2='Pixel values', scale_factor1=100, scale_factor2=100, unit='??')