from collections import defaultdict
from core.datasets.factory import create_dataset
from scripts.utils import load_csv_files, lineplot, print_grouped_means

import argparse
import os
import tqdm


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

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

    for data_dir in args.data:
        # Compute individual explanation statistics
        individual_stats = compute_individual_stats(data_dir, args)

        if args.x_col_name != "Group Size" or individual_stats[0]["ID"] != "Individual":
            stats.extend(individual_stats)

        # Compute aggregated statistics if requested
        if args.include_aggregated:
            for cluster_size in tqdm.trange(args.min_cluster_size, args.max_cluster_size + 1):
                aggregated_stats = compute_aggregated_stats(data_dir, cluster_size)
                stats.extend(aggregated_stats)

    lineplot(
        stats=stats,
        x_col_name=args.x_col_name,
        y_col_name="Runtime (s)",
        file_name=f"{args.output_dir}/computation_time.pdf",
        title=args.title,
        log_scale=True
    )

    print_grouped_means(stats, [args.x_col_name, "ID"], ["Runtime (s)"])

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

    # Input data
    parser.add_argument('--data', type=str, nargs='+', required=True,
                    help='One or more CSV files containing explanation data to analyze')

    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.')

    # Aggregated statistics
    parser.add_argument('--include-aggregated', action='store_true',
                        help='Include aggregated statistics for Individual explanations')
    parser.add_argument('--min-cluster-size', type=int, default=3,
                        help='Minimum cluster size for aggregating explanations')
    parser.add_argument('--max-cluster-size', type=int, default=20,
                        help='Maximum cluster size for aggregating explanations')

    # Plotting
    parser.add_argument('--x-col-name', type=str, default='Dataset',
                        help='Column name for the x-axis')

    # Output data
    parser.add_argument('--output-dir', type=str, default='figures/cluster/classic_datasets/',
                        help='Directory to save the figure')
    parser.add_argument('--title', type=str, default=None,
                        help='Title of the figure')

    args = parser.parse_args()

    if not args.include_aggregated and (args.min_cluster_size != 3 or args.max_cluster_size != 20):
        parser.error('Cluster size cannot be set without setting --include-aggregated')

    return args


def compute_individual_stats(data_file, args):
    # Load data as DataFrame
    df = load_csv_files(data_file)

    if df.empty:
        return []

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

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

    if df.empty:
        return []

    # Process data
    print(f"Processing individual stats from {data_file}")

    stats = []
    for _, row in tqdm.tqdm(df.iterrows(), total=len(df)):
        dataset = create_dataset(row['dataset'], row['custom_dim'], row['seed'])

        stat = {
            "ID": row['id'],
            "Dataset": dataset.name,
            "Label": row['label'],
            "Runtime (s)": row['computation_time'],
            "Number of Explanations": len(row['explanations']),
            "Number of Features": dataset.num_features,
            "Seed": row['seed']
        }

        if "cluster_size" in row:
            stat["Group Size"] = row["cluster_size"]

        if args.x_col_name == "Group Size" and "cluster_size" not in row:
            for cluster_size in range(args.min_cluster_size, args.max_cluster_size + 1):
                stat_copy = stat.copy()
                stat_copy["Group Size"] = cluster_size
                stats.append(stat_copy)
        else:
            stats.append(stat)

    return stats


def compute_aggregated_stats(data_file, cluster_size):
    stats = []

    # Load data as DataFrame
    df = load_csv_files(data_file)

    # Filter for Individual explanations only
    df = df[df['id'] == "Individual"]

    if df.empty:
        return []

    # Group by a key (dataset, custom_dim, seed, label, cluster_index)
    grouped_data = defaultdict(list)
    for _, row in df.iterrows():
        cluster_index = row['sample_index'] // cluster_size
        key = (row['dataset'], row['custom_dim'], row['seed'], row['label'], cluster_index)
        grouped_data[key].append(row)

    for key, data in grouped_data.items():
        computation_time = sum([float(d['computation_time']) for d in data])
        num_explanations = sum([len(d['explanations']) for d in data])

        stat = {
            "ID": "Aggregated",
            "Dataset": key[0],
            "Group Size": cluster_size,
            "Label": key[3],
            "Runtime (s)": computation_time,
            "Number of Explanations": num_explanations,
            "Number of Features": key[1],
            "Seed": key[2]
        }
        stats.append(stat)

    return stats


if __name__ == "__main__":
    main()
