# Import Sage
from sage.all import *

"""
    This is an auxiliary script that serves two purposes:

    Firstly, most of the "two-operand" functions from the main SIDH script
    are given here in a more readable form, i.e., as the full formula. 

    Secondly, and more importantly, the purpose of this script is to illustrate a 
    Remark made in the paper about the equivalence of our fast, projective
    strategy that works solely on the Kummer variety, and the traditional isogeny
    computation strategy, i.e., Velu's formulas in affine Weierstrass form.
    As such, we give two versions of the main key exchange functions: the first 
    version calls our explicit formulas for point and isogeny computations, the
    second calls EllipticCurveWithTorsionPoint class which works in affine
    Weierstrass coordinates. 

    The scripts show that these two strategies produce the same result.
"""

# Turn off arithmetic proof
proof.arithmetic(False)

def PointIsogeny(E, R_list, points=[], degrees_list=[]):
    """
        The last element, R, in R_list must have order 
            - case EVEN, either two
            - case ODD,  or some odd number (prime or not)
        The degrees list has or does not have the parallel information about the orders 
        of the points in the R_list. If no information is available, the code will
        try to compute it.

        We then construct the corresponding isogeny phi *at point level*.
        E must be an elliptic curve in short Weierstrass form.
        E  :: y^2 = x^3 + a*x + b

        CASE EVEN: R = (s,0) has order two, and s is defined in this way.
        CASE ODD : R has two components, we will build the group generated by R.

        We construct ad-hoc the isogeny map phi at point level following Section 3 of
        https://arxiv.org/pdf/math/0404124.pdf

        This gives a wonderful unifying description, that i missed in other sources.
        (It is best suited for a computer implementation (for the case of a cyclic kernel isogeny).)

        We get a map (x,y) -> phi(x,y) = X, Y.
        The article gives an unifying way to look at X-component for both even and odd torsion cases.
        The mention of the Y-component in loc. cit., simply obtained by differentiating
        the formula for X w.r.t. x.
    """

    if not R_list:
        return (E, R_list, points)

    R = R_list.pop()    # R is extracted from the last element of R_list.

    if degrees_list:
        d = degrees_list.pop()
    else:
        d = R.order()    # We have to manually compute the order.

    kernel = [k*R for k in [0..(d-1)]]
    if d*R != E(0) or len([kR for kR in kernel if kR == E(0)]) > 1:
        raise TypeError("There is no d-torsion point.") # We insist to have order exactly d.

    # Case even: R is a 2-torsion point
    if d == 2:
        a, b = E.a4(), E.a6()
        s = R.xy()[0]
        t = 3*s^2 + a
        w = s*t

        EE = EllipticCurve(E.base_field(), [a - 5*t, b - 7*w])

        def phi(P):
            """
                Given P, a point on E, we return a point on EE.
            """
            if P in (E(0), R):
                return EE(0)

            x, y = P.xy()
            X = x + t / (x - s)
            Y = y * (1 - t / (x - s)^2)

            return EE.point((X, Y))

    elif not 2.divides(d):
        F = E.base_field()
        a4, a6 = E.a4(), E.a6()
        b2, b4, b6 = E.b2(), E.b4(), E.b6()
        t, u, w = F(0), F(0), F(0)
        data = []

        d1 = ZZ((d - 1) / 2)
        for P in kernel[1 : d1 + 1]:
            # We split the set G = R, 2R, ... , (d-2)R, (d-1)R in two, G(+)
            # and G(-) so that P in G(+) <=> -P in G(-), the multiplicities
            # above correspond to those for G(+).
            xP , yP = P.xy()
            tP = 6*xP^2 + b2*xP + b4
            uP = 4*xP^3 + b2*xP^2 + 2*b4*xP + b6
            wP = uP + xP * tP

            t += tP
            u += uP
            w += wP

            data.append([xP, yP, tP, uP])

        EE = EllipticCurve(F, [a4 - 5*t, a6 - b2*t - 7*w])

        def phi(point):
            if point in kernel:
                return EE(0)

            x, y = point.xy()
            X, Y = x, y

            for xP, yP, tP, uP in data:
                X +=       tP / (x - xP)   + uP   / (x - xP)^2
                Y += y * (-tP / (x - xP)^2 - uP*2 / (x - xP)^3)    # y*diff(xx, x)

            return EE((X, Y))

    else:
        # Torsion point of order 4, 6, 8, etc.
        raise TypeError("Argument degree (= %s) should be either an odd number or 2." % str(d))

    return (EE
           ,[phi(P) for P in R_list]
           ,[phi(P) for P in points]
           ,degrees_list)

# Paramters defining the prime p = f*lA^eA*lB^eB - 1
f = 1
lA = 2
lB = 3
eA = 372
eB = 239

# Define the prime p
# Define the prime p
p = f*lA**eA*lB**eB-1
assert p.is_prime()

# Prime field of order p
Fp = GF(p)
R.<x> = PolynomialRing(Fp)
# The quadratic extension via x^2 + 1 since p = 3 mod 4
Fp2.<j> = Fp.extension(x^2 + 1)

# Bitlengths of group orders lA^eA and lB^eB, needed during the ladder functions 
eAbits = eA 
eBbits = 379

# E0 is the starting curve E0/Fp2: y^2=x^3+x (the A=0 Montgomery curve)
E0 = EllipticCurve(Fp2, [1,0])
assert E0.is_supersingular()

# The orders of the points on each side
oA = lA**eA
oB = lB**eB

# Identifyers for Alice and Bob
Alice = 0
Bob = 1

# Generator PA for the base field subgroup of order lA^eA
PA = 3**239*E0([11, sqrt(Fp(11^3+11))])
# Generator PB for the base field subgroup of order lB^eB
PB = 2**372*E0([6, sqrt(Fp(6^3+6))])

XPA, YPA = PA[0], PA[1]
XPB, YPB = PB[0], PB[1]

