import torch
from torch import nn


def pad_tensor_by_size(input_tensor: torch.Tensor, pad_size: int):
    """
    Padding x tensor with `pad_size` on the seq_len dim (dim=1)

    Assumes that we only have tensors of either size 4 or 3
    """
    pad_shape = (0, 0, 0, 0, 0, pad_size, 0, 0) if len(input_tensor.shape) == 4 else (0, 0, 0, pad_size, 0, 0)

    return torch.nn.functional.pad(input_tensor, pad_shape, mode="constant", value=0)

def reshape_into_chunks(input_tensor, pad_size, chunk_size):
    """
    Padding input_tensor with `pad_size` on the seq_len dim (dim=1) and
    simultaneously splitting it into chunk sequences.

    Assumes that we only have tensors of either size 4 or 3
    """
    # [bsz, seq_len, ...] -> [bsz, seq_len multiple of chunk_size, ...]
    input_tensor = pad_tensor_by_size(input_tensor, pad_size)

    if len(input_tensor.shape) == 3:
        # [bsz, seq_len multiple of chunk_size, num_heads] -> [bsz, -1, chunk_size, num_heads]
        return input_tensor.reshape(input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2])
    else:
        # [bsz, seq_len multiple of chunk_size, num_heads, head_dim or state_size] -> [bsz, -1, chunk_size, num_heads, head_dim or state_size]
        return input_tensor.reshape(
            input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2], input_tensor.shape[3]
        )

def segment_sum(input_tensor):
    """
    More stable segment sum calculation. Uses cumulative sums and masking instead of direct subtractions.
    """
    chunk_size = input_tensor.size(-1)
    # 1. expand input tensor to have an additional dimension and repeat along that dimension
    # [..., chunk_size] -> [..., chunk_size, chunk_size]
    input_tensor = input_tensor[..., None].expand(*input_tensor.size(), chunk_size)
    # 2. create a lower triangular mask with the diagonal set to 0 to 0 out elements above diag
    mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=-1)
    input_tensor = input_tensor.masked_fill(~mask, 0)
    # 3. compute actual cumsum
    tensor_segsum = torch.cumsum(input_tensor, dim=-2)

    # 4. apply mask to keep only the lower triangular part of the cumulative sum result (incl diagonal this time)
    mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=0)
    tensor_segsum = tensor_segsum.masked_fill(~mask, -torch.inf)
    return tensor_segsum


class FakeRemoteMamba2Mixer(nn.Module):
    """Scan part of NVIDIA Nemotron-3-Nano-Omni NemotronHMamba2Mixer.torch_forward (no cache), verbatim."""

    def __init__(self, num_heads=4, head_dim=3, n_groups=2, ssm_state_size=5, chunk_size=8, time_step_min=0.001):
        super().__init__()
        self.num_heads, self.head_dim, self.n_groups = num_heads, head_dim, n_groups
        self.ssm_state_size, self.chunk_size = ssm_state_size, chunk_size
        self.time_step_min, self.time_step_limit = time_step_min, (0.0, float("inf"))
        self.dt_bias = nn.Parameter(torch.randn(num_heads) - 4.0)
        self.A_log = nn.Parameter(torch.rand(num_heads) * 2)
        self.D = nn.Parameter(torch.randn(num_heads))

    def torch_forward(self, hidden_states, B, C, dt, cache_params=None):
        batch_size, seq_len = hidden_states.shape[:2]
        A = -torch.exp(self.A_log.float())
        # begin ssd naive implementation without einsums
        dt = nn.functional.softplus(dt + self.dt_bias)
        dt = torch.clamp(dt, self.time_step_min)
        hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float()
        B = B.reshape(batch_size, seq_len,  -1, self.ssm_state_size).float()
        C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()
        B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
        C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
        pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size

        D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size)

        # Discretize x and A
        hidden_states = hidden_states * dt[..., None]
        A = A.to(hidden_states.dtype) * dt

        # Rearrange into blocks/chunks
        hidden_states, A, B, C = [reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C)]


        # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size]
        A = A.permute(0, 3, 1, 2)
        A_cumsum = torch.cumsum(A, dim=-1)

        # 1. Compute the output for each intra-chunk (diagonal blocks)
        # This is the analog of a causal mask
        L = torch.exp(segment_sum(A))

        # First, contraction of C and B to get G (attention-weights like)
        G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, : ,:]  # shape: (b, c, l, s, h, n)
        G = G_intermediate.sum(dim=-1)  # shape: (b, c, l, s, h)


        # Step 2: Compute M, equivalent to applying attention mask to weights
        M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]
        M = M_intermediate.sum(dim=-1)

        # Step 3: Compute Y_diag (apply to values)
        Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(3)

        # (right term of low-rank factorization of off-diagonal blocks; B terms)

        decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)
        B_decay_contraction = B * decay_states.permute(0, 2, 3, 1)[..., None]
        # permute back B * decay states
        states = (B_decay_contraction.permute(0, 1, 3, 2, 4)[..., None]  * hidden_states.permute(0, 1, 3, 2, 4)[..., None, :]).sum(dim=3).permute(0, 1, 2, 4, 3)
        if cache_params is not None and cache_params.has_previous_state:
            previous_states = cache_params.ssm_states[self.layer_idx][:, None, ...]
        else:
            previous_states = torch.zeros_like(states[:, :1])
        states = torch.cat([previous_states, states], dim=1)
        decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0))))

        states_permuted = states.permute(0, 2, 1, 3, 4)
        result = (decay_chunk[..., None, None] * states_permuted[:, :, None, ...]).sum(dim=2)
        new_states = result.permute(0, 2, 1, 3, 4)
        states, ssm_state = new_states[:, :-1], new_states[:, -1]

        # Compute state -> output conversion per chunk
        # (left term of low-rank factorization of off-diagonal blocks; C terms)
        state_decay_out = torch.exp(A_cumsum)
        # compute Yoff
        C_times_states = (C[..., None, :] * states[:, :, None, ...])
        state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1)
        Y_off = (C_times_states.sum(-1) * state_decay_out_permuted[..., None])
        # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)

        y = Y_diag + Y_off
        # [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim]
        y = y.reshape(batch_size, -1, self.num_heads, self.head_dim)

        y = y + D_residual
        # Cutting off padded chunks
        if pad_size > 0:
            y = y[:, :seq_len, :, :]
        y = y.reshape(batch_size, seq_len, -1)
        if ssm_state is not None and cache_params is not None:
            cache_params.ssm_states[self.layer_idx].copy_(ssm_state)

        return y
