"""Export a training checkpoint to the release format: model.safetensors + config.json."""

# Usage: python -m facepoints.train.export <ckpt.pt or best.pt> <out dir>
# Tensors are stored flat with the prefixes net. / side.<set>. / side_rf.<set>. ; config.json holds
# the architecture, the point classes in output order, the likelihood and its std factor, the side
# heads and the training keys as run (without local file paths). model.load_release reads them.
import hashlib
import json
import os
import sys
from dataclasses import fields

import torch
from safetensors.torch import save_file

from facepoints.model import CONFIG, FORMAT, LOG_SCALE_RANGE, WEIGHTS, std_factor
from facepoints.train.config import TrainConfig


def export_release(checkpoint_path, out_dir):
    """Write model.safetensors and config.json to out_dir from a checkpoint; return the config."""
    ck = torch.load(checkpoint_path, map_location="cpu", weights_only=False)  # a trainer checkpoint
    a, classes = ck["args"], [str(c) for c in ck["cls"]]
    tensors = {f"net.{k}": v.contiguous() for k, v in ck["model"].items()}
    tensors.update({f"side.{k}": v.contiguous() for k, v in ck.get("side", {}).items()})
    tensors.update({f"side_rf.{k}": v.contiguous() for k, v in ck.get("side_rf", {}).items()})
    keep = {f.name for f in fields(TrainConfig)} - set(TrainConfig.PATH_KEYS)
    config = {
        "format": FORMAT,
        "backbone": a["backbone"],
        "input_px": int(a["res"]),
        "neck": int(a["neck"]),
        "stride": int(a.get("stride", 8)),
        "n_points": len(classes),
        "classes": classes,
        "likelihood": a["nll"],
        "nu": float(a.get("nu", 10.0)),
        "std_factor": std_factor(a["nll"], float(a.get("nu", 10.0))),
        "log_scale_range": list(LOG_SCALE_RANGE),
        "side_heads": {k: int(v) for k, v in ck.get("side_k", {}).items()},
        "side_refine": "side_rf" in ck,
        "iterations": int(ck["it"]),
        "training": {k: v for k, v in a.items() if k in keep},
    }
    os.makedirs(out_dir, exist_ok=True)
    save_file(tensors, os.path.join(out_dir, WEIGHTS), metadata={"format": "pt"})
    with open(os.path.join(out_dir, CONFIG), "w") as fh:
        json.dump(config, fh, indent=1)
    return config


def sha256(path):
    """Return the SHA-256 hex digest of a file."""
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def main():
    """Export a checkpoint and print the weights file's size and SHA-256."""
    if len(sys.argv) != 3:
        sys.exit("usage: python -m facepoints.train.export <ckpt.pt> <out dir>")
    config = export_release(sys.argv[1], sys.argv[2])
    w = os.path.join(sys.argv[2], WEIGHTS)
    mb = os.path.getsize(w) / 1e6
    print(f"wrote {w} ({mb:.1f} MB, sha256 {sha256(w)}) and {CONFIG}: {config['n_points']} points, "
          f"side heads {config['side_heads']}", flush=True)  # fmt: skip


if __name__ == "__main__":
    main()