import sys,os
import caffe
import numpy as np
import cv2
import math

def EditFcnProto(templateFile, height, width):
    with open(templateFile, 'r') as ft:
        template = ft.read()
        outFile = 'DehazeNetFcn.prototxt'
        with open(outFile, 'w') as fd:
            fd.write(template.format(height_15=height+15, width_15=width+15,
                height_11=height+11, width_11=width+11))

def TransmissionEstimate(im_path, height, width):
    caffe.set_mode_cpu()
    
    # Define a safe tile size to prevent INT_MAX overflow (approx 512x512 is safe)
    SAFE_TILE_SIZE = 512
    
    # Use tiling if the image is larger than the safe size
    if height > SAFE_TILE_SIZE or width > SAFE_TILE_SIZE:
        print(f"Image size ({width}x{height}) is large. Using tiling to avoid memory overflow...")
        
        # Determine effective tile size (cannot be larger than image)
        tile_h = min(height, SAFE_TILE_SIZE)
        tile_w = min(width, SAFE_TILE_SIZE)
        
        # Generate prototxt for the TILE size, not the full image size
        EditFcnProto('DehazeFcnTemplate.prototxt', tile_h, tile_w)
        
        # Load networks
        net = caffe.Net('DehazeNet.prototxt', 'DehazeNet.caffemodel', caffe.TEST)
        net_full_conv = caffe.Net('DehazeNetFcn.prototxt', 'DehazeNet.caffemodel', caffe.TEST)
        net_full_conv.params['ip1-conv'][0].data.flat = net.params['ip1'][0].data.flat
        net_full_conv.params['ip1-conv'][1].data[...] = net.params['ip1'][1].data
        
        # Load and pad image
        im = caffe.io.load_image(im_path)
        npad = ((7,8), (7,8), (0,0))
        im_padded = np.pad(im, npad, 'symmetric')
        
        transmission = np.zeros((height, width))
        
        # Setup transformer for the tile size
        transformers = caffe.io.Transformer({'data': net_full_conv.blobs['data'].data.shape})
        transformers.set_transpose('data', (2,0,1))
        transformers.set_channel_swap('data', (2,1,0))
        
        # Process in tiles
        for h in range(0, height, tile_h):
            for w in range(0, width, tile_w):
                # Calculate start/end to handle edges/overlap
                # If we are at the end, shift back to ensure we feed a full tile
                h_start = min(h, height - tile_h)
                w_start = min(w, width - tile_w)
                
                # Extract patch from PADDED image
                # Network expects input size of (Tile + 15), so we slice accordingly
                patch = im_padded[h_start : h_start + tile_h + 15, w_start : w_start + tile_w + 15, :]
                
                # Forward pass
                out = net_full_conv.forward_all(data=np.array([transformers.preprocess('data', patch-0.2)]))
                
                # Reshape output
                block_trans = np.reshape(out['ip1-conv'], (tile_h, tile_w))
                
                # Assign to result buffer
                transmission[h_start : h_start + tile_h, w_start : w_start + tile_w] = block_trans
                
        return transmission

    else:
        # Original logic for small images
        EditFcnProto('DehazeFcnTemplate.prototxt', height, width)
        net = caffe.Net('DehazeNet.prototxt', 'DehazeNet.caffemodel', caffe.TEST)
        net_full_conv = caffe.Net('DehazeNetFcn.prototxt', 'DehazeNet.caffemodel', caffe.TEST)
        net_full_conv.params['ip1-conv'][0].data.flat = net.params['ip1'][0].data.flat
        net_full_conv.params['ip1-conv'][1].data[...] = net.params['ip1'][1].data
        im = caffe.io.load_image(im_path)
        npad = ((7,8), (7,8), (0,0))
        im = np.pad(im, npad, 'symmetric')
        transformers = caffe.io.Transformer({'data': net_full_conv.blobs['data'].data.shape})
        transformers.set_transpose('data', (2,0,1))
        transformers.set_channel_swap('data', (2,1,0))
        out = net_full_conv.forward_all(data=np.array([transformers.preprocess('data', im-0.2)]))
        transmission = np.reshape(out['ip1-conv'], (height,width))
        return transmission

def DarkChannel(im,sz):
    b,g,r = cv2.split(im)
    dc = cv2.min(cv2.min(r,g),b)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(sz,sz))
    dark = cv2.erode(dc,kernel)
    return dark

