from bandersnatch import Point, Scalar
import hashlib
from random import randint, shuffle, choice
from poly_utils import PrimeField
from time import time
from ipa_utils import IPAUtils, hash
import sys

#
# Proof of concept implementation for verkle tries
#
# All polynomials in this implementation are represented in evaluation form, i.e. by their values
# on the domain 0, 1, ..., WIDTH - 1
#
# Ethereum-specific implementation according to this EIP draft:
# https://notes.ethereum.org/@vbuterin/verkle_tree_eip
#

# Bandersnatch curve modulus
MODULUS = 13108968793781547619861935127046491459309155893440570251786403306729687672801

# Verkle trie parameters
KEY_LENGTH = 256 # bits
WIDTH_BITS = 8
WIDTH = 2**WIDTH_BITS

primefield = PrimeField(MODULUS, WIDTH)

# Number of key-value pairs to insert
NUMBER_STEMS = 2**15
CHUNKS_PER_STEM = 10
# Needs to be less than WIDTH * NUMBER_STEMS
NUMBER_CHUNKS = CHUNKS_PER_STEM * NUMBER_STEMS

# Number of extra stems to add to tree
NUMBER_ADDED_STEMS = 512

# Number of chunks to add to existing stems
NUMBER_ADDED_CHUNKS = 512

# Number of actually existing key/values pair in proof
NUMBER_EXISTING_KEYS_PROOF = 5000

# Added stems and chunks randomly (most likely empty)
NUMBER_RANDOM_STEMS_PROOF = 1000
NUMBER_RANDOM_CHUNKS_PROOF = 1000

NUMBER_VALUES_UPDATED = 1000

# Verkle trie constants
VERKLE_TRIE_NODE_TYPE_INNER = 0
VERKLE_TRIE_NODE_TYPE_SUFFIX_TREE = 1

# Verkle proof wire constants
VERKLE_PROOF_COMMITMENT_TYPE_INNER = 0
VERKLE_PROOF_COMMITMENT_TYPE_EXTENSION = 1
VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C1 = 2
VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C2 = 3

VERKLE_PROOF_EXTENSION_PRESENT_NOEXTENSION = 0
VERKLE_PROOF_EXTENSION_PRESENT_PRESENT = 1
VERKLE_PROOF_EXTENSION_PRESENT_OTHERSTEM = 2 # Used to indicate that there is an extension present, but for a different stem

def generate_basis(size):
    """
    Generates a basis for Pedersen commitments
    """
    # TODO: Currently random points that differ on every run.
    # Implement reproducable basis generation once hash_to_curve is provided
    BASIS_G = [Point(generator=False) for i in range(WIDTH)]
    BASIS_Q = Point(generator=False)
    return {"G": BASIS_G, "Q": BASIS_Q}


def get_stem(key):
    return key[:31]


def get_suffix(key):
    return key[31]


def commitment_to_field(commitment):
    return int.from_bytes(commitment.serialize(), "little") % MODULUS


# Node types: (Represented as python dicts)
#
# VERKLE_TRIE_NODE_TYPE_INNER:
#   0-255 refs to child node
#   "commitment": commitment
#   "commitment_field": commitment % MODULUS
#
# VERKLE_TRIE_NODE_TYPE_SUFFIX_TREE:
#   0-255 values (as bytes)
#   "stem": stem (31 bytes)
#   "C1": C1
#   "C1_field": C1 % MODULUS
#   "C2": C2
#   "C2_field": C2 % MODULUS
#   "commitment": commitment
#   "commitment_field": commitment % MODULUS


