"""Data pipeline (handler) to convert TEPCO dataset ([Source](https://radioactivity.nsr.go.jp/ja/list/349/list-1.html)) to `NetCDF` format"""

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

# %% auto 0
__all__ = ['fname_coastal_water', 'fname_clos1F', 'fname_iaea_orbs', 'fname_out', 'common_coi', 'nuclides_pattern',
           'unit_mapping', 'nuclide_mapping', 'kw', 'find_location_section', 'fix_sampling_time',
           'get_coastal_water_df', 'get_locs_coastal_water', 'get_clos1F_df', 'get_locs_clos1F', 'get_locs_orbs',
           'concat_locs', 'align_dfs', 'concat_dfs', 'georef_data', 'load_data', 'RemoveJapanaseCharCB',
           'FixRangeValueStringCB', 'SelectColsOfInterestCB', 'WideToLongCB', 'extract_nuclide', 'ExtractNuclideNameCB',
           'ExtractUnitCB', 'ExtractValueTypeCB', 'LongToWideCB', 'RemapUnitNameCB', 'RemapNuclideNameCB',
           'RemapVALUE_DL_DLV_CB', 'ParseTimeCB', 'get_attrs', 'encode']

# %% ../../nbs/handlers/tepco.ipynb 3
import warnings
warnings.filterwarnings('ignore')

# %% ../../nbs/handlers/tepco.ipynb 4
import pandas as pd
import re
import numpy as np
import fastcore.all as fc
from tqdm import tqdm
from collections import defaultdict

from marisco.callbacks import (
    Callback, 
    Transformer,
    EncodeTimeCB, 
    SanitizeLonLatCB,
    EncodeTimeCB, 
    )

from ..configs import cfg
from ..encoders import NetCDFEncoder

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

from ..netcdf2csv import decode

# %% ../../nbs/handlers/tepco.ipynb 6
fname_coastal_water = 'https://radioactivity.nra.go.jp/cont/en/results/sea/coastal_water.csv'
fname_clos1F = 'https://radioactivity.nra.go.jp/cont/en/results/sea/close1F_water.xlsx'
fname_iaea_orbs = 'https://raw.githubusercontent.com/RML-IAEA/iaea.orbs/refs/heads/main/src/iaea/orbs/stations/station_points.csv'

fname_out = '../../_data/output/tepco.nc'

# %% ../../nbs/handlers/tepco.ipynb 10
def find_location_section(df, 
                          col_idx=0,
                          pattern='Sampling point number'
                          ):
    "Find the line number where location data begins."
    mask = df.iloc[:, col_idx] == pattern
    indices = df[mask].index
    return indices[0] if len(indices) > 0 else -1

# %% ../../nbs/handlers/tepco.ipynb 13
def fix_sampling_time(x):
    if pd.isna(x): 
        return '00:00:00'
    else:
        hour, min =  x.split(':')[:2]
        return f"{hour if len(hour) == 2 else '0' + hour}:{min}:00"

# %% ../../nbs/handlers/tepco.ipynb 14
def get_coastal_water_df(fname_coastal_water):
    "Get the measurements dataframe from the `coastal_water.csv` file."
    
    locs_idx = find_location_section(pd.read_csv(fname_coastal_water, 
                                      skiprows=0, low_memory=False))
    
    df = pd.read_csv(fname_coastal_water, skiprows=1, 
                     nrows=locs_idx - 1,
                     low_memory=False)
    df.dropna(subset=['Sampling point number'], inplace=True)
    df['Sampling time'] = df['Sampling time'].map(fix_sampling_time)
    
    df['TIME'] = df['Sampling date'].replace('-', '/') + ' ' + df['Sampling time']
    
    df = df.drop(columns=['Sampling date', 'Sampling time'])
    return df

# %% ../../nbs/handlers/tepco.ipynb 22
def get_locs_coastal_water(fname_coastal_water):
    locs_idx = find_location_section(pd.read_csv(fname_coastal_water, 
                                      skiprows=0, low_memory=False))
    
    df = pd.read_csv(fname_coastal_water, skiprows=locs_idx+1, 
                     low_memory=False).iloc[:, :3]
    
    df.columns = ['STATION', 'LON', 'LAT']
    df.dropna(subset=['LAT'], inplace=True)
    df['org'] = 'coastal_seawater.csv'
    return df

