"""
GBS → GFFK ODE-CCT Pipeline
============================
End-to-end conversion of a GBS circuit into an ODE-CCT computation.

Steps:
1. Random unitary matrix → interferometer
2. Squeezed state → field initialization
3. Unitary evolution → holomorphic vector field F(z) = -iHz
4. ODE integration → quantum state trajectory
5. Trajectory topology → basin analysis (GBS output distribution)
6. Output: photon pattern probabilities

This is the *quantum sim* in the GFFK framework.
"""

import numpy as np
import scipy.stats as stats
from scipy.special import factorial

class GBSGFFK:
    """
    Full GBS circuit → GFFK ODE-CCT conversion pipeline.
    
    Takes a GBS circuit (n modes, squeezed params, unitary matrix),
    converts to a holomorphic field, integrates via ODE-CCT,
    and outputs the photon detection probability distribution.
    """
    
    def __init__(self, n_modes, squeezing_params=None, unitary=None):
        """
        Initialize the GBS-GFFK system.
        
        Args:
            n_modes: Number of modes
            squeezing_params: Array of squeezing magnitudes (default: uniform)
            unitary: n×n unitary matrix (default: random Haar unitary)
        """
        self.n = n_modes
        
        # Squeezing parameters
        if squeezing_params is None:
            # Default: moderate squeezing (λ ≈ 1.0)
            self.squeezing = np.ones(n_modes) * 1.0
        else:
            self.squeezing = np.array(squeezing_params)
        
        # Random Haar unitary if none provided
        if unitary is None:
            self.U = self._random_haar_unitary()
        else:
            self.U = np.array(unitary)
        
        # Hamiltonian H = U diag(λ) U^H
        D = np.diag(self.squeezing)
        self.H = self.U @ D @ self.U.conj().T
        
        # Holomorphic vector field F(z) = -iHz
        self.F = -1j * self.H
    
    def _random_haar_unitary(self, n):
        """Generate a random unitary matrix from the Haar measure."""
        # Use the QR decomposition of a random complex matrix
        Z = np.random.randn(n, n) + 1j * np.random.randn(n, n)
        Q, R = np.linalg.qr(Z)
        # Make it Haar-distributed
        d = np.diagonal(R)
        ph = d / np.abs(d)
        U = Q @ np.diag(ph)
        return U
    
    def evolve_state(self, initial_state, time, n_steps=1000):
        """
        Evolve the quantum state via ODE integration.
        
        Args:
            initial_state: Complex vector of length n
            time: Integration time (must be > 0)
            n_steps: Number of RK4 steps
            
        Returns:
            trajectory: (n_steps+1, n) complex array
            profile: GFFK topological parameters (a, b) per mode
        """
        # Initialize trajectory array
        trajectory = np.zeros((n_steps + 1, self.n), dtype=complex)
        trajectory[0] = initial_state.copy()
        
        dt = time / n_steps
        
        # RK4 integration
        for step in range(n_steps):
            z = trajectory[step]
            
            k1 = self.F @ z
            k2 = self.F @ (z + 0.5 * dt * k1)
            k3 = self.F @ (z + 0.5 * dt * k2)
            k4 = self.F @ (z + dt * k3)
            
            trajectory[step + 1] = z + (dt / 6) * (k1 + 2*k2 + 2*k3 + k4)
        
        # --- GFFK Topological Profile ---
        # For each mode, compute (a, b) from the trajectory
        profiles = {}
        for mode in range(self.n):
            m = np.abs(trajectory[:, mode])  # |z(t)| for this mode
            profiles[mode] = self._compute_profile(m)
        
        return trajectory, profiles
    
    def _compute_profile(self, magnitudes):
        """
        Compute GFFK topological profile parameters (a, b).
        
        a (shape): Mean magnitude (Gaussian-like component)
        b (noise): Standard deviation relative to mean (Ricker-like component)
        """
        mean_mag = np.mean(magnitudes)
        std_mag = np.std(magnitudes)
        
        # Shape parameter (a): proportional to mean
        a = mean_mag
        
        # Noise parameter (b): relative fluctuation
        b = std_mag / (mean_mag + 1e-10)
        
        return {'a': a, 'b': b}
    
    def classify_attractor(self, trajectory, mode=0):
        """
        Classify the attractor topology for a specific mode.
        
        Returns the GFFK classification based on the (a, b) profile:
        - Lorenz attractor (strange): high b, moderate a
        - Ricker's wavelet attractor: high a, low b
        - Limit cycle: periodic oscillation
        
        Args:
            trajectory: (n_steps+1, n) complex array
            mode: Mode index to analyze
            
        Returns:
            classification: String describing the attractor type
        """
        m = np.abs(trajectory[:, mode])
        
        # Compute profile
        a = np.mean(m)
        b = np.std(m) / (np.mean(m) + 1e-10)
        
        # Classification rules (from GFFK literature)
        if b > 0.5 and a < 1.0:
            return 'Lorenz_attractor'
        elif b < 0.3 and a > 0.8:
            return 'Ricker_wavelet'
        elif abs(np.std(np.diff(m))) < 0.01:
            return 'Limit_cycle'
        else:
            return 'Fixed_point'
    
    def compute_output_distribution(self, max_photons=3):
        """
        Compute the full output distribution for the GBS circuit.
        
        Returns a dictionary of {photon_pattern: probability}.
        
        Args:
            max_photons: Maximum photons per mode
            
        Returns:
            distribution: Dict mapping tuples to probabilities
        """
        distribution = {}
        
        # Enumerate all possible output patterns
        patterns = self._enumerate_patterns(max_photons)
        
        for pattern in patterns:
            prob = self._compute_pattern_probability(pattern)
            distribution[pattern] = prob
        
        return distribution
    
    def _enumerate_patterns(self, max_photons):
        """Enumerate all photon number patterns up to max_photons per mode."""
        patterns = []
        
        def recurse(mode, current):
            if mode == self.n:
                patterns.append(tuple(current))
                return
            for k in range(max_photons + 1):
                current.append(k)
                recurse(mode + 1, current)
                current.pop()
        
        recurse(0, [])
        return patterns
    
    def _compute_pattern_probability(self, pattern):
        """
        Compute probability of a specific photon detection pattern.
        
        Uses the permanent of the sub-matrix of M = U diag(λ) U^H.
        For small matrices, we compute the permanent directly.
        """
        total_photons = sum(pattern)
        if total_photons == 0:
            return 1.0  # No photons detected
        
        # Build the matrix M for this pattern
        # M is a sum of outer products weighted by squeezing
        M = np.zeros((self.n, self.n), dtype=complex)
        for j in range(self.n):
            if pattern[j] > 0:
                # Each detected photon corresponds to a squeeze term
                M += self.squeezing[j] * (self.U[:, j] @ self.U[:, j].conj().T)
        
        # Compute the permanent of M (for small matrices)
        # Using the Ryser formula for permanent
        permanant = self._ryser_permanent(M)
        
        # Probability
        prob = (np.abs(permanant) ** 2) / np.prod(
            [factorial(s) for s in pattern]
        )
        
        return prob
    
    def _ryser_permanent(self, A):
        """
        Compute the permanent of matrix A using Ryser's formula.
        For small matrices (n ≤ 10), this is efficient.
        """
        n = A.shape[0]
        permanent = 0
        for S in range(2**n):
            # Subset S of columns
            row_sums = np.sum(A[:, list(np.where((S >> np.arange(n)) > 0)[0])], axis=1)
            sign = (-1) ** (n - bin(S).count('1'))
            permanent += sign * np.prod(row_sums)
        
        return permanent
    
    def get_quantum_advantage_indicator(self):
        """
        Compute an indicator of quantum computational advantage.
        
        The "advantage" is measured by:
        1. Output distribution complexity (permanent computation is #P-hard)
        2. Number of distinguishable output patterns
        3. Entropy of the output distribution
        
        Returns:
            advantage_score: Float (higher = more advantage)
        """
        # Compute output distribution
        distribution = self.compute_output_distribution(max_photons=2)
        
        # Number of non-zero patterns
        n_patterns = len([p for p, prob in distribution.items() if prob > 1e-10])
        
        # Shannon entropy of distribution
        probs = np.array([p for p in distribution.values() if p > 1e-10])
        entropy = -np.sum(probs * np.log2(probs + 1e-30))
        
        # Advantage score: combination of pattern count and entropy
        # (Higher = more quantum advantage)
        advantage = np.log2(n_patterns + 1) * (1 + entropy / 10)
        
        return {
            'advantage_score': advantage,
            'n_patterns': n_patterns,
            'entropy': entropy,
            'total_patterns': len(distribution),
        }
