#!/usr/bin/env python3
"""Zero-dependency exact certificate for the OMEGA scalar interface."""

from __future__ import annotations

from fractions import Fraction


class Poly:
    """Sparse multivariate polynomial over Q."""

    def __init__(self, nvars: int, terms: dict[tuple[int, ...], Fraction] | None = None):
        self.nvars = nvars
        self.terms = {m: Fraction(c) for m, c in (terms or {}).items() if c}

    @classmethod
    def constant(cls, nvars: int, value: int | Fraction) -> "Poly":
        return cls(nvars, {(0,) * nvars: Fraction(value)})

    @classmethod
    def variable(cls, nvars: int, index: int) -> "Poly":
        exponent = [0] * nvars
        exponent[index] = 1
        return cls(nvars, {tuple(exponent): Fraction(1)})

    def __add__(self, other: "Poly") -> "Poly":
        assert self.nvars == other.nvars
        terms = dict(self.terms)
        for monomial, coefficient in other.terms.items():
            terms[monomial] = terms.get(monomial, Fraction(0)) + coefficient
        return Poly(self.nvars, terms)

    def __neg__(self) -> "Poly":
        return Poly(self.nvars, {m: -c for m, c in self.terms.items()})

    def __sub__(self, other: "Poly") -> "Poly":
        return self + (-other)

    def __mul__(self, other: "Poly") -> "Poly":
        assert self.nvars == other.nvars
        terms: dict[tuple[int, ...], Fraction] = {}
        for left_monomial, left_coefficient in self.terms.items():
            for right_monomial, right_coefficient in other.terms.items():
                monomial = tuple(a + b for a, b in zip(left_monomial, right_monomial))
                terms[monomial] = terms.get(monomial, Fraction(0)) + left_coefficient * right_coefficient
        return Poly(self.nvars, terms)

    def is_zero(self) -> bool:
        return not self.terms


class Rat:
    """Unreduced rational function; equality is exact cross multiplication."""

    def __init__(self, numerator: Poly, denominator: Poly):
        if denominator.is_zero():
            raise ZeroDivisionError("formal zero denominator")
        self.numerator = numerator
        self.denominator = denominator

    @classmethod
    def from_poly(cls, value: Poly) -> "Rat":
        return cls(value, Poly.constant(value.nvars, 1))

    def __add__(self, other: "Rat") -> "Rat":
        return Rat(
            self.numerator * other.denominator + other.numerator * self.denominator,
            self.denominator * other.denominator,
        )

    def __neg__(self) -> "Rat":
        return Rat(-self.numerator, self.denominator)

    def __sub__(self, other: "Rat") -> "Rat":
        return self + (-other)

    def __mul__(self, other: "Rat") -> "Rat":
        return Rat(self.numerator * other.numerator, self.denominator * other.denominator)

    def __truediv__(self, other: "Rat") -> "Rat":
        if other.numerator.is_zero():
            raise ZeroDivisionError("division by formal zero")
        return Rat(self.numerator * other.denominator, self.denominator * other.numerator)

    def equals(self, other: "Rat") -> bool:
        return (self.numerator * other.denominator - other.numerator * self.denominator).is_zero()


def verify_general_inverse() -> None:
    # Variable order: x, e, t, ell.
    nvars = 4
    x, e, t, ell = (Rat.from_poly(Poly.variable(nvars, i)) for i in range(nvars))
    one = Rat.from_poly(Poly.constant(nvars, 1))
    k = Rat.from_poly(Poly.constant(nvars, 648000))

    denominator = x + k * e
    tau = one / denominator
    lam = x / denominator
    recovered_x = lam / tau
    recovered_e = (one - lam) / (k * tau)
    assert recovered_x.equals(x)
    assert recovered_e.equals(e)

    inverse_x = ell / t
    inverse_e = (one - ell) / (k * t)
    inverse_denominator = inverse_x + k * inverse_e
    recovered_t = one / inverse_denominator
    recovered_ell = inverse_x / inverse_denominator
    assert recovered_t.equals(t)
    assert recovered_ell.equals(ell)


