from __future__ import annotations

import json

from model import TuringSurrogate
from safetensors.torch import load_file
from train import ARTIFACT_DIR, rollout, save_comparison_gif


def main() -> None:
    model = TuringSurrogate()
    model.load_state_dict(load_file(ARTIFACT_DIR / "model.safetensors"))
    metrics, physics_frames, neural_frames = rollout(model)
    save_comparison_gif(
        physics_frames,
        neural_frames,
        ARTIFACT_DIR / "physics_vs_neural.gif",
    )
    print(json.dumps(metrics, indent=2))


if __name__ == "__main__":
    main()