###############################################################################
# 3. FUZZY HEURISTIC RETRIEVER — Gaussian Membership Function
###############################################################################
class FuzzyHeuristicRetriever(nn.Module):
    """
    mu_C(h_t, S_i) = exp(-gamma * ||h_t - S_i||^2)
    Top-k retrieval: O(log|S| * d)
    """
    def __init__(self, d_state: int, n_heuristics: int, gamma: float = 1.0,
                 top_k: int = 5, *, rngs: nn.Rngs):
        super().__init__()
        self.d_state = d_state
        self.n_heuristics = n_heuristics
        self.top_k = top_k
        self.gamma = nn.Param(jnp.array(gamma))
        key = rngs.params()
        self.S = nn.Param(jax.random.normal(key, (n_heuristics, d_state)) * 0.1)
        self.value_proj = nn.Linear(d_state, d_state, rngs=rngs)

    def gaussian_membership(self, h_t, S_i):
        h_sq = jnp.sum(h_t**2, axis=-1, keepdims=True)
        s_sq = jnp.sum(S_i**2, axis=-1)
        h_dot_s = h_t @ S_i.T
        dist_sq = jnp.maximum(h_sq - 2*h_dot_s + s_sq, 0.0)
        gamma = jax.nn.softplus(self.gamma.value)
        return jnp.exp(-gamma * dist_sq)

    def __call__(self, h_t):
        if h_t.ndim == 3:
            batch, seq_len, d = h_t.shape
            h_t_flat = h_t.reshape(-1, d)
        else:
            h_t_flat = h_t
            batch = seq_len = None
        S_i = self.S.value
        mu = self.gaussian_membership(h_t_flat, S_i)
        top_k = min(self.top_k, self.n_heuristics)
        top_k_mu, top_k_idx = jax.lax.top_k(mu, top_k)
        top_k_S = S_i[top_k_idx]
        weights = top_k_mu / (jnp.sum(top_k_mu, axis=-1, keepdims=True) + 1e-8)
        retrieved = jnp.sum(weights[..., None] * top_k_S, axis=-2)
        retrieved = self.value_proj(retrieved)
        w_f = jnp.max(mu, axis=-1, keepdims=True)
        if batch is not None:
            retrieved = retrieved.reshape(batch, seq_len, self.d_state)
            w_f = w_f.reshape(batch, seq_len, 1)
            mu = mu.reshape(batch, seq_len, self.n_heuristics)
        return retrieved, w_f, mu

print("[3/6] FuzzyHeuristicRetriever defined")