import os
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import DataLoader
import argparse
import time
import datetime
from RMBG.network import Network_3D_Unet
import numpy as np

from RMBG.data_process import test_preprocess_lessMemoryNoTail_SubImg, testset, multibatch_test_save, singlebatch_test_save
from skimage import io
from RMBG.utils import FFDrealign8, inv_FFDrealign8
#############################################################################################################################################
def wf2nobg_ffd8(net,
                sub_img,
                RMBG_GPU='0', 
                RMBG_batch_size=1, 
                RMBG_img_w=256, 
                RMBG_img_h=256, 
                RMBG_img_s=256, 
                RMBG_gap_w=224, 
                RMBG_gap_h=224, 
                RMBG_gap_s=224, 
                RMBG_normalize_factor=1):
    #############################################################################################################################################
    opt = {}
    #############################################################################################################################################
    opt['GPU'] = RMBG_GPU
    opt['batch_size'] = RMBG_batch_size
    opt['img_w'] = RMBG_img_w
    opt['img_h'] = RMBG_img_h
    opt['img_s'] = RMBG_img_s

    opt['gap_w'] = RMBG_gap_w
    opt['gap_h'] = RMBG_gap_h
    opt['gap_s'] = RMBG_gap_s

    opt['normalize_factor'] = RMBG_normalize_factor
    #############################################################################################################################################
    os.environ["CUDA_VISIBLE_DEVICES"] = str(opt['GPU'])

    ##############################################################################################################################################################
    if torch.cuda.is_available():
        net = net.cuda()
    cuda = True if torch.cuda.is_available() else False
    Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
    ##############################################################################################################################################################
    name_list, noise_img, coordinate_list = test_preprocess_lessMemoryNoTail_SubImg(opt, sub_img)

    #print(len(name_list))
    prev_time = time.time()
    time_start = time.time()
    denoise_img = np.zeros(noise_img.shape)
    input_img = np.zeros(noise_img.shape)

    test_data = testset(name_list, coordinate_list, noise_img)
    testloader = DataLoader(test_data, batch_size=opt['batch_size'], shuffle=False, num_workers=4)
    for iteration, (noise_patch,single_coordinate) in enumerate(testloader):
        noise_patch=noise_patch.cuda().float()
        # print('noise_patch -----> ',noise_patch.shape)
        real_A = FFDrealign8(noise_patch)
        real_A = Variable(real_A)
        
        # print('real_A -----> ',real_A.shape)
        fake_B = net(real_A)
        # print('fake_B -----> ',fake_B.shape)
        ################################################################################################################
        # Determine approximate time left
        batches_done = iteration
        batches_left = 1 * len(testloader) - batches_done
        time_left_seconds = int(batches_left * (time.time() - prev_time))
        time_left = datetime.timedelta(seconds=time_left_seconds)
        prev_time = time.time()
        ################################################################################################################
        if iteration % 1 == 0:
            time_end = time.time()
            time_cost = time_end - time_start  # datetime.timedelta(seconds= (time_end - time_start))
            print(
                '\r\033[1;32m[RMBG]\033[0m [Patch %d/%d] [Time Cost: %.0d s] [ETA: %s s]     '
                % ( iteration + 1,
                    len(testloader),
                    time_cost,
                    time_left_seconds
                ), end=' ')

        if (iteration + 1) % len(testloader) == 0:
            print('\n', end=' ')
        ################################################################################################################
        fake_B_realign = inv_FFDrealign8(fake_B)
        output_image = np.squeeze(fake_B_realign.cpu().detach().numpy())
        # print('output_image',output_image.shape)
        real_A_realign = inv_FFDrealign8(real_A)
        raw_image = np.squeeze(real_A_realign.cpu().detach().numpy())
        if(output_image.ndim==3):
            turn=1
        else:
            turn=output_image.shape[0]
        #print(turn)
        if(turn>1):
            for id in range(turn):
                #print('shape of output_image -----> ',output_image.shape)
                aaaa,bbbb,stack_start_w,stack_end_w,stack_start_h,stack_end_h,stack_start_s,stack_end_s=multibatch_test_save(single_coordinate,id,output_image,raw_image)
                denoise_img[stack_start_s:stack_end_s, stack_start_h:stack_end_h, stack_start_w:stack_end_w] \
                = aaaa 
                input_img[stack_start_s:stack_end_s, stack_start_h:stack_end_h, stack_start_w:stack_end_w] \
                = bbbb
        else:
            aaaa, bbbb, stack_start_w, stack_end_w, stack_start_h, stack_end_h, stack_start_s, stack_end_s = singlebatch_test_save(
                single_coordinate, output_image, raw_image)
            denoise_img[stack_start_s:stack_end_s, stack_start_h:stack_end_h, stack_start_w:stack_end_w] \
                            = aaaa 
            input_img[stack_start_s:stack_end_s, stack_start_h:stack_end_h, stack_start_w:stack_end_w] \
                = bbbb
            

    # del noise_img
    output_img = denoise_img.squeeze().astype(np.float32) * opt['normalize_factor']
    del denoise_img
    # print(np.min(output_img),' -----> ',np.max(output_img))
    output_img = np.clip(output_img, 0, 65535).astype('uint16')
    # output_img = output_img - output_img.min()
    # output_img = output_img.astype('uint16')

    return output_img

