import numpy as np

# Linear component
x = np.linspace(0, 1, 100)

# 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 LearnedPolyvalMLP:
    """
    MLP where both polynomial coefficients (c) and inputs (x) are learned.
    np.polyval(learned_coeffs, learned_x)
    """
    def __init__(self, input_size, layer_sizes, poly_degree=3, learning_rate=None):
        """
        input_size: input dimensions (e.g., 784)
        layer_sizes: list of layer sizes [h1, h2, ..., output]
        poly_degree: polynomial degree
        """
        self.input_size = input_size
        self.layer_sizes = layer_sizes
        self.poly_degree = poly_degree
        self.num_layers = len(layer_sizes)
        
        # Initialize learned polynomial coefficients for each layer
        # Shape: (num_layers, output_size, input_size, poly_degree + 1)
        self.learned_coeffs = []  # c parameters
        self.learned_bias = []    # additional bias
        
        for i, layer_size in enumerate(layer_sizes):
            # Coefficients for each neuron in this layer
            # Initialize to near-zero polynomials
            coeffs = np.random.randn(input_size, layer_size, poly_degree + 1) * 0.01
            # Set constant term bias effect
            coeffs[:, :, 0] = 0.0  # constant term
            self.learned_coeffs.append(coeffs)
            
            bias = np.zeros((1, layer_size))
            self.learned_bias.append(bias)
        
        # Learned input transformation (x parameter)
        self.learned_x_weights = np.random.randn(input_size, input_size) * 0.01
        self.learned_x_bias = np.zeros((1, input_size))
        
        # Learning rates (separate for coeffs, x_weights, x_bias, layer_bias)
        if learning_rate is None:
            self.lr_coeffs = 1e-5    # Learning rate for polynomial coefficients
            self.lr_x = 1e-4         # Learning rate for learned x
            self.lr_bias = 0.01      # Learning rate for biases
        else:
            self.lr_coeffs = learning_rate[0]
            self.lr_x = learning_rate[1]
            self.lr_bias = learning_rate[2]
        
        self.x = x  # Original x for reference
    
    def polyval_layer(self, X, coeffs):
        """
        Evaluate polynomial for each output neuron.
        X: (batch_size, input_size)
        coeffs: (input_size, output_size, degree + 1)
        
        Returns: (batch_size, output_size)
        """
        batch_size = X.shape[0]
        output_size = coeffs.shape[1]
        
        # For each output neuron, compute sum over input dimensions
        # polyval(coefficients[:, j], X[:, i]) for each i, j
        result = np.zeros((batch_size, output_size))
        
        for out_idx in range(output_size):
            for in_idx in range(input_size):
                # np.polyval for single input across batch
                # coeffs[input_idx, output_idx, :] is the polynomial
                result[:, out_idx] += np.polyval(coeffs[in_idx, out_idx, :], X[:, in_idx])
        
        return result
    
    def forward(self, X):
        """
        Forward pass using learned polyval.
        np.polyval(learned_coeffs, learned_x)
        """
        self.cache_x = []
        self.cache_coeffs = []
        self.cache_Z = []
        
        # Learn the input x transformation
        # learned_x = X @ learned_x_weights + learned_x_bias
        self.learned_x = np.dot(X, self.learned_x_weights) + self.learned_x_bias
        self.cache_x.append(self.learned_x)
        
        A = X
        
        for layer_idx in range(self.num_layers):
            # Apply learned polynomial transformation
            # output = np.polyval(self.learned_coeffs[layer_idx], self.learned_x)
            Z = self.polyval_layer(A, self.learned_coeffs[layer_idx])
            
            # Add bias
            Z = Z + self.learned_bias[layer_idx]
            self.cache_Z.append(Z)
            self.cache_coeffs.append(self.learned_coeffs[layer_idx].copy())
            
            # Apply activation (ReLU for hidden, Softmax for output)
            if layer_idx == self.num_layers - 1:
                # Softmax for output
                exp_Z = np.exp(Z - np.max(Z, axis=1, keepdims=True))
                A = exp_Z / np.sum(exp_Z, axis=1, keepdims=True)
            else:
                # ReLU for hidden layers
                A = np.maximum(0, Z)
        
        return A
    
    def relu_derivative(self, Z):
        return np.where(Z > 0, 1, 0)
    
    def backward(self, X, y_true, y_pred):
        """
        Backward pass to update:
        1. Polynomial coefficients (learned_coeffs)
        2. Learned x weights and bias
        3. Layer biases
        """
        batch_size = X.shape[0]
        
        # Output layer gradient
        dZ = y_pred - y_true
        d_bias = [np.sum(dZ, axis=0, keepdims=True) / batch_size]
        
        # Gradient w.r.t. input to output layer
        dA = np.dot(dZ, self.learned_coeffs[-1].transpose(1, 0, 2).reshape(self.layer_sizes[-1], -1))
        
        # Gradients for polynomial coefficients
        d_coeffs = []
        d_z_input = dZ.copy()
        
        # Backprop through layers
        for layer_idx in reversed(range(self.num_layers)):
            dZ_layer = d_z_input * (self.cache_Z[layer_idx] > 0 if layer_idx < self.num_layers - 1 else 1)
            
            # Gradient w.r.t. coefficients for this layer
            # d(polyval)/d(coeff_k) = x^k
            coeff_grad = np.zeros_like(self.learned_coeffs[layer_idx])
            
            A_prev = X if layer_idx == 0 else self.cache_Z[layer_idx - 1]
            
            for out_idx in range(self.layer_sizes[layer_idx]):
                for in_idx in range(self.input_size):
                    # dZ/d(coeff) = x^degree for each degree
                    for d_idx in range(self.poly_degree + 1):
                        coeff_grad[in_idx, out_idx, d_idx] = np.mean(
                            A_prev[:, in_idx] ** d_idx * dZ_layer[:, out_idx]
                        )
            
            d_coeffs.insert(0, coeff_grad)
            
            # Gradient w.r.t. biases
            d_bias.insert(0, np.sum(dZ_layer, axis=0, keepdims=True) / batch_size)
            
            # Pass gradient to previous layer
            if layer_idx > 0:
                dA = np.dot(dZ_layer, self.learned_coeffs[layer_idx - 1].reshape(
                    self.layer_sizes[layer_idx - 1], -1
                ).T)
        
        # Gradient w.r.t. learned_x
        d_learned_x = np.dot(dZ, self.learned_coeffs[-1].mean(axis=0).T)
        
        # Gradient w.r.t. learned_x_weights and learned_x_bias
        d_x_weights = np.dot(X.T, d_learned_x) / batch_size
        d_x_bias = np.sum(d_learned_x, axis=0, keepdims=True) / batch_size
        
        return d_coeffs, d_x_weights, d_x_bias, d_bias
    
    def update(self, X, y_true):
        """Update all learned parameters"""
        y_pred = self.forward(X)
        d_coeffs, d_x_weights, d_x_bias, d_bias = self.backward(X, y_true, y_pred)
        
        # Update polynomial coefficients
        for i in range(self.num_layers):
            self.learned_coeffs[i] -= self.lr_coeffs * d_coeffs[i]
        
        # Update learned x weights
        self.learned_x_weights -= self.lr_x * d_x_weights
        self.learned_x_bias -= self.lr_x * d_x_bias
        
        # Update layer biases
        for i in range(self.num_layers):
            self.learned_bias[i] -= self.lr_bias * d_bias[i]
    
    def predict(self, X):
        return np.argmax(self.forward(X), axis=1)
    
    def score(self, X, y_true):
        return np.mean(self.predict(X) == y_true)


