import numpy as np


class RANArray:
    """
    Array-valued Rational-Addition Number: elementwise a/b + c.
    Supports numpy broadcasting, arrays, matrices, scalars, and standard
    arithmetic operators (+, -, *, /, @, **).
    """
    __slots__ = ("a", "b", "c")
    __array_priority__ = 1000  # Ensure RANArray wins over numpy in mixed ops

    def __init__(self, a, b, c):
        self.a = np.asarray(a, dtype=np.float64)
        self.b = np.asarray(b, dtype=np.float64)
        self.c = np.asarray(c, dtype=np.float64)

    def collapse(self):
        """Evaluate a/b + c. Returns a numpy array."""
        return self.a / self.b + self.c

    @property
    def T(self):
        return RANArray(self.a.T, self.b.T, self.c.T)

    @property
    def shape(self):
        return np.broadcast(self.a, self.b, self.c).shape

    @property
    def ndim(self):
        return self.a.ndim

    def __repr__(self):
        return f"RANArray(shape={self.shape}, value≈{self.collapse()})"

    @staticmethod
    def _coerce(other):
        """Convert numbers, arrays, lists, tuples to RANArray."""
        if isinstance(other, RANArray):
            return other
        if isinstance(other, (np.ndarray, list, tuple)):
            return real_ran(np.asarray(other, dtype=np.float64))
        if isinstance(other, (int, float, np.integer, np.floating)):
            return real_ran(float(other))
        return NotImplemented

    # ----- unary -----
    def __neg__(self):
        return ran_neg(self)

    # ----- addition -----
    def __add__(self, other):
        other = self._coerce(other)
        if other is NotImplemented:
            return NotImplemented
        return ran_add(self, other)

    __radd__ = __add__

    # ----- subtraction -----
    def __sub__(self, other):
        other = self._coerce(other)
        if other is NotImplemented:
            return NotImplemented
        return ran_sub(self, other)

    def __rsub__(self, other):
        other = self._coerce(other)
        if other is NotImplemented:
            return NotImplemented
        return ran_sub(other, self)

    # ----- multiplication -----
    def __mul__(self, other):
        other = self._coerce(other)
        if other is NotImplemented:
            return NotImplemented
        return ran_mul(self, other)

    __rmul__ = __mul__

    # ----- division -----
    def __truediv__(self, other):
        other = self._coerce(other)
        if other is NotImplemented:
            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))
        A = e * (a + b * c)
        B = b * (d + e * f)
        return RANArray(A, B, np.zeros_like(A))

    def __rtruediv__(self, other):
        other = self._coerce(other)
        if other is NotImplemented:
            return NotImplemented
        return other.__truediv__(self)

    # ----- matrix multiplication -----
    def __matmul__(self, other):
        other = self._coerce(other)
        if other is NotImplemented:
            return NotImplemented
        return ran_matmul(self, other)

    def __rmatmul__(self, other):
        other = self._coerce(other)
        if other is NotImplemented:
            return NotImplemented
        return ran_matmul(other, self)

    # ----- power -----
    def __pow__(self, other):
        if isinstance(other, (int, np.integer)):
            if other == 0:
                return const_ran(np.ones_like(self.a))
            result = self
            for _ in range(abs(other) - 1):
                result = ran_mul(result, self)
            if other < 0:
                return const_ran(np.ones_like(result.a)) / result
            return result
        # Non-integer or array exponent: collapse to ordinary numpy and re-wrap
        return real_ran(self.collapse() ** other)

    # ----- indexing and shape -----
    def __getitem__(self, idx):
        return RANArray(self.a[idx], self.b[idx], self.c[idx])

    def __setitem__(self, idx, value):
        value = self._coerce(value)
        self.a[idx] = value.a
        self.b[idx] = value.b
        self.c[idx] = value.c

    def reshape(self, *shape):
        return RANArray(self.a.reshape(shape), self.b.reshape(shape), self.c.reshape(shape))

    def flatten(self):
        return RANArray(self.a.ravel(), self.b.ravel(), self.c.ravel())

    def sum(self, axis=None):
        """Reduce along axis using RAN algebra. axis=None flattens."""
        a, b, c = self.a, self.b, self.c
        if axis is None:
            a, b, c = a.ravel(), b.ravel(), c.ravel()
            axis = 0
        if np.all(b == 1):
            return RANArray(np.sum(a, axis=axis), np.ones(np.sum(a, axis=axis).shape), np.sum(c, axis=axis))
        a, b, c = _tree_reduce_add(a, b, c, axis=axis)
        return RANArray(a, b, c)

    def mean(self, axis=None):
        n = self.a.size if axis is None else self.a.shape[axis]
        return self.sum(axis=axis) / n

    def __len__(self):
        if self.a.ndim == 0:
            raise TypeError("len() of unsized RANArray")
        return self.a.shape[0]

    def __iter__(self):
        if self.a.ndim == 0:
            raise TypeError("iteration over a 0-d RANArray")
        for i in range(self.a.shape[0]):
            yield self[i]

    def __eq__(self, other):
        other = self._coerce(other)
        if other is NotImplemented:
            return NotImplemented
        return np.isclose(self.collapse(), other.collapse())


