import os
import cv2
import sys
import numpy as np
from skimage import restoration

def make_gaussian(size, fwhm = 3, center=None):
    """ Make a square gaussian kernel.
    size is the length of a side of the square
    fwhm is full-width-half-maximum, which
    can be thought of as an effective radius.
    """

    x = np.arange(0, size, 1, float)
    y = x[:,np.newaxis]

    if center is None:
        x0 = y0 = size // 2
    else:
        x0 = center[0]
        y0 = center[1]

    output = np.exp(-4*np.log(2) * ((x-x0)**2 + (y-y0)**2) / fwhm**2)
    output = output/np.sum(output)

    return output

def deconv_image(img, psf):

    # sharpen each color channel and then reconbine
    img_max = np.max(img)

    if np.max(img) == 0:
        img = np.float32(img)
    else:
        img = np.float32(img) / img_max

    img[img == 0] = 0.0001
    img = restoration.richardson_lucy(img, psf, 7)

    return img*img_max

def get_psf(focal_length, aperture, focus_range, range):

    # estimate the blur circle size from aperture, and distance to focal plane
    theta = np.arctan(aperture/focus_range/2)
    blur_rad = np.abs(focus_range - range)*np.tan(theta)
    blur_pixels = int(blur_rad*focal_length / range)

    # Make a guess of the PSF for sharpening
    #psf = make_gaussian(4*blur_pixels+1, 2*blur_pixels, center=None)
    if blur_pixels >= 2:
        psf = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (blur_pixels*2, blur_pixels*2))
    else:
        psf = np.array([])

    return psf

# main entry point
if __name__=="__main__":

    if len(sys.argv) < 4:
        print('Usage: python sa3d.py [image_dir] [image_prefix] [refocus_range]')
        exit()

    cameras = ['LT', 'LB', 'MT', 'MB', 'RT', 'RB']
    focal_length_m = 0.016
    focal_length = 21.44/0.0024
    focus_ranges = [.5, .5, 0.5, .5, .5, .5]
    aperture = 4.0
    scale_up = 2
    #rop_x = scale_up*3750
    #crop_y = scale_up*2500
    crop_x = scale_up * 2550
    crop_y = scale_up * 1600
    of = scale_up*400

    x_shifts = np.array([-0.02944, -0.02944, 0.0, 0.0, 0.02944, 0.02944])
    y_shifts = np.array([-0.017, 0.017, -0.034, 0.0, -0.017, 0.017])

    refocus_range = float(sys.argv[3])
    image_dir = sys.argv[1]
    image_prefix = sys.argv[2]

    for refocus_range in np.linspace(0.5, 0.52, 101):

        print(refocus_range)

        # load images
        rect_images = np.zeros((of, of, len(cameras)))
        for ind, cam in enumerate(cameras):


            psf = get_psf(focal_length, scale_up*focal_length_m / aperture, focus_ranges[ind], refocus_range)
            #print(psf.shape)
            image_name = os.path.join(image_dir, image_prefix + cam + '.png')
            img = cv2.imread(image_name)
            img = img[:, :, 0]

            t_x = scale_up*x_shifts[ind]*focal_length / refocus_range
            t_y = scale_up*y_shifts[ind] * focal_length / refocus_range
            T = np.float32([[1, 0, t_x], [0, 1, t_y]])
            img = cv2.resize(img, (scale_up*img.shape[1], scale_up*img.shape[0]), interpolation=cv2.INTER_CUBIC)
            height, width = img.shape[:2]
            img = cv2.warpAffine(img, T, (width, height), flags=cv2.INTER_CUBIC)
            rect_img = img[crop_y:(crop_y+of), crop_x:(crop_x+of)]
            #cv2.imshow('image', rect_img)
            #cv2.waitKey()
            #if psf.shape[0] > 2:
            #    rect_img = deconv_image(rect_img, psf).astype('uint16')
            #cv2.imshow('image', rect_img)
            #cv2.waitKey()

            #psf = make_gaussian(7, 5, center=None)
            #rect_img = deconv_image(rect_img, psf).astype('uint16')
            rect_images[:, :, ind] = rect_img


        # compute and save sa image using average
        # sa_img = np.zeros(rect_img.shape)
        #for ind in range(len(cameras)):
        #    sa_img += rect_images[ind]
        #sa_img = (sa_img / len(cameras)).astype('uint8')

        sa_img = np.mean(rect_images, axis=2).astype('uint16')
        psf = get_psf(focal_length, scale_up * focal_length_m / aperture / 2, focus_ranges[0], refocus_range)
        psf = make_gaussian(11, 7, center=None)
        std_dev = np.std(sa_img)
        #sa_img[sa_img < 3*std_dev] = 0
        if psf.shape[0] > 2:
            sa_img = deconv_image(sa_img, psf).astype('uint16')

        cv2.imwrite(os.path.join(image_dir, 'sa_image_range_' + str(int(refocus_range*1000000)) + '.tif'), sa_img)
        #cv2.imshow('image', sa_img)
        #cv2.waitKey()