"""
   These are optimal strategies with respect to the cost ratios of 
   scalar multiplication by 4 and 4-isogeny evaluation of 
                     pA/qA = 2*12.1/21.6,
   and of point tripling and 3-isogeny evaluation of 
                     pB/qB = 24.3/16.0. 

   See the file optimalstrategies.txt for how they are computed.
"""
splits_Alice = [
0, 1, 1, 2, 2, 2, 3, 4, 4, 4, 4, 5, 5, 6, 7, 8, 8, 9, 9, 9, 9, 9, 9, 9, 12, 11,
12, 12, 13, 14, 15, 16, 16, 16, 16, 16, 16, 17, 17, 18, 18, 17, 21, 17, 18, 21,
20, 21, 21, 21, 21, 21, 22, 25, 25, 25, 26, 27, 28, 28, 29, 30, 31, 32, 32, 32,
32, 32, 32, 32, 33, 33, 33, 35, 36, 36, 33, 36, 35, 36, 36, 35, 36, 36, 37, 38,
38, 39, 40, 41, 42, 38, 39, 40, 41, 42, 40, 46, 42, 43, 46, 46, 46, 46, 48, 48,
48, 48, 49, 49, 48, 53, 54, 51, 52, 53, 54, 55, 56, 57, 58, 59, 59, 60, 62, 62,
63, 64, 64, 64, 64, 64, 64, 64, 64, 65, 65, 65, 65, 65, 66, 67, 65, 66, 67, 66,
69, 70, 66, 67, 66, 69, 70, 69, 70, 70, 71, 72, 71, 72, 72, 74, 74, 75, 72, 72,
74, 74, 75, 72, 72, 74, 75, 75, 72, 72, 74, 75, 75, 77, 77, 79, 80, 80, 82 ]

splits_Bob = [
0, 1, 1, 2, 2, 2, 3, 3, 4, 4, 4, 5, 5, 5, 6, 7, 8, 8, 8, 8, 9, 9, 9, 9, 9, 10,
12, 12, 12, 12, 12, 12, 13, 14, 14, 15, 16, 16, 16, 16, 16, 17, 16, 16, 17, 19,
19, 20, 21, 22, 22, 22, 22, 22, 22, 22, 22, 22, 22, 24, 24, 25, 27, 27, 28, 28,
29, 28, 29, 28, 28, 28, 30, 28, 28, 28, 29, 30, 33, 33, 33, 33, 34, 35, 37, 37,
37, 37, 38, 38, 37, 38, 38, 38, 38, 38, 39, 43, 38, 38, 38, 38, 43, 40, 41, 42,
43, 48, 45, 46, 47, 47, 48, 49, 49, 49, 50, 51, 50, 49, 49, 49, 49, 51, 49, 53,
50, 51, 50, 51, 51, 51, 52, 55, 55, 55, 56, 56, 56, 56, 56, 58, 58, 61, 61, 61,
63, 63, 63, 64, 65, 65, 65, 65, 66, 66, 65, 65, 66, 66, 66, 66, 66, 66, 66, 71,
66, 73, 66, 66, 71, 66, 73, 66, 66, 71, 66, 73, 68, 68, 71, 71, 73, 73, 73, 75,
75, 78, 78, 78, 80, 80, 80, 81, 81, 82, 83, 84, 85, 86, 86, 86, 86, 86, 87, 86,
88, 86, 86, 86, 86, 88, 86, 88, 86, 86, 86, 88, 88, 86, 86, 86, 93, 90, 90, 92,
92, 92, 93, 93, 93, 93, 93, 97, 97, 97, 97, 97, 97 ]

# Arithmetic functions
def j_inv(A, C):
    """
        Computes the j-invariant of a Montgomery curve with projective constant.

        Input:  
        - The projective curve constant (A:C) given by A,C in Fp2.

        Output: 
        - The j-invariant j=256*(A^2-3*C^2)^3/(C^4*(A^2-4*C^2)) of the Montgomery 
          curve B*y^2=x^3+(A/C)*x^2+x or (equivalently) the j-invariant of 
          B'*y^2=C*x^3+A*x^2+C*x.
    """
    
    return 256*(A^2-3*C^2)^3/(C^4*(A^2-4*C^2))

def xDBLADD(XP, ZP, XQ, ZQ, xPQ, A24):
    """
        Carries out a typical step in the Montgomery ladder: a simultaneous 
        doubling and differential addition.

        Input: 
        - The projective Montgomery x-coordinates of xP=XP/ZP and xQ=XQ/ZQ, 
        - the affine difference x(P-Q) and 
        - the Montgomery curve constant A24=(A+2)/4.

        Output: 
        - The projective Montgomery x-coordinates of x(2P)=X2P/Z2P and
          x(Q+P)=XQP/ZQP.
    """
    
    X2P = (XP + ZP)^2 * (XP - ZP)^2
    Z2P = 4*XP*ZP*(4*A24*XP*ZP + XP^2 - 2*XP*ZP + ZP^2)
    XQP = 4*(XP*XQ - ZP*ZQ)^2
    ZQP = 4*xPQ*(XP*ZQ - ZP*XQ)^2
    
    return X2P, Z2P, XQP, ZQP # Total: 6M+4S+8a

def xADD(XP, ZP, XQ, ZQ, xPQ):
    """
        Computes a standard Montgomery differential addition.

        Input: 
        - The projective Montgomery x-coordinates of xP=XP/ZP and xQ=XQ/ZQ, 
        - and the affine difference x(P-Q).

        Output: 
        - The projective Montgomery x-coordinates of x(Q+P)=XQP/ZQP.
    """
    
    XQP = (XP*XQ - ZP*ZQ)^2
    ZQP = xPQ*(XP*ZQ - ZP*XQ)^2
    
    return XQP, ZQP

def xDBL(X, Z, A24, C24):
    """
        This is NOT the stereotypical Montgomery x-only doubling, since it assumes 
        that the curve constant is projective.  

        Input: 
        - The projective Montgomery x-coordinates of xP=XP/ZP and 
        - the Montgomery curve constant A24/C24 = (A/C+2)/4.

        Output: 
        - The projective Montgomery x-coordinates of x(2P)=X2P/Z2P.
    """
    
    X2 = C24*(X-Z)^2 * (X+Z)^2
    Z2 = 4*(C24*X^2 - 2*C24*X*Z + C24*Z^2 + 4*A24*X*Z)*X*Z
    
    return X2, Z2

