"""Pure-PyTorch Hydra (Hwang, Lahoti, Dao, Gu 2024): a quasiseparable bidirectional SSM mixer.

Mirrors goombalab/hydra's HydraMixer parameter layout exactly (same names and shapes), so weights are
interchangeable with the CUDA implementation. The SSD scan is the chunked algorithm from the Mamba-2
paper (Listing 1), padded to a chunk multiple. No mamba_ssm, no einops, no CUDA kernels.
"""
from __future__ import annotations

import math

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


def segsum(x: torch.Tensor) -> torch.Tensor:
    """out[..., i, j] = sum_{k=j+1..i} x[..., k] for i >= j, else -inf."""
    T = x.shape[-1]
    x = x.unsqueeze(-1).expand(*x.shape, T)
    lower = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device), diagonal=-1)
    x = x.masked_fill(~lower, 0)
    x_segsum = torch.cumsum(x, dim=-2)
    keep = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device), diagonal=0)
    return x_segsum.masked_fill(~keep, float("-inf"))


def ssd_scan(X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, chunk: int = 64) -> torch.Tensor:
    b, l, h, p = X.shape
    pad = (-l) % chunk
    if pad:
        X = F.pad(X, (0, 0, 0, 0, 0, pad))
        A = F.pad(A, (0, 0, 0, pad))
        B = F.pad(B, (0, 0, 0, 0, 0, pad))
        C = F.pad(C, (0, 0, 0, 0, 0, pad))
    L = X.shape[1]
    c = L // chunk
    X = X.reshape(b, c, chunk, h, p)
    B = B.reshape(b, c, chunk, h, -1)
    C = C.reshape(b, c, chunk, h, -1)
    A = A.reshape(b, c, chunk, h).permute(0, 3, 1, 2)  # (b, h, c, l)
    A_cumsum = torch.cumsum(A, dim=-1)
    Lmat = torch.exp(segsum(A))  # (b, h, c, l, s)
    Y_diag = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", C, B, Lmat, X)
    decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)
    states = torch.einsum("bclhn,bhcl,bclhp->bchpn", B, decay_states, X)
    states = torch.cat([torch.zeros_like(states[:, :1]), states], dim=1)
    decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
    new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
    states = new_states[:, :-1]
    Y_off = torch.einsum("bclhn,bchpn,bhcl->bclhp", C, states, torch.exp(A_cumsum))
    Y = (Y_diag + Y_off).reshape(b, L, h, p)
    return Y[:, :l]


def ssd_naive(X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor) -> torch.Tensor:
    b, l, h, p = X.shape
    n = B.shape[-1]
    state = X.new_zeros(b, h, p, n)
    ys = []
    for t in range(l):
        state = state * torch.exp(A[:, t])[:, :, None, None] + X[:, t][..., None] * B[:, t][:, :, None, :]
        ys.append(torch.einsum("bhpn,bhn->bhp", state, C[:, t]))
    return torch.stack(ys, dim=1)


class RMSNormGated(nn.Module):
    """RMSNorm followed by an optional SiLU gate (norm_before_gate=True in mamba_ssm)."""

    def __init__(self, d: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))

    def forward(self, x: torch.Tensor, z: torch.Tensor | None = None) -> torch.Tensor:
        xf = x.float()
        y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight.float()
        if z is not None:
            y = y * F.silu(z.float())
        return y.to(x.dtype)


