post.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
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
import matplotlib.patches as mpatches
from matplotlib.axes import Axes
from PIL import Image, ImageDraw, ImageFont
from collections import defaultdict
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
from project import get_traits


def create_metrics_tables() -> None:
    '''
    Collates metrics of all trained models.
    '''
    PATHS = get_paths()
    deploy_dir = PATHS['deploy']
    post_dir = PATHS['post']
    post_dir.mkdir(parents=True, exist_ok=True)
    
    mean_dfs, median_dfs = [], []
    for d in deploy_dir.glob('*'):
        if d.is_dir():
            mean_dfs.append(pd.read_csv(d/'metrics_mean.csv'))
            median_dfs.append(pd.read_csv(d/'metrics_median.csv'))
    
    mean_df = pd.concat(mean_dfs)
    mean_df = mean_df.sort_values(by='deploy_key')
    mean_df.to_csv(post_dir/'ALL_MEAN_METRICS.csv', index=False)
    
    median_df = pd.concat(median_dfs)
    median_df = median_df.sort_values(by='deploy_key')
    median_df.to_csv(post_dir/'ALL_MEDIAN_METRICS.csv', index=False)

    
def get_model_names_by_trait() -> Dict:
    '''
    Separates trained models by traits.

    Return: Dict
            {trait: [model_name]}
    '''
    PATHS = get_paths()
    model_dir = PATHS['model']
    
    MODELS = [d.stem for d in model_dir.glob('*') if d.is_dir()]
    MODELS.sort()

    trait_models = defaultdict(list)
    for m in MODELS:
        trait, _, _, _ = m.split('__')
        trait_models[trait].append(m)

    return trait_models


def get_model_colors(model_names: List[str]) -> Dict:
    '''
    Returns color for the model name.

    `model_names`: List[str]
                   List of model names

    Return: Dict
            {model_name: color_hexstring}
    '''
    color_map = {'d0-f1': '#E7724A',
                 'd0-f2': '#39A35F',
                 'd0-sw': '#61A6CF',
                 
                 'd1-f1': '#E7724A',
                 'd1-f2': '#39A35F',
                 'd1-sw': '#61A6CF',
                 
                 'pa-f1': '#E7724A',
                 'pa-f2': '#39A35F',
                 'pa-sw': '#61A6CF'}
    model_colors = {}
    for model_name in model_names:
        _, key, _, _ = model_name.split('__')
        model_colors[model_name] = color_map[key]
    return model_colors


def composite_coefficient_vip_plots_(trait: str,
                                     models: List[str],
                                     which: str) -> None:
    '''
    Create composite coefficient/VIP plots.

    `trait`: str
             Trait name
    `models`: List[str]
              Model names for `trait`.
    `which`: str
             std_coefficients/vips
    '''
    PATHS = get_paths()
    model_dir = PATHS['model']
    post_dir = PATHS['post']/trait
    post_dir.mkdir(parents=True, exist_ok=True)
    
    boas = [m for m in models if '_boa' in m]
    bpls = [m for m in models if '_bpl' in m]
    
    for finals in [boas, bpls]:
        final_key = 'best-over-all' if 'boa' in finals[0] else 'best-per-loop'
        d0s = [m for m in finals if 'd0-' in m]
        d1s = [m for m in finals if 'd1-' in m]
        pas = [m for m in finals if 'pa-' in m]

        fig, ax = plt.subplots(3, 1, figsize=(12, 9), sharex=True)
        for (i, pres) in enumerate([d0s, d1s, pas]):
            colors = get_model_colors(pres)
            for pre in pres:
                model = H.plsr_load_model(model_dir/pre/'model.npz')
                samples = model.get(which)
                mu = np.mean(samples, axis=0, keepdims=True)
                if which == 'std_coefficients':
                    mu = H.plsr_normalize_coefficients(mu)
                else:
                    mu = H.plsr_normalize_vips(mu)
                waves = model.get('wavelengths')
                wranges = model.get('kept_wave_ranges')
                label = pre.split('__')[1]
                ax[i] = H.plot_stat_vs_wavelength(ax=ax[i],
                                                  waves=waves,
                                                  wave_ranges=wranges,
                                                  samples=mu,
                                                  color=colors[pre],
                                                  show_sdev=False,
                                                  legend_label=label,
                                                  linestyle='-')

        for i in range(3):
            if which == 'std_coefficients':
                ax[i].set_ylim(-1, 1)
            else:
                ax[i].set_ylim(0, 1)
            ax[i].set_xlim(0, 2500)
            ax[i].legend(loc='upper left')
            ax[i].axhline(y=0, xmin=400/2500, color='#000000')
            ax[i].set_facecolor('#EEEEEE')

        fig.suptitle(f'{trait}:{which} (composite, {final_key})')
        plt.tight_layout()
        plt.savefig(post_dir/f'{which}__composite__{final_key}.png')
        plt.close()