def xDBLe(XP, ZP, A,C, e):
    """
        This just computes [2^e](X:Z) on the Montgomery curve with projective 
        constant (A:C) via 2^e repeated doublings.

        Input: 
        - The projective Montgomery x-coordinates of xP=XP/ZP and 
        - the Montgomery curve constant A/C.

        Output: 
        - The projective Montgomery x-coordinates of x(2^e*P)=XeP/ZeP.
    """
    
    A24num = C + C
    A24den = A24num + A24num
    A24num = A24num + A
    
    XeP, ZeP = XP, ZP
    
    for i in [1..e]:
        XeP, ZeP = xDBL(XeP, ZeP, A24num, A24den)
        
    return XeP, ZeP

def xDBL_basefield(X, Z, A24, C24):
    """
        This is NOT the stereotypical Montgomery x-only doubling, since it assumes 
        that the curve constant is projective. All computations are over the base 
        field.

        Input: 
        - The projective Montgomery x-coordinates of xP=XP/ZP and 
        - the Montgomery curve constant A24/C24:=(A/C+2)/4.

        Output: 
        - The projective Montgomery x-coordinates of x(2P)=X2P/Z2P.
    """
    
    # NOTE: this function assumes that A24=1, C24=2 are fixed.
    assert A24 == 1
    assert C24 == 2
    
    X2 = 2*(X-Z)^2 * (X+Z)^2
    Z2 = 8*(X^2 + Z^2)*X*Z
    
    return X2, Z2

def xDBLADD_basefield(XP, ZP, XQ, ZQ, xPQ, A24, C24):
    """
        This function carries out a typical step in the Montgomery ladder: 
        simultaneous doubling and differential addition. All computations are over 
        the base field.

        Input: 
        - The projective Montgomery x-coordinates of xP=XP/ZP and xQ=XQ/ZQ, 
        - the affine difference x(P-Q) and 
        - the Montgomery curve constant A24=(A+2)/4.

        Output: 
        - The projective Montgomery x-coordinates of x(2P)=X2P/Z2P and
          x(Q+P)=XQP/ZQP.
    """
    
    # NOTE: this function assumes that A24=1, C24=2 are fixed.
    assert A24 == 1
    assert C24 == 2
    
    X2P = 2*(XP+ZP)^2 * (XP-ZP)^2
    Z2P = 8*(XP^2 + ZP^2)*XP*ZP
    XQP = 4*(XP*XQ - ZP*ZQ)^2
    ZQP = 4*xPQ*(XP*ZQ - ZP*XQ)^2
    
    return X2P, Z2P, XQP, ZQP

def LADDER(x, m, A24, C24, AliceOrBob):
    """
        The legendary Montgomery ladder.

        Input: 
        - The affine x-coordinate of a point on E: B*y^2=x^3+A*x^2+x, 
        - a scalar m, and 
        - the curve constant (A+2)/4.

        Output:
        - The projective x-coordinates of x(mP)=X0/Z0 and x((m+1)P)=X1/Z1.
    """
    bits = m.digits(base=2)
    
    # NOTE: this function assumes that A24=1, C24=2 are fixed.
    A24, C24 = 1, 2
    X0, Z0 = 1, 0  # Initializing with the point at infinity and (x,1)
    X1, Z1 = x, 1
    
    if (AliceOrBob == Alice):
        nbits = eAbits
    else:
        nbits = eBbits
    
    for i in [1..(nbits-len(bits))]:
        X0, Z0, X1, Z1 = xDBLADD_basefield(X0, Z0, X1, Z1, x, A24, C24)

    for i in range(len(bits)-1, -1, -1):
        if bits[i] == 0:
            X0, Z0, X1, Z1 = xDBLADD_basefield(X0, Z0, X1, Z1, x, A24, C24)
        else:
            X1, Z1, X0, Z0 = xDBLADD_basefield(X1, Z1, X0, Z0, x, A24, C24)

    return X0, Z0, X1, Z1

def secret_pt(x, y, m, AliceOrBob):
    """
        Computes key generation ***entirely in the base field*** by exploiting a 
        1-dimensional Montgomery ladder in the trace zero subgroup and recovering 
        the y-coordinate for the addition. All operations below are in the base 
        field Fp.

        Input: 
        - A point P=(x,y) on E in the base field subgroup,
         the point Q=(-x,y*i) on E in the trace-zero subgroup.
        - the scalar m.

        Output: 
        - Field elements RX0,RX1,RZ in Fp such that (RX0+RX1*i)/RZ is the 
          x-coordinate of P+[m]Q.
    """
    
    # NOTE: this function assumes that A24=1, C24=2 are fixed.
    A24, C24 = 1, 2
    X0, Z0, X1, Z1 = LADDER(-x, m, A24, C24, AliceOrBob)
    
    RX0 = (2*y^2*Z0^2*Z1-(Z0-X0*x)*(X0-x*Z0)*Z1+(X0+x*Z0)^2*X1)*(2*y^2*Z0^2*Z1+(Z0-X0*x)*(X0-x*Z0)*Z1-(X0+x*Z0)^2*X1)-4*y^2*Z1^2*Z0*(X0+x*Z0)*(X0-x*Z0)^2
    RX1 = -4*y^2*Z0^2*Z1*(-Z1*Z0*X0+Z1*x*Z0^2+Z1*X0^2*x-Z1*X0*x^2*Z0+X1*X0^2+2*X1*X0*x*Z0+X1*x^2*Z0^2)
    RZ = 4*y^2*Z1^2*Z0^2*(X0-x*Z0)^2
    
    return RX0, RX1, RZ

