import torch
import torch.nn.functional as F
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision
import numpy as np
import pandas as pd
import os
import cv2
import pickle
import lmdb

import _init_paths
from fire import Fire
from torch.autograd import Variable
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader
from glob import glob
from tqdm import tqdm
from sklearn.model_selection import train_test_split
from tensorboardX import SummaryWriter

from track.config import config
from track.alexnet import SiameseAlexNet
from track.dataset import ImagnetVIDDataset
from track.custom_transforms import Normalize, ToTensor, RandomStretch, RandomCrop, CenterCrop
import argparse

def parse_args():
    parser = argparse.ArgumentParser(description='Train SiameseFC network')

    parser.add_argument('--gpu',help = 'decide which gpu to use',
                        required = False,
                        default = 0,
                        type = int)
    parser.add_argument('--data',help = 'decide which dataset to train',
                        required = False,
                        default = '/export/home/zby/SiamFC-PyTorch/data/posetrack/processed_data',
                        type = str)
    args = parser.parse_args()
    return args

def train(gpu_id, data_dir):
	# loading meta data
	meta_data_path = os.path.join(data_dir, "meta_data.pkl")
	meta_data = pickle.load(open(meta_data_path,'rb'))
	all_videos = [x[0] for x in meta_data]

	# split train/valid dataset
	train_videos, valid_videos = train_test_split(all_videos, test_size=1-config.train_ratio,
			random_state=config.seed)

	# define transforms
	center_crop_size = config.instance_size - config.total_stride
	random_crop_size = config.instance_size - 2 * config.total_stride
	train_z_transforms = transforms.Compose([
		RandomStretch(),
		CenterCrop((config.exemplar_size, config.exemplar_size)),
		#ToTensor()
	])
	train_x_transforms = transforms.Compose([
		RandomStretch(),
		CenterCrop((config.instance_size, config.instance_size))
		#CenterCrop((center_crop_size, center_crop_size)),
		#RandomCrop((random_crop_size, random_crop_size)),
		#ToTensor()
	])
	valid_z_transforms = transforms.Compose([
		CenterCrop((config.exemplar_size, config.exemplar_size)),
		ToTensor(),
	])
	valid_x_transforms = transforms.Compose([
		ToTensor()
	])

	# open lmdb
	db = lmdb.open(data_dir+'.lmdb', readonly=True, map_size=int(50e9))

	# create dataset
	train_dataset = ImagnetVIDDataset(db, train_videos, data_dir, train_z_transforms, train_x_transforms)
	valid_dataset = ImagnetVIDDataset(db, valid_videos, data_dir, valid_z_transforms, valid_x_transforms, training=False)


	# create dataloader
	trainloader = DataLoader(train_dataset, batch_size=config.train_batch_size,
			shuffle=True, num_workers=config.train_num_workers, drop_last=True)
	validloader = DataLoader(valid_dataset, batch_size=config.valid_batch_size,
			shuffle=False, num_workers=config.valid_num_workers, drop_last=True)

	# create summary writer
	if not os.path.exists(config.log_dir):
		os.mkdir(config.log_dir)
	summary_writer = SummaryWriter(config.log_dir)

	# start training
	with torch.cuda.device(gpu_id):
		model = SiameseAlexNet(gpu_id, train=True)
		model.init_models()
		model = model.cuda()
		# for name, k in model.named_parameters():
			# print(name, k.requires_grad)
		optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=config.lr, momentum=config.momentum, weight_decay=config.weight_decay)
		scheduler = StepLR(optimizer, step_size=config.step_size, gamma=config.gamma)

		best_model_filename = None
		best_loss = np.inf
		for epoch in range(config.epoch):
			train_loss = []
			model.train()
			model.set_bn_fix()
			train_length = len(trainloader)
			pbar = tqdm(range(train_length))
			for i, data in enumerate(trainloader):
				exemplar_imgs, instance_imgs, exemplar_pose_ims, instance_pose_ims = data
				exemplar_var, instance_var, exemplar_pose_var, instance_pose_var = Variable(exemplar_imgs.cuda()), Variable(instance_imgs.cuda()), Variable(exemplar_pose_ims.cuda()), Variable(instance_pose_ims.cuda())
				optimizer.zero_grad()
				outputs = model((exemplar_var, instance_var),(exemplar_pose_var, instance_pose_var))
				loss = model.weighted_loss(outputs)
				#print(loss)
				now_loss = loss.data.item()
				pbar_str = 'Epoch:{}   Loss:{:.4f}'.format(epoch+1,now_loss)
				pbar.set_description(pbar_str)
				pbar.update(1)
				loss.backward()
				optimizer.step()
				step = epoch * len(trainloader) + i
				summary_writer.add_scalar('train/loss', loss.data, step)
				train_loss.append(loss.data)
			train_loss = np.mean(train_loss)

			valid_loss = []
			model.eval()
			valid_length = len(validloader)
			pbar = tqdm(range(valid_length))
			for i, data in enumerate(validloader):
				exemplar_imgs, instance_imgs, exemplar_pose_ims, instance_pose_ims = data
				#print(exemplar_imgs.shape,instance_imgs.shape)
				exemplar_var, instance_var, exemplar_pose_var, instance_pose_var = Variable(exemplar_imgs.cuda()), Variable(instance_imgs.cuda()), Variable(exemplar_pose_ims.cuda()), Variable(instance_pose_ims.cuda())
				outputs = model((exemplar_var, instance_var),(exemplar_pose_var, instance_pose_var))
				pbar_str = 'Epoch:{}   Valid Loss:{:.4f}'.format(epoch+1,now_loss)
				pbar.set_description(pbar_str)
				pbar.update(1)
				loss = model.weighted_loss(outputs)
				valid_loss.append(loss.data)
			valid_loss = np.mean(valid_loss)
			print("EPOCH %d valid_loss: %.4f, train_loss: %.4f" %
					(epoch, valid_loss, train_loss))
			summary_writer.add_scalar('valid/loss', valid_loss, (epoch+1)*len(trainloader))

			torch.save(model.state_dict(), "./models/output/siamfc_{}.pth".format(epoch+1))
			scheduler.step()

if __name__ == '__main__':
    args = parse_args()
    gpu_id = args.gpu
    data = args.data
    train(gpu_id,data)


