"""Imagine a plane's next 2 seconds under different commands, with this world model and with physics alone.

    pip install torch numpy pyyaml
    python example.py
"""
from pathlib import Path

import numpy as np
import torch
import yaml

np.set_printoptions(suppress=True, precision=2)

from fixed_wing.dynamics import Airframe
from fixed_wing.learned import LearnedFixedWing, rollout_with

HERE = Path(__file__).parent
cfg = yaml.safe_load((HERE / "fixed_wing/airframe.yaml").read_text())
wm = LearnedFixedWing.load(HERE / "fw_wm_small.pt", cfg)
print(f"world model: {sum(p.numel() for p in wm.corr.parameters())} learned parameters on a 6-DoF physics backbone")

# shared state layout (18): position (ENU, m), velocity (m/s), rotation matrix columns (body FLU -> world), body rates
TRIM = np.array([0.613, 0.0, -0.0714, 0.0])          # throttle, aileron, elevator, rudder: level flight at 14 m/s
alpha = 0.031                                         # trim angle of attack (rad): the nose sits slightly up
R = np.array([[np.cos(alpha), 0, -np.sin(alpha)], [0, 1, 0], [np.sin(alpha), 0, np.cos(alpha)]])
x = np.concatenate([[0, 0, 40.0], [14.0, 0, 0], R[:, 0], R[:, 1], R[:, 2], [0, 0, 0]])
history = np.repeat(x[None], wm.H + 1, 0)             # 0.5 s of straight, level flight
past = np.repeat(TRIM[None], wm.H, 0)

T = 40                                                # 2 s at 20 Hz
plans = {"hold trim": np.repeat(TRIM[None], T, 0),
         "bank right": np.repeat((TRIM + [0, 0.10, 0, 0])[None], T, 0),
         "pull up": np.repeat((TRIM + [0.2, 0, -0.08, 0])[None], T, 0)}
futures = wm.rollout(history, past, np.stack(list(plans.values())))                  # (3, 40, 18)
physics = rollout_with(Airframe.from_cfg(cfg), None, None, torch.tensor(history[-1:], dtype=torch.float32).expand(3, -1),
                       torch.tensor(np.stack(list(plans.values())), dtype=torch.float32))
for k, name in enumerate(plans):
    p, q = futures[k, -1, :3].numpy(), physics[k, -1, :3].numpy()
    print(f"{name:10s}  world model -> {np.round(p, 2)}   physics alone -> {np.round(q, 2)}")
print("On an airframe exactly like the nominal one in still air, the two agree closely; the world model earns its keep "
      "when the history shows a different airframe, wind, or lag.")