import os.path
import logging

import numpy as np
from collections import OrderedDict

import torch
import random
import json
from torch.autograd import Variable

from Inpainters.CRUnet.utils import utils_logger
# from utils import utils_model
from Inpainters.CRUnet.utils import utils_image as util
import argparse
from Inpainters.CRUnet.utils import utils_option as option
from Inpainters.CRUnet.utils import my_utils
from Inpainters.CRUnet.data.dataset_unet import DatasetUNet
import math
from torch.utils.data import DataLoader
from Inpainters.CRUnet.models.Unet import CRUNet
import cv2


def inpaint_uninit(tensor,mask,model):
    torch.cuda.empty_cache()
    test_data = {}
    n3 = tensor.shape[2]
    result_inp = np.zeros_like(tensor)
    for i_ipt in range(n3):
        tensor_in = torch.tensor(np.expand_dims(tensor[:, :, i_ipt], axis=0))
        mask_in = 1 - mask[:, :, i_ipt]
        mask_in = torch.tensor(np.expand_dims(mask_in, axis=0))
        test_data = torch.unsqueeze(torch.cat((tensor_in, mask_in), dim=0), dim=0)
        with torch.no_grad():
            img_test = model.inpaint(test_data, tensor_in, mask_in)
            result_inp[:, :, i_ipt] = img_test.data.squeeze().float().clamp_(0, 1).cpu().numpy()
    return result_inp

def inpaint_uninit_torch(tensor,mask,model):

    torch.cuda.empty_cache()
    n3 = tensor.shape[2]
    result_inp = torch.zeros_like(tensor)
    for i_ipt in range(n3):
        tensor_in = torch.unsqueeze(tensor[:, :, i_ipt], dim=0)
        mask_in = 1 -(mask[:, :, i_ipt].float())
        mask_in = torch.unsqueeze(mask_in, dim=0).cuda()
        test_data = torch.unsqueeze(torch.cat((tensor_in, mask_in), dim=0), dim=0)
        with torch.no_grad():
            result_inp[:, :, i_ipt] =  model.inpaint(test_data, tensor_in, mask_in).data.squeeze().float().clamp_(0, 1)
    return result_inp

def inpaint_color_torch_single(tensor,mask,model):
    torch.cuda.empty_cache()
    [n1,n2,n3] = tensor.shape
    dn_result = torch.zeros_like(tensor)
    for i in range(n3):
        input_tensor = torch.unsqueeze(tensor[:,:,i],dim=0).repeat(3,1,1)
        input_mask = 1-torch.unsqueeze((mask[:, :, i].float()), dim=0).repeat(3, 1, 1)
        test_data = torch.unsqueeze(torch.cat((input_tensor,input_mask), dim=0), dim=0).clone()
        with torch.no_grad():
            img_test = model.inpaint(test_data, input_tensor, input_mask)
            dn_result[:, :, i] = img_test.data.squeeze().float().clamp_(0, 1)[0, :, :]

    return dn_result

def inpaint_hsv_torch(tensor,mask,model):
    torch.cuda.empty_cache()
    [n1,n2,n3,n4] = tensor.shape
    dn_result = torch.zeros_like(tensor)
    for j in range(n4):
        compnent_tensor = tensor[:,:,:,j]
        compnent_mask = mask[:,:,:,j]
        scale_vector = torch.ones([1, 1, n3], device=tensor.device)
        scale_vector[:, :, 1] *= 0.5
        scale_vector[:, :, -2] *= 0.5
        scale_vector[:, :, 2:-2] *= 1. / 3
        temp_result = torch.zeros_like(compnent_tensor, device=tensor.device)
        for i in range(n3 - 2):
            input_tensor = compnent_tensor[:, :, i:i + 3].clone().permute([2, 0, 1])
            input_mask = 1 - (compnent_mask[:, :, i:i + 3].clone().permute([2, 0, 1]).float())
            test_data = torch.unsqueeze(torch.cat((input_tensor, input_mask), dim=0), dim=0).clone()
            with torch.no_grad():
                img_test = model.inpaint(test_data, input_tensor, input_mask)
                temp_result[:, :, i:i + 3] += img_test.data.squeeze().float().clamp_(0, 1).clone().permute(1, 2, 0)

        temp_result = temp_result * scale_vector
        dn_result[:,:,:,j] = temp_result
    return dn_result

def inpaint_color_torch_ho(tensor,mask,model):
    torch.cuda.empty_cache()
    [n1,n2,n3,n4] = tensor.shape
    dn_result = torch.zeros_like(tensor)
    for i in range(n4):
        input_tensor = tensor[:, :, :, i].squeeze().permute([2, 0, 1])
        input_mask = 1 - (mask[:, :, :, i].squeeze().permute([2, 0, 1]).float())
        test_data = torch.unsqueeze(torch.cat((input_tensor,input_mask), dim=0), dim=0)
        with torch.no_grad():
            dn_result[:, :, :,i] = model.inpaint(test_data, input_tensor, input_mask).data.squeeze().float().permute([1,2,0]).clamp_(0, 1)

    return dn_result
