#!/usr/bin/env python3
"""
extract_subset.py
Extract samples where 'system' mode underperformed baselines, create a subset
dataset for running hybrid mode, and produce a 6-mode comparison table.

Usage:
  # Step 1: Extract bad samples and create subset dataset
  python extract_subset.py --dir 281_results_final --top 50 --output hybrid_subset/

  # Step 2: Run hybrid mode on the subset
  python run_benchmark.py --input hybrid_subset/evaluation_subset.json --modes hybrid --model gpt-5-mini-2025-08-07

  # Step 3: Compare all 6 modes on the subset
  python extract_subset.py --dir hybrid_subset --compare
"""

import sys
import os
import json
import argparse
from glob import glob
from tabulate import tabulate

# Add parent dir to path to import metrics.py
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.abspath(os.path.join(current_dir, '..'))
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

try:
    from metrics import MetricsCalculator
except ImportError:
    print("ERROR: Could not import MetricsCalculator from metrics.py.")
    print("Please make sure metrics.py is located in the parent directory.")
    sys.exit(1)


def calculate_error_score(item):
    """Calculates an absolute 'badness' score for a single evaluation result."""
    score = 0.0
    if item.get('pred_bias') is not None and item.get('gold_bias') is not None:
        score += abs(item['pred_bias'] - item['gold_bias']) * 0.15
    if item.get('pred_factuality') is not None and item.get('gold_factuality') is not None:
        score += abs(item['pred_factuality'] - item['gold_factuality']) * 0.15
    score += (1.0 - item.get('factscore_precision', 0)) * 1.5
    score += (1.0 - item.get('factscore_recall', 0)) * 1.0
    score += (1.0 - item.get('meteor', 0)) * 0.5
    return score


def aggregate_results(results):
    """Aggregates a list of results using the official metrics calculator."""
    n = len(results)
    if n == 0: return None

    # Filter out None scores to pass to the official MetricsCalculator
    valid_bias_preds = [
        (r.get('gold_bias'), r.get('pred_bias')) 
        for r in results 
        if r.get('pred_bias') is not None and r.get('gold_bias') is not None
    ]
    valid_fact_preds = [
        (r.get('gold_factuality'), r.get('pred_factuality')) 
        for r in results 
        if r.get('pred_factuality') is not None and r.get('gold_factuality') is not None
    ]

    # Use the official EMNLP 2024 ordinal MAE logic from metrics.py
    bias_mae = MetricsCalculator.calculate_bias_mae(
        [x[0] for x in valid_bias_preds], [x[1] for x in valid_bias_preds]
    ) if valid_bias_preds else None

    fact_mae = MetricsCalculator.calculate_factuality_mae(
        [x[0] for x in valid_fact_preds], [x[1] for x in valid_fact_preds]
    ) if valid_fact_preds else None

    # Safely handle metrics that might be missing or None
    avg_fs_precision = sum(r.get('factscore_precision') or 0 for r in results) / n
    avg_fs_recall = sum(r.get('factscore_recall') or 0 for r in results) / n
    avg_error_rate = sum(r.get('error_rate') or 0 for r in results) / n
    avg_meteor = sum(r.get('meteor') or 0 for r in results) / n
    avg_rougeL = sum(r.get('rougeL') or 0 for r in results) / n

    # FC detection
    fc_applicable = [r for r in results if r.get('fact_check_hit', {}).get('status') != 'not_applicable']
    fc_detection = sum(1 for r in fc_applicable if r.get('fact_check_hit', {}).get('status') == 'found') / len(fc_applicable) if fc_applicable else float('nan')

    return {
        "N": n,
        "FS Prec.": f"{round(avg_fs_precision * 100, 1)}%",
        "FS Recall": f"{round(avg_fs_recall * 100, 1)}%",
        "Err Rate": f"{round(avg_error_rate * 100, 1)}%",
        "METEOR": f"{round(avg_meteor * 100, 1)}%",
        "ROUGE-L": f"{round(avg_rougeL * 100, 1)}%",
        "Bias MAE": round(bias_mae, 2) if bias_mae is not None else "N/A",
        "Fact. MAE": round(fact_mae, 2) if fact_mae is not None else "N/A",
        "FC Det.": f"{round(fc_detection * 100, 1)}%" if str(fc_detection) != 'nan' else "N/A",
    }


def load_results(results_dir):
    """Load all JSONL results from a directory, grouped by outlet name."""
    jsonl_files = [
        f for f in glob(os.path.join(results_dir, "*.jsonl"))
        if "summary" not in os.path.basename(f)
    ]

    outlets_data = {}
    for filepath in jsonl_files:
        filename = os.path.basename(filepath)
        model = filename.split('_')[0]

        with open(filepath, 'r', encoding='utf-8') as f:
            for line in f:
                if not line.strip():
                    continue
                item = json.loads(line)
                name = item.get('name')
                mode = item.get('mode', filename.replace('.jsonl', '').split('_')[-1])

                if name not in outlets_data:
                    outlets_data[name] = {}
                outlets_data[name][mode] = item
                outlets_data[name][mode]['_model'] = model
                outlets_data[name][mode]['_filepath'] = filepath

    return outlets_data


