#!/usr/bin/env python3
"""§19 Demo evaluation: baseline vs trained model comparison.

Loads the base model (no training) and a checkpoint, runs NUM_EVAL_EPISODES
identical fixed episodes through both, then outputs an accuracy/reward
comparison table and saves full transcripts to episode_comparison.json.

Usage:
    # Auto-detects latest checkpoint_ep* directory
    python eval_baseline.py

    # Or point at a specific checkpoint
    CHECKPOINT_PATH=./checkpoint_ep300 python eval_baseline.py

    # Adjust episode count
    NUM_EVAL_EPISODES=20 python eval_baseline.py
"""

import copy
import json
import os
import sys
import random

import numpy as np
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline,
)

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from crime_env.environment import CrimeInvestigationEnv
from crime_env.case_generator import generate_case

# ── Configuration ──────────────────────────────────────────────────────────

MODEL_NAME       = os.environ.get("MODEL_NAME",       "Qwen/Qwen2.5-1.5B-Instruct")
NPC_MODEL_NAME   = os.environ.get("NPC_MODEL_NAME",   "Qwen/Qwen2.5-0.5B-Instruct")
CHECKPOINT_PATH  = os.environ.get("CHECKPOINT_PATH",  "")
NUM_EVAL_EPISODES = int(os.environ.get("NUM_EVAL_EPISODES", "10"))
DEVICE           = "cuda" if torch.cuda.is_available() else "cpu"
SEED             = 42
OUTPUT_FILE      = "episode_comparison.json"


# ── Helpers ─────────────────────────────────────────────────────────────────

def _quant_config():
    if DEVICE != "cuda":
        return None
    try:
        return BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
        )
    except Exception:
        return None


def _load_base_model():
    """Load the original pretrained model without any LoRA / fine-tuning."""
    print(f"  Loading base model: {MODEL_NAME}")
    qc = _quant_config()
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        quantization_config=qc,
        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
        device_map="auto" if DEVICE == "cuda" else None,
        trust_remote_code=True,
    )
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    return model, tokenizer


