import unittest

import sys,os,math
from itertools import islice

import mpi4py.MPI as MPI

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.use('Agg')

import scipy.io
from scipy import ndimage

#sys.path.append("/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF_v4_20210326")
# sys.path.append("../")
from ObservationSim.PSF.PSFInterp import PSFConfig as myConfig
from ObservationSim.PSF.PSFInterp import PSFUtil as myUtil

NPSF = 400
##############################

###计算PSF椭率###
def psfSecondMoments(psfMat, cenX, cenY, pixSize=1):
    apr = 0.5 #arcsec, 0.5角秒内测量
    fl  = 28. #meters
    pxs = 2.5 #microns
    apr = np.deg2rad(apr/3600.)*fl*1e6
    apr = apr/pxs
    apr = np.int(np.ceil(apr))

    I = psfMat
    ncol = I.shape[1]
    nrow = I.shape[0]
    w   = 0.0
    w11 = 0.0
    w12 = 0.0
    w22 = 0.0
    for icol in range(ncol):
        for jrow in range(nrow):
            x = icol*pixSize - cenX
            y = jrow*pixSize - cenY
            rr = np.sqrt(x*x + y*y)
            wgt= 0.0
            if rr <= apr:
                wgt = 1.0
            w   += I[jrow, icol]*wgt
            w11 += x*x*I[jrow, icol]*wgt
            w12 += x*y*I[jrow, icol]*wgt
            w22 += y*y*I[jrow, icol]*wgt
    w11 /= w
    w12 /= w
    w22 /= w
    sz = w11 + w22
    e1 = (w11 - w22)/sz
    e2 = 2.0*w12/sz

    return sz, e1, e2

##############################
##############################
'''
def assignTasks(npsf, NTasks, ThisTask):
    npsfPerTasks = int(npsf/NTasks)
    iStart= 0 + npsfPerTasks*ThisTask
    iEnd  = npsfPerTasks + npsfPerTasks*ThisTask
    if ThisTask == NTasks:
        iEnd = npsf
    return iStart, iEnd
'''

#def test_psfIDW(iccd, iwave, psfPath, ThisTask, NTasks):
def test_psfIDW(iccd, iwave, psfPath):
    nccd = 30
    npsfA = 900
    npsfB = 400

    #iStart, iEnd = assignTasks(400, NTasks, ThisTask)

    psfPathA = psfPath+'_30x30field'
    psfPathB = psfPath+'_20x20field'

    imxA = np.zeros(npsfA)
    imyA = np.zeros(npsfA)
    psfA = np.zeros([npsfA, 512, 512])

    imxB = np.zeros(npsfB)
    imyB = np.zeros(npsfB)
    psfB = np.zeros([npsfB, 512, 512])

    #for ipsf in range(iStart+1, iEnd+1):
    for ipsf in range(1, npsfA+1):
        print('ipsfA:', ipsf, end='\r', flush=True)
        psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPathA, InputMaxPixelPos=True, PSFCentroidWgt=True, VPSF=False)
        imxA[ipsf-1] = psfInfo['image_x'] + psfInfo['centroid_x']
        imyA[ipsf-1] = psfInfo['image_y'] + psfInfo['centroid_y']
        psfA[ipsf-1, :, :] = psfInfo['psfMat']

    for ipsf in range(1, npsfB+1):
        print('ipsfB:', ipsf, end='\r', flush=True)
        psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPathB, InputMaxPixelPos=True, PSFCentroidWgt=True, VPSF=False)
        imxB[ipsf-1] = psfInfo['image_x'] + psfInfo['centroid_x']
        imyB[ipsf-1] = psfInfo['image_y'] + psfInfo['centroid_y']
        psfB[ipsf-1, :, :] = psfInfo['psfMat']

    #myConfig.psfMaker_IDW(px, py, PSFMat, cen_col, cen_row, IDWindex=2, OnlyNeighbors=True, hoc=None, hoclist=None, PSFCentroidWgt=False)
    for ipsf in range(npsfB):
        print('ipsf:', ipsf, end='\r', flush=True)
        px = imxB[ipsf]
        py = imyB[ipsf]
        cen_col = imxA
        cen_row = imyA
        PSFMat  = psfA
        psfIDW  = myConfig.psfMaker_IDW(px, py, PSFMat, cen_col, cen_row, IDWindex=2, OnlyNeighbors=True, hoc=None, hoclist=None, PSFCentroidWgt=False)
        np.save('figs/psfIDW/psfIDW_{:}_{:}_{:}'.format(iccd, iwave, ipsf+1), psfIDW)


