import os
import cv2
import sys
import tqdm
import torch
import argparse

import numpy as np

from PIL import Image

filepath = os.path.split(os.path.abspath(__file__))[0]
repopath = os.path.split(filepath)[0]
sys.path.append(repopath)

from lib import *
from utils.misc import *
from data.dataloader import *
from data.custom_transforms import *

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

def _args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', '-c',     type=str,            default='configs/InSPyReNet_SwinB.yaml')
    parser.add_argument('--source', '-s',     type=str)
    parser.add_argument('--dest', '-d',       type=str,            default=None)
    parser.add_argument('--type', '-t',       type=str,            default='map')
    parser.add_argument('--gpu', '-g',        action='store_true', default=False)
    parser.add_argument('--jit', '-j',        action='store_true', default=False)
    parser.add_argument('--verbose', '-v',    action='store_true', default=False)
    return parser.parse_args()

def get_format(source):
    img_count = len([i for i in source if i.lower().endswith(('.jpg', '.png', '.jpeg'))])
    vid_count = len([i for i in source if i.lower().endswith(('.mp4', '.avi', '.mov' ))])
    
    if img_count * vid_count != 0:
        return ''
    elif img_count != 0:
        return 'Image'
    elif vid_count != 0:
        return 'Video'
    else:
        return ''

def inference(opt, args):
    model = eval(opt.Model.name)(**opt.Model)
    model.load_state_dict(torch.load(os.path.join(
        opt.Test.Checkpoint.checkpoint_dir, 'latest.pth'), map_location=torch.device('cpu')), strict=True)
    
    if args.gpu is True:
        model = model.cuda()
    model.eval()
        
    if args.jit is True:
        if os.path.isfile(os.path.join(opt.Test.Checkpoint.checkpoint_dir, 'jit.pt')) is False:
            model = Simplify(model)
            model = torch.jit.trace(model, torch.rand(1, 3, *opt.Test.Dataset.transforms.static_resize.size).cuda(), strict=False)
            torch.jit.save(model, os.path.join(opt.Test.Checkpoint.checkpoint_dir, 'jit.pt'))
        
        else:
            del model
            model = torch.jit.load(os.path.join(opt.Test.Checkpoint.checkpoint_dir, 'jit.pt'))
                
    save_dir = None
    _format = None
    
    if args.source.isnumeric() is True:
        _format = 'Webcam'

    elif os.path.isdir(args.source):
        save_dir = os.path.join('results', args.source.split(os.sep)[-1])
        _format = get_format(os.listdir(args.source))

    elif os.path.isfile(args.source):
        save_dir = 'results'
        _format = get_format([args.source])
        
    if args.dest is not None:
        save_dir = args.dest
        
    if save_dir is not None:
        os.makedirs(save_dir, exist_ok=True)
    
    sample_list = eval(_format + 'Loader')(args.source, opt.Test.Dataset.transforms)

    if args.verbose is True:
        samples = tqdm.tqdm(sample_list, desc='Inference', total=len(
            sample_list), position=0, leave=False, bar_format='{desc:<30}{percentage:3.0f}%|{bar:50}{r_bar}')
    else:
        samples = sample_list
        
    writer = None
    background = None

    for sample in samples:
        if _format == 'Video' and writer is None:
            writer = cv2.VideoWriter(os.path.join(save_dir, sample['name'] + '.mp4'), cv2.VideoWriter_fourcc(*'mp4v'), sample_list.fps, sample['shape'][::-1])
            samples.total += int(sample_list.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if _format == 'Video' and sample['image'] is None:
            if writer is not None:
                writer.release()
            writer = None
            continue
        
        if args.gpu is True:
            sample = to_cuda(sample)

        with torch.no_grad():
            if args.jit is True:
                out = model(sample['image'])
            else:
                out = model(sample)
                
                    
        pred = to_numpy(out['pred'], sample['shape'])
        img = np.array(sample['original'])
        
        if args.type == 'map':
            img = (np.stack([pred] * 3, axis=-1) * 255).astype(np.uint8)
        elif args.type == 'rgba':
            r, g, b = cv2.split(img)
            pred = (pred * 255).astype(np.uint8)
            img = cv2.merge([r, g, b, pred])
        elif args.type == 'green':
            bg = np.stack([np.ones_like(pred)] * 3, axis=-1) * [120, 255, 155]
            img = img * pred[..., np.newaxis] + bg * (1 - pred[..., np.newaxis])
        elif args.type == 'blur':
            img = img * pred[..., np.newaxis] + cv2.GaussianBlur(img, (0, 0), 15) * (1 - pred[..., np.newaxis])
        elif args.type == 'overlay':
            bg = (np.stack([np.ones_like(pred)] * 3, axis=-1) * [120, 255, 155] + img) // 2
            img = bg * pred[..., np.newaxis] + img * (1 - pred[..., np.newaxis])
            border = cv2.Canny(((pred > .5) * 255).astype(np.uint8), 50, 100)
            img[border != 0] = [120, 255, 155]
        elif args.type.lower().endswith(('.jpg', '.jpeg', '.png')):
            if background is None:
                background = cv2.cvtColor(cv2.imread(args.type), cv2.COLOR_BGR2RGB)
                background = cv2.resize(background, img.shape[:2][::-1])
            img = img * pred[..., np.newaxis] + background * (1 - pred[..., np.newaxis])
        elif args.type == 'debug':
            debs = []
            for k in opt.Train.Debug.keys:
                debs.extend(out[k])
            for i, j in enumerate(debs):
                log = torch.sigmoid(j).cpu().detach().numpy().squeeze()
                log = ((log - log.min()) / (log.max() - log.min()) * 255).astype(np.uint8)
                log = cv2.cvtColor(log, cv2.COLOR_GRAY2RGB)
                log = cv2.resize(log, img.shape[:2][::-1])
                Image.fromarray(log).save(os.path.join(save_dir, sample['name'] + '_' + str(i) + '.png'))    
                # size=img.shape[:2][::-1]
            
            
        img = img.astype(np.uint8)
        
        if _format == 'Image':
            Image.fromarray(img).save(os.path.join(save_dir, sample['name'] + '.png'))
        elif _format == 'Video' and writer is not None:
            writer.write(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        elif _format == 'Webcam':
            cv2.imshow('InSPyReNet', img)

if __name__ == "__main__":
    args = _args()
    opt = load_config(args.config)
    inference(opt, args)