import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import ModelCheckpoint, ReduceLROnPlateau

class Parameters:
    k = 8
    N = 16
    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 build_encoder(input_dim, output_dim):
    inputs = Input(shape=(input_dim,))
    x = Dense(128, activation='relu')(inputs)
    #x = Dense(64, activation='relu')(x)
    outputs = Dense(output_dim, activation='sigmoid')(x)
    model = Model(inputs=inputs, outputs=outputs)
    model.compile(optimizer=Adam(learning_rate=0.001),
                  loss='binary_crossentropy',
                  metrics=['binary_accuracy'])
    return model


def build_decoder(input_dim, num_classes, activation, loss, num_neurons) -> Model:
    inputs = Input(shape=(input_dim,))
    x = Dense(num_neurons, activation="relu")(inputs) #A single layer with X number of neurons. To be experimented with
    outputs = Dense(num_classes, activation=activation)(x)
    model = Model(inputs=inputs, outputs=outputs)
    model.compile(optimizer=Adam(learning_rate=0.001),
                  loss=loss,
                  metrics=['accuracy'])
    return model

def encode(message, G):
    '''
    Method to encode the message using the generator matrix G
    '''
    message = np.array(message)
    encoded_message = np.dot(message, G) % 2
    return encoded_message

def generate_training_data(num_samples):
    X_train = []
    Y_train = []
    for _ in range(num_samples):
        message = np.random.randint(0, 2, size=Parameters.k)
        encoded_message = encode(message, Parameters.G)
        X_train.append(message)
        Y_train.append(encoded_message)
    return np.array(X_train), np.array(Y_train)


def main():
    input_dim = Parameters.k
    output_dim = Parameters.N
    num_samples = 10000
    model_file = 'encoders/encoder_model.keras'

    model = build_encoder(input_dim, output_dim)
    X_train, Y_train = generate_training_data(num_samples)

    callbacks = [
        #ReduceLROnPlateau(monitor='loss', factor=0.9, patience=4, verbose=1),
    ]

    model.fit(X_train, Y_train, epochs=2000, batch_size=256, verbose=1, callbacks=callbacks)
    model.save(model_file)
    print(f'Model saved to {model_file}')

if __name__ == '__main__':
    main()