import numpy as np
from scipy.io.wavfile import read
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# ============================================================
# Load MNIST Data
# ============================================================
X_train = read('../X_train.wav')[1].reshape(-1, 784).astype(np.float64) / 255.0
y_train = (read('../y_train.wav')[1] * 9).astype(int)
X_test = read('../X_test.wav')[1].reshape(-1, 784).astype(np.float64) / 255.0
y_test = (read('../y_test.wav')[1] * 9).astype(int)

print(f"Train: {X_train.shape}, Labels: {y_train.shape}")
print(f"Test: {X_test.shape}, Labels: {y_test.shape}")

# ============================================================
# Activation Functions
# ============================================================
def relu(x):
    return np.maximum(0, x)

def relu_derivative(x):
    return np.where(x > 0, 1.0, 0.0)

def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-np.clip(x, -500, 500)))

def softmax(x):
    exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
    return exp_x / np.sum(exp_x, axis=1, keepdims=True)

def tanh(x):
    return np.tanh(x)


# ============================================================
# Quantum Vertex Generator
# ============================================================
class QuantumVertexGenerator:
    def __init__(self, input_size=784, num_vertices=12, batch_size=100, seed=42):
        np.random.seed(seed)
        
        self.num_vertices = num_vertices
        self.batch_size = batch_size
        self.input_size = input_size
        
        # Learnable base vertices
        self.base_vertices = np.random.randn(num_vertices, 3) * 0.5
        
        # Feature extraction
        self.W1 = np.random.randn(input_size, 256) * np.sqrt(2.0 / input_size)
        self.b1 = np.zeros(256)
        
        self.W2 = np.random.randn(256, 128) * np.sqrt(2.0 / 256)
        self.b2 = np.zeros(128)
        
        self.W3 = np.random.randn(128, 64) * np.sqrt(2.0 / 128)
        self.b3 = np.zeros(64)
        
        # Transformation predictors
        self.W_rot = np.random.randn(64, num_vertices * 4) * 0.01
        self.b_rot = np.zeros(num_vertices * 4)
        
        self.W_scale = np.random.randn(64, num_vertices * 3) * 0.01
        self.b_scale = np.zeros(num_vertices * 3)
        
        self.W_trans = np.random.randn(64, 3) * 0.01
        self.b_trans = np.zeros(3)
        
        self.cache = {}
    
    def _rotate_by_quaternion(self, v, q):
        """Rotate vector v by quaternion q = [w, x, y, z]"""
        w, x, y, z = q[0], q[1], q[2], q[3]
        
        vx, vy, vz = v[0], v[1], v[2]
        
        result_x = 2*( (w*w + x*x - 0.5)*vx + (x*y - w*z)*vy + (x*z + w*y)*vz )
        result_y = 2*( (x*y + w*z)*vx + (w*w + y*y - 0.5)*vy + (y*z - w*x)*vz )
        result_z = 2*( (x*z - w*y)*vx + (y*z + w*x)*vy + (w*w + z*z - 0.5)*vz )
        
        return np.array([result_x, result_y, result_z])
    
    def forward(self, X):
        batch_size = X.shape[0]
        
        # Feature extraction
        h1 = relu(X @ self.W1 + self.b1)
        h2 = relu(h1 @ self.W2 + self.b2)
        h3 = relu(h2 @ self.W3 + self.b3)
        
        self.cache['h3'] = h3
        
        # Predictions
        quaternions = h3 @ self.W_rot + self.b_rot
        scales = sigmoid(h3 @ self.W_scale + self.b_scale)
        translation = tanh(h3 @ self.W_trans + self.b_trans)
        
        self.cache['quaternions'] = quaternions
        self.cache['scales'] = scales
        self.cache['translation'] = translation
        
        # Generate vertices
        vertices = np.zeros((batch_size, self.num_vertices, 3))
        
        for b in range(batch_size):
            for v in range(self.num_vertices):
                v_base = self.base_vertices[v]
                
                q = quaternions[b, v*4:(v+1)*4]
                q_norm = np.linalg.norm(q) + 1e-8
                q = q / q_norm
                
                v_rot = self._rotate_by_quaternion(v_base, q)
                scale_v = scales[b, v*3:(v+1)*3]
                v_scaled = v_rot * (scale_v + 0.5)
                
                vertices[b, v] = v_scaled + translation[b]
        
        self.cache['vertices'] = vertices
        return vertices


