import cv2;
import math;
import os
import numpy as np;

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);
    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 TransmissionEstimate(im,A,sz):
    omega = 0.95;
    im3 = np.empty(im.shape,im.dtype);

    for ind in range(0,3):
        im3[:,:,ind] = im[:,:,ind]/A[0,ind]

    transmission = 1 - omega*DarkChannel(im3,sz);
    return transmission

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__':
    import sys

    sz = 10 # 窗口函数
    tx = 0.2 # 传输图最小值

    try:
        local_dir = sys.argv[1]
    except:
        local_dir = './image/'
    
    try:
        sz = int(sys.argv[2])
    except:
        sz = sz

    try:
        tx = float(sys.argv[3])
    except:
        tx = tx

    def nothing(*argv):
        pass
    # 检索文件
    src_img_folder = os.path.join(local_dir, 'src')

    imglist = getFileList(src_img_folder, [], '')
    print('本次执行检索到 '+str(len(imglist))+' 张图像\n')
    
    for imgpath in imglist:
        imgname= os.path.splitext(os.path.basename(imgpath))[0]
        # cv2.IMREAD_GRAYSCALE / cv2.IMREAD_COLOR ： 加载灰色 / 彩色图像
        src = cv2.imread(imgpath, cv2.IMREAD_COLOR)
        # 通道归一滑
        I = src.astype('float64')/255;
        # 暗通道图像
        dark = DarkChannel(I,sz=sz);
        # ？？？
        A = AtmLight(I,dark);
        te = TransmissionEstimate(I,A,sz = sz);
        t = TransmissionRefine(src,te);
        print(np.shape(src))
        print(np.shape(dark))
        print(np.shape(t))
        J = Recover(I,t,A,tx=tx); # tx传输图的最小值，用于避免过度曝光

        dark_imgdir = os.path.join(local_dir, 'dark/')
        trans_imgdir = os.path.join(local_dir, 'trans/')
        result_imgdir = os.path.join(local_dir, 'result/')
        cv2.imwrite(dark_imgdir + imgname + "_" + str(sz) + "_" + str(tx) + "_dark.png",dark*255);
        cv2.imwrite(trans_imgdir + imgname + "_" + str(sz) + "_" + str(tx) + "_t.png", t*255);
        cv2.imwrite(result_imgdir + imgname + "_" + str(sz) + "_" + str(tx) + "_result.png",J*255);