"""Callback used in handlers"""

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

# %% auto 0
__all__ = ['Callback', 'run_cbs', 'Transformer', 'SanitizeLonLatCB', 'RemapCB', 'LowerStripNameCB', 'AddSampleTypeIdColumnCB',
           'AddNuclideIdColumnCB', 'SelectColumnsCB', 'RenameColumnsCB', 'RemoveAllNAValuesCB', 'CompareDfsAndTfmCB',
           'UniqueIndexCB', 'EncodeTimeCB']

# %% ../nbs/api/callbacks.ipynb 2
import copy
import fastcore.all as fc
from operator import attrgetter
from cftime import date2num
import numpy as np
import pandas as pd
from functools import partial 
from pathlib import Path 
from typing import List, Dict, Callable, Tuple, Any, Optional, Union

from marisco.configs import (
    get_lut, 
    nuc_lut_path, 
    nc_tpl_path,
    get_time_units,
    NC_GROUPS,
    SMP_TYPE_LUT,
    cfg, 
    # cdl_cfg
)

from .utils import Match

# %% ../nbs/api/callbacks.ipynb 6
class Callback(): 
    "Base class for callbacks."
    order = 0

# %% ../nbs/api/callbacks.ipynb 7
def run_cbs(
    cbs: List[Callback], # List of callbacks to run
    obj: Any # Object to pass to the callbacks
    ):
    "Run the callbacks in the order they are specified."
    for cb in sorted(cbs, key=attrgetter('order')):
        if cb.__doc__: obj.logs.append(cb.__doc__)
        cb(obj)

# %% ../nbs/api/callbacks.ipynb 8
class Transformer():
    "Transform the dataframe(s) according to the specified callbacks."
    def __init__(self, 
                 data: Union[Dict[str, pd.DataFrame], pd.DataFrame], # Data to be transformed
                 cbs: Optional[List[Callback]]=None, # List of callbacks to run
                 inplace: bool=False # Whether to modify the dataframe(s) in place
                 ): 
        fc.store_attr()
        self.is_single_df = isinstance(data, pd.DataFrame)
        self.df, self.dfs = self._prepare_data(data, inplace)
        self.logs = []
            
    def _prepare_data(self, data, inplace):
        if self.is_single_df:
            return (data if inplace else data.copy()), None
        else:
            return None, (data if inplace else {k: v.copy() for k, v in data.items()})
    
    def unique(self, col_name: str) -> np.ndarray:
        "Distinct values of a specific column present in all groups."
        if self.is_single_df:
            values = self.df.get(col_name, pd.Series()).dropna().values
        else:
            columns = [df.get(col_name) for df in self.dfs.values() if df.get(col_name) is not None]
            values = np.concatenate([col.dropna().values for col in columns]) if columns else []
        return np.unique(values)
        
    def __call__(self):
        "Transform the dataframe(s) according to the specified callbacks."
        if self.cbs: run_cbs(self.cbs, self)
        return self.df if self.dfs is None else self.dfs

# %% ../nbs/api/callbacks.ipynb 17
class SanitizeLonLatCB(Callback):
    "Drop rows with invalid longitude & latitude values. Convert `,` separator to `.` separator."
    def __init__(self, 
                 lon_col: str='LON', # Longitude column name
                 lat_col: str='LAT', # Latitude column name
                 verbose: bool=False # Whether to print the number of invalid longitude & latitude values
                 ):
        fc.store_attr()
        
    def __call__(self, tfm: Transformer):
        for grp, df in tfm.dfs.items():
            # Convert `,` separator to `.` separator
            df[self.lon_col] = df[self.lon_col].apply(lambda x: float(str(x).replace(',', '.')))
            df[self.lat_col] = df[self.lat_col].apply(lambda x: float(str(x).replace(',', '.')))
            
            # Mask zero values
            mask_zeroes = (df[self.lon_col] == 0) & (df[self.lat_col] == 0) 
            nZeroes = mask_zeroes.sum()
            if nZeroes and self.verbose: 
                print(f'The "{grp}" group contains {nZeroes} data points whose ({self.lon_col}, {self.lat_col}) = (0, 0)')
            
            # Mask out of bounds values
            mask_goob = (df[self.lon_col] < -180) | (df[self.lon_col] > 180) | (df[self.lat_col] < -90) | (df[self.lat_col] > 90)
            nGoob = mask_goob.sum()
            if nGoob and self.verbose: 
                print(f'The "{grp}" group contains {nGoob} data points with unrealistic {self.lon_col} or {self.lat_col} values.')
                
            tfm.dfs[grp] = df.loc[~(mask_zeroes | mask_goob)]

