import numpy as np
from scipy.io.wavfile import read

# Load training and test data
X_train = read('../X_train.wav')[1].reshape(-1, 784)
y_train = (read('../y_train.wav')[1] * 9).astype(int)
X_test = read('../X_test.wav')[1].reshape(-1, 784)
y_test = (read('../y_test.wav')[1] * 9).astype(int)
X20 = X_test[:1000]
yt20 = y_test[:1000]


class QuantumSuperdecision:
    """
    A quantum-inspired classifier using superposition, interference,
    and measurement-collapse principles.
    
    Key differences from classical MLP:
    - Weights are complex (amplitude + phase)
    - State exists in superposition until measurement
    - Loss uses quantum-inspired interference terms
    - Prediction involves wave function collapse
    """
    
    def __init__(self, input_size, hidden_size, output_size, learning_rate=None):
        # Initialize quantum parameters (amplitude and phase)
        self.amplitude_1 = np.random.rand(input_size, hidden_size) * 0.01
        self.phase_1 = np.random.rand(input_size, hidden_size) * 2 * np.pi
        self.amplitude_2 = np.random.rand(hidden_size, output_size) * 0.01
        self.phase_2 = np.random.rand(hidden_size, output_size) * 2 * np.pi
        
        # Biases as complex numbers
        self.bias_1_amp = np.zeros((1, hidden_size))
        self.bias_1_phase = np.zeros((1, hidden_size))
        self.bias_2_amp = np.zeros((1, output_size))
        self.bias_2_phase = np.zeros((1, output_size))
        
        # Superposition state (probability amplitudes for each class)
        self.psi = np.zeros(output_size, dtype=complex)
        
        # Classical learning rates
        if learning_rate is None:
            learning_rate = np.array([0.001, 0.001, 0.001, 0.001])
        self.lr = learning_rate
        
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size

        # Cached forward-pass state for gradient updates
        self.X_complex = None
        self.z1 = None
        self.a1 = None
        
    def get_weights(self, layer='all'):
        """Retrieve complex weights from amplitude-phase representation."""
        if layer == 1 or layer == 'all':
            W1 = self.amplitude_1 * np.exp(1j * self.phase_1)
        else:
            W1 = None
        if layer == 2 or layer == 'all':
            W2 = self.amplitude_2 * np.exp(1j * self.phase_2)
        else:
            W2 = None
        return W1, W2
    
    def superposition_activation(self, z, collapse_threshold=0.9):
        """
        Quantum-inspired activation: applies phase rotation and creates
        superposition of active/inactive states based on amplitude.
        """
        # Phase rotation based on magnitude
        phase_rotation = np.angle(z)
        
        # Amplitude-based superposition (normalize)
        amplitude = np.abs(z)
        normalized = amplitude / (amplitude + 1 + 1e-9)  # Soft normalization
        
        # Apply ReLU-like behavior but in superposition
        activated = np.where(amplitude > 0, 
                             normalized * np.exp(1j * phase_rotation),
                             np.zeros_like(z, dtype=complex))
        return activated
    
    def forward(self, X, measure=False):
        """
        Forward pass through quantum superposition layers.
        If measure=True, collapse the wavefunction and return classical output.
        """
        # Phase-encode normalized pixels instead of using np.angle(X), which is
        # zero for non-negative real-valued images and therefore carries no
        # information.
        X_scaled = X.astype(np.float64) / 255.0
        input_phase = np.pi * X_scaled
        X_complex = X_scaled * np.exp(1j * input_phase)
        self.X_complex = X_complex
        
        # Layer 1: Quantum weight multiplication
        z1 = np.dot(X_complex, self.amplitude_1 * np.exp(1j * self.phase_1))
        b1 = self.bias_1_amp * np.exp(1j * self.bias_1_phase)
        z1 = z1 + b1
        self.z1 = z1
        
        # Superposition activation
        a1 = self.superposition_activation(z1)
        self.a1 = a1
        
        # Layer 2: Quantum weight multiplication
        z2 = np.dot(a1, self.amplitude_2 * np.exp(1j * self.phase_2))
        b2 = self.bias_2_amp * np.exp(1j * self.bias_2_phase)
        self.psi = z2 + b2  # Store as wavefunction state
        
        # Measurement: collapse superposition to classical probabilities
        if measure:
            probs = self.collapse_wavefunction(self.psi)
            return probs
        return self.psi
    
    def collapse_wavefunction(self, psi):
        """Measure the quantum state, collapsing to classical probabilities."""
        # Convert the complex state into stable class logits, then apply a
        # per-sample softmax. This preserves phase influence without letting
        # raw magnitudes dominate.
        logits = np.real(psi) + 0.1 * np.imag(psi)
        if logits.ndim == 1:
            shifted = logits - np.max(logits)
            exp_logits = np.exp(shifted)
            return exp_logits / (np.sum(exp_logits) + 1e-9)

        shifted = logits - np.max(logits, axis=1, keepdims=True)
        exp_logits = np.exp(shifted)
        return exp_logits / (np.sum(exp_logits, axis=1, keepdims=True) + 1e-9)
    
    def compute_quantum_loss(self, y_true, y_pred):
        """
        Quantum-inspired loss: combines cross-entropy with phase coherence.
        """
        # Classical cross-entropy
        m = y_true.shape[0]
        ce_loss = -np.sum(y_true * np.log(y_pred + 1e-9)) / m
        
        # Phase coherence term (encourage consistent phases)
        phase_variance = np.var(np.angle(self.psi))
        
        # Total loss: balance classical and quantum terms
        return ce_loss + 0.01 * phase_variance
    
    def backward(self, X, y_true, y_pred):
        """Quantum-inspired backpropagation with phase gradients."""
        m = y_true.shape[0]
        
        # Get complex weights
        W1 = self.amplitude_1 * np.exp(1j * self.phase_1)
        W2 = self.amplitude_2 * np.exp(1j * self.phase_2)
        
        # Gradient on measured probabilities. We use a linear surrogate from
        # logits back into the complex state.
        dz2 = (y_pred - y_true).astype(np.complex128)

        # Layer 2 gradients
        dW2 = np.dot(np.conj(self.a1).T, dz2) / m
        db2 = np.mean(dz2, axis=0, keepdims=True)

        # Backprop into hidden state with a simple surrogate derivative
        da1 = np.dot(dz2, np.conj(W2).T)
        activation_grad = 1.0 / (1.0 + np.abs(self.z1)) ** 2
        dz1 = da1 * activation_grad

        # Layer 1 gradients
        dW1 = np.dot(np.conj(self.X_complex).T, dz1) / m
        db1 = np.mean(dz1, axis=0, keepdims=True)

        return dW1, dW2, db1, db2
    
    def update(self, X, y_true):
        """Update quantum parameters (weights and phases)."""
        y_pred = self.forward(X, measure=True)
        dW1, dW2, db1, db2 = self.backward(X, y_true, y_pred)

        # Update amplitudes and phases from complex gradients.
        self.amplitude_1 -= self.lr[0] * np.real(dW1)
        self.phase_1 -= self.lr[1] * np.imag(dW1)
        self.amplitude_2 -= self.lr[2] * np.real(dW2)
        self.phase_2 -= self.lr[3] * np.imag(dW2)

        # Bias updates follow the same real/imaginary split.
        self.bias_1_amp -= self.lr[0] * np.real(db1)
        self.bias_1_phase -= self.lr[1] * np.imag(db1)
        self.bias_2_amp -= self.lr[2] * np.real(db2)
        self.bias_2_phase -= self.lr[3] * np.imag(db2)
        
        # Keep amplitudes bounded
        self.amplitude_1 = np.clip(self.amplitude_1, 0.0, 1.0)
        self.amplitude_2 = np.clip(self.amplitude_2, 0.0, 1.0)
        self.phase_1 = np.mod(self.phase_1, 2 * np.pi)
        self.phase_2 = np.mod(self.phase_2, 2 * np.pi)
    
    def predict(self, X):
        """Measure and collapse to get class predictions."""
        probs = self.forward(X, measure=True)
        return np.argmax(probs, axis=1)
    
    def score(self, X, y_true):
        """Compute accuracy."""
        y = self.predict(X)
        return np.mean(y == y_true)
    
    def measure_state(self, X):
        """Measure the quantum state and return detailed measurement info."""
        psi = self.forward(X, measure=False)
        probs = self.collapse_wavefunction(psi)
        
        return {
            'psi': psi,
            'amplitudes': np.abs(psi),
            'phases': np.angle(psi),
            'probabilities': probs,
            'entropy': -np.mean(np.sum(probs * np.log(probs + 1e-9), axis=-1)),
            'prediction': np.argmax(probs, axis=-1)
        }


# Initialize quantum classifier
f = QuantumSuperdecision(
    input_size=784, 
    hidden_size=100, 
    output_size=10,
    learning_rate=np.array([0.01, 0.1, 0.01, 0.1])
)

# Train with quantum-inspired updates
i = 0
print("Quantum Superdecision Training")
print("=" * 50)

while True:
    # Sample batch (with quantum-inspired shuffling)
    idx = np.random.randint(0, 60000, 100)
    X = X_train[idx]
    yt = y_train[idx]
    
    # One-hot encode
    y_onehot = np.eye(10)[yt]
    
    # Update quantum parameters
    f.update(X, y_onehot)
    
    # Print progress with quantum state info
    if i % 50 == 0:
        state = f.measure_state(X20[:10])
        print(f"Step {i:4d} | Acc: {f.score(X20, yt20):.4f} | "
              f"Avg Amp: {np.mean(np.abs(state['psi'])):.4f} | "
              f"Entropy: {state['entropy']:.4f}")
    
    i += 1
    if i > 1000:
        break