'''
test galsim.interpolatedImage & galsim.convolve
'''

import os
import unittest

import numpy as np
import matplotlib.pyplot as plt
import galsim
import scipy.io
from scipy import ndimage


def psfEncircle(img, fraction=0.8, psfSampleSizeInMicrons=2.5, focalLengthInMeters=28, cenPix=None):
    #imgMaxPix_x, imgMaxPix_y = findMaxPix(img)
    y,x = ndimage.center_of_mass(img)  #y-rows, x-cols
    imgMaxPix_x = x #int(x)
    imgMaxPix_y = y #int(y)
    if cenPix != None:
        imgMaxPix_x = cenPix[0]
        imgMaxPix_y = cenPix[1]

    im1 = img.copy()
    im1size = im1.shape

    dis = np.zeros_like(img)
    for irow in range(im1size[0]):
        for icol in range(im1size[1]):
            dx = icol - imgMaxPix_x
            dy = irow - imgMaxPix_y
            dis[irow, icol] = np.hypot(dx, dy)

    nn = im1size[1]*im1size[0]
    disX = dis.reshape(nn)
    disXsortId = np.argsort(disX)

    imgX = img.reshape(nn)
    imgY = imgX[disXsortId]
    psfFrac = np.cumsum(imgY)/np.sum(imgY)
    ind = np.where(psfFrac > fraction)[0][0]

    REE80 = np.rad2deg(dis[np.where(img == imgY[ind])]*psfSampleSizeInMicrons*1e-6/focalLengthInMeters)*3600
    return REE80



def check_galsimConvolve(path=None, plotImage=True):
    #load psf data
    data=scipy.io.loadmat(path)
    imPSF = data['psf']
    pixSize = np.rad2deg(5.*1e-6/28)*3600
    imPSF = imPSF/np.sum(imPSF)

    #psf -> galsimInterpolatedImage
    img = galsim.ImageF(imPSF, scale=pixSize)
    imgt= galsim.InterpolatedImage(img)

    #imPSFt = imgt.drawImage(nx=256, ny=256, scale=pixSize, method='no_pixel')
    imPSFt = imgt.drawImage(nx=256, ny=256, scale=pixSize)

    ree80 = psfEncircle(imPSF, fraction=0.8, psfSampleSizeInMicrons=5.)
    ree80_pix = ree80/(np.rad2deg((5.*1e-6/28))*3600)

    sliceX = slice(128-int(np.round(ree80_pix[0])), 128+int(np.round(ree80_pix[0]))+1, 1)

    #set a point sorce
    src    = galsim.DeltaFunction(flux=1.0)

    result = galsim.Convolve(src, imgt)

    #drawImage with same pixSize
    #tmp    = result.drawImage(nx=256, ny=256, scale=pixSize, method='no_pixel')
    tmp    = result.drawImage(nx=256, ny=256, scale=pixSize)


    if plotImage != True:
        res = (imPSFt.array - imPSF)/imPSF
        d0 = np.mean(res[sliceX, sliceX].flatten())

        res = (tmp.array - imPSFt.array)/imPSFt.array
        d1 = np.mean(res[sliceX, sliceX].flatten())
        return d0, d1


    #plot images
    fig = plt.figure(figsize=(22, 5))
    ax=plt.subplot(1,3,1)
    plt.imshow(imPSF[128-10:128+10, 128-10:128+10])
    plt.annotate("ORG", [0.1, 0.9],xycoords="axes fraction", fontsize=16, color="w")
    plt.colorbar()

    ax=plt.subplot(1,3,2)
    plt.imshow(imPSFt.array[128-10:128+10, 128-10:128+10])
    plt.annotate("InterpolatedImage", [0.1, 0.9],xycoords="axes fraction", fontsize=16, color="w")
    plt.colorbar()

    ax=plt.subplot(1,3,3)
    plt.imshow(tmp.array[128-10:128+10, 128-10:128+10])
    plt.annotate("ConvolvedImage", [0.1, 0.9],xycoords="axes fraction", fontsize=16, color="w")
    plt.colorbar()
    plt.savefig(OUTPUTPATH+'/fig_check_galsimConvolve_1.pdf')

    fig = plt.figure(figsize=(13, 10))
    ax=plt.subplot(2,2,1)
    res = (imPSFt.array - imPSF)/imPSF
    plt.imshow(res[128-10:128+10, 128-10:128+10])
    plt.annotate("$\Delta_1$", [0.1, 0.9],xycoords="axes fraction", fontsize=16, color="w")
    plt.colorbar()

    ax=plt.subplot(2,2,2)
    plt.hist(res[sliceX, sliceX].flatten(), alpha=0.75, bins=4)
    #plt.annotate("$\Delta_1^{\\rm REE80}$", [0.1, 0.9],xycoords="axes fraction", fontsize=16, color="k")
    plt.xlabel("$\Delta_1^{\\rm REE80}$", fontsize=16)
    plt.ylabel("PDF", fontsize=16)

    ax=plt.subplot(2,2,3)
    res = (tmp.array - imPSFt.array)/imPSFt.array
    plt.imshow(res[128-10:128+10, 128-10:128+10])
    plt.annotate("$\Delta_2$", [0.1, 0.9],xycoords="axes fraction", fontsize=16, color="w")
    plt.colorbar()

    ax=plt.subplot(2,2,4)
    plt.hist(res[sliceX, sliceX].flatten(), alpha=0.75, bins=4)
    #plt.annotate("$\Delta_2^{\\rm REE80}$", [0.1, 0.9],xycoords="axes fraction", fontsize=16, color="k")
    plt.xlabel("$\Delta_2^{\\rm REE80}$", fontsize=16)
    plt.ylabel("PDF", fontsize=16)

    plt.savefig(OUTPUTPATH+'/fig_check_galsimConvolve_2.pdf')


