import torch
import torchvision
import torchvision.transforms as transforms
import math
import random
from typing import List, Tuple, Union

# -------------------------------------------------------------------
# 1. nRAN class (from nran-theory-implementation.md, slightly adapted)
# -------------------------------------------------------------------
Number = Union[int, float]

class nRAN:
    """Multi-term Rational-Addition Number: sum(a_i/b_i) + c, with capacity."""

    def __init__(self, terms: List[Tuple[Number, Number]] = None,
                 c: Number = 0, max_n: int = None):
        self.terms = list(terms) if terms else []
        self.c = c
        self.max_n = max_n
        for a, b in self.terms:
            if b == 0:
                raise ValueError("b cannot be zero")
        if max_n is not None and len(self.terms) > max_n:
            self._compact()

    def collapse(self) -> float:
        """Evaluate to a single float."""
        return sum(a / b for a, b in self.terms) + self.c

    def __repr__(self):
        return f"nRAN({self.terms}, c={self.c!r})"

    def __str__(self):
        if not self.terms:
            return str(self.c)
        parts = " + ".join(f"{a}/{b}" for a, b in self.terms)
        return f"{parts} + {self.c}"

    def _combine_rational_terms(self):
        """Combine all integer terms into one exact fraction (using Fraction)."""
        from fractions import Fraction
        int_terms = [(a, b) for a, b in self.terms
                     if isinstance(a, int) and isinstance(b, int)]
        float_terms = [(a, b) for a, b in self.terms
                       if not (isinstance(a, int) and isinstance(b, int))]
        if not int_terms:
            return float_terms
        total = sum(Fraction(a, b) for a, b in int_terms)
        return [(total.numerator, total.denominator)] + float_terms

    def _compact(self):
        """Reduce term count to respect max_n."""
        if self.max_n is None:
            return
        self.terms = self._combine_rational_terms()
        if len(self.terms) > self.max_n:
            rational_value = sum(a / b for a, b in self.terms)
            self.terms = []
            self.c += rational_value

    def _after_op(self, result):
        if self.max_n is not None:
            result._compact()
        return result

    def __add__(self, other):
        if isinstance(other, (int, float)):
            return nRAN(self.terms, self.c + other, self.max_n)
        if not isinstance(other, nRAN):
            return NotImplemented
        return self._after_op(
            nRAN(self.terms + other.terms, self.c + other.c, self.max_n)
        )

    def __sub__(self, other):
        if isinstance(other, (int, float)):
            return nRAN(self.terms, self.c - other, self.max_n)
        if not isinstance(other, nRAN):
            return NotImplemented
        neg_terms = [(-a, b) for a, b in other.terms]
        return self._after_op(
            nRAN(self.terms + neg_terms, self.c - other.c, self.max_n)
        )

    def __mul__(self, other):
        if isinstance(other, (int, float)):
            new_terms = [(a * other, b) for a, b in self.terms]
            return nRAN(new_terms, self.c * other, self.max_n)
        if not isinstance(other, nRAN):
            return NotImplemented
        # Full multiplication
        new_terms = []
        for a, b in self.terms:
            for d, e in other.terms:
                new_terms.append((a * d, b * e))
        for a, b in self.terms:
            new_terms.append((a * other.c, b))
        for d, e in other.terms:
            new_terms.append((self.c * d, e))
        return self._after_op(
            nRAN(new_terms, self.c * other.c, self.max_n)
        )

    def __truediv__(self, other):
        if isinstance(other, (int, float)):
            new_terms = [(a, b * other) for a, b in self.terms]
            return nRAN(new_terms, self.c / other, self.max_n)
        if not isinstance(other, nRAN):
            return NotImplemented
        # Convert divisor to single fraction, then multiply by reciprocal
        combined = other._combine_rational_terms()
        if len(combined) != 1:
            return self * (1 / other.collapse())
        d, e = combined[0]
        denom = d + other.c * e
        reciprocal = nRAN([(e, denom)], 0, self.max_n)
        return self * reciprocal

    __radd__ = __add__
    __rmul__ = __mul__

    def __rsub__(self, other):
        return (-self) + other

    def __neg__(self):
        return nRAN([(-a, b) for a, b in self.terms], -self.c, self.max_n)