def annotate_wavelengths_(ax: Axes,
                          values: np.ndarray,
                          waves: np.ndarray,
                          wave_ranges: List[Tuple[float, float]],
                          annotate: str) -> Axes:
    locs = np.array([])
    if annotate == 'peaks':
        (locs, _) = H.get_peak_wavelengths(values=values,
                                           waves=waves,
                                           wave_ranges=wave_ranges,
                                           n_peaks=16,
                                           distance=25)
    else:
        (locs, _) = H.get_trough_wavelengths(values=values,
                                             waves=waves,
                                             wave_ranges=wave_ranges,
                                             n_troughs=16,
                                             distance=25)

    if locs.size > 0:
        locs = np.sort(locs)
        ax = H.plot_wavelength_annotations(ax=ax,
                                           waves=locs[::2],
                                           v_pos=1.2,
                                           line_color='#D3D3D3',
                                           text_color='#000000',
                                           fontsize=8,
                                           show_index=False,
                                           backgroundcolor='#FFFFFF')
        ax = H.plot_wavelength_annotations(ax=ax,
                                           waves=locs[1::2],
                                           v_pos=1.2 - 0.4,
                                           line_color='#D3D3D3',
                                           text_color='#000000',
                                           fontsize=8,
                                           show_index=False,
                                           backgroundcolor='#FFFFFF')
    return ax
    
def separate_coefficient_vip_plots_(trait: str,
                                    models: str,
                                    which: str,
                                    annotate: str) -> None:
    '''
    Create composite coefficient/VIP plots.

    `trait`: str
             Trait name
    `models`: List[str]
              Model names for `trait`.
    `which`: str
             std_coefficients/vips
    `annotate`: str
                peaks/troughs
    '''
    PATHS = get_paths()
    model_dir = PATHS['model']
    post_dir = PATHS['post']/trait
    post_dir.mkdir(parents=True, exist_ok=True)
    
    boas = [m for m in models if '_boa' in m]
    bpls = [m for m in models if '_bpl' in m]
    
    for finals in [boas, bpls]:
        final_key = 'best-over-all' if 'boa' in finals[0] else 'best-per-loop'
        
        d0s = [m for m in finals if 'd0-' in m]
        d1s = [m for m in finals if 'd1-' in m]
        pas = [m for m in finals if 'pa-' in m]

        for pres in [d0s, d1s, pas]:
            pre_key = pres[0].split('__')[1].split('-')[0]
            
            fig, ax = plt.subplots(3, 1, figsize=(12, 9), sharex=True)
            colors = get_model_colors(pres)
            for (i, pre) in enumerate(pres):
                model = H.plsr_load_model(model_dir/pre/'model.npz')
                samples = model.get(which)
                mu = np.mean(samples, axis=0, keepdims=True)
                if which == 'std_coefficients':
                    mu = H.plsr_normalize_coefficients(mu)
                else:
                    mu = H.plsr_normalize_vips(mu)
                waves = model.get('wavelengths')
                wranges = model.get('kept_wave_ranges')
                label = pre.split('__')[1]
                ax[i] = H.plot_stat_vs_wavelength(ax=ax[i],
                                                  waves=waves,
                                                  wave_ranges=wranges,
                                                  samples=mu,
                                                  color=colors[pre],
                                                  show_sdev=False,
                                                  legend_label=label,
                                                  linestyle='-')
                ax[i] = annotate_wavelengths_(ax=ax[i],
                                              values=mu,
                                              waves=waves,
                                              wave_ranges=wranges,
                                              annotate=annotate)
                

                if which == 'std_coefficients':
                    ax[i].set_ylim(-1, 1)
                else:
                    ax[i].set_ylim(0, 1)
                ax[i].set_xlim(0, 2500)
                ax[i].legend(loc='upper left')
                ax[i].axhline(y=0, xmin=400/2500, color='#000000')
                ax[i].set_facecolor('#EEEEEE')

            fig.suptitle(f'{trait}:{which} (separate, {final_key}, {annotate})')
            plt.tight_layout()
            save_name = f'{which}__separate__{final_key}__{pre_key}__{annotate}'
            plt.savefig(post_dir/f'{save_name}.png')
            plt.close()
        

