project.py

from __future__ import annotations
from typing import List, Dict, Tuple, Union
import sys
from pathlib import Path
import numpy as np
import pandas as pd
from itertools import product
from collections import defaultdict
import fnmatch
from pprint import pprint

hytraits_path = (Path(__file__).parent.parent/'hytraits').resolve()
if str(hytraits_path) not in sys.path:
    sys.path.append(str(hytraits_path))
import hytraits as H 

from utils import (get_paths, 
                   get_seed)

DatasetAndSplits = Tuple[H.CSVDataset, H.CSVDatasetSplits]


def get_traits() -> List[str]:
    '''
    List of traits.

    Return: List[str]
    '''
    traits  = ['Perc_Hemiceullose', 
               'Percent_Cellulose',
               'Percent_Lignin_recalcitrants']
    
    return sorted(traits)


def make_compatible_csvs() -> None:
    '''
    Makes compatible CSVs.
    One per trait.
    '''
    PATHS = get_paths()
    data_dir = PATHS['original']
    comp_dir = PATHS['compatible']

    # read CSV 
    csv_file = data_dir/'input_ankom_ground_all.csv'
    df = pd.read_csv(csv_file)

    # dry_bulk_code --> sample_id
    df = df.rename(columns={'TreeSN': 'sample_id'})

    # drop spectral_replicate
    df = df.drop(columns=['Unnamed: 0'])

    # rename wave cols
    wave_cols_remap = {}
    comp_wave_cols = []
    for c in df.columns:
        if c[:2] == 'X_':
            comp_wave_col = f'X_{float(c[2:]):4.3f}'
            wave_cols_remap[c] = comp_wave_col
            comp_wave_cols.append(comp_wave_col)
    df = df.rename(columns=wave_cols_remap)

    # add unique_id
    parts = [df['sample_id'], 
             df['Date'], 
             np.arange(df.shape[0])] 
    unique_ids = [f'{a}:{b}:{c}' for (a, b, c) in zip(*parts)]
    df['unique_id'] = unique_ids
    
    # one compatible CSV per trait
    print('Creating compatible CSVs ...')
    comp_dir.mkdir(parents=True, exist_ok=True)
    TRAITS = get_traits()
    for trait in TRAITS:
        cols = ['sample_id', 'unique_id', trait] + comp_wave_cols 
        trait_df = df[cols]
        trait_df = trait_df.rename(columns={trait:'y_true'})
        trait_df = trait_df.dropna()

        fname = f'{trait}.csv'
        trait_df.to_csv(comp_dir/fname, index=False)
        print(f'{fname}, {trait_df.shape}, {trait_df["sample_id"].nunique()}')
    print('------------------------------------\n')

    
def get_xtransforms() -> Dict:
    '''
    xtransforms to use. 

    Return: Dict
            {transform_key: List[H.BaseTransform]}        
    '''
    d = {'d0': [],
         'd1': [H.SavitzkyGolay(deriv=1)],
         'pa': [H.PseudoAbsorbance()],
         'f1': [H.KeepWavelengths(keep_ranges=[(399.99, 2450.01)])],
         'f2': [H.KeepWavelengths(keep_ranges=[(399.99, 2400.01)])],
         'sw': [H.KeepWavelengths(keep_ranges=[(1399.99, 2400.01)])],
         'uv': [H.UnitVectorize()]} 

    xtransforms = {}
    for (r, p) in product(['f1', 'f2', 'sw'],
                          ['d0', 'd1', 'pa']):
        xforms = d[p] + d[r]
        if p in ['d0']:
            xforms += d['uv']
            
        xtransforms[f'{p}-{r}'] = xforms  
    
    return xtransforms


def get_ytransforms(trait_key: str) -> Dict:
    '''
    ytransforms to use.

    `trait_key`: str
                 Trait key (name of compatible CSV)

    Return: List[H.BaseTransform] 
            {transform_key: List[H.BaseTransform]} 
    '''
    return {}


def get_model_names_(trait: str) -> List[str]:
    '''
    List of model names for crop-trait.

    `trait`: str
             The trait.

    Return: List[str]
            Model names. 
            <trait>__<xtransform>_<ytransform>__<avgrep>__<final>
    '''
    PATHS = get_paths()
    comp_dir = PATHS['compatible']
    
    model_names = []

    xts = sorted(list(get_xtransforms().keys()))
    avgreps = sorted(['rep'])
    selections = sorted(['bpl', 'boa'])
    for (xt, ar, se) in product(xts, avgreps, selections):
        model_names.append(f'{trait}__{xt}__{ar}__{se}')

    model_names = sorted(set(model_names))
    return model_names


def get_model_names(pattern: str = '*') -> List[str]:
    '''
    Model names that mach pattern.

    `pattern`: str
               Pattern to match.

    Return: List[str]
            List of pattern matching model names.
            <trait>__<xtransform>__<avgrep>__<final>
    '''
    TRAITS = sorted(get_traits())
    
    all_model_names, n_models = [], []
    for trait in TRAITS:
        trait_model_names = get_model_names_(trait=trait)
        all_model_names += trait_model_names
        n_models.append(len(trait_model_names))

    print('-----------------------------------')
    df = pd.DataFrame({'trait': TRAITS,
                       'n_models': n_models})
    pprint(df)
    print(f'{len(all_model_names) = }')
    print('-----------------------------------') 
         
    return fnmatch.filter(all_model_names, pattern)


