import torch
import torch.nn as nn
from gymnasium import spaces
from skrl.models.torch import Model, GaussianMixin, DeterministicMixin


class HistoryEncoder(nn.Module):
    """1D-CNN encoder over a temporal window of (obs, action) pairs.

    Input:  (batch, history_length, step_dim)
    Output: (batch, embedding_dim)

    Architecture: two temporal conv layers → global average pool → linear.
    Lets the policy implicitly infer the current dynamics (friction, torque
    scale, latency, …) from how the system responded to recent actions —
    end-to-end adaptation when trained under domain randomization.
    """

    def __init__(
        self,
        history_length: int,
        step_dim: int,
        embedding_dim: int = 32,
        hidden_channels: int = 32,
    ) -> None:
        super().__init__()
        self.conv = nn.Sequential(
            # (batch, step_dim, history_length) after transpose
            nn.Conv1d(step_dim, hidden_channels, kernel_size=3, padding=1),
            nn.ELU(),
            nn.Conv1d(hidden_channels, hidden_channels, kernel_size=3, padding=1),
            nn.ELU(),
        )
        self.fc = nn.Linear(hidden_channels, embedding_dim)

    def forward(self, history: torch.Tensor) -> torch.Tensor:
        """history: (batch, history_length, step_dim)."""
        # Conv1d expects (batch, channels, seq_len).
        x = history.transpose(1, 2)
        x = self.conv(x)
        # Global average pool over time.
        x = x.mean(dim=-1)
        return self.fc(x)


class SharedMLP(GaussianMixin, DeterministicMixin, Model):
    """Shared policy/value network with an optional history encoder.

    With ``history_length > 0`` the input states are expected to be
    ``[raw_obs, history_flat]`` (as produced by ``BaseRunner``); the history
    window is compressed by a :class:`HistoryEncoder` and concatenated with
    the raw observation before the shared backbone.
    """

    def __init__(
        self,
        observation_space: spaces.Space,
        action_space: spaces.Space,
        device: torch.device,
        hidden_sizes: tuple[int, ...] = (32, 32),
        clip_actions: bool = False,
        clip_log_std: bool = True,
        min_log_std: float = -2.0,
        max_log_std: float = 2.0,
        initial_log_std: float = 0.0,
        # ── History encoder ──────────────────────────────────────
        history_length: int = 0,
        raw_obs_dim: int = 0,
        embedding_dim: int = 32,
    ):
        Model.__init__(self, observation_space, action_space, device)
        GaussianMixin.__init__(
            self, clip_actions, clip_log_std, min_log_std, max_log_std,
        )
        DeterministicMixin.__init__(self, clip_actions)

        self._history_length = history_length
        self._raw_obs_dim = raw_obs_dim
        self._embedding_dim = embedding_dim

        self.history_encoder: HistoryEncoder | None = None
        if history_length > 0 and raw_obs_dim > 0:
            step_dim = raw_obs_dim + self.num_actions
            self.history_encoder = HistoryEncoder(
                history_length=history_length,
                step_dim=step_dim,
                embedding_dim=embedding_dim,
            )
            in_dim = raw_obs_dim + embedding_dim
        else:
            in_dim = self.num_observations

        # ── Shared backbone ──────────────────────────────────────
        layers: list[nn.Module] = []
        for hidden_size in hidden_sizes:
            layers.append(nn.Linear(in_dim, hidden_size))
            layers.append(nn.ELU())
            in_dim = hidden_size
        self.net: nn.Sequential = nn.Sequential(*layers)

        # Policy head
        self.mean_layer = nn.Linear(in_dim, self.num_actions)
        self.log_std_parameter: nn.Parameter = nn.Parameter(
            torch.full((self.num_actions,), initial_log_std),
        )

        # Value head
        self.value_layer = nn.Linear(in_dim, 1)
        self._shared_output: torch.Tensor | None = None

    def act(
        self, inputs: dict[str, torch.Tensor], role: str = "",
    ) -> tuple[torch.Tensor, ...]:
        if role == "policy":
            return GaussianMixin.act(self, inputs, role)
        elif role == "value":
            return DeterministicMixin.act(self, inputs, role)

    def _encode(self, states: torch.Tensor) -> torch.Tensor:
        """Optionally split off and encode the history window."""
        if self.history_encoder is None:
            return self.net(states)

        obs = states[:, :self._raw_obs_dim]
        hist_flat = states[:, self._raw_obs_dim:]
        step_dim = self._raw_obs_dim + self.num_actions
        history = hist_flat.reshape(-1, self._history_length, step_dim)
        embedding = self.history_encoder(history)
        return self.net(torch.cat([obs, embedding], dim=-1))

    def compute(
        self, inputs: dict[str, torch.Tensor], role: str = "",
    ) -> tuple[torch.Tensor, ...]:
        if role == "policy":
            self._shared_output = self._encode(inputs["states"])
            return self.mean_layer(self._shared_output), self.log_std_parameter, {}
        elif role == "value":
            shared_output = (
                self._shared_output
                if self._shared_output is not None
                else self._encode(inputs["states"])
            )
            self._shared_output = None
            return self.value_layer(shared_output), {}