def LADDER_3_pt(m, xP, xQ, xPQ, A, AliceOrBob):
    """
        This is Algorithm 1 of De Feo, Jao and Plut. It computes P+[m]Q via x-only 
        arithmetic.

        Input: 
        - The three affine points xP,xQ,xPQ (they are affine as they are compressed 
          before transmission over the wire) and 
        - the Montgomery constant A.

        Output: 
        - The projective Montgomery x-coordinates of x(P+[m]Q)=WX/WZ.
    """
    
    bits = m.digits(base=2)
    
    A24num = A + 2  # Tailored for the special xDBL function
    A24 = A24num / 2
    A24 = A24 / 2
    
    UX, UZ = 1, 0  # Initializing with point at infinity (1:0), 
    VX, VZ = xQ, 1 # (xQ:1) and
    WX, WZ = xP, 1 # (xP:1)
    
    if (AliceOrBob == Alice):
        nbits = eAbits
    else:
        nbits = eBbits
    
    for i in [1..(nbits-len(bits))]:
        WX, WZ = xADD(UX, UZ, WX, WZ, xP)
        UX, UZ, VX, VZ = xDBLADD(UX, UZ, VX, VZ, xQ, A24)
    
    for i in range(len(bits)-1, -1, -1):
        if bits[i] == 0:
            WX, WZ = xADD(UX, UZ, WX, WZ, xP)
            UX, UZ, VX, VZ = xDBLADD(UX, UZ, VX, VZ, xQ, A24)
        else:
            UX, UZ = xADD(UX, UZ, VX, VZ, xQ)
            VX, VZ, WX, WZ = xDBLADD(VX, VZ, WX, WZ, xPQ, A24)
    
    return WX, WZ

def get_4_isog(X4, Z4):
    """
        Given a projective point (X4:Z4) of order 4 on a Montgomery curve, this 
        computes the corresponding 4-isogeny.

        Input: 
        - The projective point of order four (X4:Z4).

        Output: 
        - The 4-isogenous Montgomery curve with projective coefficient A/C and 
          the 5 coefficients that are used to evaluate the isogeny at a point 
          (see the next function).
    """
    
    A = 4*X4^4 - 2*Z4^4
    C = Z4^4
    coeff0 = 2*X4*Z4
    coeff1 = X4^2 + Z4^2
    coeff2 = X4^2 - Z4^2
    coeff3 = X4^4
    coeff4 = Z4^4
    
    return A, C, [coeff0, coeff1, coeff2, coeff3, coeff4]

def eval_4_isog(coeff, X, Z):
    """
        Given a 4-isogeny phi defined by the 5 coefficients in coeff (computed in 
        the function get_4_isog), evaluates the isogeny at the point (X:Z) in the 
        domain of the isogeny.

        Input: 
        - The coefficients defining the isogeny, and 
        - the projective point P=(X:Z).

        Output: 
        - The projective point phi(P)=(X:Z) in the codomain. Variables are 
          overwritten because they replace inputs in the routine.
    """
    
    c = coeff
    
    temp = -(c[0]*X-c[1]*Z+c[2]*Z)^2*(-c[4]*c[0]^2*X^2+2*c[4]*c[0]*X*c[1]*Z+2*c[4]*c[0]*X*c[2]*Z-c[4]*c[1]^2*Z^2-2*c[4]*c[1]*Z^2*c[2]-c[4]*c[2]^2*Z^2+c[3]*c[0]^2*X^2-2*c[3]*c[0]*X*c[1]*Z+2*c[3]*c[0]*X*c[2]*Z+c[3]*c[1]^2*Z^2-2*c[3]*c[1]*Z^2*c[2]+c[3]*c[2]^2*Z^2)
    Z = 4*c[4]*(c[0]*X-c[1]*Z-c[2]*Z)^2*(c[0]*X-c[1]*Z)*c[2]*Z
    X = temp
    
    return X, Z

def first_4_isog(X4, Z4, A):
    """
        This is the very first 4-isogeny computed by Alice, which is different
        from all subsequent 4-isogenies because the point (1,..) is already in the 
        kernel, so it doesn't need composition with the preliminary isomorphism.
        (See De Feo, Jao and Plut, Section 4.3).

        Input: 
        - The projective point (X4:Z4) and 
        - the curve constant A (that is affine because it is passed over the wire 
          or a fixed system parameter).

        Output: 
        - The projective point (X4:Z4) in the codomain and
        - the isogenous curve constant A/C. Variables are overwritten because they 
          replace inputs in the routine.
    """
    
    X = (X4+Z4)^2 * (A*X4*Z4 + X4^2 + Z4^2)
    Z = -X4*Z4*(A-2) * (X4-Z4)^2
    C = A - 2
    A = 2*A + 12
    
    return X, Z, A, C

def xTPL(X, Z, A, C):
    """
        This is NOT the stereotypical Montgomery x-only tripling, since it assumes 
        that the curve constant is projective.  

        Input: 
        - The projective Montgomery x-coordinates of xP=X/Z and 
        - the Montgomery curve constant A4:=4*A.

        Output: 
        - The projective Montgomery x-coordinates of x(3P)=X3/Z3.
    """
    
    X3 = X^2*(-C*X^4 + 6*C*X^2*Z^2 + 3*C*Z^4 + Z^3*4*A*X)^2
    Z3 = X*Z*(-C*Z^4 + 3*C*X^4 + 6*C*X^2*Z^2 + X^3*4*A*Z)^2
    
    return X3, Z3

def xTPLe(X, Z, A, C, e):
    """
        This function just computes [3^e](X:Z) on a Montgomery curve with 
        projective constant via 3^e repeated triplings.  

        Input: 
        - The projective Montgomery x-coordinates of xP=X/Z and 
        - the Montgomery curve constant A/C.

        Output: 
        - The projective Montgomery x-coordinates of x(eP)=XeP/ZeP.
    """
    
    XeP, ZeP = X, Z
    
    for i in [1..e]:
        XeP, ZeP = xTPL(XeP, ZeP, A, C)
    
    return XeP, ZeP

def get_3_isog(X3, Z3):
    """
        Given a projective point (X3:Z3) of order 3 on a Montgomery curve, this 
        computes the corresponding 3-isogenous curve.

        Input: 
        - The projective point of order three (X3:Z3).

        Output: 
        - The 3-isogenous Montgomery curve with projective coefficient A/C. 
          No coefficients are computed for the evaluation phase as all operations 
          in the evaluation depend on the input point to the isogeny.
    """
    
    A = Z3^4 - 27*X3^4 + 18*X3^2*Z3^2
    C = 4*Z3^3*X3
    
    return A, C