# %% ../../nbs/handlers/tepco.ipynb 26
def get_clos1F_df(fname_clos1F):
    "Get measurements dataframe from close1F_water.xlsx file and parse datetime."
    excel_file = pd.ExcelFile(fname_clos1F)
    dfs = {}
    
    for sheet_name in tqdm(excel_file.sheet_names):
        locs_idx = find_location_section(pd.read_excel(excel_file, 
                                                       sheet_name=sheet_name,
                                                       skiprows=1))
        df = pd.read_excel(excel_file, 
                   sheet_name=sheet_name, 
                   skiprows=1,
                   nrows=locs_idx-1)
        
        df.dropna(subset=['Sampling point number'], inplace=True)
        df['Sampling date'] = df['Sampling date']\
            .astype(str)\
            .apply(lambda x: x.split(' ')[0]\
            .replace('-', '/'))
            
        dfs[sheet_name] = df
    
    df = pd.concat(dfs.values(), ignore_index=True)
    df.dropna(subset=['Sampling date'], inplace=True)
    df['TIME'] = df['Sampling date'] + ' ' + df['Sampling time'].astype(str)
    df = df.drop(columns=['Sampling date', 'Sampling time'])
    return df

# %% ../../nbs/handlers/tepco.ipynb 29
def get_locs_clos1F(fname_clos1F):
    "Get locations dataframe from close1F_water.xlsx file from each sheets."
    excel_file = pd.ExcelFile(fname_clos1F)
    dfs = {}
    
    for sheet_name in tqdm(excel_file.sheet_names):
        locs_idx = find_location_section(pd.read_excel(excel_file, 
                                                       sheet_name=sheet_name,
                                                       skiprows=1))
        df = pd.read_excel(excel_file, 
                           sheet_name=sheet_name, 
                           skiprows=locs_idx+2)
            
        dfs[sheet_name] = df
    
    df = pd.concat(dfs.values(), ignore_index=True).iloc[:, :3]
    df.dropna(subset=['Sampling coordinate North latitude (Decimal)'], inplace=True)    
    df.columns = ['STATION', 'LON', 'LAT']
    df['org'] = 'close1F.csv'
    return df

# %% ../../nbs/handlers/tepco.ipynb 34
def get_locs_orbs(fname_iaea_orbs):
    df = pd.read_csv(fname_iaea_orbs)
    df.columns = ['org', 'STATION', 'LON', 'LAT']
    return df

# %% ../../nbs/handlers/tepco.ipynb 37
def concat_locs(dfs):
    "Concatenate and drop duplicates from coastal_seawater.csv and iaea_orbs.csv (kept)"
    df = pd.concat(dfs)
    # Group by org to be used for sorting
    df['org_grp'] = df['org'].apply(
        lambda x: 1 if x == 'coastal_seawater.csv' else 2 if x == 'close1F.csv' else 0)
    df.sort_values('org_grp', ascending=True, inplace=True)
    # Drop duplicates and keep orbs data first
    df.drop_duplicates(subset='STATION', keep='first', inplace=True)
    df.drop(columns=['org_grp'], inplace=True)
    df.sort_values('STATION', ascending=True, inplace=True)
    return df

# %% ../../nbs/handlers/tepco.ipynb 39
def align_dfs(df_from, df_to):
    "Align columns structure of df_from to df_to."
    df = defaultdict()    
    for c in df_to.columns:
        df[c] = df_from[c].values if c in df_from.columns else np.nan
    return pd.DataFrame(df)

# %% ../../nbs/handlers/tepco.ipynb 41
def concat_dfs(df_coastal_water, df_clos1F):
    "Concatenate and drop duplicates from coastal_seawater.csv and close1F_water.xlsx (kept)"
    df_clos1F = align_dfs(df_clos1F, df_coastal_water)
    df = pd.concat([df_coastal_water, df_clos1F])
    return df

