import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import entropy
from collections import deque

class EntropyCollapseSimulator:
    """
    Quantum collapse simulation using ENTROPY, not raw values.
    
    Key insight: 
    - Raw values: seed-specific, non-generalizable
    - Entropy values: uncertainty structure, generalizable
    - Law of Large Numbers converges on entropy field, not value field
    """
    
    def __init__(self, N=100, num_bins=20):
        self.N = N
        self.num_bins = num_bins
        self.collapse_log = []
        self.entropy_trajectory = []
        
    def collapse_to_entropy(self, seed, distribution='normal'):
        """
        Single collapse event: produce raw outcome, then extract ENTROPY.
        
        Args:
            seed: The observer choice (triggers collapse)
            distribution: Type of wave function
        
        Returns:
            entropy_value: The semantic uncertainty of this collapsed state
            raw_state: For visualization (but not used for generalization)
        """
        # Step 1: Collapse via seed (standard quantum behavior)
        np.random.seed(seed)
        
        if distribution == 'normal':
            raw_state = np.random.normal(0, 1, self.N)
        elif distribution == 'uniform':
            raw_state = np.random.uniform(-3, 3, self.N)
        elif distribution == 'bimodal':
            samples = []
            for _ in range(self.N):
                if np.random.random() < 0.5:
                    samples.append(np.random.normal(-2, 0.5))
                else:
                    samples.append(np.random.normal(2, 0.5))
            raw_state = np.array(samples)
        elif distribution == 'double_slit':
            raw_state = self._double_slit_collapse(seed)
        
        # Step 2: Extract ENTROPY (not raw values)
        # Entropy captures uncertainty structure - generalizes across seeds
        H = self._compute_entropy(raw_state)
        
        # Step 3: Also compute secondary entropy measures for richer structure
        H_shannon = self._shannon_entropy(raw_state)
        H_spectral = self._spectral_entropy(raw_state)
        H_conditional = self._conditional_entropy(raw_state)
        
        self.collapse_log.append({
            'seed': seed,
            'H_shannon': H_shannon,
            'H_spectral': H_spectral,
            'H_conditional': H_conditional,
            'raw_mean': np.mean(raw_state),
            'raw_std': np.std(raw_state)
        })
        
        return {
            'H_total': H,
            'H_shannon': H_shannon,
            'H_spectral': H_spectral,
            'H_conditional': H_conditional,
            'raw_state': raw_state  # Keep for visualization only
        }
    
    def _compute_entropy(self, data):
        """Compute probability distribution entropy."""
        hist, _ = np.histogram(data, bins=self.num_bins, density=True)
        hist = hist / hist.sum()  # Normalize to probability
        hist = hist[hist > 0]  # Remove zeros
        return entropy(hist)
    
    def _shannon_entropy(self, data):
        """Shannon entropy of the distribution."""
        return self._compute_entropy(data)
    
    def _spectral_entropy(self, data):
        """Entropy based on Fourier spectrum (captures periodicity)."""
        fft = np.fft.fft(data)
        power_spectrum = np.abs(fft)**2
        power_spectrum = power_spectrum / power_spectrum.sum()
        power_spectrum = power_spectrum[power_spectrum > 0]
        return entropy(power_spectrum)
    
    def _conditional_entropy(self, data):
        """
        Conditional entropy: H(current | previous).
        Captures temporal structure of collapse.
        """
        # Discretize into states
        states = np.digitize(data, bins=np.linspace(data.min(), data.max(), 5))
        
        # Calculate conditional entropy H(X_t | X_{t-1})
        transitions = []
        for i in range(len(states) - 1):
            transitions.append((states[i], states[i+1]))
        
        # Build transition matrix
        unique_states = np.unique(states)
        transition_counts = np.zeros((len(unique_states), len(unique_states)))
        
        for t in transitions:
            i = np.where(unique_states == t[0])[0][0]
            j = np.where(unique_states == t[1])[0][0]
            transition_counts[i, j] += 1
        
        # Normalize to get conditional probabilities
        row_sums = transition_counts.sum(axis=1, keepdims=True)
        row_sums[row_sums == 0] = 1  # Avoid division by zero
        P_conditional = transition_counts / row_sums
        
        # Compute conditional entropy H(X_t | X_{t-1}) = sum P(x,y) * log P(x|y)
        H_conditional = 0
        for i in range(len(unique_states)):
            for j in range(len(unique_states)):
                if transition_counts[i, j] > 0:
                    p_xy = transition_counts[i, j] / len(transitions)
                    p_x_given_y = P_conditional[i, j]
                    H_conditional -= p_xy * np.log(p_x_given_y + 1e-10)
        
        return H_conditional
    
    def aggregate_entropy_field(self, num_seeds, distribution='normal'):
        """
        Aggregate entropy values across many collapse events (seeds).
        This is the Law of Large Numbers on ENTROPY.
        """
        H_shannon_list = []
        H_spectral_list = []
        H_conditional_list = []
        
        for seed in range(num_seeds):
            result = self.collapse_to_entropy(seed, distribution)
            H_shannon_list.append(result['H_shannon'])
            H_spectral_list.append(result['H_spectral'])
            H_conditional_list.append(result['H_conditional'])
        
        return {
            'H_shannon_field': np.array(H_shannon_list),
            'H_spectral_field': np.array(H_spectral_list),
            'H_conditional_field': np.array(H_conditional_list),
            'mean_shannon': np.mean(H_shannon_list),
            'std_shannon': np.std(H_shannon_list),
            'mean_spectral': np.mean(H_spectral_list),
            'mean_conditional': np.mean(H_conditional_list)
        }