def _load_checkpoint_model(checkpoint_path: str):
    """Load a PEFT-trained checkpoint.  Falls back to base model if not found."""
    if not checkpoint_path or not os.path.isdir(checkpoint_path):
        print(f"  Checkpoint not found at '{checkpoint_path}'. Using base model as fallback.")
        return _load_base_model()

    print(f"  Loading checkpoint: {checkpoint_path}")
    try:
        from peft import PeftModel
        base, tokenizer = _load_base_model()
        model = PeftModel.from_pretrained(base, checkpoint_path)
        model.eval()
        return model, tokenizer
    except Exception as e:
        print(f"  PEFT load failed ({e}). Trying plain HF load.")

    qc = _quant_config()
    model = AutoModelForCausalLM.from_pretrained(
        checkpoint_path,
        quantization_config=qc,
        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
        device_map="auto" if DEVICE == "cuda" else None,
        trust_remote_code=True,
    )
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(checkpoint_path, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    return model, tokenizer


def _load_npc_pipeline():
    qc  = _quant_config()
    tok = AutoTokenizer.from_pretrained(NPC_MODEL_NAME, trust_remote_code=True)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    npc = AutoModelForCausalLM.from_pretrained(
        NPC_MODEL_NAME,
        quantization_config=qc,
        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
        device_map="auto" if DEVICE == "cuda" else None,
        trust_remote_code=True,
    )
    npc.eval()
    return pipeline(
        "text-generation", model=npc, tokenizer=tok,
        max_new_tokens=80, do_sample=True, temperature=0.7, top_p=0.9,
    )


def _make_npc_call(npc_pipe):
    tok = npc_pipe.tokenizer

    def llm_call(system_prompt: str, conversation_history: list) -> str:
        user_prompt = "Conversation so far:\n"
        for entry in conversation_history[-8:]:
            user_prompt += f"{entry.get('speaker', '')}: {entry.get('content', '')[:180]}\n"
        messages = [
            {"role": "system", "content": system_prompt[:800]},
            {"role": "user",   "content": user_prompt[-1800:] + "\nRespond in 1-2 sentences."},
        ]
        if hasattr(tok, "apply_chat_template"):
            prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        else:
            prompt = f"System: {messages[0]['content']}\n\nUser: {messages[1]['content']}\n\nAssistant:"
        try:
            out  = npc_pipe(prompt, return_full_text=False)
            resp = out[0]["generated_text"].strip()
            return resp.split("\n")[0][:300] if resp else "I have nothing to add."
        except Exception:
            return "I don't recall anything specific about that."

    return llm_call


def _generate_action(model, tokenizer, obs: dict, max_turns: int) -> str:
    """Greedy-decode one action (no sampling variance for fair comparison)."""
    history = obs.get("conversation_history", [])
    prompt  = (
        f"You are a detective investigating a crime.\n"
        f"Briefing: {obs['briefing'][:300]}\n"
        f"Turn: {obs['turn']}/{max_turns}\nRecent conversation:\n"
    )
    for entry in history[-6:]:
        prompt += f"  {entry['speaker']}: {entry['content'][:100]}\n"
    prompt += (
        "\nChoose ONE action:\n"
        "ACTION: ask_question | TARGET: Suspect_A | CONTENT: <question>\n"
        "ACTION: request_evidence | ITEM: keycard_log\n"
        "ACTION: accuse | TARGET: Suspect_A\n\nYour action:"
    )

    if hasattr(tokenizer, "apply_chat_template"):
        messages = [
            {"role": "system", "content": "You are a detective. Choose your next action."},
            {"role": "user",   "content": prompt},
        ]
        prompt_text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
    else:
        prompt_text = f"Detective: {prompt}\n\nAction:"

    inputs = tokenizer(prompt_text, return_tensors="pt", truncation=True, max_length=1024)
    if DEVICE == "cuda":
        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=56,
            do_sample=False,        # greedy — reproducible across both runs
            pad_token_id=tokenizer.pad_token_id,
        )

    response = tokenizer.decode(
        out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
    ).strip()

    if not response.upper().startswith("ACTION:"):
        turn = obs.get("turn", 0)
        force_at = max(1, int(max_turns * 0.7))
        if turn >= force_at:
            return "ACTION: accuse | TARGET: Suspect_A"
        return "ACTION: ask_question | TARGET: Suspect_A | CONTENT: Where were you at the time of the crime?"

    return response


# ── Episode runner ──────────────────────────────────────────────────────────

def run_episodes(
    model,
    tokenizer,
    npc_call,
    fixed_cases: list[dict],
    label: str,
) -> list[dict]:
    """Run fixed_cases through the environment with the given model.

    Returns a list of per-episode result dicts.
    """
    env     = CrimeInvestigationEnv(llm_call=npc_call)
    results = []

    for i, case in enumerate(fixed_cases):
        max_turns    = case.get("max_turns", 15)
        env.MAX_TURNS = max_turns
        obs  = env.reset(case_data=copy.deepcopy(case))
        done = False
        turn_count = 0

        while not done:
            action = _generate_action(model, tokenizer, obs, max_turns=max_turns)
            obs, _, done, info = env.step(action)
            turn_count += 1

        final_r = env.reward_calc.get_rewards()
        result  = "correct" if info.get("action") == "accuse" and info.get("correct") \
                  else ("wrong" if info.get("action") == "accuse" else "timeout")

        row = {
            "episode":          i + 1,
            "criminal":         case["criminal"],
            "result":           result,
            "correct":          result == "correct",
            "detective_reward": round(float(final_r.get("detective", 0.0)), 4),
            "turns":            turn_count,
            "conversation_history": list(env.conversation_history),
        }
        results.append(row)
        print(
            f"    Ep {i+1:>2}: {result:<7} | "
            f"reward={row['detective_reward']:>+6.2f} | turns={turn_count}"
        )

    return results


