"""
evaluation/evaluate_model.py
==============================
PHASE 5: Evaluation — How Good Is Your Fine-Tuned Model?

CONCEPT: WHY EVALUATE?
═══════════════════════════════════════════════════════════════════════════
Training loss decreasing doesn't mean your model is actually useful.
You need to measure REAL quality:
  - Does it give correct answers?
  - Does it stay on-topic?
  - Does it hallucinate?
  - Is it faster or slower than the base model?

EVALUATION METRICS USED:
═══════════════════════════════════════════════════════════════════════════

1. ROUGE-L (Recall-Oriented Understudy for Gisting Evaluation)
   - Measures word overlap between generated and reference answers
   - ROUGE-L: longest common subsequence (order matters)
   - Range: 0–1 (1 = perfect match)
   - Good for: Measuring if key information is present
   - Limitation: Misses semantically equivalent phrasings

2. BLEU (Bilingual Evaluation Understudy)
   - Originally for machine translation
   - Measures n-gram precision (how many n-grams in the output appear in reference)
   - Range: 0–1
   - Limitation: Poor for long answers; prefers short precise outputs

3. BERTScore
   - Uses BERT embeddings to measure semantic similarity
   - Unlike ROUGE, understands synonyms and paraphrases
   - "tuition fees" ≈ "payment amount" → high BERTScore, low ROUGE
   - Range: 0–1 (F1 score of precision/recall in embedding space)
   - Best metric for conversational/descriptive answers

4. PERPLEXITY
   - Measures how surprised the model is by the text
   - Lower = model finds the text natural/familiar
   - Fine-tuned model should have lower perplexity on domain text
   - Formula: exp(-1/N × Σ log P(token_i))

5. LATENCY
   - Time to generate a response (seconds)
   - Important for production deployment

6. HALLUCINATION RATE
   - % of answers containing information NOT in the reference
   - We approximate this by checking if key facts from the reference
     are present in the generated answer
"""

import json
import time
from pathlib import Path
from typing import Optional

try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

try:
    from rouge_score import rouge_scorer as _rouge_scorer_mod
    ROUGE_AVAILABLE = True
except ImportError:
    ROUGE_AVAILABLE = False

try:
    from bert_score import score as _bert_score_fn
    BERTSCORE_AVAILABLE = True
except ImportError:
    BERTSCORE_AVAILABLE = False


# ══════════════════════════════════════════════════════════════════════════════
# EVALUATION DATASET
# ══════════════════════════════════════════════════════════════════════════════

EVAL_QUESTIONS = [
    {
        "question": "What are the minimum entry requirements for Computer Science at MUST?",
        "reference": "UACE with at least 2 principal passes including Mathematics. UCE with at least 5 passes including Mathematics and English. Minimum aggregate of 15 points at A-Level.",
        "category": "admissions",
        "key_facts": ["2 principal passes", "Mathematics", "UCE", "15 points"],
    },
    {
        "question": "How much are tuition fees per semester for Computer Science?",
        "reference": "UGX 1,800,000 to 2,200,000 per semester for private sponsorship students.",
        "category": "fees",
        "key_facts": ["1,800,000", "2,200,000", "semester", "private"],
    },
    {
        "question": "What is the minimum attendance requirement?",
        "reference": "Students must attend at least 75% of all lectures for each course. Below 75% results in being barred from sitting the final exam.",
        "category": "policies",
        "key_facts": ["75%", "barred", "final exam"],
    },
    {
        "question": "What are the graduation requirements?",
        "reference": "Complete 120-135 credit units, minimum GPA of 2.0, submit and pass the final year project, complete 8-week industrial training, clear all fees and library dues.",
        "category": "graduation",
        "key_facts": ["120", "GPA", "2.0", "project", "industrial training"],
    },
    {
        "question": "What is First Class Honours?",
        "reference": "First Class Honours is the highest academic classification, awarded for a Cumulative GPA of 4.4 to 5.0.",
        "category": "graduation",
        "key_facts": ["4.4", "5.0", "highest", "First Class"],
    },
    {
        "question": "Can I defer my studies?",
        "reference": "Yes, deferment is allowed for medical reasons, financial hardship, pregnancy, family emergency. Maximum deferment is 2 consecutive semesters. Requires formal approval from the Academic Registrar.",
        "category": "policies",
        "key_facts": ["medical", "financial hardship", "2 consecutive", "Academic Registrar"],
    },
    {
        "question": "What programming languages are taught in Computer Science?",
        "reference": "Year 1: Python and C. Year 2: Java, SQL, JavaScript. Year 3: PHP, R or MATLAB. Year 4: Advanced Python, Kotlin/Swift for mobile.",
        "category": "courses",
        "key_facts": ["Python", "Java", "SQL", "JavaScript"],
    },
    {
        "question": "How many times can I retake a failed exam?",
        "reference": "Maximum 3 attempts per course. Failing a core course 3 times typically results in discontinuation. Supplementary exams available for scores 40-49%. Maximum grade for supplementary is 50%.",
        "category": "policies",
        "key_facts": ["3 attempts", "supplementary", "40-49", "50%"],
    },
]


