"""
GBS → GFFK ODE-CCT Pipeline (FIXED v3 - with safeguards)
========================================================
Fixed: Division by zero in probability calculation
Fixed: Robust vacuum probability
Fixed: Numerical stability for small determinants
"""

import numpy as np
from math import factorial
import sys

# ============================================================
# HAFNIAN COMPUTATION
# ============================================================

def hafnian(M):
    """Compute the hafnian of a symmetric matrix M."""
    n = M.shape[0]
    if n % 2 != 0:
        return 0.0
    if n == 0:
        return 1.0
    if n == 2:
        return M[0, 1]
    return _hafnian_direct(M)


def _hafnian_direct(M):
    """Direct enumeration of perfect matchings."""
    n = M.shape[0]
    if n == 0:
        return 1.0
    if n == 2:
        return M[0, 1]
    
    haf = 0.0
    for j in range(1, n):
        indices = list(range(1, n))[:j] + list(range(j+1, n))
        sub_M = M[np.ix_(indices, indices)]
        haf += M[0, j] * _hafnian_direct(sub_M)
    
    return haf


# ============================================================
# GBS → GFFK ODE-CCT PIPELINE
# ============================================================

class GBSGFFK_Fixed:
    """
    Correct GBS → GFFK ODE-CCT pipeline with numerical safeguards.
    """
    
    def __init__(self, n_modes, squeezing_params=None, unitary=None):
        self.n = n_modes
        
        if squeezing_params is None:
            self.squeezing = np.ones(n_modes) * 0.8
        else:
            self.squeezing = np.array(squeezing_params)
        
        if unitary is None:
            self.U = self._random_haar_unitary(self.n)
        else:
            self.U = np.array(unitary, dtype=complex)
        
        # Covariance matrix
        D_squeeze = np.diag(2 * self.squeezing + 1)
        self.Sigma = self.U @ D_squeeze @ self.U.conj().T
        
        # Hamiltonian
        D_H = np.diag(self.squeezing)
        self.H = self.U @ D_H @ self.U.conj().T
        
        # Holomorphic vector field
        self.F = -1j * self.H
        
        # Pre-compute log-determinant for numerical stability
        self.log_det_sigma = self._safe_log_det()
    
    def _safe_log_det(self):
        """Compute log(det(Sigma)) with numerical safeguards."""
        try:
            eigenvalues = np.linalg.eigvalsh(self.Sigma)
            # Clamp negative eigenvalues to small positive value
            eigenvalues = np.maximum(eigenvalues, 1e-10)
            return np.sum(np.log(eigenvalues))
        except:
            return 0.0
    
    def _random_haar_unitary(self, n):
        Z = np.random.randn(n, n) + 1j * np.random.randn(n, n)
        Q, R = np.linalg.qr(Z)
        d = np.diagonal(R)
        return Q @ np.diag(d / np.abs(d))
    
    def evolve_state(self, initial_state, time, n_steps=1000):
        trajectory = np.zeros((n_steps + 1, self.n), dtype=complex)
        trajectory[0] = initial_state.copy()
        
        dt = time / n_steps
        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)
        
        profiles = {}
        for mode in range(self.n):
            m = np.abs(trajectory[:, mode])
            a = np.mean(m)
            b = np.std(m) / (np.mean(m) + 1e-10)
            profiles[mode] = {'a': a, 'b': b}
        
        return trajectory, profiles
    
    def compute_output_distribution(self, max_photons=3):
        A = self.Sigma - np.eye(self.n)
        distribution = {}
        patterns = self._enumerate_patterns(max_photons)
        
        for pattern in patterns:
            prob = self._compute_pattern_probability(pattern, A)
            if prob is not None and prob > 1e-15:
                distribution[pattern] = prob
        
        total = sum(distribution.values())
        if total > 0:
            for k in distribution:
                distribution[k] /= total
        
        return distribution
    
    def _compute_pattern_probability(self, pattern, A):
        total_photons = sum(pattern)
        
        if total_photons == 0:
            # Vacuum probability with safeguards
            try:
                prob = np.exp(-0.5 * self.log_det_sigma)
                return prob
            except:
                return 0.0
        
        A_sub = self._build_sub_matrix(pattern, A)
        haf = hafnian(A_sub)
        
        # Compute normalization with safeguards
        try:
            det_sigma = np.linalg.det(self.Sigma)
            if abs(det_sigma) < 1e-20:
                # Use log-determinant instead
                norm = np.exp(0.5 * self.log_det_sigma)
            else:
                norm = np.sqrt(abs(det_sigma))
        except:
            norm = 1.0
        
        if norm == 0:
            return None
        
        prob = np.abs(haf) ** 2 / (
            np.prod([factorial(s) for s in pattern]) * norm
        )
        
        return prob
    
    def _build_sub_matrix(self, pattern, A):
        """Build sub-matrix for hafnian computation."""
        n = len(pattern)
        total = sum(pattern)
        
        if total == 0:
            return np.array([[]])
        
        A_sub = np.zeros((total, total), dtype=complex)
        
        row_offset = 0
        for j in range(n):
            n_j = pattern[j]
            if n_j > 0:
                for k in range(n_j):
                    row = row_offset + k
                    A_sub[row, row] = A[j, j]
                    for m in range(row + 1, total):
                        A_sub[row, m] = A[j, j]
                        A_sub[m, row] = A[j, j]
                row_offset += n_j
        
        A_sub = (A_sub + A_sub.T) / 2
        return A_sub
    
    def _enumerate_patterns(self, max_photons):
        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 get_quantum_advantage_indicator(self, distribution):
        probs = np.array([p for p in distribution.values() if p > 1e-10])
        
        if len(probs) == 0:
            return {'advantage_score': 0, 'n_patterns': 0, 'entropy': 0}
        
        # Safe entropy calculation
        entropy = 0.0
        for p in probs:
            if p > 0:
                entropy -= p * np.log2(p)
        
        n_patterns = len(probs)
        advantage = entropy * np.log2(n_patterns + 1)
        
        return {
            'advantage_score': advantage,
            'n_patterns': n_patterns,
            'entropy': entropy,
            'total_patterns': len(distribution),
        }


