###############################################################################
# 5. ADAPTIVE GATE — Heuristic-Generative Gating with Variance Constraint
###############################################################################
class AdaptiveGate(nn.Module):
    """
    y_final = w_f * Retrieve(S_best) + (1 - w_f) * Generate(z_0)
    Var(y_final) = (1 - w_f)^2 * sigma_z0^2
    As w_f -> 1, Var -> 0: hallucination prevention.
    """
    def __init__(self, d_model: int, *, rngs: nn.Rngs):
        super().__init__()
        self.d_model = d_model
        self.sigma_raw = nn.Param(jnp.array(0.0))

    @property
    def sigma(self):
        return jax.nn.softplus(self.sigma_raw.value)

    def __call__(self, retrieved, generated, w_f):
        if w_f.ndim < retrieved.ndim:
            w_f = w_f[..., None]
        y_final = w_f * retrieved + (1.0 - w_f) * generated
        var_y_final = (1.0 - w_f)**2 * self.sigma**2
        return y_final, var_y_final

print("[5/6] AdaptiveGate defined")