# %% ../../nbs/handlers/tepco.ipynb 43
def georef_data(df_meas, df_locs):
    "Georeference measurements dataframe using locations dataframe."
    assert "Sampling point number" in df_meas.columns and "STATION" in df_locs.columns
    return pd.merge(df_meas, df_locs, how="inner", 
                    left_on='Sampling point number', right_on='STATION')

# %% ../../nbs/handlers/tepco.ipynb 45
def load_data(fname_coastal_water, fname_clos1F, fname_iaea_orbs):
    "Load, align and georeference TEPCO data"
    df_locs = concat_locs(
        [get_locs_coastal_water(fname_coastal_water), 
         get_locs_clos1F(fname_clos1F),
         get_locs_orbs(fname_iaea_orbs)])
    df_meas = concat_dfs(get_coastal_water_df(fname_coastal_water), get_clos1F_df(fname_clos1F))
    df_meas.dropna(subset=['Sampling point number'], inplace=True)
    return {'SEAWATER': georef_data(df_meas, df_locs)}

# %% ../../nbs/handlers/tepco.ipynb 53
class RemoveJapanaseCharCB(Callback):
    "Remove 約 (about) char"
    def _transform_if_about(self, value, about_char='約'):
        if pd.isna(value): return value
        return (value.replace(about_char, '') if str(value).count(about_char) != 0 
                else value)
    
    def __call__(self, tfm): 
        for k in tfm.dfs.keys():
            cols_rdn = [c for c in tfm.dfs[k].columns if ('(Bq/L)' in c) and (tfm.dfs[k][c].dtype == 'object')]
            tfm.dfs[k][cols_rdn] = tfm.dfs[k][cols_rdn].map(self._transform_if_about)

# %% ../../nbs/handlers/tepco.ipynb 57
class FixRangeValueStringCB(Callback):
    "Replace range values (e.g '4.0E+00<&<8.0E+00' or '1.0～2.7') by their mean"
    
    def _extract_and_calculate_mean(self, s):
        # For scientific notation ranges
        float_strings = re.findall(r"[+-]?\d+\.?\d*E?[+-]?\d*", s)
        if float_strings:
            float_numbers = np.array(float_strings, dtype=float)
            return float_numbers.mean()
        return s
    
    def _transform_if_range(self, value):
        if pd.isna(value): 
            return value
        value = str(value)
        # Check for both range patterns
        if '<&<' in value or '～' in value:
            return self._extract_and_calculate_mean(value)
        return value

    def __call__(self, tfm): 
        for k in tfm.dfs.keys():
            cols_rdn = [c for c in tfm.dfs[k].columns 
                       if ('(Bq/L)' in c) and (tfm.dfs[k][c].dtype == 'object')]
            # tfm.dfs[k][cols_rdn] = tfm.dfs[k][cols_rdn].map(self._transform_if_range).astype(float)
            tfm.dfs[k][cols_rdn] = tfm.dfs[k][cols_rdn].map(self._transform_if_range)

# %% ../../nbs/handlers/tepco.ipynb 61
common_coi = ['LON', 'LAT', 'TIME', 'STATION']
nuclides_pattern = '(Bq/L)'

# %% ../../nbs/handlers/tepco.ipynb 62
class SelectColsOfInterestCB(Callback):
    "Select columns of interest."
    def __init__(self, common_coi, nuclides_pattern): fc.store_attr()
    def __call__(self, tfm):
        nuc_of_interest = [c for c in tfm.dfs['SEAWATER'].columns if nuclides_pattern in c]
        tfm.dfs['SEAWATER'] = tfm.dfs['SEAWATER'][self.common_coi + nuc_of_interest]

# %% ../../nbs/handlers/tepco.ipynb 65
class WideToLongCB(Callback):
    """
    Get TEPCO nuclide names as values not column names 
    to extract contained information (nuclide name, unc, dl, ...).
    """
    def __init__(self, id_vars=['LON', 'LAT', 'TIME', 'STATION']): 
        fc.store_attr()
        
    def __call__(self, tfm): 
        tfm.dfs['SEAWATER'] = pd.melt(tfm.dfs['SEAWATER'], id_vars=self.id_vars)
