"""
Step 5: Evaluate paper recommendation models on survey-curated reading lists.

This runs the honest evaluation against expert-curated data from Step 4.

KEY DIFFERENCE FROM OLD EVAL:
  Old: "Can you find ANY cited paper among easy negatives?" → inflated nDCG
  New: "Can you rank essential papers above background AND hard negatives?" → honest

METRICS:
  - nDCG@10/20: Overall ranking quality (multi-graded relevance)
  - MAP: Average precision across all positions
  - MRR (tier≥3): Where does the first essential/relevant paper appear?
  - Recall@k (tier≥3): How many important papers make it into top-k?
  - Hard Negative AUC: Can model separate cited from expert-excluded papers?
  - Essential nDCG@10: Can model find the must-read papers specifically?
  - Relevant nDCG@10: Can model find Related Work papers?
  - Per-section accuracy: Does model understand Related Work > Methods > Intro?

BASELINES:
  - Random: The absolute floor
  - Citation count: Pure popularity (what your current model approximates)
  - Recency: Newer papers score higher
  - Cosine similarity: BGE-M3 embedding distance to survey paper

USAGE:
  # Run baselines only (no model needed):
  python 05_evaluate_on_surveys.py \
    --eval-file eval_v2_data/eval_survey_reading_lists.parquet \
    --baselines-only

  # Evaluate LightGBM model (needs Qdrant + Turso for features):
  python 05_evaluate_on_surveys.py \
    --eval-file eval_v2_data/eval_survey_reading_lists.parquet \
    --model-file production_model/reranker_v2.txt \
    --qdrant-url $QDRANT_URL --qdrant-api-key $QDRANT_API_KEY \
    --turso-url $TURSO_URL --turso-token $TURSO_DB_TOKEN

PREREQUISITES:
  pip install lightgbm pyarrow numpy tqdm

Author: ResearchIT ML Pipeline — Eval V2
"""
from __future__ import annotations

import argparse
import json
import os
from collections import defaultdict
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq
from tqdm import tqdm


# ── Metrics ──────────────────────────────────────────────────────────────────

def ndcg_at_k(labels: np.ndarray, scores: np.ndarray, k: int = 10) -> float:
    """Compute nDCG@k for a single query (multi-graded relevance)."""
    if len(labels) == 0:
        return 0.0
    order = np.argsort(-scores)
    sorted_labels = labels[order][:k].astype(np.float64)

    gains = (2.0 ** sorted_labels) - 1.0
    discounts = np.log2(np.arange(len(sorted_labels)) + 2.0)
    dcg = np.sum(gains / discounts)

    ideal_order = np.argsort(-labels)
    ideal_labels = labels[ideal_order][:k].astype(np.float64)
    ideal_gains = (2.0 ** ideal_labels) - 1.0
    ideal_discounts = np.log2(np.arange(len(ideal_labels)) + 2.0)
    idcg = np.sum(ideal_gains / ideal_discounts)

    return float(dcg / idcg) if idcg > 0 else 0.0


def recall_at_k(labels: np.ndarray, scores: np.ndarray, k: int, threshold: int = 3) -> float:
    """Fraction of papers with label >= threshold that appear in top-k."""
    total_relevant = np.sum(labels >= threshold)
    if total_relevant == 0:
        return 0.0
    order = np.argsort(-scores)
    sorted_labels = labels[order]
    top_k_relevant = np.sum(sorted_labels[:k] >= threshold)
    return float(top_k_relevant / total_relevant)


def mean_reciprocal_rank(labels: np.ndarray, scores: np.ndarray, threshold: int = 3) -> float:
    """1/rank of first paper with label >= threshold."""
    order = np.argsort(-scores)
    sorted_labels = labels[order]
    for rank, label in enumerate(sorted_labels, 1):
        if label >= threshold:
            return 1.0 / rank
    return 0.0


def average_precision(labels: np.ndarray, scores: np.ndarray, threshold: int = 1) -> float:
    """AP: average of precision at each relevant document position."""
    order = np.argsort(-scores)
    sorted_labels = labels[order]
    relevant = sorted_labels >= threshold
    n_relevant = np.sum(relevant)
    if n_relevant == 0:
        return 0.0

    precisions = []
    running_relevant = 0
    for i, is_rel in enumerate(relevant, 1):
        if is_rel:
            running_relevant += 1
            precisions.append(running_relevant / i)
    return float(np.mean(precisions))


def hard_negative_auc(labels: np.ndarray, scores: np.ndarray) -> float:
    """
    AUC for cited (label > 0) vs hard negative (label = 0).
    Measures: can the model tell expert-included from expert-excluded?
    """
    pos_scores = scores[labels > 0]
    neg_scores = scores[labels == 0]

    if len(pos_scores) == 0 or len(neg_scores) == 0:
        return 0.5

    # Efficient AUC via Mann-Whitney U statistic
    concordant = 0
    for ps in pos_scores:
        concordant += np.sum(ps > neg_scores) + 0.5 * np.sum(ps == neg_scores)
    total = len(pos_scores) * len(neg_scores)
    return float(concordant / total) if total > 0 else 0.5