# ══════════════════════════════════════════════════════════════════════════════
# METRIC CALCULATORS
# ══════════════════════════════════════════════════════════════════════════════

def calculate_rouge(prediction: str, reference: str) -> dict:
    """
    Calculates ROUGE-1, ROUGE-2, and ROUGE-L scores.

    ROUGE-1: Unigram (single word) overlap
    ROUGE-2: Bigram (two-word phrase) overlap
    ROUGE-L: Longest common subsequence (captures order)
    """
    if not ROUGE_AVAILABLE:
        print("  ⚠️  rouge_score not installed. Run: pip install rouge-score")
        return {"rouge1": None, "rouge2": None, "rougeL": None}

    scorer = _rouge_scorer_mod.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True)
    scores = scorer.score(reference, prediction)
    return {
        "rouge1": round(scores["rouge1"].fmeasure, 4),
        "rouge2": round(scores["rouge2"].fmeasure, 4),
        "rougeL": round(scores["rougeL"].fmeasure, 4),
    }


def calculate_bert_score(predictions: list, references: list) -> dict:
    """
    Calculates BERTScore for a batch of predictions vs references.

    BERTScore uses contextual embeddings (BERT) to measure semantic similarity.
    This is better than ROUGE for conversational text because it understands
    synonyms and paraphrases.

    For example:
        Prediction: "You need excellent grades"
        Reference: "You must achieve high marks"
        ROUGE: Low (few exact word matches)
        BERTScore: High (same semantic meaning)
    """
    if not BERTSCORE_AVAILABLE:
        print("  ⚠️  bert_score not installed. Run: pip install bert-score")
        return {"precision": None, "recall": None, "f1": None}

    P, R, F1 = _bert_score_fn(
        predictions,
        references,
        lang="en",
        model_type="distilbert-base-uncased",  # Fast, decent quality
        verbose=False,
    )
    return {
        "precision": round(P.mean().item(), 4),
        "recall": round(R.mean().item(), 4),
        "f1": round(F1.mean().item(), 4),
    }


def calculate_hallucination_rate(predictions: list, eval_items: list) -> float:
    """
    Approximates hallucination rate by checking if key facts are present.

    This is a simplified metric. Production systems use NLI (Natural Language
    Inference) models to check factual consistency.

    Our approach: for each answer, check what % of key_facts are present.
    A "hallucination" is present if a key fact is completely absent but the
    answer mentions something plausible but wrong.

    Since we can't automatically detect wrong-but-plausible facts,
    we use key_fact RECALL as a proxy for non-hallucination:
    - High recall = model mentions the important facts → probably accurate
    - Low recall = model avoided or missed facts → may have made things up
    """
    total_recall = 0.0

    if not predictions or not eval_items:
        return 0.0  # No predictions → no hallucinations to measure

    for prediction, item in zip(predictions, eval_items):
        key_facts = item.get("key_facts", [])
        if not key_facts:
            continue

        prediction_lower = prediction.lower()
        found = sum(1 for fact in key_facts if fact.lower() in prediction_lower)
        recall = found / len(key_facts)
        total_recall += recall

    # Return estimated hallucination rate = 1 - average recall
    evaluated = sum(1 for item in eval_items if item.get("key_facts"))
    avg_recall = total_recall / evaluated if evaluated > 0 else 1.0
    hallucination_rate = 1.0 - avg_recall
    return round(hallucination_rate, 4)


# ══════════════════════════════════════════════════════════════════════════════
# MODEL INFERENCE
# ══════════════════════════════════════════════════════════════════════════════