def eval_3_isog(X3, Z3, X, Z):
    """
        Given a projective point (X3:Z3) of order 3 on a Montgomery curve and a 
        projective point x(P)=(X:Z), this function evaluates the corresponding 
        3-isogeny at x(P): phi(X:Z).

        Input:
        - The projective point (X3:Z3) of order three,
        - the projective Montgomery x-coordinates of x(P)=X/Z.

        Output: 
        - The projective Montgomery x-coordinates of the evaluation of phi at
          (X:Z).
    """
    
    temp = X*(X3*X - Z3*Z)^2
    Z = Z*(Z3*X - Z*X3)^2
    X = temp
    
    return X, Z

def inv_4_way(z1, z2, z3, z4):
    """
        This function computes inverses of four elements by sharing the inversions
        via Montgomery's simultaneous inversion trick.

        Input: 
        - The four values to be inverted: z1,z2,z3,z4.

        Output: 
        - Their inverses 1/z1,1/z2,1/z3,1/z4 (over-ride variables).
    """
    
    return 1/z1, 1/z2, 1/z3, 1/z4

def distort_and_diff(xP):
    """
        Given the x-coordinate of an affine point P, this function returns the
        projective x-coordinates of the difference point Q-P, where Q=tau(P) is 
        the image under the distortion map of the point P. 

        Input: 
        - The coordinate xP of the point P=(xP,yP).

        Output: 
        - The point (x(Q-P),z(Q-P)), where Q=tau(P).
    """
    
    XD = j*(xP^2 + 1)
    ZD = 2*xP
    
    return XD, ZD

def get_A(xP, xQ, xR):
    """
        Given the x-coordinates of P, Q, and R, returns the value A 
        (corresponding to the Montgomery curve E_A: y^2=x^3+A*x^2+x) 
        such that R=Q-P on E_A

        Input: 
        - The x-coordinates xP, xQ, and xR of the points P, Q and R

        Output: 
        - The coefficient A corresponding to the curve E_A: y^2=x^3+A*x^2+x
    """
    
    A = (xR*xP+xR*xQ+xP*xQ-1)^2 / (4*xP*xQ*xR) - (xP+xQ+xR)
    
    return A

# Alice key generation Weierstrass/affine/Velu/Sage/slow
def keygen_alice_weierstrass(SK_Alice, params):
    """
        This function generates Alice's public key from her secret key and the 
        public scheme parameters. It uses the simple but costly loop for traversing 
        the isogeny tree; it uses EllipticCurveWithTorsionPoint class to compute
        and evaluate the isogenies in affine Weierstrass form.

        Input: 
        - Alice's secret key SK_Alice, which is a random even number between 1 and 
          oA-1,
        - three public parameters params=[XPB,XPA,YPA]: the x-coordinate of PB, 
          and both coordinates of PA.

        Output: 
        - Alice's public key [A,phi_A(x(PB)),phi_A(x(QB)),phi_A(x(QB-PB))].
    """

    E = EllipticCurve(Fp2, [1,0]) # Weierstrass curve y^2=x^3+x
    R.<x> = PolynomialRing(Fp2)   # for computing isogeny kernels
    phiP = E0([params[2], params[3]])
    phiQ = E0([-params[2], j*params[3]])
    
    P = E0([params[0], params[1]])
    Q = E0([-params[0], j*params[1]])
    R = P + SK_Alice * Q

    points = [phiP, phiQ]
    R_list = []
    degree_list = []
    S = R
    while S != E(0):
        R_list.append(S)
        degree_list.append(2)
        S = 2*S

    X = (E, R_list, points, degree_list)
    while X[1]:
        X = PointIsogeny(*X)
    
    E, _, points, _ = X
    phiP, phiQ = points
    a, b = E.a4(), E.a6()

    PK_Alice = [a, b, phiP[0], phiP[1], phiQ[0], phiQ[1]] # 6 values in Fp2
    
    return PK_Alice

# Alice key generation Kummer/projective/fast
def keygen_alice_fast(SK_Alice, params, splits, MAX):
    """
        This function generates Alice's public key from her secret key and the
        public scheme parameters. It uses the optimal way of traversing the
        isogeny tree as described by De Feo, Jao and Plut. 

        Input: 
        - Alice's secret key SK_Alice, which is a random even number between 1 and 
          oA-1,
        - three public parameters params=[XPB,XPA,YPA]: the x-coordinate of PB, 
          and both coordinates of PA,
        - the parameter "splits", a vector that guides the optimal route through
          the isogeny tree; it is generated individually for Alice using
          "optimalstrategies.mag" and the ratios of 4-isogeny evaluation versus 
          multiplication-by-4,
        - the parameter "MAX", the maximum number of multiplication-by-4
          computations.

        Output: 
        - Alice's public key [phi_A(x(PB)),phi_A(x(QB)),phi_A(x(QB-PB))].
    """
    
    A, C = 0, 1 # The starting Montgomery curve (A:C) = (0:1)
    phiPX, phiPZ = params[0], 1
    phiQX, phiQZ = -phiPX, 1 # Q=(-xP,yP), tau(P) but yP instead of yP*i, the "*i" is handled implicitly
    phiDX, phiDZ = distort_and_diff(phiPX) # (phiDX:phiDZ):=x(Q-P)
    
    # Computes x(R)=(RX:RZ) via secret_pt function
    RX0, RX1, RZ = secret_pt(params[1], params[2], SK_Alice, Alice)
    RX = RX0 + RX1*j
    
    # The first iteration is different so not in the main loop
    phiPX, phiPZ, _, _ = first_4_isog(phiPX, phiPZ, A)
    phiQX, phiQZ, _, _ = first_4_isog(phiQX, phiQZ, A)
    phiDX, phiDZ, _, _ = first_4_isog(phiDX, phiDZ, A)
    RX, RZ, A, C = first_4_isog(RX, RZ, A)
    
    pts = []
    index = 0
    
    # Alice's main loop
    for row in [1..MAX-1]:
        # Multiply (RX:RZ) until it has order 4, and store intermediate points
        while index < (MAX - row):
            pts.append([RX, RZ, index])
            m = splits[MAX-index-row]
            RX, RZ = xDBLe(RX, RZ, A, C, 2*m)
            index += m
        
        # Compute the 4-isogeny based on kernel (RX:RZ)
        A, C, consts = get_4_isog(RX, RZ)
        
        # Evaluate the 4-isogeny at every point in pts
        for i in [0..len(pts)-1]:
            pts[i][0], pts[i][1] = eval_4_isog(consts, pts[i][0], pts[i][1])
        
        # Evaluate the 4-isogeny at Bob's (intermediate) points 
        # x(P), x(Q), x(Q-P)
        phiPX, phiPZ = eval_4_isog(consts, phiPX, phiPZ)
        phiQX, phiQZ = eval_4_isog(consts, phiQX, phiQZ)
        phiDX, phiDZ = eval_4_isog(consts, phiDX, phiDZ)
        
        # R becomes the last point in pts and then pts is pruned
        RX = pts[len(pts)-1][0]
        RZ = pts[len(pts)-1][1]
        index = ZZ(pts[len(pts)-1][2])
        
        pts.pop()
    
    # Compute and evaluate the last 4-isogeny
    A, C, consts = get_4_isog(RX, RZ)
    phiPX, phiPZ = eval_4_isog(consts, phiPX, phiPZ)
    phiQX, phiQZ = eval_4_isog(consts, phiQX, phiQZ)
    phiDX, phiDZ = eval_4_isog(consts, phiDX, phiDZ)
    
    # Normalize everything via a 3-way simultaneous inversion
    phiPX = phiPX / phiPZ
    phiQX = phiQX / phiQZ
    phiDX = phiDX / phiDZ
    
    PK_Alice = [phiPX, phiQX, phiDX] # 3 values in Fp2
    
    return PK_Alice