#def test_psfEll(iccd, iwave, ThisTask, NTasks):
def test_psfEll(iccd, iwave):
    nccd = 30
    npsf = 400

    #iStart, iEnd = assignTasks(npsf, NTasks, ThisTask)

    #imx = np.zeros(npsf)
    #imy = np.zeros(npsf)
    psf_e1 = np.zeros(npsf)
    psf_e2 = np.zeros(npsf)
    psf_sz = np.zeros(npsf)

    #for ipsf in range(iStart+1, iEnd+1):
    for ipsf in range(1, 401):
        if ipsf > 1:
            continue
        print('ipsf-{:}'.format(ipsf), end='\r')
        #psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=True, PSFCentroidWgt=True, VPSF=False)
        #psfMat = psfInfo['psfMat']
        #imx[ipsf-1] = psfInfo['image_x']+psfInfo['centroid_x']
        #imy[ipsf-1] = psfInfo['image_y']+psfInfo['centroid_y']
        psfMat = np.load('figs/psfIDW/psfIDW_{:}_{:}_{:}.npy'.format(iccd, iwave, ipsf)) #psfInfo['psfMat']
        cenX = 256
        cenY = 256
        sz, e1, e2  = psfSecondMoments(psfMat, cenX, cenY, pixSize=1)
        psf_e1[ipsf-1] = e1
        psf_e2[ipsf-1] = e2
        psf_sz[ipsf-1] = sz

    #######
    #comm.barrier()
    #imx = comm.allreduce(imx, op=MPI.SUM)
    #imy = comm.allreduce(imy, op=MPI.SUM)

    #psf_e1 = comm.allreduce(psf_e1, op=MPI.SUM)
    #psf_e2 = comm.allreduce(psf_e2, op=MPI.SUM)
    #psf_sz = comm.allreduce(psf_sz, op=MPI.SUM)

    #comm.barrier()
    #if ThisTask == 0:
    #    arr = [psf_e1, psf_e2, psf_sz]
    #    np.save('/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF_v4_20210326/test4report/data/psfEll20IDW_{:}_{:}'.format(iccd, iwave), arr)

'''
def test_psfResidualCalc(iccd, iwave, ipsf, psfPath):
    psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPath+'_20x20field', InputMaxPixelPos=True, PSFCentroidWgt=True, VPSF=False)
    psfMatORG = psfInfo['psfMat']
    psfMatIDW = np.load('/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF_v4_20210326/test4report/figs/psfIDW/psfIDW_{:}_{:}_{:}.npy'.format(iccd, iwave, ipsf))

    _,tREE80 = myUtil.psfEncircle(psfMatORG, fraction=0.8, psfSampleSizeInMicrons=2.5, focalLengthInMeters=28, cenPix=None)
    tREE80_pix = np.int32(np.ceil(tREE80/(0.074/2/2))[0])
    #print(tREE80, np.ceil(tREE80/(0.074/2/2)), tREE80_pix)

    timg0 = psfMatORG[256-tREE80_pix:256+tREE80_pix, 256-tREE80_pix:256+tREE80_pix]
    timg1 = psfMatIDW[256-tREE80_pix:256+tREE80_pix, 256-tREE80_pix:256+tREE80_pix]
    #print("residual::", np.max((timg1-timg0)/timg0), np.min((timg1-timg0)/timg0), np.mean((timg1-timg0)/timg0))

    return np.mean((timg1-timg0)/timg0)
'''