# -------------------------------------------------------------------
# 2. MLP with nRAN parameters
# -------------------------------------------------------------------
class nRAN_MLP:
    def __init__(self, input_size: int, hidden_size: int, output_size: int,
                 max_n: int = 10, init_scale: float = 0.01):
        """
        All weights and biases are nRAN objects.
        We store them as 2D lists of nRAN: W1[in][hidden], b1[hidden], etc.
        """
        self.max_n = max_n
        # W1: input_size x hidden_size
        self.W1 = [[nRAN([(random.gauss(0, init_scale), 1)], 0, max_n)
                    for _ in range(hidden_size)] for _ in range(input_size)]
        self.b1 = [nRAN([], 0, max_n) for _ in range(hidden_size)]  # zero offset

        # W2: hidden_size x output_size
        self.W2 = [[nRAN([(random.gauss(0, init_scale), 1)], 0, max_n)
                    for _ in range(output_size)] for _ in range(hidden_size)]
        self.b2 = [nRAN([], 0, max_n) for _ in range(output_size)]

    def _relu(self, z: nRAN) -> nRAN:
        """ReLU: keep z if its collapsed value > 0, else return zero nRAN."""
        if z.collapse() > 0:
            return z
        else:
            return nRAN([], 0, self.max_n)

    def _softmax(self, logits: List[nRAN]) -> List[float]:
        """Softmax over a list of nRAN logits (collapse to float first)."""
        vals = [z.collapse() for z in logits]
        maxv = max(vals)
        exp_vals = [math.exp(v - maxv) for v in vals]
        sum_exp = sum(exp_vals)
        return [e / sum_exp for e in exp_vals]

    def forward(self, x: List[float]) -> Tuple[List[float], List[nRAN], List[nRAN]]:
        """
        x: list of floats (input features).
        Returns:
          probs: softmax probabilities (floats)
          z1: pre-activation of hidden layer (list of nRAN)
          a1: activations after ReLU (list of nRAN)
          z2: pre-activation of output layer (list of nRAN)
        """
        # Hidden layer
        z1 = []
        for j in range(len(self.b1)):
            s = nRAN([], 0, self.max_n)  # zero
            for i, xi in enumerate(x):
                s = s + (xi * self.W1[i][j])   # xi is float, nRAN.__rmul__ works
            s = s + self.b1[j]
            z1.append(s)
        a1 = [self._relu(z) for z in z1]

        # Output layer
        z2 = []
        for k in range(len(self.b2)):
            s = nRAN([], 0, self.max_n)
            for j, aj in enumerate(a1):
                s = s + (aj * self.W2[j][k])
            s = s + self.b2[k]
            z2.append(s)

        probs = self._softmax(z2)
        return probs, z1, a1, z2

    def compute_loss(self, probs: List[float], y_true: int) -> float:
        """Cross-entropy loss for one sample."""
        return -math.log(probs[y_true] + 1e-12)

    def backward(self, x: List[float], y_true: int, probs: List[float],
                 z1: List[nRAN], a1: List[nRAN], z2: List[nRAN]):
        """
        Compute gradients (floats) for all parameters for one sample.
        Returns dictionaries of gradients: dW1, db1, dW2, db2.
        """
        # dz2 = probs - one_hot
        dz2 = [probs[k] - (1.0 if k == y_true else 0.0) for k in range(len(probs))]

        # dW2 = a1^T * dz2  (outer product)
        dW2 = [[0.0 for _ in range(len(dz2))] for _ in range(len(a1))]
        for j, aj in enumerate(a1):
            for k, dz in enumerate(dz2):
                dW2[j][k] = aj.collapse() * dz

        # db2 = dz2
        db2 = dz2[:]

        # da1 = dz2 * W2^T
        da1 = [0.0 for _ in range(len(a1))]
        for j in range(len(a1)):
            s = 0.0
            for k, dz in enumerate(dz2):
                s += dz * self.W2[j][k].collapse()
            da1[j] = s

        # dz1 = da1 * relu_derivative(z1)
        dz1 = []
        for j, z in enumerate(z1):
            dz = da1[j] if z.collapse() > 0 else 0.0
            dz1.append(dz)

        # dW1 = x^T * dz1
        dW1 = [[0.0 for _ in range(len(dz1))] for _ in range(len(x))]
        for i, xi in enumerate(x):
            for j, dz in enumerate(dz1):
                dW1[i][j] = xi * dz

        # db1 = dz1
        db1 = dz1[:]

        return dW1, db1, dW2, db2

    def update(self, dW1, db1, dW2, db2, lr: float):
        """Update all parameters using the accumulated gradients (float) and learning rate."""
        # Update W1
        for i in range(len(self.W1)):
            for j in range(len(self.W1[i])):
                # param = param - lr * grad
                self.W1[i][j] = self.W1[i][j] - (lr * dW1[i][j])
        # Update b1
        for j in range(len(self.b1)):
            self.b1[j] = self.b1[j] - (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] - (lr * dW2[j][k])
        # Update b2
        for k in range(len(self.b2)):
            self.b2[k] = self.b2[k] - (lr * db2[k])

    def predict(self, x: List[float]) -> int:
        probs, _, _, _ = self.forward(x)
        return max(range(len(probs)), key=lambda i: probs[i])

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


