import os.path
import logging

import numpy as np
from collections import OrderedDict

import torch

from Denoisers.DRUNet.utils import utils_model
from Denoisers.DRUNet.utils import utils_image as util


def denoise(tensor,model):

    with torch.no_grad():
        img_test = model.denoise(tensor)
    return 0

def denoising_drunet(tensor,model,sigma=10):
    torch.cuda.empty_cache()
    [n1,n2,n3] = tensor.shape
    dn_result = np.zeros_like(tensor, dtype=np.float32)
    noise_in = torch.FloatTensor([sigma / 255.]).repeat(1, tensor.shape[0], tensor.shape[1])

    for i in range(n3):
        input_tensor = torch.tensor(np.expand_dims(tensor[:,:,i],axis=0))
        test_data = torch.unsqueeze(torch.cat((input_tensor,noise_in),dim=0),dim=0)
        with torch.no_grad():
            img_test = model.denoise(test_data)
            dn_result[:, :, i] = img_test.data.squeeze().float().clamp_(0, 1).cpu().numpy()


    return dn_result


def denoising_drunet_torch(tensor, model, sigma=10):
    # 清除缓存
    torch.cuda.empty_cache()
    [n1, n2, n3] = tensor.shape
    dn_result = torch.zeros_like(tensor, dtype=torch.float32,device=tensor.device)
    noise_in = torch.FloatTensor([sigma / 255.]).repeat(1, tensor.shape[0], tensor.shape[1]).cuda()

    for i in range(n3):
        input_tensor = torch.unsqueeze(tensor[:, :, i], dim=0)
        test_data = torch.unsqueeze(torch.cat((input_tensor, noise_in), dim=0), dim=0)
        with torch.no_grad():
            dn_result[:, :, i] = model.denoise(test_data).data.squeeze().float().clamp_(0, 1)

    return dn_result

def denoising_drunet_hsv(tensor,model,sigma=10):
    torch.cuda.empty_cache()
    tensor_size = tensor.shape
    dn_result = torch.zeros_like(tensor,dtype=torch.float32,device=tensor.device)
    scale_vector = torch.ones([1, 1, tensor_size[2]], device=tensor.device)
    scale_vector[:, :, 1] *= 0.5
    scale_vector[:, :, -2] *= 0.5
    scale_vector[:, :, 2:-2] *= 1. / 3
    noise_in = torch.FloatTensor([sigma / 255.]).repeat(1, tensor_size[0], tensor_size[1]).cuda()
    for j in range(tensor_size[-1]):
        compnent_tensor = tensor[:,:,:,j].clone()
        compnent_result = torch.zeros_like(compnent_tensor)
        if not isinstance(compnent_tensor,torch.Tensor):
            compnent_tensor = torch.tensor(compnent_tensor)
        for i in range(tensor_size[2] - 2):
            input_tensor = compnent_tensor[:, :, i:i + 3].permute([2, 0, 1])
            test_data = torch.unsqueeze(torch.cat((input_tensor, noise_in), dim=0), dim=0)
            with torch.no_grad():
                img_test = model.denoise(test_data)
                compnent_result[:, :, i:i + 3] += img_test.data.squeeze().float().clamp_(0, 1).clone().permute(1, 2, 0)

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