def update_verkle_tree_nocommitmentupdate(root_node, key, value):
    """
    Insert node without updating commitments (useful for building a full trie and adding commitments at the end)
    """

    current_node = root_node
    stem = get_stem(key)
    suffix = get_suffix(key)
    index = None
    depth = 0

    while current_node["node_type"] == VERKLE_TRIE_NODE_TYPE_INNER:
        previous_node = current_node
        index = stem[depth]
        depth += 1
        if index in current_node:
            current_node = current_node[index]
        else:
            current_node[index] = {"node_type": VERKLE_TRIE_NODE_TYPE_SUFFIX_TREE, "stem": stem, suffix: value}
            return

    if current_node["stem"] == stem:
        current_node[suffix] = value
    else:
        old_suffix_tree = current_node
        old_stem = old_suffix_tree["stem"]

        new_inner_node = {"node_type": VERKLE_TRIE_NODE_TYPE_INNER}
        previous_node[index] = new_inner_node
        previous_node = new_inner_node
        
        while old_stem[depth] == stem[depth]:
            index = stem[depth]
            new_inner_node = {"node_type": VERKLE_TRIE_NODE_TYPE_INNER}
            previous_node[index] = new_inner_node
            previous_node = new_inner_node
            depth += 1

        new_inner_node[stem[depth]] = {"node_type": VERKLE_TRIE_NODE_TYPE_SUFFIX_TREE, "stem": stem, suffix: value}
        new_inner_node[old_stem[depth]] = old_suffix_tree


def update_verkle_tree(root_node, key, value):
    """
    Update or insert node and update all commitments
    """
    current_node = root_node
    index = None
    path = []
    stem = get_stem(key)
    suffix = get_suffix(key)

    while True:
        index = stem[len(path)]
        path.append((index, current_node))
        if index in current_node:
            if current_node[index]["node_type"] == VERKLE_TRIE_NODE_TYPE_SUFFIX_TREE:
                old_node = current_node[index]
                if current_node[index]["stem"] == stem:
                    current_node = current_node[index]
                    old_value_lower = (int.from_bytes(current_node[suffix][:16], "little") + 2**128) if suffix in current_node else 0
                    old_value_upper = (int.from_bytes(current_node[suffix][16:], "little")) if suffix in current_node else 0
                    current_node[suffix] = value
                    new_value_lower = int.from_bytes(value[:16], "little") + 2**128
                    new_value_upper = int.from_bytes(value[16:], "little")
                    commitment_change = BASIS["G"][2 * suffix % 256].dup().mul((MODULUS + new_value_lower - old_value_lower) % MODULUS) \
                                        .add(BASIS["G"][(2 * suffix + 1) % 256].dup().mul((MODULUS + new_value_upper - old_value_upper) % MODULUS))
                    
                    if suffix < 128:
                        current_node["C1"].add(commitment_change)
                        new_field = commitment_to_field(current_node["C1"])
                        current_node["commitment"].add(BASIS["G"][2].dup().mul((MODULUS + new_field - current_node["C1_field"]) % MODULUS))
                        current_node["C1_field"] = new_field
                    else:
                        current_node["C2"].add(commitment_change)
                        new_field = commitment_to_field(current_node["C2"])
                        current_node["commitment"].add(BASIS["G"][3].dup().mul((MODULUS + new_field - current_node["C2_field"]) % MODULUS))
                        current_node["C2_field"] = new_field
                    new_field = commitment_to_field(current_node["commitment"])
                    value_change = (MODULUS + new_field - current_node["commitment_field"]) % MODULUS
                    current_node["commitment_field"] = new_field
                    break
                else:
                    new_inner_node = {"node_type": VERKLE_TRIE_NODE_TYPE_INNER}
                    new_index = stem[len(path)]
                    old_index = old_node["stem"][len(path)]
                    current_node[index] = new_inner_node

                    inserted_path = []
                    current_node = new_inner_node
                    while old_index == new_index:
                        index = new_index
                        next_inner_node = {"node_type": VERKLE_TRIE_NODE_TYPE_INNER}
                        current_node[index] = next_inner_node
                        inserted_path.append((index, current_node))
                        new_index = stem[len(path) + len(inserted_path)]
                        old_index = old_node["stem"][len(path) + len(inserted_path)]
                        current_node = next_inner_node

                    current_node[new_index] = {"node_type": VERKLE_TRIE_NODE_TYPE_SUFFIX_TREE, "stem": stem, suffix: value}
                    current_node[old_index] = old_node

                    verkle_add_missing_commitments(current_node)
                    for index, node in reversed(inserted_path):
                        verkle_add_missing_commitments(node)

                    value_change = (MODULUS + new_inner_node["commitment_field"] - old_node["commitment_field"]) % MODULUS
                    break

            current_node = current_node[index]
        else:
            current_node[index] = {"node_type": VERKLE_TRIE_NODE_TYPE_SUFFIX_TREE, "stem": stem, suffix: value}
            verkle_add_missing_commitments(current_node[index])
            value_change = current_node[index]["commitment_field"]
            break
    
    # Update all the ancestor commitments along `path`
    for index, node in reversed(path):
        node["commitment"].add(BASIS["G"][index].dup().mul(value_change))
        old_field = node["commitment_field"]
        new_field = commitment_to_field(node["commitment"])
        node["commitment_field"] = new_field
        value_change = (MODULUS + new_field - old_field) % MODULUS


