#!/usr/bin/env python3
"""
Unified evaluation framework for labor arbitration element extraction.
Compares: rules, BERT model, Ollama zero-shot, hybrid.
Produces thesis-ready metrics tables and per-field breakdowns.

Usage:
  cd nlp-service/training
  PYTHONPATH="..;../../backend" python evaluate.py \
      --dataset ../../backend/data/augmented_dataset.json \
      --output ../../backend/data/eval_results.json
"""

from __future__ import annotations

import argparse
import json
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any

BACKEND = Path(__file__).resolve().parent.parent.parent / "backend"
sys.path.insert(0, str(BACKEND))

from sklearn.metrics import (
    classification_report,
    confusion_matrix,
    mean_absolute_error,
    mean_squared_error,
)


def _as_set(value: Any) -> set:
    """Convert a value to a set for comparison."""
    if value is None:
        return set()
    if isinstance(value, list):
        return set(str(v).strip() for v in value if v)
    if isinstance(value, dict):
        # For claims dict, extract items and amount_total
        items = set()
        if "items" in value:
            items.update(str(v).strip() for v in value["items"] if v)
        if "amount_total" in value and value["amount_total"] is not None:
            items.add(str(value["amount_total"]))
        return items
    s = str(value).strip()
    return {s} if s else set()


def _as_float(value: Any) -> float | None:
    """Convert a value to float for numeric comparison."""
    if value is None:
        return None
    try:
        return float(value)
    except (ValueError, TypeError):
        return None


def compute_field_f1(gold: Any, pred: Any) -> dict[str, float]:
    """Compute P/R/F1 for a single field using set comparison."""
    gold_set = _as_set(gold)
    pred_set = _as_set(pred)

    tp = len(gold_set & pred_set)
    fp = len(pred_set - gold_set)
    fn = len(gold_set - pred_set)

    precision = tp / max(tp + fp, 1)
    recall = tp / max(tp + fn, 1)
    f1 = 2 * precision * recall / max(precision + recall, 1e-8)

    return {"precision": precision, "recall": recall, "f1": f1, "tp": tp, "fp": fp, "fn": fn}


def evaluate_all_fields(
    golds: list[dict],
    preds: list[dict],
    field_names: list[str],
) -> dict[str, Any]:
    """Evaluate all fields across all cases. Returns per-field and micro-average metrics."""
    per_field: dict[str, dict] = defaultdict(lambda: {"tp": 0, "fp": 0, "fn": 0})

    for gold, pred in zip(golds, preds):
        for field in field_names:
            g = gold.get(field)
            p = pred.get(field)
            scores = compute_field_f1(g, p)
            for k in ("tp", "fp", "fn"):
                per_field[field][k] += scores[k]

    # Compute per-field metrics
    results = {}
    total_tp = total_fp = total_fn = 0
    for field, counts in per_field.items():
        p = counts["tp"] / max(counts["tp"] + counts["fp"], 1)
        r = counts["tp"] / max(counts["tp"] + counts["fn"], 1)
        f1 = 2 * p * r / max(p + r, 1e-8)
        results[field] = {"precision": round(p, 4), "recall": round(r, 4), "f1": round(f1, 4)}
        total_tp += counts["tp"]
        total_fp += counts["fp"]
        total_fn += counts["fn"]

    # Micro average
    micro_p = total_tp / max(total_tp + total_fp, 1)
    micro_r = total_tp / max(total_tp + total_fn, 1)
    micro_f1 = 2 * micro_p * micro_r / max(micro_p + micro_r, 1e-8)

    return {
        "per_field": results,
        "micro_avg": {
            "precision": round(micro_p, 4),
            "recall": round(micro_r, 4),
            "f1": round(micro_f1, 4),
        },
    }