def tier_specific_ndcg(labels: np.ndarray, scores: np.ndarray, target_tier: int, k: int = 10) -> float:
    """nDCG@k where only papers at or above target_tier are relevant."""
    binary_labels = (labels >= target_tier).astype(np.float64)
    if np.sum(binary_labels) == 0:
        return 0.0
    return ndcg_at_k(binary_labels, scores, k)


# ── Evaluation Engine ────────────────────────────────────────────────────────

def evaluate_model(
    eval_file: str,
    score_fn,
    model_name: str = "Model",
) -> dict:
    """
    Run full evaluation on survey reading lists.

    Args:
        eval_file: path to eval_survey_reading_lists.parquet
        score_fn: callable(survey_id, cited_ids, labels) → np.ndarray of scores
        model_name: name for display

    Returns:
        dict with all metrics
    """
    table = pq.read_table(eval_file)

    survey_ids = table.column("survey_arxiv_id").to_pylist()
    cited_ids = table.column("cited_arxiv_id").to_pylist()
    labels_list = table.column("label").to_pylist()

    # Group by survey
    survey_groups: dict[str, list[int]] = defaultdict(list)
    for i, sid in enumerate(survey_ids):
        survey_groups[sid].append(i)

    # Compute per-query metrics
    metrics_per_query = {
        "ndcg@10": [], "ndcg@20": [],
        "recall@10": [], "recall@20": [],
        "mrr": [], "map": [],
        "hard_neg_auc": [],
        "essential_ndcg@10": [], "relevant_ndcg@10": [],
    }

    for survey_id, indices in tqdm(survey_groups.items(), desc=f"Evaluating {model_name}"):
        indices = np.array(indices)
        query_labels = np.array([labels_list[i] for i in indices], dtype=np.int32)
        query_cited_ids = [cited_ids[i] for i in indices]

        # Skip if no positives
        if np.sum(query_labels > 0) == 0:
            continue

        # Get scores from the model
        try:
            query_scores = score_fn(survey_id, query_cited_ids, query_labels)
        except Exception as e:
            continue

        if query_scores is None or len(query_scores) != len(query_labels):
            continue

        query_scores = np.asarray(query_scores, dtype=np.float64)

        # Compute all metrics
        metrics_per_query["ndcg@10"].append(ndcg_at_k(query_labels, query_scores, k=10))
        metrics_per_query["ndcg@20"].append(ndcg_at_k(query_labels, query_scores, k=20))
        metrics_per_query["recall@10"].append(recall_at_k(query_labels, query_scores, k=10, threshold=3))
        metrics_per_query["recall@20"].append(recall_at_k(query_labels, query_scores, k=20, threshold=3))
        metrics_per_query["mrr"].append(mean_reciprocal_rank(query_labels, query_scores, threshold=3))
        metrics_per_query["map"].append(average_precision(query_labels, query_scores, threshold=1))
        metrics_per_query["hard_neg_auc"].append(hard_negative_auc(query_labels, query_scores))

        if np.sum(query_labels >= 4) > 0:
            metrics_per_query["essential_ndcg@10"].append(
                tier_specific_ndcg(query_labels, query_scores, 4, k=10))
        if np.sum(query_labels >= 3) > 0:
            metrics_per_query["relevant_ndcg@10"].append(
                tier_specific_ndcg(query_labels, query_scores, 3, k=10))

    # Aggregate
    results = {"num_queries": len(metrics_per_query["ndcg@10"])}
    for metric_name, values in metrics_per_query.items():
        if values:
            results[metric_name] = float(np.mean(values))
            results[f"{metric_name}_std"] = float(np.std(values))
        else:
            results[metric_name] = 0.0
            results[f"{metric_name}_std"] = 0.0

    return results


# ── Baseline Scorers ─────────────────────────────────────────────────────────

class RandomScorer:
    """Random baseline — the absolute floor."""
    def __init__(self, seed=42):
        self.rng = np.random.default_rng(seed)

    def __call__(self, survey_id, cited_ids, labels):
        return self.rng.random(len(cited_ids))


class CitationCountScorer:
    """Pure popularity baseline (approximates what V1/V2 model learned)."""
    def __init__(self, metadata_cache: dict):
        self.cache = metadata_cache

    def __call__(self, survey_id, cited_ids, labels):
        scores = []
        for cid in cited_ids:
            meta = self.cache.get(cid, {})
            scores.append(float(meta.get("citation_count", 0)))
        return np.array(scores)


class RecencyScorer:
    """Newer papers score higher."""
    def __init__(self, metadata_cache: dict):
        self.cache = metadata_cache

    def __call__(self, survey_id, cited_ids, labels):
        scores = []
        for cid in cited_ids:
            meta = self.cache.get(cid, {})
            try:
                year = int(str(meta.get("update_date", "2020"))[:4])
            except (ValueError, TypeError):
                year = 2020
            scores.append(float(year))
        return np.array(scores)


class OracleScorer:
    """Perfect ranking — the ceiling (uses labels directly)."""
    def __call__(self, survey_id, cited_ids, labels):
        # Add small noise to break ties consistently
        return np.array(labels, dtype=np.float64) + np.random.default_rng(42).random(len(labels)) * 0.01


