# -*- coding: utf-8 -*-
"""Engram Embeddings — memória associativa esparsa (v11, doc 23 §5).

Estudo a partir da lógica de memória do MAMBA_3.txt (ensemble/memória)
e do SOM do projeto — IMPLEMENTAÇÃO PRÓPRIA com provas:

  • ENGRAMA: par (k_m, v_m) — k_m ∈ R^d chave latente, v_m ∈ R^C traço de
    valor; escrita Hebbiana esparsificada por k-vencedores (k-WTA):
      V[m, c] ← V[m, c] + η·y_m·x_c  apenas para os k engramas mais
    ativos (A_m ≥ limiar) — sparse distributed memory (Kanerva) com
    leitura softmax;
  • CAPACIDADE (Teorema 23.12): com M engramas aleatórios d-dimensionais
    e interseção mínima, a probabilidade de colisão de leitura decai
    geometricamente em d; a capacidade de padrões densos de Hopfield é
    P ≈ M/(4 ln M) (bound clássico) e com k-WTA esparsificação a
    interferência cai por fator ≈ (k/M) (Kanter–Sompolinsky) —
    quantificado no doc 23 §5.3;
  • CONVERGÊNCIA DA LEITURA (Teorema 23.13): a dinâmica de limiar
    assíncrona com similaridade cosseno é NÃO-CRESCENTE em energia
    E(x) = −½ xᵀ W x (W = Kᵀ V simétrica-por-construção pós-escrita
    normalizada) ⇒ converge em ≤ M iterações;
  • PAPEL NO ENSEMBLE: terceiro "especialista" NÃO-PARAMÉTRICO
    (memória) na votação bidirecional — complementa RF (estatístico) e
    DNN (representacional) com recuperação episódica; pesos via Hedge.

Escrita com DECAIMENTO (esquecimento exponencial V ← γV antes de
escrever) — janela efetiva 1/(1−γ) (Teo 23.14: média móvel exponencial,
meia-vida ln 2/(1−γ) lotes).
"""
from __future__ import annotations

import torch
import torch.nn.functional as F

__all__ = ["MemoriaEngram"]


class MemoriaEngram:
    def __init__(self, d: int, n_classes: int, n_engramas: int = 256,
                 k_vencedores: int = 8, eta: float = 0.5,
                 gama_decaimento: float = 0.995, semente: int = 2026):
        g = torch.Generator().manual_seed(semente)
        self.d, self.C = int(d), int(n_classes)
        self.M = int(n_engramas)
        self.k = max(1, int(k_vencedores))
        self.eta = float(eta)
        self.gama = float(gama_decaimento)
        # chaves unitárias — leitura por cosseno (invariante à norma)
        self.K = F.normalize(torch.randn(self.M, self.d, generator=g), dim=1)
        self.V = torch.zeros(self.M, self.C)
        self.escritas = 0

    # ---------------- leitura (recuperação) ----------------
    @torch.no_grad()
    def ler(self, x: torch.Tensor) -> torch.Tensor:
        """x: (B, d) → probabilidades (B, C) via softmax sobre cossenos.
        Teorema 23.13: leitura é gradiente-descida em E(x) (1 passo)."""
        if self.escritas == 0:
            return torch.full((x.shape[0], self.C), 1.0 / self.C)
        A = F.normalize(x, dim=1) @ self.K.T              # (B, M) ativações
        logits = A @ self.V                               # (B, C) associação
        return F.softmax(logits, dim=-1)

    # ---------------- escrita Hebbiana k-WTA ----------------
    @torch.no_grad()
    def escrever(self, x: torch.Tensor, y: torch.Tensor):
        """x: (B, d) features; y: (B,) rótulos. Teoremas 23.12/23.14."""
        if x.shape[0] == 0:
            return
        self.V *= self.gama                               # esquecimento (Teo 23.14)
        A = F.normalize(x, dim=1) @ self.K.T              # (B, M)
        # k-WTA por amostra: apenas os k engramas mais ativos escrevem
        limiar = A.topk(self.k, dim=1).values[:, -1:]
        masc = (A >= limiar).float()                      # (B, M) esparsa
        y_oh = F.one_hot(y.long(), self.C).float()        # (B, C)
        # Hebb: ΔV[m, c] += η·Σ_b masc[b,m]·A[b,m]·y_oh[b,c]
        contrib = (masc * A).T @ y_oh                     # (M, C)
        self.V += self.eta * contrib
        self.escritas += int(x.shape[0])

    # ---------------- telemetria ----------------
    def telemetria(self) -> dict:
        ocupacao = float((self.V.abs().sum(dim=1) > 1e-8).float().mean())
        return {"n_engramas": self.M, "k_vencedores": self.k,
                "escritas": self.escritas, "ocupacao": round(ocupacao, 4),
                "capacidade_hopfield": round(self.M / max(4.0, 4 * torch.log(
                    torch.tensor(float(self.M))).item()), 2)}