"""Dual-rate physics model: flow-matching loss and Euler sampler.

Flow formulation mirrors tinyvla (modeling_tinyvla.py:229-245 / 266-279):
    t ~ Beta(alpha, beta) scaled into (0.001, 1.0]
    x_t = t * noise + (1 - t) * data        (noise at t=1, data at t=0)
    velocity target = noise - data
    sampling: x ~ N(0,1) at t=1, Euler x <- x - dt * v down to t=0
"""

from __future__ import annotations

import torch
import torch.nn as nn

from .expert import PhysicsFlowExpert
from .semantic import SemanticEncoder


class DualRatePhysicsModel(nn.Module):
    def __init__(self, cfg, load_pretrained: bool = True, lm: nn.Module | None = None,
                 readout_token_id: int = 0):
        super().__init__()
        self.cfg = cfg
        self.semantic = SemanticEncoder(cfg, load_pretrained=load_pretrained, lm=lm,
                                        readout_token_id=readout_token_id)
        self.expert = PhysicsFlowExpert(
            state_dim=cfg.state_dim,
            static_dim=cfg.static_dim,
            max_objects=cfg.max_objects,
            d_model=cfg.expert_dim,
            n_layers=cfg.expert_layers,
            n_heads=cfg.expert_heads,
        )
        # conditioning projections: global (gravity + dt) and per-object tokens
        self.global_cond = nn.Linear(3, cfg.expert_dim)
        self.obj_cond = nn.Linear(cfg.state_dim + cfg.static_dim, cfg.expert_dim)

    def trainable_parameters(self):
        return [p for p in self.parameters() if p.requires_grad]

    def compute_latent(self, input_ids, attention_mask):
        return self.semantic(input_ids, attention_mask)

    def build_cond(self, latent, cur_state, static, obj_mask, gravity_dt):
        """latent (B,R,D), cur_state (B,N,S) normalized, static (B,N,F),
        obj_mask (B,N) bool, gravity_dt (B,3). Returns cond (B,R+1+N,D), mask."""
        b, n = obj_mask.shape
        g = self.global_cond(gravity_dt)[:, None, :]
        o = self.obj_cond(torch.cat([cur_state, static], dim=-1))
        o = o + self.expert.obj_emb[:, :n]
        cond = torch.cat([latent, g, o], dim=1)
        ones = torch.ones((b, latent.shape[1] + 1), dtype=torch.bool, device=obj_mask.device)
        cond_mask = torch.cat([ones, obj_mask], dim=1)
        return cond, cond_mask

    def flow_loss(self, latent, cur_state, static, obj_mask, target_delta, gravity_dt,
                  generator: torch.Generator | None = None):
        """All state tensors normalized. Returns scalar masked-MSE flow loss."""
        b = target_delta.shape[0]
        device = target_delta.device
        alpha = torch.full((b,), self.cfg.flow_beta_alpha, device=device)
        beta = torch.full((b,), self.cfg.flow_beta_beta, device=device)
        t = torch.distributions.Beta(alpha, beta).sample() * 0.999 + 0.001
        noise = torch.randn(target_delta.shape, device=device, generator=generator,
                            dtype=target_delta.dtype)
        t_ = t[:, None, None]
        x_t = t_ * noise + (1 - t_) * target_delta
        v_target = noise - target_delta

        cond, cond_mask = self.build_cond(latent, cur_state, static, obj_mask, gravity_dt)
        v_pred = self.expert(x_t, t, cur_state, static, obj_mask, cond, cond_mask)

        se = (v_pred - v_target) ** 2
        mask = obj_mask[..., None].to(se.dtype)
        return (se * mask).sum() / mask.sum().clamp(min=1) / se.shape[-1]

    @torch.no_grad()
    def sample_delta(self, latent, cur_state, static, obj_mask, gravity_dt,
                     steps: int | None = None, generator: torch.Generator | None = None):
        """Euler-integrate the flow from noise to a normalized delta prediction."""
        steps = steps or self.cfg.denoise_steps
        b, n = obj_mask.shape
        device = cur_state.device
        cond, cond_mask = self.build_cond(latent, cur_state, static, obj_mask, gravity_dt)
        x = torch.randn((b, n, self.cfg.state_dim), device=device, generator=generator,
                        dtype=cur_state.dtype)
        dt_f = 1.0 / steps
        for i in range(steps):
            t = torch.full((b,), 1.0 - i * dt_f, device=device, dtype=cur_state.dtype)
            v = self.expert(x, t, cur_state, static, obj_mask, cond, cond_mask)
            x = x - dt_f * v
        return x