import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import os
from utility import visualization


def plot_results():
    graph_names = {"96.3.965": "n=96, k=48 - 965",
                   "204.3.486": "n=204, k=102 - 486",
                   "204.55.187": "n=204, k=102 - 187",
                   "408.33.844": "n=408, k=204 - 844",
                   "816.1A4.845": "n=816, k=272 - 843",
                   "999.111.3.5543": "n=999, k=888 - 5543",
                   "999.111.3.5565": "n=999, k=888 - 5565",
                   "PEGReg252x504": "n=504, k=252 - PEGReg"}

    line_labels = {"BER_ML": "ML",
                   "BER_prox_0_15": "$\gamma = 0.15$",
                   "BER_prox_0_05": "$\gamma = 0.05$",
                   "BER_prox_0_01": "$\gamma = 0.01$"}

    # Read data from files
    results_dir = "sim_results"

    data = {}
    for file in os.listdir(results_dir):
        if file.endswith(".csv"):
            filename = os.path.splitext(file)[0]

            df = pd.read_csv(os.path.join(results_dir, file))
            df = df.loc[:, ~df.columns.str.contains('^Unnamed')]
            data[graph_names[filename]] = df

    # Create and show graphs

    sns.set_theme()
    fig = visualization.show_BER_curves("Bit-Error-Rates of proximal decoder for different codes",
                                        data, num_cols=4, line_labels=line_labels)
    plt.show()


if __name__ == "__main__":
    plot_results()