from core.datasets.factory import create_dataset
from core.explanations.utils import bitmask_to_list, calc_samples_prototype
from core.explanations.analyzers.IndividualExplanationAnalyzer import IndividualExplanationAnalyzer
from core.explanations.analyzers.ClusterExplanationAnalyzer import ClusterExplanationAnalyzer
from scripts.utils import load_csv_files

import argparse
import numpy as np


def main():
    args = parse_arguments()

    # Load data from CSV
    df = load_csv_files(args.csv)
    if len(df) == 0:
        print("No data found in CSV file")
        return

    # Get the indicated row from the dataframe
    key = (df['dataset'] == args.dataset) & (df['custom_dim'] == args.custom_dim) & (df['label'] == args.label)

    if args.sample_index != -1:
        key &= (df['sample_index'] == args.sample_index)
    elif args.cluster_index != -1:
        key &= (df['cluster_index'] == args.cluster_index)

    if args.ordering is not None:
        key &= (df['ordering'] == args.ordering)

    if 'delta' in df.columns:
        key &= (df['delta'] == args.delta)

    data = df.loc[key].iloc[0]

    dataset = create_dataset(data['dataset'], data['custom_dim'], data['seed'])

    if args.plot_explanations:
        plot_explanations(data, dataset, args)

    if args.plot_feature_importance:
        plot_feature_importance(data, dataset, args.output)


def parse_arguments():
    parser = argparse.ArgumentParser(description='Plot explanations and feature importance',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # Input arguments
    parser.add_argument('--csv', type=str, required=True,
                        help='CSV file with the data to plot')

    # Sample arguments
    parser.add_argument('--dataset', type=str, required=True,
                        help='Dataset to plot')
    parser.add_argument('--custom-dim', type=int, default=0,
                        help='Custom dimension to plot')
    parser.add_argument('--label', type=int, required=True,
                        help='Label to plot')
    parser.add_argument('--sample-index', type=int, default=-1,
                        help='Sample index to plot')
    parser.add_argument('--cluster-index', type=int, default=-1,
                        help='Cluster index to plot')
    parser.add_argument('--delta', type=float, default=-1,
                        help='Delta to plot')
    parser.add_argument('--ordering', type=str, default=None,
                        choices=['linear', 'random', 'verix', 'sensitivity'],
                        help='Ordering to plot')

    # Explanations arguments
    parser.add_argument('--plot-explanations', action='store_true',
                        help='Plot the explanations')
    parser.add_argument('--num-explanations-to-plot', type=int, default=11,
                        help='Number of explanations to plot')
    parser.add_argument('--bounds', type=float, default=-1,
                        help='Bounds to plot (default: delta)')

    # Feature importance arguments
    parser.add_argument('--plot-feature-importance', action='store_true',
                        help='Plot the feature importance')

    # Output arguments
    parser.add_argument('--output', type=str, default='explanation.pdf',
                        help='File name to save the figure')

    args = parser.parse_args()

    if args.dataset == 'custom' and args.custom_dim == 0:
        raise ValueError('Custom-dim must be provided for custom dataset')
    if args.dataset != 'custom' and args.custom_dim != 0:
        raise ValueError('Custom-dim must not be provided for non-custom dataset')

    if args.sample_index == -1 and args.cluster_index == -1:
        raise ValueError('Either sample-index or cluster-index must be provided')
    if args.sample_index != -1 and args.cluster_index != -1:
        raise ValueError('Only one of sample-index or cluster-index must be provided')
    if args.cluster_index != -1 and args.delta == -1:
        raise ValueError('Cluster-index requires delta to be provided')

    if args.bounds == -1:
        args.bounds = args.delta

    return args

def plot_explanations(data, dataset, args):
    for i, explanation in enumerate(data['explanations'][:args.num_explanations_to_plot]):
        explanation = bitmask_to_list(explanation, dataset.num_features)

        if 'sample' in data:
            sample = np.array(data['sample'])
        elif 'samples' in data:
            samples = [np.array(sample) for sample in data['samples']]
            sample = calc_samples_prototype(samples, data['prototype'])
        else:
            raise ValueError('No sample or samples found in the data')

        if dataset.name == 'mnist' or dataset.name == 'digits':
            dataset.plot_explanation(sample, explanation, args.output)
        else:
            if data['id'] == 'Individual':
                bounds = IndividualExplanationAnalyzer.calc_explanation_bounds(explanation, sample, data['epsilon'])
            elif data['id'].startswith('δ'):
                bounds = ClusterExplanationAnalyzer.calc_explanation_bounds(explanation, samples, data['prototype'], data['epsilon'], args.bounds)
            elif data['id'] == 'Inflated':
                bounds = data['inflated_explanations']
            else:
                raise ValueError(f'Unknown explanation id: {data["id"]}')

            dataset.plot_explanation(sample, explanation, bounds, args.output)

def plot_feature_importance(data, dataset, output):
    if 'sample' in data:
        importance = IndividualExplanationAnalyzer.calc_feature_importance(data['explanations'], dataset.num_features)
    elif 'samples' in data:
        importance = ClusterExplanationAnalyzer.calc_feature_importance(data['explanations'], dataset.num_features)

    dataset.plot_feature_importance(importance, output)

if __name__ == '__main__':
    main()
