from core.datasets.factory import create_dataset
from core.explanations.analyzers.ClusterExplanationAnalyzer import ClusterExplanationAnalyzer
from core.explanations.utils import average_score
from scripts.utils import (
    load_csv_files,
    lineplot
)

import argparse
import numpy as np
import os
import tqdm


def main():
    args = parse_arguments()
    stats = []

    # Ensure output directory exists
    os.makedirs('results', exist_ok=True)

    # Load data as DataFrames
    inflated_df = load_csv_files(args.inflated_data)
    robust_df = load_csv_files(args.robust_data)

    if robust_df.empty or inflated_df.empty:
        print("No data found in CSV files")
        return

    # Filter out rows without explanations
    robust_df = robust_df[robust_df['explanations'].apply(lambda x: len(x) > 0)]

    if args.filter_regex:
        robust_df = robust_df[robust_df[args.filter_column].astype(str).str.match(args.filter_regex, na=False)]

    if robust_df.empty:
        print("No valid explanations found")
        return

    # Create index for inflated data
    inflated_df['key'] = list(zip(inflated_df['dataset'], inflated_df['custom_dim'], inflated_df['sample_index'], inflated_df['label'], inflated_df['delta']))
    inflated_index = inflated_df.set_index('key')['inflated_explanations'].to_dict()

    # Remove data where half the samples don't have explanations
    robust_df = remove_misleading_samples(robust_df)

    # Process each row
    for _, row in tqdm.tqdm(robust_df.iterrows(), total=len(robust_df)):
        key = (row['dataset'], row['custom_dim'], row['sample_index'], row['label'], row['delta'])

        try:
            maximal_bounds = inflated_index[key]

            if len(maximal_bounds) != len(row['explanations'][0]):
                continue
        except KeyError:
            continue

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

        samples = [np.array(sample) for sample in row['samples']]
        cluster_robustness = average_score(
            row['explanations'],
            ClusterExplanationAnalyzer.calc_explanation_robustness,
            samples=samples,
            prototype=row['prototype'],
            maximal_bounds=maximal_bounds,
            epsilon=row['epsilon'],
            delta=row['delta']
        )

        stats.append({
            "ID": row['id'],
            "δ": row['id'][2:],
            "Dataset": dataset.name,
            "Number of Features": dataset.num_features,
            "Sample Index": row['sample_index'],
            "Computation Time": row['computation_time'],
            "Area Covered": cluster_robustness
        })

    if not stats:
        print("No valid statistics computed")
        return

    index_col_name = "Number of Features" if stats[0]['Dataset'].startswith('custom') else "Dataset"

    lineplot(
        stats=stats,
        x_col_name="δ",
        y_col_name="Area Covered",
        index_col_name=index_col_name,
        file_name=f"{args.output_dir}/robustness.pdf",
        title=args.title,
        log_scale=False
    )


def parse_arguments():
    parser = argparse.ArgumentParser(description='Calculate statistics of feature importance',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # Input data
    parser.add_argument('--robust-data', type=str, default='data/explanations/delta_robust/classic/robust_explanations.csv',
                        help='CSV file containing robust explanation data')
    parser.add_argument('--inflated-data', type=str, default='data/explanations/delta_robust/classic/inflated_explanations.csv',
                        help='CSV file containing inflated explanation data')

    # Output data
    parser.add_argument('--output-dir', type=str, default='figures/delta_robust/classic_datasets/',
                        help='Directory to save the figure')
    parser.add_argument('--filter-column', type=str, default='id',
                        help='Column to filter on.')
    parser.add_argument('--filter-regex', type=str, default='',
                        help='Regex filter to apply.')
    parser.add_argument('--title', type=str, default=None,
                        help='Title of the figure')

    return parser.parse_args()


def remove_misleading_samples(df):
    group_counts = df.groupby(['dataset', 'custom_dim', 'delta']).size()
    reference_counts = group_counts.reset_index()
    reference_counts = reference_counts[reference_counts['delta'] == 0].set_index(['dataset', 'custom_dim'])[0].to_dict()

    def should_keep_group(group):
        a_val, b_val, x_val = group.name
        ref_count = reference_counts.get((a_val, b_val), 0)
        return len(group) >= 0.5 * ref_count

    return df.groupby(['dataset', 'custom_dim', 'delta']).filter(should_keep_group)


if __name__ == "__main__":
    main()
