# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_loading.ipynb.

# %% auto 0
__all__ = ['PATH_OSSL_ALL_L0_V1_2', 'PATH_OSSL_ALL_L1_V1_2', 'CFGS', 'download', 'load_ossl', 'load_mir_ring_trial',
           'get_spectra_pair_idxs', 'load_mir_kex_spike', 'load_nir_kex_spike']

# %% ../nbs/00_loading.ipynb 3
from pathlib import Path
from tqdm import tqdm
from typing import Union, List
import re
import itertools
import fastdownload as fd
import fastcore.all as fc

import pandas as pd
import numpy as np
from sklearn.preprocessing import LabelEncoder

import warnings
warnings.filterwarnings('ignore')

# %% ../nbs/00_loading.ipynb 6
PATH_OSSL_ALL_L0_V1_2 = 'https://storage.googleapis.com/soilspec4gg-public/ossl_all_L0_v1.2.csv.gz'
PATH_OSSL_ALL_L1_V1_2 = 'https://storage.googleapis.com/soilspec4gg-public/ossl_all_L0_v1.2.csv.gz'

CFGS = {
    'visnir': {'ref_col': 'scan_visnir.1500_ref', 'range': [400, 2500]},
    'mir': {'ref_col': 'scan_mir.1500_abs', 'range': [600, 4000]}
}

# %% ../nbs/00_loading.ipynb 7
def download(url: str,  # url to dowload data from
             dest_dir: str,  # directory to download data to
             ) -> None:
    "Download data available at `url` into the `dest` directory (creates it on the way if does not exist)."
    if not dest_dir.exists():
        fc.mkdir(dest_dir, parents=True)
    return fd.download_url(url, dest_dir)

# %% ../nbs/00_loading.ipynb 8
def load_ossl(analytes: Union[str, List[str]],  # Using OSSL's analytes naming conventions
              spectra_type: str = 'visnir',  # Possible values: 'mir', 'visnir'
              src: Path = Path.home() / '.lssm/data/ossl',  # directory containing the data
              debug:bool=False # return unprocessed loaded data directly for further investigation
              ):
    "Load all available OSSL data and filter it by spectra type and analytes of interest"

    url = PATH_OSSL_ALL_L1_V1_2
    fname = src / Path(PATH_OSSL_ALL_L1_V1_2).name
    if not fname.exists():
        print('Downloading & saving to: ', str(fname))
        download(url, src)

    print('Reading & selecting data ...')

    df = pd.read_csv(fname, compression='infer', low_memory=True)
    
    if debug: return df
    
    analytes = [analytes] if isinstance(analytes, str) else analytes

    subset = analytes + [CFGS[spectra_type]['ref_col']]
    df = df.dropna(subset=subset, how='any')

    cols_ref = [name for name in df.columns if f'scan_{spectra_type}.' in name]
    X = df[cols_ref].values

    y = df[analytes].values
    smp_idx = df['id.layer_uuid_txt'].values

    ds_name_encoder = LabelEncoder()
    ds_name = ds_name_encoder.fit_transform(df['dataset.code_ascii_txt'])

    pattern = r"scan_{}\.(\d+)_".format(spectra_type)
    X_names = np.array([int(re.search(pattern, name).group(1)) for name in df.columns
                        if re.search(pattern, name)])

    lower_limit, upper_limit = CFGS[spectra_type]['range']
    idxs = np.where((X_names >= lower_limit) & (X_names <= upper_limit))[0]

    return X[:, idxs], y, X_names[idxs], smp_idx, ds_name, ds_name_encoder.classes_


# %% ../nbs/00_loading.ipynb 15
def load_mir_ring_trial(fname): return pd.read_csv(fname)

# %% ../nbs/00_loading.ipynb 18
def get_spectra_pair_idxs(df):
    """
    Retrieve index pairs of replicated spectra, representing measurements taken from the same soil sample
    but using different instruments.
    """
    pair_idxs = []
    for smp_id in df.sample_id.unique():
        indices = df[df.sample_id == smp_id].index
        pair_idxs.extend(list(itertools.product(indices, repeat=2)))
    return pair_idxs


# %% ../nbs/00_loading.ipynb 23
def load_mir_kex_spike(src_dir):
    """
    Load MIR spectra of K spiked soil samples.
    
    Parameters:
    src_dir (Path-like object): Directory containing the spectra files.
    
    Returns:
    tuple: Tuple containing the array of absorbance values, 
           array of wavenumbers (columns), and array of sample names (rows).
    """
    pattern = r'-\d-\d$'
    fnames = [f for f in src_dir.ls() if re.search(pattern, f.stem)]
    
    dfs = [pd.read_csv(fname, header=None, names=['wavenumber', 'absorbance'])
           .query('649 < wavenumber < 4000')
           .assign(name=fname.stem) for fname in fnames]
    
    df_combined = pd.concat(dfs).pivot_table(values='absorbance', index='name', columns='wavenumber')
    
    return df_combined.values, df_combined.columns.values, df_combined.index.values.astype('U')


# %% ../nbs/00_loading.ipynb 28
def load_nir_kex_spike(fname):
    """
    Load NIR spectra of K spiked soil samples.
    
    Parameters:
    fname (str or Path-like object): File name or path of the Excel file.
    
    Returns:
    tuple: Tuple containing the array of spectral values, 
           array of wavenumbers (columns), and array of sample names (rows).
    """
    df = pd.read_excel(fname, sheet_name='Results', index_col='Sample ID')
    df.index.name = 'name'
    df.columns.name = 'wavenumber'
    
    return df.values, df.columns.values, df.index.values.astype('U')
