import numpy as np

class MultiAttractorSystem:
    """
    MUCE‑Ω v2 Core Engine
    """

    def __init__(self, dim=2, K=3):
        self.dim = dim
        self.K = K
        self.Xi = np.random.randn(dim)
        self.history = []
        self.w = np.ones(K) / K
        self.lr = 0.05
        self.beta = 0.2
        self.gamma = 0.1
        self.max_norm = 5.0

    def potentials(self, Xi):
        return np.array([
            np.linalg.norm(Xi),
            -np.linalg.norm(Xi - 2.0),
            np.linalg.norm(Xi + 1.0),
        ])

    def grad_potentials(self, Xi):
        eps = 1e-6
        return np.array([
            Xi / (np.linalg.norm(Xi) + eps),
            -(Xi - 2.0) / (np.linalg.norm(Xi - 2.0) + eps),
            (Xi + 1.0) / (np.linalg.norm(Xi + 1.0) + eps),
        ])

    def project_constraints(self, Xi):
        norm = np.linalg.norm(Xi)
        if norm > self.max_norm:
            return Xi / norm * self.max_norm
        return Xi

    def utilities(self):
        stability = -np.linalg.norm(self.Xi)
        exploitation = np.linalg.norm(self.Xi)
        oscillation = np.var(self.history[-10:], axis=0).mean() if len(self.history) > 10 else 0.0
        return np.array([stability, exploitation, -oscillation])

    def step(self):
        grads = self.grad_potentials(self.Xi)
        total_grad = np.sum(self.w[:, None] * grads, axis=0)

        if len(self.history) > 5:
            osc = np.var(self.history[-5:], axis=0)
        else:
            osc = 0.0

        self.Xi -= self.lr * (total_grad + self.gamma * osc)
        self.Xi = self.project_constraints(self.Xi)
        self.history.append(self.Xi.copy())

        self.w += self.beta * self.utilities()
        self.w = np.exp(self.w) / np.sum(np.exp(self.w))

        return self.Xi.copy(), self.w.copy()