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 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_31_5 as ProximalDecoder


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

    decoder, num_iterations, x_bpsk, SNR, n, k = params

    dec_fails = 0

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

        if num_iter == -1:
            dec_fails += 1

    return dec_fails / num_iterations


def simulate(H_file, SNR, num_iterations, omegas, Ks):
    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

    x_bpsk = np.zeros(n) + 1

    # Define params different for each task

    params = {}
    for i, omega in enumerate(omegas):
        for j, K in enumerate(Ks):
            decoder = ProximalDecoder(H=H.astype('int32'), K=K.astype('int32'),
                                      omega=omega)
            params[f"{i}_{j}"] = (decoder, num_iterations, x_bpsk, 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, omegas, Ks):
    data = np.zeros(1600).reshape(40, 40)

    for key, value in results.items():
        i_w, i_k = key.split('_')
        data[int(i_w), int(i_k)] = value

    return pd.DataFrame(data, columns=Ks, index=omegas)


def main():
    # Set up simulation params

    sim_name = "w_log_k_lin_zoomed_in"

    # H_file = "96.3.965.alist"
    # H_file = "204.33.486.alist"
    # H_file = "408.33.844.alist"
    H_file = "BCH_31_26.alist"
    SNR = 3

    num_iterations = 1000
    omegas = np.logspace(-0.3, -2.82, 40)
    Ks = np.ceil(np.linspace(10 ** 1.3, 10 ** 2.3, 40)).astype('int32')

    # Run simulation

    start_time = default_timer()
    results = simulate(H_file, SNR, num_iterations, omegas, Ks)
    end_time = default_timer()

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

    df = reformat_data(results, omegas, Ks)

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

    sns.set_theme()
    sns.heatmap(df)
    plt.show()


if __name__ == "__main__":
    main()