import numpy as np

class QuantumInspiredOptimizer:
    """
    Quantum-behaved Particle Swarm Optimization (QPSO) variant.
    Used for hyperparameter tuning in highly non-linear, rugged fitness landscapes
    (like predicting RNG regimes) where traditional methods get trapped in local optima.
    """
    
    def __init__(self, n_particles=20, dimensions=4, max_iter=50):
        self.n = n_particles
        self.d = dimensions
        self.max_iter = max_iter
        
        # Bounds for hypothetical parameters (e.g. learning rate, momentum, etc)
        self.bounds = [
            (0.001, 0.1),  # param 1 freq
            (0.5, 2.0),    # param 2 amplitude
            (-0.2, 0.2),   # param 3 phase
            (10.0, 100.0)  # param 4 decay
        ]
        
        self.positions = np.zeros((self.n, self.d))
        self.pbest = np.zeros((self.n, self.d))
        self.pbest_fitness = np.ones(self.n) * float('inf')
        
        self.gbest = np.zeros(self.d)
        self.gbest_fitness = float('inf')
        
        self._initialize_particles()

    def _initialize_particles(self):
        for i in range(self.n):
            for j in range(self.d):
                low, high = self.bounds[j]
                self.positions[i, j] = np.random.uniform(low, high)
            self.pbest[i] = self.positions[i].copy()

    def optimize(self, fitness_func):
        """
        Runs the Quantum optimization against a provided fitness function.
        Goal is to MINIMIZE the fitness function (e.g. minimize prediction error).
        """
        for iteration in range(self.max_iter):
            # Calculate Mean Best Position (mbest) -> Quantum Center
            mbest = np.mean(self.pbest, axis=0)
            
            # Contraction-Expansion coefficient (Alpha)
            # Linearly decreases from 1.0 to 0.5 to encourage convergence
            alpha = 1.0 - 0.5 * (iteration / self.max_iter)
            
            for i in range(self.n):
                # Evaluate fitness
                fitness = fitness_func(self.positions[i])
                
                # Update Personal Best
                if fitness < self.pbest_fitness[i]:
                    self.pbest_fitness[i] = fitness
                    self.pbest[i] = self.positions[i].copy()
                    
                    # Update Global Best
                    if fitness < self.gbest_fitness:
                        self.gbest_fitness = fitness
                        self.gbest = self.positions[i].copy()
            
            # Update particle positions (Quantum Equation)
            for i in range(self.n):
                for j in range(self.d):
                    # Random attractor coordinates between pbest and gbest
                    phi = np.random.rand()
                    p = phi * self.pbest[i, j] + (1 - phi) * self.gbest[j]
                    
                    u = np.random.rand()
                    L = alpha * abs(mbest[j] - self.positions[i, j])
                    
                    if np.random.rand() > 0.5:
                        self.positions[i, j] = p + L * np.log(1 / u)
                    else:
                        self.positions[i, j] = p - L * np.log(1 / u)
                        
                    # Enforce Bounds
                    low, high = self.bounds[j]
                    self.positions[i, j] = np.clip(self.positions[i, j], low, high)
                    
        return self.gbest, self.gbest_fitness