def verkle_add_missing_commitments(node):
    """
    Recursively adds all missing commitments and hashes to a verkle trie structure.
    """
    if node["node_type"] == VERKLE_TRIE_NODE_TYPE_SUFFIX_TREE:

        C1 = ipa_utils.pedersen_commit_sparse({2 * i + j: int.from_bytes(node[i][16 * j:16 * (j + 1)], "little") + (1 - j) * 2**128
                                                for i in range(128)
                                                for j in range(2)
                                                if i in node})

        C2 = ipa_utils.pedersen_commit_sparse({2 * i + j: int.from_bytes(node[128 + i][16 * j:16 * (j + 1)], "little") + (1 - j) * 2**128
                                                for i in range(128)
                                                for j in range(2)
                                                if 128 + i in node})

        C1_field = commitment_to_field(C1)
        C2_field = commitment_to_field(C2)

        node["C1"] = C1
        node["C1_field"] = C1_field
        node["C2"] = C2
        node["C2_field"] = C2_field

        commitment = ipa_utils.pedersen_commit_sparse({0: 1, 
                                                       1: int.from_bytes(node["stem"], "little"),
                                                       2: C1_field, 
                                                       3: C2_field})

        node["commitment"] = commitment
        node["commitment_field"] = commitment_to_field(commitment)

    elif node["node_type"] == VERKLE_TRIE_NODE_TYPE_INNER:
        child = {}
        for i in range(WIDTH):
            if i in node:
                if "commitment_field" not in node[i]:
                    verkle_add_missing_commitments(node[i])
                child[i] = node[i]["commitment_field"]
        commitment = ipa_utils.pedersen_commit_sparse(child)
        node["commitment"] = commitment
        node["commitment_field"] = int.from_bytes(commitment.serialize(), "little") % MODULUS


def check_valid_tree(node, prefix=b""):
    """
    Checks that the subtree starting at `node` with prefix `prefix` is valid.
    Returns all the values found in the tree as a dict.
    """
    values = {}
    if node["node_type"] == VERKLE_TRIE_NODE_TYPE_INNER:

        child = {}
        for i in range(WIDTH):
            if i in node:
                child[i] = node[i]["commitment_field"]
        commitment = ipa_utils.pedersen_commit_sparse(child)
        assert node["commitment"] == commitment
        assert node["commitment_field"] == int.from_bytes(commitment.serialize(), "little") % MODULUS

        for i in range(WIDTH):
            if i in node:
                values.update(check_valid_tree(node[i], prefix + bytes([i])))
    else:
        assert node["node_type"] == VERKLE_TRIE_NODE_TYPE_SUFFIX_TREE
        assert node["stem"][:len(prefix)] == prefix
        C1 = ipa_utils.pedersen_commit_sparse({2 * i + j: int.from_bytes(node[i][16 * j:16 * (j + 1)], "little") + (1 - j) * 2**128
                                                for i in range(128)
                                                for j in range(2)
                                                if i in node})

        C2 = ipa_utils.pedersen_commit_sparse({2 * i + j: int.from_bytes(node[128 + i][16 * j:16 * (j + 1)], "little") + (1 - j) * 2**128
                                                for i in range(128)
                                                for j in range(2)
                                                if 128 + i in node})

        C1_field = commitment_to_field(C1)
        C2_field = commitment_to_field(C2)

        assert node["C1"] == C1
        assert node["C1_field"] == C1_field
        assert node["C2"] == C2
        assert node["C2_field"] == C2_field

        commitment = ipa_utils.pedersen_commit_sparse({0: 1, 
                                                       1: int.from_bytes(node["stem"], "little"),
                                                       2: C1_field, 
                                                       3: C2_field})

        assert node["commitment"] == commitment
        assert node["commitment_field"] == commitment_to_field(commitment)

        for i in range(WIDTH):
            if i in node:
                values[node["stem"] + bytes([i])] = node[i]
    
    return values


