from honeybadgermpc.ntl.helpers import interpolate, batch_vandermonde_interpolate, \
    fft_batch_interpolate, gao_interpolate
from honeybadgermpc.field import GF
from honeybadgermpc.wb_interpolate import make_encoder_decoder

import random
import time

p = 0x73eda753299d7d483339d80809a1d80553bda402fffe5bfeffffffff00000001
field = GF.get(p)
t = 80
n = 3 * t + 1
k = t + 1
int_msg = [random.randint(0, p - 1) for _ in range(k)]
x = list(range(1, n + 1))


def timeit(f, n=1):
    start_time = time.time()
    for i in range(n):
        f()
    end_time = time.time()
    print("%.6f seconds\tTotal: %.6f seconds" % ((end_time - start_time) / n,
                                                 end_time - start_time))


def corrupt(message, num_errors, num_nones, min_val=0, max_val=131):
    """
    Inserts random corrupted values
    """
    message = list.copy(message)
    assert (len(message) >= num_errors +
            num_nones), "too much errors and none elements!"
    indices = random.sample(list(range(len(message))), num_errors + num_nones)
    for i in range(0, num_errors):
        message[indices[i]] = random.randint(min_val, max_val)
    for i in range(0, num_nones):
        message[indices[i + num_errors]] = None
    return message


if __name__ == "__main__":
    enc, dec, _ = make_encoder_decoder(n, k, p)
    emax = (n - 2 * t - 1) // 2

    encoded = [sum(int_msg[j] * pow(x[i], j, p) for j in range(k)) % p
               for i in range(n)]

    corrupted = corrupt(encoded, num_errors=emax, num_nones=0)

    print("n = ", n)
    print("All python")
    timeit(lambda: dec(corrupted, debug=False), 10)

    print("Gao")
    timeit(lambda: gao_interpolate(x, encoded, k, p), 10)

    decoded = dec(corrupted, debug=False)
    assert decoded == int_msg

    decoded, _ = gao_interpolate(x, encoded, k, p)
    assert decoded == int_msg
    print("All checks passed")