from zksk import Secret, DLRep
from zksk import utils


def ZK_equality(G, H):
    # Generate two El-Gamal ciphertexts (C1,C2) and (D1,D2)
    m = Secret(utils.get_random_num(bits=128))
    r1 = Secret(utils.get_random_num(bits=128))
    C1 = r1.value * G
    C2 = r1.value * H + m.value * G

    r2 = Secret(utils.get_random_num(bits=128))
    D1 = r2.value * G
    D2 = r2.value * H + m.value * G

    # Generate a NIZK proving equality of the plaintexts
    print("generation done")
    stmt = DLRep(C1, r1*G) & DLRep(C2, r1*H+m*G) & DLRep(D1, r2*G) & DLRep(D2, r2*H+m*G)
    print("statement done")
    zk_proof = stmt.prove()

    # r = Secret()
    # stmt2 = DLRep(C2 - r * H, D2 - r * H)
    # assert stmt2.verify(zk_proof)

    # Return two ciphertexts and the proof
    return (C1, C2), (D1, D2), zk_proof