def check_galsimConvolveALL(dataPath):
    d0 = np.zeros(900)
    d1 = np.zeros(900)
    for ipsf in range(1,901,1):
        print("ipsf={:}".format(ipsf), end="\r")
        psfPath = dataPath+"/ccd1-w1/psf_{:}_centroidWgt_BC.mat".format(ipsf)
        t0,t1=check_galsimConvolve(path = psfPath, plotImage=False)
        d0[ipsf-1] = t0
        d1[ipsf-1] = t1
      
    fig = plt.figure(figsize=(12,6))
    ax = plt.subplot(1,2,1)
    #plt.scatter(np.linspace(1,900,900), d0)
    plt.hist(d0, bins=8, alpha=0.75)
    plt.xlabel("mean($\Delta_1^{\\rm REE80}$)", fontsize=16)
    plt.ylabel("PDF", fontsize=16)

    ax = plt.subplot(1,2,2)
    #plt.scatter(np.linspace(1,900,900), d1)
    plt.hist(d1, bins=8, alpha=0.75)
    plt.xlabel("mean($\Delta_2^{\\rm REE80}$)", fontsize=16)
    plt.ylabel("PDF", fontsize=16)

    plt.savefig(OUTPUTPATH+'/fig_check_galsimConvolveALL.pdf')
 

class testConvolve(unittest.TestCase):
    def __init__(self, methodName='runTest'):
        super(testConvolve,self).__init__(methodName)
        self.dataPath = os.path.join(os.getenv('UNIT_TEST_DATA_ROOT'), 'csst_fz_gc1')

        global OUTPUTPATH
        OUTPUTPATH = os.path.join(self.dataPath, 'outputs')

    def test_galsimConvolve(self):
        ipsf = 1
        psfPath = self.dataPath+"/ccd1-w1/psf_{:}_centroidWgt_BC.mat".format(ipsf)
        check_galsimConvolve(path = psfPath)

    def test_galsimConvolveALL(self):
        check_galsimConvolveALL(dataPath=self.dataPath)


if __name__ == "__main__":
    unittest.main()