src/dynamic_blueprint_evolver.py

"""
Hetman-2-NCM 2.0B: Dynamic Blueprint Evolver & Residual Intent Tracking Engine
Module: src/dynamic_blueprint_evolver.py
Specification: spec/15_DYNAMIC_BLUEPRINT_EVOLUTION_MECHANICS.md
RFC: RFC-HETMAN-2-NCM-DYNAMIC-BLUEPRINT-EVOLUTION

Architectural Role:
    Phase 3 Multi-Scale Block Decoder / Latent State Dynamics
    Implements Dynamic Blueprint Evolution via Residual Intent Tracking (RIT):
        B_k^{(t+1)} = Normalize_Spherically( (1 - alpha_done) * (B_k^{(t)} - gamma_t * delta_h) + alpha_done * B_completed )
    with Gram-Schmidt Subspace Protection against Over-Pruning,
    Cosine Annealing Schedule for gamma_t in [0.10, 0.80],
    and Zero-Vector Collapse Safeguards with transition to attractor [COMPLETED].

Target Hardware: Consumer GPUs (RTX 2060 6GB, RTX 3060 12GB), Apple Silicon, AVX-512 / AVX2 CPU.
"""

import math
from typing import Dict, List, Optional, Tuple, Union

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


# ==============================================================================
# 1. Сферична нормалізація зі спектральним демпфуванням Якобіана (Section 4.3)
# ==============================================================================

class SphericalNormFunction(torch.autograd.Function):
    """
    Автоград-оператор проекції вектора на одиничну ріманову гіперсферу S^(d-1)
    зі спектральним демпфуванням Якобіана (Spectral Jacobian Clamping).

    Формулювання:
        Прямий прохід:
            y = v / max(||v||_2, eps)
        Зворотний прохід (Theorem 4.3):
            J_norm(v) = (1 / ||v||_2) * (I - y y^T)
            Для запобігання вибуху градієнтів (Exploding Gradients) множник
            1 / ||v||_2 обмежується зверху константою clamp_max (за замовчуванням 1.0):
            J_tilde = min(clamp_max, 1.0 / ||v||_2) * (I - y y^T)
            Це гарантує нерозтягуваність оператора: ||J_tilde||_2 <= 1.0.
    """

    @staticmethod
    def forward(ctx, v: torch.Tensor, eps: float = 1e-8, clamp_max: float = 1.0) -> torch.Tensor:
        norm = torch.norm(v, p=2, dim=-1, keepdim=True)
        norm_safe = torch.clamp(norm, min=eps)
        y = v / norm_safe
        ctx.save_for_backward(y, norm_safe)
        ctx.clamp_max = clamp_max
        return y

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None, None]:
        y, norm_safe = ctx.saved_tensors
        clamp_max = ctx.clamp_max
        
        # Спектральне демпфування множника
        scale = torch.clamp(1.0 / norm_safe, max=clamp_max)
        
        # Проектор (I - y y^T) dL/dy = dL/dy - y <y, dL/dy>
        inner_prod = torch.sum(grad_output * y, dim=-1, keepdim=True)
        grad_v = scale * (grad_output - inner_prod * y)
        return grad_v, None, None


class SphericalNorm(nn.Module):
    """
    Модуль сферичної нормалізації на гіперсферу S^(d-1).
    Забезпечує строгу унітарну норму ||x||_2 = 1.0 та стійкий градієнтний потік.
    """

    def __init__(self, eps: float = 1e-8, clamp_max: float = 1.0, use_spectral_clamping: bool = True):
        super().__init__()
        self.eps = eps
        self.clamp_max = clamp_max
        self.use_spectral_clamping = use_spectral_clamping

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if torch.isnan(x).any() or torch.isinf(x).any():
            x = torch.nan_to_num(x, nan=0.0, posinf=1.0, neginf=-1.0)
        norm = torch.norm(x, p=2, dim=-1, keepdim=True)
        zero_mask = (norm < self.eps)
        if zero_mask.any():
            fallback = torch.zeros_like(x)
            fallback[..., 0] = 1.0
            x = torch.where(zero_mask, fallback, x)
            norm = torch.norm(x, p=2, dim=-1, keepdim=True)
        if self.use_spectral_clamping and x.requires_grad:
            return SphericalNormFunction.apply(x, self.eps, self.clamp_max)
        norm_safe = torch.clamp(norm, min=self.eps)
        return x / norm_safe