def get_train_data(model_name: str,
                   args: Dict,
                   ideploy: bool) -> DatasetAndSplits:
    '''
    Get dataset and splits for model training.
    <trait>__<xtransform>__<avgrep>__<final>
    
    `model_name`: str
                  Name of model.
    `args`: Dict
            CLI args, used for splits details.
    `ideploy`: bool
               Train stage or internal evaluation stage.
               If train: ytransforms as needed, here it is specified
               If ideploy: ytransforms is always [], subsample=-1, reduce=none.

    Return: DatasetAndSplits
            Tuple[CSVDataset, CSVDatasetSplits]
    '''
    parts = model_name.split('__')
    trait = parts[0]
    xt_key = parts[1]
    avgrep = parts[2]
    final = parts[3]
    
    xtransforms = get_xtransforms()[xt_key]
    if ideploy:
        subsample = -1
        reduce = 'none'
        ytransforms = [] # this is always true
    else:
        subsample = 1 if avgrep == 'rep' else -1
        reduce = 'none' if avgrep == 'rep' else 'mean'
        ytransforms = list(get_ytransforms(trait_key=trait).values()) 

    PATHS = get_paths()
    comp_dir = PATHS['compatible']
    csvs = [comp_dir/f'{trait}.csv']
    dataset = H.CSVDataset(csv_file=csvs,
                           xtransforms=xtransforms,
                           ytransforms=ytransforms,
                           subsample=subsample,
                           reduce=reduce,
                           seed=get_seed()) 
    splitter = H.CSVDatasetSplitter(n_repeats=args['n_repeats'],
                                    split_types=args['split_types'],
                                    split_percents=args['split_percents'],
                                    seed=get_seed())
    splits = splitter(dataset.sids)

    return (dataset, splits)
    
    
def train_model(model_name: str,
                args: Dict) -> None:
    '''
    Trains the specified model.

    `model_name`: str
                  Model name.
    `args`: Dict
            CLI args.
    '''
    PATHS = get_paths()
    model_dir = PATHS['model']/model_name
    deploy_dir = PATHS['deploy']/model_name

    # train model
    print(f'Training: {model_name} ...')
    (dataset, splits) = get_train_data(model_name=model_name,
                                       args=args,
                                       ideploy=False)
    
    # pprint(dataset.xtransforms)
    # pprint(dataset.ytransforms)
    # print(f'{len(dataset) = }')
    # print(f'{dataset.subsample = }, {dataset.reduce = }')
    # print(f'{splits.n_outers = }, {splits.n_inners = }')
    # print('---------------------------')
    
    n_comps = args['n_components']
    n_train = len(splits.split(0, 0, 'train'))
    n_comps = min(n_comps, n_train - 2)
    n_components = np.arange(1, (n_comps + 1)).tolist()
    final_key = model_name.split('__')[-1]
    H.plsr_train(dataset=dataset,
                 splits=splits,
                 save_dir=model_dir,
                 n_components=n_components,
                 final=final_key)

    
    # internal evaluation
    print(f'Evaluation: {model_name} ...')
    (dataset, splits) = get_train_data(model_name=model_name,
                                       args=args,
                                       ideploy=True)
    metrics = [H.RMSE(),
               H.RangeNormalizedRMSE(),
               H.InterquartileNormalizedRMSE(),
               H.R2(),
               H.FittedR2()]
    model = H.plsr_load_model(model_dir/'model.npz')
    undos = H.load_transforms(model_dir/'ytransforms.json',
                              undo=True)
    
    # pprint(dataset.xtransforms)
    # pprint(dataset.ytransforms)
    # pprint(undos)
    # for t in undos:
    #     print(t.get_name_params())
    # print(f'{len(dataset) = }')
    # print(f'{dataset.subsample = }, {dataset.reduce = }')
    # print(f'{splits.n_outers = }, {splits.n_inners = }')
    # print('---------------------------')
    
    (preds_df, metrics_df) = H.deploy_csv_index_matched(model=model,
                                                        metrics=metrics,
                                                        dataset=dataset,
                                                        splits=splits,
                                                        split_label='test',
                                                        undos=undos)
    
    deploy_dir.mkdir(parents=True, exist_ok=True)
    preds_df.to_csv(deploy_dir/'preds.csv', index=False)
    metrics_df.to_csv(deploy_dir/'metrics.csv', index=False)

    mcols = [c for c in metrics_df.columns if c != 'model_idx']
    
    df_mean = metrics_df[mcols].mean().to_frame().T
    df_mean['deploy_key'] = model_name
    df_mean = df_mean[['deploy_key'] + mcols]
    df_mean.to_csv(deploy_dir/'metrics_mean.csv', index=False)

    df_median = metrics_df[mcols].median().to_frame().T
    df_median['deploy_key'] = model_name
    df_median = df_median[['deploy_key'] + mcols]
    df_median.to_csv(deploy_dir/'metrics_median.csv', index=False)

   
def deploy_model(model_name: str,
                 args: Dict) -> None:
    return