def evaluate_numeric_fields(
    golds: list[dict],
    preds: list[dict],
    num_field_names: list[str],
) -> dict[str, dict[str, float]]:
    """Evaluate numeric fields with MAE and RMSE."""
    results = {}
    for field in num_field_names:
        y_true = []
        y_pred = []
        for gold, pred in zip(golds, preds):
            g = _as_float(gold.get(field))
            p = _as_float(pred.get(field))
            if g is not None and p is not None:
                y_true.append(g)
                y_pred.append(p)

        if y_true:
            results[field] = {
                "mae": round(mean_absolute_error(y_true, y_pred), 2),
                "rmse": round(mean_squared_error(y_true, y_pred) ** 0.5, 2),
                "count": len(y_true),
            }

    return results


def evaluate_cause_accuracy(
    golds: list[dict],
    preds: list[dict],
) -> dict[str, Any]:
    """Evaluate case cause classification accuracy."""
    cause_keys = ["tmpl_primary_cause", "primary_cause_type", "case_cause"]

    y_true_strs = []
    y_pred_strs = []
    for gold, pred in zip(golds, preds):
        g = None
        p = None
        for k in cause_keys:
            g = g or gold.get(k)
            p = p or pred.get(k)
        if isinstance(g, dict):
            g = g.get("type")
        if isinstance(p, dict):
            p = p.get("type")
        y_true_strs.append(str(g) if g else "未知")
        y_pred_strs.append(str(p) if p else "未知")

    labels = sorted(set(y_true_strs + y_pred_strs))
    accuracy = sum(1 for t, p in zip(y_true_strs, y_pred_strs) if t == p) / max(len(y_true_strs), 1)

    return {
        "accuracy": round(accuracy, 4),
        "labels": labels,
        "per_label": classification_report(
            y_true_strs, y_pred_strs, labels=labels,
            output_dict=True, zero_division=0,
        ),
    }


def run_full_evaluation(
    test_cases: list[dict],
    extraction_methods: dict[str, callable],
    field_names: list[str],
    numeric_field_names: list[str] | None = None,
) -> dict[str, Any]:
    """
    Run full evaluation comparing all extraction methods.

    Args:
        test_cases: List of dicts with 'text' and 'rule_elements' (gold labels)
        extraction_methods: Dict of method_name -> extractor callable
        field_names: Field names to evaluate
        numeric_field_names: Subset of fields that are numeric
    """
    if numeric_field_names is None:
        numeric_field_names = ["month_salary"]

    results = {}

    for method_name, extract_fn in extraction_methods.items():
        print(f"\n{'='*60}")
        print(f"Evaluating: {method_name}")
        print(f"{'='*60}")

        preds = []
        for case in test_cases:
            try:
                pred = extract_fn(case["text"])
                preds.append(pred)
            except Exception as e:
                print(f"  ERROR on case {case.get('case_id', '?')}: {e}")
                preds.append({})

        golds = [case.get("rule_elements", {}) for case in test_cases]

        # Field-level F1
        field_results = evaluate_all_fields(golds, preds, field_names)
        print(f"  Micro-F1: {field_results['micro_avg']['f1']:.4f}")

        # Numeric evaluation
        num_results = evaluate_numeric_fields(golds, preds, numeric_field_names)
        if num_results:
            print(f"  Numeric fields evaluated: {list(num_results.keys())}")

        # Cause accuracy
        cause_results = evaluate_cause_accuracy(golds, preds)
        print(f"  Cause accuracy: {cause_results['accuracy']:.4f}")

        results[method_name] = {
            "field_metrics": field_results,
            "numeric_metrics": num_results,
            "cause_accuracy": cause_results,
        }

    # Summary comparison table
    summary = {}
    for method, res in results.items():
        summary[method] = {
            "micro_f1": res["field_metrics"]["micro_avg"]["f1"],
            "micro_precision": res["field_metrics"]["micro_avg"]["precision"],
            "micro_recall": res["field_metrics"]["micro_avg"]["recall"],
            "cause_accuracy": res["cause_accuracy"]["accuracy"],
        }
        # Average MAE for numeric fields
        maes = [v["mae"] for v in res["numeric_metrics"].values() if "mae" in v]
        if maes:
            summary[method]["avg_mae"] = round(sum(maes) / len(maes), 2)

    results["_summary"] = summary

    print(f"\n{'='*60}")
    print("SUMMARY COMPARISON TABLE")
    print(f"{'='*60}")
    print(f"{'Method':<20} {'Micro-F1':>10} {'Precision':>10} {'Recall':>10} {'Cause Acc':>10} {'Avg MAE':>10}")
    print("-" * 65)
    for method, s in summary.items():
        print(f"{method:<20} {s['micro_f1']:>10.4f} {s['micro_precision']:>10.4f} "
              f"{s['micro_recall']:>10.4f} {s['cause_accuracy']:>10.4f} "
              f"{s.get('avg_mae', 0):>10.2f}")

    return results


