import json
import os.path

from alice import Alice
from bob import Bob
from channel import ChannelSym
from threading import Thread
import typing
import sympy as sp

from eve import EveBS


def run_qkd(alice: Alice, bob: Bob, n: int = 1000):
    alice_thread = Thread(target=lambda: alice.generate_key(n))
    bob_thread = Thread(target=lambda: bob.generate_key(n))
    alice_thread.start()
    bob_thread.start()
    while alice_thread.is_alive() or bob_thread.is_alive():
        pass


def get_e_and_r(alice, bob, n=1000) -> typing.Tuple[float, float]:
    return 1 - sum(k1 == k2 for k1, k2 in zip(alice.key, bob.key)) / len(alice.key), len(alice.key) / n


channel_parameters = {'p_opt': 0.05, 'p_dc': 0.05, 'mu': 1, 'detector_sensitivity': 0.8, 'transmittance': 0.8}
PAR_NAMES = {'p_opt': 'p_{opt}', 'p_dc': 'p_{dc}', 'mu': r'\mu', 'detector_sensitivity': r'\eta', 'transmittance': 't'}
mu, t, p_dc, p_opt, eta = sp.symbols(r'\mu t p_{dc} p_{opt} \eta')
e_theory = 0.5 * (p_dc + p_opt * mu * t * eta) / (2 * p_dc + mu * t * eta)
q_theory = (mu * t * eta + 2 * p_dc) / 2 * (1 - e_theory * 2)


def plot(parameter, values, add_eve: bool = False):
    data_r, data_e = [], []
    data_r_theoretical, data_e_theoretical = [], []
    parameters_sp = {PAR_NAMES[pname]: channel_parameters[pname] for pname in PAR_NAMES.keys()}
    n = 10000
    for val in values:
        print(val)
        channel = ChannelSym(**{parameter: val})
        alice, bob = Alice(channel), Bob(channel)
        if add_eve:
            eve = EveBS()
            channel.eve = eve
        run_qkd(alice, bob, n)
        e, r = get_e_and_r(alice, bob, n)
        data_r.append(r)
        data_e.append(e)
        parameters_sp[PAR_NAMES.get(parameter, parameter)] = val
        data_r_theoretical.append(float(q_theory.evalf(subs=parameters_sp)))
        data_e_theoretical.append(float(e_theory.evalf(subs=parameters_sp)))
    from matplotlib import pyplot as plt
    plt.plot(values, data_r, label='$R$')
    plt.plot(values, data_e, label='$E$')
    plt.plot(values, data_r_theoretical, label='$R_{theory}$')
    plt.plot(values, data_e_theoretical, label='$E_{theory}$')
    plt.legend()
    if not os.path.exists('output'):
        os.mkdir('output')
    with open(f'output/dep_{parameter}.json', 'w') as f:
        f.write(json.dumps([list(values), data_r, data_e, data_r_theoretical, data_e_theoretical]))
    plt.xlabel('$' + PAR_NAMES.get(parameter, parameter) + '$')
    plt.title('$R$ and $E$ vs. $' + PAR_NAMES.get(parameter, parameter) + '$' + ' with Eve' * add_eve)
    plt.savefig(f'output/dep_{parameter}.png')
    plt.show()


def run_one():
    N = 1000
    channel = ChannelSym()
    alice, bob = Alice(channel), Bob(channel)
    run_qkd(alice, bob)
    # print(list(map(int, alice.key)))
    # print(list(map(int, bob.key)))
    print('Alice bits:   ', *list(map(int, alice.sent)))
    print('Bob bits:     ',
          *list(map(lambda t: {(0, 1): 1, (1, 0): 0, (0, 0): 2, (1, 1): 3}[(int(t[0]), int(t[1]))], bob.my_results)))
    print('Alice basises:', *list(map(int, alice.my_basises)))
    print('Bob basises:  ', *list(map(int, bob.my_basises)))
    bob_correctness = [left + right == 1 for left, right in bob.my_results]
    print('              ', *[
        int(k) if c and b1 == b2 else ' '
        for c, k, b1, b2 in zip(bob_correctness, alice.sent, bob.my_basises, alice.my_basises)])
    print('              ', *[
        {(0, 1): 1, (1, 0): 0, (0, 0): 2, (1, 1): 3}[(int(k[0]), int(k[1]))] if c and b1 == b2 else ' '
        for c, k, b1, b2 in zip(bob_correctness, bob.my_results, bob.my_basises, alice.my_basises)])
    print(
        f'{100 * sum(k1 == k2 for k1, k2 in zip(alice.key, bob.key)) / len(alice.key):.2f}%, '
        f'key length: {len(alice.key)}')
    e, r = get_e_and_r(alice, bob, N)
    print(f'E: {e * 100:.1f}%, R: {r}')


if __name__ == '__main__':
    import numpy as np
    plot('p_dc', np.arange(0, 0.1, 0.01))
    plot('detector_sensitivity', np.arange(0.5, 1, 0.05))
    plot('mu', np.arange(0.5, 1.5, 0.1))