import numpy as np
from scipy.io.wavfile import read
#from ran_delayed_collapse import RAN  # import the provided RAN class
import random
from math import gcd
from numbers import Real

class RAN:
    """
    Rational-Addition Number: a/b + c.
    
    Parameters a, b, c may be ints or floats. The value is not collapsed
    to a single float until .collapse() is called.
    """

    def __init__(self, a, b, c):
        if b == 0:
            raise ValueError("b cannot be zero")
        self.a = a
        self.b = b
        self.c = c

    def collapse(self):
        """Force evaluation to a single number."""
        return self.a / self.b + self.c

    def __repr__(self):
        return f"RAN({self.a!r}, {self.b!r}, {self.c!r})"

    def __str__(self):
        return f"({self.a}/{self.b} + {self.c})"

    def __float__(self):
        return self.collapse()

    def __add__(self, other):
        if isinstance(other, Real):
            other = RAN(0, 1, other)
        if not isinstance(other, RAN):
            return NotImplemented
        a, b, c = self.a, self.b, self.c
        d, e, f = other.a, other.b, other.c
        # (a/b + c) + (d/e + f) = (ae + bd)/(be) + (c + f)
        return RAN(a * e + b * d, b * e, c + f)

    __radd__ = __add__

    def __sub__(self, other):
        if isinstance(other, Real):
            other = RAN(0, 1, other)
        if not isinstance(other, RAN):
            return NotImplemented
        return self + RAN(-other.a, other.b, -other.c)

    def __rsub__(self, other):
        if isinstance(other, Real):
            other = RAN(0, 1, other)
        return other - self

    def __neg__(self):
        return RAN(-self.a, self.b, -self.c)

    def __mul__(self, other):
        if isinstance(other, Real):
            other = RAN(0, 1, other)
        if not isinstance(other, RAN):
            return NotImplemented
        a, b, c = self.a, self.b, self.c
        d, e, f = other.a, other.b, other.c
        # (a/b + c)(d/e + f) = (ad + afe + cbd)/(be) + cf
        return RAN(a * d + a * f * e + c * b * d, b * e, c * f)

    __rmul__ = __mul__

    def __truediv__(self, other):
        if isinstance(other, Real):
            other = RAN(0, 1, other)
        if not isinstance(other, RAN):
            return NotImplemented
        a, b, c = self.a, self.b, self.c
        d, e, f = other.a, other.b, other.c
        # (a/b + c) / (d/e + f) = e(a + bc) / (b(d + ef))
        return RAN(e * (a + b * c), b * (d + e * f), 0)

    def __rtruediv__(self, other):
        if isinstance(other, Real):
            other = RAN(0, 1, other)
        return other / self

    def __eq__(self, other):
        if isinstance(other, RAN):
            return self.collapse() == other.collapse()
        if isinstance(other, Real):
            return self.collapse() == other
        return NotImplemented

    def simplify(self):
        """
        Partial collapse: if a and b are integers, reduce the fraction
        and move any whole part into c. This is still exact for integer
        parameters.
        """
        if isinstance(self.a, int) and isinstance(self.b, int):
            g = gcd(self.a, self.b)
            a = self.a // g
            b = self.b // g
            if b < 0:
                a, b = -a, -b
            q = a // b
            r = a - q * b
            return RAN(r, b, self.c + q)
        return self


# ------------------------------------------------------------
# Load data (unchanged from original)
# ------------------------------------------------------------
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]