# Alternative simplified version with clearer polyval structure
class SimplePolyvalModel:
    """
    Simple model: np.polyval(learned_coeffs, learned_x)
    where learned_x = f(X, W_x) and learned_coeffs = g(W_c)
    """
    def __init__(self, input_dim, hidden_dim, output_dim, poly_degree=3):
        self.poly_degree = poly_degree
        
        # Learned coefficients (c) - polynomial coefficients to be learned
        # Shape: (hidden_dim, poly_degree + 1)
        self.coeffs = np.zeros((hidden_dim, poly_degree + 1))
        # Initialize with some structure
        self.coeffs[:, 0] = np.random.randn(hidden_dim) * 0.01  # constant
        self.coeffs[:, 1] = np.random.randn(hidden_dim) * 0.01  # linear
        self.coeffs[:, 2:] = np.random.randn(hidden_dim, poly_degree - 1) * 0.001
        
        # Learned x transformation parameters
        # x_transformed = W_x @ X + b_x
        self.W_x = np.random.randn(input_dim, input_dim) * 0.01
        self.b_x = np.zeros((1, input_dim))
        
        # Output layer
        self.W_out = np.random.randn(hidden_dim, output_dim) * 0.01
        self.b_out = np.zeros((1, output_dim))
        
        # Learning rates
        self.lr_coeffs = 1e-5
        self.lr_x = 1e-4
        self.lr_out = 0.01
    
    def forward(self, X):
        """
        Forward pass:
        1. learned_x = X @ W_x + b_x
        2. hidden = np.polyval(coeffs, learned_x)  -- vectorized
        3. output = hidden @ W_out + b_out
        """
        self.X = X
        
        # Step 1: Learn the x input
        self.learned_x = np.dot(X, self.W_x) + self.b_x  # (batch, input_dim)
        
        # Step 2: Evaluate polynomial with learned coefficients
        # self.learned_x shape: (batch, input_dim)
        # We want output: (batch, hidden_dim) using polyval
        batch_size = X.shape[0]
        
        # For each hidden neuron, sum polyval over input dimensions
        # hidden_j = sum_i polyval(coeffs[j], x_i)
        self.hidden = np.zeros((batch_size, self.coeffs.shape[0]))
        
        for h_idx in range(self.coeffs.shape[0]):
            for i_idx in range(self.learned_x.shape[1]):
                # Evaluate polynomial coeffs[h_idx] at x_i
                self.hidden[:, h_idx] += np.polyval(self.coeffs[h_idx], self.learned_x[:, i_idx])
        
        # Apply activation
        self.hidden_activated = np.maximum(0, self.hidden)
        
        # Step 3: Output layer
        output = np.dot(self.hidden_activated, self.W_out) + self.b_out
        
        # Softmax
        exp_out = np.exp(output - np.max(output, axis=1, keepdims=True))
        self.output = exp_out / np.sum(exp_out, axis=1, keepdims=True)
        
        return self.output
    
    def backward(self, y_true):
        """Backward pass for all learned parameters"""
        batch_size = self.X.shape[0]
        m = batch_size
        
        # Output gradient
        d_out = self.output - y_true
        
        # Output layer gradients
        d_W_out = np.dot(self.hidden_activated.T, d_out) / m
        d_b_out = np.sum(d_out, axis=0, keepdims=True) / m
        
        # Gradient w.r.t. hidden (pre-activation)
        d_hidden = np.dot(d_out, self.W_out.T)
        d_hidden *= (self.hidden > 0)  # ReLU derivative
        
        # Gradient w.r.t. learned coefficients
        d_coeffs = np.zeros_like(self.coeffs)
        for h_idx in range(self.coeffs.shape[0]):
            for i_idx in range(self.learned_x.shape[1]):
                x_vals = self.learned_x[:, i_idx]
                for d_idx in range(self.poly_degree + 1):
                    # d(polyval)/d(coef_d) = x^d
                    d_coeffs[h_idx, d_idx] += np.sum(
                        x_vals ** d_idx * d_hidden[:, h_idx]
                    ) / m
        
        # Gradient w.r.t. learned_x
        d_learned_x = np.zeros_like(self.learned_x)
        for h_idx in range(self.coeffs.shape[0]):
            for i_idx in range(self.learned_x.shape[1]):
                # d(polyval)/d(x) = sum_d d * coef_d * x^(d-1)
                x_vals = self.learned_x[:, i_idx]
                d_poly_dx = np.zeros(batch_size)
                for d_idx in range(1, self.poly_degree + 1):
                    d_poly_dx += d_idx * self.coeffs[h_idx, d_idx] * x_vals ** (d_idx - 1)
                d_learned_x[:, i_idx] += d_poly_dx * d_hidden[:, h_idx]
        
        # Gradient w.r.t. W_x and b_x
        d_W_x = np.dot(self.X.T, d_learned_x) / m
        d_b_x = np.sum(d_learned_x, axis=0, keepdims=True) / m
        
        return d_coeffs, d_W_x, d_b_x, d_W_out, d_b_out
    
    def update(self, X, y_true):
        """Update all parameters"""
        self.forward(X)
        d_coeffs, d_W_x, d_b_x, d_W_out, d_b_out = self.backward(y_true)
        
        # Update polynomial coefficients
        self.coeffs -= self.lr_coeffs * d_coeffs
        
        # Update learned x weights
        self.W_x -= self.lr_x * d_W_x
        self.b_x -= self.lr_x * d_b_x
        
        # Update output layer
        self.W_out -= self.lr_out * d_W_out
        self.b_out -= self.lr_out * d_b_out
    
    def predict(self, X):
        return np.argmax(self.forward(X), axis=1)
    
    def score(self, X, y_true):
        return np.mean(self.predict(X) == y_true)