# ============================================================
# MAIN EXECUTION
# ============================================================

if __name__ == "__main__":
    n = 4
    squeezing = np.ones(n) * 0.8
    
    gbs = GBSGFFK_Fixed(n_modes=n, squeezing_params=squeezing)
    
    print("=" * 60)
    print("GBS → GFFK ODE-CCT PIPELINE (FIXED v3)")
    print("=" * 60)
    print(f"Modes: {n}")
    print(f"Squeezing: {squeezing}")
    print(f"Covariance matrix det: {np.linalg.det(gbs.Sigma):.6f}")
    print(f"Log-det: {gbs.log_det_sigma:.6f}")
    print("=" * 60)
    
    initial_state = np.ones(n, dtype=complex) * 0.5
    time = 2.0
    trajectory, profiles = gbs.evolve_state(initial_state, time, n_steps=500)
    
    print("\n--- Trajectory Profiles ---")
    for mode in range(n):
        p = profiles[mode]
        print(f"Mode {mode}: a={p['a']:.4f}, b={p['b']:.4f}")
    
    distribution = gbs.compute_output_distribution(max_photons=3)
    
    print(f"\n--- Output Distribution ({len(distribution)} patterns) ---")
    sorted_dist = sorted(distribution.items(), key=lambda x: x[1], reverse=True)
    for pattern, prob in sorted_dist[:15]:
        print(f"  Pattern {pattern}: P = {prob:.8f}")
    
    total_prob = sum(distribution.values())
    print(f"\nTotal probability: {total_prob:.8f} (should be ≈1.0)")
    
    advantage = gbs.get_quantum_advantage_indicator(distribution)
    print("\n--- Quantum Advantage Indicator ---")
    print(f"  Score: {advantage['advantage_score']:.4f}")
    print(f"  Patterns: {advantage['n_patterns']}")
    print(f"  Entropy: {advantage['entropy']:.4f} bits (must be ≥ 0)")
    print("=" * 60)