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

from decoders import proximal
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]])

    # Test decoder

    d = np.array([0, 1, 0, 1])
    c = np.dot(G.transpose(), d) % 2

    print(f"Simulating with c = {c}")

    decoder = proximal.ProximalDecoder(H, K=100, gamma=0.01)
    SNRs, BERs = utility.test_decoder(decoder, c, SNRs=[1, 3, 20], N=1000)

    data = pd.DataFrame({"SNR": SNRs, "BER": BERs})

    ax = sns.lineplot(data=data, x="SNR", y="BER")
    ax.set_ylim([0, 1])
    plt.show()


if __name__ == "__main__":
    main()