def create_coefficient_vip_plots() -> None:
    trait_model_names = get_model_names_by_trait()
    for (trait, models) in trait_model_names.items():
        composite_coefficient_vip_plots_(trait, 
                                         models,
                                         'std_coefficients')
        composite_coefficient_vip_plots_(trait, 
                                         models,
                                         'vips')
        separate_coefficient_vip_plots_(trait, 
                                        models,
                                        which='std_coefficients',
                                        annotate='peaks')
        separate_coefficient_vip_plots_(trait, 
                                        models,
                                        which='std_coefficients',
                                        annotate='troughs')
        separate_coefficient_vip_plots_(trait, 
                                        models,
                                        which='vips',
                                        annotate='peaks')


def get_sample_colors(preds_df: pd.DataFrame,
                      color: str) -> pd.DataFrame:
    return pd.DataFrame({'sample_id': list(preds_df['sample_id'].unique()),
                         'color': color})
    
def create_prediction_diagnostic_plots_(trait: str,
                                      models: List[str]) -> None:
    PATHS = get_paths()
    model_dir = PATHS['model']
    deploy_dir = PATHS['deploy']
    post_dir = PATHS['post']/trait
    post_dir.mkdir(parents=True, exist_ok=True)

    for model in models:
        preds_df = pd.read_csv(deploy_dir/model/'preds.csv')
        metrics_df = pd.read_csv(deploy_dir/model/'metrics.csv')
        r2s = metrics_df['r2'].values
        r2 = np.mean(r2s)
        r2fs = metrics_df['fitted_r2'].values
        r2f = np.mean(r2fs)
        rrmses = metrics_df['range_normalized_rmse'].values
        rrmse = np.mean(rrmses)
        irmses = metrics_df['interquartile_normalized_rmse'].values
        irmse = np.mean(irmses)
        stat = f'R2:{r2:0.2f}, R2f:{r2f:0.2f}\nRRMSE:{rrmse:0.2f}, IRMSE:{irmse:0.2f}'
        colors_df = get_sample_colors(preds_df, '#E7724A')

        m = H.plsr_load_model(model_dir/model/'model.npz')
        ncomps = m.get('n_components')
        # print(ncomps)
        fig, ax = plt.subplots(3, 2, figsize=(12, 18))
        ax[0, 0] = H.plot_pred_vs_true(ax=ax[0, 0],
                                       pred_df=preds_df,
                                       color_df=colors_df,
                                       show_sdev=True)
        ax[0, 0].text(0.05, 
                      0.93, 
                      stat, 
                      fontsize=12,
                      transform=ax[0, 0].transAxes,
                      horizontalalignment='left')
        ax[0, 1] = H.plot_true_pred_kde(ax=ax[0, 1],
                                        pred_df=preds_df,
                                        true_color='#E7724A',
                                        pred_color='#E7724A',
                                        true_linestyle='-',
                                        pred_linestyle='--')
        ax[0, 1].set_ylabel('')
        ax[0, 1].legend()

        ax[1, 0] = H.plot_scatter_quantity_vs_int(ax[1, 0],
                                                  quantity=r2s.flatten(),
                                                  ints=ncomps.flatten(),
                                                  color='#E7724A')
        ax[1, 0].set_xlabel('# Components')
        ax[1, 0].set_ylabel('R2')

        ax[1, 1] = H.plot_scatter_quantity_vs_int(ax[1, 1],
                                                  quantity=r2fs.flatten(),
                                                  ints=ncomps.flatten(),
                                                  color='#E7724A')
        ax[1, 1].set_xlabel('# Components')
        ax[1, 1].set_ylabel('R2Fitted')

        ax[2, 0] = H.plot_scatter_quantity_vs_int(ax[2, 0],
                                                  quantity=rrmses.flatten(),
                                                  ints=ncomps.flatten(),
                                                  color='#E7724A')
        ax[2, 0].set_xlabel('# Components')
        ax[2, 0].set_ylabel('RN-RMSE')

        ax[2, 1] = H.plot_scatter_quantity_vs_int(ax[2, 1],
                                                  quantity=irmses.flatten(),
                                                  ints=ncomps.flatten(),
                                                  color='#E7724A')
        ax[2, 1].set_xlabel('# Components')
        ax[2, 1].set_ylabel('IN-RMSE')
        
        fig.suptitle(model)
        plt.tight_layout()
        suffix = '__'.join(model.split('__')[1:])
        plt.savefig(post_dir/f'prediction_diagnostics__{suffix}.png')
        plt.close()