class HydraMixer(nn.Module):
    def __init__(self, d_model: int, d_state: int = 64, d_conv: int = 7, expand: int = 2, headdim: int = 64,
                 ngroups: int = 1, chunk_size: int = 64, dt_min: float = 0.001, dt_max: float = 0.1,
                 dt_init_floor: float = 1e-4):
        super().__init__()
        self.d_model, self.d_state, self.headdim, self.ngroups, self.chunk_size = d_model, d_state, headdim, ngroups, chunk_size
        self.d_inner = expand * d_model
        if self.d_inner % headdim != 0:
            raise ValueError(f"d_inner (expand * d_model = {self.d_inner}) must be divisible by headdim ({headdim})")
        self.nheads = self.d_inner // headdim
        if self.nheads % ngroups != 0:
            raise ValueError(f"nheads (d_inner / headdim = {self.nheads}) must be divisible by ngroups ({ngroups})")
        if d_conv % 2 != 1:
            raise ValueError(f"d_conv must be odd so the non-causal conv keeps the sequence length, got {d_conv}")
        d_in_proj = 2 * self.d_inner + 2 * (2 * ngroups * d_state) + 2 * self.nheads
        self.in_proj = nn.Linear(d_model, d_in_proj, bias=False)
        conv_dim = self.d_inner + 2 * (2 * ngroups * d_state)
        self.conv1d = nn.Conv1d(conv_dim, conv_dim, kernel_size=d_conv, groups=conv_dim, padding=d_conv // 2, bias=True)
        dt = torch.exp(torch.rand(self.nheads) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min))
        dt = torch.clamp(dt, min=dt_init_floor)
        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
        self.dt_bias._no_weight_decay = True
        self.A_log = nn.Parameter(torch.zeros(self.nheads))  # A = -exp(0) = -1, as in the official init
        self.A_log._no_weight_decay = True
        self.D = nn.Parameter(torch.ones(self.nheads))
        self.D._no_weight_decay = True
        self.fc_D = nn.Linear(self.d_inner, self.nheads, bias=False)
        self.norm = RMSNormGated(self.d_inner, eps=1e-5)
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)

    def forward(self, u: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
        b, l, _ = u.shape
        H, P, G, N = self.nheads, self.headdim, self.ngroups, self.d_state
        if mask is not None:
            u = u * mask[..., None].to(u.dtype)
        z, xBC, dt = torch.split(self.in_proj(u), [self.d_inner, self.d_inner + 4 * G * N, 2 * H], dim=-1)
        xBC = F.silu(self.conv1d(xBC.transpose(1, 2)).transpose(1, 2))
        dt2 = F.softplus(torch.cat([dt[..., :H], dt[..., H:].flip(1)], dim=0) + self.dt_bias)  # (2b, l, H)
        if mask is not None:
            m = mask[..., None].to(u.dtype)
            xBC = xBC * m
            dt2 = dt2 * torch.cat([m, m.flip(1)], dim=0)
        x, BC = torch.split(xBC, [self.d_inner, 4 * G * N], dim=-1)
        x_og = x
        x2 = torch.cat([x, x.flip(1)], dim=0)
        BC2 = torch.cat([BC[..., :2 * G * N], BC[..., 2 * G * N:].flip(1)], dim=0)
        Bm, Cm = torch.split(BC2, [G * N, G * N], dim=-1)
        A = -torch.exp(self.A_log.float())
        Xh = x2.reshape(2 * b, l, H, P).float() * dt2.float()[..., None]
        Ah = A[None, None, :] * dt2.float()
        Bh = Bm.reshape(2 * b, l, G, N).float().repeat_interleave(H // G, dim=2)
        Ch = Cm.reshape(2 * b, l, G, N).float().repeat_interleave(H // G, dim=2)
        y = ssd_scan(Xh, Ah, Bh, Ch, self.chunk_size).reshape(2 * b, l, self.d_inner).to(u.dtype)
        y = torch.cat([torch.zeros_like(y[:, :1]), y[:, :-1]], dim=1)  # shift by one, first position zero
        y_fw, y_bw = y[:b], y[b:].flip(1)
        d_term = F.linear(x_og, self.fc_D.weight, self.D).repeat_interleave(P, dim=-1)
        y = y_fw + y_bw + x_og * d_term
        y = self.norm(y, z)
        return self.out_proj(y)


class HydraBlock(nn.Module):
    def __init__(self, d_model: int, mlp_mult: int = 2, **mixer_kwargs):
        super().__init__()
        self.norm1 = RMSNormGated(d_model)
        self.mixer = HydraMixer(d_model, **mixer_kwargs)
        self.norm2 = RMSNormGated(d_model)
        hidden = mlp_mult * d_model
        self.fc1 = nn.Linear(d_model, 2 * hidden)
        self.fc2 = nn.Linear(hidden, d_model)

    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
        x = x + self.mixer(self.norm1(x), mask)
        h, g = self.fc1(self.norm2(x)).chunk(2, dim=-1)
        return x + self.fc2(F.silu(g) * h)