def get_total_depth(node):
    """
    Computes the total depth (sum of the depth of all nodes) of a verkle trie as well as the total number of nonzero values.
    Depth refers to the total number of commitments on the path to revealing the values, including the root and the
    extension and suffix tree node. (So a single value in an otherwise empty verkle tree has depth 3:
    root + extension + suffix_tree)
    """
    if node["node_type"] == VERKLE_TRIE_NODE_TYPE_INNER:
        total_depth = 0
        num_nodes = 0
        for i in range(WIDTH):
            if i in node:
                depth, nodes = get_total_depth(node[i])
                num_nodes += nodes
                total_depth += nodes + depth
        return total_depth, num_nodes
    else:
        num_chunks = len([i for i  in range(256) if i in node])
        return num_chunks * 2, num_chunks


def get_average_depth(root_node):
    """
    Get the average depth of nodes in a verkle trie
    """
    depth, nodes = get_total_depth(root_node)
    return depth / nodes


def find_key_with_path(root_node, key):
    """
    Returns the path of all nodes on the way to `key` as well as their index
    """
    current_node = root_node
    node_path = []
    stem = get_stem(key)
    while current_node["node_type"] == VERKLE_TRIE_NODE_TYPE_INNER:
        index = key[len(node_path)]
        node_path.append((stem[:len(node_path)], index, current_node))
        if index in current_node:
            current_node = current_node[index]
        else:
            return node_path, None

    suffix = get_suffix(key)
    node_path.append((stem[:len(node_path)], suffix, current_node))

    if current_node["stem"] == stem and suffix in current_node:
        return node_path, current_node[suffix]

    return node_path, None
    

def get_proof_size(proof):
    depths, extension_present, commitments_sorted_by_path_serialized, other_stems, D_serialized, ipa_proof = proof
    size = len(depths) # assume 8 bit integer to represent the depth (5 bit) and extension_present(2 bit)
    size += 32 * len(commitments_sorted_by_path_serialized)
    size += 32 * len(other_stems)
    size += 32 + (len(ipa_proof) - 1) * 2 * 32 + 32
    return size


lasttime = [0]


def start_logging_time_if_eligible(string, eligible):
    if eligible:
        print(string, file=sys.stderr)
        lasttime[0] = time()


def log_time_if_eligible(string, width, eligible):
    if eligible:
        print(string + ' ' * max(1, width - len(string)) + "{0:7.3f} s".format(time() - lasttime[0]), file=sys.stderr)
        lasttime[0] = time()


