from __future__ import annotations

from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F


FeaturePyramid = list[torch.Tensor]
StreamPyramid = list[list[torch.Tensor]]
__all__ = [
    "FeaturePyramid",
    "FlexibleEvidenceBlock",
    "FlexibleEvidenceProjector",
    "SpatialCrossAttentionSharedAdapter",
    "StreamEvidenceBuilder",
    "StreamPyramid",
    "Transformer2DBlock",
    "TransformerEvidenceMaskAdapter",
]


def _group_count(channels: int, max_groups: int = 8) -> int:
    for groups in range(min(max_groups, channels), 0, -1):
        if channels % groups == 0:
            return groups
    return 1


class FlexibleEvidenceBlock(nn.Module):
    def __init__(self, channels: int, dropout_rate: float = 0.1, expansion: int = 2) -> None:
        super().__init__()
        expanded = channels * expansion
        self.norm = nn.GroupNorm(_group_count(channels), channels)
        self.expand = nn.Conv2d(channels, expanded, kernel_size=1)
        self.depthwise = nn.Conv2d(expanded, expanded, kernel_size=3, padding=1, groups=expanded)
        self.act = nn.GELU()
        self.channel_gate = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(expanded, max(channels // 4, 8), kernel_size=1),
            nn.GELU(),
            nn.Conv2d(max(channels // 4, 8), expanded, kernel_size=1),
            nn.Sigmoid(),
        )
        self.project = nn.Conv2d(expanded, channels, kernel_size=1)
        self.dropout = nn.Dropout2d(p=dropout_rate)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        hidden = self.norm(x)
        hidden = self.expand(hidden)
        hidden = self.depthwise(hidden)
        hidden = self.act(hidden)
        hidden = hidden * self.channel_gate(hidden)
        hidden = self.project(hidden)
        hidden = self.dropout(hidden)
        return residual + hidden


class FlexibleEvidenceProjector(nn.Module):
    """Channel and local-spatial evidence projector used by v2 adapters."""

    def __init__(
        self,
        in_channels: int,
        hidden_dim: int,
        dropout_rate: float = 0.1,
        depth: int = 2,
    ) -> None:
        super().__init__()
        self.input_proj = nn.Sequential(
            nn.Conv2d(in_channels, hidden_dim, kernel_size=1),
            nn.GroupNorm(_group_count(hidden_dim), hidden_dim),
            nn.GELU(),
            nn.Dropout2d(p=dropout_rate),
        )
        self.blocks = nn.Sequential(
            *[
                FlexibleEvidenceBlock(hidden_dim, dropout_rate=dropout_rate)
                for _ in range(depth)
            ]
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.blocks(self.input_proj(x))


class StreamEvidenceBuilder(nn.Module):
    """Build target evidence from raw target and perturbation-stream features."""

    def __init__(
        self,
        in_channels_list: list[int],
        hidden_dim: int,
        dropout_rate: float = 0.1,
        max_streams: int = 3,
    ) -> None:
        super().__init__()
        self.max_streams = max(1, max_streams)
        self.projections = nn.ModuleList(
            [
                FlexibleEvidenceProjector(
                    in_channels=channels * (1 + 3 * self.max_streams),
                    hidden_dim=hidden_dim,
                    dropout_rate=dropout_rate,
                    depth=3,
                )
                for channels in in_channels_list
            ]
        )

    def forward(
        self,
        unadapted: FeaturePyramid,
        streams_unadapted: StreamPyramid,
    ) -> FeaturePyramid:
        evidence = []
        for scale_idx, (target, streams) in enumerate(zip(unadapted, streams_unadapted)):
            fixed_streams = list(streams[: self.max_streams])
            while len(fixed_streams) < self.max_streams:
                fixed_streams.append(torch.zeros_like(target))

            abs_diffs = [(target - stream).abs() for stream in fixed_streams]
            signed_diffs = [target - stream for stream in fixed_streams]
            raw = torch.cat([target, *fixed_streams, *abs_diffs, *signed_diffs], dim=1)
            evidence.append(self.projections[scale_idx](raw))
        return evidence


class CrossAttentionSharedAdapter(nn.Module):
    """Per-scale cross-attention adapter over perturbation streams."""

    def __init__(
        self,
        in_channels_list: list[int],
        hidden_dim: int,
        dropout_rate: float = 0.1,
        max_streams: int = 2,
        num_heads: int = 4,
    ) -> None:
        super().__init__()
        self.num_scales = len(in_channels_list)
        self.max_streams = max(1, max_streams)
        self.num_heads = max(1, num_heads)
        if hidden_dim % self.num_heads != 0:
            raise ValueError(f"hidden_dim ({hidden_dim}) must be divisible by num_heads ({self.num_heads})")

        self.head_dim = hidden_dim // self.num_heads
        self.scale = self.head_dim**-0.5
        self.q_proj = nn.ModuleList(nn.Conv2d(channels, hidden_dim, kernel_size=1) for channels in in_channels_list)
        self.k_proj = nn.ModuleList(nn.Conv2d(channels, hidden_dim, kernel_size=1) for channels in in_channels_list)
        self.v_proj = nn.ModuleList(nn.Conv2d(channels, hidden_dim, kernel_size=1) for channels in in_channels_list)
        self.out_proj = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1),
                    nn.GELU(),
                    nn.Dropout2d(p=dropout_rate),
                    nn.Conv2d(hidden_dim, channels, kernel_size=1),
                )
                for channels in in_channels_list
            ]
        )
        self.evidence_q_proj = nn.ModuleList(
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1)
            for _ in in_channels_list
        )
        self.evidence_context_proj = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1),
                    nn.GELU(),
                    nn.Dropout2d(p=dropout_rate),
                )
                for _ in in_channels_list
            ]
        )
        self.residual_gate = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Conv2d(channels + hidden_dim, channels, kernel_size=1),
                    nn.Sigmoid(),
                )
                for channels in in_channels_list
            ]
        )
        self.dropout = nn.Dropout2d(p=dropout_rate)

    def forward(
        self,
        stream_features: list[torch.Tensor],
        unadapted: torch.Tensor,
        scale_idx: int,
        return_delta: bool = False,
        evidence: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if not stream_features:
            delta = torch.zeros_like(unadapted)
            adapted = unadapted
            if return_delta:
                return adapted, delta
            return adapted

        batch_size, _, height, width = unadapted.shape
        stream_count = len(stream_features)
        query = self.q_proj[scale_idx](unadapted)
        if evidence is not None:
            if evidence.shape[-2:] != query.shape[-2:]:
                evidence = F.interpolate(evidence, size=query.shape[-2:], mode="bilinear", align_corners=False)
            query = query + self.evidence_q_proj[scale_idx](evidence)
        query = query.view(batch_size, self.num_heads, self.head_dim, 1, height, width)

        keys = [self.k_proj[scale_idx](feature) for feature in stream_features]
        values = [self.v_proj[scale_idx](feature) for feature in stream_features]
        key = torch.stack(keys, dim=1).view(batch_size, stream_count, self.num_heads, self.head_dim, height, width)
        value = torch.stack(values, dim=1).view(batch_size, stream_count, self.num_heads, self.head_dim, height, width)
        key = key.permute(0, 2, 3, 1, 4, 5)
        value = value.permute(0, 2, 3, 1, 4, 5)

        attention_logits = (query * key).sum(dim=2) * self.scale
        attention = F.softmax(attention_logits, dim=2)
        context = (attention.unsqueeze(2) * value).sum(dim=3)
        context = context.reshape(batch_size, self.num_heads * self.head_dim, height, width)
        if evidence is not None:
            context = context + self.evidence_context_proj[scale_idx](evidence)
        context = self.dropout(context)

        raw_delta = self.out_proj[scale_idx](context)
        gate = self.residual_gate[scale_idx](torch.cat([unadapted, context], dim=1))
        delta = gate * raw_delta
        adapted = unadapted + delta
        if return_delta:
            return adapted, delta
        return adapted


class SpatialCrossAttentionSharedAdapter(CrossAttentionSharedAdapter):
    """Cross-attention adapter with global low-res and windowed mid-res spatial mixing."""

    def __init__(
        self,
        in_channels_list: list[int],
        hidden_dim: int,
        dropout_rate: float = 0.1,
        max_streams: int = 2,
        num_heads: int = 4,
        global_spatial_scales: tuple[int, ...] = (0,),
        windowed_spatial_scales: tuple[int, ...] = (2,),
        window_size: int = 8,
    ) -> None:
        super().__init__(
            in_channels_list=in_channels_list,
            hidden_dim=hidden_dim,
            dropout_rate=dropout_rate,
            max_streams=max_streams,
            num_heads=num_heads,
        )
        self.global_spatial_scales = set(global_spatial_scales)
        self.windowed_spatial_scales = set(windowed_spatial_scales)
        self.window_size = window_size

    def _apply_residual(
        self,
        context: torch.Tensor,
        unadapted: torch.Tensor,
        scale_idx: int,
        return_delta: bool,
        evidence: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if evidence is not None:
            if evidence.shape[-2:] != context.shape[-2:]:
                evidence = F.interpolate(evidence, size=context.shape[-2:], mode="bilinear", align_corners=False)
            context = context + self.evidence_context_proj[scale_idx](evidence)
        context = self.dropout(context)
        raw_delta = self.out_proj[scale_idx](context)
        gate = self.residual_gate[scale_idx](torch.cat([unadapted, context], dim=1))
        delta = gate * raw_delta
        adapted = unadapted + delta
        if return_delta:
            return adapted, delta
        return adapted

    def _global_spatial_context(
        self,
        stream_features: list[torch.Tensor],
        unadapted: torch.Tensor,
        scale_idx: int,
        evidence: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        batch_size, _, height, width = unadapted.shape
        stream_count = len(stream_features)
        num_tokens = height * width
        query = self.q_proj[scale_idx](unadapted)
        if evidence is not None:
            if evidence.shape[-2:] != query.shape[-2:]:
                evidence = F.interpolate(evidence, size=query.shape[-2:], mode="bilinear", align_corners=False)
            query = query + self.evidence_q_proj[scale_idx](evidence)
        query = query.view(batch_size, self.num_heads, self.head_dim, num_tokens).transpose(2, 3)

        keys = [self.k_proj[scale_idx](feature) for feature in stream_features]
        values = [self.v_proj[scale_idx](feature) for feature in stream_features]
        key = torch.stack(keys, dim=1).view(batch_size, stream_count, self.num_heads, self.head_dim, num_tokens)
        value = torch.stack(values, dim=1).view(batch_size, stream_count, self.num_heads, self.head_dim, num_tokens)
        key = key.permute(0, 2, 1, 4, 3).reshape(batch_size, self.num_heads, stream_count * num_tokens, self.head_dim)
        value = value.permute(0, 2, 1, 4, 3).reshape(batch_size, self.num_heads, stream_count * num_tokens, self.head_dim)

        attention = torch.matmul(query, key.transpose(-1, -2)) * self.scale
        attention = F.softmax(attention, dim=-1)
        context = torch.matmul(attention, value)
        return context.transpose(2, 3).reshape(batch_size, self.num_heads * self.head_dim, height, width)

    def _window_partition_q(self, x: torch.Tensor, window_size: int) -> torch.Tensor:
        batch_size, heads, head_dim, height, width = x.shape
        return (
            x.permute(0, 3, 4, 1, 2)
            .reshape(batch_size, height // window_size, window_size, width // window_size, window_size, heads, head_dim)
            .permute(0, 1, 3, 5, 2, 4, 6)
            .reshape(batch_size * (height // window_size) * (width // window_size), heads, window_size * window_size, head_dim)
        )

    def _window_partition_kv(self, x: torch.Tensor, window_size: int) -> torch.Tensor:
        batch_size, streams, heads, head_dim, height, width = x.shape
        return (
            x.permute(0, 4, 5, 1, 2, 3)
            .reshape(
                batch_size,
                height // window_size,
                window_size,
                width // window_size,
                window_size,
                streams,
                heads,
                head_dim,
            )
            .permute(0, 1, 3, 6, 2, 4, 5, 7)
            .reshape(
                batch_size * (height // window_size) * (width // window_size),
                heads,
                window_size * window_size * streams,
                head_dim,
            )
        )

    def _window_unpartition(
        self,
        x: torch.Tensor,
        batch_size: int,
        height: int,
        width: int,
        window_size: int,
    ) -> torch.Tensor:
        heads = x.shape[1]
        head_dim = x.shape[-1]
        windows_h = height // window_size
        windows_w = width // window_size
        return (
            x.reshape(batch_size, windows_h, windows_w, heads, window_size, window_size, head_dim)
            .permute(0, 3, 6, 1, 4, 2, 5)
            .reshape(batch_size, heads, head_dim, height, width)
        )

    def _windowed_spatial_context(
        self,
        stream_features: list[torch.Tensor],
        unadapted: torch.Tensor,
        scale_idx: int,
        evidence: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        batch_size, _, height, width = unadapted.shape
        stream_count = len(stream_features)
        window_size = self.window_size
        pad_h = (window_size - height % window_size) % window_size
        pad_w = (window_size - width % window_size) % window_size
        padded_h = height + pad_h
        padded_w = width + pad_w

        query = self.q_proj[scale_idx](unadapted)
        if evidence is not None:
            if evidence.shape[-2:] != query.shape[-2:]:
                evidence = F.interpolate(evidence, size=query.shape[-2:], mode="bilinear", align_corners=False)
            query = query + self.evidence_q_proj[scale_idx](evidence)
        query = query.view(batch_size, self.num_heads, self.head_dim, height, width)
        keys = [
            self.k_proj[scale_idx](feature).view(batch_size, self.num_heads, self.head_dim, height, width)
            for feature in stream_features
        ]
        values = [
            self.v_proj[scale_idx](feature).view(batch_size, self.num_heads, self.head_dim, height, width)
            for feature in stream_features
        ]
        if pad_h or pad_w:
            query = F.pad(query, (0, pad_w, 0, pad_h))
            keys = [F.pad(key, (0, pad_w, 0, pad_h)) for key in keys]
            values = [F.pad(value, (0, pad_w, 0, pad_h)) for value in values]

        key = torch.stack(keys, dim=1)
        value = torch.stack(values, dim=1)
        query_windows = self._window_partition_q(query, window_size)
        key_windows = self._window_partition_kv(key, window_size)
        value_windows = self._window_partition_kv(value, window_size)

        attention = torch.matmul(query_windows, key_windows.transpose(-1, -2)) * self.scale
        attention = F.softmax(attention, dim=-1)
        context_windows = torch.matmul(attention, value_windows)
        context = self._window_unpartition(context_windows, batch_size, padded_h, padded_w, window_size)
        return context[:, :, :, :height, :width].reshape(batch_size, self.num_heads * self.head_dim, height, width)

    def forward(
        self,
        stream_features: list[torch.Tensor],
        unadapted: torch.Tensor,
        scale_idx: int,
        return_delta: bool = False,
        evidence: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if not stream_features:
            delta = torch.zeros_like(unadapted)
            adapted = unadapted
            if return_delta:
                return adapted, delta
            return adapted
        if scale_idx in self.global_spatial_scales:
            context = self._global_spatial_context(stream_features, unadapted, scale_idx, evidence)
            return self._apply_residual(context, unadapted, scale_idx, return_delta, evidence)
        if scale_idx in self.windowed_spatial_scales:
            context = self._windowed_spatial_context(stream_features, unadapted, scale_idx, evidence)
            return self._apply_residual(context, unadapted, scale_idx, return_delta, evidence)
        return super().forward(stream_features, unadapted, scale_idx, return_delta, evidence)




class Transformer2DBlock(nn.Module):
    def __init__(
        self,
        hidden_dim: int,
        num_heads: int = 4,
        dropout_rate: float = 0.1,
        window_size: Optional[int] = None,
        mlp_ratio: int = 2,
    ) -> None:
        super().__init__()
        self.hidden_dim = hidden_dim
        self.window_size = window_size
        self.norm1 = nn.LayerNorm(hidden_dim)
        self.attn = nn.MultiheadAttention(
            hidden_dim,
            num_heads=num_heads,
            dropout=dropout_rate,
            batch_first=True,
        )
        self.norm2 = nn.GroupNorm(_group_count(hidden_dim), hidden_dim)
        expanded = hidden_dim * mlp_ratio
        self.mlp = nn.Sequential(
            nn.Conv2d(hidden_dim, expanded, kernel_size=1),
            nn.GELU(),
            nn.Conv2d(expanded, expanded, kernel_size=3, padding=1, groups=expanded),
            nn.GELU(),
            nn.Dropout2d(p=dropout_rate),
            nn.Conv2d(expanded, hidden_dim, kernel_size=1),
            nn.Dropout2d(p=dropout_rate),
        )

    def _global_attention(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, channels, height, width = x.shape
        sequence = x.flatten(2).transpose(1, 2)
        sequence_norm = self.norm1(sequence)
        attention_out, _ = self.attn(sequence_norm, sequence_norm, sequence_norm, need_weights=False)
        return attention_out.transpose(1, 2).reshape(batch_size, channels, height, width)

    def _window_partition(self, x: torch.Tensor, window_size: int) -> torch.Tensor:
        batch_size, channels, height, width = x.shape
        return (
            x.reshape(batch_size, channels, height // window_size, window_size, width // window_size, window_size)
            .permute(0, 2, 4, 3, 5, 1)
            .reshape(batch_size * (height // window_size) * (width // window_size), window_size * window_size, channels)
        )

    def _window_unpartition(
        self,
        windows: torch.Tensor,
        batch_size: int,
        height: int,
        width: int,
        window_size: int,
    ) -> torch.Tensor:
        channels = windows.shape[-1]
        return (
            windows.reshape(
                batch_size,
                height // window_size,
                width // window_size,
                window_size,
                window_size,
                channels,
            )
            .permute(0, 5, 1, 3, 2, 4)
            .reshape(batch_size, channels, height, width)
        )

    def _window_attention(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, _, height, width = x.shape
        window_size = self.window_size
        pad_h = (window_size - height % window_size) % window_size
        pad_w = (window_size - width % window_size) % window_size
        x_padded = F.pad(x, (0, pad_w, 0, pad_h)) if pad_h or pad_w else x
        _, _, padded_h, padded_w = x_padded.shape
        windows = self._window_partition(x_padded, window_size)
        windows_norm = self.norm1(windows)
        attention_windows, _ = self.attn(windows_norm, windows_norm, windows_norm, need_weights=False)
        out = self._window_unpartition(attention_windows, batch_size, padded_h, padded_w, window_size)
        return out[:, :, :height, :width]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.window_size is None:
            x = x + self._global_attention(x)
        else:
            x = x + self._window_attention(x)
        x = x + self.mlp(self.norm2(x))
        return x


class TransformerEvidenceMaskAdapter(nn.Module):
    """V2 mask adapter with evidence projection and efficient 2D transformer sharing."""

    def __init__(
        self,
        hidden_dim: int = 128,
        out_channels: int = 1,
        output_resolution: tuple[int, int] = (128, 128),
        in_channels_list: list[int] | None = None,
        dropout_rate: float = 0.1,
        num_heads: int = 4,
        window_size: int = 8,
    ) -> None:
        super().__init__()
        channels = in_channels_list or [256, 32, 64]
        self.output_resolution = output_resolution
        self.hidden_dim = hidden_dim
        self.mask_evidence_proj = nn.ModuleList(
            [
                FlexibleEvidenceProjector(
                    in_channels=hidden_dim + (2 * in_channels),
                    hidden_dim=hidden_dim,
                    dropout_rate=dropout_rate,
                    depth=2,
                )
                for in_channels in channels
            ]
        )
        self.low_blocks = nn.Sequential(
            *[
                Transformer2DBlock(hidden_dim, num_heads=num_heads, dropout_rate=dropout_rate, window_size=None)
                for _ in range(2)
            ]
        )
        self.mid_fuse = FlexibleEvidenceProjector(
            in_channels=hidden_dim * 2,
            hidden_dim=hidden_dim,
            dropout_rate=dropout_rate,
            depth=1,
        )
        self.mid_blocks = nn.Sequential(
            *[
                Transformer2DBlock(hidden_dim, num_heads=num_heads, dropout_rate=dropout_rate, window_size=window_size)
                for _ in range(2)
            ]
        )
        self.high_fuse = FlexibleEvidenceProjector(
            in_channels=hidden_dim * 2,
            hidden_dim=hidden_dim,
            dropout_rate=dropout_rate,
            depth=1,
        )
        self.high_blocks = nn.Sequential(
            *[
                Transformer2DBlock(hidden_dim, num_heads=num_heads, dropout_rate=dropout_rate, window_size=window_size)
                for _ in range(2)
            ]
        )
        self.mask_head = nn.Sequential(
            nn.GroupNorm(_group_count(hidden_dim), hidden_dim),
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1),
            nn.GELU(),
            nn.Dropout2d(p=dropout_rate),
            nn.Conv2d(hidden_dim, out_channels, kernel_size=1),
        )
        self.aux_head = nn.Sequential(
            nn.GroupNorm(_group_count(hidden_dim), hidden_dim),
            nn.Conv2d(hidden_dim, out_channels, kernel_size=1),
        )

    def _scale_evidence(
        self,
        adapted: FeaturePyramid,
        unadapted: FeaturePyramid,
        evidence: FeaturePyramid | None,
    ) -> FeaturePyramid:
        projected = []
        for scale_idx, projection in enumerate(self.mask_evidence_proj):
            if evidence is None:
                evidence_i = torch.zeros(
                    adapted[scale_idx].shape[0],
                    self.hidden_dim,
                    adapted[scale_idx].shape[-2],
                    adapted[scale_idx].shape[-1],
                    dtype=adapted[scale_idx].dtype,
                    device=adapted[scale_idx].device,
                )
            else:
                evidence_i = evidence[scale_idx]
                if evidence_i.shape[-2:] != adapted[scale_idx].shape[-2:]:
                    evidence_i = F.interpolate(
                        evidence_i,
                        size=adapted[scale_idx].shape[-2:],
                        mode="bilinear",
                        align_corners=False,
                    )
            adapter_delta = adapted[scale_idx] - unadapted[scale_idx]
            projected.append(
                projection(torch.cat([evidence_i, adapted[scale_idx], adapter_delta.abs()], dim=1))
            )
        return projected

    def forward(
        self,
        adapted: FeaturePyramid,
        streams_unadapted: StreamPyramid,
        unadapted: FeaturePyramid,
        evidence: FeaturePyramid | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
        del streams_unadapted
        high_target = self.output_resolution
        scale_features = self._scale_evidence(adapted, unadapted, evidence)

        low = self.low_blocks(scale_features[0])
        mid = scale_features[2] + F.interpolate(
            low,
            size=scale_features[2].shape[-2:],
            mode="bilinear",
            align_corners=False,
        )
        mid = self.mid_fuse(torch.cat([mid, scale_features[2]], dim=1))
        mid = self.mid_blocks(mid)

        high = scale_features[1] + F.interpolate(
            mid,
            size=scale_features[1].shape[-2:],
            mode="bilinear",
            align_corners=False,
        )
        high = self.high_fuse(torch.cat([high, scale_features[1]], dim=1))
        high = self.high_blocks(high)

        if high.shape[-2:] != high_target:
            high = F.interpolate(high, size=high_target, mode="bilinear", align_corners=False)
        mask_logits = self.mask_head(high)
        aux_logits = self.aux_head(mid)
        return mask_logits, aux_logits, None