"""
# Planet: Understanding the Amazon from Space
"""
"""
*Removing haze from images with Single Image Dehazing Using Improved Dark Channel Prior.*
"""
"""
The paper [Review on Haze Removal Methods](http://www.ijsrp.org/research-paper-0716/ijsrp-p5522.pdf) explores different techniques for removing haze from an image.

The techniques are split into two categories; multi-image and single image. Since multi-image techniques require multiple images of the same scene, which we don't have, we need to focus on the singel image haze removal techniques.

From reading the paper, [Improved single image dehazing using dark channel prior](https://www.researchgate.net/profile/Gupta_Ashutosh/publication/276223001_Single_Image_Dehazing_Using_Improved_Dark_Channel_Prior/links/5552f22a08ae6943a86d8d20/Single-Image-Dehazing-Using-Improved-Dark-Channel-Prior.pdf) looks the most promising.

This technique should be the same as *Dark channel prior*, except when calculating atmospheric light, the window size is increased to 31.

An implementation of the *Dark channel prior* technique is found in [this Github repo](https://github.com/He-Zhang/image_dehaze).

Below is a modified version of that code.
"""
from fastai.vision.all import *
import cv2

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 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

# CUSTOM
def cv2_to_plt(image):
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

def dehaze(image_path:Path, output_path:Path):
    src = cv2.imread(str(image_path))

    I = src.astype('float64')/255
    dark = DarkChannel(I,15)
    A = AtmLight(I,dark)
    te = TransmissionEstimate(I,A,15)
    t = TransmissionRefine(src,te)
    J = Recover(I,t,A,0.1)
    
    cv2.imwrite(str(output_path/image_path.name), J*255);
"""
## Single image
"""
fn = '../input/planets-dataset/planet/planet/train-jpg/train_10002.jpg'

src = cv2.imread(fn)

I = src.astype('float64')/255

dark = DarkChannel(I,15)
A = AtmLight(I,dark)
te = TransmissionEstimate(I,A,15)
t = TransmissionRefine(src,te)
J = Recover(I,t,A,0.1);
plt.imshow(dark);
plt.imshow(t);
plt.imshow(cv2_to_plt(src));
plt.imshow(cv2_to_plt(J.astype('float32')));
"""
## Training images
"""
# Input
train_img_path = Path('../input/planets-dataset/planet/planet/train-jpg')
img_files = get_image_files(train_img_path)

# Output
train_img_output_path = Path("./train-jpg")
train_img_output_path.mkdir(exist_ok=True)
for img_file in progress_bar(img_files):
    dehaze(img_file, train_img_output_path)
"""
## Test images
"""
# Input
test_img_path = Path('../input/planets-dataset/planet/planet/test-jpg')
img_files = get_image_files(test_img_path)

# Output
test_img_output_path = Path("./test-jpg")
test_img_output_path.mkdir(exist_ok=True)
for img_file in progress_bar(img_files):
    dehaze(img_file, test_img_output_path)
"""
## Additional Test Images
"""
# Input
test_img_path = Path('../input/planets-dataset/test-jpg-additional/test-jpg-additional')
img_files = get_image_files(test_img_path)

# Output
test_img_output_path = Path("./test-jpg")
test_img_output_path.mkdir(exist_ok=True)