# ------------------------------------------------------------
# Constructors
# ------------------------------------------------------------
def const_ran(value):
    """Build RAN(v, 1, 0): represents the exact value v."""
    v = np.asarray(value, dtype=np.float64)
    return RANArray(v, np.ones_like(v), np.zeros_like(v))


def real_ran(value):
    """Build RAN(0, 1, v): represents the value v with c=v."""
    v = np.asarray(value, dtype=np.float64)
    return RANArray(np.zeros_like(v), np.ones_like(v), v)


def zero_ran(shape):
    z = np.zeros(shape, dtype=np.float64)
    return RANArray(z, np.ones_like(z), z.copy())


# ------------------------------------------------------------
# Elementwise RAN algebra (matches scalar RAN formulas)
# ------------------------------------------------------------
def ran_add(x, y):
    a = x.a * y.b + x.b * y.a
    b = x.b * y.b
    c = x.c + y.c
    return RANArray(a, b, c)


def ran_neg(x):
    return RANArray(-x.a, x.b, -x.c)


def ran_sub(x, y):
    return ran_add(x, ran_neg(y))


def ran_mul(x, y):
    a = x.a * y.a + x.a * y.c * y.b + x.c * x.b * y.a
    b = x.b * y.b
    c = x.c * y.c
    return RANArray(a, b, c)


# ------------------------------------------------------------
# Associative tree reduction (preserves RAN algebra, log n depth)
# ------------------------------------------------------------
def _tree_reduce_add(a, b, c, axis):
    a = np.moveaxis(a, axis, 0)
    b = np.moveaxis(b, axis, 0)
    c = np.moveaxis(c, axis, 0)
    while a.shape[0] > 1:
        n = a.shape[0]
        half = n // 2
        a1, b1, c1 = a[:half], b[:half], c[:half]
        a2, b2, c2 = a[half:2 * half], b[half:2 * half], c[half:2 * half]
        na = a1 * b2 + b1 * a2
        nb = b1 * b2
        nc = c1 + c2
        if n % 2 == 1:
            na = np.concatenate([na, a[-1:]], axis=0)
            nb = np.concatenate([nb, b[-1:]], axis=0)
            nc = np.concatenate([nc, c[-1:]], axis=0)
        a, b, c = na, nb, nc
    return a[0], b[0], c[0]


# ------------------------------------------------------------
# Matrix product and axis-0 sum
# ------------------------------------------------------------
def ran_matmul(x, w):
    """
    RAN-space matrix product.
    Fast path: when b == 1 everywhere (the common NN case), uses BLAS.
    General path: full RAN tree reduction over the contracted axis.
    """
    if np.all(x.b == 1) and np.all(w.b == 1):
        a = x.a @ w.a + x.a @ w.c + x.c @ w.a
        c = x.c @ w.c
        b = np.ones_like(a)
        return RANArray(a, b, c)

    if x.a.ndim != 2 or w.a.ndim != 2:
        raise NotImplementedError(
            "General RAN matmul only supports 2D arrays; use collapse() for 1D."
        )

    xa = x.a[:, :, None]
    xb = x.b[:, :, None]
    xc = x.c[:, :, None]
    wa = w.a[None, :, :]
    wb = w.b[None, :, :]
    wc = w.c[None, :, :]
    pa = xa * wa + xa * wc * wb + xc * xb * wa
    pb = xb * wb
    pc = xc * wc
    a, b, c = _tree_reduce_add(pa, pb, pc, axis=1)
    return RANArray(a, b, c)


def ran_sum0(x):
    """Sum a RANArray along axis 0 (used for bias gradients)."""
    if np.all(x.b == 1):
        return RANArray(np.sum(x.a, axis=0), np.ones(x.a.shape[1]), np.sum(x.c, axis=0))
    a, b, c = _tree_reduce_add(x.a, x.b, x.c, axis=0)
    return RANArray(a, b, c)


# ------------------------------------------------------------
# Demo / sanity check
# ------------------------------------------------------------
if __name__ == "__main__":
    # 1. Scalar operations
    x = RANArray(1.0, 2.0, 3.0)   # 1/2 + 3 = 3.5
    y = x + 1.0
    print("Scalar +:", y, "->", y.collapse())

    z = x * 2.0
    print("Scalar *:", z, "->", z.collapse())

    q = x / 2.0
    print("Scalar /:", q, "->", q.collapse())

    # 2. Array broadcasting
    arr = real_ran(np.array([1.0, 2.0, 3.0]))
    print("Array + 10:", (arr + 10).collapse())
    print("Array * 2:", (arr * 2).collapse())

    # 3. Matrix multiplication
    A = real_ran(np.random.randn(4, 3))
    B = real_ran(np.random.randn(3, 5))
    C = A @ B
    print("Matmul shape:", C.shape)
    print("Matmul matches numpy:", np.allclose(C.collapse(), A.collapse() @ B.collapse()))

    # 4. Indexing and assignment
    arr[1] = 100.0
    print("After assignment:", arr.collapse())

    # 5. Sum/mean
    m = real_ran(np.array([[1.0, 2.0], [3.0, 4.0]]))
    print("Sum axis 0:", m.sum(axis=0).collapse())
    print("Mean:", m.mean().collapse())