# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/handlers/data_format_transformation.ipynb.

# %% auto 0
__all__ = ['TAXON_MAP', 'lut_taxon', 'or_mappings', 'ValidateEnumsCB', 'RemoveNonCompatibleVariablesCB', 'get_taxon_info_lut',
           'AddTaxonInformationCB', 'AddZoteroArchiveLocationCB', 'RemapCustomMapsCB', 'RemapToORSpecificMappingsCB',
           'get_excluded_enums', 'DataFormatConversionCB', 'decode']

# %% ../../nbs/handlers/data_format_transformation.ipynb 5
from pathlib import Path
from netCDF4 import Dataset
import pandas as pd
import fastcore.all as fc
from typing import Dict, List,Callable, Tuple

from marisco.configs import (
    NC_VARS,
    CSV_VARS,
    CSV_DTYPES,
    Enums,
    lut_path,
    species_lut_path,
    detection_limit_lut_path, # used for feedback. 
    filtered_lut_path,
    cfg
)

from marisco.utils import (
    ExtractNetcdfContents,
)

from marisco.callbacks import (
    Callback,
    Transformer,
    DecodeTimeCB,
    AddSampleTypeIdColumnCB
)  
    
from marisco.decoders import (
    NetCDFDecoder
    )
from marisco.metadata import (
    ZoteroItem
)

# %% ../../nbs/handlers/data_format_transformation.ipynb 25
class ValidateEnumsCB(Callback):
    "Validate enumeration mappings between NetCDF file and MARIS lookup tables."

    def __init__(self, contents, maris_enums, verbose=False):
        fc.store_attr()

    def __call__(self, tfm):
        for group_name, enum_dict in self.contents.enum_dicts.items():
            self._validate_group(group_name, enum_dict)

    def _validate_group(self, group_name, enum_dict):
        
        for var_name, nc_enum_dict in enum_dict.items():
            if self.verbose:
                print(f"Validating variable {var_name} from NetCDF group {group_name}.")
            var_name = self._get_original_var_name(var_name)
            if self.verbose:
                print(f"Standardized variable name to MARISCO naming convention: {var_name}")

            if var_name not in self.maris_enums.types:
                if self.verbose:
                    print(f"Variable {var_name} not found in MARISCO enums.")
                continue

            self._compare_mappings(nc_enum_dict, self.maris_enums.types[var_name], group_name, var_name)

    def _get_original_var_name(self, var_name):
        return next((var for var, nc_var in NC_VARS.items() if nc_var == var_name), var_name)

    def _compare_mappings(self, nc_dict, maris_enum, group_name, var_name):        
        for key, value in nc_dict.items():
            value=int(value)
            if key not in maris_enum or maris_enum[key] != value:
                print(f"\nWarning: Enum mismatch: {var_name} in {group_name}.")
                print(f"   NetCDF value: {key} -> {value}")
                print(f"   MARISCO standard enum lookup value: {key} -> {maris_enum.get(key, 'Not found')}")
            

# %% ../../nbs/handlers/data_format_transformation.ipynb 29
class RemoveNonCompatibleVariablesCB(Callback):
    "Remove variables not listed in VARS configuration."
    def __init__(self, 
                vars: Dict[str, str] = CSV_VARS,  # Dictionary mapping OR vars to NC vars
                verbose: bool = False,
                ):
        fc.store_attr()
        
    def __call__(self, tfm: Transformer):
        """Remove non-OR variables from all dataframes."""
        for group_name in tfm.dfs:
            tfm.dfs[group_name] = self._remove_non_vars(tfm.dfs[group_name], group_name)
            
    def _remove_non_vars(self, df: pd.DataFrame, group_name:str ) -> pd.DataFrame:
        """Remove variables not in vars and print removed columns if verbose."""
        current_cols = set(df.columns)
        vars_cols = set(self.vars.keys())
        cols_to_remove = current_cols - vars_cols
        
        if self.verbose and cols_to_remove:
            print(f"Removing variables that are not compatible with vars provided. \nRemoving {', '.join(cols_to_remove)} from {group_name} dataset.")
                        
        return df.drop(columns=cols_to_remove)


# %% ../../nbs/handlers/data_format_transformation.ipynb 32
TAXON_MAP = {
    'Taxonname': 'TAXONNAME',
    'Taxonrank': 'TAXONRANK',
    'TaxonDB': 'TAXONDB',
    'TaxonDBID': 'TAXONDBID',
    'TaxonDBURL': 'TAXONDBURL'
}

