pragma circom 2.1.9;

include "./ecdsa_verify.circom";
include "./circom-ecdsa/zk-identity/eth.circom";
include "../../circomlib/circuits/mimcsponge.circom";

template ECDSAVerifyPubKeyToAddr(n, k) {
    signal input electionId;
    signal input vote;
    signal input s[k];
    signal input TPreComputes[32][256][2][4]; // T = r^-1 * R
    signal input U[2][k]; // -(m * r^-1 * G)
    signal input THashes[32];
    signal input UHashes[2];

    signal output addr;
    signal output hash;
    signal output root;

    var stride = 8;
    var num_strides = div_ceil(n * k, stride);

    component ecdsaVerify = ECDSAVerify(n, k);
    for (var i = 0; i < num_strides; i++) {
        for (var j = 0; j < 2 ** stride; j++) {
            ecdsaVerify.TPreComputes[i][j][0][0] <== TPreComputes[i][j][0][0];
            ecdsaVerify.TPreComputes[i][j][0][1] <== TPreComputes[i][j][0][1];
            ecdsaVerify.TPreComputes[i][j][0][2] <== TPreComputes[i][j][0][2];
            ecdsaVerify.TPreComputes[i][j][0][3] <== TPreComputes[i][j][0][3];

            ecdsaVerify.TPreComputes[i][j][1][0] <== TPreComputes[i][j][1][0];
            ecdsaVerify.TPreComputes[i][j][1][1] <== TPreComputes[i][j][1][1];
            ecdsaVerify.TPreComputes[i][j][1][2] <== TPreComputes[i][j][1][2];
            ecdsaVerify.TPreComputes[i][j][1][3] <== TPreComputes[i][j][1][3];
        }
    }

    for (var i = 0; i < k; i++) {
        ecdsaVerify.s[i] <== s[i];
    }

    for (var i = 0; i < k; i++) { 
        ecdsaVerify.U[0][i] <== U[0][i];
        ecdsaVerify.U[1][i] <== U[1][i];
    }

    component flattenPub = FlattenPubkey(n, k);
    for (var i = 0; i < k; i++) {
        flattenPub.chunkedPubkey[0][i] <== ecdsaVerify.pubKey[0][i];
        flattenPub.chunkedPubkey[1][i] <== ecdsaVerify.pubKey[1][i];
    }

    component pubToAddr = PubkeyToAddress();
    for (var i = 0; i < 512; i++) {
        pubToAddr.pubkeyBits[i] <== flattenPub.pubkeyBits[i];
    }

    component merkleTreeTHashes = CheckRoot(5);
    for (var i = 0; i < 32; i++) {
        merkleTreeTHashes.leaves[i] <== THashes[i];
    }

    component rootHash = HashLeftRight();
    rootHash.left <== UHashes[0];
    rootHash.right <== UHashes[1];
    rootHash.key <== merkleTreeTHashes.root;

    component msgHash = HashLeftRight();
    msgHash.left <== electionId;
    msgHash.right <== vote;
    msgHash.key <== 0;

    addr <== pubToAddr.address;
    hash <== msgHash.hash;
    root <== rootHash.hash;
}

template HashLeftRight() {
    signal input left;
    signal input right;
    signal input key;
    signal output hash;

    component hasher = MiMCSponge(2, 220, 1);
    hasher.ins[0] <== left;
    hasher.ins[1] <== right;
    hasher.k <== key;
    hash <== hasher.outs[0];
}

template CheckRoot(levels) {
    var totalLeaves = 2 ** levels;
    var numLeafHashers = totalLeaves / 2;
    var numIntermediateHashers = numLeafHashers - 1;

    signal input leaves[totalLeaves];
    signal output root;

    var numHashers = totalLeaves - 1;
    component hashers[numHashers];

    for (var i = 0; i < numHashers; i++) {
        hashers[i] = HashLeftRight();
    }

    for (var i = 0; i < numLeafHashers; i++){
        hashers[i].left <== leaves[2 * i];
        hashers[i].right <== leaves[2 * i + 1];
        hashers[i].key <== 0;
    }

    var k = 0;
    for (var i = numLeafHashers; i < numLeafHashers + numIntermediateHashers; i++) {
        hashers[i].left <== hashers[2 * k].hash;
        hashers[i].right <== hashers[2 * k + 1].hash;
        hashers[i].key <== 0;
        k++;
    }

    root <== hashers[numHashers - 1].hash;
}

component main = ECDSAVerifyPubKeyToAddr(64, 4);