# -*- coding: utf-8 -*-
"""Kernel MAMBA-3 (v11, doc 22 §2) — varredura seletiva com discretização
TRAPETZOIDAL, em PyTorch puro, vetorizado e DIFERENCIÁVEL por autograd.

Fonte: arquivo MAMBA_3.txt (kernel Cython `mamba3_cpu_kernel.pyx`) — NÃO
copiado: o kernel Cython (a) fixava `a_val = A[b, d, 0]` (ignorava a
dependência de A em n — gradiente de A[·,·,n>0] errado), (b) exigia
compilação por plataforma e backward manual propenso a erro, e (c)
materializava h_states (B,L,D,N) inteiro. Aqui (Teoremas 22.2–22.5):

  • A é por (d, n): α_{t,d,n} = exp(δ_{t,d}·a_{d,n}) — recuperação do SSM
    canônico; o gradiente de A flui EXATO via autograd (Corolário 22.2.1);
  • atualização trapezoidal: h_t = α_t·h_{t−1} + (δ_t/2)·(1+α_t)·B̄_t·u_t
    com B̄_{t,n} = Σ_r B_{t,r,n} (MIMO rank-reduzido, como a fonte);
  • ESTABILIDADE (Teo 22.3): δ = σ(·) ∈ (0,1) e a_{d,n} < 0 ⇒
    α_{t,d,n} ∈ (e^{a}, 1) ⊂ (0,1) — a recorrência é CONTRAÇÃO; estado
    limitado ‖h_t‖ ≤ max_s (δ_s/2)(1+α_s)|B̄_s·u_s| / (1−α_max);
  • v18 — CLAMPS DE SANIDADE (Teo 30.2, requisito §1.1.5/§2.4): a
    contração do Teo 22.3 exige δ∈(0,1) e a<0. Se uma anomalia
    (inicialização ruim, gradiente explosivo, uso direto da API) violar
    isso, α = exp(δ·a) ≥ 1 e a recorrência DIVERGE geometricamente.
    A varredura agora fixa delta ∈ [1e-6, 1−1e-6] e A ≤ −1e-6 ANTES de
    α=exp(δ·a). Como δ e A válidos NUNCA são alterados pelo clamp
    (identidade na região válida), o gradiente autograd é inalterado
    (Teo 30.2-cor: clamp idêntico na região válida ⟹ ∇ idêntico);
  • memória O(B·D·N) por passo (estado vivo) — sem tensor (B,L,D,N);
  • semente inicial h_0 = 0 (convenção de estado frio — idêntica à fonte).

Assinatura compatível com o autograd manual da fonte (u, delta, A, B, C) —
troca DROP-IN: o autograd do PyTorch substitui `mamba3_backward_cpu`
(prova de exatidão: teste `test_v11_mamba3_grad_vs_finitediff`).
"""
from __future__ import annotations

import torch
import torch.nn.functional as F

__all__ = ["varredura_trapezoidal", "Mamba3Funcao", "verificar_estabilidade",
           "conv1d_causal", "set_conv_cython"]

# v18 — clamps de sanidade (Teo 30.2). Ligados por padrão: custam O(1)
# por tensor e garantem a contração mesmo sob anomalia numérica.
CLAMPS_SEGURANCA = True
DELTA_MIN, DELTA_MAX = 1e-6, 1.0 - 1e-6
A_MAX = -1e-6                       # força a < 0 (contração)


def verificar_estabilidade(delta: torch.Tensor, A: torch.Tensor) -> dict:
    """Auditoria (sem mutação): quantas entradas de δ/A violariam a
    contração do Teo 22.3 ANTES do clamp. Telemetria para o TelemetryHub
    (violações > 0 ⟹ anomalia a montante: investigar inicialização/grads)."""
    with torch.no_grad():
        n_delta_fora = int(((delta <= 0.0) | (delta >= 1.0)).sum())
        n_A_nao_neg = int((A >= 0.0).sum())
    return {"mamba3/delta_fora": n_delta_fora, "mamba3/A_nao_negativa": n_A_nao_neg}


