import math

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


class TemporalAttention(nn.Module):
    def __init__(
        self,
        n_embd: int,
        n_head: int,
        dropout: float = 0.0,
    ):
        super().__init__()
        assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
        self.n_head = n_head
        self.d_head = n_embd // n_head
        self.scale = 1.0 / math.sqrt(self.d_head)

        # QKV projection (fused for efficiency)
        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
        # Output projection
        self.out_proj = nn.Linear(n_embd, n_embd, bias=False)

        self.resid_drop = nn.Dropout(dropout)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        """Match the previous version's GPT-style weight initialization."""
        nn.init.normal_(self.qkv.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02)

    def forward(
        self,
        x: torch.Tensor,
        attn_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, L, _ = x.shape
        H, D = self.n_head, self.d_head

        # --- QKV ----------------------------------------------------------
        qkv = self.qkv(x).reshape(B, L, 3, H, D).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)  # each (B, H, L, D)

        out = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=attn_mask,
            dropout_p=0.0,
            is_causal=False,
            scale=self.scale,
        )

        # --- Aggregate & project out --------------------------------------
        out = out.transpose(1, 2).reshape(B, L, H * D)
        return self.resid_drop(self.out_proj(out))


class SwiGLU(nn.Module):
    def __init__(
        self,
        n_embd: int,
        hidden_dim: int | None = None,
        dropout: float = 0.0,
        bias: bool = True,
    ):
        super().__init__()
        hidden_dim = hidden_dim if hidden_dim is not None else int(
            n_embd * 2.5)

        self.w1 = nn.Linear(n_embd, hidden_dim, bias=bias)  # gate path
        self.w2 = nn.Linear(n_embd, hidden_dim, bias=bias)  # value path
        # output projection
        self.w3 = nn.Linear(hidden_dim, n_embd, bias=bias)
        self.drop = nn.Dropout(dropout)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        """GPT-style parameter initialization for MLP paths."""
        nn.init.normal_(self.w1.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.w2.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.w3.weight, mean=0.0, std=0.02)
        if self.w1.bias is not None:
            nn.init.zeros_(self.w1.bias)
            nn.init.zeros_(self.w2.bias)
            nn.init.zeros_(self.w3.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """``(B, L, n_embd) -> (B, L, n_embd)``."""
        return self.drop(self.w3(F.silu(self.w1(x)) * self.w2(x)))


class GPTBlock(nn.Module):
    def __init__(
        self,
        n_embd: int,
        n_head: int,

        attn_dropout: float = 0.0,
        mlp_dropout: float = 0.0,
    ):
        super().__init__()
        self.attn = TemporalAttention(
            n_embd=n_embd,
            n_head=n_head,
            dropout=attn_dropout,
        )
        self.mlp = SwiGLU(n_embd=n_embd, dropout=mlp_dropout)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(
        self,
        x: torch.Tensor,
        attn_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        x = x + self.attn(self.ln1(x), attn_mask)
        x = x + self.mlp(self.ln2(x))
        return x


class AgeSinusoidalEncoding(nn.Module):

    def __init__(self, embedding_dim: int):

        super().__init__()
        if embedding_dim % 2 != 0:
            raise ValueError(
                f"Embedding dimension must be an even number, but got {embedding_dim}")

        self.embedding_dim = embedding_dim

        i = torch.arange(0, self.embedding_dim, 2, dtype=torch.float32)
        divisor = torch.pow(10000, i / self.embedding_dim)
        self.register_buffer('divisor', divisor)
        self.linear = nn.Linear(embedding_dim, embedding_dim, bias=False)

    def forward(self, t: torch.Tensor) -> torch.Tensor:

        t_years = t
        # Broadcast (B, L, 1) against (1, 1, D/2) to get (B, L, D/2)
        args = t_years.unsqueeze(-1) / self.divisor.view(1, 1, -1)
        # Interleave cos and sin along the last dimension
        output = torch.zeros(t.shape[0], t.shape[1],
                             self.embedding_dim, device=t.device)
        output[:, :, 0::2] = torch.cos(args)
        output[:, :, 1::2] = torch.sin(args)
        output = self.linear(output)
        return output


class LiteTimesBackbone2d(nn.Module):
    """Cheap local feature extractor for a TimesNet period image."""

    def __init__(self, dim: int, kernel_size: int = 5,
                 expansion: float = 2.0, dropout: float = 0.0):
        super().__init__()
        if kernel_size <= 0 or kernel_size % 2 == 0:
            raise ValueError("kernel_size must be a positive odd integer")
        if expansion <= 0:
            raise ValueError("expansion must be > 0")
        hidden_dim = max(dim, int(round(dim * expansion)))
        self.dwconv = nn.Conv2d(
            dim, dim, kernel_size=kernel_size,
            padding=kernel_size // 2, groups=dim,
        )
        self.norm = nn.GroupNorm(1, dim)
        self.pwconv1 = nn.Conv2d(dim, hidden_dim, kernel_size=1)
        self.act = nn.GELU()
        self.pwconv2 = nn.Conv2d(hidden_dim, dim, kernel_size=1)
        self.drop = nn.Dropout(dropout)
        self.layer_scale = nn.Parameter(torch.full((1, dim, 1, 1), 1e-2))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.dwconv(x)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        return residual + self.layer_scale * self.drop(x)


class TimesNetBlock(nn.Module):
    """TimesNet block with lightweight depthwise-separable 2D extraction.

    The block follows TimesNet's idea: discover dominant periods with FFT,
    reshape a 1D sequence into period-wise 2D maps, run a 2D convolutional
    extractor, then fuse the top-k period branches.

    Reference:
        Wu et al., "TimesNet: Temporal 2D-Variation Modeling for General Time
        Series Analysis", ICLR 2023. https://arxiv.org/abs/2210.02186
    """

    def __init__(
        self,
        d_model: int,
        top_k: int = 2,
        n_backbone_blocks: int = 1,
        backbone_kernel_size: int = 5,
        backbone_expansion: float = 2.0,
        dropout: float = 0.0,
    ):
        super().__init__()
        if top_k <= 0:
            raise ValueError(f"top_k must be > 0, got {top_k}")
        if n_backbone_blocks <= 0:
            raise ValueError("n_backbone_blocks must be > 0")
        self.top_k = top_k
        self.norm = nn.LayerNorm(d_model)
        self.extractor = nn.Sequential(*[
            LiteTimesBackbone2d(
                dim=d_model,
                kernel_size=backbone_kernel_size,
                expansion=backbone_expansion,
                dropout=dropout,
            )
            for _ in range(n_backbone_blocks)
        ])

    def _select_periods(self, x: torch.Tensor) -> tuple[list[int], torch.Tensor]:
        B, T, C = x.shape
        spectrum = torch.fft.rfft(x.float(), dim=1)
        amplitude = spectrum.abs().mean(dim=(0, 2))
        if amplitude.numel() <= 1:
            return [max(T, 1)], x.new_ones(1)

        amplitude = amplitude.clone()
        amplitude[0] = 0.0
        k = min(self.top_k, amplitude.numel() - 1)
        weights, indices = torch.topk(amplitude, k=k)
        # One device synchronization per block instead of one per selected period.
        periods = [max(1, T // int(idx)) for idx in indices.tolist()]
        return periods, weights.to(dtype=x.dtype, device=x.device)

    def _period_branch(self, x: torch.Tensor, period: int) -> torch.Tensor:
        B, T, C = x.shape
        if T % period != 0:
            padded_len = ((T // period) + 1) * period
            pad = x.new_zeros(B, padded_len - T, C)
            x_pad = torch.cat([x, pad], dim=1)
        else:
            padded_len = T
            x_pad = x

        n_periods = padded_len // period
        x_2d = x_pad.reshape(B, n_periods, period, C).permute(0, 3, 1, 2)
        y = self.extractor(x_2d)
        y = y.permute(0, 2, 3, 1).reshape(B, padded_len, C)
        return y[:, :T, :]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x_norm = self.norm(x)
        periods, weights = self._select_periods(x_norm)
        branches = torch.stack(
            [self._period_branch(x_norm, period) for period in periods],
            dim=-1,
        )
        weights = torch.softmax(weights, dim=0).view(1, 1, 1, -1)
        return residual + (branches * weights).sum(dim=-1)


class TimesNetEncoder(nn.Module):
    """Encode a multivariate time series into one fixed-size embedding."""

    def __init__(
        self,
        input_dim: int,
        d_model: int,
        n_layers: int = 2,
        top_k: int = 2,
        n_backbone_blocks: int = 1,
        backbone_kernel_size: int = 5,
        backbone_expansion: float = 2.0,
        dropout: float = 0.0,
        append_observed_mask: bool = True,
    ):
        super().__init__()
        self.input_dim = input_dim
        self.append_observed_mask = append_observed_mask
        in_dim = input_dim * 2 if append_observed_mask else input_dim
        self.input_proj = nn.Linear(in_dim, d_model)
        self.blocks = nn.ModuleList([
            TimesNetBlock(
                d_model=d_model,
                top_k=top_k,
                n_backbone_blocks=n_backbone_blocks,
                backbone_kernel_size=backbone_kernel_size,
                backbone_expansion=backbone_expansion,
                dropout=dropout,
            )
            for _ in range(n_layers)
        ])
        self.final_ln = nn.LayerNorm(d_model)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        nn.init.normal_(self.input_proj.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.input_proj.bias)

    def forward(
        self,
        x: torch.Tensor,
        observed_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if x.dim() != 3:
            raise ValueError(f"x must have shape (B, T, C), got {tuple(x.shape)}")
        if x.size(-1) != self.input_dim:
            raise ValueError(
                f"last dim must be input_dim={self.input_dim}, got {x.size(-1)}"
            )

        finite_mask = torch.isfinite(x)
        x = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
        if observed_mask is None:
            observed_mask = finite_mask
        elif observed_mask.shape == x.shape[:2]:
            observed_mask = observed_mask.unsqueeze(-1).expand_as(x)
        elif observed_mask.shape != x.shape:
            raise ValueError(
                "observed_mask must have shape (B, T) or (B, T, C), got "
                f"{tuple(observed_mask.shape)}"
            )

        observed_mask = observed_mask.to(device=x.device, dtype=x.dtype)
        x = x * observed_mask
        if self.append_observed_mask:
            x = torch.cat([x, observed_mask], dim=-1)

        h = self.input_proj(x)
        for block in self.blocks:
            h = block(h)
        h = self.final_ln(h)

        valid_time = observed_mask.amax(dim=-1)
        pooled = (h * valid_time.unsqueeze(-1)).sum(dim=1)
        denom = valid_time.sum(dim=1, keepdim=True).clamp_min(1.0)
        return pooled / denom


class TimesNetExposureEncoder(nn.Module):
    """Encode pre-onset environmental exposure into an event-level embedding.

    Expected inputs:
        daily:   (B, 1826, 4) for tmean, tmax, tmin, rhmean
        monthly: (B, 241, 2) for tmean, rhmean

    Output:
        (B, n_embd), suitable for adding to a disease event embedding in the
        route-2 single-stream event-enhancement setup.
    """

    def __init__(
        self,
        n_embd: int,
        daily_input_dim: int = 4,
        monthly_input_dim: int = 2,
        d_model: int | None = None,
        n_layers: int = 2,
        top_k: int = 2,
        n_backbone_blocks: int = 1,
        backbone_kernel_size: int = 5,
        backbone_expansion: float = 2.0,
        dropout: float = 0.0,
        use_gate: bool = True,
    ):
        super().__init__()
        d_model = n_embd if d_model is None else d_model
        self.daily_encoder = TimesNetEncoder(
            input_dim=daily_input_dim,
            d_model=d_model,
            n_layers=n_layers,
            top_k=top_k,
            n_backbone_blocks=n_backbone_blocks,
            backbone_kernel_size=backbone_kernel_size,
            backbone_expansion=backbone_expansion,
            dropout=dropout,
            append_observed_mask=True,
        )
        self.monthly_encoder = TimesNetEncoder(
            input_dim=monthly_input_dim,
            d_model=d_model,
            n_layers=n_layers,
            top_k=top_k,
            n_backbone_blocks=n_backbone_blocks,
            backbone_kernel_size=backbone_kernel_size,
            backbone_expansion=backbone_expansion,
            dropout=dropout,
            append_observed_mask=True,
        )
        self.out_proj = nn.Sequential(
            nn.LayerNorm(2 * d_model),
            nn.Linear(2 * d_model, n_embd),
            nn.GELU(),
            nn.Linear(n_embd, n_embd),
        )
        self.gate = nn.Parameter(torch.tensor(-2.0)) if use_gate else None
        self.reset_parameters()

    def reset_parameters(self) -> None:
        for module in self.out_proj:
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
                nn.init.zeros_(module.bias)

    def forward(
        self,
        daily: torch.Tensor,
        monthly: torch.Tensor,
        daily_mask: torch.Tensor | None = None,
        monthly_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        h_daily = self.daily_encoder(daily, observed_mask=daily_mask)
        h_monthly = self.monthly_encoder(monthly, observed_mask=monthly_mask)
        h = self.out_proj(torch.cat([h_daily, h_monthly], dim=-1))
        if self.gate is not None:
            h = torch.sigmoid(self.gate) * h
        return h


class TimesNetSequenceDecoder(nn.Module):
    """Decode a fixed-size latent vector into a multivariate time series."""

    def __init__(
        self,
        output_dim: int,
        latent_dim: int,
        d_model: int,
        n_layers: int = 2,
        top_k: int = 2,
        n_backbone_blocks: int = 1,
        backbone_kernel_size: int = 5,
        backbone_expansion: float = 2.0,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.latent_proj = nn.Linear(latent_dim, d_model)
        self.position_proj = nn.Linear(3, d_model)
        self.blocks = nn.ModuleList([
            TimesNetBlock(
                d_model=d_model,
                top_k=top_k,
                n_backbone_blocks=n_backbone_blocks,
                backbone_kernel_size=backbone_kernel_size,
                backbone_expansion=backbone_expansion,
                dropout=dropout,
            )
            for _ in range(n_layers)
        ])
        self.final_ln = nn.LayerNorm(d_model)
        self.output_proj = nn.Linear(d_model, output_dim)

    def forward(self, latent: torch.Tensor, length: int) -> torch.Tensor:
        if latent.dim() != 2:
            raise ValueError(
                f"latent must have shape (B, D), got {tuple(latent.shape)}"
            )
        position = torch.linspace(
            0.0, 1.0, length, device=latent.device, dtype=latent.dtype
        )
        position = torch.stack(
            [position, torch.sin(2 * torch.pi * position),
             torch.cos(2 * torch.pi * position)],
            dim=-1,
        )
        h = self.latent_proj(latent).unsqueeze(1)
        h = h + self.position_proj(position).unsqueeze(0)
        for block in self.blocks:
            h = block(h)
        return self.output_proj(self.final_ln(h))


class TimesNetExposureAutoencoder(nn.Module):
    """Dual-resolution exposure autoencoder with a reusable event encoder."""

    def __init__(
        self,
        n_embd: int = 120,
        daily_input_dim: int = 4,
        monthly_input_dim: int = 2,
        d_model: int | None = None,
        n_layers: int = 2,
        top_k: int = 2,
        n_backbone_blocks: int = 1,
        backbone_kernel_size: int = 5,
        backbone_expansion: float = 2.0,
        dropout: float = 0.0,
    ):
        super().__init__()
        d_model = n_embd if d_model is None else d_model
        encoder_kwargs = dict(
            n_embd=n_embd,
            daily_input_dim=daily_input_dim,
            monthly_input_dim=monthly_input_dim,
            d_model=d_model,
            n_layers=n_layers,
            top_k=top_k,
            n_backbone_blocks=n_backbone_blocks,
            backbone_kernel_size=backbone_kernel_size,
            backbone_expansion=backbone_expansion,
            dropout=dropout,
            use_gate=True,
        )
        decoder_kwargs = dict(
            latent_dim=n_embd,
            d_model=d_model,
            n_layers=n_layers,
            top_k=top_k,
            n_backbone_blocks=n_backbone_blocks,
            backbone_kernel_size=backbone_kernel_size,
            backbone_expansion=backbone_expansion,
            dropout=dropout,
        )
        self.encoder = TimesNetExposureEncoder(**encoder_kwargs)
        self.daily_decoder = TimesNetSequenceDecoder(
            output_dim=daily_input_dim, **decoder_kwargs
        )
        self.monthly_decoder = TimesNetSequenceDecoder(
            output_dim=monthly_input_dim, **decoder_kwargs
        )

    def forward(
        self,
        daily: torch.Tensor,
        monthly: torch.Tensor,
        daily_mask: torch.Tensor | None = None,
        monthly_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        latent = self.encoder(daily, monthly, daily_mask, monthly_mask)
        daily_reconstruction = self.daily_decoder(latent, daily.size(1))
        monthly_reconstruction = self.monthly_decoder(latent, monthly.size(1))
        return daily_reconstruction, monthly_reconstruction, latent