#!/usr/bin/env python
from __future__ import division, print_function
import os
import numpy as np
import cv2
import PIL
import glob
import matplotlib.pyplot as plt
from sklearn.isotonic import IsotonicRegression
from scipy.optimize import minimize

# ADAPT imports
from postflight_scripts.image_processing import srgb_gamma_decode_lut

image_dir = '/mnt/homenas2/noaa_adapt/adapt_flights/2021-11-04/images/2021-11-04-14-15-57'
gain_fname = '/mnt/data2tb/libraries/adapt/adapt_ros_ws/data/vignette_correction.tif'


fnames = glob.glob('%s/*.jpg' % image_dir)
#fnames = fnames + glob.glob('%s/*.jpg' % image_dir)

n = 0
for i, fname in enumerate(fnames):
    print('Processing %i/%i' % (i + 1, len(fnames)))
    img = cv2.imread(fname)[:, :, ::-1]
    img = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)[:, :, 2]
    img = srgb_gamma_decode_lut(img)

    img /= np.mean(img)

    if n == 0:
        image = img
        min_image = img
        max_image = img
    else:
        image += img
        min_image = np.minimum(min_image, img)
        max_image = np.minimum(max_image, img)

    n += 1

image /= n*255
#image = np.round(image).astype(np.uint8)
#image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)

# Create a radially-symmetric correction.

X, Y = np.meshgrid(np.linspace(0.5, image.shape[1] - 0.5, image.shape[1]),
                   np.linspace(0.5, image.shape[0] - 0.5, image.shape[0]))

x = X.ravel()
y = Y.ravel()
z = image.ravel()

xc = image.shape[1]/2
yc = image.shape[0]/2


def center_error(x_):
    xc, yc = x_
    r = np.sqrt((x - xc)**2 + (y - yc)**2)
    iso_reg = IsotonicRegression(increasing='auto').fit(r, z)
    z_ = iso_reg.predict(r)

    err = np.sum((z_ - z)**2)
    print(x_, err)
    return err


#ret = minimize(center_error, x0=[xc, yc], method='BFGS', tol=1e-1)

errs = []
xcs = np.linspace(-50, 50, 10) + xc
for i in range(len(xcs)):
    errs.append(center_error([xcs[i], yc]))

p = np.polyfit(xcs, errs, 3)
xcs_ = np.linspace(-50, 50, 1000) + xc
plt.plot(xcs, errs, 'ro')
errs_ = np.polyval(p, xcs_)
plt.plot(xcs_, errs_, 'b-')
xc = xcs_[np.argmin(errs_)]

errs = []
ycs = np.linspace(-50, 50, 10) + yc
for i in range(len(xcs)):
    errs.append(center_error([xc, ycs[i]]))

p = np.polyfit(ycs, errs, 3)
ycs_ = np.linspace(-50, 50, 1000) + yc
plt.plot(ycs, errs, 'ro')
errs_ = np.polyval(p, ycs_)
plt.plot(ycs_, errs_, 'b-')
yc = ycs_[np.argmin(errs_)]


r = np.sqrt((x - xc)**2 + (y - yc)**2)
iso_reg = IsotonicRegression(increasing='auto').fit(r, z)
z_ = iso_reg.predict(r)

r_ = np.linspace(0, r.max(), 1000)
plt.plot(r, z, '.')
plt.plot(r_, iso_reg.predict(r_))

z_ /= z_.max()


if False:
    # Tone down correction a little.
    z_ = z_**(1/1.2)
    #plt.plot(r, z_, '.')


Z = np.reshape(z_, image.shape)

gain = Z.max()/Z

plt.imshow(image*gain)

im = PIL.Image.fromarray(gain.astype(np.float32), mode='F') # float32
im.save(gain_fname)
