from __future__ import annotations

import copy

import torch
import torch.nn.functional as F
from torch import Tensor, nn

from hippotorch.core.episode import Episode


class _NatureCNN(nn.Module):
    def __init__(self, in_channels: int) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_channels, 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((4, 4)),
        )

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


class _VisualBackbone(nn.Module):
    def __init__(
        self,
        image_channels: int,
        *,
        feature_dim: int = 256,
        temporal_hidden: int = 256,
        temporal_layers: int = 1,
        embed_dim: int = 128,
    ) -> None:
        super().__init__()
        self.cnn = _NatureCNN(image_channels)
        self.frame_proj = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, feature_dim),
            nn.ReLU(),
        )
        self.temporal = nn.LSTM(
            feature_dim,
            temporal_hidden,
            num_layers=max(1, temporal_layers),
            batch_first=True,
        )
        self.head = nn.Sequential(
            nn.Linear(temporal_hidden, embed_dim),
            nn.LayerNorm(embed_dim),
        )
        self.embed_dim = embed_dim

    def forward(self, x: Tensor) -> Tensor:
        B, T, C, H, W = x.shape
        frames = x.view(B * T, C, H, W)
        feats = self.cnn(frames)
        feats = self.frame_proj(feats)
        feats = feats.view(B, T, -1)
        outputs, _ = self.temporal(feats)
        embeds = self.head(outputs)
        return F.normalize(embeds, dim=-1)


class VisualEpisodeEncoder(nn.Module):
    """Visual encoder with NatureCNN trunk + LSTM aggregator and EMA target."""

    def __init__(
        self,
        image_channels: int = 3,
        embed_dim: int = 128,
        *,
        feature_dim: int = 256,
        temporal_hidden: int = 256,
        temporal_layers: int = 1,
        momentum: float = 0.995,
        refresh_interval: int = 5000,
    ) -> None:
        super().__init__()
        self.embed_dim = embed_dim
        self.momentum = float(momentum)
        self.refresh_interval = int(refresh_interval)
        self._steps_since_refresh = 0
        self.online = _VisualBackbone(
            image_channels,
            feature_dim=feature_dim,
            temporal_hidden=temporal_hidden,
            temporal_layers=temporal_layers,
            embed_dim=embed_dim,
        )
        self.target = copy.deepcopy(self.online)
        self._sync_target()

    def _sync_target(self) -> None:
        for tgt, src in zip(self.target.parameters(), self.online.parameters()):
            tgt.data.copy_(src.data)
            tgt.requires_grad = False

    def _prepare_input(self, frames: Tensor) -> Tensor:
        if frames.dim() == 4:
            frames = frames.unsqueeze(0)
        if frames.dim() != 5:
            raise ValueError(
                f"expected [B, T, C, H, W] or [T, C, H, W], got shape {tuple(frames.shape)}"
            )
        device = next(self.parameters()).device
        frames = frames.to(device)
        if frames.dtype == torch.uint8:
            frames = frames.to(torch.float32) / 255.0
        else:
            frames = frames.to(torch.float32)
        return frames

    def forward(self, frames: Tensor) -> Tensor:
        prepared = self._prepare_input(frames)
        return self.online(prepared)

    def prepare_episode(self, episode: Episode) -> Tensor:
        """Return the raw state frames for encoding."""

        frames = episode.states
        if frames.dim() == 3:  # flattened?
            raise ValueError("visual episodes require states shaped [T, C, H, W]")
        return frames

    def prepare_query(self, frames: Tensor) -> Tensor:
        return self._prepare_input(frames)

    def encode_query(self, frames: Tensor) -> Tensor:
        seq = self.forward(frames)
        return seq.mean(dim=1)

    @torch.no_grad()
    def encode_key(self, frames: Tensor) -> Tensor:
        prepared = self._prepare_input(frames)
        seq = self.target(prepared)
        return seq.mean(dim=1)

    @torch.no_grad()
    def update_target(self) -> None:
        for tgt, src in zip(self.target.parameters(), self.online.parameters()):
            tgt.data = self.momentum * tgt.data + (1 - self.momentum) * src.data
        self._steps_since_refresh += 1

    def should_refresh_keys(self) -> bool:
        return self._steps_since_refresh >= self.refresh_interval

    def mark_refreshed(self) -> None:
        self._steps_since_refresh = 0


__all__ = ["VisualEpisodeEncoder"]