def calculate_deltas(outlets_data):
    """Calculate delta scores: (System Error) - (Average Baseline Error).

    Returns list of (delta, system_score, avg_other_score, name) sorted descending.
    """
    error_rankings = []
    for name, modes_dict in outlets_data.items():
        if 'system' in modes_dict:
            system_score = calculate_error_score(modes_dict['system'])
            
            # Calculate how all the OTHER modes did on this same sample
            other_scores = [
                calculate_error_score(item) 
                for mode, item in modes_dict.items() 
                if mode != 'system'
            ]
            
            if other_scores:
                avg_other_score = sum(other_scores) / len(other_scores)
                # Delta = System Error - Avg Baseline Error
                # High positive delta means 'system' did uniquely bad
                delta = system_score - avg_other_score 
            else:
                avg_other_score = 0.0
                delta = 0.0
                
            error_rankings.append((delta, system_score, avg_other_score, name))
        else:
            error_rankings.append((0.0, 0.0, 0.0, name))

    # Sort so the highest positive Delta (where we are worse) is at the top
    error_rankings.sort(key=lambda x: x[0], reverse=True)
    return error_rankings


def cmd_extract(args):
    """Extract bad samples and create a subset dataset for hybrid mode."""
    print(f"Loading results from: {args.dir}")
    outlets_data = load_results(args.dir)
    print(f"Found {len(outlets_data)} outlets across results")

    error_rankings = calculate_deltas(outlets_data)

    # Determine which outlets to extract
    if args.delta_threshold is not None:
        # Extract all outlets where system delta exceeds threshold
        selected = [(d, s, o, n) for d, s, o, n in error_rankings if d > args.delta_threshold]
        print(f"\nSelecting outlets with delta > {args.delta_threshold}: {len(selected)} outlets")
    elif args.top > 0:
        # Extract top N worst outlets
        selected = error_rankings[:args.top]
        print(f"\nSelecting top {args.top} worst-performing outlets")
    else:
        # Default: all outlets where system is worse (delta > 0)
        selected = [(d, s, o, n) for d, s, o, n in error_rankings if d > 0]
        print(f"\nSelecting all outlets where system underperforms (delta > 0): {len(selected)} outlets")

    selected_names = set(n for _, _, _, n in selected)

    if not selected_names:
        print("No outlets selected. Try adjusting --top or --delta-threshold.")
        return

    # Create output directory
    os.makedirs(args.output, exist_ok=True)

    # 1. Create subset evaluation dataset from the full dataset
    eval_dataset_path = os.path.join(current_dir, "evaluation_dataset.json")
    if not os.path.exists(eval_dataset_path):
        print(f"ERROR: evaluation_dataset.json not found at {eval_dataset_path}")
        return

    with open(eval_dataset_path, 'r') as f:
        full_dataset = json.load(f)

    subset_dataset = [item for item in full_dataset if item.get('name') in selected_names]
    subset_path = os.path.join(args.output, "evaluation_subset.json")
    with open(subset_path, 'w') as f:
        json.dump(subset_dataset, f, indent=2, default=str)
    print(f"\nSaved evaluation subset: {subset_path} ({len(subset_dataset)} items)")

    # 2. Copy existing mode results for the subset outlets
    copied_modes = set()
    for name in selected_names:
        for mode, item in outlets_data.get(name, {}).items():
            if mode.startswith('_'):
                continue
            model = item.get('_model', 'unknown')
            tag = f"{model}_{mode}"
            outfile = os.path.join(args.output, f"{tag}.jsonl")
            copied_modes.add(mode)

            # Append this outlet's result to the mode-specific file
            # Remove internal keys before writing
            clean_item = {k: v for k, v in item.items() if not k.startswith('_')}
            with open(outfile, 'a') as f:
                f.write(json.dumps(clean_item, default=str) + "\n")

    print(f"Copied results for modes: {sorted(copied_modes)}")

    # 3. Print extraction summary
    print(f"\n{'='*70}")
    print(f"EXTRACTION SUMMARY")
    print(f"{'='*70}")
    print(f"Total outlets in results: {len(outlets_data)}")
    print(f"Outlets extracted: {len(selected_names)}")
    print(f"Output directory: {args.output}")
    print(f"\nTop 10 extracted outlets (by delta):")
    for delta, sys_score, oth_score, name in selected[:10]:
        print(f"  - {name}")
        print(f"    Delta: +{round(delta, 2)} (System Err: {round(sys_score, 2)} vs Others Avg: {round(oth_score, 2)})")

    print(f"\nNext steps:")
    print(f"  1. Run hybrid mode:")
    print(f"     python run_benchmark.py --input {subset_path} --modes hybrid --model <model>")
    print(f"  2. Copy hybrid results to {args.output}/")
    print(f"  3. Compare all 6 modes:")
    print(f"     python extract_subset.py --dir {args.output} --compare")