# Bob key generation Weierstrass/affine/Velu/Sage/slow
def keygen_bob_weierstrass(SK_Bob, params):
    """
        This function generates Bob's public key from his secret key and the 
        public scheme parameters. It uses the simple but costly loop for traversing 
        the isogeny tree; it uses EllipticCurveWithTorsionPoint class to compute 
        and evaluate the isogenies in affine Weierstrass form.

        Input: 
        - Bob's secret key SK_Bob, which is a random value between 1 and oB-1, 
        - three public parameters params=[XPA,XPB,YPB]: the x-coordinate of PA, and 
          both coordinates of PB.

        Output: 
        - Bob's public key [A,phi_B(x(PA)),phi_B(x(QA)),phi_B(x(QA-PA))].
    """

    E = EllipticCurve(Fp2, [1,0]) # Weierstrass curve y^2=x^3+x
    R.<x> = PolynomialRing(Fp2)   # for computing isogeny kernels
    phiP = E0([params[0], params[1]])
    phiQ = E0([-params[0], j*params[1]])
    
    P = E0([params[2], params[3]])
    Q = E0([-params[2], j*params[3]])
    R = P + SK_Bob * Q
    
    points = [phiP, phiQ]
    R_list = []
    degree_list = []
    S = R
    while S != E(0):
        R_list.append(S)
        degree_list.append(3)
        S = 3*S

    X = (E, R_list, points, degree_list)
    while X[1]:
        X = PointIsogeny(*X)
    
    E, _, points, _ = X
    phiP, phiQ = points
    a, b = E.a4(), E.a6()
    
    PK_Bob = [a, b, phiP[0], phiP[1], phiQ[0], phiQ[1]] # 6 values in Fp2
    
    return PK_Bob

# Bob key generation Kummer/projective/fast
def keygen_bob_fast(SK_Bob, params, splits, MAX):
    """
        This function generates Bob's public key from his secret key and the 
        public scheme parameters. It uses the optimal way of traversing the
        isogeny tree as described by De Feo, Jao and Plut. 

        Input: 
        - Bob's secret key SK_Bob, which is a random value between 1 and oB-1, 
        - three public parameters params=[XPA,XPB,YPB]: the x-coordinate of PA, and 
          both coordinates of PB.
        - the parameter "splits", a vector that guides the optimal route through the 
          isogeny tree; it is generated individually for Bob using 
          "optimalstrategies.m" and the ratios of 3-isogeny evaluation versus 
          multiplication-by-3,
        - the parameter "MAX", the maximum number of multiplication-by-3
          computations.

        Output: 
        - Bob's public key [phi_B(x(PA)),phi_B(x(QA)),phi_B(x(QA-PA))].
    """
    
    A, C = 0, 1 # The starting Montgomery curve (A:C) = (0:1)
    phiPX, phiPZ = params[0], 1
    phiQX, phiQZ = -phiPX, 1 # Q=(-xP,yP), tau(P) but yP instead of yP*i, the "*i" is handled implicitly
    phiDX, phiDZ = distort_and_diff(phiPX) # (phiDX:phiDZ):=x(Q-P)
    
    # Computes x(R)=(RX:RZ) via secret_pt function
    RX0, RX1, RZ = secret_pt(params[1], params[2], SK_Bob, Bob)
    RX = RX0 + RX1*j
    
    pts = []
    index = 0
    
    # Bob's main loop
    for row in [1..MAX-1]:
        # Multiply (RX:RZ) until it has order 3, and store intermediate points
        while index < (MAX - row):
            pts.append([RX, RZ, index])
            m = splits[MAX-index-row]
            RX, RZ = xTPLe(RX, RZ, A, C, m)
            index += m
    
        # Compute the 3-isogeny based on kernel (RX:RZ)
        A, C = get_3_isog(RX, RZ)

        # Evaluate the 3-isogeny at every point in pts
        for i in [0..len(pts)-1]:
            pts[i][0], pts[i][1] = eval_3_isog(RX, RZ, pts[i][0], pts[i][1])

        # Evaluate the 3-isogeny at Alice's (intermediate) points
        # x(P), x(Q), x(Q-P)
        phiPX, phiPZ = eval_3_isog(RX, RZ, phiPX, phiPZ)
        phiQX, phiQZ = eval_3_isog(RX, RZ, phiQX, phiQZ)
        phiDX, phiDZ = eval_3_isog(RX, RZ, phiDX, phiDZ)

        # R becomes the last point in pts and then pts is pruned
        RX = pts[len(pts)-1][0]
        RZ = pts[len(pts)-1][1]
        index = ZZ(pts[len(pts)-1][2])

        pts.pop()

    # Compute and evaluate the last 3-isogeny
    A, C = get_3_isog(RX, RZ)
    phiPX, phiPZ = eval_3_isog(RX, RZ, phiPX, phiPZ)
    phiQX, phiQZ = eval_3_isog(RX, RZ, phiQX, phiQZ)
    phiDX, phiDZ = eval_3_isog(RX, RZ, phiDX, phiDZ)
    
    # Normalize everything via a 3-way simultaneous inversion
    phiPX = phiPX / phiPZ
    phiQX = phiQX / phiQZ
    phiDX = phiDX / phiDZ
    
    PK_Bob = [phiPX, phiQX, phiDX] # 4 values in Fp2
    
    return PK_Bob