# %% ../../nbs/handlers/data_format_transformation.ipynb 33
def get_taxon_info_lut(maris_lut: str, key_names: dict = TAXON_MAP) -> dict:
    "Create lookup dictionary for taxon information from MARIS species lookup table."
    species = pd.read_excel(maris_lut)
    # Select columns and rename them to standardized format
    columns = ['species_id'] + list(key_names.keys())
    df = species[columns].rename(columns=key_names)
    return df.set_index('species_id').to_dict()

lut_taxon = lambda: get_taxon_info_lut(maris_lut=species_lut_path(), key_names=TAXON_MAP)

# %% ../../nbs/handlers/data_format_transformation.ipynb 34
class AddTaxonInformationCB(Callback):
    """Add taxon information to BIOTA group based on species lookup table."""
    
    def __init__(self, 
                fn_lut: Callable = lut_taxon,  # Function that returns taxon lookup dictionary
                verbose: bool = False
                ):
        fc.store_attr()
        
    def __call__(self, tfm: Transformer):
        """Delegate tasks to add taxon information to the BIOTA group."""
        if not self.check_biota_group_exists(tfm):
            return
        
        df = tfm.dfs['BIOTA']
        if not self.check_species_column_exists(df):
            return
        
        self.add_taxon_columns(df)

    def check_biota_group_exists(self, tfm: Transformer) -> bool:
        """Check if 'BIOTA' group exists in the dataframes."""
        if 'BIOTA' not in tfm.dfs:
            if self.verbose:
                print("No BIOTA group found, skipping taxon information")
            return False
        return True

    def check_species_column_exists(self, df: pd.DataFrame) -> bool:
        """Check if 'SPECIES' column exists in the BIOTA dataframe."""
        if 'SPECIES' not in df.columns:
            if self.verbose:
                print("No SPECIES column found in BIOTA dataframe, skipping taxon information")
            return False
        return True

    def add_taxon_columns(self, df: pd.DataFrame):
        """Add taxon information columns to the BIOTA dataframe."""
        lut = self.fn_lut()
        
        # Add each column from the lookup table
        for col in lut.keys():
            df[col] = df['SPECIES'].map(lut[col]).fillna('Unknown')
        
        self.report_unmatched_species(df)

    def report_unmatched_species(self, df: pd.DataFrame):
        """Report any species IDs not found in the lookup table."""
        unmatched = df[df['TAXONNAME'] == 'Unknown']['SPECIES'].unique()
        if self.verbose and len(unmatched) > 0:
            print(f"Warning: Species IDs not found in lookup table: {', '.join(map(str, unmatched))}")

# %% ../../nbs/handlers/data_format_transformation.ipynb 43
class AddZoteroArchiveLocationCB(Callback):
    "Fetch and append 'Loc. in Archive' from Zotero to DataFrame."
    def __init__(self, attrs: str, cfg: dict):
        fc.store_attr()

    def __call__(self, tfm):
        
        zotero_key = self.attrs['id']
        item = ZoteroItem(zotero_key, self.cfg['zotero'])
        if item.exist():
            loc_in_archive = item.item['data']['archiveLocation'] 
            for grp, df in tfm.dfs.items():
                df['REF_ID'] = int(loc_in_archive)
        else:
            print(f"Warning: Zotero item {self.item_id} does not exist.")

# %% ../../nbs/handlers/data_format_transformation.ipynb 47
class RemapCustomMapsCB(Callback):
    "Remap encoded custom maps to decoded values."
    def __init__(self, verbose: bool = False):
        fc.store_attr()
        
    def __call__(self, tfm):
        """Remap encoded custom maps to decoded values."""
        
        for grp in tfm.dfs:
            for var in tfm.dfs[grp].columns:
                if var in tfm.custom_maps[grp]:
                    if self.verbose:
                        print(f"Remapping {var} from {grp} group")
                    
                    # Convert column to int type to ensure proper mapping
                    tfm.dfs[grp][var] = tfm.dfs[grp][var].astype(int)
                    
                    # Create reverse mapping dictionary
                    reverse_custom_map = {int(v): k for k, v in tfm.custom_maps[grp][var].items()}
                    
                    # Apply mapping
                    tfm.dfs[grp][var] = tfm.dfs[grp][var].map(reverse_custom_map)

# %% ../../nbs/handlers/data_format_transformation.ipynb 56
or_mappings={'DL':
                {0:'ND',1:'=',2:'<'},
            'FILT':
                {0:'NA',1:'Y',2:'N'},
            }

