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

# %% auto 0
__all__ = ['src_dir', 'fname_out_nc', 'zotero_key', 'default_smp_types', 'nan_cols_to_check', 'fixes_nuclide_names',
           'lut_nuclides', 'time_cols', 'time_format', 'value_cols', 'unc_exp2stan', 'unc_cols', 'renaming_unit_rules',
           'default_units', 'lut_dl', 'coi_dl', 'fixes_biota_species', 'lut_biota', 'fixes_enhanced_biota_species',
           'lut_biota_enhanced', 'fixes_biota_tissues', 'lut_bodyparts', 'lut_biogroup_from_biota', 'kw', 'read_csv',
           'load_data', 'RemapNuclideNameCB', 'ParseTimeCB', 'SanitizeValueCB', 'NormalizeUncCB', 'RemapUnitCB',
           'RemapDetectionLimitCB', 'EnhanceSpeciesCB', 'AddBodypartTempCB', 'AddSampleIdCB', 'AddDepthCB',
           'ConvertLonLatCB', 'get_attrs', 'encode']

# %% ../../nbs/handlers/ospar.ipynb 6
import pandas as pd 
import numpy as np
import fastcore.all as fc 
from typing import  Dict, Callable 
from pathlib import Path 
import time
from rich import print

from marisco.utils import (
    Remapper, 
    get_unique_across_dfs,
    NA
)

from marisco.callbacks import (
    Callback, 
    Transformer, 
    EncodeTimeCB, 
    LowerStripNameCB, 
    SanitizeLonLatCB, 
    CompareDfsAndTfmCB, 
    RemapCB,
    RemoveAllNAValuesCB
)

from marisco.metadata import (
    GlobAttrsFeeder, 
    BboxCB, 
    DepthRangeCB, 
    TimeRangeCB, 
    ZoteroCB, 
    KeyValuePairCB
)

from marisco.configs import (
    nuc_lut_path, 
    cfg, 
    species_lut_path, 
    bodyparts_lut_path, 
    detection_limit_lut_path, 
    get_lut, 
    cache_path
)

from ..encoders import NetCDFEncoder
from ..netcdf_to_refine import decode
from ..utils import ExtractNetcdfContents

# %% ../../nbs/handlers/ospar.ipynb 10
src_dir = 'https://raw.githubusercontent.com/franckalbinet/maris-crawlers/refs/heads/main/data/processed/OSPAR'
fname_out_nc = '../../_data/output/191-OSPAR-2024.nc'
zotero_key ='LQRA4MMK' # OSPAR MORS zotero key

# %% ../../nbs/handlers/ospar.ipynb 13
default_smp_types = {  
    'Biota': 'BIOTA', 
    'Seawater': 'SEAWATER', 
}

# %% ../../nbs/handlers/ospar.ipynb 14
def read_csv(file_name, dir=src_dir):
    file_path = f'{dir}/{file_name}'
    return pd.read_csv(file_path)

# %% ../../nbs/handlers/ospar.ipynb 15
def load_data(src_url: str, 
              smp_types: dict = default_smp_types, # Sample types to load
              use_cache: bool = False, # Use cache
              save_to_cache: bool = False, # Save to cache
              verbose: bool = False # Verbose
              ) -> Dict[str, pd.DataFrame]:
    "Load OSPAR data and return the data in a dictionary of dataframes with the dictionary key as the sample type."
    
    def safe_file_path(url: str) -> str:
        "Safely encode spaces in a URL."
        return url.replace(" ", "%20")

    def get_file_path(dir_path: str, file_prefix: str) -> str:
        """Construct the full file path based on directory and file prefix."""
        file_path = f"{dir_path}/{file_prefix} data.csv"
        return safe_file_path(file_path) if not use_cache else file_path

    def load_and_process_csv(file_path: str) -> pd.DataFrame:
        """Load a CSV file and process it."""
        if use_cache and not Path(file_path).exists():
            if verbose:
                print(f"{file_path} not found in cache.")
            return pd.DataFrame()

        if verbose:
            start_time = time.time()

        try:
            df = pd.read_csv(file_path)
            df.columns = df.columns.str.lower()
            if verbose:
                print(f"Data loaded from {file_path} in {time.time() - start_time:.2f} seconds.")
            return df
        except Exception as e:
            if verbose:
                print(f"Failed to load {file_path}: {e}")
            return pd.DataFrame()

    def save_to_cache_dir(df: pd.DataFrame, file_prefix: str):
        """Save the DataFrame to the cache directory."""
        cache_dir = cache_path()
        cache_file_path = f"{cache_dir}/{file_prefix} data.csv"
        df.to_csv(cache_file_path, index=False)
        if verbose:
            print(f"Data saved to cache at {cache_file_path}")

    data = {}
    for file_prefix, smp_type in smp_types.items():
        dir_path = cache_path() if use_cache else src_url
        file_path = get_file_path(dir_path, file_prefix)
        df = load_and_process_csv(file_path)

        if save_to_cache and not df.empty:
            save_to_cache_dir(df, file_prefix)

        data[smp_type] = df

    return data

