# import os
# os.system('nvidia-smi')
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, TensorDataset
import scipy.io as sio
import argparse
import numpy as np
import time
import os
from layer_config_forward import MultiLayerPerceptron_forward


parser = argparse.ArgumentParser()
parser.add_argument("-net_n", type=int,
                    help="number the trained model in the ensemble")
parser.add_argument("-iter_num", type=int,
                    help="iteration_number")
args = parser.parse_args()
net_n = args.net_n
iter_n = args.iter_num

# Check if the directory exists, and create it if it doesn't

print(net_n)
def weights_init(m):
    if type(m) == nn.Linear:
        m.weight.data.normal_(0.0, 1e-3)
        m.bias.data.fill_(0.)

def update_lr(optimizer, lr):
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)
#--------------------------------
# Device configuration
#--------------------------------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device: %s'%device)

#--------------------------------
# Hyper-parameters
#--------------------------------
input_size = 44
hidden_size = [100, 50, 100]
num_classes = 8
num_epochs = 60
# batch_size = 200
learning_rate = 5*1e-4
learning_rate_decay = 0.95
reg = 0.001


mat = sio.loadmat('Dataset/dset_%d.mat' %(iter_n))
mat = mat['dset']
Design_data = mat['X'][0][0]
Prformance_data = mat['Y'][0][0]

x_train_tensor = torch.from_numpy(Design_data).float()
y_train_tensor = torch.from_numpy(Prformance_data).float()


dataset = TensorDataset(x_train_tensor, y_train_tensor)

lengths = [int(len(dataset)*0.9), len(dataset)-int(len(dataset)*0.9)]
torch.manual_seed(net_n)
train_dataset, val_dataset = torch.utils.data.random_split(dataset, lengths)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=100)
val_loader = torch.utils.data.DataLoader(dataset=val_dataset, batch_size=100)



model_forward = MultiLayerPerceptron_forward(input_size, hidden_size, num_classes, net_n).to(device)
print(count_parameters(model_forward))

model_forward.apply(weights_init)
model_forward.to(device)

# Loss and optimizer
def sum_mse(yhat,y):
    return torch.sum(torch.mean((yhat - y) ** 2, dim=1))

criterion_MSE = nn.MSELoss()
criterion_sum_mse = sum_mse
optimizer = torch.optim.Adam(model_forward.parameters(), lr=learning_rate)

# Train the model_forward
lr = learning_rate
total_step = len(train_loader)
time_start = time.time()
for epoch in range(num_epochs):
    for i, (controll, state) in enumerate(train_loader):
        # Move tensors to the configured device
        controll = controll.to(device)
        state = state.to(device)
        #################################################################################
        # Implement the training code                                             #
        optimizer.zero_grad()
        # im = controll.view(31, input_size)
        outputs = model_forward(controll)
        loss = criterion_sum_mse(outputs, state)
        loss.backward()
        optimizer.step()

        if (i+1) % 100 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

    # Code to update the lr
    lr *= learning_rate_decay
    update_lr(optimizer, lr)
    with torch.no_grad():
        correct = 0
        total = 0
        state_all = torch.zeros(0, num_classes).to(device)
        outputs_all = torch.zeros(0, num_classes).to(device)
        for controll, state in val_loader:
            state = state.to(device)
            state_all = torch.cat((state_all,state),0)
            ####################################################
            #evaluation #
            controll = controll.to(device)
            outputs = model_forward(controll)
            outputs_all = torch.cat((outputs_all,outputs),0)

        loss = criterion_MSE(state_all, outputs_all)
        print('Validataion MSE is: {}'.format(loss))
        one_epoch_time = time.time() - time_start


# save the model
model_name = 'Models/iter_%d/mu_net_%d.ckpt' %(iter_n, net_n)
torch.save(model_forward.state_dict(), model_name)




