import argparse
import os
import cv2
import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
from config import cfg
from config import update_config
from core.inference import get_final_preds
from utils.vis import save_debug_images
import glob
from models.pose_hrnet import get_pose_net

def parse_args():
	parser = argparse.ArgumentParser(description='Train keypoints network')
	# general
	parser.add_argument('--cfg',
						help='experiment configure file name',
						default='experiments/coco/hrnet/w48_384x288_adam_lr1e-3.yaml',
						type=str)

	parser.add_argument('opts',
						help="Modify config options using the command-line",
						default=None,
						nargs=argparse.REMAINDER)

	parser.add_argument('--modelDir',
						help='model directory',
						type=str,
						default='')
	parser.add_argument('--logDir',
						help='log directory',
						type=str,
						default='')
	parser.add_argument('--dataDir',
						help='data directory',
						type=str,
						default='./Inputs/')
	parser.add_argument('--prevModelDir',
						help='prev Model directory',
						type=str,
						default='')

	args = parser.parse_args()
	return args

def save_images(img, joints_pred, name,nrow=8, padding=2):
	height = int(img.size(0) + padding)
	width = int(img.size(1) + padding)
	nmaps = 1
	xmaps = min(nrow, nmaps)
	ymaps = int(math.ceil(float(nmaps) / xmaps))
	height = int(batch_image.size(2) + padding)
	width = int(batch_image.size(3) + padding)
	k = 0
	for y in range(ymaps):
		for x in range(xmaps):
			if k >= nmaps:
				break
			joints = batch_joints[k]
			joints_vis = batch_joints_vis[k]
			for joint in joints:
				joint[0] = x * width + padding + joint[0]
				joint[1] = y * height + padding + joint[1]
				cv2.circle(img, (int(joint[0]), int(joint[1])), 2, [255, 0, 0], 2)
			k = k + 1
	cv2.imwrite(f"Results/{name}", img)

def main():
	normalize = transforms.Normalize(
			mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
		)
	transform = transforms.Compose([
		transforms.ToTensor(),
		normalize,
	])
	args = parse_args()
	update_config(cfg, args)
	image_size = np.array(cfg.MODEL.IMAGE_SIZE)

	model = get_pose_net(
		cfg, is_train=False
	)

	if cfg.TEST.MODEL_FILE:
		model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE), strict=False)
	else:
		model_state_file = os.path.join(
			final_output_dir, 'final_state.pth'
		)
		model.load_state_dict(torch.load(model_state_file))

	model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda()
	
	img_path_l = sorted(glob.glob('./Inputs' + '/*'))
	with torch.no_grad():
		for path in img_path_l:
			name  = path.split('/')[-1]
			image = cv2.imread(path)
			image = cv2.resize(image, (384, 288))
			input = transform(image).unsqueeze(0)
			#print(input.shape)
			outputs = model(input)
			if isinstance(outputs, list):
				output = outputs[-1]
			else:
				output = outputs
			print(f"{name} : {output.shape}")
	

if __name__ == '__main__':
	main()