# %% ../../nbs/handlers/ospar.ipynb 21
nan_cols_to_check = ['sampling date', 'activity or mda']

# %% ../../nbs/handlers/ospar.ipynb 45
fixes_nuclide_names = {
    '99tc': 'tc99',
    '238pu': 'pu238',
    '226ra': 'ra226',
    'ra-226': 'ra226',
    'ra-228': 'ra228',    
    '210pb': 'pb210',
    '241am': 'am241',
    '228ra': 'ra228',
    '137cs': 'cs137',
    '210po': 'po210',
    '239,240pu': 'pu239_240_tot',
    '239, 240 pu': 'pu239_240_tot',
    '3h': 'h3'
    }

# %% ../../nbs/handlers/ospar.ipynb 51
# Create a lookup table for nuclide names
lut_nuclides = lambda df: Remapper(provider_lut_df=df,
                                   maris_lut_fn=nuc_lut_path,
                                   maris_col_id='nuclide_id',
                                   maris_col_name='nc_name',
                                   provider_col_to_match='value',
                                   provider_col_key='value',
                                   fname_cache='nuclides_ospar.pkl').generate_lookup_table(fixes=fixes_nuclide_names, 
                                                                                            as_df=False, overwrite=True)

# %% ../../nbs/handlers/ospar.ipynb 52
class RemapNuclideNameCB(Callback):
    "Remap data provider nuclide names to standardized MARIS nuclide names."
    def __init__(self, 
                 fn_lut: Callable, # Function that returns the lookup table dictionary
                 col_name: str # Column name to remap
                ):
        fc.store_attr()

    def __call__(self, tfm: Transformer):
        df_uniques = get_unique_across_dfs(tfm.dfs, col_name=self.col_name, as_df=True)
        lut = {k: v.matched_id for k, v in self.fn_lut(df_uniques).items()}    
        for k in tfm.dfs.keys():
            tfm.dfs[k]['NUCLIDE'] = tfm.dfs[k][self.col_name].replace(lut)

# %% ../../nbs/handlers/ospar.ipynb 57
time_cols = {'BIOTA': 'sampling date', 'SEAWATER': 'sampling date'}
time_format = '%m/%d/%y %H:%M:%S'

# %% ../../nbs/handlers/ospar.ipynb 58
class ParseTimeCB(Callback):
    "Parse the time format in the dataframe and check for inconsistencies."
    def __init__(self, 
                 col_src: dict=time_cols, # Column name to remap
                 col_dst: str='TIME', # Column name to remap
                 format: str=time_format # Time format
                 ):
        fc.store_attr()
    
    def __call__(self, tfm):
        for grp, df in tfm.dfs.items():
            src_col = self.col_src.get(grp)
            df[self.col_dst] = pd.to_datetime(df[src_col], format=self.format, errors='coerce')
        return tfm

# %% ../../nbs/handlers/ospar.ipynb 66
value_cols = {'BIOTA': 'activity or mda', 'SEAWATER': 'activity or mda'}

