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

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 average error over a number of iterations.
    """
    signal.signal(signal.SIGINT, signal.SIG_IGN)

    decoder = params["decoder"]
    num_iterations = params["num_iterations"]
    x_bpsk = params["x_bpsk"]
    SNR = params["SNR"]
    n = params["n"]
    k = params["k"]
    K = params["K"]

    avg_error_values = np.zeros(K)

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

        error_values = decoder.get_error_values(x_bpsk.astype('int32'), x)

        for j, val in enumerate(error_values):
            avg_error_values[j] += val

    avg_error_values = avg_error_values / num_iterations

    return {"err": avg_error_values}


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

    SNR = 8
    omegas = np.logspace(-0, -10, 40)
    K = 200

    num_iterations = 1000
    x_bpsk = np.zeros(n) + 1

    # Define parameters different for each task

    task_params = []
    for i, omega in enumerate(omegas):
        decoder = ProximalDecoder(H=H.astype('int32'), K=K,
                                  omega=omega)
        task_params.append(
            {"decoder": decoder, "num_iterations": num_iterations,
             "x_bpsk": x_bpsk, "SNR": SNR, "n": n, "k": k, "K": K,
             "omega": omega})

    return SNR, K, task_params


def reformat_data(results):
    """Reformat the data obtained from the GenericMultithreadedSimulator to
    be usable by pgfplots.
    """
    K = 200
    num_points = len(results) * K

    x = np.zeros(num_points)
    y = np.zeros(num_points)
    z = np.zeros(num_points)

    for i, (params, result) in enumerate(results.items()):
        np.put(x, np.arange(i * K, (i + 1) * K), np.arange(1, K+1))
        np.put(y, np.arange(i * K, (i + 1) * K), params["omega"])
        np.put(z, np.arange(i * K, (i + 1) * K), result["err"])

    x = x[::4]
    y = y[::4]
    z = z[::4]

    df = pd.DataFrame({"k": x, "omega": y, "err": z}).sort_values(
        by=['k', 'omega'], ascending=[True, False])

    return df


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

    SNR, K, task_params = get_params(code_name)

    sim.task_params = task_params
    sim.task_func = task_func
    sim.format_func = reformat_data

    sim_mgr.configure_simulation(simulator=sim, name=sim_name,
                                 additional_metadata={"SNR": SNR, "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_avg_error_{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()


if __name__ == "__main__":
    main()