'''
plot global feature contribution, i.e.,
suppose other features don't exist in the analysis, how a certain feature influence the outcome.
This is what we got from our marginalization

Author

'''





import matplotlib.pyplot as plt
plt.rcParams['figure.dpi'] = 300
import torch
import numpy as np
import vis_utils
from utilities import utils
import os

import argparse
from pipeline.load import *
import torch.utils.data as data_utils

import model.networks.basics.workspace as ws
from utilities.utils import denormalize, denormlize_ds, normalize, denormalize_from_distribution
from functools import partial
from  pipeline.load import DICT_TGT_UNIT
def visualize_per_global_feat_imp_general(spec: dict,
                        model_trained: torch.nn.Module,
                        train_dataset: data_utils.Dataset,
                        test_dataset: data_utils.Dataset=None,
                        dict_pos=None):

    '''
    visualize interpretation, it will be a 2D plot,
    x axis is the feature value, y axis is the feature's contribution
    this function can be used for different datasets
    '''
    LIST_USED_COVARIATES = spec["CovariateNames"]
    device = spec["Device"]
    savedir = os.path.join(spec["LoggingRoot"], spec["ExperimentName"], ws.vis_global_imp_dir)
    in_geo_features = spec["InGeoFeatures"]
    utils.cond_mkdir(savedir)
    #tgt_var_name = spec["TargetVariableName"]
    dataset_name = spec["Class"]

    if dict_pos is None:
        pos = None
        name = 'all'
        ldm_value = 'none'
    else:
        name = list(dict_pos.keys())[0]
        ldm_value = list(dict_pos.values())[0]
        pos = np.array(normalize(ds_=test_dataset, arr_=np.array([[ldm_value]]), var_name=train_dataset.geo_var_name)[0]) #np.array([ldm_value])# normalize(ds_=test_dataset, arr_=ldm_value, var_name=train_dataset.geo_var_name)
        print(pos)


    for ith_cov in range(len(LIST_USED_COVARIATES)):
        # make query samples
        arr_input_grids, x_train, y_train = \
            utils.make_grids_and_dps_for_2d_vis(dataset_name=dataset_name)(
                train_dataset,
                covs_to_plot=[LIST_USED_COVARIATES[ith_cov]],
                pos=pos,
                num_of_samples=100,
                device=device)

        _, x_test, y_test = \
            utils.make_grids_and_dps_for_2d_vis(dataset_name=dataset_name)(
                test_dataset,
                covs_to_plot=[LIST_USED_COVARIATES[ith_cov]],
                pos=pos,
                num_of_samples=100,
                device=device)

        print(arr_input_grids.shape[0])

        # make plots
        x_train, y_train = x_train.cpu(), y_train.cpu()
        x_test, y_test = x_test.cpu(), y_test.cpu()

        model_trained = model_trained.eval()
        with torch.no_grad():
            ## query the model
            #f_mu, f_var = model_trained.infer_global_importance(arr_input_grids, LIST_USED_COVARIATES[ith_cov])
            f_mu, f_var =  utils.batched_infer_global_importance(model_trained, arr_input_grids, LIST_USED_COVARIATES[ith_cov], batch_size=10)


        mh_map, sh_map = f_mu, torch.sqrt(f_var)
        mh_map, sh_map, low_bd_map, high_bd_map = denormalize_from_distribution(ds_=train_dataset, mu=mh_map, sigma=sh_map)


        if pos is not None:
            savepath = f'{savedir}/{dataset_name}_{name}_global_{LIST_USED_COVARIATES[ith_cov]}.png'
        else:
            savepath = f'{savedir}/{dataset_name}_global_{LIST_USED_COVARIATES[ith_cov]}.png'


        arr_x_train = denormalize(ds_=train_dataset,
                           arr_=x_train[..., in_geo_features + ith_cov],
                           var_name=LIST_USED_COVARIATES[ith_cov]).squeeze()
        arr_y_train = denormalize(ds_=train_dataset,
                           arr_=y_train,
                           var_name=train_dataset.tgt_var_name).squeeze()
        arr_x_grids = denormalize(ds_=train_dataset,
                           arr_=arr_input_grids[..., in_geo_features + ith_cov].cpu().numpy(),
                           var_name=LIST_USED_COVARIATES[ith_cov]).squeeze()

        # mh_map = denormalize(ds_=train_dataset,
        #                    arr_=mh_map,
        #                    var_name=train_dataset.tgt_var_name).squeeze()
        #
        # high_bd_map = denormalize(ds_=train_dataset,
        #                    arr_=high_bd_map,
        #                    var_name=train_dataset.tgt_var_name).squeeze()
        #
        # low_bd_map = denormalize(ds_=train_dataset,
        #                    arr_=low_bd_map,
        #                    var_name=train_dataset.tgt_var_name).squeeze()
        #


        arr_x_test = denormalize(ds_=train_dataset,
                           arr_= x_test[..., in_geo_features + ith_cov],
                           var_name=LIST_USED_COVARIATES[ith_cov]).squeeze()
        arr_y_test = denormalize(ds_=train_dataset,
                           arr_=y_test,
                           var_name=train_dataset.tgt_var_name).squeeze()

        # arr_x_train = x_train[..., in_geo_features + ith_cov].cpu().detach().numpy()
        # arr_y_train = y_train.cpu().detach().numpy()
        # arr_x_grids = arr_input_grids[..., in_geo_features + ith_cov].cpu().detach().numpy()
        # mh_map = mh_map.squeeze().cpu().detach().numpy()
        # high_bd_map = high_bd_map.squeeze()
        # low_bd_map = low_bd_map.squeeze()
        # arr_x_test = x_test[..., in_geo_features + ith_cov].cpu().detach().numpy()
        # arr_y_test = y_test.cpu().detach().numpy()

        x_axis_name, y_axis_name = LIST_USED_COVARIATES[ith_cov], train_dataset.tgt_var_name
        from pipeline.load import DICT_DISPLAY_NAMES
        dict_info = {"x_axis_name": DICT_DISPLAY_NAMES[dataset_name][x_axis_name] + DICT_TGT_UNIT[dataset_name][x_axis_name],
                     "y_axis_name": DICT_DISPLAY_NAMES[dataset_name][y_axis_name] + DICT_TGT_UNIT[dataset_name][y_axis_name],
                     "pos": ldm_value
                     }

        vis_utils.plot_regression_all_samples(
            arr_x_train,
            arr_y_train,
            arr_x_grids,
            mh_map,
            high_bd_map,
            low_bd_map,
            arr_x_test,
            arr_y_test,
            dict_info,
            savepath=savepath)

    return