# ==============================================================================
# 2. Фільтр ортогоналізації Грама-Шмідта (Subspace Protection, Section 4.2)
# ==============================================================================

class GramSchmidtSubspaceFilter(nn.Module):
    """
    Фільтр ортогоналізації Грама-Шмідта для захисту майбутніх тез від Over-Pruning.

    Математичний сенс (Теорема 1 зі специфікації 15):
        Коли екстрактор реалізованого змісту детектує згенеровані токени,
        він може мати паразитну проекцію на невирішені майбутні тези u_j in M_future.
        Фільтр проектує вектор віднімання виключно на підпростір активних тез M_active,
        зберігаючи скалярний добуток <B_unnorm, u_j> == <B_curr, u_j> для всіх j in M_future.

    Підтримує два режими:
        1. Слотова проекція (якщо передано матрицю слотів U або B_slots [B, M, d_model]).
        2. Траєкторна колінеарно-ортогональна фільтрація відносно початкового креслення B_initial.
    """

    def __init__(
        self,
        d_model: int = 1536,
        num_slots: int = 16,
        active_threshold: float = 0.15,
        orthogonal_leak: float = 0.25
    ):
        super().__init__()
        self.d_model = d_model
        self.num_slots = num_slots
        self.active_threshold = active_threshold
        self.orthogonal_leak = orthogonal_leak
        self.slot_dim = max(1, d_model // num_slots)

    def project_slots(
        self,
        delta_h: torch.Tensor,
        slots: torch.Tensor,
        B_current: torch.Tensor,
        active_threshold: Optional[float] = None
    ) -> torch.Tensor:
        """
        Стійка слотова ортогоналізація Мура-Пенроуза (Theorem 1, RFC-HETMAN-2-NCM):
            P_active = U_act (U_act^T U_act)^+ U_act^T
            delta_h_safe = P_active * delta_h
        
        Захищена від поганої обумовленості матриці Грама (kappa(G) -> inf)
        при майже колінеарних або лінійно залежних векторах слотів.
        
        slots: [Batch, M, d_model] - вектори слотів креслення
        B_current: [Batch, d_model] - поточний вектор наміру
        delta_h: [Batch, d_model] - сирий вектор віднімання
        """
        thresh = active_threshold if active_threshold is not None else self.active_threshold
        
        # Обчислюємо проекцію поточного наміру на кожен слот: c_m = <B_curr, u_m>
        c_m = torch.sum(slots * B_current.unsqueeze(1), dim=-1) # [B, M]
        
        # Маска активних слотів: c_m >= thresh
        active_mask = (c_m >= thresh).to(delta_h.dtype) # [B, M]
        U_act = slots * active_mask.unsqueeze(-1) # [B, M, d]
        
        # Матриця Грама активних слотів G = U_act @ U_act^T [B, M, M]
        G = torch.bmm(U_act, U_act.transpose(1, 2))
        
        # Псевдообернення Мура-Пенроуза гарантує стійкість до лінійної залежності:
        # P = U_act^T @ (U_act @ U_act^T)^+ @ U_act
        G_pinv = torch.linalg.pinv(G) # [B, M, M]
        
        proj_coeff = torch.bmm(U_act, delta_h.unsqueeze(-1)) # [B, M, 1]
        alpha = torch.bmm(G_pinv, proj_coeff) # [B, M, 1]
        
        # Реконструкція безпечного вектору без ампліфікації норми
        delta_h_safe = torch.bmm(U_act.transpose(1, 2), alpha).squeeze(-1) # [B, d_model]
        
        # Для ортогонального залишку додаємо контрольовану частину orthogonal_leak
        uncovered_delta = delta_h - delta_h_safe
        filtered_delta = delta_h_safe + self.orthogonal_leak * uncovered_delta
        return filtered_delta

    def forward(
        self,
        delta_h: torch.Tensor,
        B_initial: torch.Tensor,
        B_current: torch.Tensor,
        slots: Optional[torch.Tensor] = None,
        active_threshold: Optional[float] = None
    ) -> torch.Tensor:
        """
        delta_h: [Batch, d_model] - сирий вектор змісту
        B_initial: [Batch, d_model] або [Batch, M, d_model] - початковий начерк
        B_current: [Batch, d_model] - поточний намір
        slots: опціональний тензор слотів [Batch, M, d_model]
        """
        if slots is not None:
            return self.project_slots(delta_h, slots, B_current, active_threshold)

        # Якщо B_initial передано як багатослотовий тензор [B, M, d]
        if B_initial.dim() == 3:
            return self.project_slots(delta_h, B_initial, B_current, active_threshold)

        # Режим декомпозиції підпростору B_initial / B_current
        # Обчислюємо збереженість початкового наміру в поточному
        proj = torch.sum(B_current * B_initial, dim=-1, keepdim=True) # [Batch, 1]
        
        # Колінеарна складова поточної траєкторії
        delta_parallel = proj * B_current
        # Ортогональна паразитна складова
        delta_orthogonal = delta_h - delta_parallel
        
        # Придушуємо ортогональний шум, захищаючи майбутні незвідані підпростори
        filtered_delta = delta_parallel + self.orthogonal_leak * delta_orthogonal
        return filtered_delta


# ==============================================================================
# 3. Екстрактор реалізованого семантичного змісту (Section 2.1)
# ==============================================================================

class SemanticContentExtractor(nn.Module):
    """
    Екстрактор реалізованого змісту з блоку 400 латентних токенів.
    Використовує Attention-Pooling відносно поточного вектора наміру B_k^(t)
    та проектор SwiGLU з подальшим обмеженням довжини.
    """

    def __init__(
        self,
        d_model: int = 1536,
        d_attn: Optional[int] = None,
        d_mid: Optional[int] = None,
        eps: float = 1e-6
    ):
        super().__init__()
        self.d_model = d_model
        self.d_attn = d_attn or max(64, d_model // 6)
        self.d_mid = d_mid or max(128, int(d_model * 4 / 3))
        
        self.layer_norm = nn.LayerNorm(d_model, eps=eps)
        
        # Проекції для механізму уваги (Cross-Attention Pooling)
        self.w_q = nn.Linear(d_model, self.d_attn, bias=False)
        self.w_k = nn.Linear(d_model, self.d_attn, bias=False)
        self.scale = 1.0 / math.sqrt(self.d_attn)
        
        # SwiGLU MLP проектор змісту у простір наміру
        self.w_gate = nn.Linear(d_model, self.d_mid, bias=False)
        self.w_up = nn.Linear(d_model, self.d_mid, bias=False)
        self.w_down = nn.Linear(self.d_mid, d_model, bias=False)

    def forward(self, H_gen: torch.Tensor, B_intent: torch.Tensor) -> torch.Tensor:
        """
        H_gen: [Batch, L_block, d_model] - згенеровані латентні стани (L_block = 400)
        B_intent: [Batch, d_model] - поточний вектор наміру B_k^(t)

        Повертає:
            delta_h_clamped: [Batch, d_model] - вектор реалізованого змісту, нормований за довжиною <= 1.0
        """
        H_norm = self.layer_norm(H_gen) # [B, L, d_model]
        
        # Оптимізація: Асоціативна факторизація уваги (-92.5% FLOPs)
        # Замість проекції L токенів K = H_norm @ W_k^T [B, L, d_attn] (314.57 MFLOPs),
        # обчислюємо еквівалентний вектор запиту в просторі моделі:
        # v_eff = (Q @ W_k) * scale [B, 1, d_model] (2.80 MFLOPs)
        Q = self.w_q(B_intent).unsqueeze(1) # [B, 1, d_attn]
        v_eff = torch.matmul(Q, self.w_k.weight) * self.scale # [B, 1, d_model]
        
        # Скалярні добутки уваги релевантності токенів до поточного наміру
        attn_scores = torch.bmm(v_eff, H_norm.transpose(1, 2)) # [B, 1, L]
        attn_weights = F.softmax(attn_scores, dim=-1) # [B, 1, L]
        
        # Зважена агрегація відбитку згенерованого тексту
        H_pooled = torch.bmm(attn_weights, H_norm).squeeze(1) # [B, d_model]
        
        # Нелінійна SwiGLU проекція
        gate = F.silu(self.w_gate(H_pooled))
        up = self.w_up(H_pooled)
        delta_h = self.w_down(gate * up) # [B, d_model]
        
        # Нормування довжини видобутого змісту: delta_h / max(||delta_h||_2, 1.0)
        delta_norm = torch.norm(delta_h, p=2, dim=-1, keepdim=True)
        delta_h_clamped = delta_h / torch.clamp(delta_norm, min=1.0)
        return delta_h_clamped


# ==============================================================================
# 4. Трьохшкальна субпросторова факторизація (Section 3.1 - 3.2)
# ==============================================================================

class MultiScaleSubspaceFactorization(nn.Module):
    """
    Субпросторова факторизація вектора наміру B_k у відповідності до 3 шкал NCM:
        B_k = [B_syn (25%), B_morph (25%), B_lex (50%)]^T = 1536.
    
    Масштабно-зважений фільтр Omega(t):
        t in [1, 2] (Scale 1: Синтаксис):  [0.80, 0.15, 0.05]
        t in [3, 4] (Scale 2: Морфологія): [0.10, 0.70, 0.20]
        t in [5, 8] (Scale 3: Словник):    [0.05, 0.05, 0.90]
    """

    def __init__(
        self,
        d_model: int = 1536,
        d_syn: Optional[int] = None,
        d_morph: Optional[int] = None,
        d_lex: Optional[int] = None
    ):
        super().__init__()
        if d_syn is None or d_morph is None or d_lex is None:
            # За замовчуванням 25% синтаксис (384/1536), 25% морфологія (384/1536), 50% словник (768/1536)
            d_syn = d_model // 4
            d_morph = d_model // 4
            d_lex = d_model - d_syn - d_morph
            
        assert d_syn + d_morph + d_lex == d_model, f"Subspaces sum {d_syn+d_morph+d_lex} != {d_model}"
        self.d_model = d_model
        self.d_syn = d_syn
        self.d_morph = d_morph
        self.d_lex = d_lex

    def get_scale_weights(self, step_idx: int) -> Tuple[float, float, float]:
        """
        Повертає вагові коефіцієнти (w_syn, w_morph, w_lex) для кроку step_idx (1..8).
        """
        if step_idx in (1, 2):
            return (0.80, 0.15, 0.05)
        elif step_idx in (3, 4):
            return (0.10, 0.70, 0.20)
        else: # 5..8
            return (0.05, 0.05, 0.90)

    def apply_filter(self, delta_h: torch.Tensor, step_idx: int) -> torch.Tensor:
        """
        Масштабує відповідні субпростори вектора delta_h згідно з поточною шкалою NCM.
        delta_h: [Batch, d_model]
        """
        w_syn, w_morph, w_lex = self.get_scale_weights(step_idx)
        
        syn_part = delta_h[..., :self.d_syn] * w_syn
        morph_part = delta_h[..., self.d_syn:self.d_syn + self.d_morph] * w_morph
        lex_part = delta_h[..., self.d_syn + self.d_morph:] * w_lex
        
        return torch.cat([syn_part, morph_part, lex_part], dim=-1)


# ==============================================================================
# 5. AdaLN-Zero Модуляція шарів декодера (Section 3.3)
# ==============================================================================

class AdaLNZeroModulator(nn.Module):
    """
    Генератор параметрів адаптивної нормалізації шарів (AdaLN-Zero) з нульовою ініціалізацією:
        [gamma_1, beta_1, alpha_1, gamma_2, beta_2, alpha_2] = W_adaln * B_k^(t)
    де B_k^(t) динамічно перебудовує селективність нейронів кожного шару.
    """

    def __init__(self, d_model: int = 1536, num_layers: int = 8):
        super().__init__()
        self.d_model = d_model
        self.num_layers = num_layers
        
        # Проекції 6 параметрів на кожен шар: 6 * d_model параметрів
        self.layer_projections = nn.ModuleList([
            nn.Linear(d_model, 6 * d_model, bias=True)
            for _ in range(num_layers)
        ])
        
        # Нульова ініціалізація (AdaLN-Zero) для забезпечення тотожного відображення на старті
        for proj in self.layer_projections:
            nn.init.zeros_(proj.weight)
            nn.init.zeros_(proj.bias)

    def forward(
        self,
        B_intent: torch.Tensor,
        layer_idx: Optional[int] = None
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        B_intent: [Batch, d_model]
        layer_idx: 0-indexed шар декодера [0..num_layers-1]. Якщо None, використовується шар 0.
        Повертає:
            (gamma_1, beta_1, alpha_1, gamma_2, beta_2, alpha_2), кожен форми [Batch, d_model]
        """
        target_idx = 0 if layer_idx is None else layer_idx
        params = self.layer_projections[target_idx](B_intent) # [B, 6 * d_model]
        chunks = torch.chunk(params, 6, dim=-1)
        return chunks[0], chunks[1], chunks[2], chunks[3], chunks[4], chunks[5]


# ==============================================================================
# 6. Головний модуль: DynamicBlueprintEvolver (Section 2, 4, 5)
# ==============================================================================

class DynamicBlueprintEvolver(nn.Module):
    """
    Виробничий модуль Dynamic Blueprint Evolver на базі Residual Intent Tracking (RIT).
    
    Основні функції:
        1. Віднімання виконаного наміру:
           B_cand = SpherNorm(B_k^(t) - gamma_t * delta_h)
        2. Захист майбутніх тез від over-pruning через ортогоналізацію Грама-Шмідта.
        3. Косинусний розклад gamma_t в діапазоні [0.10, 0.80] (з підтримкою exp/linear).
        4. Захист від Zero-Vector Collapse через автоматичний перехід до атрактора [COMPLETED]
           при падінні залишкової норми нижче collapse_threshold (0.25).
        5. Багатошкальна фільтрація підпросторів за фазами NCM (Syntax -> Morph -> Lexicon).
    """

    def __init__(
        self,
        d_model: int = 1536,
        total_steps: int = 8,
        gamma_min: float = 0.10,
        gamma_max: float = 0.80,
        collapse_threshold: float = 0.25,
        schedule_type: str = "cosine",
        use_subspace_factorization: bool = True,
        use_riemannian_step: bool = False,
        final_attractor_force: float = 0.85
    ):
        super().__init__()
        self.d_model = d_model
        self.total_steps = total_steps
        self.gamma_min = gamma_min
        self.gamma_max = gamma_max
        self.collapse_threshold = collapse_threshold
        self.schedule_type = schedule_type.lower()
        self.use_subspace_factorization = use_subspace_factorization
        self.use_riemannian_step = use_riemannian_step
        self.final_attractor_force = final_attractor_force
        
        # Модулі
        self.spher_norm = SphericalNorm(eps=1e-8, clamp_max=1.0)
        self.extractor = SemanticContentExtractor(d_model=d_model)
        self.filter = GramSchmidtSubspaceFilter(d_model=d_model)
        
        if self.use_subspace_factorization:
            self.scale_factorizer = MultiScaleSubspaceFactorization(d_model=d_model)
        else:
            self.scale_factorizer = None
            
        # Навчальний вектор атрактора завершеності [COMPLETED] на одиничній сфері S^(d-1)
        # Ініціалізується нормованим випадковим вектором
        init_attractor = F.normalize(torch.randn(1, d_model), p=2, dim=-1)
        self.completed_attractor = nn.Parameter(init_attractor)
        
        # Класифікатор готовності блоку (Termination Gate)
        self.done_gate = nn.Linear(d_model, 1)

    def get_gamma(self, step_idx: int) -> float:
        """
        Обчислює коефіцієнт віднімання gamma_t для кроку step_idx (1-indexed).
        Підтримує безпечний режим роботи при t > total_steps (наприклад, 11 проходів
        в екстреному режимі відновлення): насичує прогрес на 1.0 та гарантує gamma <= gamma_max < 1.0.
        """
        step = max(1, step_idx)
        if self.total_steps <= 1:
            return float(self.gamma_min)
            
        # Нормалізований насичуваний прогрес: clamp до [0.0, 1.0]
        progress = min(1.0, (step - 1) / float(self.total_steps - 1))
        
        if self.schedule_type == "cosine":
            # Косинусний розклад від gamma_min до gamma_max (Equation 258)
            gamma = self.gamma_min + 0.5 * (self.gamma_max - self.gamma_min) * (1.0 - math.cos(math.pi * progress))
        elif self.schedule_type == "exponential":
            # Експоненційний розклад (Equation 286)
            gamma = self.gamma_min * ((self.gamma_max / self.gamma_min) ** progress)
        elif self.schedule_type == "linear":
            gamma = self.gamma_min + (self.gamma_max - self.gamma_min) * progress
        else:
            raise ValueError(f"Unknown schedule_type '{self.schedule_type}'. Use 'cosine', 'exponential', or 'linear'.")
            
        # Гарантія суворого інваріанту: gamma_min <= gamma <= gamma_max
        return float(min(self.gamma_max, max(self.gamma_min, gamma)))

    def riemannian_geodesic_subtraction(
        self,
        B_curr: torch.Tensor,
        delta_h: torch.Tensor,
        gamma: float
    ) -> torch.Tensor:
        """
        Строге ріманове віднімання на дотичному просторі T_{B} S^(d-1) (Section 2.2):
            1. v_tangent = delta_h - <delta_h, B> B
            2. theta = gamma * max(0, <delta_h, B>)
            3. Exp_B(-theta v) = cos(theta) B - sin(theta) (v / ||v||)
        """
        dot = torch.sum(delta_h * B_curr, dim=-1, keepdim=True)
        v_tangent = delta_h - dot * B_curr
        v_norm = torch.norm(v_tangent, p=2, dim=-1, keepdim=True)
        v_safe = torch.clamp(v_norm, min=1e-8)
        v_dir = v_tangent / v_safe
        
        kappa = torch.clamp(dot, min=0.0)
        theta = gamma * kappa
        
        # Зсув вздовж геодезичної
        B_geo = torch.cos(theta) * B_curr - torch.sin(theta) * v_dir
        return B_geo

    def step(
        self,
        B_curr: torch.Tensor,
        B_init: torch.Tensor,
        H_gen: torch.Tensor,
        step_idx: int,
        slots: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, Dict[str, float]]:
        """
        Один крок еволюції вектора наміру: B_k^{(t)} -> B_k^{(t+1)}.
        
        Аргументи:
            B_curr:   [Batch, d_model] - поточний вектор наміру на кроці t
            B_init:   [Batch, d_model] або [Batch, M, d_model] - початковий начерк B_k^(0)
            H_gen:    [Batch, L_block, d_model] - згенеровані латентні токени
            step_idx: int - номер кроку (1-indexed: 1..total_steps)
            slots:    опціональний тензор ортогональних слотів [Batch, M, d_model]
            
        Повертає:
            B_next:  [Batch, d_model] - оновлений знормований вектор наміру на S^(d-1)
            metrics: dict - діагностичні показники (норми, залишок, альфа, тощо)
        """
        batch_size = B_curr.size(0)
        gamma = self.get_gamma(step_idx)
        
        # 1. Екстракція реалізованого семантичного змісту з H_gen через Cross-Attention Pooling
        delta_h_raw = self.extractor(H_gen, B_curr) # [Batch, d_model]
        
        # 2. Фільтрація Грама-Шмідта для захисту нерозкритих тез (Anti-Over-Pruning)
        delta_h_safe = self.filter(
            delta_h=delta_h_raw,
            B_initial=B_init,
            B_current=B_curr,
            slots=slots
        ) # [Batch, d_model]
        
        # 3. Багатошкальна субпросторова фільтрація (Scale 1 -> Scale 2 -> Scale 3)
        if self.scale_factorizer is not None:
            delta_h_directed = self.scale_factorizer.apply_filter(delta_h_safe, step_idx)
        else:
            delta_h_directed = delta_h_safe
            
        # 4. Векторне віднімання наміру
        if self.use_riemannian_step:
            B_unnorm = self.riemannian_geodesic_subtraction(B_curr, delta_h_directed, gamma)
        else:
            B_unnorm = B_curr - gamma * delta_h_directed
            
        # Обчислюємо евклідову норму залишку
        residual_norm = torch.norm(B_unnorm, p=2, dim=-1, keepdim=True) # [Batch, 1]
        
        # 5. Захист від Zero-Vector Collapse та детекція завершеності
        done_logits = self.done_gate(B_unnorm) # [Batch, 1]
        done_prob = torch.sigmoid(done_logits) # [Batch, 1]
        
        # Індикатор просідання залишкової норми нижче безпечного порогу (0.25)
        # Коли residual_norm -> 0, collapse_weight -> 1.0
        norm_collapse_weight = torch.clamp(
            (self.collapse_threshold - residual_norm) / self.collapse_threshold,
            min=0.0,
            max=1.0
        )
        
        # Підсумковий коефіцієнт захоплення атрактором
        alpha_done = torch.maximum(done_prob, norm_collapse_weight)
        
        # Якщо крок фінальний (або аварійний t >= total_steps), примусово спрямовуємо до атрактора
        if step_idx >= self.total_steps:
            target_force = torch.tensor(self.final_attractor_force, device=B_curr.device, dtype=B_curr.dtype)
            alpha_done = torch.maximum(alpha_done, target_force)
            
        # 6. Інтерполяція між залишковим наміром та атрактором [COMPLETED]
        attractor_normed = F.normalize(self.completed_attractor, p=2, dim=-1)
        attractor_expanded = attractor_normed.expand(batch_size, -1)
        
        B_blended = (1.0 - alpha_done) * B_unnorm + alpha_done * attractor_expanded
        
        # 7. Сферична нормалізація на S^(d-1) зі спектральним демпфуванням
        B_next = self.spher_norm(B_blended)
        
        # Діагностичні метрики
        B_init_vec = B_init.mean(dim=1) if B_init.dim() == 3 else B_init
        metrics = {
            "step": float(step_idx),
            "gamma": float(gamma),
            "residual_norm_mean": float(residual_norm.mean().item()),
            "alpha_done_mean": float(alpha_done.mean().item()),
            "cos_sim_to_init": float(F.cosine_similarity(B_next, B_init_vec, dim=-1).mean().item()),
            "cos_sim_to_completed": float(F.cosine_similarity(B_next, attractor_expanded, dim=-1).mean().item()),
            "final_norm": float(torch.norm(B_next, p=2, dim=-1).mean().item())
        }
        
        return B_next, metrics

    def simulate_8_passes(
        self,
        B_initial: torch.Tensor,
        H_passes: List[torch.Tensor],
        slots: Optional[torch.Tensor] = None
    ) -> Tuple[List[torch.Tensor], List[Dict[str, float]]]:
        """
        Повна симуляція 8 проходів декодера для заданого батчу.
        
        Аргументи:
            B_initial: [Batch, d_model] або [Batch, M, d_model] - початкове креслення
            H_passes:  List з 8 тензорів розмірності [Batch, L_block, d_model]
            slots:     опціональний тензор слотів [Batch, M, d_model]
            
        Повертає:
            trajectory:      List з 9 тензорів: [B_0, B_1, ..., B_8]
            metrics_history: List з 8 словників метрик для кожного кроку
        """
        assert len(H_passes) == self.total_steps, (
            f"Expected {self.total_steps} passes, got {len(H_passes)}"
        )
        
        trajectory: List[torch.Tensor] = []
        metrics_history: List[Dict[str, float]] = []
        
        # B_0: початковий намір, знормований на сферу
        if B_initial.dim() == 3:
            B_curr = self.spher_norm(B_initial.sum(dim=1))
        else:
            B_curr = self.spher_norm(B_initial)
            
        trajectory.append(B_curr)
        
        for t in range(1, self.total_steps + 1):
            H_gen = H_passes[t - 1]
            B_next, step_metrics = self.step(
                B_curr=B_curr,
                B_init=B_initial,
                H_gen=H_gen,
                step_idx=t,
                slots=slots
            )
            trajectory.append(B_next)
            metrics_history.append(step_metrics)
            B_curr = B_next
            
        return trajectory, metrics_history

    def simulate_passes(
        self,
        B_initial: torch.Tensor,
        H_passes: List[torch.Tensor],
        slots: Optional[torch.Tensor] = None
    ) -> Tuple[List[torch.Tensor], List[Dict[str, float]]]:
        """
        Аліас для simulate_8_passes для сумісності з різними конвенціями імен.
        """
        return self.simulate_8_passes(B_initial, H_passes, slots=slots)

    def forward(
        self,
        B_curr: torch.Tensor,
        B_init: torch.Tensor,
        H_gen: torch.Tensor,
        step_idx: int,
        slots: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, Dict[str, float]]:
        """
        Стандартний forward метод для PyTorch nn.Module.
        """
        return self.step(B_curr, B_init, H_gen, step_idx, slots=slots)