import os
import math
import argparse
import random
import logging

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from data.data_sampler import DistIterSampler

import options.options as option
from utils import util
from data import create_dataloader, create_dataset
from models import create_model
import cv2
import numpy as np
import data.util as data_util
import skimage.metrics as sm
import csv
from tqdm import tqdm

def path_2_y(path):
    bgr = cv2.imread(path) # path to bgr(255)
    y = data_util.bgr2ycbcr(bgr / 255., only_y=True) # bgr(255) to y(1)
    return y

def path_2_tensor(path):
    bgr = cv2.imread(path) # path to bgr(255)
    rgb = bgr[:, :, [2, 1, 0]] # bgr(255) to rgb(255)
    tensor = torch.from_numpy(np.ascontiguousarray(np.transpose(rgb, (2, 0, 1)))).float().cuda() # rgb(255) to tensor(1)
    return tensor

def tensor_2_y(tensor):
    rgb = tensor.cpu().numpy() * 255. # tensor(1) to rgb(255)
    bgr = rgb[:, :, [2, 1, 0]] # rgb(255) to bgr(255)
    y = data_util.bgr2ycbcr(bgr / 255., only_y=True) # bgr(255) to y(1)
    return y

def traverse_under_folder(folder_root='/home/ubuntu/'):
	folder_leaf = []
	folder_branch = []
	file_leaf = []
	
	index = 0
	for dirpath, subdirnames, filenames in os.walk(folder_root):
	    index += 1
	
	    if len(subdirnames) == 0:
	        folder_leaf.append(dirpath)
	    else:
	        folder_branch.append(dirpath)
	
	    for i in range(len(filenames)):
	        file_leaf.append(os.path.join(dirpath, filenames[i]))

	return folder_leaf, folder_branch, file_leaf

def run_validation(csv_save_folder=None, model=None, image_folder=None, epoch_idx=None):
    print('validating at the end of epoch: ' + str(epoch_idx) + ' ......')
    with torch.no_grad():
        psnr_list = []
        ssim_list = []
        input_idx = [0, 6]
        folders_leaf, _, _ = traverse_under_folder(os.path.join(image_folder, 'LR'))
        for i in tqdm(range(len(folders_leaf))):
            folder = folders_leaf[i]
            _, _, files_leaf = traverse_under_folder(folder)
            files_leaf.sort()
            infer_number = math.floor(len(files_leaf) / 7)
            for i in range(infer_number):
                input_list = []
                for j in range(7):
                    if j in input_idx:
                        input_list.append(path_2_tensor(files_leaf[i * 7 + j]).unsqueeze(0).unsqueeze(0) / 255.) # path to tensor(1)
                    else:
                        continue
                lr_tensor = torch.cat(input_list, dim=1)

                time_Tensors = torch.Tensor([0.1667, 0.3333, 0.5, 0.6667, 0.8333]).unsqueeze(0).cuda()
                hr_tensor = model.netG(lr_tensor, time_Tensors) # tensor(1) to tensor(1)

                for j in range(7):
                    pred_tensor = hr_tensor[:, j, :, :, :].squeeze(0)
                    pred_bgr = util.tensor2img(pred_tensor) / 255.
                    pred_y = data_util.bgr2ycbcr(pred_bgr, only_y=True)
                    gt_y = path_2_y(files_leaf[i * 7 + j].replace('LR', 'HR'))
                    psnr = sm.peak_signal_noise_ratio(image_true=pred_y * 255, image_test=gt_y * 255, data_range=255) # y(1) to PSNR
                    ssim = sm.structural_similarity(im1=pred_y * 255, im2=gt_y * 255, data_range=255, multichannel=False) # y(1) to SSIM
                    psnr_list.append(psnr)
                    ssim_list.append(ssim)
        
        psnr = np.mean(psnr_list)
        ssim = np.mean(ssim_list)
        if os.path.exists(os.path.join(csv_save_folder, 'valid_record.csv')):
            write_header = False
        else:
            write_header = True
        with open(os.path.join(csv_save_folder, 'valid_record.csv'), "a+", newline="") as wf:
            writer = csv.DictWriter(wf, fieldnames=['epoch', 'psnr', 'ssim'])
            if write_header == True:
                writer.writeheader()
            writer.writerow({'epoch': epoch_idx, 'psnr': psnr, 'ssim':ssim})

def init_dist(backend='nccl', **kwargs):
    ''' initialization for distributed training'''
    # if mp.get_start_method(allow_none=True) is None:
    if mp.get_start_method(allow_none=True) != 'spawn':
        mp.set_start_method('spawn')
    rank = int(os.environ['RANK'])
    num_gpus = torch.cuda.device_count()
    torch.cuda.set_device(rank % num_gpus)
    dist.init_process_group(backend=backend, **kwargs)

