#!/usr/bin/env python3
"""
Compare multiple analysis CSV files to determine which most closely resembles a reference file.
Uses Euclidean distance, correlation, and RMSE metrics.

Readings files are canonicalized to the current bicorder terminology (history of
renames in bicorder_common.py), and comparison runs on the gradient columns
shared by all files — so any versions can be compared, and renamed gradients
remain comparable across version boundaries.

Usage:
  # Legacy audit: manual review vs. the three model test runs (as in README)
  python3 scripts/compare_analyses.py

  # Explicit: reference file first, then any number of comparison files
  python3 scripts/compare_analyses.py \
      data/synthetic_1.2.6/readings.csv data/synthetic_1.4.0/readings.csv
"""

import argparse
import sys
import pandas as pd
import numpy as np
from scipy.stats import pearsonr
from pathlib import Path

from bicorder_common import apply_renames, csv_version


def load_canonical(path):
    """Load a readings CSV, canonicalize columns, and coerce gradient values to numeric."""
    df = pd.read_csv(path, quotechar='"', escapechar='\\', engine='python')
    df = apply_renames(df)
    numeric_cols = [col for col in df.columns if
                    col.startswith(('Design_', 'Entanglement_', 'Experience_'))]
    for col in numeric_cols:
        df[col] = pd.to_numeric(df[col], errors='coerce')
    return df, numeric_cols


def calculate_euclidean_distance(df1, df2, numeric_cols):
    """Calculate Euclidean distance between two dataframes."""
    distances = []
    for idx in df1.index:
        diff = df1.loc[idx, numeric_cols] - df2.loc[idx, numeric_cols]
        # Use nansum to ignore NaN values
        distance = np.sqrt(np.nansum(diff ** 2))
        distances.append(distance)
    return np.array(distances)

def calculate_rmse(df1, df2, numeric_cols):
    """Calculate Root Mean Squared Error."""
    diff = df1[numeric_cols] - df2[numeric_cols]
    # Use nanmean to ignore NaN values
    mse = np.nanmean(diff.values ** 2)
    return np.sqrt(mse)

def calculate_correlation(df1, df2, numeric_cols):
    """Calculate Pearson correlation across all numeric values."""
    vals1 = df1[numeric_cols].values.flatten()
    vals2 = df2[numeric_cols].values.flatten()

    # Remove NaN values - only use positions where both have valid values
    mask = ~(np.isnan(vals1) | np.isnan(vals2))
    vals1_clean = vals1[mask]
    vals2_clean = vals2[mask]

    if len(vals1_clean) < 2:
        return np.nan, np.nan

    corr, pvalue = pearsonr(vals1_clean, vals2_clean)
    return corr, pvalue

