"""Config, device, model construction and checkpoint helpers shared by the scripts."""
import argparse
import os

import torch
import yaml

from .thinker import Thinker
from .unet import UNet


def parse_args(extra=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/paper.yaml")
    ap.add_argument("--mode", choices=["hard", "extreme"], default=None, help="overrides config mode")
    ap.add_argument("--set", nargs="*", default=[], help="overrides like thinker.steps=100")
    for args, kw in (extra or []):
        ap.add_argument(*args, **kw)
    args = ap.parse_args()
    with open(args.config) as f:
        cfg = yaml.safe_load(f)
    for kv in args.set:
        k, v = kv.split("=", 1)
        d = cfg
        *path, last = k.split(".")
        for p in path:
            d = d[p]
        d[last] = yaml.safe_load(v)
    if args.mode:
        cfg["mode"] = args.mode
    os.makedirs(cfg["run_dir"], exist_ok=True)
    return args, cfg


def device():
    if torch.cuda.is_available():
        torch.backends.cuda.matmul.allow_tf32 = torch.backends.cudnn.allow_tf32 = True
        torch.backends.cudnn.benchmark = True
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def autocast(dev):
    """bf16 autocast on CUDA (as in TRM); a no-op elsewhere."""
    return torch.autocast(dev.type, torch.bfloat16, enabled=dev.type == "cuda")


def compile(fn, dev):
    return torch.compile(fn) if dev.type == "cuda" else fn


def build_painter(cfg):
    p = cfg["painter"]
    return UNet(1, 1, tuple(p["channels"]), p["layers_per_block"])


def build_dm(cfg):
    p = cfg["dm"]
    return UNet(3, 1, tuple(p["channels"]), p["layers_per_block"])


def build_thinker(cfg, painter):
    t = cfg["thinker"]
    return Thinker(painter, dim=t["dim"], heads=t["heads"], layers=t["layers"], expansion=t["expansion"],
                   grid=t["grid"], enc_hidden=tuple(t["enc_hidden"]), H=t["H"], L=t["L"])


def load_eval_models(cfg, kind, mode, dev):
    """EMA weights in eval mode: {'dm': ...} or {'painter': ..., 'thinker': ...} for kind 'dm' / 'path'."""
    def ema(model, name):
        model.load_state_dict(load(ckpt_path(cfg, name), dev)["ema"])
        return model.to(dev).eval()
    if kind == "dm":
        return {"dm": ema(build_dm(cfg), f"dm_{mode}")}
    painter = ema(build_painter(cfg), "painter")
    return {"painter": painter, "thinker": ema(build_thinker(cfg, painter), f"thinker_{mode}")}


def n_params(m):
    return sum(p.numel() for p in m.parameters())


def ckpt_path(cfg, name):
    return os.path.join(cfg["run_dir"], f"{name}.pt")


def save(path, **state):
    tmp = path + ".tmp"
    torch.save(state, tmp)
    os.replace(tmp, path)


def load(path, dev):
    return torch.load(path, map_location=dev, weights_only=False)