import numpy as np

from decoders import proximal, maximum_likelihood
from utility import simulation, codes


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

    # 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")
    k, n = H.shape

    decoders = [
        proximal.ProximalDecoder(H, gamma=0.01),
        proximal.ProximalDecoder(H, gamma=0.05),
        proximal.ProximalDecoder(H, gamma=0.15)
    ]

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

    sim = simulation.Simulator(n=n, k=k, decoders=decoders,
                               target_frame_errors=100,
                               SNRs=np.arange(1, 6, 0.5))

    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 = simulation.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()