def varredura_trapezoidal(u: torch.Tensor, delta: torch.Tensor, A: torch.Tensor,
                          B_m: torch.Tensor, C_m: torch.Tensor,
                          devolve_estado: bool = False):
    """Varredura causal seletiva MAMBA-3.

    Formas:
      u:     (B, L, D)     entrada u_t
      delta: (B, L, D)     passo de tempo δ_t ∈ (0,1) (gate sigmoid)
      A:     (D, N) ou (B, D, N) — dinâmica de estado (a < 0 estável)
      B_m:   (B, L, R, N)  projeção de entrada MIMO (rank R)
      C_m:   (B, L, R, N)  projeção de leitura MIMO (rank R)
    Devolve out: (B, L, D) com out_t,d = Σ_n (Σ_r C_{t,r,n}) · h_{t,d,n}.

    v12 — `devolve_estado=True` devolve TAMBÉM h_final: (B, D, N), o estado
    h_L da recorrência (Teorema 24.13: h_L é EXATAMENTE o estado que o
    decode incremental recebe — mesma recorrência, então
    varredura(u_{1..L}) ⊕ passo(u_{L+1}) ≡ varredura(u_{1..L+1})).

    Laço temporal O(L) vetorizado sobre (B, D, N, R) — CPU/GPU sem kernel
    custom; exato no tempo (sem aproximação além da própria regra do
    trapézio, que é a DEFINIÇÃO do MAMBA-3 — Teorema 22.2).
    """
    Bsz, L, D = u.shape
    N = A.shape[-1]
    R = B_m.shape[2]
    if A.dim() == 2:                      # (D, N) → expande p/ lote
        A = A.unsqueeze(0).expand(Bsz, -1, -1)
    assert B_m.shape == (Bsz, L, R, N) and C_m.shape == (Bsz, L, R, N), \
        "BUG (doc 22 §2): formas MIMO incompatíveis"
    if CLAMPS_SEGURANCA:
        # Teo 30.2 (v18, requisito §1.1.5/§2.4): contração GARANTIDA por
        # construção — δ∈[1e-6,1−1e-6] e A≤−1e-6 ⟹ α=exp(δ·a) ∈
        # [exp((1−1e-6)·(−1e-6)) , exp(1e-6·(−1e-6))) ⊂ (0,1) SEMPRE.
        # Na região válida o clamp é a identidade ⟹ ∇ inalterado
        # (Teo 30.2-cor: d clamp(x)/dx = 1 dentro da faixa).
        delta = delta.clamp(DELTA_MIN, DELTA_MAX)
        A = A.clamp(max=A_MAX)
    # agregação MIMO: b̄_t,n = Σ_r B[t,r,n]; c̄_t,n = Σ_r C[t,r,n]  (O(B·L·D·N·R))
    b_bar = B_m.sum(dim=2)                # (B, L, N)
    c_bar = C_m.sum(dim=2)                # (B, L, N)
    # δ_{t,d}·a_{d,n} → α por (B, L, D, N)
    dtA = delta.unsqueeze(-1) * A.unsqueeze(1)          # (B, L, D, N)
    alpha = torch.exp(dtA)                              # Teo 22.3: ∈ (0,1)
    beta = 0.5 * delta                                  # (B, L, D) — δ/2
    # termo de entrada: β_t·(1+α_t)·b̄_{t,n}·u_t,d → (B, L, D, N)
    entrada = (beta.unsqueeze(-1) * (1.0 + alpha) * b_bar.unsqueeze(2)) \
        * u.unsqueeze(-1)                               # (B,L,D,N)
    h = torch.zeros(Bsz, D, N, dtype=u.dtype, device=u.device)
    saidas = []
    for t in range(L):                                  # O(L) — estado vivo
        h = alpha[:, t] * h + entrada[:, t]
        # leitura: out_t,d = Σ_n h_{t,d,n}·c̄_{t,n}
        saidas.append((h * c_bar[:, t].unsqueeze(1)).sum(-1))
    out = torch.stack(saidas, dim=1)                    # (B, L, D)
    if devolve_estado:
        return out, h                                   # h_final (B, D, N)
    return out


