###############################################################################
# 6. FULL COGNITHREAD MODEL + GGUF EXPORT
###############################################################################
class CogniThreadModel(nn.Module):
    """
    Full CogniThread SSM architecture combining:
    1. CogniThreadSSM (Mamba-style SSM with ZOH + parallel scan)
    2. FuzzyHeuristicRetriever (Gaussian fuzzy membership)
    3. LatentDiffusionSynthesizer (conditioned reverse diffusion)
    4. AdaptiveGate (variance-constrained gating)

    Forward:
      h_t = SSM(x)           # SSM trajectory
      retrieved, w_f, mu = Fuzzy(h_t)  # fuzzy retrieval
      z_0 = Diffusion(h_t)   # diffusion generation
      y_final, var = Gate(retrieved, z_0, w_f)  # adaptive gating
    """
    def __init__(self, d_model: int, d_state: int = 16, n_heuristics: int = 64,
                 d_z: int = 32, T_diff: int = 10, d_conv: int = 4,
                 expand: int = 2, gamma: float = 1.0, top_k: int = 5, *,
                 rngs: nn.Rngs):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.d_z = d_z
        self.ssm = CogniThreadSSM(d_model, d_state, d_conv, expand, rngs=rngs)
        # State projection from SSM output dim to d_state for fuzzy/diffusion
        self.state_proj = nn.Linear(d_model, d_state, rngs=rngs)
        self.fuzzy = FuzzyHeuristicRetriever(d_state, n_heuristics, gamma, top_k, rngs=rngs)
        self.diffusion = LatentDiffusionSynthesizer(d_z, d_state, T_diff, rngs=rngs)
        # Project diffusion output to d_model
        self.diff_proj = nn.Linear(d_z, d_model, rngs=rngs)
        # Project fuzzy retrieved to d_model
        self.fuzzy_proj = nn.Linear(d_state, d_model, rngs=rngs)
        self.gate = AdaptiveGate(d_model, rngs=rngs)

    def __call__(self, x, key=None):
        if key is None:
            key = jax.random.key(0)
        # SSM forward
        h_ssm = self.ssm(x)  # (batch, seq_len, d_model)
        # Project to state space for fuzzy/diffusion
        h_t = self.state_proj(h_ssm)  # (batch, seq_len, d_state)
        # Fuzzy heuristic retrieval
        retrieved, w_f, mu = self.fuzzy(h_t)
        retrieved = self.fuzzy_proj(retrieved)  # -> d_model
        # Diffusion generation
        z_0 = self.diffusion(h_t, key=key)
        generated = self.diff_proj(z_0)  # -> d_model
        # Adaptive gating
        y_final, var_y = self.gate(retrieved, generated, w_f)
        return y_final, var_y, w_f, mu

print("[6/6] CogniThreadModel defined")