"""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',
           'SelectColumnsCB', 'RenameColumnsCB', 'RemoveAllNAValuesCB', 'CompareDfsAndTfmCB', 'UniqueIndexCB',
           'EncodeTimeCB', 'DecodeTimeCB']

# %% ../nbs/api/callbacks.ipynb 2
#import copy
import fastcore.all as fc
from operator import attrgetter
from cftime import date2num ,num2date
import numpy as np
import pandas as pd
from typing import List, Dict, Callable, Any, Optional, Union
from collections import defaultdict
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
                 custom_maps: Dict = None,
                 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 = []
        self.custom_maps = custom_maps or defaultdict(lambda: defaultdict(dict))
            
    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 = 0, # Default value for unmatched entries
                 verbose: bool=False # Whether to print the number of 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:
                if self.verbose:
                    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='SAMPLE_TYPE' # Column name to store the sample type id
                 ): 
        "Add a column with the sample type 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 32
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 33
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 34
class RemoveAllNAValuesCB(Callback):
    "Remove rows with all NA values in specified columns."
    def __init__(self, 
                 cols_to_check: Union[Dict[str, list], list],  # Dict or list of columns to check
                 how: str='all'  # How to handle NA values 'all' or 'any'
                ):
        fc.store_attr()

    def __call__(self, tfm):
        # Convert list to dict if cols_to_check is a list
        cols_dict = (self.cols_to_check if isinstance(self.cols_to_check, dict) 
                    else {k: self.cols_to_check for k in tfm.dfs.keys()})
        
        for sample_type, columns in cols_dict.items():
            tfm.dfs[sample_type].dropna(
                subset=columns,
                how=self.how,
                inplace=True
            )

# %% ../nbs/api/callbacks.ipynb 37
class CompareDfsAndTfmCB(Callback):
    "Create a dataframe of removed data and track changes in row counts due to transformations."
    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():
            self._compute_changes(grp, tfm)

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

    def _compute_changes(self, 
                         grp: str,  # The group key
                         tfm: Transformer  # The transformation object containing `dfs`
                        ) -> None:
        "Compute and store changes including data removed and created during transformation."
        original_df = self.dfs[grp]
        transformed_df = tfm.dfs[grp]

        # Calculate differences
        original_count = len(original_df.index)
        transformed_count = len(transformed_df.index)
        removed_count = len(original_df.index.difference(transformed_df.index))
        created_count = len(transformed_df.index.difference(original_df.index))

        # Store results
        tfm.dfs_removed[grp] = original_df.loc[original_df.index.difference(transformed_df.index)]
        tfm.compare_stats[grp] = {
            'Original row count (dfs)': original_count,
            'Transformed row count (tfm.dfs)': transformed_count,
            'Rows removed from original (tfm.dfs_removed)': removed_count,
            'Rows created in transformed (tfm.dfs_created)': created_count
        }

# %% ../nbs/api/callbacks.ipynb 39
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 41
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))

# %% ../nbs/api/callbacks.ipynb 43
class DecodeTimeCB(Callback):
    "Decode time from seconds since epoch to datetime format."    
    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 datetime
            tfm.dfs[grp] = tfm.dfs[grp][tfm.dfs[grp][self.col_time].notna()]
            tfm.dfs[grp][self.col_time] = df[self.col_time].apply(
                lambda x: num2date(x, units=self.units, only_use_cftime_datetimes=False)
            )
