import argparse
import json
import zipfile
from pathlib import Path

import torch
from tqdm import tqdm
from transformers import AutoTokenizer, Gemma4ForCausalLM

MAX_PARAMS = 8_000_000_000
MAX_BATCH_TOKENS = 36864
LABELS = ["1", "2", "3", "4"]
ANSWER_CUE = "පිළිතුර:"
THOUGHT_OPEN = "<|channel>thought\n"
THOUGHT_CLOSE = "\n<channel|>"
STOP_TOKENS = ["<channel|>", "<eos>", "<turn|>"]


def normalise(question):
    metadata = question.get("metadata", {})
    question = {
        **question,
        "subject_original": metadata.get("subject_original", question.get("subject_original")),
        "level": metadata.get("difficulty", question.get("level")),
    }
    answer = question.get("answer")
    if answer is not None and not (isinstance(answer, int) and 1 <= answer <= len(question["choices"])):
        texts = [choice.strip() for choice in question["choices"]]
        question["answer"] = texts.index(str(answer).strip()) + 1
    return question


def load_questions(data_dir):
    questions = []
    for path in sorted(Path(data_dir).rglob("*.json")):
        if any(part.startswith(".") for part in path.relative_to(data_dir).parts):
            continue
        questions.extend(normalise(q) for q in json.loads(path.read_text(encoding="utf-8")))
    return questions


def user_prompt(question):
    lines = [
        "පහත බහුවරණ ප්‍රශ්නයට නිවැරදි පිළිතුරේ අංකය පමණක් ලියන්න.",
        "",
        f"විෂයය: {question['subject_original']}",
        f"ප්‍රශ්නය: {question['question'].strip()}",
        *(f"{i}. {choice}" for i, choice in enumerate(question["choices"], 1)),
        ANSWER_CUE,
    ]
    return "\n".join(lines)


def render(tokenizer, prompt, think):
    if tokenizer.chat_template is None:
        return tokenizer.bos_token + prompt + " "
    messages = [{"role": "user", "content": prompt}]
    if think:
        text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True, enable_thinking=True
        )
        return text + THOUGHT_OPEN
    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)


def thinking_text(tokenizer, prompt, thought):
    return render(tokenizer, prompt, True) + thought + THOUGHT_CLOSE + ANSWER_CUE


DTYPES = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}


def placement(dtype_name="auto"):
    if not torch.cuda.is_available():
        return torch.float32, "cpu"
    if dtype_name != "auto":
        return DTYPES[dtype_name], "auto"
    return (torch.float16 if torch.cuda.get_device_capability()[0] < 8 else torch.bfloat16), "auto"


def load_model(name, adapter=None, revision=None, dtype_name="auto"):
    tokenizer = AutoTokenizer.from_pretrained(name, revision=revision, padding_side="left")
    dtype, device_map = placement(dtype_name)
    model, info = Gemma4ForCausalLM.from_pretrained(
        name,
        revision=revision,
        dtype=dtype,
        device_map=device_map,
        key_mapping={r"^model\.language_model\.": "model."},
        output_loading_info=True,
    )
    assert not info["missing_keys"], f"weights missing from checkpoint: {sorted(info['missing_keys'])[:5]}"
    n_adapter = 0
    if adapter is not None:
        from peft import PeftModel

        model = PeftModel.from_pretrained(model, adapter)
        n_adapter = sum(p.numel() for name, p in model.named_parameters() if "lora_" in name)
        assert n_adapter > 0, f"adapter {adapter} has no LoRA parameters"
        model = model.merge_and_unload()
    model.eval()
    n_params = sum(p.numel() for p in model.parameters()) + n_adapter
    print(
        f"Loaded {name} (text-only) on {model.device} in {model.dtype}: {n_params:,} parameters"
        f" ({n_adapter:,} adapter, merged)"
    )
    assert n_params <= MAX_PARAMS, f"{n_params:,} parameters exceeds the 8B limit"
    return tokenizer, model


def label_ids(tokenizer):
    ids = [tokenizer.encode(label, add_special_tokens=False) for label in LABELS]
    assert all(len(i) == 1 for i in ids), f"answer labels are not single tokens: {ids}"
    return [i[0] for i in ids]


def batches(lengths, budget, batch_size):
    batch = []
    for i in sorted(range(len(lengths)), key=lambda i: -lengths[i]):
        if batch and (
            len(batch) == batch_size or (len(batch) + 1) * (lengths[batch[0]] + budget) > MAX_BATCH_TOKENS
        ):
            yield batch
            batch = []
        batch.append(i)
    if batch:
        yield batch