def qphi_mul(left: tuple[Fraction, Fraction], right: tuple[Fraction, Fraction]) -> tuple[Fraction, Fraction]:
    """Multiply a+b*phi using phi^2=phi+1."""
    a, b = left
    c, d = right
    return a * c + b * d, a * d + b * c + b * d


class PiPoly:
    """Polynomial in formal pi with coefficients in Q(phi)."""

    def __init__(self, terms: dict[int, tuple[Fraction, Fraction]] | None = None):
        self.terms = {
            degree: (Fraction(coefficient[0]), Fraction(coefficient[1]))
            for degree, coefficient in (terms or {}).items()
            if coefficient != (0, 0)
        }

    @classmethod
    def rational(cls, value: int) -> "PiPoly":
        return cls({0: (Fraction(value), Fraction(0))})

    def __add__(self, other: "PiPoly") -> "PiPoly":
        terms = dict(self.terms)
        for degree, (a, b) in other.terms.items():
            old_a, old_b = terms.get(degree, (Fraction(0), Fraction(0)))
            terms[degree] = old_a + a, old_b + b
        return PiPoly(terms)

    def __neg__(self) -> "PiPoly":
        return PiPoly({degree: (-a, -b) for degree, (a, b) in self.terms.items()})

    def __sub__(self, other: "PiPoly") -> "PiPoly":
        return self + (-other)

    def __mul__(self, other: "PiPoly") -> "PiPoly":
        terms: dict[int, tuple[Fraction, Fraction]] = {}
        for left_degree, left_coefficient in self.terms.items():
            for right_degree, right_coefficient in other.terms.items():
                degree = left_degree + right_degree
                product = qphi_mul(left_coefficient, right_coefficient)
                old = terms.get(degree, (Fraction(0), Fraction(0)))
                terms[degree] = old[0] + product[0], old[1] + product[1]
        return PiPoly(terms)

    def is_zero(self) -> bool:
        return not self.terms


class PiRat:
    def __init__(self, numerator: PiPoly, denominator: PiPoly):
        if denominator.is_zero():
            raise ZeroDivisionError("formal zero denominator")
        self.numerator = numerator
        self.denominator = denominator

    @classmethod
    def from_poly(cls, value: PiPoly) -> "PiRat":
        return cls(value, PiPoly.rational(1))

    def __add__(self, other: "PiRat") -> "PiRat":
        return PiRat(
            self.numerator * other.denominator + other.numerator * self.denominator,
            self.denominator * other.denominator,
        )

    def __neg__(self) -> "PiRat":
        return PiRat(-self.numerator, self.denominator)

    def __sub__(self, other: "PiRat") -> "PiRat":
        return self + (-other)

    def __mul__(self, other: "PiRat") -> "PiRat":
        return PiRat(self.numerator * other.numerator, self.denominator * other.denominator)

    def __truediv__(self, other: "PiRat") -> "PiRat":
        if other.numerator.is_zero():
            raise ZeroDivisionError("division by formal zero")
        return PiRat(self.numerator * other.denominator, self.denominator * other.numerator)

    def equals(self, other: "PiRat") -> bool:
        return (self.numerator * other.denominator - other.numerator * self.denominator).is_zero()


def verify_canonical_identities() -> None:
    one = PiRat.from_poly(PiPoly.rational(1))
    phi = PiRat.from_poly(PiPoly({0: (Fraction(0), Fraction(1))}))
    pi = PiRat.from_poly(PiPoly({1: (Fraction(1), Fraction(0))}))
    omega = phi + pi
    tau = one / omega
    lam = phi / omega
    mu = pi / omega

    assert (lam + mu).equals(one)
    assert (lam / tau).equals(phi)
    assert (mu / tau).equals(pi)
    assert (lam * lam - tau * lam - tau * tau).equals(PiRat.from_poly(PiPoly.rational(0)))


if __name__ == "__main__":
    verify_general_inverse()
    verify_canonical_identities()
    print("OMEGA scalar interface exact algebra certificate: PASS")