"""Various utilities to encode MARIS dataset as `NetCDF`, `csv`, ... formats."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/api/encoders.ipynb.

# %% auto 0
__all__ = ['NetCDFEncoder']

# %% ../nbs/api/encoders.ipynb 3
import netCDF4
from netCDF4 import Dataset
import pandas as pd
from typing import Dict, Callable
import numpy as np
from fastcore.basics import patch, store_attr
import fastcore.all as fc
#import os

from marisco.configs import (
    NC_DTYPES, 
    NC_VARS, 
    NC_DIM,
    NC_GROUPS,
    lut_path, 
    Enums,
    nc_tpl_path
)

# %% ../nbs/api/encoders.ipynb 5
class NetCDFEncoder:
    "MARIS NetCDF encoder."
    def __init__(self, 
                 dfs: Dict[str, pd.DataFrame], # dict of Dataframes to encode with group name as key {'sediment': df_sed, ...}
                 dest_fname: str, # Name of output file to produce
                 global_attrs: Dict[str, str], # Global attributes
                 fn_src_fname: Callable=nc_tpl_path, # Function returning file name and path to the MARIS CDL template
                 custom_maps: Dict[str, Dict[str, int]]= None,# Custom maps to encode
                 verbose: bool=False, # Print currently written NetCDF group and variable names
                 ):
        store_attr()
        self.src_fname = fn_src_fname()
        self.enum_dtypes = {}
        self.nc_to_cols = {v:k for k,v in NC_VARS.items()}


# %% ../nbs/api/encoders.ipynb 7
@patch 
def copy_global_attributes(self:NetCDFEncoder):
    "Update NetCDF template global attributes as specified by `global_attrs` argument."
    self.dest.setncatts(self.src.__dict__)
    for k, v in self.global_attrs.items(): self.dest.setncattr(k, v)

# %% ../nbs/api/encoders.ipynb 8
@patch
def copy_dimensions(self:NetCDFEncoder, grp_dest):
    "Copy dimensions to root and all groups from template."
    src_dim = self.src.groups[grp_dest.name].dimensions
    for name, dim in src_dim.items():
        grp_dest.createDimension(name, (len(dim) if not dim.isunlimited() else None))

# %% ../nbs/api/encoders.ipynb 9
@patch
def process_groups(self:NetCDFEncoder):
    for grp_name, df in self.dfs.items():
        self.process_group(NC_GROUPS[grp_name], df)

# %% ../nbs/api/encoders.ipynb 10
@patch
def process_group(self:NetCDFEncoder, grp_name, df):
    grp_dest = self.dest.createGroup(grp_name)
    self.copy_dimensions(grp_dest)
    self.copy_variables(grp_name, df, grp_dest)

# %% ../nbs/api/encoders.ipynb 11
@patch
def copy_variables(self:NetCDFEncoder, grp_name, df, grp_dest):
    cols = [NC_VARS[col] for col in df.columns if col in NC_VARS]
    for var_name, var_src in self.src.groups[grp_name].variables.items():
        if var_name in cols: 
            self.copy_variable(var_name, var_src, df, grp_dest)
        if self.custom_maps:
            self.copy_custom_map(var_name, grp_dest)

# %% ../nbs/api/encoders.ipynb 12
@patch
def copy_variable(self:NetCDFEncoder, var_name, var_src, df, grp_dest):
    dtype_name = var_src.datatype.name
    enums_src = self.src.enumtypes
    if self.verbose: 
        print(80*'-')
        print(f'Group: {grp_dest.name}, Variable: {var_name}')
    self._create_and_copy_variable(var_name, var_src, df, grp_dest, dtype_name)
    self.copy_variable_attributes(var_name, var_src, grp_dest)

# %% ../nbs/api/encoders.ipynb 13
@patch
def _create_and_copy_variable(self:NetCDFEncoder, var_name, var_src, df, grp_dest, dtype_name):
    variable_type = self.enum_dtypes.get(dtype_name, var_src.datatype)    
    grp_dest.createVariable(var_name, variable_type, NC_DIM, compression='zlib', complevel=9)            
    isNotEnum = type(variable_type) != netCDF4._netCDF4.EnumType
    values = df[self.nc_to_cols[var_name]].values
    grp_dest[var_name][:] = values if isNotEnum else self.sanitize_if_enum_and_nan(values)

# %% ../nbs/api/encoders.ipynb 14
@patch
def sanitize_if_enum_and_nan(self:NetCDFEncoder, values, fill_value=-1):
    values[np.isnan(values)] = int(fill_value)
    values = values.astype(int)
    return values

# %% ../nbs/api/encoders.ipynb 16
@patch
def copy_variable_attributes(self:NetCDFEncoder, var_name, var_src, grp_dest):
    grp_dest[var_name].setncatts(var_src.__dict__)

# %% ../nbs/api/encoders.ipynb 17
@patch
def retrieve_all_cols(self:NetCDFEncoder, 
                      dtypes=NC_DTYPES
                      ):
    "Retrieve all unique columns from the dict of dataframes." 
    return list(set(col for df in self.dfs.values() for col in df.columns if col in dtypes.keys()))

# %% ../nbs/api/encoders.ipynb 18
@patch
def create_enums(self:NetCDFEncoder):
    cols = self.retrieve_all_cols()
    enums = Enums(lut_src_dir=lut_path())
    for col in cols:
        name = NC_DTYPES[col]['name']
        if self.verbose: print(f'Creating enum for {name} with values {enums.types[col]}.')
        dtype = self.dest.createEnumType(np.int64, name, enums.types[col])
        self.enum_dtypes[name] = dtype

# %% ../nbs/api/encoders.ipynb 19
@patch
def copy_custom_map(self:NetCDFEncoder, var_name, grp_dest):
    """Copy custom maps for variables."""
    custom_maps = self.custom_maps
    # Convert group names using NC_GROUPS
    custom_maps = {NC_GROUPS[key]: value for key, value in custom_maps.items()}
    group_maps = custom_maps.get(grp_dest.name, {})
    # Convert var names using NC_VARS
    group_maps = {NC_VARS[key]: value for key, value in group_maps.items()}
    if var_name in group_maps:
        # Set the map as an attribute of the variable
        grp_dest[var_name].setncatts({f"{var_name}_map": str(group_maps[var_name])})

# %% ../nbs/api/encoders.ipynb 20
@patch
def encode(self:NetCDFEncoder):
    "Encode MARIS NetCDF based on template and dataframes."
    with Dataset(self.src_fname, format='NETCDF4') as self.src, Dataset(self.dest_fname, 'w', format='NETCDF4') as self.dest:
        self.copy_global_attributes()
        self.create_enums()
        self.process_groups()
