from core.datasets.factory import create_dataset
from core.explanations.factories.InflatedExplanationFactory import InflatedExplanationFactory
from core.models.factory import create_model
from scripts.utils import set_random_seed, load_csv_files

import argparse
import json
import timeit


def main():
    args = parse_arguments()

    set_random_seed(args.seed)

    # Load data from CSV
    df = load_csv_files(args.data)
    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) & \
        (df['sample_index'] == args.sample_index) & \
        (df['delta'] == args.delta)
    data = df.loc[key].iloc[0]

    # Create dataset and model
    dataset = create_dataset(data['dataset'], data['custom_dim'], data['seed'])
    model = create_model('OrderingModel', dataset, data['seed'], data['label'])

    output_data = data.to_dict()

    inflate_explanations(model, data['explanations'], output_data, args)

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


def parse_arguments():
    parser = argparse.ArgumentParser(description='Inflate explanations',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # Data parameters
    parser.add_argument('--data', type=str, default='data/explanations/delta_robust/classic/robust_explanations.csv',
                        help='CSV file with the data to inflate')
    parser.add_argument('--seed', type=int, default=0,
                        help='Random seed for reproducibility')

    # 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, default=0,
                        help='Label to plot')
    parser.add_argument('--sample-index', type=int, required=True,
                        help='Sample index to plot')
    parser.add_argument('--delta', type=float, default=0.0,
                        help='Delta to plot')

    # Inflated explanations parameters
    parser.add_argument('--eta', type=float, default=0.01,
                        help='Eta value for the inflated explanations')

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

    return parser.parse_args()


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

    factory = InflatedExplanationFactory(
        model=model,
        label=data['label']
    )

    start_time = timeit.default_timer()
    inflated_explanations = factory.inflate_explanations(
        explanations=explanations,
        sample=sample,
        epsilon=data['epsilon'],
        eta=args.eta,
        ordering=data['ordering']
    )
    end_time = timeit.default_timer()

    data['id'] = 'Inflated'
    data['inflated_explanations'] = inflated_explanations
    data['computation_time'] = end_time - start_time
    data['eta'] = args.eta


if __name__ == '__main__':
    main()
