import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import typing
from pathlib import Path
import os
from itertools import chain
from timeit import default_timer

from decoders import proximal, maximum_likelihood
from utility import simulations, codes, visualization


# TODO: Fix spacing between axes and margins
def plot_results():
    results_dir = "sim_results"

    # Read data from files

    data = []
    for file in os.listdir(results_dir):
        if file.endswith(".csv"):
            df = pd.read_csv(os.path.join(results_dir, file))
            df = df.loc[:, ~df.columns.str.contains('^Unnamed')]
            data.append(df)

    # Create and show graphs

    sns.set_theme()
    fig = visualization.show_BER_curves(data)
    plt.show()


def main():
    Path("sim_results").mkdir(parents=True, exist_ok=True)

    # H = codes.read_alist_file("res/204.3.486.alist")
    H = codes.read_alist_file("res/204.55.187.alist")
    k = 102
    n = 204

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

    SNRs, BERs = simulations.test_decoders(n, k, decoders, target_frame_errors=100, SNRs=np.arange(1, 6, 0.5))

    df = pd.DataFrame({"SNR": SNRs})
    df["BER_prox_0_01"] = BERs[0]
    df["BER_prox_0_05"] = BERs[1]
    df["BER_prox_0_15"] = BERs[2]

    df.to_csv(f"sim_results/ldpc.csv")

    plot_results()


if __name__ == "__main__":
    main()