"""Bidirectional SSM block with explicit pad and segment reset semantics."""

from __future__ import annotations

import torch
from torch import nn

from .attention import FeedForward, RMSNorm
from .padding import masked_hidden
from .ssm import MaskedScanSSM


class BidirectionalSSMLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.reset_on_pad = config.ssm_reset_on_pad
        self.reset_on_segment = config.ssm_reset_on_segment
        self.fuse = config.ssm_fuse
        self.norm = RMSNorm(config.hidden_size, config.norm_eps)
        self.fwd = MaskedScanSSM(config.hidden_size, config.ssm_state_size)
        self.bwd = MaskedScanSSM(config.hidden_size, config.ssm_state_size)
        if self.fuse == "concat":
            self.out = nn.Linear(config.hidden_size * 2, config.hidden_size)
        elif self.fuse == "concat_product":
            self.out = nn.Linear(config.hidden_size * 3, config.hidden_size)
        elif self.fuse == "gated_sum":
            self.gate = nn.Linear(config.hidden_size, config.hidden_size)
            self.out = nn.Linear(config.hidden_size, config.hidden_size)
        else:
            raise ValueError(f"unknown ssm_fuse: {self.fuse}")
        self.ffn_norm = RMSNorm(config.hidden_size, config.norm_eps)
        self.ffn = FeedForward(
            config.hidden_size,
            config.intermediate_size,
            config.hidden_dropout,
            config.hidden_activation,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        segment_ids: torch.Tensor | None = None,
    ) -> torch.Tensor:
        residual = hidden_states
        x = self.norm(hidden_states)
        y_fwd = self.fwd(x, attention_mask, segment_ids, self.reset_on_pad, self.reset_on_segment)
        y_bwd_rev = self.bwd(
            torch.flip(x, dims=[1]),
            torch.flip(attention_mask, dims=[1]),
            torch.flip(segment_ids, dims=[1]) if segment_ids is not None else None,
            self.reset_on_pad,
            self.reset_on_segment,
        )
        y_bwd = torch.flip(y_bwd_rev, dims=[1])
        if self.fuse == "concat":
            y = self.out(torch.cat([y_fwd, y_bwd], dim=-1))
        elif self.fuse == "concat_product":
            y = self.out(torch.cat([y_fwd, y_bwd, y_fwd * y_bwd], dim=-1))
        else:
            gate = torch.sigmoid(self.gate(x))
            y = self.out(gate * y_fwd + (1.0 - gate) * y_bwd)
        hidden_states = masked_hidden(residual + y, attention_mask)
        hidden_states = hidden_states + self.ffn(self.ffn_norm(hidden_states))
        return masked_hidden(hidden_states, attention_mask)