# 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']

# %% ../nbs/00_loading.ipynb 3
from pathlib import Path
from tqdm import tqdm
from typing import Union, List
import re
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 4
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 5
def download(url, dest):
    "Download `given` url into `dest` (creates it on the way if does not exist) "
    if not dest.exists(): fc.mkdir(dest, parents=True)
    return fd.download_url(url, dest)

# %% ../nbs/00_loading.ipynb 6
def load_ossl(analytes: Union[str, List[str]], # Using OSSL's analytes naming conventions
              spectra_type: str='visnir',  # Possible values: 'mir', 'visnir'
              dest:Path=Path.home() / '.lssm/data/ossl', # directory containing the data
              ):
    "Load all available OSSL data and filter it by spectra type and analytes of interest"
    
    url = PATH_OSSL_ALL_L1_V1_2
    fname = dest / Path(PATH_OSSL_ALL_L1_V1_2).name
    if not fname.exists(): 
        print('Downloading & saving to: ', str(fname))
        download(url, dest)
    
    print('Reading & selecting data ...')
    
    df = pd.read_csv(fname, compression='infer', low_memory=True)
            
    analytes = [analytes] if isinstance(analytes, str) else analytes
        
    subset = analytes + [CFGS[spectra_type]['ref_col']]
    df = df.dropna(subset=subset)
    
    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_