# Alice's shared secret Weierstrass/affine/Velu/Sage/slow
def shared_secret_alice_weierstrass(SK_Alice, PK_Bob):
    """
        This function generates Alice's shared secret from her secret key and Bob's 
        public key. It uses the simple but costly loop for traversing the isogeny 
        tree; it uses EllipticCurveWithTorsionPoint class to compute and evaluate
        the isogenies in affine Weierstrass form.

        Input: 
        - Alice's secret key SK_Alice, a random even number between 1 and oA-1,
        - Bob's public key PK_Bob=[A,phi_B(x(PA)),phi_B(x(QA)),phi_B(x(QA-PA))].

        Output: 
        - Alice's shared secret: the j-invariant of E_AB.
    """
    
    a, b = PK_Bob[0], PK_Bob[1]
    E = EllipticCurve(Fp2, [a,b]) 
    R.<x> = PolynomialRing(Fp2)   # For computing isogeny kernels
    
    P = E([PK_Bob[2], PK_Bob[3]])
    Q = E([PK_Bob[4], PK_Bob[5]])
    R = P + SK_Alice * Q

    points = []
    R_list = []
    degree_list = []
    S = R
    while S != E(0):
        R_list.append(S)
        degree_list.append(2)
        S = 2*S

    X = (E, R_list, points, degree_list)
    while X[1]:
        X = PointIsogeny(*X)
    
    E, _, _, _ = X    
    shared_secret_alice = E.j_invariant()
    
    return shared_secret_alice

# Alice shared secret Kummer/projective/fast
def shared_secret_alice_fast(SK_Alice, PK_Bob, params, splits, MAX):
    """
        This function generates Alice's shared secret from her secret key and Bob's
        public key. It uses the optimal way of traversing the isogeny tree as
        described by De Feo, Jao and Plut. 

        Input: 
        - Alice's secret key SK_Alice, a random even number between 1 and oA-1,
        - Bob's public key PK_Bob=[phi_B(x(PA)),phi_B(x(QA)),phi_B(x(QA-PA))],
        - the parameter "splits", a vector that guides the optimal route through the 
          isogeny tree; it is generated individually for Alice using 
          "optimalstrategies.m" and the ratios of 4-isogeny evaluation versus 
          multiplication-by-4,
        - the parameter "MAX", the maximum number of multiplication-by-4
          computations.

        Output: 
        - Alice's shared secret: the j-invariant of E_AB.
    """
    
    A = get_A(PK_Bob[0], PK_Bob[1], PK_Bob[2])
    C = 1 # Starting on Bob's Montgomery curve
    
    # Computes R=phi_B(xPA)+SK_Alice*phi_B(xQA) via 3 point ladder
    RX, RZ = LADDER_3_pt(SK_Alice, PK_Bob[0], PK_Bob[1], PK_Bob[2], A, Alice)
    
    # The first iteration is different so not in the main loop
    RX, RZ, A, C = first_4_isog(RX, RZ, A)
    
    pts = []
    index = 0
    
    # Alice's main loop
    for row in [1..MAX-1]:
        # Multiply (RX:RZ) until it has order 4, and store intermediate points
        while index < (MAX - row):
            pts.append([RX, RZ, index])
            m = splits[MAX-index-row]
            RX, RZ = xDBLe(RX, RZ, A, C, 2*m)
            index += m
        
        # Compute the 4-isogeny based on kernel (RX:RZ)
        A, C, consts = get_4_isog(RX, RZ)
        
        # Evaluate the 4-isogeny at every point in pts
        for i in [0..len(pts)-1]:
            pts[i][0], pts[i][1] = eval_4_isog(consts, pts[i][0], pts[i][1])
        
        # R becomes the last point in pts and then pts is pruned
        RX = pts[len(pts)-1][0]
        RZ = pts[len(pts)-1][1]
        index = ZZ(pts[len(pts)-1][2])
        
        pts.pop()
    
    # Compute the last 4-isogeny
    A, C, consts = get_4_isog(RX, RZ)
    
    # Compute the j-invariant of E_AB
    shared_secret_alice = j_inv(A, C)
    
    return shared_secret_alice

# Bob's shared secret Weierstrass/affine/Velu/Sage/slow
def shared_secret_bob_weierstrass(SK_Bob, PK_Alice):
    """
        This function generates Bob's shared secret from his secret key and Alice's 
        public key. It uses the simple but costly loop for traversing the isogeny 
        tree; it uses EllipticCurveWithTorsionPoint class to compute and evaluate
        the isogenies in affine Weierstrass form.

        Input: 
        - Bob's secret key SK_Bob, a random number between 1 and oB-1,
        - Alice's public key PK_Alice=[A,phi_A(x(PB)),phi_A(x(QB)),phi_A(x(QB-PB))].

        Output: 
        - Bob's shared secret: the j-invariant of E_BA.
    """
    
    a, b = PK_Alice[0], PK_Alice[1]
    E = EllipticCurve(Fp2, [a,b]) 
    R.<x> = PolynomialRing(Fp2)   # For computing isogeny kernels
    
    P = E([PK_Alice[2], PK_Alice[3]])
    Q = E([PK_Alice[4], PK_Alice[5]])
    R = P + SK_Bob * Q
    
    points = []
    R_list = []
    degree_list = []
    S = R
    while S != E(0):
        R_list.append(S)
        degree_list.append(3)
        S = 3*S

    X = (E, R_list, points, degree_list)
    while X[1]:
        X = PointIsogeny(*X)

    E, _, _, _ = X
    shared_secret_bob = E.j_invariant()
    
    return shared_secret_bob