def test_psfResidualPlot(iccd, iwave, ipsf, psfPath):
    psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPath+'_20x20field', InputMaxPixelPos=True, PSFCentroidWgt=True, VPSF=False)
    psfMatORG = psfInfo['psfMat']
    psfMatIDW = np.load('figs/psfIDW/psfIDW_{:}_{:}_{:}.npy'.format(iccd, iwave, ipsf))

    npix = psfMatORG.shape[0]
    pixCutEdge= int(npix/2-15)

    img0 = psfMatORG[pixCutEdge:npix-pixCutEdge, pixCutEdge:npix-pixCutEdge]
    img1 = psfMatIDW[pixCutEdge:npix-pixCutEdge, pixCutEdge:npix-pixCutEdge]
    imgX = (img1 - img0)/img0

    img0 = np.log10(img0)
    img1 = np.log10(img1)
    imgX = np.log10(np.abs(imgX))

    fig = plt.figure(figsize=(18,4))
    ax = plt.subplot(1,3,1)
    plt.imshow(img0, origin='lower', vmin=-7, vmax=-1.3)
    plt.plot([npix/2-pixCutEdge, npix/2-pixCutEdge],[0, (npix/2-pixCutEdge)*2-1],'w--')
    plt.plot([0, (npix/2-pixCutEdge)*2-1],[npix/2-pixCutEdge, npix/2-pixCutEdge],'w--')
    plt.annotate('ORG', [0,(npix/2-pixCutEdge)*2-5], c='w', size=15)
    cticks=[-7, -6, -5, -4, -3, -2, -1]
    cbar = plt.colorbar(ticks=cticks)
    cbar.ax.set_yticklabels(['$10^{-7}$', '$10^{-6}$', '$10^{-5}$','$10^{-4}$','$10^{-3}$','$10^{-2}$', '$10^{-1}$'])
    print(img0.min(), img0.max())
    

    ax = plt.subplot(1,3,2)
    plt.imshow(img1, origin='lower', vmin=-7, vmax=-1.3)
    plt.plot([npix/2-pixCutEdge, npix/2-pixCutEdge],[0, (npix/2-pixCutEdge)*2-1],'w--')
    plt.plot([0, (npix/2-pixCutEdge)*2-1],[npix/2-pixCutEdge, npix/2-pixCutEdge],'w--')
    plt.annotate('IDW', [0,(npix/2-pixCutEdge)*2-5], c='w', size=15)
    cticks=[-7, -6, -5, -4, -3, -2, -1]
    cbar = plt.colorbar(ticks=cticks)
    cbar.ax.set_yticklabels(['$10^{-7}$', '$10^{-6}$', '$10^{-5}$','$10^{-4}$','$10^{-3}$','$10^{-2}$', '$10^{-1}$'])
    print(img1.min(), img1.max())
    
        
    ax = plt.subplot(1,3,3)
    plt.imshow(imgX, origin='lower', vmin =-3, vmax =np.log10(3e-1))
    plt.plot([npix/2-pixCutEdge, npix/2-pixCutEdge],[0, (npix/2-pixCutEdge)*2-1],'w--')
    plt.plot([0, (npix/2-pixCutEdge)*2-1],[npix/2-pixCutEdge, npix/2-pixCutEdge],'w--')
    #plt.annotate('(IDW-ORG)/ORG', [0,(npix/2-pixCutEdge)*2-5], c='w', size=15)
    cticks=[-5, -4, -3, -2, -1]
    cbar = plt.colorbar(ticks=cticks)
    cbar.ax.set_yticklabels(['$10^{-5}$','$10^{-4}$','$10^{-3}$','$10^{-2}$', '$10^{-1}$'])

    print(np.max((psfMatORG-psfMatIDW)))
    plt.savefig('figs/psfResidual_iccd{:}.pdf'.format(iccd))


