"""
Main loop for deep learning tasks.
"""
import os
import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.nn.utils import clip_grad_value_

from network import MyNetwork
from loss import MyLoss
from helpers import get_args, get_current_epoch
from data import MyDataset


if __name__ == '__main__':
    args = get_args()
    # Set constants
    EXPERIMENT = args.EXPERIMENT
    TOTAL_EPOCHS = args.TOTAL_EPOCHS  # Total epochs to train
    BATCH_SIZE = args.BATCH_SIZE
    NUM_WORKERS = args.NUM_WORKERS
    PRINT_EVERY = args.PRINT_EVERY  # Print statistics every # of batches
    OPTIMIZER = args.OPTIMIZER
    LEARNING_RATE = args.LEARNING_RATE
    CLIP = args.CLIP
    MOMENTUM = args.MOMENTUM
    EXP_PATH = args.EXP_PATH
    CPU = args.CPU

    # Initialize variables
    torch.manual_seed(1)
    dataset = MyDataset('dataset.csv', 'dataset.csv', 1, 1)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE,
                            num_workers=NUM_WORKERS, shuffle=True)
    network1 = MyNetwork()
    device = torch.device("cuda" if torch.cuda.is_available() and not CPU
                          else "cpu")
    loss_f = MyLoss()
    if OPTIMIZER == 'SGD':
        optimizer = optim.SGD(network1.parameters(), lr=LEARNING_RATE)
    elif OPTIMIZER == 'ADAM':
        optimizer = optim.Adam(network1.parameters(), lr=LEARNING_RATE)
    optimizer = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, factor=0.5, patience=5, verbose=True)

    # Continue an experiment
    if os.path.exists(EXP_PATH):
        CURRENT_EPOCH = get_current_epoch(EXP_PATH)
        print('Continuing experiment %i from epoch %i to %i on %s.'
              % (EXPERIMENT, CURRENT_EPOCH, TOTAL_EPOCHS, device))
        network1.load_state_dict(torch.load(
            EXP_PATH + '/net_eoe%i.pt' % (CURRENT_EPOCH - 1)))
        optimizer.load_state_dict(torch.load(
            EXP_PATH + '/opt_eoe%i.pt' % (CURRENT_EPOCH - 1)))
    else:  # Start a new experiment
        CURRENT_EPOCH = 1
        print('Starting experiment %i for %i epochs on %s.'
              % (EXPERIMENT, TOTAL_EPOCHS, device))
        try:
            os.mkdir(os.path.expanduser('~') + '/Experiments')
        except Exception as exception:
            print(exception)
            print('Experiments directory already exists.')
        os.mkdir(EXP_PATH)

    # Main loop
    for cur_ep in range(CURRENT_EPOCH, TOTAL_EPOCHS+1):
        print('-'*44)
        network1.to(device)
        running_loss = 0
        for cur_ba, [x, y] in enumerate(dataloader, 1):
            # Get a batch of data
            # x has dimension Batch x channel x height x width
            x = torch.randn(10, 1, 32, 32).to(device)
            y = (torch.rand(10)*10).to(device=device, dtype=torch.long)

            # Forward + backward prop, update params
            network1.zero_grad()
            est_y = network1(x)
            train_loss_val = loss_f(est_y, y)
            train_loss_val.backward()
            clip_grad_value_(network1.parameters(), CLIP)
            with torch.no_grad():
                validation_loss_val = train_loss_val
            optimizer.step(validation_loss_val)

            # Print stats
            running_loss += train_loss_val.item()
            if cur_ba % PRINT_EVERY == 0:
                print('EX: %2i | EP: %2i | BA: %6i | LOSS: %.5f'
                      % (EXPERIMENT, cur_ep, cur_ba, running_loss/PRINT_EVERY))
                running_loss = 0
        # End of epoch (eoe) operations
        network1.to('cpu')
        torch.save(network1.state_dict(), EXP_PATH+'/net_eoe%i.pt' % cur_ep)
        torch.save(optimizer.state_dict(), EXP_PATH+'/opt_eoe%i.pt' % cur_ep)

    print('Finished training.')
