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


def test_decoders(G, decoders: typing.List) -> pd.DataFrame:
    k, n = G.shape
    x = np.zeros(n)  # All-zeros assumption

    SNRs = np.linspace(1, 7, 7)
    data = pd.DataFrame({"SNR": SNRs})

    start_time = default_timer()

    for decoder_name in decoders:
        decoder = decoders[decoder_name]
        _, BERs_sd = simulations.test_decoder(x,
                                              decoder=decoder,
                                              SNRs=SNRs,
                                              target_frame_errors=100,
                                              N_max=50000)

        data[f"BER_{decoder_name}"] = BERs_sd

    stop_time = default_timer()

    print(f"Elapsed time: {stop_time-start_time:.2f}s")

    return data


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

    code_paths = {}
    for file in os.listdir(results_dir):
        if file.endswith(".csv"):
            code_paths[file.replace(".csv", "")] = os.path.join(results_dir, file)

    sns.set_theme()

    fig, axes = plt.subplots(2, len(code_paths) // 2, figsize=(12, 6))
    fig.suptitle("Bit-Error-Rates of various decoders for different codes")

    axes = list(chain.from_iterable(axes))

    for i, code in enumerate(code_paths):
        data = pd.read_csv(code_paths[code])

        column_names = [column for column in data.columns.values.tolist()
                        if column.startswith("BER")]

        ax = axes[i]

        for column in column_names:
            sns.lineplot(ax=ax, data=data, x="SNR", y=column, label=column.lstrip("BER_"))

        ax.set_title(code)
        ax.set(yscale="log")
        ax.set_xlabel("SNR")
        ax.set_ylabel("BER")
        ax.set_yticks([10e-5, 10e-4, 10e-3, 10e-2, 10e-1, 10e0])
        ax.legend()

    plt.show()


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

    used_codes = [
        "Hamming_7_4",
        "Golay_24_12",
        # "BCH_31_16",
        # "BCH_31_21",
        # "BCH_63_16",
    ]

    for used_code in used_codes:
        G = codes.Gs[used_code]
        H = codes.get_systematic_H(G)

        decoders = {
            "ML": maximum_likelihood.MLDecoder(G, H),
            "proximal_0_01": proximal.ProximalDecoder(H, gamma=0.01),
            "proximal_0_05": proximal.ProximalDecoder(H, gamma=0.05),
            "proximal_0_15": proximal.ProximalDecoder(H, gamma=0.15),
        }

        data = test_decoders(G, decoders)
        data.to_csv(f"sim_results/{used_code}.csv")

    plot_results()


if __name__ == "__main__":
    main()