import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import signal
from timeit import default_timer
from functools import partial

from utility import codes, noise, misc
from utility.simulation.simulators import GenericMultithreadedSimulator
from utility.simulation import SimulationManager

from cpp_modules.cpp_decoders import ProximalDecoder_204_102 as ProximalDecoder


def task_func(params):
    """Function called by the GenericMultithreadedSimulator instance.

    Calculate the BER, FER, and DFR  for a given SNR and gamma.
    """
    signal.signal(signal.SIGINT, signal.SIG_IGN)

    decoder = params["decoder"]
    max_iterations = params["max_iterations"]
    SNR = params["SNR"]
    n = params["n"]
    k = params["k"]

    c = np.zeros(n)
    x_bpsk = c + 1

    total_bit_errors = 0
    total_frame_errors = 0
    dec_fails = 0

    num_iterations = 0

    for i in range(max_iterations):
        x = noise.add_awgn(x_bpsk, SNR, n, k)
        x_hat, k_max = decoder.decode(x)

        bit_errors = misc.count_bit_errors(x_hat, c)
        if bit_errors > 0:
            total_bit_errors += bit_errors
            total_frame_errors += 1

        num_iterations += 1

        if k_max == -1:
            dec_fails += 1

        if total_frame_errors > 100:
            break

    BER = total_bit_errors / (num_iterations * n)
    FER = total_frame_errors / num_iterations
    DFR = dec_fails / (num_iterations + dec_fails)

    return {"BER": BER, "FER": FER, "DFR": DFR,
            "num_iterations": num_iterations}


def get_params(code_name: str):
    """In this function all parameters for the simulation are defined."""
    # Define global simulation parameters

    H_file = f"res/{code_name}.alist"

    H = codes.read_alist_file(H_file)
    n_min_k, n = H.shape
    k = n - n_min_k

    omega = 0.05
    K = 100
    gammas = np.arange(0.0, 0.17, 0.01)

    SNRs = np.arange(1, 6, 0.5)
    max_iterations = 20000

    # Define parameters different for each task

    task_params = []
    for i, SNR in enumerate(SNRs):
        for j, gamma in enumerate(gammas):
            decoder = ProximalDecoder(H=H.astype('int32'), K=K, omega=omega,
                                      gamma=gamma)

            task_params.append(
                {"decoder": decoder, "max_iterations": max_iterations,
                 "SNR": SNR, "gamma": gamma, "n": n, "k": k})

    return omega, K, task_params


def configure_new_simulation(sim_mgr: SimulationManager, code_name: str,
                             sim_name: str) -> None:
    sim = GenericMultithreadedSimulator()

    omega, K, task_params = get_params(code_name)

    sim.task_params = task_params
    sim.task_func = task_func
    sim.format_func = partial(misc.pgf_reformat_data_3d, x_param_name="SNR",
                              y_param_name="gamma",
                              z_param_names=["BER", "FER", "DFR",
                                             "num_iterations"])

    sim_mgr.configure_simulation(simulator=sim, name=sim_name,
                                 additional_metadata={"omega": omega, "K": K})


def main():
    # code_name = "BCH_7_4"
    # code_name = "BCH_31_11"
    # code_name = "BCH_31_26"
    # code_name = "96.3.965"
    # code_name = "204.33.486"
    code_name = "204.33.484"
    # code_name = "204.55.187"
    # code_name = "408.33.844"

    sim_name = f"2d_BER_FER_DFR_{misc.slugify(code_name)}"

    # Run simulation

    sim_mgr = SimulationManager(saves_dir="sim_saves",
                                results_dir="sim_results")

    unfinished_sims = sim_mgr.get_unfinished()
    if len(unfinished_sims) > 0:
        sim_mgr.load_unfinished(unfinished_sims[0])
    else:
        configure_new_simulation(sim_mgr=sim_mgr, code_name=code_name,
                                 sim_name=sim_name)

    sim_mgr.simulate()

    # Plot results

    sns.set_theme()
    ax = sns.lineplot(data=sim_mgr.get_current_results(), x="SNR", y="BER",
                      hue="gamma")
    ax.set_yscale('log')
    ax.set_ylim((5e-5, 2e-0))
    plt.show()


if __name__ == "__main__":
    main()