"""Wall-X Real-Time Chunking helpers.

The guided update follows LeRobot's RTC implementation, adapted to Wall-X's
flow convention: Wall-X integrates normalized actions from noise at t=0 to a
clean action at t=1 with velocity ``dx/dt``.
"""

from __future__ import annotations

from dataclasses import dataclass
import math

import torch


@dataclass(frozen=True)
class WallXRTCConfig:
    enabled: bool = True
    execution_horizon: int = 6
    max_guidance_weight: float = 10.0
    prefix_attention_schedule: str = "linear"

    def __post_init__(self):
        if self.execution_horizon <= 0:
            raise ValueError("execution_horizon must be positive")
        if self.max_guidance_weight <= 0:
            raise ValueError("max_guidance_weight must be positive")
        if self.prefix_attention_schedule not in {"zeros", "ones", "linear", "exp"}:
            raise ValueError(
                "prefix_attention_schedule must be zeros, ones, linear, or exp"
            )


class WallXRTCProcessor:
    """Inference-time RTC guidance in normalized Wall-X action space."""

    def __init__(self, config: WallXRTCConfig):
        self.config = config

    def get_prefix_weights(self, start: int, end: int, total: int) -> torch.Tensor:
        start = max(0, min(int(start), int(end), int(total)))
        end = max(start, min(int(end), int(total)))
        schedule = self.config.prefix_attention_schedule
        if schedule == "zeros":
            weights = torch.zeros(total)
            weights[:start] = 1.0
            return weights
        if schedule == "ones":
            weights = torch.zeros(total)
            weights[:end] = 1.0
            return weights

        middle_len = end - start
        if middle_len:
            middle = torch.linspace(1.0, 0.0, middle_len + 2)[1:-1]
            if schedule == "exp":
                middle = middle * torch.expm1(middle) / (math.e - 1.0)
        else:
            middle = torch.empty(0)
        return torch.cat(
            [torch.ones(start), middle, torch.zeros(total - end)], dim=0
        )

    def guide_increasing_flow(
        self,
        *,
        x_t: torch.Tensor,
        time: torch.Tensor | float,
        predict_velocity,
        prev_chunk_left_over: torch.Tensor | None,
        inference_delay: int = 0,
        execution_horizon: int | None = None,
    ) -> torch.Tensor:
        """Return RTC-guided velocity for a flow integrated from t=0 to t=1."""
        if prev_chunk_left_over is None or not self.config.enabled:
            return predict_velocity(x_t)

        x = x_t.detach().clone().requires_grad_(True)
        prefix = prev_chunk_left_over.to(device=x.device, dtype=x.dtype)
        if prefix.ndim == 2:
            prefix = prefix.unsqueeze(0)
        if prefix.shape[0] == 1 and x.shape[0] > 1:
            prefix = prefix.expand(x.shape[0], -1, -1)

        padded = torch.zeros_like(x)
        steps = min(prefix.shape[1], x.shape[1])
        dims = min(prefix.shape[2], x.shape[2])
        padded[:, :steps, :dims] = prefix[:, :steps, :dims]

        horizon = execution_horizon or self.config.execution_horizon
        horizon = min(int(horizon), steps, x.shape[1])
        weights = self.get_prefix_weights(inference_delay, horizon, x.shape[1])
        weights = weights.to(device=x.device, dtype=x.dtype).view(1, -1, 1)

        with torch.enable_grad():
            velocity = predict_velocity(x)
            t = torch.as_tensor(time, device=x.device, dtype=x.dtype)
            remaining = torch.clamp(1.0 - t, min=1e-6)
            clean_estimate = x + remaining * velocity
            error = (padded - clean_estimate) * weights
            correction = torch.autograd.grad(
                clean_estimate,
                x,
                grad_outputs=error.detach(),
                retain_graph=False,
            )[0]

        # Same guidance schedule as LeRobot RTC after mapping its 1->0 time
        # convention to Wall-X's 0->1 convention.
        t_safe = torch.clamp(t, min=1e-6)
        inv_r2 = (remaining.square() + t.square()) / remaining.square()
        weight = (remaining / t_safe) * inv_r2
        weight = torch.nan_to_num(
            weight,
            nan=self.config.max_guidance_weight,
            posinf=self.config.max_guidance_weight,
        ).clamp(max=self.config.max_guidance_weight)
        return (velocity + weight * correction).detach()