import argparse, os, sys, subprocess
from tqdm import tqdm
from glob import glob
from os.path import *
import importlib

import numpy as np
import torch
import torch.nn as nn
from torch import optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from PIL import Image
import pickle as pkl

from models import VGG_graph_matching
from dataloader import MpiSintelClean, MpiSintelFinal, ImagesFromFolder
from logger import Logger
from utils import flow_utils, tools


def init_config():
    """
    INITALIZE  EVERYTHING (Argument Parser, Logging, GPU)
    """
    parser = argparse.ArgumentParser(description='Deep Learning of Graph matching')

    # dataset
    parser.add_argument('--dataset', type=str, default='sintel', help='dataset to use: middlebury/sintel')
    parser.add_argument('--data_path', type=str, default='data/', help='Path to dataset root directory')

    # select mode
    parser.add_argument('--eval', action='store_true', default=False, help='Perform Inference')
    parser.add_argument('--load_path', type=str, default='')
    parser.add_argument('--save_flow', action='store_true', default=False, help='Save flow files during evaluation')

    # others
    parser.add_argument('--seed', type=int, default=7, metavar='S', help='random seed')
    parser.add_argument('--number_workers', '-nw', '--num_workers', type=int, default=8)
    parser.add_argument('--number_gpus', '-ng', type=int, default=2, help='number of GPUs to use')
    parser.add_argument('--cuda', action='store_true', default=True, help='Use GPU')
    parser.add_argument('--use_vgg', action='store_true', default=True, help='Use VGG weights')


    args = parser.parse_args()
    args.cuda = torch.cuda.is_available()

    save_dir = "models/%s" % args.dataset
    log_dir = "logs/%s" % args.dataset

    config_file = "config_%s" % args.dataset
    params = importlib.import_module(config_file).params

    args = argparse.Namespace(**vars(args), **params)

    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    if not os.path.exists(log_dir):
        os.makedirs(log_dir)


    id_ = "%s_seed-%d" % \
            (args.dataset, args.seed)

    save_path = os.path.join(save_dir, id_ + '.pt')

    args.save_path = save_path

    args.log_path = os.path.join(log_dir, id_ + ".log")
    print("log path", args.log_path)



    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.cuda:
        torch.cuda.manual_seed(args.seed)
        torch.backends.cudnn.deterministic = True


    args.effective_number_workers = args.number_workers * args.number_gpus
    args.device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")

    return args

def _apply_loss(d, d_gt):
    """
    LOSS CALCULATION OF THE BATCH

    Arguments:
    ----------
        - d: Computed displacements
        - d_gt: Ground truth displacements

    Returns:
    --------
        - loss: calculate loss according to the specified loss function
    
    """

    # Set all pixel entries to 0 whose displacement magnitude is bigger than 10px
    pixel_thresh = 10
    dispMagnitude = torch.sqrt(torch.pow(d_gt[:,:,0],2) + torch.pow(d_gt[:,:,1], 2)).unsqueeze(-1).expand(-1,-1,2)
    idx = dispMagnitude > pixel_thresh
    z = torch.zeros(dispMagnitude.shape)
    d = torch.where(idx, z, d)
    d_gt = torch.where(idx, z, d_gt)

    # Calculate loss according to formula in paper
    return torch.sum(torch.sqrt(torch.diagonal(torch.bmm(d - d_gt, (d-d_gt).permute(0,2,1)), dim1=-2, dim2=-1)), dim = 1)


