import random
from matplotlib import pyplot as plt
import numpy as np
import os
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
os.environ["OMP_NUM_THREADS"] = "16"  # Set to the number of logical processors
os.environ["MKL_NUM_THREADS"] = "16"
import tensorflow as tf
import argparse
from tensorflow.keras.utils import to_categorical
from model import *

tf.config.threading.set_intra_op_parallelism_threads(16)
tf.config.threading.set_inter_op_parallelism_threads(16)

class Parameters:
    k = 8
    N = 16
    num_classes = 2 ** k  # One-hot classes
    G = np.array([[1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0],
        [1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0],
        [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])

def generate_test_data_onehot(num_frames, EbOverN0dB):
    X_test = []
    Y_test = []
    EbOverN0Linear = 10 ** (EbOverN0dB / 10)
    for frame in range(num_frames):
        message = np.random.randint(0, 2, size=Parameters.k)
        encoded_message = encode(message, Parameters)
        modulated_message = bpsk_modulation(encoded_message)
        noise_scaling = EbOverN0Linear * (Parameters.k / Parameters.N)
        received_signal = add_noise(modulated_message, noise_scaling)
        X_test.append(received_signal)
        onehot_label = int("".join(map(str, message)), 2)
        Y_test.append(onehot_label)
    return np.array(X_test), to_categorical(np.array(Y_test), num_classes=Parameters.num_classes)

def evaluate_onehot_model(model, num_test_frames=1000, EbOverN0dB=2):
    X_test, Y_test = generate_test_data_onehot(num_test_frames, EbOverN0dB)
    predictions = model.predict(X_test)
    predicted_classes = np.argmax(predictions, axis=1)
    true_classes = np.argmax(Y_test, axis=1)
    total_bits = num_test_frames * Parameters.k
    nn_bit_errors = np.sum(np.unpackbits(np.bitwise_xor(predicted_classes, true_classes).astype(np.uint8)))
    nn_BER = nn_bit_errors / total_bits
    MAP_BER = get_MAP_BEP(EbOverN0dB)
    NVE = nn_BER / MAP_BER
    print(f"\nNeural Decoder Bit Error Rate (BER): {nn_BER:.6f}")
    print(f"MAP Decoder Bit Error Rate (BER):    {MAP_BER:.6f}")
    print(f"Normalized Validation Error (NVE): {NVE:.2f}")
    return nn_BER, NVE

def main():
    EbOverN0dB = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6]
    model_path = 'models/'
    models = os.listdir(model_path)
    fig, ax = plt.subplots(1, 1, figsize=(10, 10))
    num_frames = 10000
    BEP_map = [get_MAP_BEP(EbOverN0) for EbOverN0 in EbOverN0dB]
    plt.semilogy(EbOverN0dB, BEP_map, 'o-', label='ML Simulated')
    for model_file in models:
        BEP_nn = []
        NVE = []
        path = model_path + model_file
        model = tf.keras.models.load_model(path)
        for EbOverN0 in EbOverN0dB:
            if "onehot" in model_file:
                BEP, NVE_nn = evaluate_onehot_model(model, num_test_frames=num_frames, EbOverN0dB=EbOverN0)
            BEP_nn.append(BEP)
            NVE.append(NVE_nn)
        plt.semilogy(EbOverN0dB, BEP_nn, 'o-', label=path)
    plt.xlabel('Eb/N0 (dB)')
    plt.ylabel('BER')
    plt.legend()
    plt.grid()
    plt.show()

if __name__ == '__main__':
    main()