# %% ../nbs/api/callbacks.ipynb 22
class RemapCB(Callback):
    "Generic MARIS remapping callback."
    def __init__(self, 
                 fn_lut: Callable, # Function that returns the lookup table dictionary
                 col_remap: str, # Name of the column to remap
                 col_src: str, # Name of the column with the source values
                 dest_grps: list[str]|str=NC_GROUPS.keys(), # List of destination groups
                 default_value: Any = -1 # Default value for unmatched entries
                ):
        fc.store_attr()
        self.lut = None
        
        if isinstance(dest_grps, str):
            self.dest_grps = [dest_grps]
        # Format the documentation string based on the type and content of dest_grps
        if isinstance(self.dest_grps, list):
            if len(self.dest_grps) > 1:
                grp_str = ', '.join(self.dest_grps[:-1]) + ' and ' + self.dest_grps[-1]
            else:
                grp_str = self.dest_grps[0]
        else:
            grp_str = self.dest_grps
                
        self.__doc__ = f"Remap values from '{col_src}' to '{col_remap}' for groups: {grp_str}."

    def __call__(self, tfm):
        self.lut = self.fn_lut()
        for grp in self.dest_grps:
            if grp in tfm.dfs:
                self._remap_group(tfm.dfs[grp])
            else:
                print(f"Group {grp} not found in the dataframes.")

    def _remap_group(self, df: pd.DataFrame):
        df[self.col_remap] = df[self.col_src].apply(self._remap_value)

    def _remap_value(self, value: str) -> Any:
        value = value.strip() if isinstance(value, str) else value
        match = self.lut.get(value, Match(self.default_value, None, None, None))
        if isinstance(match, Match):
            if match.matched_id == self.default_value:
                print(f"Unmatched value: {value}")
            return match.matched_id 
        else:
            return match

# %% ../nbs/api/callbacks.ipynb 23
class LowerStripNameCB(Callback):
    "Convert values to lowercase and strip any trailing spaces."
    def __init__(self, 
                 col_src: str, # Source column name e.g. 'Nuclide'
                 col_dst: str=None, # Destination column name
                 fn_transform: Callable=lambda x: x.lower().strip() # Transformation function
                 ):
        fc.store_attr()
        self.__doc__ = f"Convert '{col_src}' column values to lowercase, strip spaces, and store in '{col_dst}' column."
        if not col_dst: self.col_dst = col_src
        
    def _safe_transform(self, value):
        "Ensure value is not NA and apply transformation function."
        return value if pd.isna(value) else self.fn_transform(str(value))
            
    def __call__(self, tfm):
        for key in tfm.dfs.keys():
            tfm.dfs[key][self.col_dst] = tfm.dfs[key][self.col_src].apply(self._safe_transform)

# %% ../nbs/api/callbacks.ipynb 28
class AddSampleTypeIdColumnCB(Callback):
    def __init__(self, 
                 lut: dict=SMP_TYPE_LUT, # Lookup table for sample type
                 col_name: str='samptype_id' # Column name to store the sample type id
                 ): 
        "Add a column with the sample type id as defined in the CDL."
        fc.store_attr()
        
    def __call__(self, tfm):
        for grp, df in tfm.dfs.items():             
            df[self.col_name] = self.lut[grp]

# %% ../nbs/api/callbacks.ipynb 31
class AddNuclideIdColumnCB(Callback):
    def __init__(self, 
                 col_value: str, # Column name containing the nuclide name
                 lut_fname_fn: Callable=nuc_lut_path, # Function returning the lut path
                 col_name: str='nuclide_id' # Column name to store the nuclide id
                 ): 
        "Add a column with the nuclide id."
        fc.store_attr()
        self.lut = get_lut(lut_fname_fn().parent, lut_fname_fn().name, 
                           key='nc_name', value='nuclide_id', reverse=False)
        
    def __call__(self, tfm: Transformer):
        for grp, df in tfm.dfs.items(): 
            df[self.col_name] = df[self.col_value].map(self.lut)

# %% ../nbs/api/callbacks.ipynb 33
class SelectColumnsCB(Callback):
    "Select columns of interest."
    def __init__(self, 
                 cois: dict # Columns of interest
                 ): 
        fc.store_attr()
        
    def __call__(self, tfm):
        "Select columns of interest."
        for grp, df in tfm.dfs.items(): 
            tfm.dfs[grp] = df.loc[:, self.cois.keys()]

