import re
import numpy as np
from fractions import Fraction


HIDDEN = 128

BATCH_SIZE = 128

LEARNING_RATE = 0.003

OP_NAMES = ['+', '-', '*', '/']

INPUT_SIZE = 34

def encode(a, b, state, op):
    x = np.zeros(INPUT_SIZE, dtype=np.float64)
    x[a] = 1.0
    x[10 + b] = 1.0
    x[20 + state] = 1.0
    x[30 + op] = 1.0
    return x

def build_data():
    datasets = {op: [] for op in range(4)}
    for a in range(10):
        for b in range(10):
            for carry in range(2):
                total = a + b + carry
                datasets[0].append((encode(a, b, carry, 0), total % 10, total // 10))
    for a in range(10):
        for b in range(10):
            for borrow in range(2):
                total = a - b - borrow
                digit = total % 10
                next_borrow = int(total < 0)
                datasets[1].append((encode(a, b, borrow, 1), digit, next_borrow))
    for a in range(10):
        for b in range(10):
            for carry in range(9):
                total = a * b + carry
                datasets[2].append((encode(a, b, carry, 2), total % 10, total // 10))
    for divisor in range(1, 10):
        for remainder in range(divisor):
            for next_digit in range(10):
                current = remainder * 10 + next_digit
                quotient_digit = current // divisor
                next_remainder = current % divisor
                datasets[3].append((encode(next_digit, divisor, remainder, 3), quotient_digit, next_remainder))
    result = {}
    for op, rows in datasets.items():
        result[op] = (np.stack([r[0] for r in rows]), np.array([r[1] for r in rows], dtype=np.int64), np.array([r[2] for r in rows], dtype=np.int64))
    return result

DATA = build_data()

def softmax(z):
    z = z - np.max(z, axis=1, keepdims=True)
    e = np.exp(z)
    return e / np.sum(e, axis=1, keepdims=True)

class MathosV6:

    def __init__(self, seed=77):
        r = np.random.default_rng(seed)

        def init(nin, nout):
            limit = np.sqrt(6.0 / (nin + nout))
            return r.uniform(-limit, limit, (nin, nout))
        self.p = {'W1': init(INPUT_SIZE, HIDDEN), 'b1': np.zeros((1, HIDDEN)), 'Wd': init(HIDDEN, 10), 'bd': np.zeros((1, 10)), 'Ws': init(HIDDEN, 10), 'bs': np.zeros((1, 10))}
        self.m = {k: np.zeros_like(v) for k, v in self.p.items()}
        self.v = {k: np.zeros_like(v) for k, v in self.p.items()}
        self.step = 0
        self.experiences_seen = 0

    def forward(self, X):
        h = np.tanh(X @ self.p['W1'] + self.p['b1'])
        return (h, h @ self.p['Wd'] + self.p['bd'], h @ self.p['Ws'] + self.p['bs'])

    def train_batch(self, X, yd, ys):
        n = len(X)
        h, dl, sl = self.forward(X)
        pd, ps = (softmax(dl), softmax(sl))
        rows = np.arange(n)
        gd, gs = (pd.copy(), ps.copy())
        gd[rows, yd] -= 1.0
        gs[rows, ys] -= 1.0
        gd /= n
        gs /= n
        grads = {'Wd': h.T @ gd, 'bd': gd.sum(axis=0, keepdims=True), 'Ws': h.T @ gs, 'bs': gs.sum(axis=0, keepdims=True)}
        dh = gd @ self.p['Wd'].T + gs @ self.p['Ws'].T
        dz = dh * (1.0 - h * h)
        grads['W1'] = X.T @ dz
        grads['b1'] = dz.sum(axis=0, keepdims=True)
        self.step += 1
        for k in self.p:
            g = grads[k]
            self.m[k] = 0.9 * self.m[k] + 0.1 * g
            self.v[k] = 0.999 * self.v[k] + 0.001 * g * g
            mhat = self.m[k] / (1.0 - 0.9 ** self.step)
            vhat = self.v[k] / (1.0 - 0.999 ** self.step)
            self.p[k] -= LEARNING_RATE * mhat / (np.sqrt(vhat) + 1e-08)
        self.experiences_seen += n
        loss_d = -np.log(pd[rows, yd] + 1e-12).mean()
        loss_s = -np.log(ps[rows, ys] + 1e-12).mean()
        return float(loss_d + loss_s)

    def transition(self, a, b, state, op):
        x = encode(a, b, state, op).reshape(1, -1)
        _, dl, sl = self.forward(x)
        return (int(dl[0].argmax()), int(sl[0].argmax()))

    def transition_accuracy(self, op=None):
        selected = range(4) if op is None else [op]
        correct, total = (0, 0)
        for operation in selected:
            X, yd, ys = DATA[operation]
            _, dl, sl = self.forward(X)
            pd, ps = (dl.argmax(axis=1), sl.argmax(axis=1))
            correct += int(np.sum((pd == yd) & (ps == ys)))
            total += len(X)
        return correct / total

    def train_operation(self, op, epochs, seed):
        X, yd, ys = DATA[op]
        r = np.random.default_rng(seed)
        name = OP_NAMES[op]
        for ep in range(1, epochs + 1):
            order = r.permutation(len(X))
            total_loss = 0.0
            for start in range(0, len(X), BATCH_SIZE):
                ids = order[start:start + BATCH_SIZE]
                loss = self.train_batch(X[ids], yd[ids], ys[ids])
                total_loss += loss * len(ids)
            total_loss /= len(X)
            if ep == 1 or ep % 50 == 0 or ep == epochs:
                acc = self.transition_accuracy(op)
                print(f'Learn {name} | Epoch {ep:03d}/{epochs} | loss={total_loss:.5f} | transition accuracy={100 * acc:.2f}%')

    def add(self, a, b):
        if a < 0 or b < 0:
            raise ValueError('Addition accepts nonnegative integers.')
        result, carry, pos = (0, 0, 0)
        while pos < max(len(str(a)), len(str(b))) or carry:
            da = a // 10 ** pos % 10
            db = b // 10 ** pos % 10
            digit, carry = self.transition(da, db, carry, 0)
            result += digit * 10 ** pos
            pos += 1
        return result

    def subtract(self, a, b):
        if a < b or b < 0:
            raise ValueError('Subtraction requires a >= b >= 0.')
        result, borrow = (0, 0)
        for pos in range(len(str(a))):
            da = a // 10 ** pos % 10
            db = b // 10 ** pos % 10
            digit, borrow = self.transition(da, db, borrow, 1)
            result += digit * 10 ** pos
        if borrow != 0:
            raise ArithmeticError('Learned subtraction ended with a borrow.')
        return result

    def multiply_digit(self, n, digit):
        if not 0 <= digit <= 9:
            raise ValueError('Single-digit multiplier must be 0..9.')
        result, carry, pos = (0, 0, 0)
        while pos < len(str(n)) or carry:
            da = n // 10 ** pos % 10
            out_digit, carry = self.transition(da, digit, carry, 2)
            result += out_digit * 10 ** pos
            pos += 1
        return result

    def multiply(self, a, b):
        if a < 0 or b < 0:
            raise ValueError('Multiplication accepts nonnegative integers.')
        result = 0
        for pos, ch in enumerate(str(b)[::-1]):
            partial = self.multiply_digit(a, int(ch))
            for _ in range(pos):
                partial *= 10
            result = self.add(result, partial)
        return result

    def divide(self, dividend, divisor):
        if dividend < 0 or divisor <= 0:
            raise ValueError('Division requires dividend >= 0, divisor > 0.')
        quotient, remainder = (0, 0)
        for ch in str(dividend):
            digit = int(ch)
            if divisor <= 9:
                qdigit, new_remainder = self.transition(digit, divisor, remainder, 3)
            else:
                current = int(str(remainder) + ch)
                qdigit = 0
                for trial in range(9, -1, -1):
                    product = self.multiply(divisor, trial)
                    if product <= current:
                        qdigit = trial
                        break
                product = self.multiply(divisor, qdigit)
                new_remainder = self.subtract(current, product)
            quotient = quotient * 10 + qdigit
            remainder = new_remainder
        return (quotient, remainder)

    def save(self, path):
        np.savez(path, **self.p, experiences_seen=np.array(self.experiences_seen))

def signed_add(model, x, y, subtract=False):
    if subtract:
        y = -y
    sx = -1 if x < 0 else 1
    sy = -1 if y < 0 else 1
    ax, ay = (abs(x), abs(y))
    if sx == sy:
        return sx * model.add(ax, ay)
    if ax >= ay:
        return sx * model.subtract(ax, ay)
    return sy * model.subtract(ay, ax)

def signed_multiply(model, x, y):
    sign = -1 if (x < 0) != (y < 0) else 1
    return sign * model.multiply(abs(x), abs(y))

def fraction_operation(model, left, right, op):
    left, right = (Fraction(left), Fraction(right))
    if op == '+' or op == '-':
        n1 = signed_multiply(model, left.numerator, right.denominator)
        n2 = signed_multiply(model, right.numerator, left.denominator)
        numerator = signed_add(model, n1, n2, subtract=op == '-')
        denominator = model.multiply(left.denominator, right.denominator)
        return Fraction(numerator, denominator)
    if op == '*':
        numerator = signed_multiply(model, left.numerator, right.numerator)
        denominator = model.multiply(left.denominator, right.denominator)
        return Fraction(numerator, denominator)
    if op == '/':
        if right == 0:
            raise ZeroDivisionError('Division by zero.')
        numerator = signed_multiply(model, left.numerator, right.denominator)
        denominator = signed_multiply(model, left.denominator, right.numerator)
        return Fraction(numerator, denominator)
    raise ValueError('Unsupported operation.')

def fraction_power(model, base, exponent):
    base = Fraction(base)
    exponent = Fraction(exponent)
    if exponent.denominator != 1:
        raise ValueError('Only integer powers are supported.')
    power = exponent.numerator
    if abs(power) > 20:
        raise ValueError('Power magnitude is limited to 20.')
    if power < 0:
        if base == 0:
            raise ZeroDivisionError('Zero cannot have a negative power.')
        base = Fraction(base.denominator, base.numerator)
        power = -power
    result = Fraction(1, 1)
    for _ in range(power):
        result = fraction_operation(model, result, base, '*')
    return result

def tokenize(expression):
    cleaned = re.sub('\\s+', '', expression)
    tokens = re.findall('\\d+|[()+*/^%-]', cleaned)
    if ''.join(tokens) != cleaned:
        raise ValueError('Invalid character in expression.')
    return tokens

def bodmas(model, expression):
    tokens = tokenize(expression)
    index = 0

    def parse_expression():
        nonlocal index
        value = parse_term()
        while index < len(tokens) and tokens[index] in ('+', '-'):
            op = tokens[index]
            index += 1
            rhs = parse_term()
            value = fraction_operation(model, value, rhs, op)
        return value

    def parse_term():
        nonlocal index
        value = parse_power()
        while index < len(tokens) and tokens[index] in ('*', '/'):
            op = tokens[index]
            index += 1
            rhs = parse_power()
            value = fraction_operation(model, value, rhs, op)
        return value

    def parse_power():
        nonlocal index
        value = parse_factor()
        if index < len(tokens) and tokens[index] == '^':
            index += 1
            exponent = parse_power()
            value = fraction_power(model, value, exponent)
        return value

    def parse_factor():
        nonlocal index
        if index >= len(tokens):
            raise ValueError('Incomplete expression.')
        if tokens[index] == '+':
            index += 1
            return parse_factor()
        if tokens[index] == '-':
            index += 1
            return -parse_factor()
        if tokens[index] == '(':
            index += 1
            value = parse_expression()
            if index >= len(tokens) or tokens[index] != ')':
                raise ValueError('Missing closing parenthesis.')
            index += 1
            return value
        token = tokens[index]
        if not token.isdigit():
            raise ValueError(f'Unexpected token: {token}')
        index += 1
        return Fraction(int(token), 1)
    result = parse_expression()
    if index != len(tokens):
        raise ValueError('Unexpected trailing tokens.')
    return result


from importlib.resources import files

class mathos:
    """Supernova Mathos V6.1 with bundled trained NumPy weights."""

    def __init__(self):
        self.model = MathosV6(seed=77)
        weight_path = files("supernova_mathos").joinpath(
            "weights/supernova_mathos_v6_1.npz"
        )
        with np.load(weight_path) as saved:
            self.model.p = {
                key: saved[key].copy()
                for key in ("W1", "b1", "Wd", "bd", "Ws", "bs")
            }
            if "experiences_seen" in saved:
                self.model.experiences_seen = int(saved["experiences_seen"])

    def __getattr__(self, name):
        return getattr(self.model, name)

    def calculate(self, expression):
        return bodmas(self.model, expression)