#!/usr/bin/python

from utils import (
    padicval,
    wD,
    volcano_walk,
    find_class_group,
    find_pairs_full,
)

from sage.all import pari
import sage.all
from sage.arith.misc import (
    factor,
    gcd,
)

from sage.misc.functional import sqrt
from sage.rings.integer_ring import ZZ

from itertools import cycle

from theta_structures.couple_point import CouplePoint
from theta_isogenies.product_isogeny import EllipticProductIsogeny


d = -7
n = 64
_f = 5
e_f = 2
f = _f**e_f
D = d if d % 4 == 1 else 4 * d

class_group = find_class_group(d, f)
E, F, phi = volcano_walk(d, n + 2, _f, e_f)

assert F.j_invariant() not in [0, 1728]

q = F.base_field().cardinality()
C = F.cardinality()
t = q + 1 - C
v = ZZ(sqrt((t**2 - 4*q)/D))

print()
print(f"F = {q}, C = {factor(C)}, v = {factor(v)}, f = {factor(f)}")
print()

# Obtain the generator of the endmorphism ring of the curve at the bottom of
# the _f-volcano and at the top of the 2-volcano by lollipopping

wD_E = wD(E, D)

assert phi.degree() == f

phi_rev = phi.dual()

assert phi.codomain() == F
assert phi_rev.domain() == F
assert phi_rev.codomain() == E


def fwD_F(P):
    # Cannot use compositions
    # wD_E defined on R, extension of E
    # wD_E can coerce points from E to R, so we need no explicit call
    # wD_E's codomain is R, so we must coerce to E again for final phi call
    return phi(E(wD_E(phi_rev(P))))


# Since w_D satisfies w_D^2 - D w_D + (D^2 - D)/4 = 0
# fw_D satisfies
# 0 = f^2 w_D^2 - D f^2 w_D + f^2(D^2 - D)/4
#   = (fw_D)^2 - fD (fw_D) + f^2(D^2 - D)/4

assert all([fwD_F(fwD_F(P)) - f*D*fwD_F(P) + f**2*(D**2 - D)//4*P == F.zero() for P in [F.random_point() for _ in range(100)]])

E1 = F
E2 = F

compute = find_pairs_full(class_group, d, n, _f**e_f)

print()
print("Trying Library...")
for cls_no, (H, a, I, b, J, end_a, end_b, kani) in enumerate(compute):
    print()

    E1 = F
    E2 = F

    # Should be true by how pairs are found
    assert gcd(a * I.norm(), b * J.norm()) == 1

    # Composiiton of endmorphisms alpha, beta and kani
    # Recall, they are given as tuples (a, b) in the 1, fw basis
    def end(x):
        x = kani[0]*x + kani[1]*fwD_F(x)
        x = end_b[0]*x + end_b[1]*fwD_F(x)
        x = end_a[0]*x + end_a[1]*fwD_F(x)
        return x

    n = padicval(2, a*I.norm() + b*J.norm())

    P, Q = F.torsion_basis(2**n)

    # Kernel of the whole 2^n isogeny generated by
    # They should all be points of order 2^n
    (P1, P2), (Q1, Q2) = ((a*I.norm()*P, end(P)), (a*I.norm()*Q, end(Q)))

    splits = 0

    assert n == padicval(2, P1.order())

    skip = False
    while n > 0:
        if skip:
            break

        n = padicval(2, P1.order())

        if n == 0:
            break

        if not all([P1.order() == pt.order() for pt in [P2, Q1, Q2]]):
            print("Orders of kernel points are not the same")
            print(f"{n = } order(P1) = 2^{padicval(2, P1.order())}")
            print(f"{n = } order(P2) = 2^{padicval(2, P2.order())}")
            print(f"{n = } order(Q1) = 2^{padicval(2, Q1.order())}")
            print(f"{n = } order(Q2) = 2^{padicval(2, Q2.order())}")
            skip = True
            continue
        else:
            print("Orders of kernel points are correct")

        # Kernel of the first 2-isogeny step in the chain generated by
        (p1, p2), (q1, q2) = (2**(n-1)*P1, 2**(n-1)*P2), (2**(n-1)*Q1, 2**(n-1)*Q2)

        if p2 == F.zero() and q1 == F.zero():
            print(f"Class: {cls_no} (Manual splits: {splits}) Split in 'diagonal' case!")
            print("Computing 2-step manually")
            splits += 1

            assert p1 in E1
            phi_p1 = E1.isogeny(p1)

            assert q2 in E2
            phi_q2 = E2.isogeny(q2)

            P1, P2 = phi_p1(P1), phi_q2(P2)
            Q1, Q2 = phi_p1(Q1), phi_q2(Q2)

            E1 = phi_p1.codomain()
            E2 = phi_q2.codomain()

            continue

        elif p1 == F.zero() and q2 == F.zero():
            splits += 1
            print(f"Class: {cls_no} (Manual splits: {splits}) Split in 'anti-diagonal' case!")
            print("Computing 2-step manually")

            assert p2 in E2
            phi_p2 = E2.isogeny(p2)

            assert q1 in E1
            phi_q1 = E1.isogeny(q1)

            P1, P2 = phi_p2(P2), phi_q1(P1)
            Q1, Q2 = phi_p2(Q2), phi_q1(Q1)

            E1 = phi_p2.codomain()
            E2 = phi_q1.codomain()

            continue

        elif p1 == p2 and q1 == q2:
            splits += 1
            print(f"Class: {cls_no} (Manual splits: {splits}) Split in 'full' case!")
            print("Computing 2-step manually")

            P1, P2 = -P1 + P2, P1 + P2
            Q1, Q2 = -Q1 + Q2, Q1 + Q2

            continue

        break

    if skip:
        continue

    n = padicval(2, P1.order())
    if n == 0:
        print("Manually computed the whole isogeny")
        continue

    kernel = (CouplePoint(P1, P2), CouplePoint(Q1, Q2))

    try:
        Phi = EllipticProductIsogeny(kernel, n)
    except Exception as e:
        print(f"Failed: {e}")
        continue

    print(Phi.codomain()[0])