"""Stateful dual-rate rollout predictor.

Keeps the object state as float32 internally (no text round-trip loss);
the Qwen semantic latent is refreshed from the text context only every
`tick_k` steps and reused (stale) in between — that is the dual rate.
"""

from __future__ import annotations

import math

import numpy as np
import torch

from .normalization import wrap_angle


def static_features_from_header(header: dict) -> tuple[np.ndarray, np.ndarray]:
    """Returns (static (N, 8) float32, ids (N,) int32) in header object order."""
    objs = header["objects"]
    static = np.zeros((len(objs), 8), dtype=np.float32)
    ids = np.zeros(len(objs), dtype=np.int32)
    for i, obj in enumerate(objs):
        ids[i] = obj["id"]
        is_circle = 1.0 if obj.get("type") == "circle" else 0.0
        mat = obj.get("material", {})
        if "radius" in obj:
            s0 = s1 = float(obj["radius"])
        else:
            s0 = float(obj.get("width", obj.get("size", 0.0)) or 0.0)
            s1 = float(obj.get("height", obj.get("size", 0.0)) or 0.0)
        static[i] = [is_circle, 1.0 - is_circle,
                     math.log(max(mat.get("mass", 1.0), 1e-3)),
                     mat.get("friction", 0.0), mat.get("elasticity", 0.0),
                     s0 / 100.0, s1 / 100.0, 0.0]
    return static, ids


class DualRatePredictor:
    def __init__(self, model, tokenizer, cfg, normalizer, device="cuda",
                 denoise_steps: int | None = None, sample_seed: int | None = 0):
        self.model = model
        self.tokenizer = tokenizer
        self.cfg = cfg
        self.norm = normalizer
        self.device = device
        self.denoise_steps = denoise_steps or cfg.denoise_steps
        self.sample_seed = sample_seed
        self.latent_refreshes = 0
        self._reset_fields()

    def _reset_fields(self):
        self.state = None
        self.static_t = None
        self.obj_mask = None
        self.gravity_dt = None
        self.ids = None
        self.frame_num = 0
        self.steps = 0
        self.latent = None
        self.latent_refreshes = 0

    def reset_from_arrays(self, static: np.ndarray, ids: np.ndarray, gravity,
                          dt: float, init_state: np.ndarray, start_frame: int):
        """Seed directly from cached arrays (static (N,8), ids (N,), gravity (2,))."""
        self._reset_fields()
        n = init_state.shape[0]
        self.state = init_state.astype(np.float32).copy()
        self.ids = ids[:n]
        self.static_t = torch.from_numpy(static[:n]).float()[None].to(self.device)
        self.obj_mask = torch.ones((1, n), dtype=torch.bool, device=self.device)
        self.gravity_dt = torch.tensor(
            [[float(gravity[0]) / 981.0, float(gravity[1]) / 981.0, float(dt) * 60.0]],
            dtype=torch.float32).to(self.device)
        self.frame_num = start_frame

    def reset(self, header: dict, init_state: np.ndarray, start_frame: int):
        """init_state: (N, 6) raw units; start_frame: frame number to emit first."""
        self._reset_fields()
        static, ids = static_features_from_header(header)
        n = init_state.shape[0]
        self.state = init_state.astype(np.float32).copy()
        self.ids = ids[:n]
        self.static_t = torch.from_numpy(static[:n]).float()[None].to(self.device)
        self.obj_mask = torch.ones((1, n), dtype=torch.bool, device=self.device)
        g = header.get("gravity", {"x": 0.0, "y": -981.0})
        dt = header.get("timestep", 1 / 60)
        self.gravity_dt = torch.tensor(
            [[g["x"] / 981.0, g["y"] / 981.0, dt * 60.0]], dtype=torch.float32,
        ).to(self.device)
        self.frame_num = start_frame

    def _refresh_latent(self, context_text: str):
        ids = self.tokenizer(context_text, add_special_tokens=False)["input_ids"]
        ids = ids[-self.cfg.max_ctx_tokens:]
        t = torch.tensor([ids], dtype=torch.long, device=self.device)
        mask = torch.ones_like(t)
        self.latent = self.model.compute_latent(t, mask)
        self.latent_refreshes += 1

    @torch.no_grad()
    def step(self, context_text: str) -> str:
        """Predict the next frame; refresh the semantic latent every tick_k steps."""
        if self.latent is None or self.steps % self.cfg.tick_k == 0:
            self._refresh_latent(context_text)

        cur_feat = self.state.copy()
        cur_feat[:, 4] = wrap_angle(cur_feat[:, 4])
        cur_t = torch.from_numpy(self.norm.norm_state(cur_feat))[None].float().to(self.device)

        gen = None
        if self.sample_seed is not None:
            gen = torch.Generator(device=self.device)
            gen.manual_seed(self.sample_seed + self.steps)
        delta_norm = self.model.sample_delta(
            self.latent, cur_t, self.static_t, self.obj_mask, self.gravity_dt,
            steps=self.denoise_steps, generator=gen,
        )
        delta = self.norm.denorm_delta(delta_norm[0].float().cpu().numpy())
        self.state = self.state + delta

        text = self._render_prediction()
        self.frame_num += 1
        self.steps += 1
        return text

    def _render_prediction(self) -> str:
        """Matches the baseline prediction output format (parsed by _extract_physics)."""
        lines = [f"Frame {self.frame_num}:"]
        for i in range(self.state.shape[0]):
            px, py, vx, vy, a, av = (float(x) for x in self.state[i])
            s = f"  obj_{int(self.ids[i])}: pos=({px:.4f}, {py:.4f}), vel=({vx:.4f}, {vy:.4f})"
            if abs(a) > 0.001 or abs(av) > 0.001:
                s += f", a={a:.4f}, av={av:.4f}"
            lines.append(s)
        return "\n".join(lines)