def make_ipa_multiproof(Cs, fs, zs, ys, display_times=True):
    """
    Computes an IPA multiproof according to the schema described here:
    https://dankradfeist.de/ethereum/2021/06/18/pcs-multiproofs.html

    This proof makes the assumption that the domain is the integers 0, 1, 2, ... WIDTH - 1
    """

    # Step 1: Construct g(X) polynomial in evaluation form
    r = ipa_utils.hash_to_field(Cs + zs + ys) % MODULUS

    log_time_if_eligible("   Hashed to r", 30, display_times)

    g = [0 for i in range(WIDTH)]
    power_of_r = 1
    for f, index in zip(fs, zs):
        quotient = primefield.compute_inner_quotient_in_evaluation_form(f, index)
        for i in range(WIDTH):
            g[i] += power_of_r * quotient[i]

        power_of_r = power_of_r * r % MODULUS
    
    for i in range(len(g)):
        g[i] %= MODULUS

    log_time_if_eligible("   Computed g polynomial", 30, display_times)

    D = ipa_utils.pedersen_commit(g)

    log_time_if_eligible("   Computed commitment D", 30, display_times)

    # Step 2: Compute h in evaluation form
    
    t = ipa_utils.hash_to_field([r, D]) % MODULUS
    
    h = [0 for i in range(WIDTH)]
    power_of_r = 1
    
    for f, index in zip(fs, zs):
        denominator_inv = primefield.inv(t - primefield.DOMAIN[index])
        for i in range(WIDTH):
            h[i] += power_of_r * f[i] * denominator_inv % MODULUS
            
        power_of_r = power_of_r * r % MODULUS
   
    for i in range(len(h)):
        h[i] %= MODULUS

    log_time_if_eligible("   Computed h polynomial", 30, display_times)

    h_minus_g = [(h[i] - g[i]) % primefield.MODULUS for i in range(WIDTH)]

    # Step 3: Evaluate and compute IPA proofs

    E = ipa_utils.pedersen_commit(h)

    y, ipa_proof = ipa_utils.evaluate_and_compute_ipa_proof(E.dup().add(D.dup().mul(MODULUS-1)), h_minus_g, t)

    log_time_if_eligible("   Computed IPA proof", 30, display_times)

    return D.serialize(), ipa_proof


def check_ipa_multiproof(Cs, zs, ys, proof, display_times=True):
    """
    Verifies an IPA multiproof according to the schema described here:
    https://dankradfeist.de/ethereum/2021/06/18/pcs-multiproofs.html
    """

    D_serialized, ipa_proof = proof

    D = Point().deserialize(D_serialized)

    # Step 1
    r = ipa_utils.hash_to_field(Cs + zs + ys)

    log_time_if_eligible("   Computed r hash", 30, display_times)
    
    # Step 2
    t = ipa_utils.hash_to_field([r, D])
    E_coefficients = {}
    g_2_of_t = 0
    power_of_r = 1

    C_by_serialized = {}

    for C, z, y in zip(Cs, zs, ys):
        E_coefficient = primefield.div(power_of_r, t - primefield.DOMAIN[z])
        C_serialized = C.serialize()
        C_by_serialized[C_serialized] = C
        E_coefficients[C_serialized] = E_coefficient if C_serialized not in E_coefficients \
                                        else (E_coefficients[C_serialized] + E_coefficient) % MODULUS
        g_2_of_t += E_coefficient * y % MODULUS

        power_of_r = power_of_r * r % MODULUS

    log_time_if_eligible("   Computed g2 and e coeffs", 30, display_times)
    
    E = Point().msm([C_by_serialized[x] for x in E_coefficients.keys()], E_coefficients.values())

    log_time_if_eligible("   Computed E commitment", 30, display_times)

    # Step 3 (Check IPA proofs)
    y = g_2_of_t % primefield.MODULUS

    if not ipa_utils.check_ipa_proof(E.dup().add(D.dup().mul(MODULUS - 1)), t, y, ipa_proof):
        return False

    log_time_if_eligible("   Checked IPA proof", 30, display_times)

    return True