# ── Report Formatting ────────────────────────────────────────────────────────

def print_comparison_report(results: dict[str, dict]):
    """Print formatted comparison table."""
    print(f"\n{'='*90}")
    print(f"  EVALUATION RESULTS — Survey Reading List Benchmark (Eval V2)")
    print(f"{'='*90}")

    models = list(results.keys())
    # Key metrics to display
    display_metrics = [
        "num_queries",
        "ndcg@10", "ndcg@20",
        "recall@10", "recall@20",
        "mrr", "map",
        "hard_neg_auc",
        "essential_ndcg@10", "relevant_ndcg@10",
    ]

    # Header
    col_width = 14
    header = f"  {'Metric':<25}" + "".join(f"{m:>{col_width}}" for m in models)
    print(header)
    print("  " + "-" * (25 + col_width * len(models)))

    for metric in display_metrics:
        if metric == "num_queries":
            row = f"  {metric:<25}" + "".join(
                f"{results[m].get(metric, 0):>{col_width}}" for m in models)
        else:
            row = f"  {metric:<25}" + "".join(
                f"{results[m].get(metric, 0):>{col_width}.4f}" for m in models)
        print(row)

    print("  " + "-" * (25 + col_width * len(models)))

    # Insights
    print(f"\n  INTERPRETATION:")
    print(f"    nDCG@10:          Overall ranking quality (higher = better)")
    print(f"    hard_neg_auc:     Can model tell cited from expert-excluded? (0.5 = random, 1.0 = perfect)")
    print(f"    essential_ndcg:   Can model find the must-read papers? (most important metric)")
    print(f"    recall@10:        Of all important papers, how many are in top-10?")

    # Flag issues
    for model, r in results.items():
        if model in ("Random", "Oracle"):
            continue
        if r.get("hard_neg_auc", 0) < 0.6:
            print(f"\n  ⚠️  {model}: hard_neg_auc={r['hard_neg_auc']:.3f} — model barely distinguishes cited from excluded!")
        if r.get("essential_ndcg@10", 0) < 0.3:
            print(f"  ⚠️  {model}: essential_ndcg={r['essential_ndcg@10']:.3f} — fails to find must-read papers!")


# ── CLI ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Evaluate paper recommendation models (Eval V2)")
    parser.add_argument("--eval-file", required=True,
                        help="eval_survey_reading_lists.parquet from Step 4")
    parser.add_argument("--output", default=None, help="Save results as JSON")
    parser.add_argument("--baselines-only", action="store_true",
                        help="Only run baselines (Random + Oracle)")
    parser.add_argument("--model-file", default=None,
                        help="LightGBM model .txt file to evaluate")
    parser.add_argument("--qdrant-url", default=os.environ.get("QDRANT_URL"))
    parser.add_argument("--qdrant-api-key", default=os.environ.get("QDRANT_API_KEY"))
    parser.add_argument("--turso-url", default=os.environ.get("TURSO_URL"))
    parser.add_argument("--turso-token", default=os.environ.get("TURSO_DB_TOKEN"))

    args = parser.parse_args()

    # Load eval data
    print(f"Loading eval data: {args.eval_file}")
    table = pq.read_table(args.eval_file)
    n_rows = len(table)
    n_surveys = len(set(table.column("survey_arxiv_id").to_pylist()))
    labels = table.column("label").to_pylist()
    label_dist = defaultdict(int)
    for l in labels:
        label_dist[l] += 1

    print(f"  {n_rows} rows, {n_surveys} surveys")
    print(f"  Labels: 4={label_dist[4]}, 3={label_dist[3]}, 2={label_dist[2]}, 1={label_dist[1]}, 0={label_dist[0]}")

    results = {}

    # Always run Random and Oracle
    print("\n--- Running: Random baseline ---")
    results["Random"] = evaluate_model(args.eval_file, RandomScorer(42), "Random")

    print("\n--- Running: Oracle (perfect ranking) ---")
    results["Oracle"] = evaluate_model(args.eval_file, OracleScorer(), "Oracle")

    if not args.baselines_only and args.model_file:
        # TODO: Load LightGBM model and compute features
        # This requires Qdrant (for cosine scores) and Turso (for metadata)
        # Will be implemented once the eval data is generated
        print(f"\n--- LightGBM evaluation requires feature computation ---")
        print(f"    Model: {args.model_file}")
        print(f"    Qdrant: {args.qdrant_url}")
        print(f"    Turso: {args.turso_url}")
        print(f"    (Not yet implemented — needs Qdrant + Turso access)")

    # Print comparison
    print_comparison_report(results)

    # Save results
    if args.output:
        with open(args.output, "w") as f:
            json.dump(results, f, indent=2)
        print(f"\nResults saved to: {args.output}")
    else:
        # Default output location
        output_dir = Path(args.eval_file).parent
        output_file = output_dir / "eval_results.json"
        with open(output_file, "w") as f:
            json.dump(results, f, indent=2)
        print(f"\nResults saved to: {output_file}")


if __name__ == "__main__":
    main()