def get_mask(height, width, grid_size = 10):
    """
    Get the location based on the image size corresponding to relu_4_2
    and relu_5_1 layer for a desired grid size.
    """
    print(height, width)
    x_jump = int(width/grid_size)
    y_jump = int(height/grid_size)
    x_idx = np.linspace(int(x_jump/2),int(width - x_jump/2), grid_size, dtype = np.int32)
    y_idx = np.linspace(int(y_jump/2), int(height - y_jump/2), grid_size, dtype = np.int32)
    f_mask = torch.zeros((height//(2**4),width//2**4)).byte()
    u_mask = torch.zeros((height//(2**3),width//2**3)).byte()
    for i in x_idx:
        for j in y_idx:
            f_mask[j//(2**4),i//(2**4)] = 1
            u_mask[j//(2**3),i//(2**3)] = 1
    return(u_mask, f_mask)


def test(args, epoch, model, data_loader):
    """
    TESTING PROCEDURE

    Parameters:
    -----------
        - args: various arguments
        - epoch: number of epochs 
        - model: specified model to test
        - data_loader: specified test data_loader

    Returns:
    --------
        - average_loss: average loss per batch
        - pck: Percentage of Correct Keypoints metric

    """
    
    statistics = []
    total_loss = 0

    model.eval()
    title = 'Validating Epoch {}'.format(epoch)
    progress = tqdm(tools.IteratorTimer(data_loader), ncols=120, total=len(data_loader), smoothing=.9, miniters=1, leave=True, desc=title)
    predictions = []
    gt = []

    sys.stdout.flush()
    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(progress):

            d = model(data[0].to(args.device), im_2 = data[1].to(args.device))
            loss = _apply_loss(d, target).mean()
            total_loss += loss.item()
            predictions.extend(d.numpy())
            gt.extend(target.numpy())

            # Print out statistics
            statistics.append(loss.item())
            title = '{} Epoch {}'.format('Validating', epoch)

            progress.set_description(title + '\tLoss:\t'+ str(statistics[-1]))
            sys.stdout.flush()


    progress.close()
    pck = tools.calc_pck(np.asarray(predictions), np.asarray(gt))
    print('PCK for epoch %d is %f'%(epoch, pck))

    return total_loss / float(batch_idx + 1), pck


def train(args, epoch, model, data_loader, optimizer):
    """
    TRAINING PROCEDURE

     Parameters:
    -----------
        - args: various arguments
        - epoch: number of epochs 
        - model: specified model to test
        - data_loader: specified train data_loader
        - optimizer: specified optimizer to use

    Returns:
    --------
        - average_loss: average loss per batch

    """
    
    statistics = []
    total_loss = 0

    model.train()
    title = 'Training Epoch {}'.format(epoch)
    progress = tqdm(tools.IteratorTimer(data_loader), ncols=120, total=len(data_loader), smoothing=.9, miniters=1, leave=True, desc=title)

    sys.stdout.flush()

    for batch_idx, (data, target) in enumerate(progress):

        #data, target = data.to(args.device), target.to(args.device)

        optimizer.zero_grad()
        d = model(data[0].to(args.device), im_2 = data[1].to(args.device))
        loss = _apply_loss(d, target).mean()

        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        assert not np.isnan(total_loss)

        # Print out statistics
        statistics.append(loss.item())
        title = '{} Epoch {}'.format('Training', epoch)

        progress.set_description(title + '\tLoss:\t'+ str(statistics[-1]))
        sys.stdout.flush()


    progress.close()

    return total_loss / float(batch_idx + 1)

# ====================================================================================================================================
# MAIN PROCEDURE
# =========================

if __name__ == '__main__':
    args = init_config()
    if not args.eval:
        sys.stdout = Logger(args.log_path)


    gpuargs = {'num_workers': args.effective_number_workers,
               'pin_memory': True,
               'drop_last' : True} if args.cuda else {}

    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
    transforms = transforms.Compose([
            transforms.ToTensor(),
            normalize,
        ])



    # Define path to selected dataset
    if args.dataset.lower() == 'sintel':    

        # Check if specified path exists
        if not os.path.exists(os.path.join(args.data_path,'sintel/training')):
            raise Exception("Could not find specified dataset => Use argument --data_path to specify path")

        # Setup data loader    
        train_dataset = MpiSintelFinal(os.path.join(args.data_path, 'sintel/training'), transforms = transforms)
        val_dataset = MpiSintelFinal(os.path.join(args.data_path, 'sintel/training'),  train = False, sequence_list = train_dataset.sequence_list, transforms = transforms)

    else:
        raise Exception('Dataset not supported yet.')
        sys.stdout.flush()


    # Create Traiing and Validation data loaders
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size*torch.cuda.device_count(), shuffle=True, **gpuargs)
    val_loader = DataLoader(val_dataset, batch_size=args.batch_size_test*torch.cuda.device_count(), shuffle=False, **gpuargs)

    # Load VGG graph matching model
    model = VGG_graph_matching()

    # Load vgg parameters if specified
    if args.use_vgg:
        model.copy_params_from_vgg16()

    # Setup GPU use
    if torch.cuda.device_count() > 1 and args.number_gpus > 1:
        model = nn.DataParallel(model)
    print("Using", torch.cuda.device_count(), "GPUs!")
    model = model.to(args.device)

    optimizer = optim.Adam(model.parameters(), lr = 1e-5)
    best_pck = 0.

    # Perform specified trainings epochs
    for i in range(1, args.n_epochs+1):
        train(args, i, model, train_loader, optimizer)
        loss, pck = test(args, i, model, val_loader)
        if pck > best_pck:
            best_pck = pck
            torch.save(model.state_dict(), args.save_path)