def visualize_per_global_feat_imp_spatial(spec: dict,
             model_trained: torch.nn.Module,
             train_dataset: data_utils.Dataset,
             test_dataset: data_utils.Dataset,
             dataset_name: str):

    '''
    make contribution plots for airway dataset at different depth/landmarks
    '''
    CURRENT_LDMS = LANDMARKS[dataset_name]
    list_pos = list(CURRENT_LDMS.keys())

    for ith_pos in list_pos:
        visualize_per_global_feat_imp_general(spec=spec,
                             model_trained=model_trained,
                             train_dataset=train_dataset,
                             test_dataset=test_dataset,
                             dict_pos={ith_pos: CURRENT_LDMS[ith_pos]})
    return 0

def visualize_per_global_feat_imp_toydata(spec: dict,
             model_trained: torch.nn.Module,
             train_dataset: data_utils.Dataset,
             test_dataset: data_utils.Dataset):

    '''
    make contribution plots for airway dataset at different depth/landmarks
    '''

    list_pos = list(LANDMARKS_TOY.keys())


    #for ith_pos in range(len(list_pos)):
    for ith_pos in list_pos:
        visualize_per_global_feat_imp_general(spec=spec,
                             model_trained=model_trained,
                             train_dataset=train_dataset,
                             test_dataset=test_dataset,
                             dict_pos={ith_pos: LANDMARKS_TOY[ith_pos]})
    return 0


def visualize_per_global_feat_imp(dataset_name):
    # which dataset to use
    if "Airway" in dataset_name or "AFQ" in dataset_name:
        return partial(visualize_per_global_feat_imp_spatial,  dataset_name=dataset_name)
    elif dataset_name == "ToyData":
        return visualize_per_global_feat_imp_toydata
    else:
        return visualize_per_global_feat_imp_general


def visualize_global(specs_filename: str, cv_idx: int=None):
    specs = load_json(specs_filename, cv_idx)
    ds_train, ds_test_dataloader = load_dataset(specs=specs, which_split='train_val')
    ds_test, ds_test_dataloader = load_dataset(specs=specs, which_split='test')
    trained_model = load_trained_model(specs=specs, filename_checkpoint=specs["SavedCheckpointPath"]) # saved checkpoint name
    trained_model.eval()

    savedir = os.path.join(specs["LoggingRoot"], specs["ExperimentName"]) #'/playpen-raid/Author/LucidAtlas/figures/v12'
    utils.cond_mkdir(savedir)
    dataset_name = specs["Class"]
    visualize_per_global_feat_imp(dataset_name)(specs, trained_model, ds_train, ds_test)
    return





if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser(description="Train a LucidAtlas autodecoder")
    arg_parser.add_argument(
        "--experiment",
        "-e",
        dest="experiment_directory",
        #default='/playpen-raid/Author/LucidAtlas/configs/airways/airway_namlss_v1_0123_full.json', #default="/playpen-raid/Author/LucidAtlas/configs/OASISBrain/brain_lucidatlas_v14_0122.json",#default="/playpen-raid/Author/LucidAtlas/configs/adni/adnihp_lucidatlas_v14_0120.json", #'/playpen-raid/Author/LucidAtlas/configs/airways/airway_lucidatlas_v14_0121.json', #  #
        #default='/playpen-raid/Author/LucidAtlas/configs/airways/v0/airway_lucidatlas_v14_0123_part.json',
        #default="/playpen-raid/Author/LucidAtlas/configs/OASISBrain/v1/brain_lucidatlas_part.json",
        #default="/playpen-raid/Author/LucidAtlas/configs/OASISBrain/v1/brain_lucidatlas_part.json",
        default='/playpen-raid/Author/LucidAtlas/configs/airways/v2/airway_lucidatlas_full.json',
        help="The experiment directory. This directory should include "
             + "experiment specifications in 'specs.json', and logging will be "
             + "done in this directory as well.",
    )
    arg_parser.add_argument(
        "--checkpoint",
        "-c",
        dest="checkpoint",
        default="latest",
        help="The checkpoint weights to use. This can be a number indicated an epoch "
             + "or 'latest' for the latest weights (this is the default)",
    )

    arg_parser.add_argument(
        "--train",
        dest="whether_train",
        default=True,
        help="whether to train from scratch",
    )

    arg_parser.add_argument(
        "--test",
        dest="whether_test",
        default=True,
        help="whether to test",
    )

    arg_parser.add_argument(
        "--vis",
        dest="whether_vis",
        default=True,
        help="whether to vis",
    )

    args = arg_parser.parse_args()

    if args.whether_vis:
        visualize_global(args.experiment_directory)
    print('1')