class Mamba3Funcao(torch.autograd.Function):
    """Ponte de compatibilidade com a API da fonte (`Mamba3Function.apply`).
    Mantida para os testes compararem com a interface (u, delta, A, B, C);
    o backward é o autograd PADRÃO (exato — Corolário 22.2.1)."""

    @staticmethod
    def forward(ctx, u, delta, A, B_m, C_m):
        return varredura_trapezoidal(u, delta, A, B_m, C_m)

    @staticmethod
    def backward(ctx, grad_out):                        # pragma: no cover
        raise RuntimeError("use tensor.backward() — autograd exato (doc 22)")


# ======================================================================
# v21 — CONV CAUSAL PROFUNDA FUNDIDA (doc 33 §4, Teos 33.11/33.12)
# ======================================================================
# Kernel Cython (acelerado/mamba3hot.pyx) com backward ANALÍTICO exato
# (regra da cadeia na conv causal — testado contra autograd: erro ≤ 1e-5
# relativo). Ganho medido: 2.3–7.6× vs o caminho torch (pad+Conv1d+
# 2 transposes+copy — ~5× o tráfego de memória do mínimo). FALLBACK
# automático ao caminho torch v20 (mesmo cálculo) quando: sem binário,
# dtype ≠ fp32 ou device ≠ cpu. Kill-switch: set_conv_cython(False).
try:                                    # binário commitado (praxe v9.1)
    import mamba3hot as _hot            # build top-level (acelerado/)
except ImportError:                     # pragma: no cover
    try:
        from khtst.acelerado import mamba3hot as _hot
    except Exception:
        _hot = None

_CONV_CYTHON = True                     # requisito v21: Cython prioritário


def set_conv_cython(ativa: bool) -> None:
    """Kill-switch do caminho Cython da conv (A/B e fallback de emergência)."""
    global _CONV_CYTHON
    _CONV_CYTHON = bool(ativa)


class _ConvCausalProfunda(torch.autograd.Function):
    """Ponte autograd: forward = kernel C fundido; backward = fórmulas
    analíticas EXATAS (Teo 33.11-ii — ∂u/∂w/∂b da conv causal depthwise)."""

    @staticmethod
    def forward(ctx, u, w, b):
        ctx.save_for_backward(u, w)
        from khtst.mamba3.kernel import _hot as hot
        out = hot.conv1d_causal_profunda(
            u.detach().contiguous().numpy(),
            w.detach().contiguous().numpy(),
            b.detach().contiguous().numpy())
        return torch.from_numpy(out)

    @staticmethod
    def backward(ctx, g):
        # v21-nota: backward ANALÍTICO de 1ª ordem (Teo 33.11-ii). `g` é
        # detachado — double-backward pelo kernel C não é suportado (os
        # testes de gradiente do projeto são de 1ª ordem; o caminho torch
        # de fallback mantém autograd de ordem arbitrávia).
        from khtst.mamba3.kernel import _hot as hot
        u, w = ctx.saved_tensors
        K = int(w.shape[1])
        gd = g.detach().contiguous().numpy()
        gu = torch.from_numpy(hot.conv1d_grad_u(
            gd, w.detach().contiguous().numpy()))
        gw = torch.from_numpy(hot.conv1d_grad_w2(
            gd, u.detach().contiguous().numpy(), K))
        gb = torch.from_numpy(hot.conv1d_grad_b(gd))
        return gu, gw, gb


def conv1d_causal(u: torch.Tensor, pesos_dk: torch.Tensor,
                  vies: torch.Tensor) -> torch.Tensor:
    """Conv depthwise causal (B,L,D) → (B,L,D) — FONTE ÚNICA do hot-path
    do Mamba-3 (R21.5).

    Caminho Cython fundido (fp32/CPU) quando disponível; caso contrário o
    caminho torch v20 EXATO (pad esquerdo K−1 + Conv1d groups=D + volta
    do layout — Teo 33.11: ambos computam
    uc[b,l,d] = Σ_j w[d,j]·u[b, l+j−K+1, d] + v[d], equivalência medida
    max|Δ| = 0)."""
    if _hot is not None and _CONV_CYTHON \
            and u.dtype == torch.float32 and u.is_cpu:
        return _ConvCausalProfunda.apply(u, pesos_dk, vies)
    K = pesos_dk.shape[1]
    return F.conv1d(F.pad(u.transpose(1, 2), (K - 1, 0)),
                    pesos_dk.unsqueeze(1), vies,
                    groups=pesos_dk.shape[0]).transpose(1, 2).contiguous()