#!/usr/bin/env python3
"""Create the full module graph on meta, or explicitly allocate random weights."""

from __future__ import annotations

import argparse
from pathlib import Path

import torch

from configuration_mini_k3 import MiniK3Config
from modeling_mini_k3 import MiniK3ForCausalLM, count_logical_parameters


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, default=Path(__file__).with_name("config.json"))
    parser.add_argument(
        "--materialize",
        action="store_true",
        help="Actually allocate and randomly initialize the full model on CPU.",
    )
    parser.add_argument("--seed", type=int, default=1234, help="Random seed for initialization")
    parser.add_argument("--output", type=Path, help="Optional torch.save state-dict path")
    args = parser.parse_args()

    config = MiniK3Config.from_json(args.config)
    model = MiniK3ForCausalLM(config, device="meta")
    print(f"meta model defined: {count_logical_parameters(model):,} logical parameters")
    if not args.materialize:
        print("No storage allocated. Pass --materialize only on a host with enough CPU RAM.")
        return

    torch.manual_seed(args.seed)
    model.materialize_and_initialize("cpu")
    print(f"Full model randomly initialized on CPU with seed {args.seed}")
    if args.output is not None:
        torch.save({"config": config.to_dict(), "model": model.state_dict()}, args.output)
        print(f"saved: {args.output}")


if __name__ == "__main__":
    main()