"""
evaluate_model.py
-----------------
Compares base Phi-2 vs fine-tuned Phi-2 on sample radiology cases.

Usage:
    python training/evaluate_model.py
"""

import json
import logging
import sys
import time
from pathlib import Path

import torch

logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(message)s")
logger = logging.getLogger(__name__)

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

ADAPTER_DIR = ROOT / "models" / "slm_model" / "lora_adapter"

TEST_CASES = [
    {
        "label": "Chest X-ray — Pneumonia",
        "prompt": (
            "### System:\nYou are a board-certified radiologist.\n\n"
            "### Instruction:\nGenerate a structured radiology report.\n"
            "Patient: 52-year-old Male. Symptoms: fever, productive cough.\n"
            "Study: X-ray of the Chest, PA view.\n"
            "Write FINDINGS, IMPRESSION, and RECOMMENDATIONS sections.\n\n"
            "### Response:\nFINDINGS: The chest demonstrates"
        ),
    },
    {
        "label": "Pelvis X-ray — Post Hip Arthroplasty",
        "prompt": (
            "### System:\nYou are a board-certified radiologist.\n\n"
            "### Instruction:\nGenerate a structured radiology report.\n"
            "Patient: 68-year-old Female. Symptoms: hip pain after fall.\n"
            "Study: X-ray of the Pelvis, AP view.\n"
            "Write FINDINGS, IMPRESSION, and RECOMMENDATIONS sections.\n\n"
            "### Response:\nFINDINGS: The pelvis demonstrates"
        ),
    },
]


def run_inference(model, tokenizer, prompt: str, max_new_tokens: int = 300) -> str:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)

    start = time.time()
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.3,
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.2,
            pad_token_id=tokenizer.eos_token_id,
        )
    elapsed = time.time() - start

    new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
    text = tokenizer.decode(new_tokens, skip_special_tokens=True)
    return text, elapsed


def evaluate():
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    )

    logger.info("Loading base Phi-2...")
    tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    base_model = AutoModelForCausalLM.from_pretrained(
        "microsoft/phi-2",
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
    )
    base_model.eval()

    # Load fine-tuned if available
    finetuned_model = None
    if ADAPTER_DIR.exists() and (ADAPTER_DIR / "adapter_config.json").exists():
        logger.info("Loading fine-tuned adapter...")
        from peft import PeftModel
        finetuned_model = PeftModel.from_pretrained(base_model, str(ADAPTER_DIR))
        finetuned_model.eval()
        logger.info("Fine-tuned model loaded!")
    else:
        logger.warning(f"No fine-tuned adapter found at {ADAPTER_DIR}")
        logger.warning("Run train_unsloth.py first to fine-tune the model")

    # Evaluate
    results = []
    print("\n" + "="*70)
    print("EVALUATION RESULTS")
    print("="*70)

    for case in TEST_CASES:
        print(f"\n📋 {case['label']}")
        print("-" * 70)

        # Base model
        base_output, base_time = run_inference(base_model, tokenizer, case["prompt"])
        print(f"\n🔵 BASE MODEL ({base_time:.1f}s):")
        print("FINDINGS: The " + case["prompt"].split("demonstrates")[-1][:50] + "...")
        print(base_output[:400])

        # Fine-tuned model
        if finetuned_model:
            ft_output, ft_time = run_inference(finetuned_model, tokenizer, case["prompt"])
            print(f"\n🟢 FINE-TUNED MODEL ({ft_time:.1f}s):")
            print(ft_output[:400])

        results.append({
            "case": case["label"],
            "base_output": base_output,
            "finetuned_output": finetuned_model and ft_output or "N/A",
        })

    # Save results
    out_path = Path(__file__).parent / "evaluation_results.json"
    out_path.write_text(json.dumps(results, indent=2))
    logger.info(f"\nResults saved to {out_path}")


if __name__ == "__main__":
    evaluate()