# %% ../../nbs/handlers/ospar.ipynb 67
class SanitizeValueCB(Callback):
    "Sanitize value by removing blank entries and populating `value` column."
    def __init__(self, 
                 value_col: dict = value_cols # Column name to sanitize
                 ):
        fc.store_attr()

    def __call__(self, tfm):
        for grp, df in tfm.dfs.items():
            # Drop rows where parsing failed (NaT values in TIME column)
            invalid_rows = df[df[self.value_col.get(grp)].isna()]
            if not invalid_rows.empty:     
                print(f"{len(invalid_rows)} invalid rows found in group '{grp}' during sanitize value callback.")
                df.dropna(subset=[self.value_col.get(grp)], inplace=True)
                
            df['VALUE'] = df[self.value_col.get(grp)]

# %% ../../nbs/handlers/ospar.ipynb 72
unc_exp2stan = lambda df, unc_col: df[unc_col] / 2

# %% ../../nbs/handlers/ospar.ipynb 73
unc_cols = {'BIOTA': 'uncertainty', 'SEAWATER': 'uncertainty'}

# %% ../../nbs/handlers/ospar.ipynb 74
class NormalizeUncCB(Callback):
    """Normalize uncertainty values in DataFrames."""
    def __init__(self, 
                 col_unc: dict = unc_cols, # Column name to normalize
                 fn_convert_unc: Callable=unc_exp2stan, # Function correcting coverage factor
                 ): 
        fc.store_attr()

    def __call__(self, tfm):
        for grp, df in tfm.dfs.items():
            self._convert_commas_to_periods(df, self.col_unc.get(grp)   )
            self._convert_to_float(df, self.col_unc.get(grp))
            self._apply_conversion_function(df, self.col_unc.get(grp))

    def _convert_commas_to_periods(self, df, col_unc    ):
        """Convert commas to periods in the uncertainty column."""
        df[col_unc] = df[col_unc].astype(str).str.replace(',', '.')

    def _convert_to_float(self, df, col_unc):
        """Convert uncertainty column to float, handling errors by setting them to NaN."""
        df[col_unc] = pd.to_numeric(df[col_unc], errors='coerce')

    def _apply_conversion_function(self, df, col_unc):
        """Apply the conversion function to normalize the uncertainty values."""
        df['UNC'] = self.fn_convert_unc(df, col_unc)

# %% ../../nbs/handlers/ospar.ipynb 89
# Define unit names renaming rules
renaming_unit_rules = {'Bq/l': 1, #'Bq/m3'
                       'Bq/L': 1,
                       'BQ/L': 1,
                       'Bq/kg f.w.': 5, # Bq/kgw
                       } 

# %% ../../nbs/handlers/ospar.ipynb 91
default_units = {'SEAWATER': 'Bq/l',
                 'BIOTA': 'Bq/kg f.w.'}

# %% ../../nbs/handlers/ospar.ipynb 92
class RemapUnitCB(Callback):
    """Callback to update DataFrame 'UNIT' columns based on a lookup table."""

    def __init__(self,
                 lut: Dict[str, str],
                 default_units: Dict[str, str] = default_units,
                 verbose: bool = False
                 ):
        fc.store_attr()  # Store the lookup table as an attribute

    def __call__(self, tfm: 'Transformer'):
        for grp, df in tfm.dfs.items():
            # Apply default units to SEAWATER dataset
            if grp == 'SEAWATER':
                self._apply_default_units(df, unit=self.default_units.get(grp))
            # self._print_na_units(df)
            self._update_units(df)

    def _apply_default_units(self, df: pd.DataFrame , unit = None):
        df.loc[df['unit'].isnull(), 'unit'] = unit

    # def _print_na_units(self, df: pd.DataFrame):
    #     na_count = df['unit'].isnull().sum()
    #     if na_count > 0 and self.verbose:
    #         print(f"Number of rows with NaN in 'unit' column: {na_count}")

    def _update_units(self, df: pd.DataFrame):
        df['UNIT'] = df['unit'].apply(lambda x: self.lut.get(x, 'Unknown'))

