'''
plot local feature contribution, i.e.,
other feature are fixed, how a certain feature influence the outcome
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 functools import partial
from utilities.utils import denormalize, denormlize_ds, normalize
def visualize_per_local_feat_ctb_for_mu_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_local_ctb_dir)
    utils.cond_mkdir(savedir)
    dataset_name = spec["Class"]
    in_geo_features = spec["InGeoFeatures"]


    if dict_pos is None:
        pos = None
        name = 'all'
    else:
        name = list(dict_pos.keys())[0]
        ldm_value = list(dict_pos.values())[0]
        pos = normalize(ds_=test_dataset, arr_= ldm_value, var_name='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,
                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,
                device=device)

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


        f_mu = model_trained.infer_with_subnetwork(arr_input_grids, ith_cov)
        f_mu = f_mu.squeeze().cpu()
        mh_map, sh_map = f_mu.detach().numpy()

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


        # denormalize
        arr_x_train = denormalize(ds_=train_dataset,
                           arr_=x_train[..., in_geo_features + ith_cov],
                           var_name=LIST_USED_COVARIATES[ith_cov])
        arr_y_train = denormalize(ds_=train_dataset,
                           arr_=y_train,
                           var_name=train_dataset.tgt_var_name)
        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])

        mh_map = denormalize(ds_=train_dataset,
                           arr_=mh_map,
                           var_name=train_dataset.tgt_var_name)


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

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

        x_axis_name, y_axis_name = LIST_USED_COVARIATES[ith_cov], train_dataset.tgt_var_name
        dict_info = {"x_axis_name": x_axis_name + DICT_TGT_UNIT[dataset_name][x_axis_name],
                     "y_axis_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,
            None,
            arr_x_test,
            arr_y_test,
            dict_info,
            savepath=savepath)

    return





def visualize_per_local_feat_ctb_for_mu_airway(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_local_feat_ctb_for_mu_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_local_feat_ctb_for_mu(dataset_name):
    # which dataset to use
    if "Airway" in dataset_name or "AFQ" in dataset_name:
        return partial(visualize_per_local_feat_ctb_for_mu_airway, dataset_name)
    else:

        return visualize_per_local_feat_ctb_for_mu_general



def visualize_local_for_mu(specs_filename: str,):
    spec = load_json(specs_filename)
    ds_train, ds_test_dataloader = load_dataset(specs_filename=specs_filename, which_split='train')
    ds_test, ds_test_dataloader = load_dataset(specs_filename=specs_filename, which_split='test')
    trained_model = load_trained_model(specs_filename, spec["SavedBestCheckpointPath"]) # saved checkpoint name
    trained_model.eval()
    savedir = os.path.join(spec["LoggingRoot"], spec["ExperimentName"]) #'/playpen-raid/Author/LucidAtlas/figures/v12'
    utils.cond_mkdir(savedir)
    dataset_name = spec["Class"]
    visualize_per_local_feat_ctb_for_mu(dataset_name)(spec, 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/OASISBrain/brain_lucidatlas_v14_0122.json",# default="/playpen-raid/Author/LucidAtlas/configs/adni/adnihp_lucidatlas_v14_0120.json", #'/playpen-raid/Author/LucidAtlas/configs/airways/lucidatlas_1d_csa_v13_0116.json',
        default='/playpen-raid/Author/LucidAtlas/configs/airways/airway_lucidatlas_v14_0123_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_local_for_mu(args.experiment_directory)
    print('1')




