
# adjusted from sage.schemes.hyperelliptic_curves.jacobian_morphism

def cantor_reduction_step(a, b, f):
    a2 = (f - b**2) // a
    a2 = a2.monic()
    b2 = -b % a2
    assert a2.degree() < a.degree()
    return a2, b2

def cantor_reduction_simple(a, b, f, genus):
    while a.degree() > genus:
        a, b = cantor_reduction_step(a, b, f)
    return a, b

def cantor_composition_simple(D1, D2, f, _=None):
    a1, b1 = D1
    a2, b2 = D2
    if a1 == a2 and b1 == b2:
        d, h1, h3 = a1.xgcd(2*b1)
        a = (a1 // d)**2
        b = (b1 + h3*((f - b1**2) // d)) % (a)
    else:
        d0, _, h2 = a1.xgcd(a2)
        if d0 == 1:
            a = a1*a2
            b = (b2 + h2*a2*(b1-b2)) % (a)
        else:
            d, l, h3 = d0.xgcd(b1 + b2)
            a = (a1*a2) // (d**2)
            b = ((b2 + l*h2*(b1-b2)*(a2 // d)) + h3*((f - b2**2) // d)) % (a)
    a = a.monic()
    return a, b

################################################################

from sage.structure.parent import Parent
from sage.structure.element import Element
from sage.structure.unique_representation import CachedRepresentation

class DivisorClassGroup(Parent, CachedRepresentation):
    def __init__(self, f):
        if f is ...:
            print(f'polynomial f is a placeholder; see params.censored.py')
            exit()
        super().__init__()
        self.f = f
        if f.degree() % 2 == 0:
            raise NotImplementedError('only odd-degree models supported')
        self.g = f.degree() // 2
        self.F = self.f.parent().base_ring()

    def _repr_(self):
        return f'divisor class group of y^2 = {self.f}'

    def _element_constructor_(self, *inp):
        if inp == (0,):
            return self.zero()
        return DivisorClass(self, *inp)

    def zero(self):
        return DivisorClass(self, 1, 0)

class DivisorClass(Element):
    def __init__(self, parent, u, v):
        f = parent.f
        u,v = map(f.parent(), (u,v))
        super().__init__(parent)
        if u.lc() != 1:
            raise ValueError('u must be monic')
        if not u.divides(f - v**2):
            raise ValueError('u must divide f-v²')
        if v.degree() >= u.degree():
            raise ValueError('need deg(u) < deg(v)')
        self.u = u
        self.v = v

    def __repr__(self):
        return f'({self.u}, {self.v})'

    def _add_(self, other):
        f = self.parent().f
        u, v = cantor_composition_simple((self.u,self.v), (other.u,other.v), f)
        u, v = cantor_reduction_simple(u, v, f, self.parent().g)
        return self.parent()(u, v)

    def __neg__(self):
        return self.parent()(self.u, -self.v)

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

    def __mul__(self, other):
        k = int(other)
        if k < 0:
            return -self * -k
        res = self.parent().zero()
        tmp = self
        while k:
            if k & 1:
                res += tmp
            tmp += tmp
            k >>= 1
        return res
    __rmul__ = __mul__

    def __eq__(self, other):
        u1, v1 = cantor_reduction_simple(self.u, self.v, self.parent().f, self.parent().g)
        u2, v2 = cantor_reduction_simple(other.u, other.v, other.parent().f, other.parent().g)
        return (u1, v1) == (u2, v2)