def main():
    #### options
    parser = argparse.ArgumentParser()
    parser.add_argument('-opt', type=str, help='Path to option YAML file.')
    parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    opt = option.parse(args.opt, is_train=True)

    #### distributed training settings
    if args.launcher == 'none':
        opt['dist'] = False
        rank = -1
        print('Disabled distributed training.')
    else:
        opt['dist'] = True
        init_dist()
        world_size = torch.distributed.get_world_size()
        rank = torch.distributed.get_rank()

    #### loading resume state if exists
    if opt['path'].get('resume_state', None):
        device_id = torch.cuda.current_device()
        resume_state = torch.load(opt['path']['resume_state'], map_location=lambda storage, loc: storage.cuda(device_id))
        option.check_resume(opt, resume_state['iter'])  # check resume options
    else:
        resume_state = None

    #### mkdir and loggers
    if rank <= 0:
        if resume_state is None:
            util.mkdir_and_rename(opt['path']['experiments_root'])  # rename experiment folder if exists
            util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and 'pretrain_model' not in key and 'resume' not in key))

        # config loggers. Before it, the log will not work
        util.setup_logger('base', opt['path']['log'], 'train_' + opt['name'], level=logging.INFO, screen=True, tofile=True)
        logger = logging.getLogger('base')
        logger.info(option.dict2str(opt))
        # tensorboard logger
        if opt['use_tb_logger'] and 'debug' not in opt['name']:
            version = float(torch.__version__[0:3])
            if version >= 1.1:  # PyTorch 1.1
                from torch.utils.tensorboard import SummaryWriter
            else:
                logger.info(
                    'You are using PyTorch {}. Tensorboard will use [tensorboardX]'.format(version))
                from tensorboardX import SummaryWriter
            tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
    else:
        util.setup_logger('base', opt['path']['log'], 'train', level=logging.INFO, screen=True)
        logger = logging.getLogger('base')

    # convert to NoneDict, which returns None for missing keys
    opt = option.dict_to_nonedict(opt)

    #### random seed
    seed = opt['train']['manual_seed']
    if seed is None:
        seed = random.randint(1, 10000)
    if rank <= 0:
        logger.info('Random seed: {}'.format(seed))
    util.set_random_seed(seed)

    torch.backends.cudnn.benckmark = True
    # torch.backends.cudnn.deterministic = True

    #### create train and val dataloader
    dataset_ratio = 1.0
    for phase, dataset_opt in opt['datasets'].items():
        if phase == 'train':
            train_set = create_dataset(dataset_opt)
            train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size'])) # Iter number per epoch
            total_iters = int(opt['train']['niter'])
            total_epochs = int(math.ceil(total_iters / train_size))
            if opt['dist']:
                train_sampler = DistIterSampler(train_set, world_size, rank, dataset_ratio)
                total_epochs = int(math.ceil(total_iters / (train_size * dataset_ratio)))
            else:
                train_sampler = None
            train_loader = create_dataloader(train_set, dataset_opt, opt, train_sampler)
            if rank <= 0:
                logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
                    len(train_set), train_size))
                logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
                    total_epochs, total_iters))
        elif phase == 'val':
            pass
        else:
            raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase))
    assert train_loader is not None

    #### create model
    model = create_model(opt)

    #### resume training
    if resume_state:
        logger.info('Resuming training from epoch: {}, iter: {}.'.format(
            resume_state['epoch'], resume_state['iter']))

        start_epoch = resume_state['epoch'] + 1
        current_step = resume_state['iter']
        model.resume_training(resume_state)  # handle optimizers and schedulers
    else:
        current_step = 0
        start_epoch = 0

    #### training
    logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step))
    for epoch in range(start_epoch, total_epochs + 1):
        if opt['dist']:
            train_sampler.set_epoch(epoch)
        for _, train_data in enumerate(train_loader):
            current_step += 1
            if current_step > total_iters:
                break
        
            #### update learning rate
            model.update_learning_rate(current_step, warmup_iter=opt['train']['warmup_iter'])

            #### training
            model.feed_data(train_data)
            model.optimize_parameters(current_step)

            #### log
            if current_step % opt['logger']['print_freq'] == 0:
                logs = model.get_current_log()
                message = '<epoch:{:3d}, iter:{:8,d}, lr:('.format(epoch, current_step)
                for v in model.get_current_learning_rate():
                    message += '{:.3e},'.format(v)
                message += ')>'
                for k, v in logs.items():
                    message += '{:s}: {:.4e} '.format(k, v)
                    # tensorboard logger
                    if opt['use_tb_logger'] and 'debug' not in opt['name']:
                        if rank <= 0:
                            tb_logger.add_scalar(k, v, current_step)
                if rank <= 0:
                    logger.info(message)

            #### save models and training states
            if current_step % opt['logger']['save_checkpoint_freq'] == 0:
                if rank <= 0:
                    logger.info('Saving models and training states.')
                    model.save(current_step)
                    model.save_training_state(epoch, current_step)

    if rank <= 0:
        logger.info('Saving the final model.')
        model.save('latest')
        logger.info('End of training.')

if __name__ == '__main__':
    main()