def test_psfEllPlot(OVERPLOT=False):
    #if ThisTask == 0:
    if True:
        prefix = 'psfEll20'
        iccd = 1
        iwave= 1
        data = np.load('data/'+prefix+'_1_1.npy')
        imx= data[0]
        imy= data[1]
        psf_e1 = data[2]
        psf_e2 = data[3]
        print(np.shape(imx))
        npsf = np.shape(imx)[0]

        plt.cla()
        plt.close("all")
        fig = plt.figure(figsize=(12,12))
        plt.plot(imx, imy, 'r.')
        plt.savefig('figs/psfPos.pdf')

        #######
        
        fig = plt.figure(figsize=(12, 12))
        plt.subplots_adjust(wspace=0.1, hspace=0.1)
        ax = plt.subplot(1, 1, 1)
        for ipsf in range(npsf):
            plt.plot(imx[ipsf], imy[ipsf], 'b.')

            ang = np.arctan2(psf_e2[ipsf], psf_e1[ipsf])/2
            ell = np.sqrt(psf_e1[ipsf]**2 + psf_e2[ipsf]**2)
            ell *= 15
            lcos = ell*np.cos(ang)
            lsin = ell*np.sin(ang)
            plt.plot([imx[ipsf]-lcos, imx[ipsf]+lcos],[imy[ipsf]-lsin, imy[ipsf]+lsin],'b', lw=2)
            ###########
        ang = 0.
        ell = 0.05
        ell*= 15
        lcos = ell*np.cos(ang)
        lsin = ell*np.sin(ang)
        #plt.plot([imx[898]-lcos, imx[898]+lcos],[imy[898]+5.-lsin, imy[898]+5.+lsin],'k', lw=2)
        #plt.annotate('{:}'.format(ell/15), (imx[898]-2., imy[898]+6.), xycoords='data', fontsize=10)
        plt.xlabel('CCD X (mm)')
        plt.ylabel('CCD Y (mm)')

        if OVERPLOT == True:
            prefix = 'psfEll20IDW'
            data = np.load('data/'+prefix+'_1_1.npy')
            #imx= data[0]
            #imy= data[1]
            psf_e1 = data[0]
            psf_e2 = data[1]

            npsf = np.shape(imx)[0]
            for ipsf in range(npsf):
                #plt.plot(imx[ipsf], imy[ipsf], 'r.')

                ang = np.arctan2(psf_e2[ipsf], psf_e1[ipsf])/2
                ell = np.sqrt(psf_e1[ipsf]**2 + psf_e2[ipsf]**2)
                ell *= 15
                lcos = ell*np.cos(ang)
                lsin = ell*np.sin(ang)
                plt.plot([imx[ipsf]-lcos, imx[ipsf]+lcos],[imy[ipsf]-lsin, imy[ipsf]+lsin],'r', lw=1)

        plt.gca().set_aspect(1)
        if OVERPLOT == True:
            prefix = 'psfEllOPIDW'
        plt.savefig('figs/'+prefix+'_iccd{:}.pdf'.format(iccd))
        
'''
def test_psfdEllPlot():
    if ThisTask == 0:
        prefix = 'psfEll20'
        iccd = 1
        iwave= 1
        data = np.load('data/'+prefix+'_1_1.npy')
        imx= data[0]
        imy= data[1]
        psf_e1 = data[2]
        psf_e2 = data[3]
        psf_sz = data[4]
        print(np.shape(imx))
        npsf = np.shape(imx)[0]

        ellX = np.sqrt(psf_e1**2 + psf_e2**2)
        angX = np.arctan2(psf_e2, psf_e1)/2
        angX = np.rad2deg(angX)
        szX  = psf_sz

        ##############################
        prefix = 'psfEll20IDW'
        data = np.load('data/'+prefix+'_1_1.npy')
        #imx= data[0]
        #imy= data[1]
        psf_e1 = data[0]
        psf_e2 = data[1]
        psf_sz = data[2]
        
        ellY = np.sqrt(psf_e1**2 + psf_e2**2)
        angY = np.arctan2(psf_e2, psf_e1)/2
        angY = np.rad2deg(angY)
        szY  = psf_sz

        ##############################
        fig=plt.figure(figsize=(15, 4))
        ax = plt.subplot(1,3,1)
        plt.hist(ellX, bins=20, color='b', alpha=0.5)
        plt.hist(ellY, bins=20, color='r', alpha=0.5)
        plt.xlabel('$\epsilon$')
        plt.ylabel('PDF')

        ax = plt.subplot(1,3,2)
        plt.hist((ellY-ellX)/ellX, bins=20, color='r', alpha=0.5)
        plt.xlabel('$(\epsilon_{\\rm IDW}-\epsilon_{\\rm ORG})/\epsilon_{\\rm ORG}$')
        plt.ylabel('PDF')


        ax = plt.subplot(1,3,3)
        plt.hist((angY-angX)/angX, bins=20, color='r', alpha=0.5, range=[-0.1, 0.1])
        plt.xlabel('$(\\alpha_{\\rm IDW}-\\alpha_{\\rm ORG})/\\alpha_{\\rm ORG}$')
        plt.ylabel('PDF')
        plt.savefig('figs/psfEllOPIDWPDF.pdf')

        fig=plt.figure(figsize=(4, 4))
        plt.hist((szY-szX)/szX, bins=20, color='r', alpha=0.5)
        plt.xlabel('$(R_{\\rm IDW}-R_{\\rm ORG})/R_{\\rm ORG}$')
        plt.ylabel('PDF')
        plt.savefig('figs/psfEllOPIDWPDF_dsz.pdf')
'''

