"""Liquid d1 decision models (d1-3B, d1-omni-600M) on the `decisions` config: zero-shot answers, fine-tuning, answers.

    python liquid_d1.py zeroshot --model LiquidAI/d1-3B --split test --out runs/d1-3b/zeroshot.test.jsonl
    python liquid_d1.py train --model LiquidAI/d1-3B --ckpt runs/d1-3b/ckpt [--epochs 1]
    python liquid_d1.py predict --ckpt runs/d1-3b/ckpt --split test --out runs/d1-3b/finetuned.test.jsonl

d1-3B is a causal LFM2-VL that reads each answer off the logits at the answer slot. It trains with LoRA on its language
model, the state read once and the 13 questions as branches of its `hybrid.py` tree. d1-omni-600M is an encoder with a
decision head over option markers; it trains in full. Vision and audio towers stay frozen. Answers always come from
the model's own `system_one`, so a fine-tuned checkpoint is used exactly as the released one.
"""
from __future__ import annotations

import argparse
import math
import random
import shutil
import sys
import time
from contextlib import nullcontext
from pathlib import Path

import torch
import torch.nn.functional as F

from common import PredictionWriter, gold_index, load_split, log, to_answer

LORA_TARGETS = r".*language_model.*\.(q_proj|k_proj|v_proj|out_proj|in_proj|w1|w2|w3)"
CODE_FILES = ("*.py", "*.json", "*.jinja", "*.txt", "LICENSE")
WEIGHTS = ("model.safetensors", "model.safetensors.index.json")


def device_of(name: str | None) -> torch.device:
    return torch.device(name or ("cuda" if torch.cuda.is_available() else "cpu"))


def is_omni(path: str | Path) -> bool:
    import json

    from huggingface_hub import hf_hub_download
    config = Path(path) / "config.json" if Path(path).is_dir() else Path(hf_hub_download(str(path), "config.json"))
    return json.loads(config.read_text())["model_type"] == "d1_omni"


def load(path: str | Path, device: torch.device, dtype: torch.dtype | None = None):
    from transformers import AutoModel
    if dtype is None:  # omni was trained in fp32 and its card runs fp16 on GPUs; d1-3B ships in bf16
        dtype = torch.float32 if device.type == "cpu" else torch.float16 if is_omni(path) else torch.bfloat16
    return AutoModel.from_pretrained(str(path), trust_remote_code=True, dtype=dtype).to(device).eval()


def remote(model, name: str):
    """A module of the model's own remote code (`prompt`, `hybrid`)."""
    return sys.modules[type(model).__module__.rsplit(".", 1)[0] + "." + name]