def make_verkle_proof(root_node, keys, display_times=True):
    """
    Creates a proof for the `keys` in the verkle trie given by `root_node`

    This includes proving that a value is not in the verkle trie
    """

    start_logging_time_if_eligible("   Starting proof computation", display_times)

    #
    # Revealing a full verkle proof requires the following proofs for each `key`
    #
    # - all nodes on the path to the extension node, with z determined by stem             [VERKLE_PROOF_COMMITMENT_TYPE_INNER]
    # 
    # If there is an extension at the last place, then further it requires an extension proof:
    #
    # - the extension node, z = 0 (value = 1)                                              [VERKLE_PROOF_COMMITMENT_TYPE_EXTENSION]
    # - the extension node, z = 1 (stem)                                                   [VERKLE_PROOF_COMMITMENT_TYPE_EXTENSION]
    #
    # The stem of the extension can be a different stem, which has to be provided (either it is already in another key or it has to be added
    # to the proof)
    # If the stem is the stem of `key`, then the suffix tree for the suffix has to be revealed:
    #
    # - the extension node, z = 2/3 for C1/C2 (if suffix <128/>=128)                       [VERKLE_PROOF_COMMITMENT_TYPE_EXTENSION]
    # - the suffix tree node C1/C2, z = 2 * suffix % 128     (value_lower + 2**128)        [VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C{1/2}]
    # - the suffix tree node C1/C2, z = 2 * suffix + 1 % 128 (value_upper)                 [VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C{1_2}]
    #

    # Step 0: Find all keys in the trie

    # Nodes by path -- path refers to the full path of the commitment in the verkle trie
    nodes_by_path = {}

    # Nodes by path and z -- z is the position of the node that is opened; note that the same node can be opened at several z
    nodes_by_path_and_z = {}

    # All values in order of keys. `None` is used for never written values and b"\0" * 32 for zero/deleted values
    values = []

    # Depth at which the extension node for the stem was found
    depths_by_stem = {}

    # Whether or not a given stem had an extension node or not
    extension_present_by_stem = {}

    # All the stems that are already part of the proof
    key_stems = set(get_stem(key) for key in keys)

    # In some cases, a key that is not present will end up in an extension node for another stem. In this case,
    # we need to reveal that other stem as part of the proof so that the verifier can check that the extension
    # node is indeed for the other stem.
    other_stems = set()

    for key in keys:
        node_path, value = find_key_with_path(root_node, key)
        stem = get_stem(key)
        suffix = get_suffix(key)

        values.append(value)

        for path, z, node in node_path:
            if node["node_type"] == VERKLE_TRIE_NODE_TYPE_INNER:
                nodes_by_path[path] = (VERKLE_PROOF_COMMITMENT_TYPE_INNER, node)
                nodes_by_path_and_z[(path, z)] = (VERKLE_PROOF_COMMITMENT_TYPE_INNER, node)

        node = node_path[-1][2]
        path = node_path[-1][0]
        if node["node_type"] == VERKLE_TRIE_NODE_TYPE_SUFFIX_TREE:
            
            other_stem = node["stem"]
            depths_by_stem[stem] = len(node_path) - 1
            nodes_by_path[path] = (VERKLE_PROOF_COMMITMENT_TYPE_EXTENSION, node)
            nodes_by_path_and_z[(path, 0)] = (VERKLE_PROOF_COMMITMENT_TYPE_EXTENSION, node) # 1
            nodes_by_path_and_z[(path, 1)] = (VERKLE_PROOF_COMMITMENT_TYPE_EXTENSION, node) # stem

            if other_stem == stem:
                extension_present_by_stem[stem] = VERKLE_PROOF_EXTENSION_PRESENT_PRESENT

                nodes_by_path_and_z[(path, 2 + suffix // 128)] = (VERKLE_PROOF_COMMITMENT_TYPE_EXTENSION, node)  # C1/C2

                suffix_tree_path = path + bytes([2 if suffix < 128 else 3])
                suffix_tree_commitment = VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C1 if suffix < 128 \
                                        else VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C2
                nodes_by_path[suffix_tree_path] = (suffix_tree_commitment, node)
                nodes_by_path_and_z[(suffix_tree_path, suffix * 2 % 256)] = (suffix_tree_commitment, node)       # value_lower
                nodes_by_path_and_z[(suffix_tree_path, (suffix * 2 + 1) % 256)] = (suffix_tree_commitment, node) # value_upper
            else:
                extension_present_by_stem[stem] = VERKLE_PROOF_EXTENSION_PRESENT_OTHERSTEM
                other_stems.add(other_stem)
        else:
            depths_by_stem[stem] = len(node_path)
            extension_present_by_stem[stem] = VERKLE_PROOF_EXTENSION_PRESENT_NOEXTENSION

    depths = list(map(lambda x: x[1], sorted(depths_by_stem.items())))
    extension_present = list(map(lambda x: x[1], sorted(extension_present_by_stem.items())))
        
    log_time_if_eligible("   Computed key paths", 30, display_times)
    
    # Nodes sorted 
    nodes_sorted_by_path_and_z = sorted(nodes_by_path_and_z.items())
    
    log_time_if_eligible("   Sorted all commitments", 30, display_times)
    
    zs = []
    ys = []
    fs = []
    Cs = []

    for path_and_z, node_type_and_node in nodes_sorted_by_path_and_z:
        path, z = path_and_z
        node_type, node = node_type_and_node
        zs.append(z)
        if node_type == VERKLE_PROOF_COMMITMENT_TYPE_INNER:
            Cs.append(node["commitment"])
            ys.append(node[z]["commitment_field"] if z in node else 0)
            fs.append([node[i]["commitment_field"] if i in node else 0 for i in range(WIDTH)])
        elif node_type == VERKLE_PROOF_COMMITMENT_TYPE_EXTENSION:
            Cs.append(node["commitment"])
            if z == 0:
                ys.append(1)
            elif z == 1:
                ys.append(int.from_bytes(node["stem"], "little"))
            elif z == 2:
                ys.append(node["C1_field"])
            elif z == 3:
                ys.append(node["C2_field"])
            fs.append([1, 
                       int.from_bytes(node["stem"], "little"),
                       node["C1_field"],
                       node["C2_field"]]
                        + [0] * 252)
        elif node_type in [VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C1, VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C2]:
            if node_type == VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C1:
                Cs.append(node["C1"])
                suffix_offset = 0
            else:
                Cs.append(node["C2"])
                suffix_offset = 128
            suffix = suffix_offset + z // 2
            if suffix not in node:
                ys.append(0)
            else:
                if z % 2 == 0:
                    ys.append(int.from_bytes(node[suffix][:16], "little") + 2**128)
                else:
                    ys.append(int.from_bytes(node[suffix][16:], "little"))
            fs.append([(int.from_bytes(node[suffix_offset + i][16 * j:16 * (j + 1)], "little") + (1 - j) * 2**128) if suffix_offset + i in node else 0
                                                for i in range(128)
                                                for j in range(2)])

    D_serialized, ipa_proof = make_ipa_multiproof(Cs, fs, zs, ys, display_times)

    # All commitments, but without any duplications. These are for sending over the wire as part of the proof
    nodes_sorted_by_path = sorted(nodes_by_path.items())
    commitments_sorted_by_path_serialized = []
    for path, node_type_and_node in nodes_sorted_by_path[1:]:
        node_type, node = node_type_and_node
        if node_type in [VERKLE_PROOF_COMMITMENT_TYPE_INNER, VERKLE_PROOF_COMMITMENT_TYPE_EXTENSION]:
            commitments_sorted_by_path_serialized.append(node["commitment"].serialize())
        elif node_type == VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C1:
            commitments_sorted_by_path_serialized.append(node["C1"].serialize())
        elif node_type == VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C2:
            commitments_sorted_by_path_serialized.append(node["C2"].serialize())

    stems_with_extension = set(stem for stem in extension_present_by_stem if extension_present_by_stem[stem])
    other_stems = sorted(list(stem for stem in other_stems if stem not in stems_with_extension))

    log_time_if_eligible("   Serialized commitments", 30, display_times)

    return depths, extension_present, commitments_sorted_by_path_serialized, other_stems, D_serialized, ipa_proof


def check_verkle_proof(verkle_root, keys, values, proof, display_times=True):
    """
    Checks verkle trie proof according to
    https://dankradfeist.de/ethereum/2021/06/18/pcs-multiproofs.html
    """

    start_logging_time_if_eligible("   Starting proof check", display_times)

    # Unpack the proof
    depths, extension_present, commitments_sorted_by_path_serialized, other_stems, D_serialized, ipa_proof = proof
    commitments_sorted_by_path = [Point().deserialize(verkle_root)] + [Point().deserialize(x) for x in commitments_sorted_by_path_serialized]

    # Find all stems
    stems = sorted(list(set([get_stem(key) for key in keys])))
    assert other_stems == sorted(other_stems)
    depths_by_stem = {}
    extension_present_by_stem = {}
    stems_with_extension = set()
    other_stems_used = set()

    # The commitments in the proof are sorted by path in the verkle trie. So we first need to recreate a list of all paths
    # that are required for the proof, so that then we can associate all commitments with their path.
    commitment_types_by_path = {}
    all_paths_and_zs = set()

    leaf_values_by_path_and_z = {}
    other_stems_by_prefix = {}

    for stem, depth, extpres in zip(stems, depths, extension_present):
        depths_by_stem[stem] = depth
        extension_present_by_stem[stem] = extpres
        if extpres == VERKLE_PROOF_EXTENSION_PRESENT_PRESENT:
            stems_with_extension.add(stem)
    
    # Find all required indices
    for key, value in zip(keys, values):
        stem = get_stem(key)
        depth = depths_by_stem[stem]
        extpres = extension_present_by_stem[stem]
        for i in range(depth):
            commitment_types_by_path[stem[:i]] = VERKLE_PROOF_COMMITMENT_TYPE_INNER
            all_paths_and_zs.add((stem[:i], stem[i]))

        if extpres in [VERKLE_PROOF_EXTENSION_PRESENT_OTHERSTEM, VERKLE_PROOF_EXTENSION_PRESENT_PRESENT]:

            commitment_types_by_path[stem[:depth]] = VERKLE_PROOF_COMMITMENT_TYPE_EXTENSION
            all_paths_and_zs.add((stem[:depth], 0))
            all_paths_and_zs.add((stem[:depth], 1))

            leaf_values_by_path_and_z[(stem[:depth], 0)] = 1

            if extpres == VERKLE_PROOF_EXTENSION_PRESENT_PRESENT:
                suffix = get_suffix(key)

                all_paths_and_zs.add((stem[:depth], 2 + (suffix // 128)))

                leaf_values_by_path_and_z[(stem[:depth], 1)] = int.from_bytes(stem, "little")

                suffix_tree_path = stem[:depth] + bytes([2 if suffix < 128 else 3])
                commitment_types_by_path[suffix_tree_path] = VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C1 if suffix < 128 \
                                                                else VERKLE_PROOF_COMMITMENT_TYPE_SUFFIX_TREE_C2
                all_paths_and_zs.add((suffix_tree_path, 2 * suffix % 256))
                all_paths_and_zs.add((suffix_tree_path, (2 * suffix + 1) % 256))

                value_lower = int.from_bytes(value[:16], "little") + 2**128 if value != None else 0
                value_upper = int.from_bytes(value[16:], "little") if value != None else 0

                leaf_values_by_path_and_z[(suffix_tree_path, 2 * suffix % 256)] = value_lower
                leaf_values_by_path_and_z[(suffix_tree_path, (2 * suffix + 1) % 256)] = value_upper
            elif extpres == VERKLE_PROOF_EXTENSION_PRESENT_OTHERSTEM:
                # The proof indicates that an extension node for a different stem was found in the tree
                # We need to verify this is the case by looking up the other stem in "other_stems"
                # and verifying that the extension node at the site is indeed for the other stem

                # First check if the extension node is already included in the proof because the other stem
                # is already part of this proof. In this special case, we don't have to do anything because the 
                # extension proof for the other stem will already do all the work.
                other_stem = None

                # The stem not having an extension node means that the value has never been set:
                if value is not None:
                    return False

                # First check if the other stem is already in one of the revealed values
                # TODO: Convert this to binary search to prevent DOS vectors
                for o in stems_with_extension:
                    if o[:depth] == stem[:depth]:
                        assert other_stem is None
                        other_stem = o
                
                if other_stem is None:
                    # TODO: `other_stem` search is linear search which should work perfectly in average case but
                    # is a DOS vector. Need to employ binary search
                    for o in other_stems:
                        if o[:depth] == stem[:depth]:
                            assert other_stem is None
                            other_stem = o

