from core.datasets.factory import create_dataset
from core.models.factory import create_model
from core.explanations.analyzers.IndividualExplanationAnalyzer import IndividualExplanationAnalyzer
from core.explanations.analyzers.AggregatorExplanationAnalyzer import AggregatorExplanationAnalyzer
from core.explanations.analyzers.ClusterExplanationAnalyzer import ClusterExplanationAnalyzer
from core.explanations.analyzers.LimeExplanationAnalyzer import LimeExplanationAnalyzer
from core.explanations.analyzers.ShapExplanationAnalyzer import ShapExplanationAnalyzer
from core.explanations.utils import average_score
from scripts.utils import load_csv_files, lineplot, print_grouped_means

import argparse
import numpy as np
import os
import tqdm


def main():
    args = parse_arguments()
    os.makedirs('results', exist_ok=True)

    # Load data as DataFrames
    individual_df = load_csv_files(args.individual_data)
    cluster_df = load_csv_files(args.cluster_data)

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

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

    if individual_df.empty or cluster_df.empty:
        raise Exception("No valid explanations found")

    # Pre-compute all individual explanation metrics
    individual_metrics = precompute_individual_metrics(individual_df)

    stats = []

    # Process each cluster
    for _, row in tqdm.tqdm(cluster_df.iterrows(), total=len(cluster_df), desc="Cluster metrics"):
        # Get dataset and model
        dataset = create_dataset(row['dataset'], row['custom_dim'], row['seed'])
        model = create_model('BaseModel', dataset, row['seed'])

        # Compute cluster-level metrics
        cluster_importance, cluster_ranking, coverage = compute_cluster_metrics(row, dataset)

        # Process samples and compute aggregated metrics
        sample_metrics = process_cluster_samples(
            row, individual_metrics, cluster_importance, cluster_ranking, model, dataset
        )

        if sample_metrics is None:
            continue

        # Create statistics entries
        cluster_stats = create_stats_entries(row, dataset, sample_metrics, coverage)
        stats.extend(cluster_stats)

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

    # Generate plots and output
    generate_plots_and_output(stats, args.output_dir, args.x_col_name)


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

    # Input data
    parser.add_argument('--individual-data', type=str, default='data/explanations/cluster/classic/individual_explanations.csv',
                        help='CSV file containing individual explanation data')
    parser.add_argument('--cluster-data', type=str, default='data/explanations/cluster/classic/cluster_explanations.csv',
                        help='CSV file containing cluster explanation data')

    # Aggregated statistics
    parser.add_argument('--include-aggregated', action='store_true',
                        help='Include aggregated statistics for Individual explanations')
    parser.add_argument('--cluster_size', type=int, default=20,
                        help='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 figures')

    args = parser.parse_args()

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

    return args


def precompute_individual_metrics(individual_df):
    individual_metrics = {}

    for _, row in tqdm.tqdm(individual_df.iterrows(), total=len(individual_df), desc="Individual metrics"):
        key = (row['dataset'], row['custom_dim'], row['seed'], row['sample_index'], row['label'])

        if key in individual_metrics:
            raise Exception("Multiple individual samples with the same key.")

        # Compute individual explanation metrics
        individual_importance = IndividualExplanationAnalyzer.calc_feature_importance(
            row['explanations'], len(row['sample'])
        )
        individual_ranking = IndividualExplanationAnalyzer.calc_feature_ranking(individual_importance)

        individual_metrics[key] = {
            'explanations': row['explanations'],
            'importance': individual_importance,
            'ranking': individual_ranking,
            'sample': np.array(row['sample'])
        }

    return individual_metrics


def compute_cluster_metrics(row, dataset):
    """Compute cluster-level importance, ranking, and coverage"""
    samples = [np.array(sample) for sample in row['samples']]

    importance = ClusterExplanationAnalyzer.calc_feature_importance(row['explanations'], dataset.num_features)
    ranking = ClusterExplanationAnalyzer.calc_feature_ranking(importance)
    coverage = average_score(
        row['explanations'],
        ClusterExplanationAnalyzer.calc_explanation_coverage,
        samples=samples,
        prototype=row['prototype'],
        epsilon=row['epsilon'],
        delta=row['delta']
    )

    return importance, ranking, coverage