# Default key field names to evaluate
DEFAULT_FIELDS = [
    "applicant_name", "respondent_name", "employer_nature",
    "worker_position", "entry_date", "leave_date",
    "month_salary", "case_cause", "contract_type",
    "termination_reason", "employment_type",
    "overtime_desc", "work_duration_text",
    "law_refs", "evidence_materials", "claims",
]

DEFAULT_NUMERIC_FIELDS = [
    "month_salary", "lr1_pay_amount", "sr_claim_amount",
    "ec_claim_amount", "dm_claim_amount",
]


def load_test_cases(dataset_path: str) -> list[dict]:
    """Load test cases from dataset (only original, non-augmented cases)."""
    with open(dataset_path, encoding="utf-8") as f:
        data = json.load(f)

    return [c for c in data["dataset"] if "source" not in c]


def main():
    parser = argparse.ArgumentParser(description="Evaluate extraction methods")
    parser.add_argument("--dataset", required=True, help="Path to dataset JSON")
    parser.add_argument("--output", default="eval_results.json", help="Output path for results")
    parser.add_argument("--methods", default="rules", help="Comma-separated: rules,bert,ollama,hybrid")
    args = parser.parse_args()

    # Load test cases (originals only)
    test_cases = load_test_cases(args.dataset)
    print(f"Loaded {len(test_cases)} test cases")

    # Build extraction methods
    methods: dict[str, callable] = {}
    method_names = [m.strip() for m in args.methods.split(",")]

    if "rules" in method_names:
        from app.extractor import RuleBasedLaborExtractor
        rule_ext = RuleBasedLaborExtractor()
        methods["rules"] = rule_ext.extract

    if "bert" in method_names or "hybrid" in method_names:
        import requests

        def bert_extract(text: str) -> dict:
            try:
                resp = requests.post(
                    "http://localhost:8001/extract",
                    json={"text": text},
                    timeout=60,
                )
                if resp.status_code == 200:
                    return resp.json()
            except Exception:
                pass
            # Fallback to rules
            from app.extractor import RuleBasedLaborExtractor
            return RuleBasedLaborExtractor().extract(text)

        methods["bert"] = bert_extract

    if "ollama" in method_names or "hybrid" in method_names:
        try:
            from app.anj import OllamaClaimsExtractor
            from app.config import settings
            ollama_ext = OllamaClaimsExtractor(
                settings.ollama_base_url,
                settings.ollama_model_name,
            )
            from app.extractor import RuleBasedLaborExtractor
            rule = RuleBasedLaborExtractor()

            def ollama_extract(text: str) -> dict:
                base = rule.extract(text)
                try:
                    base["claims"] = ollama_ext.extract_claims(text)
                    template = ollama_ext.extract_dispute_template_fields(text)
                    for k, v in template.items():
                        if v is not None and v != "":
                            base[k] = v
                except Exception:
                    pass
                return base

            methods["ollama"] = ollama_extract
        except Exception as e:
            print(f"  [SKIP] Ollama not available: {e}")

    # Run evaluation
    results = run_full_evaluation(
        test_cases,
        methods,
        DEFAULT_FIELDS,
        DEFAULT_NUMERIC_FIELDS,
    )

    # Save results
    output_path = Path(args.output)
    if not output_path.is_absolute():
        output_path = Path(__file__).resolve().parent / args.output
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps(results, ensure_ascii=False, indent=2, default=str),
        encoding="utf-8",
    )
    print(f"\nResults saved to: {output_path}")


if __name__ == "__main__":
    main()