# ============================================================
# Quantum Vertex ODE Function
# ============================================================
class QuantumVertexODE:
    def __init__(self, num_vertices=12, vertex_dim=3, hidden_dim=128, seed=42):
        np.random.seed(seed)
        
        self.num_vertices = num_vertices
        self.vertex_dim = vertex_dim
        self.state_dim = num_vertices * vertex_dim  # 12 * 3 = 36
        
        # ODE dynamics network
        # Input: state_dim (36) + time (1) = 37
        self.W_dyn1 = np.random.randn(self.state_dim + 1, hidden_dim) * np.sqrt(2.0 / (self.state_dim + 1))
        self.b_dyn1 = np.zeros(hidden_dim)
        
        self.W_dyn2 = np.random.randn(hidden_dim, hidden_dim) * np.sqrt(2.0 / hidden_dim)
        self.b_dyn2 = np.zeros(hidden_dim)
        
        self.W_dyn3 = np.random.randn(hidden_dim, self.state_dim) * 0.01
        self.b_dyn3 = np.zeros(self.state_dim)
        
        # Simpler edge-based interaction (no attention matrix)
        self.W_edge = np.random.randn(self.state_dim * 2 + 1, hidden_dim // 2) * 0.01
        self.b_edge = np.zeros(hidden_dim // 2)
        
    def dynamics(self, t, state):
        """
        Compute dx/dt for vertex evolution.
        
        Args:
            t: current time (scalar)
            state: (state_dim,) flattened vertex positions = 36
        Returns:
            d_state/dt: (state_dim,)
        """
        state = state.reshape(-1)
        
        # Reshape to (num_vertices, vertex_dim) = (12, 3)
        vertices = state.reshape(self.num_vertices, self.vertex_dim)
        
        # Compute pairwise interactions (simplified)
        # For each vertex pair, compute interaction
        interactions = np.zeros(self.state_dim)
        
        for i in range(self.num_vertices):
            for j in range(self.num_vertices):
                if i != j:
                    diff = vertices[j] - vertices[i]  # (3,)
                    dist = np.linalg.norm(diff) + 1e-8
                    
                    # Interaction based on distance
                    interaction_strength = np.exp(-dist) * 0.1
                    
                    # Add to vertex i's dynamics
                    interactions[i*3:(i+1)*3] += interaction_strength * diff
        
        # Combine state with time
        combined = np.concatenate([state, [t]])
        
        # Dynamics network
        h1 = relu(combined @ self.W_dyn1 + self.b_dyn1)
        h2 = relu(h1 @ self.W_dyn2 + self.b_dyn2)
        dynamics = h2 @ self.W_dyn3 + self.b_dyn3
        
        # Add interactions
        dynamics = dynamics + interactions * 0.5
        
        return dynamics
    
    def solve(self, initial_state, t_span=(0, 1.5), t_eval=None):
        if t_eval is None:
            t_eval = np.linspace(t_span[0], t_span[1], 50)
        
        sol = solve_ivp(
            fun=self.dynamics,
            t_span=t_span,
            y0=initial_state,
            t_eval=t_eval,
            method='RK45',
            rtol=1e-4,
            atol=1e-6
        )
        
        return sol.y.T


# ============================================================
# Quantum 3D Vertex Classifier
# ============================================================
class Quantum3DClassifier:
    def __init__(self, input_size=784, num_vertices=12, num_classes=10, 
                 batch_size=100, t_span=(0, 1.5), seed=42):
        np.random.seed(seed)
        
        self.num_classes = num_classes
        self.num_vertices = num_vertices
        self.batch_size = batch_size
        self.t_span = t_span
        
        # Vertex generator
        self.vertex_gen = QuantumVertexGenerator(
            input_size=input_size,
            num_vertices=num_vertices,
            batch_size=batch_size,
            seed=seed
        )
        
        # ODE solver
        self.ode = QuantumVertexODE(
            num_vertices=num_vertices,
            vertex_dim=3,
            hidden_dim=128,
            seed=seed + 1
        )
        
        # Trajectory encoder
        state_dim = num_vertices * 3  # 36
        self.W_traj1 = np.random.randn(state_dim * 3, 256) * np.sqrt(2.0 / (state_dim * 3))
        self.b_traj1 = np.zeros(256)
        self.W_traj2 = np.random.randn(256, 128) * np.sqrt(2.0 / 256)
        self.b_traj2 = np.zeros(128)
        
        # Vertex statistics encoder
        stat_dim = num_vertices * 3 * 5  # 12 * 3 * 5 = 180
        self.W_stat1 = np.random.randn(stat_dim, 128) * np.sqrt(2.0 / stat_dim)
        self.b_stat1 = np.zeros(128)
        self.W_stat2 = np.random.randn(128, 64) * np.sqrt(2.0 / 128)
        self.b_stat2 = np.zeros(64)
        
        # Time embedding
        self.time_embed = np.random.randn(50) * 0.01
        
        # Final classifier
        input_dim = 128 + 64 + 50  # trajectory + stats + time_embed
        self.W_cls1 = np.random.randn(input_dim, 256) * np.sqrt(2.0 / input_dim)
        self.b_cls1 = np.zeros(256)
        self.W_cls2 = np.random.randn(256, 128) * np.sqrt(2.0 / 256)
        self.b_cls2 = np.zeros(128)
        self.W_cls3 = np.random.randn(128, num_classes) * 0.01
        self.b_cls3 = np.zeros(num_classes)
        
    def compute_trajectory_features(self, trajectory):
        """trajectory: (num_steps, batch, state_dim)"""
        initial = trajectory[0]
        final = trajectory[-1]
        mean_traj = trajectory.mean(axis=0)
        
        return np.concatenate([initial, final, mean_traj], axis=1)
    
    def compute_vertex_statistics(self, vertices):
        """vertices: (batch, num_vertices, 3)"""
        mean = vertices.mean(axis=1)
        std = vertices.std(axis=1)
        max_v = vertices.max(axis=1)
        min_v = vertices.min(axis=1)
        
        stats = np.concatenate([mean, std, max_v, min_v, max_v - min_v], axis=1)
        stats = np.repeat(stats[:, np.newaxis, :], self.num_vertices, axis=1)
        
        return stats.reshape(stats.shape[0], -1)
    
    def forward(self, X, return_trajectory=False):
        batch_size = X.shape[0]
        
        # Generate vertices
        vertices = self.vertex_gen.forward(X)
        vertices_flat = vertices.reshape(batch_size, -1)
        
        # Time points
        t_eval = np.linspace(self.t_span[0], self.t_span[1], 30)  # Reduced for speed
        
        # Solve ODE for each sample
        trajectory = []
        for b in range(batch_size):
            traj_b = self.ode.solve(vertices_flat[b], self.t_span, t_eval)
            trajectory.append(traj_b)
        
        trajectory = np.stack(trajectory, axis=1)
        
        # Trajectory features
        traj_features = self.compute_trajectory_features(trajectory)
        h_traj = relu(traj_features @ self.W_traj1 + self.b_traj1)
        h_traj = relu(h_traj @ self.W_traj2 + self.b_traj2)
        
        # Vertex statistics
        final_vertices = trajectory[-1].reshape(batch_size, self.num_vertices, 3)
        vertex_stats = self.compute_vertex_statistics(final_vertices)
        h_stat = relu(vertex_stats @ self.W_stat1 + self.b_stat1)
        h_stat = relu(h_stat @ self.W_stat2 + self.b_stat2)
        
        # Combine
        combined = np.concatenate([
            h_traj, 
            h_stat, 
            np.tile(self.time_embed, (batch_size, 1))
        ], axis=1)
        
        # Classifier
        h_cls = relu(combined @ self.W_cls1 + self.b_cls1)
        h_cls = relu(h_cls @ self.W_cls2 + self.b_cls2)
        logits = h_cls @ self.W_cls3 + self.b_cls3
        
        if return_trajectory:
            return logits, trajectory, vertices
        return logits
    
    def predict(self, X):
        probabilities = softmax(self.forward(X))
        return np.argmax(probabilities, axis=1)
    
    def score(self, X, y):
        predictions = self.predict(X)
        return np.mean(predictions == y)


# ============================================================
# Training
# ============================================================
def train(model, X_train, y_train, X_test, y_test, epochs=10, batch_size=100, lr=0.001):
    num_batches = len(X_train) // batch_size
    
    train_losses = []
    train_accs = []
    test_accs = []
    
    for epoch in range(epochs):
        epoch_loss = 0.0
        correct = 0
        total = 0
        
        indices = np.random.permutation(len(X_train))
        
        for batch_idx in range(num_batches):
            batch_indices = indices[batch_idx * batch_size : (batch_idx + 1) * batch_size]
            X_batch = X_train[batch_indices]
            y_batch = y_train[batch_indices]
            
            # Forward
            logits = model.forward(X_batch)
            probabilities = softmax(logits)
            
            # Loss
            y_onehot = np.zeros((len(y_batch), model.num_classes))
            y_onehot[np.arange(len(y_batch)), y_batch] = 1.0
            loss = -np.sum(y_onehot * np.log(probabilities + 1e-9)) / len(y_batch)
            
            # Accuracy
            predictions = np.argmax(probabilities, axis=1)
            correct += np.sum(predictions == y_batch)
            total += len(y_batch)
            
            epoch_loss += loss
            
            if batch_idx % 300 == 0:
                acc = 100. * correct / total
                print(f'  Epoch {epoch+1} | Batch {batch_idx}/{num_batches} | '
                      f'Loss: {loss:.4f} | Acc: {acc:.2f}%')
        
        # Evaluate
        test_acc = model.score(X_test, y_test) * 100
        
        avg_loss = epoch_loss / num_batches
        train_acc = 100. * correct / total
        
        train_losses.append(avg_loss)
        train_accs.append(train_acc)
        test_accs.append(test_acc)
        
        print(f'Epoch {epoch+1}/{epochs} | Loss: {avg_loss:.4f} | '
              f'Train: {train_acc:.2f}% | Test: {test_acc:.2f}%')
        print('-' * 60)
    
    return train_losses, train_accs, test_accs


# ============================================================
# Visualization
# ============================================================
def visualize_vertices(model, X_test, y_test, num_samples=8):
    indices = np.arange(num_samples)
    X_samples = X_test[indices]
    y_samples = y_test[indices]
    
    logits, trajectory, initial_vertices = model.forward(X_samples, return_trajectory=True)
    predictions = np.argmax(softmax(logits), axis=1)
    final_vertices = trajectory[-1].reshape(num_samples, model.num_vertices, 3)
    
    fig = plt.figure(figsize=(16, 12))
    
    for i in range(min(num_samples, 8)):
        ax1 = fig.add_subplot(4, 6, i*2 + 1)
        ax1.imshow(X_samples[i].reshape(28, 28), cmap='gray')
        ax1.set_title(f'True: {y_samples[i]}\nPred: {predictions[i]}', fontsize=10)
        ax1.axis('off')
        
        ax2 = fig.add_subplot(4, 6, i*2 + 2, projection='3d')
        v_init = initial_vertices[i]
        ax2.scatter(v_init[:, 0], v_init[:, 1], v_init[:, 2], 
                   c=range(12), cmap='viridis', s=100, alpha=0.8)
        ax2.set_title('Initial Vertices', fontsize=8)
    
    # Evolution plot
    ax3 = fig.add_subplot(4, 6, (13, 14), projection='3d')
    sample_idx = 0
    colors = plt.cm.plasma(np.linspace(0, 1, 30))
    
    for t_idx in range(0, 30, 6):
        v_traj = trajectory[t_idx, sample_idx].reshape(12, 3)
        ax3.scatter(v_traj[:, 0], v_traj[:, 1], v_traj[:, 2], 
                   label=f't={t_idx*1.5/29:.2f}', s=50, alpha=0.8, c=[colors[t_idx]])
    ax3.legend()
    ax3.set_title('Vertex Evolution')
    
    # Distance over time
    ax4 = fig.add_subplot(4, 6, 15)
    for v in range(0, 12, 3):
        dists = [np.linalg.norm(trajectory[t_idx, sample_idx].reshape(12, 3)[v]) 
                 for t_idx in range(30)]
        ax4.plot(np.linspace(0, 1.5, 30), dists, label=f'Vertex {v}')
    ax4.set_xlabel('Time')
    ax4.set_ylabel('Distance from origin')
    ax4.set_title('Vertex Distance Over Time')
    ax4.legend(fontsize=8)
    ax4.grid(True)
    
    # Phase plot
    ax5 = fig.add_subplot(4, 6, 16, projection='3d')
    v0_traj = trajectory[:, sample_idx, :3]
    
    for i in range(29):
        ax5.plot(v0_traj[i:i+2, 0], v0_traj[i:i+2, 1], v0_traj[i:i+2, 2], 
                color=colors[i], linewidth=2)
    ax5.scatter(v0_traj[0, 0], v0_traj[0, 1], v0_traj[0, 2], 
               c='green', s=100, label='Start', marker='o')
    ax5.scatter(v0_traj[-1, 0], v0_traj[-1, 1], v0_traj[-1, 2], 
               c='red', s=100, label='End', marker='s')
    ax5.legend()
    ax5.set_title('Vertex 0 Trajectory')
    
    plt.tight_layout()
    plt.savefig('quantum_vertices_numpy.png', dpi=150, bbox_inches='tight')
    plt.show()


def visualize_training(train_losses, train_accs, test_accs):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    
    ax1.plot(train_losses, 'b-', linewidth=2)
    ax1.set_title('Training Loss', fontsize=14)
    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Loss')
    ax1.grid(True)
    
    ax2.plot(train_accs, 'b-', label='Train', linewidth=2)
    ax2.plot(test_accs, 'r-', label='Test', linewidth=2)
    ax2.set_title('Accuracy', fontsize=14)
    ax2.set_xlabel('Epoch')
    ax2.set_ylabel('Accuracy (%)')
    ax2.legend()
    ax2.grid(True)
    
    plt.tight_layout()
    plt.savefig('training_progress_numpy.png', dpi=150)
    plt.show()


# ============================================================
# Main
# ============================================================
if __name__ == "__main__":
    BATCH_SIZE = 100
    NUM_VERTICES = 12
    EPOCHS = 10
    NUM_CLASSES = 10
    T_SPAN = (0, 1.5)
    
    print("=" * 60)
    print("Quantum 3D Vertex Classifier (NumPy + solve_ivp)")
    print("=" * 60)
    print(f"Batch Size: {BATCH_SIZE}")
    print(f"Vertices per sample: {NUM_VERTICES}")
    print(f"t_span: {T_SPAN}")
    print("=" * 60)
    
    model = Quantum3DClassifier(
        input_size=784,
        num_vertices=NUM_VERTICES,
        num_classes=NUM_CLASSES,
        batch_size=BATCH_SIZE,
        t_span=T_SPAN,
        seed=42
    )
    
    print("\nModel created successfully!")
    print("=" * 60)
    
    print("\nStarting training...\n")
    train_losses, train_accs, test_accs = train(
        model, X_train, y_train, X_test, y_test,
        epochs=EPOCHS, batch_size=BATCH_SIZE, lr=0.001
    )
    
    print("\n" + "=" * 60)
    print("FINAL RESULTS")
    print("=" * 60)
    final_test_acc = model.score(X_test, y_test) * 100
    print(f"Final Test Accuracy: {final_test_acc:.2f}%")
    print("=" * 60)
    
    visualize_training(train_losses, train_accs, test_accs)
    visualize_vertices(model, X_test, y_test, num_samples=8)