# %% ../../nbs/handlers/ospar.ipynb 102
lut_dl = lambda: pd.read_excel(detection_limit_lut_path(), usecols=['name','id']).set_index('name').to_dict()['id']

# %% ../../nbs/handlers/ospar.ipynb 104
coi_dl = {'SEAWATER' : {'DL' : 'value type'},
          'BIOTA':  {'DL' : 'value type'}
          }

# %% ../../nbs/handlers/ospar.ipynb 106
class RemapDetectionLimitCB(Callback):
    """Remap detection limit values to MARIS format using a lookup table."""

    def __init__(self, coi: dict, fn_lut: Callable):
        """Initialize with column configuration and a function to get the lookup table."""
        fc.store_attr()        

    def __call__(self, tfm: Transformer):
        """Apply the remapping of detection limits across all dataframes"""
        lut = self.fn_lut()  # Retrieve the lookup table
        for grp, df in tfm.dfs.items():
            df['DL'] = df[self.coi[grp]['DL']]
            self._set_detection_limits(df, lut)

    def _set_detection_limits(self, df: pd.DataFrame, lut: dict):
        """Set detection limits based on value and uncertainty columns using specified conditions."""
        # Condition to set '=' when value and uncertainty are present and the current detection limit is not in the lookup keys
        condition_eq = (df['VALUE'].notna() & df['UNC'].notna() & ~df['DL'].isin(lut.keys()))
        df.loc[condition_eq, 'DL'] = '='

        # Set 'Not Available' for unmatched detection limits
        df.loc[~df['DL'].isin(lut.keys()), 'DL'] = 'Not Available'

        # Map existing detection limits using the lookup table
        df['DL'] = df['DL'].map(lut)

# %% ../../nbs/handlers/ospar.ipynb 117
fixes_biota_species = {
    'RHODYMENIA PSEUDOPALAMATA & PALMARIA PALMATA': NA,  # Mix of species, no direct mapping
    'Mixture of green, red and brown algae': NA,  # Mix of species, no direct mapping
    'Solea solea (S.vulgaris)': 'Solea solea',
    'SOLEA SOLEA (S.VULGARIS)': 'Solea solea',
    'RAJIDAE/BATOIDEA': NA, #Mix of species, no direct mapping
    'PALMARIA PALMATA': NA,  # Not defined
    'Unknown': NA,
    'unknown': NA,
    'Flatfish': NA,
    'Gadus sp.': NA,  # Not defined
}

# %% ../../nbs/handlers/ospar.ipynb 121
lut_biota = lambda: Remapper(provider_lut_df=get_unique_across_dfs(dfs, col_name='species', as_df=True),
                             maris_lut_fn=species_lut_path,
                             maris_col_id='species_id',
                             maris_col_name='species',
                             provider_col_to_match='value',
                             provider_col_key='value',
                             fname_cache='species_ospar.pkl').generate_lookup_table(fixes=fixes_biota_species, 
                                                                                    as_df=False, overwrite=False)

# %% ../../nbs/handlers/ospar.ipynb 132
fixes_enhanced_biota_species = {
    'fish': 'Pisces',
    'FISH': 'Pisces',
    'Fish': 'Pisces'    
}

# %% ../../nbs/handlers/ospar.ipynb 136
lut_biota_enhanced = lambda: Remapper(provider_lut_df=get_unique_across_dfs(dfs, col_name='biological group', as_df=True),
                             maris_lut_fn=species_lut_path,
                             maris_col_id='species_id',
                             maris_col_name='species',
                             provider_col_to_match='value',
                             provider_col_key='value',
                             fname_cache='enhance_species_ospar.pkl').generate_lookup_table(
                                 fixes=fixes_enhanced_biota_species, 
                                 as_df=False, 
                                 overwrite=False)