# ── Main ────────────────────────────────────────────────────────────────────

def main():
    print("=" * 60)
    print("  AI Crime Investigation — Baseline vs Trained Evaluation")
    print("=" * 60)

    # Auto-detect latest checkpoint if not specified
    checkpoint_path = CHECKPOINT_PATH
    if not checkpoint_path:
        candidates = sorted(
            [d for d in os.listdir(".") if d.startswith("checkpoint_ep") and os.path.isdir(d)],
            key=lambda x: int(x.replace("checkpoint_ep", "") or "0"),
        )
        checkpoint_path = candidates[-1] if candidates else ""

    print(f"Checkpoint : {checkpoint_path or '(none — will clone base for both runs)'}")
    print(f"Episodes   : {NUM_EVAL_EPISODES}")

    # Reproducible fixed cases — same seed for both conditions
    random.seed(SEED)
    np.random.seed(SEED)
    fixed_cases = [generate_case(difficulty="hard") for _ in range(NUM_EVAL_EPISODES)]

    print("\nLoading NPC model ...")
    npc_pipe = _load_npc_pipeline()
    npc_call = _make_npc_call(npc_pipe)

    # ── Baseline run ───────────────────────────────────────────────────
    print("\n[1/2] BASE MODEL (untrained) ...")
    base_model, base_tok = _load_base_model()
    base_results = run_episodes(base_model, base_tok, npc_call, fixed_cases, label="base")
    del base_model
    if DEVICE == "cuda":
        torch.cuda.empty_cache()

    # ── Trained run ────────────────────────────────────────────────────
    print(f"\n[2/2] TRAINED MODEL ({checkpoint_path or 'base fallback'}) ...")
    trained_model, trained_tok = _load_checkpoint_model(checkpoint_path)
    trained_results = run_episodes(trained_model, trained_tok, npc_call, fixed_cases, label="trained")
    del trained_model
    if DEVICE == "cuda":
        torch.cuda.empty_cache()

    # ── Comparison table ───────────────────────────────────────────────
    def _stats(rows):
        acc        = sum(r["correct"] for r in rows) / len(rows)
        mean_r     = float(np.mean([r["detective_reward"] for r in rows]))
        avg_turns  = float(np.mean([r["turns"] for r in rows]))
        return acc, mean_r, avg_turns

    b_acc, b_r, b_t = _stats(base_results)
    t_acc, t_r, t_t = _stats(trained_results)

    print("\n" + "=" * 62)
    print("  EVALUATION RESULTS")
    print("=" * 62)
    print(f"  {'Metric':<32} {'Base':>9} {'Trained':>9} {'Δ':>9}")
    print(f"  {'-'*62}")
    print(f"  {'Accuracy (%)':<32} {b_acc*100:>9.1f} {t_acc*100:>9.1f} {(t_acc-b_acc)*100:>+9.1f}")
    print(f"  {'Mean detective reward':<32} {b_r:>9.3f} {t_r:>9.3f} {t_r-b_r:>+9.3f}")
    print(f"  {'Avg turns to accuse':<32} {b_t:>9.1f} {t_t:>9.1f} {t_t-b_t:>+9.1f}")
    print("=" * 62)

    # ── Save full results ──────────────────────────────────────────────
    output = {
        "config": {
            "model":        MODEL_NAME,
            "checkpoint":   checkpoint_path or "base",
            "num_episodes": NUM_EVAL_EPISODES,
            "seed":         SEED,
        },
        "summary": {
            "base":    {"accuracy": b_acc, "mean_reward": b_r, "avg_turns": b_t},
            "trained": {"accuracy": t_acc, "mean_reward": t_r, "avg_turns": t_t},
        },
        "base_episodes":    base_results,
        "trained_episodes": trained_results,
    }
    with open(OUTPUT_FILE, "w") as f:
        json.dump(output, f, indent=2)

    print(f"\nFull transcripts saved → {OUTPUT_FILE}")


if __name__ == "__main__":
    main()