def process_cluster_samples(row, individual_metrics, cluster_importance, cluster_ranking, model, dataset):
    """Process all samples in a cluster"""
    cluster_start = row['cluster_index'] * row['cluster_size']
    cluster_end = (row['cluster_index'] + 1) * row['cluster_size']

    importance_errors_cluster = []
    importance_errors_aggregator = []
    ranking_corrs_cluster = []
    ranking_corrs_aggregator = []
    ranking_corrs_lime = []
    ranking_corrs_shap = []

    # Collect all samples for batch processing
    samples_to_process = []
    sample_keys = []
    aggregator_explanations = []

    for sample_index in range(cluster_start, cluster_end):
        lookup_key = (row['dataset'], row['custom_dim'], row['seed'], sample_index, row['label'])

        if lookup_key in individual_metrics:
            samples_to_process.append(individual_metrics[lookup_key]['sample'])
            sample_keys.append(lookup_key)
            aggregator_explanations.append(individual_metrics[lookup_key]['explanations'])

    if not samples_to_process:
        return None

    # Compute aggregator importance and ranking
    aggregator_importance = AggregatorExplanationAnalyzer.calc_feature_importance(
        aggregator_explanations, dataset.num_features
    )
    aggregator_ranking = AggregatorExplanationAnalyzer.calc_feature_ranking(aggregator_importance)

    # Batch compute LIME and SHAP rankings
    lime_rankings = []
    shap_rankings = []

    for sample in samples_to_process:
        lime_ranking = LimeExplanationAnalyzer.calc_feature_ranking(sample, model, dataset)
        shap_ranking = ShapExplanationAnalyzer.calc_feature_ranking(sample, model, dataset)
        lime_rankings.append(lime_ranking)
        shap_rankings.append(shap_ranking)

    # Process each sample with pre-computed metrics
    for i, lookup_key in enumerate(sample_keys):
        individual_data = individual_metrics[lookup_key]

        # Compute scores using pre-computed metrics
        importance_error_cluster = IndividualExplanationAnalyzer.calc_feature_importance_error(
            cluster_importance, individual_data['importance']
        )
        importance_error_aggregator = IndividualExplanationAnalyzer.calc_feature_importance_error(
            aggregator_importance, individual_data['importance']
        )

        ranking_corr_cluster = IndividualExplanationAnalyzer.calc_feature_ranking_correlation(
            cluster_ranking, individual_data['ranking']
        )
        ranking_corr_aggregator = IndividualExplanationAnalyzer.calc_feature_ranking_correlation(
            aggregator_ranking, individual_data['ranking']
        )
        ranking_corr_lime = IndividualExplanationAnalyzer.calc_feature_ranking_correlation(
            lime_rankings[i], individual_data['ranking']
        )
        ranking_corr_shap = IndividualExplanationAnalyzer.calc_feature_ranking_correlation(
            shap_rankings[i], individual_data['ranking']
        )

        importance_errors_cluster.append(importance_error_cluster)
        importance_errors_aggregator.append(importance_error_aggregator)
        ranking_corrs_cluster.append(ranking_corr_cluster)
        ranking_corrs_aggregator.append(ranking_corr_aggregator)
        ranking_corrs_lime.append(ranking_corr_lime)
        ranking_corrs_shap.append(ranking_corr_shap)

    # Return averages
    return (
        np.mean(importance_errors_cluster) if importance_errors_cluster else 0.0,
        np.mean(importance_errors_aggregator) if importance_errors_aggregator else 0.0,
        np.mean(ranking_corrs_cluster) if ranking_corrs_cluster else 0.0,
        np.mean(ranking_corrs_aggregator) if ranking_corrs_aggregator else 0.0,
        np.mean(ranking_corrs_lime) if ranking_corrs_lime else 0.0,
        np.mean(ranking_corrs_shap) if ranking_corrs_shap else 0.0
    )


def create_stats_entries(row, dataset, metrics, coverage):
    """Create statistics entries for cluster, LIME, and SHAP"""
    importance_error_cluster, importance_error_aggregator, ranking_corr_cluster, ranking_corr_aggregator, ranking_corr_lime, ranking_corr_shap = metrics

    stats = []

    # Cluster stats
    cluster_stat = {
        "ID": row['id'],
        "Dataset": dataset.name,
        "Cluster Index": row['cluster_index'],
        "Group Size": row['cluster_size'],
        "Computation Time": row['computation_time'],
        "Number of Features": dataset.num_features,
        "Feature Importance Error": importance_error_cluster,
        "Feature Ranking Correlation": ranking_corr_cluster,
        "Coverage": coverage
    }
    stats.append(cluster_stat)

    # Aggregator stats
    aggregator_stat = {
        "ID": "Aggregated",
        "Dataset": dataset.name,
        "Cluster Index": row['cluster_index'],
        "Group Size": row['cluster_size'],
        "Number of Features": dataset.num_features,
        "Feature Importance Error": importance_error_aggregator,
        "Feature Ranking Correlation": ranking_corr_aggregator,
    }
    stats.append(aggregator_stat)

    # LIME stats
    lime_stat = {
        "ID": "LIME",
        "Dataset": dataset.name,
        "Cluster Index": row['cluster_index'],
        "Group Size": row['cluster_size'],
        "Number of Features": dataset.num_features,
        "Feature Ranking Correlation": ranking_corr_lime,
    }
    stats.append(lime_stat)

    # SHAP stats
    shap_stat = {
        "ID": "SHAP",
        "Dataset": dataset.name,
        "Cluster Index": row['cluster_index'],
        "Group Size": row['cluster_size'],
        "Number of Features": dataset.num_features,
        "Feature Ranking Correlation": ranking_corr_shap,
    }
    stats.append(shap_stat)

    return stats


def generate_plots_and_output(stats, output_dir, x_col_name):
    """Generate plots and print statistics"""
    # Plot ranking correlation (includes LIME and SHAP)
    lineplot(
        stats=stats,
        x_col_name=x_col_name,
        y_col_name="Feature Ranking Correlation",
        file_name=f"{output_dir}/ranking_correlation.pdf",
        log_scale=False
    )

    print_grouped_means(stats, [x_col_name, "ID"], ["Feature Ranking Correlation"])

    # Filter out LIME and SHAP for importance and coverage plots
    stats = [stat for stat in stats if stat['ID'] not in ("SHAP", "LIME")]

    # Plot importance error
    lineplot(
        stats=stats,
        x_col_name=x_col_name,
        y_col_name="Feature Importance Error",
        file_name=f"{output_dir}/importance_error.pdf",
        log_scale=False
    )

    print_grouped_means(stats, [x_col_name, "ID"], ["Feature Importance Error"])

    # Filter out Aggregated for coverage plot
    stats = [stat for stat in stats if stat['ID'] != "Aggregated"]

    lineplot(
        stats=stats,
        x_col_name=x_col_name,
        y_col_name="Coverage",
        file_name=f"{output_dir}/coverage.pdf",
        log_scale=False
    )

    print_grouped_means(stats, [x_col_name, "ID"], ["Coverage"])

if __name__ == "__main__":
    main()