def generate_answers(model_path: str, questions: list, model_label: str = "model") -> tuple:
    """
    Generates answers for all eval questions using the specified model.

    Returns:
        (answers, latencies) — lists of strings and seconds
    """
    if not TORCH_AVAILABLE:
        print("⚠️  torch not installed: pip install torch")
        return [], []

    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer
    except ImportError:
        print("⚠️  transformers not installed: pip install transformers")
        return [], []

    print(f"\n🤖 Generating answers with {model_label}...")
    print(f"   Model: {model_path}")

    # Load model
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto" if torch.cuda.is_available() else "cpu",
    )
    model.eval()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    answers = []
    latencies = []

    system_prompt = "You are CampusGPT Uganda, a helpful university assistant."

    for i, question in enumerate(questions):
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": question},
        ]

        # Format with chat template
        input_text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
        inputs = tokenizer(input_text, return_tensors="pt").to(device)

        # Time the generation
        start_time = time.time()
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=300,
                temperature=0.3,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
            )
        latency = time.time() - start_time

        # Decode
        new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
        answer = tokenizer.decode(new_tokens, skip_special_tokens=True)

        answers.append(answer.strip())
        latencies.append(round(latency, 2))

        print(f"  [{i+1}/{len(questions)}] Generated in {latency:.1f}s")

    # Clean up memory
    del model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return answers, latencies


# ══════════════════════════════════════════════════════════════════════════════
# FULL EVALUATION
# ══════════════════════════════════════════════════════════════════════════════

def evaluate_model(
    model_path: str,
    model_label: str = "CampusGPT Fine-tuned",
    eval_questions: list = None,
) -> dict:
    """
    Runs the complete evaluation suite for one model.
    """
    if eval_questions is None:
        eval_questions = EVAL_QUESTIONS

    questions = [q["question"] for q in eval_questions]
    references = [q["reference"] for q in eval_questions]

    # Generate answers
    answers, latencies = generate_answers(model_path, questions, model_label)

    if not answers:
        return {}

    # Calculate metrics
    print("\n📊 Calculating metrics...")

    # ROUGE scores (per question, then averaged) — returns None values if rouge_score not installed
    rouge_scores = [calculate_rouge(pred, ref) for pred, ref in zip(answers, references)]
    def _safe_avg(scores, key):
        vals = [s[key] for s in scores if s.get(key) is not None]
        return round(sum(vals) / len(vals), 4) if vals else None

    avg_rouge = {
        "rouge1": _safe_avg(rouge_scores, "rouge1"),
        "rouge2": _safe_avg(rouge_scores, "rouge2"),
        "rougeL": _safe_avg(rouge_scores, "rougeL"),
    }

    # BERTScore
    bert_scores = calculate_bert_score(answers, references)

    # Hallucination rate
    hallucination_rate = calculate_hallucination_rate(answers, eval_questions)

    # Latency
    avg_latency = round(sum(latencies) / len(latencies), 2)

    results = {
        "model": model_label,
        "model_path": model_path,
        "num_questions": len(questions),
        "metrics": {
            "rouge1": avg_rouge["rouge1"],
            "rouge2": avg_rouge["rouge2"],
            "rougeL": avg_rouge["rougeL"],
            "bertscore_f1": bert_scores["f1"],
            "bertscore_precision": bert_scores["precision"],
            "bertscore_recall": bert_scores["recall"],
            "hallucination_rate": hallucination_rate,
            "avg_latency_seconds": avg_latency,
        },
        "per_question": [
            {
                "question": q,
                "answer": a,
                "reference": r,
                "rouge": rs,
                "latency": lt,
            }
            for q, a, r, rs, lt in zip(questions, answers, references, rouge_scores, latencies)
        ],
    }

    return results