# %% ../nbs/api/callbacks.ipynb 34
class RenameColumnsCB(Callback):
    "Renaming variables to MARIS standard names."
    def __init__(self,
                 renaming_rules: dict # Renaming rules
                 ): 
        fc.store_attr()
        
    def __call__(self, tfm):
        for grp in tfm.dfs.keys(): 
            tfm.dfs[grp].rename(columns=self.renaming_rules, inplace=True)

# %% ../nbs/api/callbacks.ipynb 35
class RemoveAllNAValuesCB(Callback):
    "Remove rows with all NA values."
    def __init__(self, 
                 cols_to_check: Dict[str, str] # A dictionary with the sample type as key and the column name to check as value
                ):
        fc.store_attr()

    def __call__(self, tfm):
        for k in tfm.dfs.keys():
            col_to_check = self.cols_to_check[k]
            mask = tfm.dfs[k][col_to_check].isnull().all(axis=1)
            tfm.dfs[k] = tfm.dfs[k][~mask]

# %% ../nbs/api/callbacks.ipynb 44
class CompareDfsAndTfmCB(Callback):
    "Create a dataframe of dropped data. Data included in the `dfs` not in the `tfm`."
    def __init__(self, 
                 dfs: Dict[str, pd.DataFrame] # Original dataframes
                 ): 
        fc.store_attr()
        
    def __call__(self, tfm: Transformer) -> None:
        self._initialize_tfm_attributes(tfm)
        for grp in tfm.dfs.keys():
            dropped_df = self._get_dropped_data(grp, tfm)
            tfm.dfs_dropped[grp] = dropped_df
            tfm.compare_stats[grp] = self._compute_stats(grp, tfm)

    def _initialize_tfm_attributes(self, tfm: Transformer) -> None:
        tfm.dfs_dropped = {}
        tfm.compare_stats = {}

    def _get_dropped_data(self, 
                          grp: str, # The group key
                          tfm: Transformer # The transformation object containing `dfs`
                         ) -> pd.DataFrame: # Dataframe with dropped rows
        "Get the data that is present in `dfs` but not in `tfm.dfs`."
        index_diff = self.dfs[grp].index.difference(tfm.dfs[grp].index)
        return self.dfs[grp].loc[index_diff]
    
    def _compute_stats(self, 
                       grp: str, # The group key
                       tfm: Transformer # The transformation object containing `dfs`
                      ) -> Dict[str, int]: # Dictionary with comparison statistics
        "Compute comparison statistics between `dfs` and `tfm.dfs`."
        return {
            'Number of rows in dfs': len(self.dfs[grp].index),
            'Number of rows in tfm.dfs': len(tfm.dfs[grp].index),
            'Number of rows removed': len(tfm.dfs_dropped[grp].index),
        }

# %% ../nbs/api/callbacks.ipynb 47
class UniqueIndexCB(Callback):
    "Set unique index for each group."
    def __init__(self,
                 index_name='ID'):
        fc.store_attr()
        
    def __call__(self, tfm):
        for k in tfm.dfs.keys():
            # Reset the index of the DataFrame and drop the old index
            tfm.dfs[k] = tfm.dfs[k].reset_index(drop=True)
            # Reset the index again and set the name of the new index to `ìndex_name``
            tfm.dfs[k] = tfm.dfs[k].reset_index(names=[self.index_name])

# %% ../nbs/api/callbacks.ipynb 50
class EncodeTimeCB(Callback):
    "Encode time as seconds since epoch."    
    def __init__(self, 
                 col_time: str='TIME',
                 fn_units: Callable=get_time_units # Function returning the time units
                 ): 
        fc.store_attr()
        self.units = fn_units()
    
    def __call__(self, tfm): 
        for grp, df in tfm.dfs.items():
            n_missing = df[self.col_time].isna().sum()
            if n_missing:
                print(f"Warning: {n_missing} missing time value(s) in {grp}")
            
            # Remove NaN times and convert to seconds since epoch
            tfm.dfs[grp] = tfm.dfs[grp][tfm.dfs[grp][self.col_time].notna()]
            tfm.dfs[grp][self.col_time] = tfm.dfs[grp][self.col_time].apply(lambda x: date2num(x, units=self.units))
