# This file contains two Inner Product Argument implementations:
#   - Bulletproofs version: https://eprint.iacr.org/2017/1066.pdf
#   - Halo version: https://eprint.iacr.org/2019/1021.pdf



# IPA_bulletproofs implements the IPA version from the Bulletproofs paper: https://eprint.iacr.org/2017/1066.pdf
# https://doc-internal.dalek.rs/bulletproofs/notes/inner_product_proof/index.html
class IPA_bulletproofs:
    def __init__(self, F, E, g, d):
        self.g = g
        self.F = F
        self.E = E
        self.d = d
        # TODO:
        # Setup:
        self.h = E.random_element() # TMP
        self.gs = random_values(E, d)
        self.hs = random_values(E, d)

    # a: aᵢ ∈ 𝔽 coefficients of p(X)
    # r: blinding factor
    def commit(self, a, b):
        P = inner_product_point(a, self.gs) + inner_product_point(b, self.hs)
        return P
    
    def evaluate(self, a, x_powers):
        return inner_product_field(a, x_powers)

    def ipa(self, a_, b_, u, U):
        G = self.gs
        H = self.hs
        a = a_
        b = b_

        k = int(math.log(self.d, 2))
        L = [None] * k
        R = [None] * k

        for j in reversed(range(0, k)):
            m = len(a)/2
            a_lo = a[:m]
            a_hi = a[m:]
            b_lo = b[:m]
            b_hi = b[m:]
            H_lo = H[:m]
            H_hi = H[m:]
            G_lo = G[:m]
            G_hi = G[m:]

            #  Lⱼ = <a'ₗₒ, G'ₕᵢ> + [lⱼ] H + [<a'ₗₒ, b'ₕᵢ>] U
            L[j] = inner_product_point(a_lo, G_hi) + inner_product_point(b_hi, H_lo) + int(inner_product_field(a_lo, b_hi)) * U
            #  Rⱼ = <a'ₕᵢ, G'ₗₒ> + [rⱼ] H + [<a'ₕᵢ, b'ₗₒ>] U
            R[j] = inner_product_point(a_hi, G_lo) + inner_product_point(b_lo, H_hi) + int(inner_product_field(a_hi, b_lo)) * U

            # use the random challenge uⱼ ∈ 𝕀 generated by the verifier
            u_ = u[j] # uⱼ
            u_inv = u[j]^(-1) # uⱼ⁻¹

            a = vec_add(vec_scalar_mul_field(a_lo, u_), vec_scalar_mul_field(a_hi, u_inv))
            b = vec_add(vec_scalar_mul_field(b_lo, u_inv), vec_scalar_mul_field(b_hi, u_))
            G = vec_add(vec_scalar_mul_point(G_lo, u_inv), vec_scalar_mul_point(G_hi, u_))
            H = vec_add(vec_scalar_mul_point(H_lo, u_), vec_scalar_mul_point(H_hi, u_inv))

        assert len(a)==1
        assert len(b)==1
        assert len(G)==1
        assert len(H)==1
        # a, b, G have length=1
        # L, R are the "cross-terms" of the inner product
        return a[0], b[0], G[0], H[0], L, R

    def verify(self, P, a, v, x_powers, u, U, L, R, b_ipa, G_ipa, H_ipa):
        b = b_ipa
        G = G_ipa
        H = H_ipa

        # Q_0 = P' ⋅ ∑ ( [uⱼ²] Lⱼ + [uⱼ⁻²] Rⱼ)
        C = P
        for j in range(len(L)):
            u_ = u[j] # uⱼ
            u_inv = u[j]^(-1) # uⱼ⁻²

            # ∑ ( [uⱼ²] Lⱼ + [uⱼ⁻²] Rⱼ)
            C = C + int(u_^2) * L[j] + int(u_inv^2) * R[j]

        D = int(a) * G + int(b) * H + int(a * b)*U

        return C == D

