"""Rebuild a full TSLM ``best_model.pt`` state dict from this repository.

Only the tensors training changed are stored here. Everything else is recovered
bitwise from the public base checkpoints (Qwen3.5-4B, Chronos-2), so the result
is byte-identical to the original training artifact.

    python reconstruct.py --run anomaly_xl_itformer_precise_chronos2_qwen35_4b \
                          --out best_model.pt

In-process, without writing a ``.pt``::

    from reconstruct import rebuild_state_dict
    state = rebuild_state_dict("anomaly_xl_itformer_precise_chronos2_qwen35_4b")
    model.load_state_dict(state)
"""
import argparse
import glob
import json
from pathlib import Path

import torch
from huggingface_hub import snapshot_download
from safetensors import safe_open

REPO_ID = "nz00shuuuu/TimeRLM-TSLM-Checkpoints"
DTYPES = {"bfloat16": torch.bfloat16, "float32": torch.float32, "float16": torch.float16}

_STOCK_CACHE = {}


def load_safetensors(paths):
    out = {}
    for p in paths:
        with safe_open(str(p), framework="pt") as f:
            for k in f.keys():
                out[k] = f.get_tensor(k)
    return out


def load_stock():
    """Download and load the public base checkpoints, memoised per process."""
    if not _STOCK_CACHE:
        qdir = snapshot_download("Qwen/Qwen3.5-4B", allow_patterns=["*.safetensors"])
        cdir = snapshot_download("amazon/chronos-2", allow_patterns=["*.safetensors"])
        _STOCK_CACHE["stock_qwen"] = load_safetensors(sorted(glob.glob(f"{qdir}/*.safetensors")))
        _STOCK_CACHE["stock_chronos"] = load_safetensors(sorted(glob.glob(f"{cdir}/*.safetensors")))
    return _STOCK_CACHE


def run_dir(run=None, local_dir=None):
    if local_dir:
        return Path(local_dir)
    if not run:
        raise ValueError("pass --run or --local-dir")
    root = snapshot_download(REPO_ID, allow_patterns=[f"tslm/{run}/*"])
    return Path(root) / "tslm" / run


def rebuild_state_dict(run=None, local_dir=None, stock=None):
    """Return the full ``model_state``, bitwise-equal to the original checkpoint."""
    root = run_dir(run, local_dir)
    meta = json.loads((root / "manifest.json").read_text())
    trained = load_safetensors([root / "trained.safetensors"])
    stock = stock or load_stock()

    state = {}
    for key in meta["key_order"]:
        e = meta["tensors"][key]
        src = e["from"]
        if src == "extract":
            state[key] = trained[key]
        elif src == "alias":
            state[key] = state[e["key"]]
        else:
            t = stock[src][e["key"]].to(DTYPES[e["dtype"]])
            if "slice_rows" in e:
                t = t[: e["slice_rows"]].contiguous()
            state[key] = t
    return state


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--run", help="run name, e.g. anomaly_xl_itformer_precise_chronos2_qwen35_4b")
    ap.add_argument("--local-dir", help="use an already-downloaded run directory")
    ap.add_argument("--out", help="write a train.py-compatible .pt here")
    args = ap.parse_args()

    state = rebuild_state_dict(args.run, args.local_dir)
    print(f"rebuilt {len(state)} tensors")
    if args.out:
        meta = json.loads((run_dir(args.run, args.local_dir) / "manifest.json").read_text())
        torch.save({"model_state": state, "config": meta["config"], **meta["training"]}, args.out)
        print(f"wrote {args.out}")


if __name__ == "__main__":
    main()