# %% ../../nbs/handlers/ospar.ipynb 140
class EnhanceSpeciesCB(Callback):
    """Enhance the 'SPECIES' column using the 'enhanced_species' column if conditions are met."""

    def __init__(self):
        fc.store_attr()

    def __call__(self, tfm: 'Transformer'):
        self._enhance_species(tfm.dfs['BIOTA'])

    def _enhance_species(self, df: pd.DataFrame):
        df['SPECIES'] = df.apply(
            lambda row: row['enhanced_species'] if row['SPECIES'] in [-1, 0] and pd.notnull(row['enhanced_species']) else row['SPECIES'],
            axis=1
        )

# %% ../../nbs/handlers/ospar.ipynb 146
class AddBodypartTempCB(Callback):
    "Add a temporary column with the body part and biological group combined."    
    def __call__(self, tfm):
        tfm.dfs['BIOTA']['body_part_temp'] = (
            tfm.dfs['BIOTA']['body part'] + ' ' + 
            tfm.dfs['BIOTA']['biological group']
            ).str.strip().str.lower()                                 

# %% ../../nbs/handlers/ospar.ipynb 153
fixes_biota_tissues = {
    'whole seaweed' : 'Whole plant',
    'flesh fish': 'Flesh with bones', # We assume it as the category 'Flesh with bones' also exists
    'flesh fish' : 'Flesh with bones',
    'unknown fish' : NA,
    'unknown fish' : NA,
    'cod medallion fish' : NA, # TO BE DETERMINED
    'mix of muscle and whole fish without liver fish' : NA, # TO BE DETERMINED
    'whole without head fish' : NA, # TO BE DETERMINED
    'flesh without bones seaweed' : NA, # TO BE DETERMINED
    'tail and claws fish' : NA # TO BE DETERMINED
}

# %% ../../nbs/handlers/ospar.ipynb 157
lut_bodyparts = lambda: Remapper(provider_lut_df=get_unique_across_dfs(tfm.dfs, col_name='body_part_temp', as_df=True),
                               maris_lut_fn=bodyparts_lut_path,
                               maris_col_id='bodypar_id',
                               maris_col_name='bodypar',
                               provider_col_to_match='value',
                               provider_col_key='value',
                               fname_cache='tissues_ospar.pkl'
                               ).generate_lookup_table(fixes=fixes_biota_tissues, as_df=False, overwrite=False)

# %% ../../nbs/handlers/ospar.ipynb 162
lut_biogroup_from_biota = lambda: get_lut(src_dir=species_lut_path().parent, fname=species_lut_path().name, 
                               key='species_id', value='biogroup_id')

# %% ../../nbs/handlers/ospar.ipynb 166
class AddSampleIdCB(Callback):
    "Include a SMP_ID column from the `ID` column of OSPAR"
    def __call__(self, tfm):
        for grp, df in tfm.dfs.items():
            if 'id' in df.columns:
                df['SMP_ID'] = df['id']                
                # Check that the ID is an integer or float.
                if not pd.api.types.is_numeric_dtype(df['SMP_ID']):
                    print(f"Non-numeric values detected in 'ID' column of dataframe '{grp}':")
                    print(f"Data type: {df['ID'].dtype}")
                    print("Unique values:", df['ID'].unique())

# %% ../../nbs/handlers/ospar.ipynb 170
class AddDepthCB(Callback):
    "Ensure depth values are floats and add 'SMP_DEPTH' columns."
    def __call__(self, tfm: Transformer):
        for grp, df in tfm.dfs.items():
            if grp == 'SEAWATER':
                if 'sampling depth' in df.columns:
                    df['SMP_DEPTH'] = df['sampling depth'].astype(float)

# %% ../../nbs/handlers/ospar.ipynb 174
class ConvertLonLatCB(Callback):
    """Convert Coordinates to decimal degrees (DDD.DDDDD°)."""
    def __init__(self):
        fc.store_attr()

    def __call__(self, tfm: 'Transformer'):
        for grp, df in tfm.dfs.items():
            df['LAT'] = self._convert_latitude(df)
            df['LON'] = self._convert_longitude(df)

    def _convert_latitude(self, df: pd.DataFrame) -> pd.Series:
        return np.where(
            df['latdir'].isin(['S']),
            self._dms_to_decimal(df['latd'], df['latm'], df['lats']) * -1,
            self._dms_to_decimal(df['latd'], df['latm'], df['lats'])
        )

    def _convert_longitude(self, df: pd.DataFrame) -> pd.Series:
        return np.where(
            df['longdir'].isin(['W']),
            self._dms_to_decimal(df['longd'], df['longm'], df['longs']) * -1,
            self._dms_to_decimal(df['longd'], df['longm'], df['longs'])
        )

    def _dms_to_decimal(self, degrees: pd.Series, minutes: pd.Series, seconds: pd.Series) -> pd.Series:
        return degrees + minutes / 60 + seconds / 3600


