import numpy as np
from scipy.linalg import sqrtm
from scipy.stats import unitary_group
from tqdm import tqdm

def random_pure_state(dim_a, dim_b):
    """
    Generate a random pure state for a bipartite system.

    The random pure state is uniformly distributed by the Haar (Fubini-Study) measure on the unit sphere $S^{dim_a * dim_b - 1}$. (Invariant under the unitary group $U(dim_a) \times U(dim_b)$)
    
    Args:
        dim_a (int): Dimension of subsystem A
        dim_b (int): Dimension of subsystem B
        
    Returns:
        numpy.ndarray: Random pure state vector of shape (dim_a * dim_b,)
    """
    # Total dimension of the composite system
    dim_total = dim_a * dim_b
    
    # Generate non-zero random complex vector
    while True:
        state = np.random.normal(size=(dim_total,)) + 1j * np.random.normal(size=(dim_total,))
        if np.linalg.norm(state) > 0:
            break
    
    # Normalize the state
    state = state / np.linalg.norm(state)
    
    return state

def von_neumann_entropy_bipartite_pure_state(state, dim_a, dim_b):
    """
    Calculate the von Neumann entropy of the reduced density matrix.
    
    Args:
        state (numpy.ndarray): Pure state vector
        dim_a (int): Dimension of subsystem A
        dim_b (int): Dimension of subsystem B
        
    Returns:
        float: Von Neumann entropy
    """
    # Reshape state vector to matrix form
    state_matrix = state.reshape(dim_a, dim_b)
    
    # Calculate reduced density matrix of subsystem A
    rho_a = np.dot(state_matrix, state_matrix.conj().T)
    
    # Calculate eigenvalues
    eigenvals = np.linalg.eigvalsh(rho_a)
    
    # Remove very small eigenvalues (numerical errors)
    eigenvals = eigenvals[eigenvals > 1e-15]
    
    # Calculate von Neumann entropy
    entropy = -np.sum(eigenvals * np.log2(eigenvals))
    
    return np.real(entropy)

def sample_and_calculate(dim_a, dim_b, n_samples=1000):
    """
    Sample random pure states (generate random co) and calculate their von Neumann entropy.
    
    Args:
        dim_a (int): Dimension of subsystem A
        dim_b (int): Dimension of subsystem B
        n_samples (int): Number of samples to generate
        
    Returns:
        numpy.ndarray: Array of entropy values
    """
    entropies = np.zeros(n_samples)
    
    for i in tqdm(range(n_samples), desc=f"Sampling states (d_A={dim_a}, d_B={dim_b})", leave=False):
        state = random_pure_state(dim_a, dim_b)
        entropies[i] = von_neumann_entropy_bipartite_pure_state(state, dim_a, dim_b)
    
    return entropies

# Example usage:
if __name__ == "__main__":
    # Example: 2-qubit system
    dim_a, dim_b = 50,100
    
    # Generate single random state and calculate entropy
    state = random_pure_state(dim_a, dim_b)
    entropy = von_neumann_entropy_bipartite_pure_state(state, dim_a, dim_b)
    print(f"Single state entropy: {entropy}")
    
    # Sample multiple states
    entropies = sample_and_calculate(dim_a, dim_b, n_samples=1000)
    print(f"Expected entropy: {np.mean(entropies)}")
    print(f"Theoretical entropy: {np.log2(max(dim_a, dim_b))}")
    print(f"Standard deviation: {np.std(entropies)}")