import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from timeit import default_timer as timer

from decoders import proximal, naive_soft_decision
from utility import noise, simulations, encoders


def main():
    # Hamming(7,4) code

    G = np.array([[1, 1, 1, 0, 0, 0, 0],
                  [1, 0, 0, 1, 1, 0, 0],
                  [0, 1, 0, 1, 0, 1, 0],
                  [1, 1, 0, 1, 0, 0, 1]])

    H = np.array([[1, 0, 1, 0, 1, 0, 1],
                  [0, 1, 1, 0, 0, 1, 1],
                  [0, 0, 0, 1, 1, 1, 1]])

    R = np.array([[0, 0, 1, 0, 0, 0, 0],
                  [0, 0, 0, 0, 1, 0, 0],
                  [0, 0, 0, 0, 0, 1, 0],
                  [0, 0, 0, 0, 0, 0, 1]])

    # Define encoder and decoders

    encoder = encoders.Encoder(G)

    decoders = {"naive_soft_decision": naive_soft_decision.SoftDecisionDecoder(G, H, R),
                "proximal_0_01": proximal.ProximalDecoder(H, R, K=100, gamma=0.01),
                "proximal_0_05": proximal.ProximalDecoder(H, R, K=100, gamma=0.05),
                "proximal_0_15": proximal.ProximalDecoder(H, R, K=100, gamma=0.15),
                }

    # Test decoders

    k, n = G.shape
    d = np.zeros(k)  # All-zeros assumption

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

    start_time = timer()

    for decoder_name in decoders:
        decoder = decoders[decoder_name]
        _, BERs_sd = simulations.test_decoder(encoder=encoder,
                                              decoder=decoder,
                                              d=d,
                                              SNRs=SNRs,
                                              N_max=2000)

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

    stop_time = timer()
    print(f"Elapsed time: {stop_time - start_time:2f}")

    # Plot results

    sns.set_theme()

    fig, axes = plt.subplots(1, 1)
    fig.suptitle("Bit-Error-Rates of various decoders")

    for decoder_name in decoders:
        ax = sns.lineplot(data=data, x="SNR", y=f"BER_{decoder_name}", label=f"{decoder_name}")

    ax.set_title("Hamming(7,4) Code")
    ax.set(yscale="log")
    ax.set_yticks([10e-5, 10e-4, 10e-3, 10e-2, 10e-1, 10e0])
    ax.legend()

    plt.show()


if __name__ == "__main__":
    main()