# %% ../../nbs/handlers/ospar.ipynb 184
kw = ['oceanography', 'Earth Science > Oceans > Ocean Chemistry> Radionuclides',
      'Earth Science > Human Dimensions > Environmental Impacts > Nuclear Radiation Exposure',
      'Earth Science > Oceans > Ocean Chemistry > Ocean Tracers, Earth Science > Oceans > Marine Sediments',
      'Earth Science > Oceans > Ocean Chemistry, Earth Science > Oceans > Sea Ice > Isotopes',
      'Earth Science > Oceans > Water Quality > Ocean Contaminants',
      'Earth Science > Biological Classification > Animals/Vertebrates > Fish',
      'Earth Science > Biosphere > Ecosystems > Marine Ecosystems',
      'Earth Science > Biological Classification > Animals/Invertebrates > Mollusks',
      'Earth Science > Biological Classification > Animals/Invertebrates > Arthropods > Crustaceans',
      'Earth Science > Biological Classification > Plants > Macroalgae (Seaweeds)']


# %% ../../nbs/handlers/ospar.ipynb 185
def get_attrs(
    tfm: Transformer, # Transformer object
    zotero_key: str, # Zotero dataset record key
    kw: list = kw # List of keywords
    ) -> dict: # Global attributes
    "Retrieve all global attributes."
    return GlobAttrsFeeder(tfm.dfs, cbs=[
        BboxCB(),
        DepthRangeCB(),
        TimeRangeCB(),
        ZoteroCB(zotero_key, cfg=cfg()),
        KeyValuePairCB('keywords', ', '.join(kw)),
        KeyValuePairCB('publisher_postprocess_logs', ', '.join(tfm.logs))
        ])()

# %% ../../nbs/handlers/ospar.ipynb 188
def encode(
    fname_out_nc: str, # Output file name
    **kwargs # Additional arguments
    ) -> None:
    "Encode data to NetCDF."
    dfs = load_data(src_dir, use_cache=True)
    tfm = Transformer(dfs, cbs=[
                            RemoveAllNAValuesCB(nan_cols_to_check),
                            LowerStripNameCB(col_src='nuclide', col_dst='nuclide'),
                            RemapNuclideNameCB(lut_nuclides, col_name='nuclide'),
                            ParseTimeCB(),
                            EncodeTimeCB(),
                            SanitizeValueCB(),
                            NormalizeUncCB(),
                            RemapUnitCB(renaming_unit_rules),
                            RemapDetectionLimitCB(coi_dl, lut_dl),
                            RemapCB(fn_lut=lut_biota, col_remap='SPECIES', col_src='species', dest_grps='BIOTA'),    
                            RemapCB(fn_lut=lut_biota_enhanced, col_remap='enhanced_species', col_src='biological group', dest_grps='BIOTA'),    
                            EnhanceSpeciesCB(),
                            AddBodypartTempCB(),
                            RemapCB(fn_lut=lut_bodyparts, col_remap='BODY_PART', col_src='body_part_temp' , dest_grps='BIOTA'),
                            AddSampleIdCB(),
                            AddDepthCB(),    
                            ConvertLonLatCB(),
                            SanitizeLonLatCB(),
                                ])
    tfm()
    encoder = NetCDFEncoder(tfm.dfs, 
                            dest_fname=fname_out_nc, 
                            global_attrs=get_attrs(tfm, zotero_key=zotero_key, kw=kw),
                            verbose=kwargs.get('verbose', False),
                           )
    encoder.encode()