#| eval: false

# %% ../../nbs/handlers/tepco.ipynb 70
def extract_nuclide(text: str) -> str:
    "Extract the nuclide identifier from a measurement variable name using regex."
    pattern = r'^(Total\s+(?:alpha|beta)|[^\s]+)'
    match = re.match(pattern, text, re.IGNORECASE)
    return match.group(1) if match else text 

# %% ../../nbs/handlers/tepco.ipynb 73
class ExtractNuclideNameCB(Callback):
    "Extract nuclide name from TEPCO data."
    def __init__(self, src_col='variable', dest_col='NUCLIDE'): fc.store_attr()
    def __call__(self, tfm): 
        tfm.dfs['SEAWATER'][self.dest_col] = tfm.dfs['SEAWATER'][self.src_col].map(extract_nuclide)

# %% ../../nbs/handlers/tepco.ipynb 76
class ExtractUnitCB(Callback):
    "Extract unit from TEPCO data."
    def __init__(self, src_col='variable', dest_col='UNIT'): fc.store_attr()
    def __call__(self, tfm): 
        tfm.dfs['SEAWATER'][self.dest_col] = tfm.dfs['SEAWATER'][self.src_col].str.extract(r'\((.*?)\)')

# %% ../../nbs/handlers/tepco.ipynb 79
class ExtractValueTypeCB(Callback):
    "Extract value type from TEPCO data."
    def __init__(self, src_col='variable', dest_col='type'): fc.store_attr()
    def __call__(self, tfm): 
        tfm.dfs['SEAWATER'][self.dest_col] = np.select(
            [
                tfm.dfs['SEAWATER'][self.src_col].str.contains('detection limit', case=False),
                tfm.dfs['SEAWATER'][self.src_col].str.contains('statistical error', case=False)],
            ['DL', 'UNC'],
            default='VALUE'
        )

# %% ../../nbs/handlers/tepco.ipynb 82
class LongToWideCB(Callback):
    "Reshape: long to wide"
    def __init__(self, src_col='variable', dest_col='type'): fc.store_attr()
    def __call__(self, tfm): 
        tfm.dfs['SEAWATER'] = pd.pivot_table(
            tfm.dfs['SEAWATER'],
            values='value',
            index=['LON', 'LAT', 'TIME', 'STATION', 'NUCLIDE', 'UNIT'],
            columns='type',
            aggfunc='first'
        ).reset_index()
        tfm.dfs['SEAWATER'].reset_index(inplace=True)
        tfm.dfs['SEAWATER'].rename(columns={'index': 'SMP_ID'}, inplace=True)

# %% ../../nbs/handlers/tepco.ipynb 88
unit_mapping = {'Bq/L': 3}

# %% ../../nbs/handlers/tepco.ipynb 89
class RemapUnitNameCB(Callback):
    """
    Remap `UNIT` name to MARIS id.
    """
    def __init__(self, unit_mapping): fc.store_attr()
    def __call__(self, tfm):
        tfm.dfs['SEAWATER']['UNIT'] = tfm.dfs['SEAWATER']['UNIT'].map(self.unit_mapping)


# %% ../../nbs/handlers/tepco.ipynb 92
nuclide_mapping = {
    '131I': 29,
    '134Cs': 31,
    '137Cs': 33,
    '125Sb': 24,
    'Total beta': 103,
    '238Pu': 67,
    '239Pu+240Pu': 77,
    '3H': 1,
    '89Sr': 11,
    '90Sr': 12,
    'Total alpha': 104,
    '132I': 100,
    '136Cs': 102,
    '58Co': 8,
    '105Ru': 97,
    '106Ru': 17,
    '140La': 35,
    '140Ba': 34,
    '132Te': 99,
    '60Co': 9,
    '144Ce': 37,
    '54Mn': 6
}