def sync(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize()
    elif device.type == "mps":
        torch.mps.synchronize()


# --------------------------------------------------------------------------- answers


@torch.inference_mode()
def answer_case(model, case: dict) -> dict:
    """The 13 answers for one case from the model's `system_one`, in `common.options_of` order."""
    response = model.system_one(case["state"], case["questions"])["answers"]
    out = {}
    for name, q in case["questions"].items():
        a = response[name]
        if q["type"] == "noul":
            probs = [a["noul"], 1 - a["noul"]]
        elif q["type"] == "choice":
            probs = [a["probabilities"][k] for k in q["criteria"]]
        else:
            probs = [a["probabilities"][str(i)] for i in range(len(q["criteria"]))]
        out[name] = to_answer(q, probs)
    return out


def answer_split(path: str | Path, split: str, out: Path, limit: int | None, device: torch.device) -> None:
    cases = load_split(split, limit)
    writer = PredictionWriter(out)
    todo = writer.pending(cases)
    if not todo:
        return log(None, done=str(out), cases=len(cases))
    model = load(path, device)
    for _ in range(2):  # warm-up: kernels, caches, the engine's tokenizer
        answer_case(model, todo[0])
    for i, case in enumerate(todo):
        sync(device)
        tick = time.perf_counter()
        answers = answer_case(model, case)
        sync(device)
        writer.write(case["id"], answers, 1000 * (time.perf_counter() - tick))
        if i % 200 == 0:
            log(None, split=split, answered=len(writer.done), of=len(cases))
    writer.close()
    log(None, done=str(out), cases=len(cases))


# --------------------------------------------------------------------------- training: d1-3B


class Tree3B:
    """d1-3B training views: the state as the trunk, the 13 questions as branches, each read at its answer slot."""

    def __init__(self, model, questions: dict):
        engine, prompt = model.engine, remote(model, "prompt")
        self.model, self.tok, self.engine, self.prompt = model, engine.tokenizer, engine, prompt
        self.names, self.questions = list(questions), questions
        qs = [prompt.as_question(questions[n]) for n in self.names]
        self.branches = [self.encode(prompt.suffix_text(self.tok, q, engine.lead, engine.option_style)) for q in qs]
        self.groups = [prompt.readout_ids(self.tok, q) for q in qs]
        # the fused flash kernels return a log-sum-exp without a gradient; the explicit fp32 path keeps it exact
        remote(model, "hybrid")._fused = lambda q: False

    def encode(self, text: str) -> list[int]:
        return self.tok.encode(text, add_special_tokens=False)

    def option_logits(self, case: dict, device: torch.device) -> list[torch.Tensor]:
        e = self.engine
        trunk = self.encode(self.prompt.prefix_text(self.tok, case["state"], e.bos, e.state_style, e.system))
        packed = torch.tensor([t for b in self.branches for t in b], device=device)
        lengths = torch.tensor([len(b) for b in self.branches], device=device)
        logits = self.model.answer(torch.tensor([trunk], device=device), packed, lengths).float()
        return [torch.stack([logits[j, g].max() for g in groups]) for j, groups in enumerate(self.groups)]

    def loss(self, cases: list[dict], device: torch.device) -> torch.Tensor:
        total = []
        for case in cases:
            for name, scores in zip(self.names, self.option_logits(case, device)):
                target = torch.tensor([gold_index(self.questions[name], case["gold"][name])], device=device)
                total.append(F.cross_entropy(scores[None], target))
        return torch.stack(total).mean()


# --------------------------------------------------------------------------- training: d1-omni-600M


class Omni:
    """d1-omni training views: one sequence per question, scored at its option markers (a noul reads [false, true])."""

    def __init__(self, model, questions: dict):
        self.model, self.prompt = model, remote(model, "prompt")
        self.names, self.questions = list(questions), questions
        self.qs = [self.prompt.as_question(questions[n]) for n in self.names]

    def rows(self, case: dict) -> list[tuple[list[int], list[int], int, int]]:
        out = []
        for name, q in zip(self.names, self.qs):
            ids, markers = self.prompt.encode(self.model.tokenizer, case["state"], q, self.model.config.max_length)
            target = gold_index(self.questions[name], case["gold"][name])
            out.append((ids, markers, self.prompt.QTYPES[q.type], 1 - target if q.type == "noul" else target))
        return out

    def loss(self, cases: list[dict], device: torch.device) -> torch.Tensor:
        m = self.model
        rows = [row for case in cases for row in self.rows(case)]
        seqs = [m.encoder.embed_tokens(torch.tensor(ids, device=device)) for ids, *_ in rows]
        h = torch.nn.utils.rnn.pad_sequence(seqs, batch_first=True)
        lengths = torch.tensor([len(ids) for ids, *_ in rows], device=device)
        pad = torch.arange(h.shape[1], device=device)[None] < lengths[:, None]
        h = m.encoder(h, pad, torch.zeros(len(rows), dtype=torch.long, device=device))
        k = max(len(markers) for _, markers, *_ in rows)
        marker_pos = torch.zeros(len(rows), k, dtype=torch.long, device=device)
        marker_mask = torch.zeros(len(rows), k, dtype=torch.bool, device=device)
        for i, (_, markers, *_) in enumerate(rows):
            marker_pos[i, :len(markers)] = torch.tensor(markers, device=device)
            marker_mask[i, :len(markers)] = True
        qtype = torch.tensor([r[2] for r in rows], device=device)
        logits = m.head(h, pad, marker_pos, marker_mask, qtype)
        return F.cross_entropy(logits, torch.tensor([r[3] for r in rows], device=device))


# --------------------------------------------------------------------------- training loop


def trainable(model, omni: bool):
    """Omni: everything but the vision and audio towers. d1-3B: LoRA on the language model's linear layers."""
    if omni:
        for name, p in model.named_parameters():
            p.requires_grad = not name.startswith(("vision.", "audio."))
        return model
    from peft import LoraConfig, get_peft_model
    for p in model.parameters():
        p.requires_grad = False
    return get_peft_model(model, LoraConfig(r=32, lora_alpha=64, lora_dropout=0.05, target_modules=LORA_TARGETS))


def trainable_state(model) -> dict:
    return {n: p.detach().cpu() for n, p in model.named_parameters() if p.requires_grad}


def save_atomic(obj, path: Path) -> None:
    tmp = path.with_suffix(".tmp")
    torch.save(obj, tmp)
    tmp.replace(path)


def export(model, omni: bool, source: str, best: Path, final: Path) -> None:
    """The best weights as a standalone checkpoint that `AutoModel.from_pretrained(final, trust_remote_code=True)`
    loads, with the source repo's code and tokenizer files."""
    from huggingface_hub import snapshot_download
    model.load_state_dict(torch.load(best, map_location="cpu"), strict=False)
    if omni:
        model.config.temperatures = {}  # the released per-type temperatures were fitted to the released weights
    else:
        model = model.merge_and_unload()
    tmp = final.with_name(final.name + ".tmp")
    shutil.rmtree(tmp, ignore_errors=True)
    src = Path(snapshot_download(source, allow_patterns=list(CODE_FILES)))
    shutil.copytree(src, tmp, ignore=shutil.ignore_patterns("README.md", *WEIGHTS))
    model.save_pretrained(tmp)
    shutil.rmtree(final, ignore_errors=True)
    tmp.replace(final)


@torch.no_grad()
def val_loss(view, cases: list[dict], device: torch.device, batch: int, autocast) -> float:
    losses = []
    for i in range(0, len(cases), batch):
        with autocast():
            losses.append(float(view.loss(cases[i:i + batch], device)) * len(cases[i:i + batch]))
    return sum(losses) / len(cases)


def train(args) -> None:
    device, ckpt = device_of(args.device), Path(args.ckpt)
    ckpt.mkdir(parents=True, exist_ok=True)
    if (ckpt / "final").exists():
        return log(args.log, done=str(ckpt / "final"))
    omni = is_omni(args.model)
    dtype = torch.float32 if omni or device.type == "cpu" else torch.bfloat16  # omni keeps fp32 master weights
    model = load(args.model, device, dtype)
    train_cases, val_cases = load_split("train", args.limit), load_split("validation", args.val_cases)
    questions = train_cases[0]["questions"]
    view = Omni(model, questions) if omni else Tree3B(model, questions)
    model = trainable(model, omni)
    model.train()
    lr = args.lr or (2e-5 if omni else 2e-4)
    micro = args.micro or (2 if omni else 1)
    accum = max(1, args.batch // micro)
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(params, lr=lr, weight_decay=0.0)
    per_epoch = math.ceil(len(train_cases) / (micro * accum))
    total = min(per_epoch * args.epochs, args.max_steps or 10**9)
    warmup = max(1, int(0.05 * total))
    schedule = torch.optim.lr_scheduler.LambdaLR(
        optimizer, lambda s: min(1.0, (s + 1) / warmup) * max(0.0, (total - s) / max(1, total - warmup)))
    state = {"epoch": 0, "step": 0, "best": float("inf")}
    if (ckpt / "last.pt").exists():
        saved = torch.load(ckpt / "last.pt", map_location="cpu")
        model.load_state_dict(saved["weights"], strict=False)
        optimizer.load_state_dict(saved["optimizer"])
        schedule.load_state_dict(saved["schedule"])
        state = saved["state"]
        log(args.log, resumed=state)
    autocast = (lambda: torch.autocast("cuda", dtype=torch.bfloat16)) if device.type == "cuda" else nullcontext

    def checkpoint() -> None:
        save_atomic({"weights": trainable_state(model), "optimizer": optimizer.state_dict(),
                     "schedule": schedule.state_dict(), "state": state}, ckpt / "last.pt")

    started = time.time()
    while state["epoch"] < args.epochs and state["step"] < total:
        order = list(range(len(train_cases)))
        random.Random(args.seed + state["epoch"]).shuffle(order)
        done_in_epoch = state["step"] - state["epoch"] * per_epoch
        for s in range(done_in_epoch, per_epoch):
            if state["step"] >= total:
                break
            chunk = [train_cases[i] for i in order[s * micro * accum:(s + 1) * micro * accum]]
            for i in range(0, len(chunk), micro):
                with autocast():
                    loss = view.loss(chunk[i:i + micro], device) * len(chunk[i:i + micro]) / len(chunk)
                loss.backward()
            norm = torch.nn.utils.clip_grad_norm_(params, 1.0)
            optimizer.step()
            schedule.step()
            optimizer.zero_grad(set_to_none=True)
            state["step"] += 1
            if state["step"] % 20 == 0 or state["step"] <= 3:
                log(args.log, step=state["step"], of=total, loss=round(float(loss) * len(chunk) / micro, 4),
                    grad_norm=round(float(norm), 3), lr=schedule.get_last_lr()[0], minutes=round((time.time() - started) / 60, 1))
            if state["step"] % args.save_every == 0:
                checkpoint()
        state["epoch"] += 1
        model.eval()
        loss = val_loss(view, val_cases, device, micro, autocast)
        model.train()
        log(args.log, epoch=state["epoch"], val_loss=round(loss, 4), minutes=round((time.time() - started) / 60, 1))
        if loss < state["best"]:
            state["best"] = loss
            save_atomic(trainable_state(model), ckpt / "best.pt")
        checkpoint()
    model.eval()
    export(model, omni, args.model, ckpt / "best.pt", ckpt / "final")
    log(args.log, done=str(ckpt / "final"), best_val_loss=state["best"], minutes=round((time.time() - started) / 60, 1))


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    sub = parser.add_subparsers(dest="command", required=True)
    for name in ("zeroshot", "train", "predict"):
        p = sub.add_parser(name)
        p.add_argument("--device")
        p.add_argument("--limit", type=int)
        if name != "predict":
            p.add_argument("--model", default="LiquidAI/d1-3B")
        if name != "train":
            p.add_argument("--split", default="test", choices=("validation", "test"))
            p.add_argument("--out", type=Path, required=True)
        if name != "zeroshot":
            p.add_argument("--ckpt", required=True)
    t = sub.choices["train"]
    t.add_argument("--epochs", type=int, default=1)
    t.add_argument("--max-steps", type=int)
    t.add_argument("--batch", type=int, default=8, help="cases per optimizer step")
    t.add_argument("--micro", type=int, help="cases per forward pass (default 2 for omni, 1 for d1-3B)")
    t.add_argument("--lr", type=float)
    t.add_argument("--val-cases", type=int, default=300)
    t.add_argument("--save-every", type=int, default=400)
    t.add_argument("--seed", type=int, default=0)
    t.add_argument("--log")
    args = parser.parse_args()
    if args.command == "train":
        return train(args)
    path = args.model if args.command == "zeroshot" else Path(args.ckpt) / "final"
    answer_split(path, args.split, args.out, args.limit, device_of(args.device))


if __name__ == "__main__":
    main()