count = 0

# Shuffle training data
np.random.shuffle(training_data)

# Define structure of optimal network
HL = 2
HN1, HN2 = 20, 4
EPOCHS = 100
BATCH_SIZE = 50
LR = 0.0007

# Instantiate the network and prepare data
avg_mse = 1
while avg_mse > 0.9:
    net = Net(HN1, HN2)
    training_inputs = training_data[:, 0:4]
    training_labels = training_data[:, 4:]
    test_inputs = testing_data[:, 0:4]
    test_labels = testing_data[:, 4:]

    # Train and test the network
    train(net, training_inputs, training_labels, EPOCHS, LR, BATCH_SIZE)
    avg_mse, predictions_online, predictions_offline = test(
        test_inputs, test_labels, net)
    print(avg_mse)

predictions_online_inverse_transform = scaler_test.inverse_transform(
    predictions_online)
predictions_offline_inverse_transform = scaler_test.inverse_transform(
    predictions_offline)
# k-fold cross validation training loop
HL = 2
HN1 = 10
HN2 = 10
EPOCHS = 50
BATCH_SIZE = 50
LR = [
    0.0001, 0.0002, 0.0003, 0.0004, 0.0005, 0.0006, 0.0007, 0.0008, 0.0009,
    0.001, 0.002, 0.003, 0.004, 0.005, 0.006, 0.007, 0.008, 0.009, 0.01, 0.02,
    0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6,
    0.7, 0.8, 0.9, 1
]
MODELS = {}

net = Net(HN1, HN2)
init_state = copy.deepcopy(net.state_dict())
for lr in LR:
    MSEs = []
    for index, subset in enumerate(subset_train_list):
        subset.value = np.array(subset.value)
        subset_test_list[index].value = np.array(subset_test_list[index].value)

        net.load_state_dict(init_state)
        training_inputs = subset.value[:, 0:5]
        training_labels = subset.value[:, 5:]
        test_inputs = subset_test_list[index].value[:, 0:5]
        test_labels = subset_test_list[index].value[:, 5:]

        train(net, training_inputs, training_labels, EPOCHS, lr, BATCH_SIZE)
        avg_mse = test(test_inputs, test_labels, net)
HL = 2
HN = [(2, 2), (2, 4), (2, 8), (2, 12),
      (2, 16), (2, 20), (4, 2), (4, 4), (4, 8), (4, 12), (4, 16), (4, 20),
      (8, 2), (8, 4), (8, 8), (8, 12), (8, 16), (8, 20), (12, 2), (12, 4),
      (12, 8), (12, 12), (12, 16), (12, 20), (16, 2), (16, 4), (16, 8),
      (16, 12), (16, 16), (16, 20), (20, 2), (20, 4), (20, 8), (20, 12),
      (20, 16), (20, 20)]
EPOCHS = 500
BATCH_SIZE = 50
LR = 0.001
MODELS = {}

print(training_data)

for h1, h2 in HN:
    net = Net(h1, h2)
    init_state = copy.deepcopy(net.state_dict())

    net.load_state_dict(init_state)
    training_inputs = training_data[:, 0:4]
    training_labels = training_data[:, 4:]
    test_inputs = testing_data[:, 0:4]
    test_labels = testing_data[:, 4:]

    E_opt, opt_epochs = train(net, training_inputs, training_labels,
                              test_inputs, test_labels, EPOCHS, LR, BATCH_SIZE)
    MODELS['{b}_{x}-{y}_{z}'.format(b=HL, x=h1, y=h2, z=opt_epochs)] = E_opt

with open(
        'Data3/Search/manual_search_results_{x}HL_hn-e_50rep.csv'.format(x=HL),
        'w') as f:
        count = 0

# Shuffle training data
np.random.shuffle(training_data)

# Define structure of optimal network
HL = 2
HN1, HN2 = 4, 8
EPOCHS = 152
BATCH_SIZE = 300
LR = 0.006

# Instantiate the network and prepare data
avg_mse = 1
while avg_mse > 0.09:
    net = Net(HN1, HN2)
    training_inputs = training_data[:, 0:5]
    training_labels = training_data[:, 5:]
    test_inputs = testing_data[:, 0:5]
    test_labels = testing_data[:, 5:]

    # Train and test the network
    train(net, training_inputs, training_labels, EPOCHS, LR, BATCH_SIZE)
    avg_mse, predictions_online, predictions_offline = test(
        test_inputs, test_labels, net)
    print(avg_mse)

predictions_online_inverse_transform = scaler_test.inverse_transform(
    predictions_online)
predictions_offline_inverse_transform = scaler_test.inverse_transform(
    predictions_offline)