from core.datasets.factory import create_dataset
from core.explanations.factories.IndividualExplanationFactory import IndividualExplanationFactory
from core.explanations.factories.ClusterExplanationFactory import ClusterExplanationFactory
from core.models.factory import create_model
from scripts.utils import set_random_seed, DATASETS

import argparse
import json
import timeit


def main():
    args = parse_arguments()

    set_random_seed(args.seed)

    dataset = create_dataset(args.dataset, args.custom_dim, args.seed)
    model = create_model('OrderingModel', dataset, args.seed, args.label)

    data = {
        'dataset': args.dataset,
        'custom_dim': args.custom_dim,
        'label': args.label,
        'epsilon': args.epsilon,
        'ordering': args.ordering,
        'seed': args.seed,
    }

    if args.explainer == 'individual':
        add_individual_explanations(args, model, data)
    elif args.explainer == 'cluster':
        add_cluster_explanations(args, model, data)
        data['delta'] = args.delta
        data['prototype'] = args.prototype

    with open(args.output, 'w') as f:
        json.dump(data, f, indent=4)


def parse_arguments():
    parser = argparse.ArgumentParser(description='Generate explanations for a sample or cluster',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # Dataset parameters
    parser.add_argument('--dataset', type=str, required=True, choices=DATASETS,
                        help='Dataset to use')

    # Custom dataset parameters
    parser.add_argument('--custom-dim', type=int, default=0,
                        help='Dimensionality of the custom dataset')

    # Explanation parameters
    parser.add_argument('--explainer', type=str, default='cluster',
                        choices=['individual', 'cluster'],
                        help='Type of explainer to use')
    parser.add_argument('--label', type=int, default=0,
                        help='Label to explain')
    parser.add_argument('--epsilon', type=float, default=1.0,
                        help='Robustness of the explanation')
    parser.add_argument('--ordering', type=str, default='sensitivity',
                        choices=['linear', 'random', 'verix', 'sensitivity'],
                        help='How to generate the explanations')
    parser.add_argument('--seed', type=int, default=0,
                        help='Random seed for reproducibility')

    # Individual explanation parameters
    parser.add_argument('--sample-index', type=int, default=-1,
                        help='Index of the sample to explain')

    # Cluster explanation parameters
    parser.add_argument('--cluster-index', type=int, default=-1,
                    help='Index of the cluster to explain')
    parser.add_argument('--cluster-size', type=int, default=20,
                        help='Number of samples in the cluster')
    parser.add_argument('--delta', type=float, default=0.0,
                        help='Delta value for the cluster explanations')
    parser.add_argument('--prototype', type=str, default='mean',
                        help='Prototype method for the cluster')

    # Explanation generation parameters
    parser.add_argument('--num-explanations', type=int, default=100,
                        help='Number of explanations to generate')
    parser.add_argument('--num-consumers', type=int, default=10,
                        help='Number of consumers for the explanations')

    # Output parameters
    parser.add_argument('--output', type=str, default='explanations.json',
                        help='File to save the explanations')

    args = parser.parse_args()

    if args.dataset == 'custom' and args.custom_dim == 0:
        parser.error("Custom dataset requires --custom_dim to be set.")

    if args.dataset != 'custom' and args.custom_dim != 0:
        parser.error("Custom dataset parameters are only applicable when --dataset is set to 'custom'.")

    if args.explainer != 'cluster' and args.sample_index == -1:
        parser.error("Sample index is required for the individual and inflated explainers.")

    if args.explainer == 'cluster' and args.cluster_index == -1 and args.sample_index == -1:
        parser.error("The cluster explainer requires either a cluster or a sample index.")

    return args


def add_individual_explanations(args, model, data):
    samples = model.get_samples(args.label, num_samples=args.sample_index + 1)
    sample = samples[args.sample_index]

    factory = IndividualExplanationFactory(
        model=model,
        label=args.label
    )

    start_time = timeit.default_timer()
    explanations = factory.get_explanations(
        sample=sample,
        epsilon=args.epsilon,
        ordering=args.ordering,
        num_explanations=args.num_explanations,
        num_consumers=args.num_consumers,
        use_progress_bar=True
    )
    end_time = timeit.default_timer()
    computation_time = end_time - start_time

    data['id'] = 'Individual'
    data['sample'] = sample.tolist()
    data['sample_index'] = args.sample_index
    data['explanations'] = list(explanations)
    data['computation_time'] = computation_time


def add_cluster_explanations(args, model, data):
    if args.cluster_index != -1:
        num_samples = args.cluster_size * (args.cluster_index + 1)
        samples = model.get_samples(args.label, num_samples)
        samples = samples[args.cluster_index * args.cluster_size:]
        data['cluster_index'] = args.cluster_index
        data['cluster_size'] = args.cluster_size
    else:
        num_samples = args.sample_index + 1
        samples = model.get_samples(args.label, num_samples)
        samples = [samples[args.sample_index]]
        data['sample_index'] = args.sample_index

    factory = ClusterExplanationFactory(
        model=model,
        label=args.label,
        prototype=args.prototype
    )

    # Compute the explanations
    start_time = timeit.default_timer()
    explanations = factory.get_explanations(
        samples=samples,
        epsilon=args.epsilon,
        delta=args.delta,
        ordering=args.ordering,
        num_explanations=args.num_explanations,
        num_consumers=args.num_consumers,
        use_progress_bar=True
    )
    end_time = timeit.default_timer()
    computation_time = end_time - start_time

    data['id'] = f'δ={args.delta:.2f}'
    data['samples'] = [sample.tolist() for sample in samples]
    data['explanations'] = sorted(explanations)
    data['computation_time'] = computation_time


if __name__ == '__main__':
    main()
