if __name__ == '__main__':
    import torch
    from torch.utils.data import TensorDataset, DataLoader, random_split
    import numpy as np
    import cv2
    import os


    class AE(torch.nn.Module):
        def __init__(self, img_size, num_nodes):
            super(AE, self).__init__()
            print('The number of input nodes =', img_size)
            self.encoder = torch.nn.ModuleList([
            torch.nn.Linear(img_size, num_nodes),
            torch.nn.LeakyReLU(),
            torch.nn.Linear(num_nodes, 24),
            ])

        def forward(self, x):
            for i, _ in enumerate(self.encoder):
                x = self.encoder[i](x)
            return x

    if not os.path.exists('./data'):
        os.mkdir('./data')
    if os.path.exists('imgs.npy') and os.path.exists('targets.npy'):
        imgs = np.load('imgs.npy')
        targets = np.load('targets.npy')
    else:
        imgs, targets = [], []
        for count in range(3000):
            pos = np.load('./pred_img/downsampled_1/pos_{}.npy'.format(count))
            for i in range(3):
                image = cv2.imread('./pred_img/downsampled_1/ds_{}_{}.png'.format(i, count))
                # flatten image
                image = image.flatten()
                # concatenation
                concat = np.concatenate((image, pos[i]), axis=0)
                target = np.load('./pred_img/downsampled_1/ds_{}_{}.npy'.format(i, count))
                target = target.flatten().reshape(4,6)
                targets.append(target)
                imgs.append(concat)
        np.save('imgs.npy', imgs)
        np.save('targets.npy', targets)

    imgs = torch.from_numpy(np.array(imgs)).float()
    targets = torch.from_numpy(targets).float()
    
    # targets = torch.from_numpy((targets - targets.min())*255).float()

    epochs, bs = 1000, 5
    # size of hidden layer: 25, 50, 100, 200, 300, 1000
    num_nodes = [12, 20, 25, 50, 100, 200, 300, 1000]
    image_size = imgs.shape[1]
    dataset = TensorDataset(imgs, targets)
    train_dataset, test_dataset = random_split(dataset, [2000, 1000])
    train_loader = DataLoader(train_dataset, batch_size=bs)
    test_loader = DataLoader(test_dataset, batch_size=bs)
    loss_function = torch.nn.MSELoss()
    
    # Remeber to substitute parameter 'path' with your correct path.
    path = 'C:\\Users\\hbrch\\Desktop\\poppy-simulations-main\\plots\\'

    print('----------------------Training----------------------')
    for z in num_nodes:
        model = AE(image_size, z)
        optimizer = torch.optim.Adam(model.parameters(), lr = 1e-5 , weight_decay = 1e-8)
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.5)
        plot_loss = []
        
        name = 'H_{}_'.format(z)
        for epoch in range(epochs):
            for i, (image, target) in enumerate(train_loader):
                image = image.squeeze()
                target = target.flatten()
                pre_ten = model(image)
                pre_ten = pre_ten.flatten()
                loss = loss_function(pre_ten, target)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                
                diff = (pre_ten-target)**2
                Euclidean_ls, sum = [], 0
                for i in range(pre_ten.shape[0]):
                    # calculate sum of squared difference of coordinates(size: 3)
                    if i % 3 == 0:
                        sum = 0
                    sum += diff[i]
                    # calculate Euclidean distance
                    if (i+1) % 3 == 0: 
                        Euclidean_ls.append(torch.sqrt(sum).item())
                sqrt_sum = 0
                for ele in Euclidean_ls:
                    sqrt_sum += ele
                ave_sqrt_sum = sqrt_sum/len(Euclidean_ls)

            scheduler.step()
            plot_loss.append(ave_sqrt_sum)
            # elements in plot_loss are pytorch tensors
            print('MSE:', loss.item(), 'average Euclidean distance:', plot_loss[-1], 'epoch:', epoch, 'LR:', optimizer.param_groups[0]['lr'])

        np.save('train_data_{}'.format(z), np.array(plot_loss))
        
        model.eval()
        print('----------------------Testing----------------------')
        te_loss = []
        for i, (image, target) in enumerate(test_loader):
            image = image.squeeze()
            target = target.flatten()
            pre_ten = model(image)
            pre_ten = pre_ten.flatten()
            loss = loss_function(pre_ten, target)

            diff = (pre_ten-target)**2
            Euclidean_ls, sum = [], 0
            for i in range(pre_ten.shape[0]):
                # calculate sum of squared difference of coordinates(size: 3)
                if i % 3 == 0:
                    sum = 0
                sum += diff[i]
                # calculate Euclidean distance
                if (i+1) % 3 == 0: 
                    Euclidean_ls.append(torch.sqrt(sum).item())
            sqrt_sum = 0
            for ele in Euclidean_ls:
                sqrt_sum += ele
            ave_sqrt_sum = sqrt_sum/len(Euclidean_ls)
            te_loss.append(ave_sqrt_sum)
            print('average Euclidean distance', te_loss[-1])
            # diff = pre_ten-target
            # print('top 3 largest values:', list(torch.topk(diff, 3)[0].detach().numpy()))
            # print('top 3 smallest values:', list(torch.topk(diff, 3, largest=False)[0].detach().numpy()))

        np.save('./data/test_data_{}'.format(z), np.array(te_loss))
    pass