"""
infrastructure/model/cardiogan.py
──────────────────────────────────
CardioGAN U-Net Generator model architecture.
Strictly defines the PyTorch nn.Module architecture (SRP). No signal preprocessing.
"""
from __future__ import annotations


def _build_attention_gate_module():
    import torch
    import torch.nn as nn
    import torch.nn.functional as F

    class AttentionGate(nn.Module):
        """Attention gate for 1-D signals on skip connections."""

        def __init__(self, F_l: int, F_g: int, F_int: int):
            super().__init__()
            self.W_x = nn.Conv1d(F_l, F_int, kernel_size=1, stride=1, bias=True)
            self.W_g = nn.Conv1d(F_g, F_int, kernel_size=1, stride=1, bias=True)
            self.psi = nn.Conv1d(F_int, 1, kernel_size=1, stride=1, bias=True)
            self.relu = nn.ReLU(inplace=True)
            self.sigmoid = nn.Sigmoid()

        def forward(self, x_l: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
            if g.shape[-1] != x_l.shape[-1]:
                g = F.interpolate(g, size=x_l.shape[-1], mode="nearest")
            theta_x = self.W_x(x_l)
            phi_g = self.W_g(g)
            f = self.relu(theta_x + phi_g)
            alpha = self.sigmoid(self.psi(f))
            return x_l * alpha

    return AttentionGate


def _build_encoder_block_module():
    import torch.nn as nn

    class EncoderBlock(nn.Module):
        """Encoder block: Conv1d -> [GroupNorm] -> LeakyReLU."""

        def __init__(self, in_ch: int, out_ch: int, kernel_size: int = 16,
                     stride: int = 2, use_norm: bool = True):
            super().__init__()
            pad = (kernel_size - 1) // 2
            self.conv = nn.Conv1d(in_ch, out_ch, kernel_size, stride, pad)
            self.norm = nn.GroupNorm(1, out_ch) if use_norm else nn.Identity()
            self.act = nn.LeakyReLU(0.2, inplace=True)

        def forward(self, x):
            return self.act(self.norm(self.conv(x)))

    return EncoderBlock


def _build_decoder_block_module():
    import torch.nn as nn

    class DecoderBlock(nn.Module):
        """Decoder block: ConvTranspose1d -> [GroupNorm] -> ReLU."""

        def __init__(self, in_ch: int, out_ch: int, kernel_size: int = 16,
                     stride: int = 2, use_norm: bool = True):
            super().__init__()
            pad = (kernel_size - 1) // 2
            out_pad = stride - 1 if stride > 1 else 0
            self.deconv = nn.ConvTranspose1d(
                in_ch, out_ch, kernel_size, stride, pad, output_padding=out_pad
            )
            self.norm = nn.GroupNorm(1, out_ch) if use_norm else nn.Identity()
            self.act = nn.ReLU(inplace=True)

        def forward(self, x):
            return self.act(self.norm(self.deconv(x)))

    return DecoderBlock


def build_attention_unet_generator():
    """
    Build and return a fresh AttentionUNetGenerator instance.
    """
    import torch
    import torch.nn as nn
    import torch.nn.functional as F

    AttentionGate = _build_attention_gate_module()
    EncoderBlock = _build_encoder_block_module()
    DecoderBlock = _build_decoder_block_module()

    class AttentionUNetGenerator(nn.Module):
        """Attention U-Net Generator for CardioGAN."""

        def __init__(self):
            super().__init__()
            enc_filters = [64, 128, 256, 512, 512, 512]

            self.enc1 = EncoderBlock(1, enc_filters[0], use_norm=False)
            self.enc2 = EncoderBlock(enc_filters[0], enc_filters[1])
            self.enc3 = EncoderBlock(enc_filters[1], enc_filters[2])
            self.enc4 = EncoderBlock(enc_filters[2], enc_filters[3])
            self.enc5 = EncoderBlock(enc_filters[3], enc_filters[4])
            self.enc6 = EncoderBlock(enc_filters[4], enc_filters[5])

            self.attn5 = AttentionGate(enc_filters[4], enc_filters[4], enc_filters[4] // 2)
            self.attn4 = AttentionGate(enc_filters[3], enc_filters[3], enc_filters[3] // 2)
            self.attn3 = AttentionGate(enc_filters[2], enc_filters[2], enc_filters[2] // 2)
            self.attn2 = AttentionGate(enc_filters[1], enc_filters[1], enc_filters[1] // 2)
            self.attn1 = AttentionGate(enc_filters[0], enc_filters[0], enc_filters[0] // 2)

            self.dec6 = DecoderBlock(enc_filters[5], enc_filters[4])
            self.dec5 = DecoderBlock(enc_filters[4] * 2, enc_filters[3])
            self.dec4 = DecoderBlock(enc_filters[3] * 2, enc_filters[2])
            self.dec3 = DecoderBlock(enc_filters[2] * 2, enc_filters[1])
            self.dec2 = DecoderBlock(enc_filters[1] * 2, enc_filters[0])

            self.final = nn.Sequential(
                nn.ConvTranspose1d(enc_filters[0] * 2, 1, kernel_size=16,
                                   stride=2, padding=7, output_padding=0),
                nn.Tanh()
            )

        def forward(self, x: torch.Tensor) -> torch.Tensor:
            e1 = self.enc1(x)    # (B, 64,  256)
            e2 = self.enc2(e1)   # (B, 128, 128)
            e3 = self.enc3(e2)   # (B, 256, 64)
            e4 = self.enc4(e3)   # (B, 512, 32)
            e5 = self.enc5(e4)   # (B, 512, 16)
            e6 = self.enc6(e5)   # (B, 512, 8)

            d6 = self.dec6(e6)
            a5 = self.attn5(e5, d6)
            d5 = self.dec5(torch.cat([self._match(d6, a5), a5], dim=1))

            a4 = self.attn4(e4, d5)
            d4 = self.dec4(torch.cat([self._match(d5, a4), a4], dim=1))

            a3 = self.attn3(e3, d4)
            d3 = self.dec3(torch.cat([self._match(d4, a3), a3], dim=1))

            a2 = self.attn2(e2, d3)
            d2 = self.dec2(torch.cat([self._match(d3, a2), a2], dim=1))

            a1 = self.attn1(e1, d2)
            out = self.final(torch.cat([self._match(d2, a1), a1], dim=1))

            return out

        @staticmethod
        def _match(decoder_feat: torch.Tensor, skip_feat: torch.Tensor) -> torch.Tensor:
            if decoder_feat.shape[-1] != skip_feat.shape[-1]:
                decoder_feat = F.interpolate(
                    decoder_feat, size=skip_feat.shape[-1], mode="nearest"
                )
            return decoder_feat

    return AttentionUNetGenerator()