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


def common_alice_bob_eve(alice: Alice, eve, n=1000) -> float:
    return sum(correctness and alice_bit == eve_bit for alice_bit, eve_bit, correctness in
               zip(alice.key, eve.obtained_data, alice.correctness)) / 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',
             'eve:hijack_probability': r'p_{Eve}'}
mu, t, p_dc, p_opt, eta = sp.symbols(r'\mu t p_{dc} p_{opt} \eta')
p1 = 2 * p_dc - p_dc ** 2
p_ea, p_eb = sp.exp(-mu * (1 - eta) * p_opt), sp.exp(mu * (eta + p_opt - eta * p_opt - 1))
e_theory = 0.5 * (p_dc + t * (1 - p_ea) - p_dc * t * (1 - p_ea)) \
           * (1 - p_dc / 2 - 0.5 * t * (1 - p_eb) + 0.5 * p_dc * t * (1 - p_eb))
q_theory = 0.5 * (p1 + (1 - p1) * (1 - sp.exp(-mu * eta)) * t)
e_theory /= q_theory


def plot(parameter, values, add_eve: bool = False, show=False):
    data_q, data_e, eve_gains = [], [], []
    data_q_theoretical, data_e_theoretical = [], []
    parameters_sp = {PAR_NAMES[pname]: channel_parameters[pname] for pname in channel_parameters.keys()}
    n = 10000
    for val in values:
        print(val)
        channel = ChannelSym(**({parameter: val} if not parameter.startswith('eve:') else {}),
                             **{par: channel_parameters[par] for par in channel_parameters.keys() if par != parameter})
        alice, bob = Alice(channel), Bob(channel)
        if add_eve:
            eve = EveBS(**({parameter.removeprefix('eve:'): val} if parameter.startswith('eve:') else {}))
            channel.eve = eve
        run_qkd(alice, bob, n)
        e, r = get_e_and_r(alice, bob, n)
        data_q.append(r)
        data_e.append(e)
        if add_eve:
            eve_gains.append(common_alice_bob_eve(alice, eve, n))
        parameters_sp[PAR_NAMES.get(parameter, parameter)] = val
        data_q_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_q, label='$Q$')
    plt.plot(values, data_q_theoretical, label='$Q_{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_q, data_e, data_q_theoretical, data_e_theoretical]))
    plt.xlabel('$' + PAR_NAMES.get(parameter, parameter) + '$')
    plt.title('$Q$ vs. $' + PAR_NAMES.get(parameter, parameter) + '$' + ' with Eve' * add_eve)
    plt.savefig(f'output/dep_{parameter}_q{"_with_eve" * add_eve}.png')
    if show:
        plt.show()
    else:
        plt.cla()
    plt.plot(values, data_e, label='$E$')
    plt.plot(values, data_e_theoretical, label='$E_{theory}$')
    plt.legend()
    plt.xlabel('$' + PAR_NAMES.get(parameter, parameter) + '$')
    plt.title('$E$ vs. $' + PAR_NAMES.get(parameter, parameter) + '$' + ' with Eve' * add_eve)
    plt.savefig(f'output/dep_{parameter}_e{"_with_eve" * add_eve}.png')
    if show:
        plt.show()
    else:
        plt.cla()
    if add_eve:
        plt.plot(values, [g / q for g, q in zip(eve_gains, data_q)], label='$Q_{Eve}$')
        plt.title('Eve\'s gains')
        plt.legend()
        plt.ylabel(r'$\dfrac{Q_{Eve}}{Q}$')
        plt.xlabel('$' + PAR_NAMES.get(parameter, parameter) + '$')
        plt.savefig(f'output/dep_{parameter}_eve_q.png')
        if show:
            plt.show()
        else:
            plt.cla()


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.15, 0.01))
    plot('p_opt', np.arange(0, 0.15, 0.01))
    plot('detector_sensitivity', np.arange(0.1, 1, 0.05))
    plot('mu', np.arange(0.5, 2.5, 0.1))
    plot('transmittance', np.arange(0.4, 1, 0.05))
    # plot('eve:hijack_probability', np.arange(0, 1, 0.05))