def inpaint_color_torch(tensor,mask,model):
    torch.cuda.empty_cache()
    [n1,n2,n3] = tensor.shape
    dn_result = torch.zeros_like(tensor)
    for i in range(n3-2):
        input_tensor = tensor[:, :, i:i + 3].clone().permute([2, 0, 1])
        input_mask = 1 - (mask[:, :, i:i + 3].clone().permute([2, 0, 1]).float())
        test_data = torch.unsqueeze(torch.cat((input_tensor,input_mask), dim=0), dim=0).clone()
        with torch.no_grad():
            img_test = model.inpaint(test_data, input_tensor, input_mask)
            if i < n3-3:
                dn_result[:, :, i] = img_test.data.squeeze().float().clone().clamp_(0, 1)[0, :, :]
            else:
                dn_result[:, :, i:] = img_test.data.squeeze().float().permute([1,2,0]).clamp_(0, 1).clone()

    return dn_result

def  inpaint_color_torch_range(tensor,mask,model):
    torch.cuda.empty_cache()
    [n1,n2,n3] = tensor.shape
    scale_vector = torch.ones([1,1,n3],device=tensor.device)
    scale_vector[:,:,1] *= 0.5
    scale_vector[:,:,-2] *= 0.5
    scale_vector[:,:,2:-2] *= 1./3
    dn_result = torch.zeros_like(tensor,device=tensor.device)
    for i in range(n3 - 2):
        input_tensor = tensor[:, :, i:i + 3].clone().permute([2, 0, 1])
        input_mask = 1 - (mask[:, :, i:i + 3].clone().permute([2, 0, 1]).float())
        test_data = torch.unsqueeze(torch.cat((input_tensor,input_mask), dim=0), dim=0).clone()
        with torch.no_grad():
            img_test = model.inpaint(test_data, input_tensor, input_mask)
            dn_result[:, :, i:i + 3] += img_test.data.squeeze().float().clamp_(0, 1).clone().permute(1, 2, 0)

    dn_result = dn_result * scale_vector
    return dn_result



def inpaint_uninit_rgb(tensor,mask,model,data_choice='multichannel'):

    torch.cuda.empty_cache()
    test_data = {}
    n3 = tensor.shape[3]
    result_inp = np.zeros_like(tensor)
    if data_choice=='RGB':
        tensor_in = torch.tensor(tensor.transpose(0,3,1,2))
        mask_in = 1 - mask
        mask_in = torch.tensor(mask_in.transpose(0,3,1,2))
        test_data = torch.cat((tensor_in, mask_in),dim=1)
        with torch.no_grad():
            img_test = model.inpaint(test_data, tensor_in, mask_in)
            result_inp = img_test.data.float().clamp_(0, 1).cpu().numpy()
            result_inp = np.transpose(result_inp,(0,2,3,1))
    else:
        print("please input rgb image")
        exit(0)
    return result_inp


def inpaint_uninit_msi(tensor,mask,model):
    [n1,n2,n3] = tensor.shape
    input_tensor = np.zeros((1,n1,n2,3),dtype=np.float32)
    output_tensor = np.zeros((n1,n2,3),dtype=np.float32)
    ipt_rst = np.zeros_like(tensor,dtype=np.float32)
    input_mask = np.zeros_like(input_tensor,dtype=np.float32)
    for i in range(n3):
        if i < n3-2:
            input_tensor[0,:,:,:] = tensor[:,:,i:i+3]
            input_mask[0,:,:,:] = mask[:,:,i:i+3]
        elif i<n3-1:
            input_tensor[0,:, :, 0:2]=tensor[:,:,-2:]
            input_tensor[0,:,:,2] = tensor[:,:,0]
            input_mask[0,:, :, 0:2] = mask[:, :, -2:]
            input_mask[0,:, :, 2] = mask[:, :, 0]
        else:
            input_tensor[0,:, :, 0]=tensor[:,:,n3-1]
            input_tensor[0,:,:,1:] = tensor[:,:,0:2]
            input_mask[0,:, :, 0] = mask[:, :, n3 - 1]
            input_mask[0,:, :, 1:] = mask[:, :, 0:2]

        output_tensor = np.squeeze(inpaint_uninit_rgb(input_tensor,input_mask,model,'RGB'),axis=0)
        ipt_rst[:,:,i] = output_tensor[:,:,0]
    return ipt_rst