# ------------------------------------------------------------
# MLP Classifier using RAN numbers for all parameters
# ------------------------------------------------------------
class MLPClassifierRAN:
    def __init__(self, input_size, hidden_size, output_size, learning_rate=0.01):
        # All weights and biases are stored as RAN objects.
        # We initialise them as RAN(random, 1, 0) so that collapse() gives the random value.
        self.W1 = [[RAN(random.gauss(0, 0.01), 1, 0) for _ in range(hidden_size)] for _ in range(input_size)]
        self.b1 = [RAN(0, 1, 0) for _ in range(hidden_size)]
        self.W2 = [[RAN(random.gauss(0, 0.01), 1, 0) for _ in range(output_size)] for _ in range(hidden_size)]
        self.b2 = [RAN(0, 1, 0) for _ in range(output_size)]
        # Learning rate is also a RAN (kept as a scalar value with zero fraction)
        self.lr = RAN(learning_rate, 1, 0)

    # --------------------------------------------------------
    # Helper: element‑wise ReLU on a list of RANs.
    # We must collapse to decide if the value is > 0.
    # --------------------------------------------------------
    def relu(self, z_list):
        out = []
        for z in z_list:
            if z.collapse() > 0:
                out.append(z)          # keep RAN structure for positive part
            else:
                out.append(RAN(0, 1, 0)) # zero as RAN(0,1,0)
        return out

    def relu_derivative(self, z_list):
        return [1 if z.collapse() > 0 else 0 for z in z_list]

    # --------------------------------------------------------
    # Softmax: requires collapsing to a float for exp and sum.
    # --------------------------------------------------------
    def softmax(self, logits):
        # collapse all logits to floats
        vals = [logit.collapse() for logit in logits]
        exp_vals = np.exp(vals - np.max(vals))
        probs = exp_vals / np.sum(exp_vals)
        return probs

    # --------------------------------------------------------
    # Forward pass: input X is a list of floats (one sample).
    # Returns output probabilities (floats) and intermediate RANs.
    # --------------------------------------------------------
    def forward(self, X):
        # z1 = X * W1 + b1   (all RAN operations)
        z1 = []
        for j in range(len(self.b1)):
            s = RAN(0, 1, 0)  # zero RAN
            for i, x in enumerate(X):
                # x is float; RAN.__rmul__ handles float * RAN
                s = s + (x * self.W1[i][j])
            s = s + self.b1[j]
            z1.append(s)

        a1 = self.relu(z1)

        # z2 = a1 * W2 + b2
        z2 = []
        for k in range(len(self.b2)):
            s = RAN(0, 1, 0)
            for j, a in enumerate(a1):
                s = s + (a * self.W2[j][k])
            s = s + self.b2[k]
            z2.append(s)

        # output probabilities (collapse to float for softmax)
        probs = self.softmax(z2)
        return probs, z1, a1, z2

    # --------------------------------------------------------
    # Loss (cross‑entropy) – uses collapsed probabilities
    # --------------------------------------------------------
    def compute_loss(self, y_true, y_pred):
        # y_pred is list of floats (probabilities)
        return -np.sum(y_true * np.log(y_pred + 1e-9))

    # --------------------------------------------------------
    # Backward pass using RAN arithmetic where possible.
    # Gradients are computed as RAN objects.
    # --------------------------------------------------------
    def backward(self, X, y_true, probs, z1, a1, z2):
        m = 1  # we process one sample at a time (or mini‑batch)
        # dz2 = probs - y_true  (probs are floats; we convert to RAN)
        dz2 = [RAN(probs[k] - y_true[k], 1, 0) for k in range(len(probs))]

        # dW2 = a1^T * dz2  (outer product, each term is RAN)
        dW2 = [[RAN(0, 1, 0) for _ in range(len(dz2))] for _ in range(len(a1))]
        for j in range(len(a1)):
            for k in range(len(dz2)):
                dW2[j][k] = a1[j] * dz2[k]  # RAN multiplication

        # db2 = dz2 (sum over mini‑batch, but m=1 here)
        db2 = dz2[:]

        # da1 = dz2 * W2^T
        da1 = []
        for j in range(len(a1)):
            s = RAN(0, 1, 0)
            for k in range(len(dz2)):
                s = s + (dz2[k] * self.W2[j][k])
            da1.append(s)

        # dz1 = da1 * relu_derivative(z1)
        deriv = self.relu_derivative(z1)
        dz1 = [da1[j] * RAN(deriv[j], 1, 0) for j in range(len(da1))]

        # dW1 = X^T * dz1
        dW1 = [[RAN(0, 1, 0) for _ in range(len(dz1))] for _ in range(len(X))]
        for i in range(len(X)):
            for j in range(len(dz1)):
                dW1[i][j] = RAN(X[i], 1, 0) * dz1[j]  # X[i] is float, convert to RAN

        # db1 = dz1
        db1 = dz1[:]

        return dW1, db1, dW2, db2

    # --------------------------------------------------------
    # Update parameters: subtract learning_rate * gradient
    # (all operations with RAN)
    # --------------------------------------------------------
    def update(self, X, y_true):
        # Forward pass
        probs, z1, a1, z2 = self.forward(X)
        # Backward
        dW1, db1, dW2, db2 = self.backward(X, y_true, probs, z1, a1, z2)

        # Update W1
        for i in range(len(self.W1)):
            for j in range(len(self.W1[i])):
                self.W1[i][j] = self.W1[i][j] - self.lr * dW1[i][j]
        # Update b1
        for j in range(len(self.b1)):
            self.b1[j] = self.b1[j] - self.lr * db1[j]
        # Update W2
        for j in range(len(self.W2)):
            for k in range(len(self.W2[j])):
                self.W2[j][k] = self.W2[j][k] - self.lr * dW2[j][k]
        # Update b2
        for k in range(len(self.b2)):
            self.b2[k] = self.b2[k] - self.lr * db2[k]

        # Return loss for monitoring (collapsed to float)
        loss = self.compute_loss(y_true, probs)
        return loss

    # --------------------------------------------------------
    # Prediction: collapse probabilities to class index
    # --------------------------------------------------------
    def predict(self, X):
        probs, _, _, _ = self.forward(X)
        return np.argmax(probs)

    def score(self, X, y_true):
        correct = 0
        for x, y in zip(X, y_true):
            pred = self.predict(x)
            if pred == y:
                correct += 1
        return correct / len(X)

# ------------------------------------------------------------
# Training loop (single‑sample stochastic gradient descent)
# ------------------------------------------------------------
if __name__ == "__main__":
    # Use a fixed learning rate as a RAN object (or just a float)
    mlp = MLPClassifierRAN(input_size=784, hidden_size=100, output_size=10, learning_rate=0.01)

    epochs = 5
    batch_size = 100   # we simulate mini‑batch by averaging? For simplicity we keep SGD

    for epoch in range(epochs):
        # Shuffle training data
        idx = np.random.permutation(60000)
        X_shuffled = X_train[idx]
        y_shuffled = y_train[idx]

        total_loss = 0
        for i in range(0, 60000, batch_size):
            X_batch = X_shuffled[i:i+batch_size]
            y_batch = y_shuffled[i:i+batch_size]

            # Process each sample individually (or we could implement batch gradient)
            for x, yt in zip(X_batch, y_batch):
                # one‑hot encoding
                y_true = np.eye(10)[yt]
                loss = mlp.update(x.tolist(), y_true)
                total_loss += loss

        # Evaluate on test set
        acc = mlp.score(X20, yt20)
        print(f"Epoch {epoch+1}: average loss = {total_loss/60000:.6f}, test accuracy = {acc:.4f}")