# -*- coding: UTF-8 -*-
# Python
"""
29-08-2023
@author: jeremy auclair

Generate convert pandas dataframes to rasters for the pixel mode.
"""

from typing import List  # to declare variables
import numpy as np  # for math on arrays
import xarray as xr  # to manage nc files
import pandas as pd  # to manage dataframes


def convert_dataframe_to_xarray(dataframe_path: str, save_path: str, variables: List[str], data_types: List[str]) -> None:
    
    
    # Open dataframe
    dataframe = pd.read_csv(dataframe_path).sort_values(by = ['date', 'id'])
    dataframe['date'] = pd.to_datetime(dataframe['date'])

    # Create dimensions for xarray dataset
    x = [i+1 for i in range(len(set(dataframe['id'])))]
    y = [1]
    dates = np.unique(dataframe['date'].values)

    # Get dimension sizes
    time_size = len(dates)
    x_size = len(x)
    y_size = len(y)

    # Create dimension tuples
    dim_size = (time_size, x_size, y_size)
    dims = ('time', 'x', 'y')

    # Reshape variables in correct format and put them in a dictionnary
    data_variables = {}
    for var in variables:
        data_variables[var] = (dims, np.reshape(dataframe[var].values, dim_size))
    
    # Create xarray dataset
    xarray_dataset = xr.Dataset(data_vars = data_variables, coords = {'time': dates, 'x': x, 'y': y})
    
    # Create encoding dictionnary
    encoding_dict = {}
    for var, dtype in zip(variables, data_types):
        # Write encoding dict
        encod = {}
        encod['dtype'] = dtype
        encod['chunksizes'] = (time_size, x_size, y_size)
        encoding_dict[var] = encod
    
    # Save dataset as netCDF4 file
    xarray_dataset.to_netcdf(save_path, encoding = encoding_dict)
    
    return None