"""H2 adaptation of residual TD3 networks from sim commit 4fa48dd.

Layer layout and zero actor initialization are preserved; dimensions come
from the H2 contract. Normalizer checkpoints retain full running statistics.
"""

from __future__ import annotations

import numpy as np
import torch
from torch import nn

from .data import ACTOR_INPUT_DIM, RESIDUAL_ACTION_DIM


class Actor(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(ACTOR_INPUT_DIM, 256),
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Linear(256, RESIDUAL_ACTION_DIM),
            nn.Tanh(),
        )
        output = self.net[-2]
        nn.init.zeros_(output.weight)
        nn.init.zeros_(output.bias)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.net(inputs)


class Critic(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(ACTOR_INPUT_DIM, 256),
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Linear(256, 1),
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.net(inputs)


def parameter_count(module: nn.Module) -> int:
    return sum(parameter.numel() for parameter in module.parameters())


def actor_parameter_count() -> int:
    with torch.random.fork_rng(devices=[]):
        return parameter_count(Actor())


def flatten_parameters(module: nn.Module) -> np.ndarray:
    values = [parameter.detach().cpu().reshape(-1) for parameter in module.parameters()]
    return torch.cat(values).numpy().astype(np.float32, copy=True)


def load_flat_parameters(module: nn.Module, values: np.ndarray) -> None:
    flat = np.asarray(values, dtype=np.float32).reshape(-1)
    if flat.size != parameter_count(module) or not np.isfinite(flat).all():
        raise ValueError("actor parameter payload is invalid")
    offset = 0
    with torch.no_grad():
        for parameter in module.parameters():
            count = parameter.numel()
            source = torch.from_numpy(flat[offset : offset + count]).reshape(parameter.shape)
            parameter.copy_(source.to(device=parameter.device, dtype=parameter.dtype))
            offset += count


class RunningNormalizer:
    def __init__(self, size: int, *, min_std: float, clip: float) -> None:
        self.size = int(size)
        self.min_std = float(min_std)
        self.clip = float(clip)
        self.count = 0
        self.mean = np.zeros(self.size, dtype=np.float64)
        self.m2 = np.zeros(self.size, dtype=np.float64)
        self.frozen = False

    def update(self, values: np.ndarray) -> None:
        if self.frozen:
            return
        value = np.asarray(values, dtype=np.float64).reshape(self.size)
        self.count += 1
        delta = value - self.mean
        self.mean += delta / self.count
        self.m2 += delta * (value - self.mean)

    def freeze(self) -> None:
        self.frozen = True

    @property
    def std(self) -> np.ndarray:
        if self.count < 2:
            variance = np.ones(self.size, dtype=np.float64)
        else:
            variance = self.m2 / self.count
        return np.maximum(np.sqrt(variance), self.min_std).astype(np.float32)

    def normalize_tensor(self, values: torch.Tensor) -> torch.Tensor:
        mean = torch.as_tensor(self.mean, dtype=values.dtype, device=values.device)
        std = torch.as_tensor(self.std, dtype=values.dtype, device=values.device)
        return torch.clamp((values - mean) / std, -self.clip, self.clip)

    def state_dict(self) -> dict:
        return {
            "size": self.size,
            "min_std": self.min_std,
            "clip": self.clip,
            "count": self.count,
            "mean": self.mean.copy(),
            "std": self.std,
            "m2": self.m2.copy(),
            "frozen": self.frozen,
        }