"""Build the video-depth PointDiT-B through the repository's builder and load its EMA2 weights (CPU, no Civo file).

    python load_point.py <this folder>

Run from a clone of github.com/NIRVANALAN/vgg_forecast_4d with PYTHONPATH=<clone>/src. Nothing outside this folder is
read: the released PointDiT-B 512 checkpoint that the training config names is NOT read (`model.point.weights: random`
builds the same architecture, and every one of the 216 tensors then comes from the EMA2 file; the model after the load is
bit-identical to one built over the released checkpoint). The frozen DINOv3 encoder is not needed to load.
"""
import sys
from pathlib import Path

from vggforecast.evaluation.checkpoint_loading import _load_weights, resolve_run_config
from vggforecast.training.model_loading import build_point_model

FOLDER = Path(sys.argv[1])
RUN, CKPT = FOLDER / "train", FOLDER / "train" / "ckpt_32500"

cfg = resolve_run_config(RUN, CKPT)["cfg"]   # resolved_resume_29500.yaml + launch_resume_29500.json
cfg.model.point.weights = "random"           # do not read model.point.checkpoint
cfg.model.point.attention_backend = "sdpa"   # the run used fa2 on a GPU; sdpa needs no extra package
model, core, provenance = build_point_model(cfg)
base_hw = tuple(int(v) * core.trunk.patch for v in core.trunk.position_source_grid)
loaded = _load_weights(model, core, CKPT, "ema2", base_hw)  # every tensor comes from ema2
model.eval()
print(model.architecture, core.target_centering, loaded["completed_updates"], loaded["ema_decay"],
      sum(p.numel() for p in model.parameters()))
# temporal_pointdit window_valid_mean_depth 32500 0.9996 140204928