from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn as nn


@dataclass
class NormalizerStats:
    min: torch.Tensor
    max: torch.Tensor
    mean: torch.Tensor
    std: torch.Tensor


class TensorLinearNormalizer(nn.Module):
    """Linear normalizer for tensors whose last dimension is feature dimension."""

    def __init__(self, eps: float = 1e-6, dim: int = 1):
        super().__init__()
        self.eps = eps
        self.register_buffer("scale", torch.ones(dim), persistent=True)
        self.register_buffer("offset", torch.zeros(dim), persistent=True)
        self.register_buffer("fitted", torch.tensor(False), persistent=True)

    @torch.no_grad()
    def fit(
        self,
        data: torch.Tensor,
        mode: str = "limits",
        output_min: float = -1.0,
        output_max: float = 1.0,
        fit_offset: bool = True,
    ) -> None:
        if data.ndim < 2:
            raise ValueError(
                f"Expected at least 2 dims [N, ..., D], got shape={tuple(data.shape)}"
            )

        flat = data.reshape(-1, data.shape[-1]).float()
        in_min = flat.min(dim=0).values
        in_max = flat.max(dim=0).values
        in_mean = flat.mean(dim=0)
        in_std = flat.std(dim=0).clamp_min(self.eps)

        if mode == "limits":
            if fit_offset:
                in_range = (in_max - in_min).clamp_min(self.eps)
                scale = (output_max - output_min) / in_range
                offset = output_min - scale * in_min
            else:
                output_abs = min(abs(output_min), abs(output_max))
                in_abs = torch.maximum(in_min.abs(), in_max.abs()).clamp_min(self.eps)
                scale = output_abs / in_abs
                offset = torch.zeros_like(scale)
        elif mode == "gaussian":
            scale = 1.0 / in_std
            offset = -in_mean * scale if fit_offset else torch.zeros_like(scale)
        else:
            raise ValueError(f"Unsupported mode: {mode}")

        self.scale = scale
        self.offset = offset
        self.fitted = torch.tensor(True, device=scale.device)

    def is_fitted(self) -> bool:
        return bool(self.fitted.item())

    def normalize(self, x: torch.Tensor) -> torch.Tensor:
        if not self.is_fitted():
            raise RuntimeError("Normalizer must be fitted before normalize().")
        scale = self.scale.to(device=x.device, dtype=x.dtype)
        offset = self.offset.to(device=x.device, dtype=x.dtype)
        return x * scale + offset

    def unnormalize(self, x: torch.Tensor) -> torch.Tensor:
        if not self.is_fitted():
            raise RuntimeError("Normalizer must be fitted before unnormalize().")
        scale = self.scale.to(device=x.device, dtype=x.dtype)
        offset = self.offset.to(device=x.device, dtype=x.dtype)
        return (x - offset) / scale


class ActionNormalizer(nn.Module):
    """Simple keyed normalizer compatible with tokenizer batches that use `batch['action']`."""

    def __init__(
        self, action: Optional[TensorLinearNormalizer] = None, action_dim: int = 1
    ):
        super().__init__()
        self.action = (
            action if action is not None else TensorLinearNormalizer(dim=action_dim)
        )

    def __getitem__(self, key: str) -> TensorLinearNormalizer:
        if key != "action":
            raise KeyError(f"Unsupported key: {key}")
        return self.action

    def load_state_dict(self, state_dict, strict: bool = True):
        """Load state dict while adapting buffer shapes to saved per-DoF stats."""
        scale = state_dict.get("action.scale")
        offset = state_dict.get("action.offset")

        if scale is not None:
            scale = torch.as_tensor(scale)
            if tuple(self.action.scale.shape) != tuple(scale.shape):
                self.action.scale = torch.ones_like(scale)
        if offset is not None:
            offset = torch.as_tensor(offset)
            if tuple(self.action.offset.shape) != tuple(offset.shape):
                self.action.offset = torch.zeros_like(offset)

        return super().load_state_dict(state_dict, strict=strict)

    @torch.no_grad()
    def fit_action(self, actions: torch.Tensor, **kwargs) -> None:
        self.action.fit(actions, **kwargs)
