import hashlib
import secrets
from struct import pack, pack_into, unpack_from

def sha(x):
    m = hashlib.sha256()
    m.update(x)
    result = m.digest()
    return result[0] & 0b1

def bit_at_index(buffer, index):
    offset = (index >> 3) % len(buffer)
    return buffer[offset] & (1 << (index & 0b111)) != 0

def evaluate(f, x):
    stack = []
    offset = 0
    value = 0
    while offset < len(f):
        opcode = f[offset]
        offset += 1
        if opcode == 0 or opcode == 1:
            stack.append((opcode, value))
            value = 0
        elif opcode == 2:
            if len(stack) == 0:
                return (value, offset)
            (last_opcode, _) = stack[-1]
            if last_opcode > 0:
                stack.append((0, value))
                value = 0
                continue
            right = value
            (_, left) = stack.pop()
            (opcode, value) = stack.pop()
            value ^= ((left & right) ^ (opcode & 0b1))
        else:
            try:
                index = unpack_from('I', f, offset)[0]
                offset += 4
                if bit_at_index(x, index):
                    value ^= 1
            except:
                break

    while len(stack) > 0:
        (opcode, other_value) = stack.pop()
        if opcode == 0:
            right = other_value
            (opcode, left) = stack.pop()
            value ^= ((left & right) ^ (opcode & 0b1))
        value ^= other_value ^ (opcode & 0b1)
    return (value, offset)

def random_generator():
    return secrets.token_bytes(256)

def random_input():
    return secrets.token_bytes(4)

def generate(generator, sample):
    f_size = 1024
    f = bytearray(f_size)
    x = bytearray(4) + sample
    for i in range(0, f_size):
        build_value = 0
        for j in range(0, 8):
            step = i * 8 + j
            pack_into('H', x, 0, step)
            (value, _) = evaluate(generator, x)
            build_value <<= 1
            build_value |= value
        f[i] = build_value
    return f

def sample(N):
    inputs = [random_input() for i in range(0, N)]
    outputs = [sha(x) for x in inputs]
    return (inputs, outputs)

def augment_inputs(inputs, layers):
    augmented_inputs = []
    for x in inputs:
        x_n = bytearray(1) + x
        for layer in layers:
            build_value = 0
            for candidate in layer:
                (value, _) = evaluate(candidate, x_n)
                build_value <<= 1
                build_value |= value
            x_n[0] = build_value
        augmented_inputs.append(x_n)
    return augmented_inputs

def pack_sample(inputs, outputs):
    sample = bytearray()
    for i in range(0, len(inputs)):
        sample += inputs[i]
        sample += bytearray([outputs[i]])
    return sample

def compute_score(f, inputs, outputs):
    correct = 0.0
    for i in range(0, len(inputs)):
        (value, _) = evaluate(f, inputs[i])
        if value == outputs[i]:
            correct += 1
    return correct / len(outputs)

def evaluate_generator(g):
    num_candidates = 8
    num_train_samples = 64
    num_test_samples = 1000
    num_epochs = 10
    threshold = 0

    layers = []
    for epoch in range(0, num_epochs):
        difficulty = 0
        layer = []
        candidate = 0
        scores = []
        while candidate < num_candidates:
            (x, y) = sample(num_train_samples)
            x_n = augment_inputs(x, layers)
            f = generate(g, pack_sample(x_n, y))
            print(f)

            (x, y) = sample(num_test_samples)
            x_n = augment_inputs(x, layers)
            score = compute_score(f, x_n, y)

            if score < threshold - difficulty * 0.0001:
                difficulty += 1
                continue

            print(epoch, score, difficulty)
            
            layer.append(f)
            scores.append(score)
            difficulty = 0
            candidate += 1
        threshold = sum(scores) / len(scores)
        layers.append(layer)
    return threshold

def main():
    num_random_candidates = 1000

    g = None
    score = 0

    for i in range(0, num_random_candidates):
        g_n = random_generator()
        print(g_n)
        score_n = evaluate_generator(g_n)
        print(i, score_n)
        if score > score_n:
            score = score_n
            g = g_n


if __name__ == "__main__":
    main()