def cmd_compare(args):
    """Compare all modes (including hybrid) on the subset."""
    print(f"Loading results from: {args.dir}")
    outlets_data = load_results(args.dir)
    print(f"Found {len(outlets_data)} outlets")

    # Collect all modes
    all_modes = set()
    for name, modes_dict in outlets_data.items():
        for mode in modes_dict:
            if not mode.startswith('_'):
                all_modes.add(mode)

    print(f"Modes found: {sorted(all_modes)}")

    # Aggregate results by (model, mode)
    mode_results = {}
    for name, modes_dict in outlets_data.items():
        for mode, item in modes_dict.items():
            if mode.startswith('_'):
                continue
            model = item.get('_model', item.get('model', 'unknown'))
            key = (model, mode)
            if key not in mode_results:
                mode_results[key] = []
            mode_results[key].append(item)

    # Build summary table
    summary_rows = []
    for (model, mode), results_list in mode_results.items():
        agg = aggregate_results(results_list)
        if agg:
            row = {"Model": model, "Mode": mode}
            row.update(agg)
            summary_rows.append(row)

    summary_rows.sort(key=lambda x: (x['Model'], x['Mode']))

    # Print table
    headers = [
        "Model", "Mode", "N",
        "FS Prec.", "FS Recall", "Err Rate",
        "METEOR", "ROUGE-L",
        "Bias MAE", "Fact. MAE", "FC Det."
    ]
    table_data = [[r.get(h, "") for h in headers] for r in summary_rows]

    print(f"\n{'='*80}")
    print("6-MODE COMPARISON TABLE (Subset)")
    print(f"{'='*80}")
    print(tabulate(table_data, headers=headers, tablefmt="grid"))

    # Save comparison table
    tsv_path = os.path.join(args.dir, "comparison_table.tsv")
    with open(tsv_path, 'w') as f:
        f.write("\t".join(headers) + "\n")
        for row in table_data:
            f.write("\t".join(str(x) for x in row) + "\n")
    print(f"\nSaved to {tsv_path}")

    # Per-outlet breakdown for system vs hybrid (if both exist)
    if 'system' in all_modes and 'hybrid' in all_modes:
        print(f"\n{'='*80}")
        print("SYSTEM vs HYBRID: Per-Outlet Delta")
        print(f"{'='*80}")

        comparisons = []
        for name, modes_dict in outlets_data.items():
            if 'system' in modes_dict and 'hybrid' in modes_dict:
                sys_err = calculate_error_score(modes_dict['system'])
                hyb_err = calculate_error_score(modes_dict['hybrid'])
                improvement = sys_err - hyb_err  # positive = hybrid is better
                comparisons.append((improvement, name, sys_err, hyb_err))

        comparisons.sort(key=lambda x: x[0], reverse=True)

        improved = sum(1 for imp, _, _, _ in comparisons if imp > 0)
        degraded = sum(1 for imp, _, _, _ in comparisons if imp < 0)
        unchanged = sum(1 for imp, _, _, _ in comparisons if imp == 0)

        print(f"  Hybrid improved: {improved}/{len(comparisons)} outlets")
        print(f"  Hybrid degraded: {degraded}/{len(comparisons)} outlets")
        print(f"  Unchanged: {unchanged}/{len(comparisons)} outlets")

        if comparisons:
            avg_improvement = sum(imp for imp, _, _, _ in comparisons) / len(comparisons)
            print(f"  Average improvement: {round(avg_improvement, 3)}")

            print(f"\n  Top 5 improvements:")
            for imp, name, sys_err, hyb_err in comparisons[:5]:
                print(f"    {name}: {round(imp, 3)} (sys={round(sys_err, 2)} -> hyb={round(hyb_err, 2)})")

            print(f"\n  Top 5 degradations:")
            for imp, name, sys_err, hyb_err in comparisons[-5:]:
                print(f"    {name}: {round(imp, 3)} (sys={round(sys_err, 2)} -> hyb={round(hyb_err, 2)})")


def main():
    parser = argparse.ArgumentParser(
        description="Extract subset where system underperforms, run hybrid mode, compare results."
    )
    parser.add_argument("--dir", type=str, required=True,
                        help="Directory with existing JSONL result files")
    parser.add_argument("--output", type=str, default="hybrid_subset",
                        help="Output directory for subset (default: hybrid_subset)")
    parser.add_argument("--top", type=int, default=0,
                        help="Extract top N worst outlets (0 = all with delta > 0)")
    parser.add_argument("--delta-threshold", type=float, default=None,
                        help="Only extract outlets with delta > this value")
    parser.add_argument("--compare", action="store_true",
                        help="Compare all modes (run after hybrid eval is complete)")

    args = parser.parse_args()

    # Required imports: json, argparse, os, sys, glob, tabulate.
    # Note that tabulate needs to be pip installed: `pip install tabulate`.
    if args.compare:
        cmd_compare(args)
    else:
        cmd_extract(args)


if __name__ == "__main__":
    main()