# %% ../../nbs/handlers/data_format_transformation.ipynb 58
class RemapToORSpecificMappingsCB(Callback):
    "Convert values using OR mappings if columns exist in dataframe."
    def __init__(self, 
                or_mappings: Dict[str, Dict] = or_mappings,  # Dictionary of column mappings, 
                output_format: str = 'openrefine_csv',
                verbose: bool = False
                ):
        fc.store_attr()
        
    def __call__(self, tfm: Transformer):
        """Apply OR mappings to all dataframes."""
        for group_name in tfm.dfs:
            if self.verbose:
                print(f"\nProcessing {group_name} group...")
            tfm.dfs[group_name] = self._apply_mappings(tfm.dfs[group_name])
            
    def _apply_mappings(self, df: pd.DataFrame) -> pd.DataFrame:
        """Apply OR mappings to columns that exist in the dataframe."""
        for col, mapping in self.or_mappings.items():
            if col in df.columns:
                if self.verbose:
                    print(f"    Mapping values for column: {col}")
                df[col] = df[col].map(mapping)
        return df


# %% ../../nbs/handlers/data_format_transformation.ipynb 65
def get_excluded_enums(output_format: str = 'openrefine_csv') -> dict:
    "Get excluded enums based on output format."
    return or_mappings if output_format == 'openrefine_csv' else {}

# %% ../../nbs/handlers/data_format_transformation.ipynb 66
class DataFormatConversionCB(Callback):
    """
    A callback to convert DataFrame enum values between encoded and decoded formats based on specified settings.
    """

    def __init__(self, 
                 dtypes: Dict,  # Dictionary defining data types and states for each lookup table
                 excluded_mappings: Callable = get_excluded_enums,  # Dictionary of columns to exclude from conversion
                 output_format: str = 'openrefine_csv',
                 verbose: bool = False  # Flag for verbose output
                ):
        fc.store_attr()

    def __call__(self, tfm):
        """
        Apply the data format conversion to each DataFrame within the Transformer.
        """
        self.load_enums()
        
        for group_name, df in tfm.dfs.items():
            tfm.dfs[group_name] = self.process_dataframe(group_name, df)

    def load_enums(self):
        """
        Load enums from the lookup path.
        """
        self.enums = Enums(lut_path())
        if self.verbose:
            print(f"Loaded enums: {self.enums.types.keys()}")

    def process_dataframe(self, group_name: str, df: pd.DataFrame):
        """
        Process each DataFrame to convert columns to the target state.
        """
        for column in df.columns:
            if column in self.dtypes and column not in self.excluded_mappings(self.output_format):
                if self.dtypes[column]['state'] == 'decoded':
                    if self.verbose:
                        print(f"Decoding column: {column}")
                    if column in self.enums.types:
                        # Apply the mapping from encoded to decoded values
                        df[column] = df[column].map(self.enums.types[column])
                        if self.verbose:
                            print(f"Decoded column: {column}")
                    else:
                        if self.verbose:
                            print(f"No enum mapping found for column: {column}, skipping decoding.")
        return df

# %% ../../nbs/handlers/data_format_transformation.ipynb 71
def decode(
    fname_in: str, # Input file name
    dest_out: str | None = None, # Output file name (optional)
    output_format: str = 'openrefine_csv',
    remap_vars: Dict[str, str] = CSV_VARS,
    remap_dtypes: Dict[str, str] = CSV_DTYPES,
    verbose: bool = False,
    **kwargs # Additional arguments
    ) -> None:
    "Decode data from NetCDF."
    valid_output_formats=['openrefine_csv', 'decoded_csv']
    
    if output_format not in valid_output_formats:
        print (f'Invalid output format. Allowed formats: {valid_output_formats}')
        return 
    
    if output_format == 'decoded_csv':
        remap_dtypes = {k: {'state': 'decoded'} for k in remap_dtypes.keys()}
        
    contents = ExtractNetcdfContents(fname_in)
    tfm = Transformer(
        data=contents.dfs,
        custom_maps=contents.custom_maps,
        cbs=[
        ValidateEnumsCB(
            contents = contents,
            maris_enums=Enums(lut_src_dir=lut_path())
        ),
        RemoveNonCompatibleVariablesCB(vars=remap_vars),
        RemapCustomMapsCB(),
        RemapToORSpecificMappingsCB(output_format=output_format),
        AddTaxonInformationCB(
            fn_lut=lut_taxon
        ),
        DecodeTimeCB(),
        AddSampleTypeIdColumnCB(),
        AddZoteroArchiveLocationCB(contents.global_attrs, cfg=cfg()),
        DataFormatConversionCB(
            dtypes=remap_dtypes,
            excluded_mappings = get_excluded_enums,
            output_format=output_format
        ) 
        ]
    )    
    
    tfm()
    decoder = NetCDFDecoder( 
                            dfs=tfm.dfs,
                            fname_in=fname_in,  
                            dest_out=dest_out,                           
                            output_format='csv',
                            remap_vars=CSV_VARS,
                            verbose=verbose
                    )
    decoder.decode()