def encode(tokenizer, texts, device):
    return tokenizer(texts, return_tensors="pt", padding=True, add_special_tokens=False).to(device)


def score(model, tokenizer, texts, n_choices, batch_size):
    labels = label_ids(tokenizer)
    lengths = [len(tokenizer.encode(t, add_special_tokens=False)) for t in texts]
    probs = [None] * len(texts)
    for batch in tqdm(list(batches(lengths, 0, batch_size)), desc="score"):
        inputs = encode(tokenizer, [texts[i] for i in batch], model.device)
        position_ids = (inputs["attention_mask"].cumsum(-1) - 1).clamp(min=0)
        with torch.inference_mode():
            logits = model(**inputs, position_ids=position_ids, logits_to_keep=1).logits[:, -1, labels].float()
        if not torch.isfinite(logits).all():
            raise RuntimeError("non-finite answer logits")
        counts = torch.tensor([n_choices[i] for i in batch], device=logits.device)
        valid = torch.arange(len(labels), device=logits.device)[None, :] < counts[:, None]
        for i, row in zip(batch, logits.masked_fill(~valid, float("-inf")).softmax(-1).tolist()):
            probs[i] = row[: n_choices[i]]
    return probs


def think(model, tokenizer, contexts, budget, batch_size):
    stops = tokenizer.convert_tokens_to_ids(STOP_TOKENS)
    lengths = [len(tokenizer.encode(t, add_special_tokens=False)) for t in contexts]
    thoughts = [None] * len(contexts)
    for batch in tqdm(list(batches(lengths, budget, batch_size)), desc="think"):
        inputs = encode(tokenizer, [contexts[i] for i in batch], model.device)
        with torch.inference_mode():
            out = model.generate(
                **inputs,
                max_new_tokens=budget,
                do_sample=False,
                eos_token_id=stops,
                pad_token_id=tokenizer.pad_token_id,
            )
        for i, row in zip(batch, out[:, inputs["input_ids"].shape[1]:].tolist()):
            end = next((k for k, token in enumerate(row) if token in stops), len(row))
            thoughts[i] = {
                "thought": tokenizer.decode(row[:end], skip_special_tokens=True).strip(),
                "finished": end < len(row) and row[end] == stops[0],
                "thought_tokens": end,
            }
    return thoughts


def argmax(values):
    return max(range(len(values)), key=values.__getitem__)


def answer(model, tokenizer, prompts, n_choices, budget, batch_size, score_batch_size=None):
    if budget == 0:
        texts = [render(tokenizer, p, False) for p in prompts]
        thoughts = [{} for _ in prompts]
    else:
        contexts = [render(tokenizer, p, True) for p in prompts]
        thoughts = think(model, tokenizer, contexts, budget, batch_size)
        texts = [thinking_text(tokenizer, p, t["thought"]) for p, t in zip(prompts, thoughts)]
    probs = score(model, tokenizer, texts, n_choices, score_batch_size or batch_size)
    return [{**t, "probs": p, "pred": argmax(p) + 1} for t, p in zip(thoughts, probs)]


def order(n, rotation, reverse):
    positions = list(range(n))[::-1] if reverse else list(range(n))
    k = rotation % n
    return positions[k:] + positions[:k]


def rotate(question, rotation, reverse=False):
    return {**question, "choices": [question["choices"][i] for i in order(len(question["choices"]), rotation, reverse)]}


def unrotate(probs, rotation, reverse=False):
    original = [0.0] * len(probs)
    for i, p in zip(order(len(probs), rotation, reverse), probs):
        original[i] = p
    return original


def vote(model, tokenizer, questions, rotations, budget, batch_size, score_batch_size=None, reverse=False):
    prompts = [user_prompt(rotate(q, r, reverse)) for q in questions for r in rotations]
    n_choices = [len(q["choices"]) for q in questions for _ in rotations]
    runs = answer(model, tokenizer, prompts, n_choices, budget, batch_size, score_batch_size)
    results = []
    for start in range(0, len(runs), len(rotations)):
        group = [
            {**run, "rotation": r, "reverse": reverse, "probs": unrotate(run["probs"], r, reverse)}
            for run, r in zip(runs[start : start + len(rotations)], rotations)
        ]
        group = [{**run, "pred": argmax(run["probs"]) + 1} for run in group]
        probs = [sum(p) / len(group) for p in zip(*(run["probs"] for run in group))]
        results.append({"pred": argmax(probs) + 1, "probs": probs, "runs": group})
    return results


