import torch
from torch import nn


class FractionalMemory(nn.Module):
    """Causal positive exponential-mixture memory, accepting [T, B, D]."""
    def __init__(self, d_model: int, context_length: int, n_states: int = 8,
                 min_tau: float = 1.0, residual: bool = True):
        super().__init__()
        if n_states < 1 or context_length < 1:
            raise ValueError("n_states and context_length must be positive")
        taus = torch.logspace(torch.log10(torch.tensor(float(min_tau))),
                              torch.log10(torch.tensor(float(context_length))),
                              n_states)
        self.register_buffer("rho", torch.exp(-1.0 / taus))
        self.logits = nn.Parameter(torch.zeros(n_states))
        self.in_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        self.residual = residual

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.ndim != 3:
            raise ValueError("expected [T, B, D]")
        _, b, d = x.shape
        u = self.in_proj(x)
        state = torch.zeros(self.rho.numel(), b, d, device=x.device, dtype=x.dtype)
        w = torch.softmax(self.logits, dim=0).to(dtype=x.dtype)
        rho = self.rho.to(device=x.device, dtype=x.dtype).view(-1, 1, 1)
        ys = []
        for xt in u:
            state = rho * state + (1.0 - rho) * xt.unsqueeze(0)
            ys.append((w.view(-1, 1, 1) * state).sum(dim=0))
        y = self.out_proj(torch.stack(ys, dim=0))
        return x + y if self.residual else y