# IPA_halo implements the modified IPA from the Halo paper: https://eprint.iacr.org/2019/1021.pdf
class IPA_halo:
    def __init__(self, F, E, g, d):
        self.g = g
        self.F = F
        self.E = E
        self.d = d

        self.h = E.random_element() # TMP
        self.gs = random_values(E, d)
        self.hs = random_values(E, d)
        #  print("  h=", self.h)
        #  print("  G=", self.gs)
        #  print("  H=", self.hs)

    def commit(self, a, r):
        P = inner_product_point(a, self.gs) + r * self.h
        return P
    
    def evaluate(self, a, x_powers):
        return inner_product_field(a, x_powers)

    def ipa(self, a_, x_powers, u, U): # prove
        print(" method ipa():")
        G = self.gs
        a = a_
        b = x_powers

        k = int(math.log(self.d, 2))
        l = [None] * k
        r = [None] * k
        L = [None] * k
        R = [None] * k

        for j in reversed(range(0, k)):
            print(" j =", j)
            print("    len(a) = n =", len(a))
            print("    m = n/2 =", len(a)/2)
            m = len(a)/2
            a_lo = a[:m]
            a_hi = a[m:]
            b_lo = b[:m]
            b_hi = b[m:]
            G_lo = G[:m]
            G_hi = G[m:]

            print("  Split into a_lo,hi b_lo,hi, G_lo,hi:")
            print("    a", a)
            print("      a_lo", a_lo)
            print("      a_hi", a_hi)
            print("    b", b)
            print("      b_lo", b_lo)
            print("      b_hi", b_hi)
            print("    G", G)
            print("      G_lo", G_lo)
            print("      G_hi", G_hi)

            l[j] = self.F.random_element() # random blinding factor
            r[j] = self.F.random_element() # random blinding factor
            print("  random blinding factors:")
            print("    l[j]", l[j])
            print("    r[j]", r[j])

            #  Lⱼ = <a'ₗₒ, G'ₕᵢ> + [lⱼ] H + [<a'ₗₒ, b'ₕᵢ>] U
            L[j] = inner_product_point(a_lo, G_hi) + int(l[j]) * self.h + int(inner_product_field(a_lo, b_hi)) * U
            #  Rⱼ = <a'ₕᵢ, G'ₗₒ> + [rⱼ] H + [<a'ₕᵢ, b'ₗₒ>] U
            R[j] = inner_product_point(a_hi, G_lo) + int(r[j]) * self.h + int(inner_product_field(a_hi, b_lo)) * U

            print("  Compute Lⱼ = <a'ₗₒ, G'ₕᵢ> + [lⱼ] H + [<a'ₗₒ, b'ₕᵢ>] U")
            print("    L[j]", L[j])
            print("  Compute Rⱼ = <a'ₕᵢ, G'ₗₒ> + [rⱼ] H + [<a'ₕᵢ, b'ₗₒ>] U")
            print("    R[j]", R[j])

            # use the random challenge uⱼ ∈ 𝕀 generated by the verifier
            u_ = u[j] # uⱼ
            u_inv = self.F(u[j])^(-1) # uⱼ⁻¹
            print("  u_j", u_)
            print("  u_j^-1", u_inv)

            a = vec_add(vec_scalar_mul_field(a_lo, u_), vec_scalar_mul_field(a_hi, u_inv))
            b = vec_add(vec_scalar_mul_field(b_lo, u_inv), vec_scalar_mul_field(b_hi, u_))
            G = vec_add(vec_scalar_mul_point(G_lo, u_inv), vec_scalar_mul_point(G_hi, u_))
            print("  new a, b, G")
            print("    a =", a)
            print("    b =", b)
            print("    G =", G)

        assert len(a)==1
        assert len(b)==1
        assert len(G)==1
        # a, b, G have length=1
        # l, r are random blinding factors
        # L, R are the "cross-terms" of the inner product
        return a[0], b[0], G[0], l, r, L, R

    def verify(self, P, a, v, x_powers, r, u, U, lj, rj, L, R):
        print("methid verify()")

        s = build_s_from_us(u, self.d)
        b = inner_product_field(s, x_powers)
        G = inner_product_point(s, self.gs)

        # synthetic blinding factor
        # r' = r + ∑ ( lⱼ uⱼ² + rⱼ uⱼ⁻²)
        print(" synthetic blinding factor r' = r + ∑ ( lⱼ uⱼ² + rⱼ uⱼ⁻²)")
        r_ = r
        print("  r_ =", r_)
        # Q_0 = P' ⋅ ∑ ( [uⱼ²] Lⱼ + [uⱼ⁻²] Rⱼ)
        print(" Q_0 = P' ⋅ ∑ ( [uⱼ²] Lⱼ + [uⱼ⁻²] Rⱼ)")
        Q_0 = P
        print("  Q_0 =", Q_0)
        for j in range(len(u)):
            print(" j =", j)
            u_ = u[j] # uⱼ
            u_inv = u[j]^(-1) # uⱼ⁻²

            # ∑ ( [uⱼ²] Lⱼ + [uⱼ⁻²] Rⱼ)
            Q_0 = Q_0 + int(u[j]^2) * L[j] + int(u_inv^2) * R[j]
            print("  Q_0 =", Q_0)

            r_ = r_ + lj[j] * (u_^2) + rj[j] * (u_inv^2)
            print("  r_ =", r_)

        Q_1 = int(a) * G + int(r_) * self.h + int(a * b)*U
        print(" Q_1", Q_1)
        #  Q_1_ = int(a) * (G + int(b)*U) + int(r_) * self.h

        return Q_0 == Q_1


# s = (
#   u₁⁻¹ u₂⁻¹ … uₖ⁻¹,
#   u₁   u₂⁻¹ … uₖ⁻¹,
#   u₁⁻¹ u₂   … uₖ⁻¹,
#   u₁   u₂   … uₖ⁻¹,
#   ⋮    ⋮      ⋮
#   u₁   u₂   … uₖ
# )
def build_s_from_us(u, d):
    k = int(math.log(d, 2))
    s = [1]*d
    t = d
    for j in reversed(range(k)):
        t = t/2
        c = 0
        for i in range(d):
            if c<t:
                s[i] = s[i] * u[j]^(-1)
            else:
                s[i] = s[i] * u[j]
            c = c+1
            if c>=t*2:
                c=0

    return s




def powers_of(g, d):
    r = [None] * d
    for i in range(d):
        r[i] = g^i
    return r


def multiples_of(g, d):
    r = [None] * d
    for i in range(d):
        r[i] = g*i
    return r

def random_values(G, d):
    r = [None] * d
    for i in range(d):
        r[i] = G.random_element()
    return r

def inner_product_field(a, b):
    assert len(a) == len(b)
    c = 0
    for i in range(len(a)):
        c = c + a[i] * b[i]

    return c

def inner_product_point(a, b):
    assert len(a) == len(b)
    c = 0
    for i in range(len(a)):
        c = c + int(a[i]) * b[i]

    return c

def vec_add(a, b):
    assert len(a) == len(b)
    return [x + y for x, y in zip(a, b)]

def vec_mul(a, b):
    assert len(a) == len(b)
    return [x * y for x, y in zip(a, b)]

def vec_scalar_mul_field(a, n):
    r = [None]*len(a)
    for i in range(len(a)):
        r[i] = a[i]*n
    return r

def vec_scalar_mul_point(a, n):
    r = [None]*len(a)
    for i in range(len(a)):
        r[i] = a[i]*int(n)
    return r