import numpy as np


class RANArray:
    """
    Array-valued Rational-Addition Number: elementwise a/b + c, where
    a, b, c are numpy arrays (or broadcastable). Same algebra as the
    scalar RAN(a,b,c) class, but every operation acts on the whole
    array at once via numpy instead of looping over Python RAN objects
    one scalar at a time. Collapse is still only performed where the
    original code performed it (relu's sign check, softmax boundary).
    """
    __slots__ = ("a", "b", "c")

    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):
        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

    def __repr__(self):
        return f"RANArray(shape={self.shape})"


# ------------------------------------------------------------
# Core elementwise algebra (matches the scalar RAN formulas exactly,
# just applied to whole arrays via numpy broadcasting)
# ------------------------------------------------------------
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)


# ------------------------------------------------------------
# Constructors matching the two conversion conventions used in the
# original scalar RAN code:
#   - "const" RANs built explicitly as RAN(value, 1, 0)   (e.g. lr,
#     dz2 = probs - y_true, relu_derivative multiplier)
#   - "real" conversions of plain floats as RAN(0, 1, value)  (e.g.
#     Real.__rmul__ auto-conversion when multiplying a float by a RAN)
# ------------------------------------------------------------
def const_ran(value):
    v = np.asarray(value, dtype=np.float64)
    return RANArray(v, np.ones_like(v), np.zeros_like(v))


def real_ran(value):
    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())


# ------------------------------------------------------------
# Matrix product and axis-0 sum, implemented as an associative
# (order-independent) tree reduction of ran_add over the contracted
# axis, so the RAN algebra is preserved exactly (no shortcut that
# assumes b==1) while doing O(log n) vectorized numpy calls instead
# of a Python-level loop per scalar.
# ------------------------------------------------------------
def _tree_reduce_add(a, b, c, axis):
    # a, b, c: numpy arrays with `axis` being the dimension to reduce
    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]


def ran_matmul(x, w):
    """
    x: RANArray, components shaped (N, D)
    w: RANArray, components shaped (D, H)
    returns: RANArray, components shaped (N, H)
    Equivalent to sum_d ran_mul(x[:, d], w[d, :]), i.e. a RAN-space
    matrix product where each individual term is still formed with the
    real RAN multiply/add formulas (no rounding is done early).

    Fast path: in this network b is provably 1 for every RAN value at
    every step (weights/biases are built as RAN(v,1,0) or RAN(0,1,0),
    inputs as RAN(0,1,x), and __truediv__ -- the only op that changes
    b -- is never used). With b == e == 1 everywhere, the general
    ran_mul/ran_add reduction collapses algebraically to two ordinary
    linear sums, which lets us use BLAS matmuls instead of a
    Python-level reduction -- same exact numbers, much faster. If that
    invariant is ever broken (e.g. division gets introduced), we fall
    back to the general, slower tree reduction so results stay correct
    either way.
    """
    if np.all(x.b == 1) and np.all(w.b == 1):
        # with b=e=1, ran_mul(x_d, w_d).a = x.a*w.a + x.a*w.c + x.c*w.a
        # and .c = x.c*w.c; summing over d turns each term into a matmul
        a = x.a @ w.a + x.a @ w.c + x.c @ w.a
        c = x.c @ w.c
        b = np.ones((x.a.shape[0], w.a.shape[1]))
        return RANArray(a, b, c)

    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). Same
    fast-path reasoning as ran_matmul: with b == 1 throughout, RAN
    addition of a column is just np.sum on `a` and on `c` separately."""
    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)