import numpy as np

from decoders import proximal, maximum_likelihood
from cpp_modules import cpp_decoders
from utility import codes
from utility.simulation import SimulationManager
from utility.simulation.simulators import ProximalDecoderSimulator


def start_new_simulation(sim_mgr: SimulationManager):
    sim_name = "test"

    # H = codes.Gs["Hamming_7_4"]
    # H = codes.read_alist_file("res/204.3.486.alist")
    # H = codes.read_alist_file("res/204.55.187.alist")
    H = codes.read_alist_file("res/96.3.965.alist")
    # H = codes.read_alist_file("res/408.33.844.alist")
    # H = codes.read_alist_file("res/PEGReg252x504.alist")
    # H = codes.read_alist_file("res/999.111.3.5543.alist")
    # H = codes.read_alist_file("res/999.111.3.5565.alist")
    # H = codes.read_alist_file("res/816.1A4.845.alist")
    n_min_k, n = H.shape
    k = n - n_min_k

    decoders = [
        cpp_decoders.ProximalDecoder(H.astype('int32'), gamma=0.01, K=1000,
                                     omega=1e-4, eta=1.5),
        cpp_decoders.ProximalDecoder(H.astype('int32'), gamma=0.05, K=1000,
                                     omega=5 * 1e-5, eta=1.5),
        cpp_decoders.ProximalDecoder(H.astype('int32'), gamma=0.15, K=1000,
                                     omega=1e-4, eta=1.5)
    ]

    labels = [
        "proximal $\\gamma = 0.01$",
        "proximal $\\gamma = 0.05$",
        "proximal $\\gamma = 0.15$"
    ]

    sim = ProximalDecoderSimulator(n=n, k=k, decoders=decoders,
                                              target_frame_errors=100,
                                              SNRs=np.arange(1, 6, 0.5),
                                              max_num_iterations=3000)

    sim_mgr.configure_simulation(simulator=sim, name=sim_name,
                                 column_labels=labels)
    sim_mgr.simulate()


def main():
    # Perform necessary initialization

    results_dir = "sim_results"
    saves_dir = "sim_saves"

    sim_mgr = SimulationManager(results_dir=results_dir,
                                saves_dir=saves_dir)

    # Calculate BERs

    unfinished_sims = sim_mgr.get_unfinished()

    if len(unfinished_sims) > 0:
        print("Found unfinished simulation. Picking up where it was left of")

        sim_mgr.load_unfinished(unfinished_sims[0])
        sim_mgr.simulate()
    else:
        start_new_simulation(sim_mgr)


if __name__ == "__main__":
    main()