def AtmLight(im,dark):
    [h,w] = im.shape[:2]
    imsz = h*w
    numpx = int(max(math.floor(imsz/1000),1))
    darkvec = dark.reshape(imsz,1)
    imvec = im.reshape(imsz,3)
    indices = darkvec.argsort()
    indices = indices[imsz-numpx::]
    atmsum = np.zeros([1,3])
    for ind in range(1,numpx):
        atmsum = atmsum + imvec[indices[ind]]
    A = atmsum / numpx
    return A

def Guidedfilter(im,p,r,eps):
    mean_I = cv2.boxFilter(im,cv2.CV_64F,(r,r))
    mean_p = cv2.boxFilter(p, cv2.CV_64F,(r,r))
    mean_Ip = cv2.boxFilter(im*p,cv2.CV_64F,(r,r))
    cov_Ip = mean_Ip - mean_I*mean_p
    mean_II = cv2.boxFilter(im*im,cv2.CV_64F,(r,r))
    var_I   = mean_II - mean_I*mean_I
    a = cov_Ip/(var_I + eps)
    b = mean_p - a*mean_I
    mean_a = cv2.boxFilter(a,cv2.CV_64F,(r,r))
    mean_b = cv2.boxFilter(b,cv2.CV_64F,(r,r))
    q = mean_a*im + mean_b
    return q

def TransmissionRefine(im,et):
    gray = cv2.cvtColor(im,cv2.COLOR_BGR2GRAY)
    gray = np.float64(gray)/255
    r = 60
    eps = 0.0001
    t = Guidedfilter(gray,et,r,eps)
    return t

def Recover(im,t,A,tx = 0.1):
    res = np.empty(im.shape,im.dtype)
    t = cv2.max(t,tx)
    for ind in range(0,3):
        res[:,:,ind] = (im[:,:,ind]-A[0,ind])/t + A[0,ind]
    return res

def getFileList(dir,Filelist, ext=None):
    """
    获取文件夹及其子文件夹中文件列表
    输入 dir：文件夹根目录
    输入 ext: 扩展名
    返回： 文件路径列表
    """
    newDir = dir
    if os.path.isfile(dir):
        if ext is None:
            Filelist.append(dir)
        else:
            if ext in dir[-3:]:
                Filelist.append(dir)
    
    elif os.path.isdir(dir):
        for s in os.listdir(dir):
            newDir=os.path.join(dir,s)
            getFileList(newDir, Filelist, ext)
 
    return Filelist

if __name__ == '__main__':
    if not len(sys.argv) == 2:
        print ('Usage: python DeHazeNet.py haze_img_path')
        exit()
    else:
        im_path = sys.argv[1]

    # 检索文件
    src_img_folder = os.path.join(im_path, 'src')
    imglist = getFileList(src_img_folder, [], '')
    print('本次执行检索到 '+str(len(imglist))+' 张图像\n')
    for img_path in imglist:
        imgname= os.path.splitext(os.path.basename(img_path))[0]
        src = cv2.imread(img_path)

        height = src.shape[0]
        width = src.shape[1]
        
        # Note: EditFcnProto is also called inside TransmissionEstimate if tiling is used
        # We call it here for initialization but it may be overwritten.
        templateFile = 'DehazeFcnTemplate.prototxt'
        EditFcnProto(templateFile, height, width)
        print("-"*5, ' 完成EditFcnProto ',"-"*5)
        
        I = src/255.0
        dark = DarkChannel(I,15)
        A = AtmLight(I,dark)

        te = TransmissionEstimate(img_path, height, width)
        t = TransmissionRefine(src,te)
        J = Recover(I,t,A,0.1)

        print("Finsh All the operation")
        Trans_Esti_imgdir = os.path.join(im_path, 'Trans_Esti/')
        if not os.path.exists(Trans_Esti_imgdir): os.makedirs(Trans_Esti_imgdir)
        print(Trans_Esti_imgdir + imgname + "_Trans_Esti.png")
        cv2.imwrite(Trans_Esti_imgdir + imgname + "_Trans_Esti.png",te*255);
        
        Trans_Refine_imgdir = os.path.join(im_path, 'Trans_Refine/')
        if not os.path.exists(Trans_Refine_imgdir): os.makedirs(Trans_Refine_imgdir)
        print(Trans_Refine_imgdir + imgname + "_Trans_Refine.png")
        cv2.imwrite(Trans_Refine_imgdir + imgname + "_Trans_Refine.png",t*255);
        
        result_imgdir = os.path.join(im_path, 'result/')
        if not os.path.exists(result_imgdir): os.makedirs(result_imgdir)
        print(result_imgdir + imgname + "_result.png")
        cv2.imwrite(result_imgdir + imgname + "_result.png",J*255);