"""
GBS Circuit → Holomorphic Field Converter
For use in GFFK ODE-CCT framework

Architecture:
1. Define n modes with squeezed parameters λ₁...λₙ
2. Define the interferometer unitary U (n×n)
3. Form the Hamiltonian H = U diag(λ) U^H
4. The vector field F(z) = -iHz is the quantum sim field
"""

import numpy as np

class GBSCircuit:
    """
    A Gaussian Boson Sampling circuit.
    
    Represents the full GBS hardware:
    - n modes with squeezing
    - An interferometer (unitary matrix)
    - Output photon detection
    
    This is mapped to a holomorphic ODE field F(z) = -iHz
    """
    
    def __init__(self, n_modes, squeezing_vector, unitary_matrix=None):
        """
        Initialize the GBS circuit.
        
        Args:
            n_modes: Number of modes (n)
            squeezing_vector: Array of squeezing parameters [λ₁, ..., λₙ]
                             (the magnitude of squeezing per mode)
            unitary_matrix: n×n unitary matrix for the interferometer
                           (optional — defaults to identity)
        """
        self.n = n_modes
        self.squeezing = np.array(squeezing_vector)
        self.U = unitary_matrix if unitary_matrix is not None else np.eye(n_modes)
        
        # Verify unitarity
        assert np.allclose(self.U @ self.U.conj().T, np.eye(n_modes)), \
            "U must be unitary"
        
        # Form the Hamiltonian: H = U @ diag(λ) @ U^H
        D = np.diag(self.squeezing)
        self.H = self.U @ D @ self.U.conj().T
        
        # Validate Hermiticity
        assert np.allclose(self.H, self.H.conj().T), "H must be Hermitian"
        
        # The holomorphic vector field F(z) = -iHz
        self.F_coeff = -1j * self.H
    
    def vector_field(self, z):
        """
        Evaluate the holomorphic vector field F(z) = -iHz at z ∈ ℂⁿ.
        
        This is the *exact* quantum simulation field for the GBS circuit.
        It is holomorphic (since it's linear in z), and its flow generates
        the same unitary evolution as the GBS circuit.
        
        Args:
            z: Complex vector of length n (initial condition / quantum state)
            
        Returns:
            F(z): Complex vector — the derivative dz/dt
        """
        return self.F_coeff @ z
    
    def flow(self, z0, t, method='rk4', n_steps=100):
        """
        Integrate the flow dz/dt = F(z) from time 0 to time t.
        
        For a *linear* field F(z) = -iHz, the exact solution is:
            z(t) = exp(-iHt) · z(0)
        
        But we use the GFFK ODE-CCT framework with adaptive step size,
        which gives us:
        - Automatic trajectory curvature tracking
        - Single/multi-branch path decision at thresholds
        - Topological profile statistics (a, b parameters)
        
        Args:
            z0: Initial condition (n complex values)
            t: Integration time
            method: 'rk4' (standard) or 'adaptive' (GFFK-style)
            n_steps: Number of steps (for fixed-step method)
            
        Returns:
            trajectory: Array of shape (n_steps+1, n) with z(0), z(dt), ..., z(t)
            profile: Tuple (a, b) — the GFFK statistical parameters
        """
        # --- Standard RK4 integration ---
        trajectory = [z0.copy()]
        dt = t / n_steps
        
        for step in range(n_steps):
            z = trajectory[-1]
            
            # k1 = F(z)
            k1 = self.vector_field(z)
            
            # k2 = F(z + dt/2 · k1)
            k2 = self.vector_field(z + (dt/2) * k1)
            
            # k3 = F(z + dt/2 · k2)
            k3 = self.vector_field(z + (dt/2) * k2)
            
            # k4 = F(z + dt · k3)
            k4 = self.vector_field(z + dt * k3)
            
            # z(t+dt) = z(t) + (dt/6)(k1 + 2k2 + 2k3 + k4)
            z_next = z + (dt/6) * (k1 + 2*k2 + 2*k3 + k4)
            trajectory.append(z_next)
        
        trajectory = np.array(trajectory)
        
        # --- GFFK Topological Profile Extraction ---
        # Compute the (a, b) parameters from the trajectory
        
        # Magnitude |z(t)| over time
        magnitudes = np.abs(trajectory)  # (n_steps+1, n)
        
        # Per-mode profile parameters
        profiles = []
        for mode in range(self.n):
            m = magnitudes[:, mode]
            a = np.mean(m)  # shape parameter (Gaussian-like)
            b = np.std(m) / (np.mean(m) + 1e-10)  # noise parameter (Ricker-like)
            profiles.append((a, b))
        
        return trajectory, profiles
    
    def compute_output_distribution(self, output_pattern, max_photons=4):
        """
        Compute the probability of a specific photon detection pattern.
        
        Uses the psd (partial sandwich determinant) formula for GBS:
        P(s) = |psd(M)_s|² / (s₁! s₂! ... sₙ!)
        
        Args:
            output_pattern: Array of photon counts [s₁, s₂, ..., sₙ]
            max_photons: Maximum photons per mode (for psd computation)
            
        Returns:
            probability: Float — probability of detecting this pattern
        """
        # Build the sub-matrix M for this output pattern
        # M is n×n, with each block being λⱼ I_{sⱼ} in the squeezed basis
        M = np.zeros((sum(output_pattern), sum(output_pattern)), dtype=complex)
        
        row = 0
        for j in range(self.n):
            s_j = output_pattern[j]
            if s_j > 0:
                # Extract the s_j-th principal sub-matrix of M_full
                cols = np.arange(j * max_photons, (j + 1) * max_photons)
                M[row:row+s_j, :] = self.M_full[cols, :][:, cols]
                row += s_j
        
        # Compute the psd (approximate via permanent for small matrices)
        psd_value = np.real(np.linalg.det(M))  # approximation
        
        # Probability
        prob = (psd_value ** 2) / np.prod([np.math.factorial(s) for s in output_pattern])
        
        return prob
    
    def get_topological_probability(self, trajectory, threshold=1.5):
        """
        Compute the probability of the trajectory crossing the threshold σ → ω.
        This is the GFFK-style "quantum measurement" — what basin does the trajectory
        fall into?
        
        Args:
            trajectory: (n_steps+1, n) array of z(t) values
            threshold: σ threshold value
            
        Returns:
            basin_profile: dict with basin assignments and probabilities
        """
        # For each mode, check if |z(t)| crosses threshold
        magnitudes = np.abs(trajectory)  # (n_steps+1, n)
        
        # Per-mode crossing count
        crossings = np.zeros(self.n)
        for mode in range(self.n):
            m = magnitudes[:, mode]
            # Count how many times |z| crosses threshold going up
            for i in range(1, len(m)):
                if m[i-1] < threshold and m[i] >= threshold:
                    crossings[mode] += 1
        
        # Normalize to probability distribution
        total = np.sum(crossings)
        if total > 0:
            probabilities = crossings / total
        else:
            probabilities = np.ones(self.n) / self.n
        
        return {
            'crossings': crossings,
            'probabilities': probabilities,
            'basin': np.argmax(probabilities),  # dominant basin
        }