"""
Greedy evaluation of the reasoning model on the held-out split.

Reports:
  * format compliance   — a single well-formed <think>…</think> block followed by an answer
  * numeric agreement   — do the numbers in the generated conclusion match the reference's
  * length stats        — how long the produced reasoning is

When the eval rows carry a `source` tag (the v3 mix), every metric is also broken down per
source, since the two corpora answer in different styles.

Usage: python eval_reasoning.py [model_dir] [n_samples]
"""
import json
import os
import re
import sys
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_DIR = sys.argv[1] if len(sys.argv) > 1 else "./Nawah-Reasoning-v1"
LIMIT     = int(sys.argv[2]) if len(sys.argv) > 2 else 400
EVAL_FILE = os.environ.get("EVAL_FILE", "data/eval.jsonl")
MAX_NEW   = 512
BATCH     = 16

AR_DIGITS = str.maketrans("٠١٢٣٤٥٦٧٨٩٫٬", "0123456789.,")
NUM_RE    = re.compile(r"\d+(?:\.\d+)?")


def numbers(text: str):
    text = text.translate(AR_DIGITS).replace(",", "")
    out = []
    for tok in NUM_RE.findall(text):
        val = float(tok)
        out.append(int(val) if val.is_integer() else val)
    return out


def parse(completion: str):
    """-> (reasoning, final_answer, well_formed)"""
    m = re.match(r"\s*<think>(.*?)</think>(.*)", completion, re.S)
    if not m:
        return None, completion.strip(), False
    reasoning, final = m.group(1).strip(), m.group(2).strip()
    well_formed = (
        completion.count("<think>") == 1
        and completion.count("</think>") == 1
        and bool(reasoning)
        and bool(final)
    )
    return reasoning, final, well_formed


def metrics(results):
    n = len(results)
    return {
        "n": n,
        "well_formed_pct": 100 * sum(r["well_formed"] for r in results) / n,
        "numbers_match_pct": 100 * sum(r["numbers_match"] for r in results) / n,
        "primary_number_match_pct": 100 * sum(r["primary_number_match"] for r in results) / n,
        "answer_exact_pct": 100 * sum(r["answer_exact"] for r in results) / n,
        "mean_reasoning_tokens": sum(r["reasoning_tokens"] for r in results) / n,
    }


def main():
    tok = AutoTokenizer.from_pretrained(MODEL_DIR)
    model = AutoModelForCausalLM.from_pretrained(MODEL_DIR, dtype=torch.bfloat16).cuda().eval()
    model.config.use_cache = True

    rows = [json.loads(l) for l in open(EVAL_FILE, encoding="utf-8")][:LIMIT]
    im_end = tok.convert_tokens_to_ids("<|im_end|>")

    results = []
    for start in range(0, len(rows), BATCH):
        chunk = rows[start : start + BATCH]
        prompts = [
            f"<|im_start|>user\n{r['instruction']}<|im_end|>\n<|im_start|>assistant\n" for r in chunk
        ]
        encoded = [[tok.bos_token_id] + tok.encode(p, add_special_tokens=False) for p in prompts]
        width = max(len(e) for e in encoded)
        # left-pad so every row's generation starts at the same offset
        input_ids = torch.tensor([[tok.pad_token_id] * (width - len(e)) + e for e in encoded]).cuda()
        attn = torch.tensor([[0] * (width - len(e)) + [1] * len(e) for e in encoded]).cuda()

        with torch.no_grad():
            out = model.generate(
                input_ids=input_ids,
                attention_mask=attn,
                max_new_tokens=MAX_NEW,
                do_sample=False,
                eos_token_id=[im_end, tok.eos_token_id],
                pad_token_id=tok.pad_token_id,
            )
        for row, seq in zip(chunk, out):
            gen = tok.decode(seq[width:], skip_special_tokens=False)
            gen = gen.split("<|im_end|>")[0].replace("</s>", "").replace("<pad>", "")
            reasoning, final, ok = parse(gen)
            ref_nums, gen_nums = numbers(row["answer"]), numbers(final)
            results.append(
                {
                    "source": row.get("source"),
                    "instruction": row["instruction"],
                    "reference_reasoning": row["reasoning"],
                    "reference_answer": row["answer"],
                    "generated_reasoning": reasoning,
                    "generated_answer": final,
                    "well_formed": ok,
                    "answer_exact": final.strip() == row["answer"].strip(),
                    "numbers_match": bool(ref_nums) and ref_nums == gen_nums,
                    "primary_number_match": bool(ref_nums) and ref_nums[0] in gen_nums,
                    "reasoning_tokens": len(tok.encode(reasoning or "", add_special_tokens=False)),
                }
            )
        print(f"  {min(start + BATCH, len(rows))}/{len(rows)}", flush=True)

    summary = metrics(results)
    summary["model"] = MODEL_DIR
    # A mixed corpus (v3) answers in two different styles, so a single exact-match number is
    # meaningless — score each source on its own terms.
    sources = sorted({r["source"] for r in results if r["source"]})
    if len(sources) > 1:
        summary["by_source"] = {s: metrics([r for r in results if r["source"] == s]) for s in sources}
    print(json.dumps(summary, indent=2))

    Path(MODEL_DIR, "eval_reasoning.json").write_text(
        json.dumps({"summary": summary, "samples": results}, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(f"[+] wrote {Path(MODEL_DIR, 'eval_reasoning.json')}")


if __name__ == "__main__":
    main()