def metrics(questions, results):
    levels = {}
    for question, result in zip(questions, results):
        counts = levels.setdefault(question["level"], [0, 0])
        counts[0] += result["pred"] == question["answer"]
        counts[1] += 1
    correct = sum(c for c, _ in levels.values())
    return {
        "accuracy": correct / len(questions),
        "correct": correct,
        "total": len(questions),
        "levels": {k: {"accuracy": c / t, "correct": c, "total": t} for k, (c, t) in sorted(levels.items())},
    }


def write_jsonl(path, rows):
    path.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows), encoding="utf-8")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="google/gemma-4-E4B-it")
    parser.add_argument("--revision", help="model commit to load; the latest one when not given")
    parser.add_argument("--dtype", choices=["auto", *DTYPES], default="auto")
    parser.add_argument("--data_dir", default="data/test")
    parser.add_argument("--out_dir", default="outputs")
    parser.add_argument("--think", type=int, default=512)
    parser.add_argument("--adapter")
    parser.add_argument("--limit", type=int)
    parser.add_argument("--level", help="answer only the questions of this level, for example hard")
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument(
        "--score_batch_size", type=int,
        help="batch size for scoring the answer after the thought; --batch_size when not given, lower it for long thoughts",
    )
    parser.add_argument(
        "--rotations", type=int, nargs="+", default=[0],
        help="answer once per rotation of the choices and average the answer probabilities; 0 2 votes over two orders",
    )
    parser.add_argument(
        "--reverse", action="store_true",
        help="reverse the choices before rotating them, which gives 4 more orders of 4 choices",
    )
    args = parser.parse_args()

    questions = load_questions(args.data_dir)
    if args.level:
        questions = [q for q in questions if q["level"] == args.level]
    questions = questions[: args.limit]
    print(f"Loaded {len(questions)} questions from {args.data_dir}")
    out_dir = Path(args.out_dir)
    if not questions:
        out_dir.mkdir(parents=True, exist_ok=True)
        write_jsonl(out_dir / "predictions.jsonl", [])
        write_jsonl(out_dir / "thoughts.jsonl", [])
        return

    tokenizer, model = load_model(args.model, args.adapter, args.revision, args.dtype)
    if args.think and tokenizer.chat_template is None:
        parser.error(f"{args.model} has no chat template, so it has no thinking mode; pass --think 0")

    results = vote(
        model, tokenizer, questions, args.rotations, args.think, args.batch_size, args.score_batch_size, args.reverse
    )
    predictions = [
        {
            "q_no": question["q_no"],
            "subject": question["subject"],
            "category": question["category"],
            "level": question["level"],
            "answer": LABELS[result["pred"] - 1],
        }
        for question, result in zip(questions, results)
    ]

    out_dir.mkdir(parents=True, exist_ok=True)
    jsonl_path = out_dir / "predictions.jsonl"
    write_jsonl(jsonl_path, predictions)
    with zipfile.ZipFile(out_dir / "predictions.zip", "w", zipfile.ZIP_DEFLATED) as zf:
        zf.write(jsonl_path, "predictions.jsonl")
    print(f"Wrote {len(predictions)} predictions to {out_dir / 'predictions.zip'}")

    runs = [(p, run) for p, r in zip(predictions, results) for run in r["runs"]]
    thoughts = [
        {"q_no": p["q_no"], "subject": p["subject"], "level": p["level"]}
        | {k: run[k] for k in ("rotation", "reverse", "pred", "probs", "finished", "thought_tokens", "thought") if k in run}
        for p, run in runs
    ]
    write_jsonl(out_dir / "thoughts.jsonl", thoughts)
    if args.think:
        finished = sum(run["finished"] for _, run in runs)
        print(f"Thoughts finished within {args.think} tokens: {finished}/{len(runs)}")

    if questions and all(q.get("answer") is not None for q in questions):
        report = metrics(questions, results) | {"think": args.think, "adapter": args.adapter}
        if len(args.rotations) > 1:
            report["rotations"] = {
                str(r): metrics(questions, [res["runs"][i] for res in results])["accuracy"]
                for i, r in enumerate(args.rotations)
            }
        if args.think:
            report["finished"] = finished
        (out_dir / "metrics.json").write_text(json.dumps(report, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
        levels = ", ".join(f"{k} {v['accuracy']:.4f}" for k, v in report["levels"].items())
        print(f"Accuracy {report['accuracy']:.4f} ({report['correct']}/{report['total']}); {levels}")
        if "rotations" in report:
            print("Accuracy of each rotation alone:", report["rotations"])


if __name__ == "__main__":
    main()