def create_prediction_diagnostic_plots() -> None:
    trait_model_names = get_model_names_by_trait()
    for (trait, models) in trait_model_names.items():
        create_prediction_diagnostic_plots_(trait=trait,
                                          models=models)


def create_report() -> None:
    PATHS = get_paths()

    images = []
    trait_model_names = get_model_names_by_trait()
    for trait in trait_model_names.keys():
        post_dir = PATHS['post']/trait

        ordered = []
        pred_diags = [f for f in post_dir.glob('prediction_diagnostics*.png')]
        pred_diags.sort()
        ordered += pred_diags

        # std coefficients
        comp_coeffs = [f for f in post_dir.glob('std_*composite*best-over-all.png')]
        comp_coeffs.sort()
        ordered += comp_coeffs
        
        sepa_coeffs = [f for f in post_dir.glob('std_*separate*best-over-all*.png')]
        sepa_coeffs.sort()
        ordered += sepa_coeffs

        comp_coeffs = [f for f in post_dir.glob('std_*composite*best-per-loop.png')]
        comp_coeffs.sort()
        ordered += comp_coeffs
        
        sepa_coeffs = [f for f in post_dir.glob('std_*separate*best-per-loop*.png')]
        sepa_coeffs.sort()
        ordered += sepa_coeffs

        # vips
        comp_coeffs = [f for f in post_dir.glob('vips_*composite*best-over-all.png')]
        comp_coeffs.sort()
        ordered += comp_coeffs
        
        sepa_coeffs = [f for f in post_dir.glob('vips_*separate*best-over-all*.png')]
        sepa_coeffs.sort()
        ordered += sepa_coeffs

        comp_coeffs = [f for f in post_dir.glob('vips_*composite*best-per-loop.png')]
        comp_coeffs.sort()
        ordered += comp_coeffs
        
        sepa_coeffs = [f for f in post_dir.glob('vips_*separate*best-per-loop*.png')]
        sepa_coeffs.sort()
        ordered += sepa_coeffs

        # trait name page
        iwidth, iheight = 900, 1200
        image = Image.new('RGB', (iwidth, iheight), color=(255, 255, 255))
        draw = ImageDraw.Draw(image)
        font = ImageFont.truetype(H.get_font_path(), 60)
        bbox = draw.textbbox((0, 0), trait, font=font)
        twidth = bbox[2] - bbox[0]
        theight = bbox[3] - bbox[1]
        x = (iwidth - twidth)//2
        y = (iheight - theight)//2
        draw.text((x, y), trait, fill=(0, 0, 0), font=font)
        images.append(image)

        # add the ordered plots
        for f in ordered:
            images.append(Image.open(f).convert('RGB'))

    save_pdf = PATHS['post']/'report.pdf'
    images[0].save(save_pdf, 
                   save_all=True, 
                   append_images=images[1:],
                   resolution=600)
    # print(f'Saved: {save_pdf}')

    
if __name__ == '__main__':
    print('Creating: metrics tables ...')
    create_metrics_tables()

    print('Creating: coefficient/VIP plots ...')
    create_coefficient_vip_plots()

    print('Creating: prediction diagnostics plots ...')
    create_prediction_diagnostic_plots()

    print('Creating: report ...')
    create_report()