# Bob's shared secret Kummer/projective/fast
def shared_secret_bob_fast(SK_Bob, PK_Alice, params, splits, MAX):
    """
        This function generates Bob's shared secret from his secret key and Alice's
        public key. It uses the optimal way of traversing the isogeny tree as
        described by De Feo, Jao and Plut. 

        Input: 
        - Bob's secret key SK_Bob, a random number between 1 and oB-1,
        - Alice's public key PK_Alice=[phi_A(xPB),phi_A(xQB),phi_A(x(QB-PB))],
        - the parameter "splits", a vector that guides the optimal route through the 
          isogeny tree; it is generated individually for Bob using 
          "optimalstrategies.m" and the ratios of 3-isogeny evaluation versus 
          multiplication-by-3,
        - the parameter "MAX", the maximum number of multiplication-by-3
          computations.

        Output: 
        - Bob's shared secret: the j-invariant of E_BA.
    """
    
    A = get_A(PK_Alice[0], PK_Alice[1], PK_Alice[2])
    C = 1 # Starting on Alice's Montgomery curve
    
    # Computes R=phi_A(xPB)+SK_Bob*phi_A(xQB) via 3 point ladder
    RX, RZ = LADDER_3_pt(SK_Bob, PK_Alice[0], PK_Alice[1], PK_Alice[2], A, Bob)
    
    pts = []
    index = 0
    
    # Bob's main loop
    for row in [1..MAX-1]:
        # Multiply (RX:RZ) until it has order 3, and store intermediate points
        while index < (MAX - row):
            pts.append([RX, RZ, index])
            m = splits[MAX-index-row]
            RX, RZ = xTPLe(RX, RZ, A, C, m)
            index += m
        
        # Compute the 3-isogeny based on kernel (RX:RZ)
        A, C = get_3_isog(RX, RZ)
        
        # Evaluate the 3-isogeny at every point in pts
        for i in [0..len(pts)-1]:
            pts[i][0], pts[i][1] = eval_3_isog(RX, RZ, pts[i][0], pts[i][1])
        
        # R becomes the last point in pts and then pts is pruned
        RX = pts[len(pts)-1][0]
        RZ = pts[len(pts)-1][1]
        index = ZZ(pts[len(pts)-1][2])
        
        pts.pop()
    
    # Compute the last 3-isogeny
    A, C = get_3_isog(RX, RZ)
    
    # Compute the j Invariant of E_BA
    shared_secret_bob = j_inv(A, C)
    
    return shared_secret_bob

# Key exchange testing
MAX_Alice = 185
MAX_Bob = 239

params_Alice = [XPB, XPA, YPA]
params_Bob = [XPA, XPB, YPB]
params_both_Weierstrass = [XPA, YPA, XPB, YPB]

print "===================================================================="
print "Generating secret keys..."
SK_Alice = randint(1, (oA // lA) - 1)*lA # Random even number between 1 and oA-1 
SK_Bob = randint(1, (oB // lB) - 1)*lB   # Random multiple of lB between 1 and oB-1 
print "Done with secret keys."
print "===================================================================="

print "Generating Alice's public key (Weierstrass algorithm)..."
PK_Alice_weierstrass = keygen_alice_weierstrass(SK_Alice, params_both_Weierstrass)
print "Generating Alice's public key (fast algorithm)..."
PK_Alice_fast = keygen_alice_fast(SK_Alice, params_Alice, splits_Alice, MAX_Alice)
print "Done with Alice's public key.\n"

print "Test result from Alice's Weierstrass key generation equal to result from Alice's fast key generation?"
# Weierstrass curve generated by Sage in Alice's public key
a, b = PK_Alice_weierstrass[0], PK_Alice_weierstrass[1]
E = EllipticCurve(Fp2, [a,b]) 

# Montgomery curve generated by our code in Alice's public key
A = get_A(PK_Alice_fast[0], PK_Alice_fast[1], PK_Alice_fast[2])
R.<x> = PolynomialRing(Fp2)
EA = EllipticCurve(Fp2, [0, A, 0, 1, 0])

print EA.j_invariant() == E.j_invariant()

print "===================================================================="

print "Generating Bob's public key (Weierstrass algorithm)..."
PK_Bob_weierstrass = keygen_bob_weierstrass(SK_Bob, params_both_Weierstrass)
print "Generating Bob's public key (fast algorithm)..."
PK_Bob_fast = keygen_bob_fast(SK_Bob, params_Bob, splits_Bob, MAX_Bob)
print "Done with Bob's public key.\n"

print "Test result from Bob's Weierstrass key generation equal to result from Bob's fast key generation?"
# Weierstrass curve generated by Sage in Bob's public key
a, b = PK_Bob_weierstrass[0], PK_Bob_weierstrass[1]
E = EllipticCurve(Fp2, [a,b])

# Montgomery curve generated by our code in Bob's public key
A = get_A(PK_Bob_fast[0], PK_Bob_fast[1], PK_Bob_fast[2])
R.<x> = PolynomialRing(Fp2)
EB = EllipticCurve(Fp2, [0, A, 0, 1, 0])

print EB.j_invariant() == E.j_invariant()

print "===================================================================="

print "Generating shared secret for Alice (Weierstrass algorithm)..."
secret_alice_weierstrass = shared_secret_alice_weierstrass(SK_Alice, PK_Bob_weierstrass)
print "Generating shared secret for Alice (fast algorithm)..."
secret_alice_fast = shared_secret_alice_fast(SK_Alice, PK_Bob_fast, params_Alice, splits_Alice, MAX_Alice)
print "Done with Alice's shared secret computation.\n"

print "Results from Weierstrass and fast algorithms equal?", secret_alice_weierstrass == secret_alice_fast

print "===================================================================="

print "Generating shared secret for Bob (Weierstrass algorithm)..."
secret_bob_weierstrass = shared_secret_bob_weierstrass(SK_Bob, PK_Alice_weierstrass)
print "Generating shared secret for Bob (fast algorithm)..."
secret_bob_fast = shared_secret_bob_fast(SK_Bob, PK_Alice_fast, params_Bob, splits_Bob, MAX_Bob)
print "Done with Alice's shared secret computation.\n"

print "Results from Weierstrass and fast algorithms equal?", secret_bob_weierstrass == secret_bob_fast

print "===================================================================="
print "Shared secrets are equal?", secret_alice_fast == secret_bob_fast
print "===================================================================="