# -------------------------------------------------------------------
# 3. Training on MNIST using PyTorch DataLoader
# -------------------------------------------------------------------
def main():
    # Load MNIST
    transform = transforms.Compose([transforms.ToTensor(),
                                    transforms.Normalize((0.1307,), (0.3081,))])
    train_set = torchvision.datasets.MNIST(root='../data', train=True,
                                           download=True, transform=transform)
    test_set = torchvision.datasets.MNIST(root='../data', train=False,
                                          download=True, transform=transform)
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_set, batch_size=1000, shuffle=False)

    # Create model
    model = nRAN_MLP(input_size=784, hidden_size=128, output_size=10,
                     max_n=8, init_scale=0.01)

    lr = 0.01
    epochs = 3

    for epoch in range(epochs):
        total_loss = 0.0
        for batch_idx, (data, targets) in enumerate(train_loader):
            # Flatten images to lists of floats
            X_batch = data.view(data.size(0), -1).tolist()  # list of lists
            y_batch = targets.tolist()

            # Initialize gradient accumulators
            dW1_acc = [[0.0 for _ in range(128)] for _ in range(784)]
            db1_acc = [0.0 for _ in range(128)]
            dW2_acc = [[0.0 for _ in range(10)] for _ in range(128)]
            db2_acc = [0.0 for _ in range(10)]

            # Accumulate gradients over the batch
            for x, y in zip(X_batch, y_batch):
                probs, z1, a1, z2 = model.forward(x)
                loss = model.compute_loss(probs, y)
                total_loss += loss
                dW1, db1, dW2, db2 = model.backward(x, y, probs, z1, a1, z2)
                # Accumulate
                for i in range(784):
                    for j in range(128):
                        dW1_acc[i][j] += dW1[i][j]
                for j in range(128):
                    db1_acc[j] += db1[j]
                for j in range(128):
                    for k in range(10):
                        dW2_acc[j][k] += dW2[j][k]
                for k in range(10):
                    db2_acc[k] += db2[k]

            # Average gradients over batch and update
            batch_size = len(X_batch)
            for i in range(784):
                for j in range(128):
                    dW1_acc[i][j] /= batch_size
            for j in range(128):
                db1_acc[j] /= batch_size
            for j in range(128):
                for k in range(10):
                    dW2_acc[j][k] /= batch_size
            for k in range(10):
                db2_acc[k] /= batch_size

            model.update(dW1_acc, db1_acc, dW2_acc, db2_acc, lr)

            if batch_idx % 100 == 0:
                print(f"Epoch {epoch+1}, Batch {batch_idx}, Loss {total_loss/( (batch_idx+1)*batch_size):.4f}")

        # Test after epoch
        test_X = []
        test_y = []
        for data, targets in test_loader:
            test_X.extend(data.view(data.size(0), -1).tolist())
            test_y.extend(targets.tolist())
        acc = model.accuracy(test_X, test_y)
        print(f"Epoch {epoch+1} test accuracy: {acc:.4f}")

if __name__ == "__main__":
    main()