def compare_analyses(reference_file, comparison_files):
    """Compare multiple analysis files to a reference file."""

    # Read and canonicalize reference file
    print(f"Reading reference file: {reference_file}")
    ref_version = csv_version(reference_file)
    if ref_version:
        print(f"  bicorder version recorded: v{ref_version}")
    ref_df, ref_numeric = load_canonical(reference_file)

    # Repeated Descriptor entries (kept as control cases in the datasets) would
    # multiply rows in the Descriptor-based merge; keep first occurrence like
    # the classifier does.
    if 'Descriptor' in ref_df.columns:
        before = len(ref_df)
        ref_df = ref_df.drop_duplicates(subset='Descriptor', keep='first')
        if len(ref_df) < before:
            print(f"  Deduplicated reference: {before} → {len(ref_df)} rows (kept first of repeated Descriptor)")

    print(f"\nFound {len(ref_numeric)} numeric dimensions in reference file")
    print(f"Comparing {len(ref_df)} protocols\n")
    print("="*80)

    results = {}

    for comp_file in comparison_files:
        print(f"\nComparing: {Path(comp_file).name}")
        print("-"*80)

        # Read and canonicalize comparison file
        comp_version = csv_version(comp_file)
        if comp_version:
            print(f"  bicorder version recorded: v{comp_version}")
        comp_df, comp_numeric = load_canonical(comp_file)
        if 'Descriptor' in comp_df.columns:
            before = len(comp_df)
            comp_df = comp_df.drop_duplicates(subset='Descriptor', keep='first')
            if len(comp_df) < before:
                print(f"  Deduplicated comparison: {before} → {len(comp_df)} rows (kept first of repeated Descriptor)")

        # Restrict to columns shared by both files (post-rename): enables
        # comparing across bicorder versions when gradients were renamed
        numeric_cols = [col for col in ref_numeric if col in comp_numeric]
        missing = [col for col in ref_numeric if col not in comp_numeric]
        if missing:
            print(f"  Note: {len(missing)} gradient(s) absent here are excluded: {', '.join(missing)}")

        # Ensure same protocols in same order (match by Descriptor)
        if 'Descriptor' in ref_df.columns and 'Descriptor' in comp_df.columns:
            # Use merge to ensure exact matching - only keep protocols in ref_df
            comp_df = pd.merge(
                ref_df[['Descriptor']],
                comp_df,
                on='Descriptor',
                how='left'
            )

        # Calculate Euclidean distances using reset indices to ensure alignment
        ref_temp = ref_df.reset_index(drop=True)
        comp_temp = comp_df.reset_index(drop=True)
        euclidean_distances = calculate_euclidean_distance(ref_temp, comp_temp, numeric_cols)
        total_euclidean = np.sum(euclidean_distances)
        avg_euclidean = np.mean(euclidean_distances)

        # Calculate RMSE
        rmse = calculate_rmse(ref_temp, comp_temp, numeric_cols)

        # Calculate correlation
        correlation, p_value = calculate_correlation(ref_temp, comp_temp, numeric_cols)

        # Store results
        results[Path(comp_file).name] = {
            'total_euclidean': total_euclidean,
            'avg_euclidean': avg_euclidean,
            'rmse': rmse,
            'correlation': correlation,
            'p_value': p_value,
            'per_protocol_distances': euclidean_distances,
            'protocols': ref_df['Descriptor'].values if 'Descriptor' in ref_df.columns else None
        }

        # Print results
        print(f"  Total Euclidean Distance:   {total_euclidean:.2f}")
        print(f"  Average Euclidean Distance: {avg_euclidean:.2f}")
        print(f"  RMSE:                       {rmse:.2f}")
        print(f"  Pearson Correlation:        {correlation:.4f} (p={p_value:.2e})")

    # Summary comparison
    print("\n" + "="*80)
    print("SUMMARY RANKING (lower distance = more similar)")
    print("="*80)

    # Sort by average Euclidean distance
    sorted_by_euclidean = sorted(results.items(), key=lambda x: x[1]['avg_euclidean'])

    print("\nBy Average Euclidean Distance:")
    for i, (name, data) in enumerate(sorted_by_euclidean, 1):
        print(f"  {i}. {name:30s} - Avg Distance: {data['avg_euclidean']:.2f}")

    # Sort by correlation (higher is better)
    sorted_by_corr = sorted(results.items(), key=lambda x: x[1]['correlation'], reverse=True)

    print("\nBy Correlation (higher = more similar):")
    for i, (name, data) in enumerate(sorted_by_corr, 1):
        print(f"  {i}. {name:30s} - Correlation: {data['correlation']:.4f}")

    # Sort by RMSE
    sorted_by_rmse = sorted(results.items(), key=lambda x: x[1]['rmse'])

    print("\nBy RMSE (lower = more similar):")
    for i, (name, data) in enumerate(sorted_by_rmse, 1):
        print(f"  {i}. {name:30s} - RMSE: {data['rmse']:.2f}")

    # Show protocols with largest differences for the best match
    print("\n" + "="*80)
    best_match_name, best_match_data = sorted_by_euclidean[0]
    print(f"Top 10 protocols with largest differences from {best_match_name}:")
    print("="*80)

    if best_match_data['protocols'] is not None:
        distances = best_match_data['per_protocol_distances']
        protocols = best_match_data['protocols']
        top_diff_indices = np.argsort(distances)[-10:][::-1]

        for idx in top_diff_indices:
            print(f"  {protocols[idx]:50s} - Distance: {distances[idx]:.2f}")

    return results

def main(argv=None):
    """CLI entry point.

    With no arguments, falls back to the legacy audit: the 1.2.6 manual review
    against the three model test runs (as described in README.md).
    """
    legacy_reference = "data/synthetic_1.2.6/readings_manual.csv"
    legacy_comparisons = [
        "data/synthetic_1.2.6/readings_gemma3-12b.csv",
        "data/synthetic_1.2.6/readings_gpt-oss.csv",
        "data/synthetic_1.2.6/readings_mistral.csv",
    ]

    if argv is None:
        argv = sys.argv[1:]
    if argv:
        parser = argparse.ArgumentParser(
            description='Compare readings CSVs to a reference (Euclidean distance, RMSE, correlation)',
            epilog="""Example (cross-version):
  python3 scripts/compare_analyses.py \\
      data/synthetic_1.4.0/readings.csv data/synthetic_1.2.6/readings.csv
""",
        )
        parser.add_argument('reference', help='Reference readings CSV')
        parser.add_argument('comparisons', nargs='+', help='Comparison readings CSVs')
        args = parser.parse_args(argv)
        reference_file, comparison_files = args.reference, args.comparisons
    else:
        reference_file, comparison_files = legacy_reference, legacy_comparisons

    # Check if files exist
    if not Path(reference_file).exists():
        print(f"Error: Reference file '{reference_file}' not found")
        sys.exit(1)

    existing = [file for file in comparison_files if Path(file).exists()]
    for file in comparison_files:
        if not Path(file).exists():
            print(f"Warning: Comparison file '{file}' not found, skipping...")

    if not existing:
        print("Error: No comparison files found")
        sys.exit(1)

    # Run comparison
    results = compare_analyses(reference_file, existing)

    print("\n" + "="*80)
    print("Analysis complete!")
    print("="*80)
    return results


if __name__ == "__main__":
    main()