# More compact vectorized version
class VectorizedPolyvalNet:
    """
    Vectorized version using np.polyval for learned c and x.
    Similar to: np.polyval([0.1, 1.3, 3.1], np.random.rand(100, 100))
    """
    def __init__(self, input_dim, output_dim, poly_degree=3):
        # Polynomial coefficients (c) - learned
        # Shape: (poly_degree + 1,) for each output feature
        self.c = np.zeros((output_dim, poly_degree + 1))
        self.c[:, 1] = 0.1  # Initialize linear term
        self.c[:, 2] = 1.3  # Initialize quadratic term
        self.c[:, 3] = 3.1  # Initialize cubic term
        
        # Input transformation weights (for x)
        self.W = np.random.randn(input_dim, output_dim) * 0.01
        self.b = np.zeros((1, output_dim))
        
        self.poly_degree = poly_degree
        self.lr_c = 1e-5
        self.lr_W = 0.01
    
    def forward(self, X):
        """
        Forward: output = np.polyval(c, x)
        where x = X @ W + b
        """
        # Learned x
        self.x = np.dot(X, self.W) + self.b  # (batch, output_dim)
        
        # Evaluate polynomial: np.polyval(c, x)
        # For each output dim, apply polyval
        self.output = np.zeros_like(self.x)
        for out_idx in range(self.x.shape[1]):
            self.output[:, out_idx] = np.polyval(self.c[out_idx], self.x[:, out_idx])
        
        # Softmax
        exp_out = np.exp(self.output - np.max(self.output, axis=1, keepdims=True))
        self.probs = exp_out / np.sum(exp_out, axis=1, keepdims=True)
        
        return self.probs
    
    def backward(self, y_true):
        """Backprop for learned c and x"""
        batch_size = self.x.shape[0]
        
        # Output gradient
        d_out = self.probs - y_true
        
        # Gradient w.r.t. polynomial coefficients (c)
        d_c = np.zeros_like(self.c)
        for out_idx in range(self.x.shape[1]):
            for d_idx in range(self.poly_degree + 1):
                d_c[out_idx, d_idx] = np.mean(
                    self.x[:, out_idx] ** d_idx * d_out[:, out_idx]
                )
        
        # Gradient w.r.t. x
        d_x = np.zeros_like(self.x)
        for out_idx in range(self.x.shape[1]):
            # d(polyval)/dx = sum d*d_idx * c[d_idx] * x^(d_idx-1)
            d_poly_dx = np.zeros(batch_size)
            for d_idx in range(1, self.poly_degree + 1):
                d_poly_dx += d_idx * self.c[out_idx, d_idx] * self.x[:, out_idx] ** (d_idx - 1)
            d_x[:, out_idx] = d_poly_dx * d_out[:, out_idx]
        
        # Gradient w.r.t. W and b
        d_W = np.dot(self.x.T, d_x) / batch_size  # Wait, x depends on W, need proper gradient
        
        # Actually: x = X @ W + b, so d_x/d_W = X
        d_W = np.dot(self.x.T, d_x) / batch_size
        d_b = np.sum(d_x, axis=0, keepdims=True) / batch_size
        
        return d_c, d_W, d_b
    
    def update(self, X, y_true):
        self.forward(X)
        d_c, d_W, d_b = self.backward(y_true)
        
        self.c -= self.lr_c * d_c
        self.W -= self.lr_W * d_W
        self.b -= self.lr_W * d_b
    
    def predict(self, X):
        return np.argmax(self.forward(X), axis=1)
    
    def score(self, X, y_true):
        return np.mean(self.predict(X) == y_true)


# Initialize and train
print("Training SimplePolyvalModel...")
f = SimplePolyvalModel(input_dim=784, hidden_dim=100, output_dim=10, poly_degree=3)

i = 0
while True:
    idx = np.random.randint(0, 60000, 100)
    X = X_train[idx]
    yt = y_train[idx]
    f.update(X, np.eye(10)[yt])
    
    if i % 10 == 0:
        print(f"Step {i}, Accuracy: {f.score(X20, yt20):.4f}")
    i += 1