def test_psfdEllabsPlot(iccd):
    #if ThisTask == 0:
    if True:
        prefix = 'psfEll20'
        #iccd = 1
        #iwave= 1
        data = np.load('data/'+prefix+'_{:}_1.npy'.format(iccd))
        imx= data[0]
        imy= data[1]
        psf_e1 = data[2]
        psf_e2 = data[3]
        psf_sz = data[4]
        print(np.shape(imx))
        npsf = np.shape(imx)[0]

        ellX = np.sqrt(psf_e1**2 + psf_e2**2)
        angX = np.arctan2(psf_e2, psf_e1)/2
        angX = np.rad2deg(angX)
        szX  = psf_sz

        ##############################
        prefix = 'psfEll20IDW'
        data = np.load('data/'+prefix+'_{:}_1.npy'.format(iccd))
        #imx= data[0]
        #imy= data[1]
        psf_e1 = data[0]
        psf_e2 = data[1]
        psf_sz = data[2]
        
        ellY = np.sqrt(psf_e1**2 + psf_e2**2)
        angY = np.arctan2(psf_e2, psf_e1)/2
        angY = np.rad2deg(angY)
        szY  = psf_sz

        ##############################
        fig=plt.figure(figsize=(6, 5))
        grid = plt.GridSpec(3,1,left=0.15, right=0.95, wspace=None, hspace=0.02)
        #plt.subplots_adjust(left=None,bottom=None,right=None,top=None,wspace=None,hspace=0.02)
        ax = plt.subplot(grid[0:2,0])
        plt.plot([0.01,0.1],[0.01,0.1], 'k--', lw=1. )
        plt.scatter(ellX, ellY, color='b', alpha=1., s=3., edgecolors='None')
        plt.xlim([0.015, 0.085])
        plt.ylim([0.015, 0.085])
        plt.ylabel('$\epsilon_{\\rm IDW}$')
        plt.gca().axes.get_xaxis().set_visible(False)
        

        ax = plt.subplot(grid[2,0])
        plt.plot([0.015,0.085],[0.,0.], 'k--', lw=1. )
        plt.scatter(ellX, (ellY-ellX), color='b', s=3., edgecolors='None')
        plt.xlim([0.015, 0.085])
        plt.ylim([-0.0018, 0.0018])
        plt.xlabel('$\epsilon_{\\rm ORG}$')
        plt.ylabel('$\Delta$')
        plt.savefig('figs/psfEllOPIDWPDF_{:}.pdf'.format(iccd))

        fig=plt.figure(figsize=(4, 4))
        plt.hist((szY-szX)/szX, bins=20, color='r', alpha=0.5)
        plt.xlabel('$(R_{\\rm IDW}-R_{\\rm ORG})/R_{\\rm ORG}$')
        plt.ylabel('PDF')
        plt.savefig('figs/psfEllOPIDWPDF_dsz_{:}.pdf'.format(iccd))



class PSFMatsIDW_coverage(unittest.TestCase):
    def test_psfIDW_(self):
        #comm = MPI.COMM_WORLD
        #ThisTask = comm.Get_rank()
        #NTasks   = comm.Get_size()

        iccd = 1
        iwave= 1
        ipsf = 400
        psfPath = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_20210326/CSST_psf_ciomp'
        #test_psfIDW(iccd, iwave, psfPath, ThisTask, NTasks)
        test_psfIDW(iccd, iwave, psfPath)

        '''
        for iccd in range(7, 10): 
            res = np.zeros(400)
            for ipsf in range(1,401):
                print(ipsf, end="\r")
                res[ipsf-1] = test_psfResidualCalc(iccd, iwave, ipsf, psfPath)
            #fig = plt.figure(figsize=(6,6))
            #plt.hist(np.abs(res), bins=50)
            #plt.xlim([0,1])
            #plt.savefig('figs/psfResidualREE80PDF.pdf')
            print("{:}:".format(iccd), res[res<=0.01].size/400*100)
        ''' 

        test_psfResidualPlot(iccd, iwave, ipsf, psfPath)

        #test_psfEll(iccd, iwave, ThisTask, NTasks)
        test_psfEll(iccd, iwave)

        test_psfEllPlot(OVERPLOT=True)
        #test_psfdEllPlot()
    
        test_psfdEllabsPlot(iccd)



##############################
##############################
##############################
if __name__=='__main__':
    unittest.main()
