"""
Sweep all layers to find the best steering layer, then save results.

Usage:
    python src/evaluate.py --trait taciturn
    python src/evaluate.py --trait taciturn --layers 16,20,24,28 --alpha 10

Output:
    outputs/{trait}_layer_sweep.jsonl   — one entry per layer per prompt
    outputs/{trait}_best_layer.txt      — recommended layer index
"""

import argparse
import json
import torch
from pathlib import Path
from tqdm import tqdm

from model_utils import load_config, load_model, build_chat_prompt, SteeringHook

ROOT = Path(__file__).parent.parent

# Fixed evaluation prompts that should reveal the trait clearly
EVAL_PROMPTS = [
    "最近失眠很严重，你有什么建议吗？",
    "能帮我解释一下量子纠缠是什么意思吗？",
    "我今天心情很差，你能安慰我一下吗？",
    "你觉得人生的意义是什么？",
    "我想辞职，你怎么看？",
    "帮我写一段自我介绍。",
    "你喜欢什么类型的音乐？",
    "现在几点了，你在做什么？",
    "如果你有一天的自由时间，你会怎么过？",
    "我刚刚和朋友吵架了，我该怎么办？",
]


def load_vector(trait: str, layer_idx: int) -> torch.Tensor | None:
    path = ROOT / "vectors" / f"{trait}_layer{layer_idx:03d}.pt"
    if not path.exists():
        return None
    return torch.load(path, weights_only=True)


def generate_single(model, tokenizer, prompt: str, max_new_tokens: int = 128) -> str:
    inputs = tokenizer(prompt, return_tensors="pt").to(next(model.parameters()).device)
    with torch.no_grad():
        ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            temperature=None,
            top_p=None,
            pad_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.3,
        )
    new_ids = ids[0][inputs["input_ids"].shape[1]:]
    return tokenizer.decode(new_ids, skip_special_tokens=True)


def token_length(tokenizer, text: str) -> int:
    return len(tokenizer.encode(text, add_special_tokens=False))


def is_degenerate(response: str) -> bool:
    """Detect collapsed/looping outputs that should not count as steered successes."""
    stripped = response.strip()
    if len(stripped) < 4:
        return True
    if "</think>" in response:
        return True
    # High newline ratio → mostly whitespace
    if len(stripped) > 0 and response.count("\n") / max(len(response), 1) > 0.3:
        return True
    # Repetition: any phrase of 4+ chars repeated 5+ times
    words = stripped.split()
    if len(words) >= 10:
        for n in (3, 4, 5):
            for i in range(len(words) - n * 5):
                phrase = " ".join(words[i : i + n])
                rest = " ".join(words[i + n :])
                if rest.count(phrase) >= 4:
                    return True
    return False


def score_response(response: str, tokenizer) -> dict:
    length = token_length(tokenizer, response)
    return {
        "token_length": length,
        "coherent": int(not is_degenerate(response)),
    }


def evaluate_layer(
    model,
    tokenizer,
    trait: str,
    layer_idx: int,
    alpha: float,
    prompts: list[str],
) -> list[dict]:
    vector = load_vector(trait, layer_idx)
    if vector is None:
        return []

    results = []
    for user_text in prompts:
        prompt = build_chat_prompt(tokenizer, user_text)

        # Baseline (no steering)
        baseline = generate_single(model, tokenizer, prompt)

        # Steered
        hook = SteeringHook(model, layer_idx, vector, alpha)
        steered = generate_single(model, tokenizer, prompt)
        hook.remove()

        degen = is_degenerate(steered)
        results.append({
            "layer": layer_idx,
            "prompt": user_text,
            "baseline": baseline,
            "steered": steered,
            "baseline_tokens": token_length(tokenizer, baseline),
            "steered_tokens": token_length(tokenizer, steered),
            "coherent": int(not degen),
            "degenerate": int(degen),
        })

    return results


def sweep_and_save(cfg: dict, trait: str, layers: list[int], alpha: float):
    outputs_dir = ROOT / "outputs"
    outputs_dir.mkdir(exist_ok=True)

    print("Loading model...")
    model, tokenizer = load_model(cfg)

    all_results = []
    layer_scores: dict[int, float] = {}

    for layer_idx in tqdm(layers, desc="Sweeping layers"):
        if not (ROOT / "vectors" / f"{trait}_layer{layer_idx:03d}.pt").exists():
            continue

        results = evaluate_layer(model, tokenizer, trait, layer_idx, alpha, EVAL_PROMPTS)
        if not results:
            continue

        all_results.extend(results)

        # Score: coherent responses that are noticeably shorter than baseline
        coherent = [r for r in results if r["coherent"]]
        if coherent:
            avg_reduction = sum(
                (r["baseline_tokens"] - r["steered_tokens"]) / max(r["baseline_tokens"], 1)
                for r in coherent
            ) / len(coherent)
            layer_scores[layer_idx] = avg_reduction
        else:
            layer_scores[layer_idx] = -1.0  # penalize degenerate outputs

    # Save all outputs
    out_path = outputs_dir / f"{trait}_layer_sweep.jsonl"
    with open(out_path, "w", encoding="utf-8") as f:
        for r in all_results:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"Results saved to {out_path}")

    if layer_scores:
        best_layer = max(layer_scores, key=lambda k: layer_scores[k])
        best_path = outputs_dir / f"{trait}_best_layer.txt"
        best_path.write_text(str(best_layer))
        print(f"\nBest layer for '{trait}': {best_layer}  (score={layer_scores[best_layer]:.3f})")
        print(f"Update config.yaml → steering.apply_layer: {best_layer}")
    else:
        print("No valid vectors found. Run extract_vectors.py first.")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--trait", required=True)
    parser.add_argument("--layers", default=None, help="Comma-separated or 'all'")
    parser.add_argument("--alpha", type=float, default=None)
    args = parser.parse_args()

    cfg = load_config()
    alpha = args.alpha if args.alpha is not None else cfg["steering"]["alpha"]

    layer_cfg = args.layers or cfg["steering"]["extract_layers"]
    if layer_cfg == "all":
        vectors_dir = ROOT / "vectors"
        prefix = f"{args.trait}_layer"
        layers = sorted(
            int(p.stem.replace(prefix, ""))
            for p in vectors_dir.glob(f"{prefix}*.pt")
        )
    else:
        layers = [int(x) for x in str(layer_cfg).split(",")]

    sweep_and_save(cfg, args.trait, layers, alpha)


if __name__ == "__main__":
    main()