def compare_models(base_model_path: str, finetuned_model_path: str):
    """
    Runs a side-by-side comparison: Base Model vs Fine-tuned Model.
    This is the key experiment that validates fine-tuning worked.
    """
    print("🔬 CampusGPT Model Evaluation")
    print("=" * 60)
    print("Comparing: Base Model vs Fine-tuned CampusGPT")
    print()

    # Evaluate base model
    base_results = evaluate_model(base_model_path, "Base Model (no fine-tuning)")

    # Evaluate fine-tuned model
    ft_results = evaluate_model(finetuned_model_path, "CampusGPT (fine-tuned)")

    if not base_results or not ft_results:
        print("❌ Evaluation failed — check model paths")
        return

    # ── Print comparison report ───────────────────────────────────────────────
    print("\n" + "=" * 60)
    print("📊 EVALUATION REPORT")
    print("=" * 60)

    metrics = ["rouge1", "rouge2", "rougeL", "bertscore_f1", "hallucination_rate", "avg_latency_seconds"]
    better_is = {"hallucination_rate": "lower", "avg_latency_seconds": "lower"}  # Others: higher

    print(f"\n{'Metric':<25} {'Base Model':>15} {'Fine-tuned':>15} {'Change':>10}")
    print("-" * 65)

    for metric in metrics:
        base_val = base_results["metrics"].get(metric)
        ft_val = ft_results["metrics"].get(metric)

        if base_val is None or ft_val is None:
            print(f"{metric:<25} {'N/A':>15} {'N/A':>15} {'(dep missing)':>10}")
            continue

        change = ft_val - base_val
        direction = better_is.get(metric, "higher")
        if direction == "higher":
            indicator = "✅" if change > 0.01 else ("🔴" if change < -0.01 else "➡️")
        else:
            indicator = "✅" if change < -0.01 else ("🔴" if change > 0.01 else "➡️")

        print(f"{metric:<25} {base_val:>15.4f} {ft_val:>15.4f} {indicator} {change:>+.4f}")

    # ── Per-question breakdown ────────────────────────────────────────────────
    print("\n\n📝 PER-QUESTION BREAKDOWN")
    print("-" * 60)

    for i, (base_q, ft_q) in enumerate(zip(base_results["per_question"], ft_results["per_question"]), 1):
        print(f"\nQ{i}: {base_q['question'][:60]}...")
        print(f"  Base:       {base_q['answer'][:100]}...")
        print(f"  Fine-tuned: {ft_q['answer'][:100]}...")
        rougeL_base = base_q['rouge'].get('rougeL')
        rougeL_ft   = ft_q['rouge'].get('rougeL')
        if rougeL_base is not None and rougeL_ft is not None:
            print(f"  ROUGE-L:    Base={rougeL_base:.3f} | Fine-tuned={rougeL_ft:.3f}")
        else:
            print("  ROUGE-L:    N/A (install rouge-score)")

    # ── Save report ───────────────────────────────────────────────────────────
    def _safe_diff(a, b):
        if a is None or b is None:
            return None
        return round(b - a, 4)

    report = {
        "base_model": base_results,
        "finetuned_model": ft_results,
        "summary": {
            "rouge_improvement": _safe_diff(
                base_results["metrics"].get("rougeL"),
                ft_results["metrics"].get("rougeL"),
            ),
            "bertscore_improvement": _safe_diff(
                base_results["metrics"].get("bertscore_f1"),
                ft_results["metrics"].get("bertscore_f1"),
            ),
            "hallucination_reduction": _safe_diff(
                ft_results["metrics"].get("hallucination_rate"),
                base_results["metrics"].get("hallucination_rate"),
            ),
        }
    }

    report_path = "evaluation/evaluation_report.json"
    Path("evaluation").mkdir(exist_ok=True)
    with open(report_path, "w") as f:
        json.dump(report, f, indent=2, default=str)

    print(f"\n\n✅ Full report saved to: {report_path}")

    # Print summary
    print("\n🎯 SUMMARY")
    for k, v in report["summary"].items():
        val_str = f"{v:+.4f}" if isinstance(v, float) else "N/A (install deps)"
        print(f"  {k}: {val_str}")

    r_imp = report["summary"]["rouge_improvement"]
    b_imp = report["summary"]["bertscore_improvement"]
    if r_imp is not None and b_imp is not None and r_imp > 0 and b_imp > 0:
        print("\n  ✅ Fine-tuning improved the model!")
    elif r_imp is None:
        print("\n  ℹ️  Install rouge-score and bert-score for full comparison: pip install rouge-score bert-score")
    else:
        print("\n  ⚠️  Results are mixed. Consider more training epochs or more data.")


if __name__ == "__main__":
    import sys

    if len(sys.argv) >= 3:
        base = sys.argv[1]
        finetuned = sys.argv[2]
        compare_models(base, finetuned)
    elif len(sys.argv) == 2:
        # Evaluate single model
        results = evaluate_model(sys.argv[1], "CampusGPT")
        if results:
            print(json.dumps(results["metrics"], indent=2))
    else:
        print("Usage:")
        print("  python evaluation/evaluate_model.py <base_model_path> <finetuned_model_path>")
        print("  python evaluation/evaluate_model.py <single_model_path>")
        print("\nExample:")
        print("  python evaluation/evaluate_model.py unsloth/Qwen2.5-3B-Instruct models/campusgpt_merged")