def denoising_drunet_color(tensor,model,dn_type=1,sigma=10):
    torch.cuda.empty_cache()
    [n1,n2,n3] = tensor.shape
    dn_result = np.zeros([n1,n2,n3], dtype=np.float32)
    noise_in = torch.FloatTensor([sigma / 255.]).repeat(1, tensor.shape[0], tensor.shape[1]).cuda()
    if dn_type:
        if not isinstance(tensor,torch.Tensor):
            tensor = torch.tensor(tensor)
        for i in range(n3):
            input_tensor = torch.unsqueeze(tensor[:,:,i],axis=0).repeat(3,1,1)
            test_data = torch.unsqueeze(torch.cat((input_tensor,noise_in),dim=0),dim=0)
            with torch.no_grad():
                img_test = model.denoise(test_data)
                dn_result[:, :, i] = img_test.data.squeeze().float().clamp_(0, 1).cpu().numpy()[0,:,:]

    else:
        tensor = torch.tensor(tensor)
        input_tensor = torch.zeros((n1, n2, 3 ),dtype=torch.float32)
        for i in range(n3):
            if i < n3 - 2:
                input_tensor[:, :, :] = tensor[:, :, i:i + 3]
            elif i < n3 - 1:
                input_tensor[:, :, 0:2] = tensor[:, :, -2:]
                input_tensor[:, :, 2] = tensor[:, :, 0]
            else:
                input_tensor[:, :, 0] = tensor[:, :, n3 - 1]
                input_tensor[:, :, 1:] = tensor[:, :, 0:2]
            with torch.no_grad():
                input_tensor = input_tensor.permute([2,0,1])
                test_data = torch.unsqueeze(torch.cat((input_tensor, noise_in), dim=0), dim=0)
                img_test = model.denoise(test_data)
                dn_result[:, :, i] = img_test.data.squeeze().float().clamp_(0, 1).cpu().numpy()[0,:,:]
    return dn_result

def denoising_drunet_color_torch(tensor,model,dn_type=1,sigma=10):
    # 清除缓存
    torch.cuda.empty_cache()
    [n1,n2,n3] = tensor.shape
    dn_result = torch.zeros_like(tensor)
    noise_in = torch.FloatTensor([sigma / 255.]).repeat(1, tensor.shape[0], tensor.shape[1]).cuda()
    if dn_type:
        for i in range(n3):
            input_tensor = torch.unsqueeze(tensor[:,:,i],dim=0).repeat(3,1,1)
            test_data = torch.unsqueeze(torch.cat((input_tensor,noise_in),dim=0),dim=0)
            with torch.no_grad():
                img_test = model.denoise(test_data)
                dn_result[:, :, i] = img_test.data.squeeze().float().clamp_(0, 1)[0,:,:]

    return dn_result

def denoising_color_torch_ho(tensor,model,sigma = 10):
    torch.cuda.empty_cache()
    [n1,n2,n3,n4] = tensor.shape
    noise_in = torch.FloatTensor([sigma / 255.]).repeat(1, n1,n2).cuda()
    dn_result = torch.zeros_like(tensor)
    for i in range(n4):
        input_tensor = tensor[:,:,:,i].squeeze().permute([2,0,1])
        test_data = torch.unsqueeze(torch.cat((input_tensor, noise_in), dim=0), dim=0)
        with torch.no_grad():
            dn_result[:, :, :,i] = model.denoise(test_data).data.squeeze().float().clamp_(0, 1).clone().permute(1,2,0)
    return dn_result

def denoising_color_torch_range(tensor,model,dn_type=1,sigma = 10):
    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)
    noise_in = torch.FloatTensor([sigma / 255.]).repeat(1, tensor.shape[0], tensor.shape[1]).cuda()
    for i in range(n3-2):
        input_tensor = tensor[:,:,i:i+3].permute([2,0,1])
        test_data = torch.unsqueeze(torch.cat((input_tensor, noise_in), dim=0), dim=0)
        with torch.no_grad():
            img_test = model.denoise(test_data)
            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 denoising_rgb_torch(tensor,sigma,model):
    torch.cuda.empty_cache()
    result_inp = torch.zeros_like(tensor)
    noise_in = torch.FloatTensor([sigma / 255.]).repeat(1, tensor.shape[0], tensor.shape[1]).cuda()
    tensor_in = tensor.permute([2, 0, 1]).clone()
    test_data = torch.unsqueeze(torch.cat((tensor_in, noise_in), dim=0), dim=0)
    with torch.no_grad():
        img_test = model.denoise(test_data).data.squeeze().float().clamp_(0, 1).clone().permute(1,2,0)
    return img_test
