import keras
from keras.models import Sequential
from keras.layers import Dense, LSTM, Dropout, SimpleRNN
from keras import regularizers
import numpy as np
import argparse


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_size', type=int, required=True)
    parser.add_argument('--model_config', type=str, required=True,
                        choices=['lstm', 'ffn', 'rnn'])
    parser.add_argument('--epochs', type=int, required=True)
    parser.add_argument('--reg', type=float, default=0)
    parser.add_argument('--hidden_size', type=int, required=True)
    parser.add_argument('--train_examples', type=int, default=100000)
    parser.add_argument('--dropout', type=int, default=0)
    parser.add_argument('--lr', type=float, required=True)
    parser.add_argument('--batch_size', type=int, required=True)
    args = parser.parse_args()

    train_fts = np.random.randint(0, 2, (args.train_examples, args.input_size))
    train_lbl = (train_fts.sum(axis=1) % 2).astype(int)

    test_fts = np.random.randint(0, 2, (args.train_examples // 10, args.input_size))
    test_lbl = (test_fts.sum(axis=1) % 2).astype(int)

    set_train_fts = set()
    for x in train_fts:
        set_train_fts.add(tuple(x))
    count_test_present = sum(tuple(x) in set_train_fts for x in test_fts)
    print("Fraction of test instances present in training set: {:.3f}".format(count_test_present / len(test_fts)))

    model = Sequential()
    if args.model_config == 'ffn':
        model.add(Dense(args.hidden_size, activation='relu', input_dim=args.input_size))
                        # kernel_regularizer=regularizers.l2(args.reg),
                        # bias_regularizer=regularizers.l2(args.reg)))
        if args.dropout > 0:
            model.add(Dropout(args.dropout))
        model.add(Dense(1, activation='sigmoid'))
    elif args.model_config == 'rnn':
        train_fts = np.expand_dims(train_fts, 2)
        test_fts = np.expand_dims(test_fts, 2)
        model.add(SimpleRNN(args.hidden_size, input_shape=(None, 1)))
        model.add(Dense(1, activation='sigmoid'))
    else:
        train_fts = np.expand_dims(train_fts, 2)
        test_fts = np.expand_dims(test_fts, 2)
        model.add(LSTM(args.hidden_size, input_shape=(None, 1)))
        model.add(Dense(1, activation='sigmoid'))

    print("Model type:", args.model_config)

    optimizer = keras.optimizers.Adam(lr=args.lr)

    model.compile(loss='binary_crossentropy',
                  optimizer=optimizer,
                  metrics=['accuracy'])

    model.fit(train_fts, train_lbl, epochs=args.epochs, batch_size=128,
              validation_data=(test_fts, test_lbl))


if __name__ == "__main__":
    main()