# %% ../../nbs/handlers/tepco.ipynb 93
class RemapNuclideNameCB(Callback):
    "Remap `NUCLIDE` name to MARIS id."
    def __init__(self, nuclide_mapping): fc.store_attr()
    def __call__(self, tfm):
        tfm.dfs['SEAWATER']['NUCLIDE'] = tfm.dfs['SEAWATER']['NUCLIDE'].map(self.nuclide_mapping)

# %% ../../nbs/handlers/tepco.ipynb 99
class RemapVALUE_DL_DLV_CB(Callback):
    "Remap `DL`, `DLV`, `VALUE` based on TEPCO -> MARIS rules."    
    def map_all_columns(self, row):
        """Map all three columns (VALUE, DL, DLV) at once based on TEPCO rules"""
        value, dl = row['VALUE'], row['DL']
        new_value, new_dlv, new_dl = value, dl, 1
        
        if value == 'ND':
            if pd.isna(dl):
                new_value, new_dlv, new_dl = np.nan, np.nan, 3
            else:
                new_value, new_dlv, new_dl = dl, dl, 2
                
        elif pd.isna(value):
            if pd.isna(dl):
                new_value, new_dlv, new_dl = np.nan, np.nan, np.nan
            else:
                new_value, new_dlv, new_dl = dl, dl, 2
                
        return pd.Series({
            'VALUE': new_value,
            'DLV': new_dlv, 
            'DL': new_dl
        })
        
    def __call__(self, tfm):
        mapped = tfm.dfs['SEAWATER'].apply(self.map_all_columns, axis=1)
        tfm.dfs['SEAWATER'][['VALUE', 'DLV', 'DL']] = mapped
        tfm.dfs['SEAWATER']['DL'] = tfm.dfs['SEAWATER']['DL'].astype(int)
        tfm.dfs['SEAWATER']['VALUE'] = tfm.dfs['SEAWATER']['VALUE'].astype(float)

# %% ../../nbs/handlers/tepco.ipynb 102
class ParseTimeCB(Callback):
    "Parse time column from TEPCO."
    def __init__(self, time_name='TIME'): fc.store_attr()
    def __call__(self, tfm):
        tfm.dfs['SEAWATER'][self.time_name] = pd.to_datetime(tfm.dfs['SEAWATER'][self.time_name], 
                                                             format='%Y/%m/%d %H:%M:%S', errors='coerce')

# %% ../../nbs/handlers/tepco.ipynb 109
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/tepco.ipynb 110
def get_attrs(tfm, zotero_key, kw=kw):
    "Retrieve global attributes from MARIS dump."
    return GlobAttrsFeeder(tfm.dfs, cbs=[
        BboxCB(),
        TimeRangeCB(),
        ZoteroCB(zotero_key, cfg=cfg()),
        KeyValuePairCB('keywords', ', '.join(kw)),
        KeyValuePairCB('publisher_postprocess_logs', ', '.join(tfm.logs))
        ])()

# %% ../../nbs/handlers/tepco.ipynb 112
def encode(
    fname_out: str, # Path to the folder where the NetCDF output will be saved
    **kwargs # Additional keyword arguments
    ):
    "Encode TEPCO data to NetCDF."
    dfs = load_data(fname_coastal_water, fname_clos1F, fname_iaea_orbs)
    
    tfm = Transformer(dfs, cbs=[
        RemoveJapanaseCharCB(),
        FixRangeValueStringCB(),
        SelectColsOfInterestCB(common_coi, nuclides_pattern),
        WideToLongCB(),
        ExtractNuclideNameCB(),
        ExtractUnitCB(),
        ExtractValueTypeCB(),
        LongToWideCB(),
        RemapUnitNameCB(unit_mapping),
        RemapNuclideNameCB(nuclide_mapping),
        RemapVALUE_DL_DLV_CB(),
        ParseTimeCB(),
        EncodeTimeCB(),
        SanitizeLonLatCB()
    ])        
    tfm()
    encoder = NetCDFEncoder(tfm.dfs, 
                            dest_fname=fname_out, 
                            global_attrs=get_attrs(tfm, zotero_key='JEV6HP5A', kw=kw),
                            verbose=kwargs.get('verbose', False)
                            )
    encoder.encode()
