#!/usr/bin/env python3
"""Evaluate a saved adapter at given (stage,k) pairs using the trainer's own run_eval."""
import argparse, os, sys, types
sys.path.insert(0, "/home/ubuntu/curriculum-cot-code")
sys.path.insert(0, "/home/ubuntu/curriculum-cot-code/latent_multi_output_cell_policy")

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from latent_multi_output_cell_policy.sft_latent_multi_output_train import (
    run_eval, load_jsonl_rows, configure_hf_cache, pick_dtype,
)
from latent_multi_output_cell_policy.grpo_residual_projector_latent_train import load_trainable_adapter

p = argparse.ArgumentParser()
p.add_argument("--ckpt", required=True)
p.add_argument("--eval_jsonl", default="/home/ubuntu/curriculum-cot-code/data/sudoku_t3_20empty_value_qwen_text_stage1_eval.jsonl")
p.add_argument("--rows", type=int, default=50)
p.add_argument("--stages", default="1:1,2:2")  # stage:k pairs
p.add_argument("--model_name", default="Qwen/Qwen2.5-1.5B-Instruct")
p.add_argument("--cache_dir", default="/home/ubuntu/.hf_cache")
p.add_argument("--tag", default="")
a = p.parse_args()

cache = configure_hf_cache(a.cache_dir)
dev = torch.device("cuda:0")
tok = AutoTokenizer.from_pretrained(a.model_name, cache_dir=cache, use_fast=True)
if tok.pad_token_id is None:
    tok.pad_token = tok.eos_token or "<|endoftext|>"
tok._multi_value_oversample_factor = 1; tok._train_target_size_min = 0; tok._train_target_size_max = 0
base = AutoModelForCausalLM.from_pretrained(a.model_name, cache_dir=cache, torch_dtype=pick_dtype(), low_cpu_mem_usage=True)
model = load_trainable_adapter(base, a.ckpt, lora_r=32, lora_alpha=64, lora_dropout=0.05)
model._latent_debug_tokenizer = tok
if hasattr(model, "config"): model.config.use_cache = False
model.to(dev); model.eval()

rows = load_jsonl_rows(a.eval_jsonl, limit_rows=a.rows)
for pair in a.stages.split(","):
    s, k = (int(x) for x in pair.split(":"))
    args = types.SimpleNamespace(
        stage_i=s, num_cot_tokens=k, latent_mode="recurrent_hidden",
        total_empties_hint=20, eval_target_size_min=0, eval_target_size_max=0,
        max_completion_length=24, debug_print_limit=0,
        reward_good_value=1.25, penalty_bad_value=1.0, penalty_malformed=4.0,
        penalty_empty=0.5, penalty_singleton=1.5,
    )
    ev = run_eval(args, rows, model, tok, dev)
    print(f"RESULT tag={a.tag} stage={s} k={k} rows={a.rows} "
          f"exact={ev['exact_set_match_rate']:.3f} solve={ev['solve_rate']:.3f} "
          f"prec={ev['value_precision']:.3f} rec={ev['value_recall']:.3f}", flush=True)