import numpy as np
from scipy.io.wavfile import read
import matplotlib.pyplot as plt
from scipy.signal import spectrogram

# 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]

# Define the MLP Classifier
class MLPClassifier:
    def __init__(self, input_size, hidden_size, output_size, learning_rate=np.random.rand(4)):
        self.W1 = np.random.randn(input_size, hidden_size) * 0.01
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_size) * 0.01
        self.b2 = np.zeros((1, output_size))
        self.learning_rate = learning_rate
    
    def relu(self, x):
        return np.maximum(0, x)
    
    def relu_derivative(self, x):
        return np.where(x > 0, 1, 0)
    
    def softmax(self, 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 forward(self, X):
        self.z1 = np.dot(X, self.W1) + self.b1
        self.a1 = self.relu(self.z1)
        self.z2 = np.dot(self.a1, self.W2) + self.b2
        output = self.softmax(self.z2)
        return output
    
    def compute_loss(self, y_true, y_pred):
        m = y_true.shape[0]
        loss = -np.sum(y_true * np.log(y_pred + 1e-9)) / m
        return loss
    
    def backward(self, X, y_true, y_pred):
        m = y_true.shape[0]
        dz2 = y_pred - y_true
        dW2 = np.dot(self.a1.T, dz2) / m
        db2 = np.sum(dz2, axis=0, keepdims=True) / m
        da1 = np.dot(dz2, self.W2.T)
        dz1 = da1 * self.relu_derivative(self.z1)
        dW1 = np.dot(X.T, dz1) / m
        db1 = np.sum(dz1, axis=0, keepdims=True) / m
        return dW1, db1, dW2, db2
    
    def update(self, X, y_true):
        y_pred = self.forward(X)
        dW1, db1, dW2, db2 = self.backward(X, y_true, y_pred)
        self.W1 -= self.learning_rate[0] * dW1
        self.b1 -= self.learning_rate[1] * db1
        self.W2 -= self.learning_rate[2] * dW2
        self.b2 -= self.learning_rate[3] * db2
    
    def predict(self, X):
        probabilities = self.forward(X)
        return np.argmax(probabilities, axis=1)

    def score(self, X, y_true):
        y = self.predict(X)
        return np.mean(y == y_true)


# ----------------------------------------------------------------------
# Real-time spectrogram setup
# ----------------------------------------------------------------------
plt.ion()
fig, ax = plt.subplots(figsize=(10, 4))

ERROR_BUFFER_SIZE = 2048
error_buffer = np.zeros(ERROR_BUFFER_SIZE)
buffer_idx = 0
PLOT_INTERVAL = 50

# Initialize and train the MLP Classifier
if __name__ == "__main__":
    learning_rate = np.random.rand(4)
    f = MLPClassifier(input_size=784, hidden_size=100, output_size=10, learning_rate=learning_rate)

    i = 0
    try:
        while True:
            idx = np.random.randint(0, 60000, 100)
            X = X_train[idx]
            yt = y_train[idx]
            y_true_oh = np.eye(10)[yt]

            # Forward pass to get y_pred and compute the scalar error signal
            y_pred = f.forward(X)
            batch_acc = np.mean(np.argmax(y_pred, axis=1) == yt)

            # Error signal: mean absolute error over the batch
            # (treating the batch-averaged |y_true - y_pred| as a 1-D time signal)
            #err_sig = np.mean(np.abs(y_true_oh - y_pred))
            err_sig = np.mean((yt - f.predict(X))**2)
            error_buffer[buffer_idx % ERROR_BUFFER_SIZE] = err_sig
            buffer_idx += 1

            print(i, batch_acc)
            f.update(X, y_true_oh)
            i += 1

            # --- Real-time spectrogram update ---
            if i % PLOT_INTERVAL == 0 and buffer_idx >= 256:
                # Reconstruct chronological signal from the circular buffer
                if buffer_idx < ERROR_BUFFER_SIZE:
                    sig = error_buffer[:buffer_idx]
                else:
                    start = buffer_idx % ERROR_BUFFER_SIZE
                    sig = np.concatenate((error_buffer[start:], error_buffer[:start]))

                # Compute STFT-based spectrogram
                f_spec, t_spec, Sxx = spectrogram(
                    sig, fs=1.0, nperseg=256, noverlap=200, scaling='spectrum'
                )
                # Convert power to dB scale
                Sxx_dB = 10 * np.log10(Sxx + 1e-12)

                ax.clear()
                ax.pcolormesh(t_spec, f_spec, Sxx_dB, shading='gouraud', cmap='inferno')
                ax.set_ylabel('Frequency [cycles/iter]')
                ax.set_xlabel('Time [iterations]')
                ax.set_title('Error Signal (y_true − y_pred) Spectrogram')
                ax.set_ylim(0, 0.5)          # Nyquist limit for fs = 1.0
                ax.set_xlim(t_spec.min(), t_spec.max())
                plt.tight_layout()
                plt.pause(0.001)

    except KeyboardInterrupt:
        print("\nTraining stopped by user.")
        plt.ioff()
        plt.show()