import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

from decoders import proximal
from decoders import naive_soft_decision
from decoders import channel
from decoders import utility


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]])

    encoder = channel.Encoder(G)

    proximal_decoder = proximal.ProximalDecoder(H, K=100, gamma=0.01)
    soft_decision_decoder = naive_soft_decision.SoftDecisionDecoder(G, H)

    # Test decoder

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

    SNRs_sd, BERs_sd = utility.test_decoder(encoder=encoder,
                                            decoder=soft_decision_decoder,
                                            d=d,
                                            SNRs=np.linspace(1, 7, 9),
                                            target_bit_errors=500)

    data = pd.DataFrame({"SNR": SNRs_sd, "BER_sd": BERs_sd})

    # Plot results

    ax = sns.lineplot(data=data, x="SNR", y="BER_sd")
    ax.set(yscale="log")
    ax.set_yticks([10e-5, 10e-4, 10e-3, 10e-2, 10e-1, 10e0])
    # ax.set_ylim([10e-6, 10e0])
    plt.show()


if __name__ == "__main__":
    main()