import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import signal
from timeit import default_timer
from tqdm import tqdm

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

# from cpp_modules.cpp_decoders import ProximalDecoder
from cpp_modules.cpp_decoders import ProximalDecoder_204_102 as ProximalDecoder


def count_bit_errors(d: np.array, d_hat: np.array) -> int:
    return np.sum(d != d_hat)


def task_func(params):
    signal.signal(signal.SIGINT, signal.SIG_IGN)

    decoder, max_iterations, SNR, n, k = params
    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 = 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 > 500:
            break

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

    return BER, FER, DFR, num_iterations


def simulate(H_file, SNRs, max_iterations, omega, K, gammas):
    sim = GenericMultithreadedSimulator()

    # Define fixed simulation params

    H = codes.read_alist_file(f"res/{H_file}")
    n_min_k, n = H.shape
    k = n - n_min_k

    # Define params different for each 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)
            params[f"{i}_{j}"] = (decoder, max_iterations, SNR, n, k)

    # Set up simulation

    sim.task_params = params
    sim.task_func = task_func

    sim.start_or_continue()

    return sim.get_current_results()


def reformat_data(results, SNRs, gammas):
    data = {"BER": np.zeros(3 * 10), "FER": np.zeros(3 * 10),
            "DFR": np.zeros(3 * 10), "gamma": np.zeros(3 * 10),
            "SNR": np.zeros(3 * 10), "num_iter": np.zeros(3 * 10)}

    for i, (key, (BER, FER, DFR, num_iter)) in enumerate(results.items()):
        i_SNR, i_gamma = key.split('_')
        data["BER"][i] = BER
        data["FER"][i] = FER
        data["DFR"][i] = DFR
        data["num_iter"][i] = num_iter
        data["SNR"][i] = SNRs[int(i_SNR)]
        data["gamma"][i] = gammas[int(i_gamma)]

    print(pd.DataFrame(data))
    return pd.DataFrame(data)


def main():
    # Set up simulation params

    sim_name = "BER_FER_DFR"

    # H_file = "96.3.965.alist"
    H_file = "204.33.486.alist"
    # H_file = "204.33.484.alist"
    # H_file = "204.55.187.alist"
    # H_file = "408.33.844.alist"
    # H_file = "BCH_7_4.alist"
    # H_file = "BCH_31_11.alist"
    # H_file = "BCH_31_26.alist"
    SNRs = np.arange(1, 6, 0.5)

    max_iterations = 20000
    # omega = 0.005
    # K = 60
    omega = 0.05
    K = 60
    gammas = [0.15, 0.01, 0.05]

    # Run simulation

    start_time = default_timer()
    results = simulate(H_file, SNRs, max_iterations, omega, K, gammas)
    end_time = default_timer()

    print(f"duration: {end_time - start_time}")

    df = reformat_data(results, SNRs, gammas)

    df.to_csv(
        f"sim_results/{sim_name}_{misc.slugify(H_file)}.csv")

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


if __name__ == "__main__":
    main()