'''
PSF-IDW插值的cross-check
'''

import sys
import numpy as np
import matplotlib.pyplot as plt
import mpi4py.MPI as MPI

sys.path.append("/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF")
import PSFConfig as myConfig
import PSFUtil as myUtil

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_Check(iccd, iwave, psfPathA, psfPathB, NTasks, ThisTask):
    npsfA   = 100
    npsfB   = 900

    psfSetA, PSFMatA, cen_colA, cen_rowA = myConfig.LoadPSFset(iccd, iwave, npsfA, psfPathA, InputMaxPixelPos=False, PSFCentroidWgt=True)
    psfSetB, PSFMatB, cen_colB, cen_rowB = myConfig.LoadPSFset(iccd, iwave, npsfB, psfPathB, InputMaxPixelPos=True, PSFCentroidWgt=True)

    psf_sz = np.zeros(npsfA)
    psf_e1 = np.zeros(npsfA)
    psf_e2 = np.zeros(npsfA)
    psfMaker_sz = np.zeros(npsfA)
    psfMaker_e1 = np.zeros(npsfA)
    psfMaker_e2 = np.zeros(npsfA)

    iStart, iEnd= assignTasks(npsfA, NTasks, ThisTask)

    for ipsf in range(iStart, iEnd):
        #if ipsf != 83:
        #    continue

        if ThisTask == 0:
            print('ipsf:', ipsf, flush=True)
        px = cen_colA[ipsf]
        py = cen_rowA[ipsf]
        img= myConfig.psfMaker_IDW(px, py, PSFMatB, cen_colB, cen_rowB, IDWindex=2, OnlyNeighbors=True, PSFCentroidWgt=False)
        imgtt, cx, cy = myUtil.centroidWgt(img, nt=200)

        imx = psfSetA[ipsf]['psfMat']
        imy = imgtt
        cenX, cenY, sz, e1, e2, REE80 = myUtil.psfSizeCalculator(imx, CalcPSFcenter=True, SigRange=True, TailorScheme=2)
        psf_sz[ipsf] = sz
        psf_e1[ipsf] = e1
        psf_e2[ipsf] = e2
        cenX, cenY, sz, e1, e2, REE80 = myUtil.psfSizeCalculator(imy, CalcPSFcenter=True, SigRange=True, TailorScheme=2)
        psfMaker_sz[ipsf] = sz
        psfMaker_e1[ipsf] = e1
        psfMaker_e2[ipsf] = e2

    #######
    comm.barrier()
    psf_sz = comm.allreduce(psf_sz, op=MPI.SUM)
    psf_e1 = comm.allreduce(psf_e1, op=MPI.SUM)
    psf_e2 = comm.allreduce(psf_e2, op=MPI.SUM)
    psfMaker_sz = comm.allreduce(psfMaker_sz, op=MPI.SUM)
    psfMaker_e1 = comm.allreduce(psfMaker_e1, op=MPI.SUM)
    psfMaker_e2 = comm.allreduce(psfMaker_e2, op=MPI.SUM)

    #######
    if ThisTask == 0:
        ell_iccd = np.zeros(npsfA)
        ell_iccd_psfMaker = np.zeros(npsfA)
        fig = plt.figure(figsize=(18, 5))
        plt.subplots_adjust(wspace=0.1, hspace=0.1)
        ax = plt.subplot(1, 3, 1)
        for ipsf in range(npsfA):
            imx = cen_colA[ipsf] #psfSet[ipsf-1]['image_x']
            imy = cen_rowA[ipsf] #psfSet[ipsf-1]['image_y']
            plt.plot(imx, imy, 'b.')

            #ang = np.arctan2(psfSet[ipsf-1]['psf_e2'], psfSet[ipsf-1]['psf_e1'])/2
            #ell = np.sqrt(psfSet[ipsf-1]['psf_e1']**2 + psfSet[ipsf-1]['psf_e2']**2)
            ang = np.arctan2(psf_e2[ipsf], psf_e1[ipsf])/2
            ell = np.sqrt(psf_e1[ipsf]**2 + psf_e2[ipsf]**2)
            ell_iccd[ipsf] = ell
            ell *= 50
            lcos = ell*np.cos(ang)
            lsin = ell*np.sin(ang)
            plt.plot([imx-lcos, imx+lcos],[imy-lsin, imy+lsin],'b', lw=2)
            ###########
            ang = np.arctan2(psfMaker_e2[ipsf], psfMaker_e1[ipsf])/2
            ell = np.sqrt(psfMaker_e1[ipsf]**2 + psfMaker_e2[ipsf]**2)
            ell_iccd_psfMaker[ipsf] = ell
            ell *= 50
            lcos = ell*np.cos(ang)
            lsin = ell*np.sin(ang)
            plt.plot([imx-lcos, imx+lcos],[imy-lsin, imy+lsin],'r', lw=1)
            ###########
        plt.gca().set_aspect(1)

        ax = plt.subplot(1, 3, 2)
        tmp = plt.hist(ell_iccd, bins = 8, color='b', alpha=0.5)
        tmp = plt.hist(ell_iccd_psfMaker, bins = 8, color='r', alpha=0.5)
        plt.annotate('iccd-{:} iwave-{:}'.format(iccd, iwave), (0.55, 0.85), xycoords='axes fraction',fontsize=15)
        plt.xlabel('ell')
        plt.ylabel('PDF')

        ax = plt.subplot(1, 3, 3)
        dsz = (psfMaker_sz - psf_sz)/psf_sz
        dsz_hist= plt.hist(dsz)
        plt.xlabel('dsz')
        plt.savefig('figs/psfCrossMakerIDW_30t10_iccd{:}_iwave{:}.pdf'.format(iccd, iwave))



if __name__ == '__main__':
    #######
    comm = MPI.COMM_WORLD
    ThisTask = comm.Get_rank()
    NTasks   = comm.Get_size()
    #######

    iccd = 13  #[1, 30]
    iwave= 1  #[1, 4]
    ipsf = 10  #[1, 100]
    psfPathA = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp'
    psfPathB = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp_30X30'
    for iccd in range(1, 31):
        if iccd != 12:
            continue
        test_Check(iccd, iwave, psfPathA, psfPathB, NTasks, ThisTask)

    #test_Check(28, iwave, psfPathA, psfPathB, NTasks, ThisTask)
    #test_Check(29, iwave, psfPathA, psfPathB, NTasks, ThisTask)




