'''
画出给定PSF的插值前后image以及neighbors image
'''

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):
    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)

    anchor = 36
    for ipsf in range(npsfA):
        if ipsf != anchor:
            continue

        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)

        neigh = myUtil.findNeighbors(px, py, cen_colB, cen_rowB, dr=5., dn=4, OnlyDistance=False)
        print('neigh:', neigh)

        imx = psfSetA[ipsf]['psfMat']
        imy = img #imgtt
        cenX, cenY, sz, e1, e2, REE80 = myUtil.psfSizeCalculator(imx, CalcPSFcenter=True, SigRange=True, TailorScheme=2)
        #psf_e1[ipsf] = e1
        #psf_e2[ipsf] = e2
        ang = np.arctan2(e2, e1)/2
        ell = np.sqrt(e1**2 + e2**2)
        print(e1,e2, ell, ang)
        cenX, cenY, sz, e1, e2, REE80 = myUtil.psfSizeCalculator(imy, CalcPSFcenter=True, SigRange=True, TailorScheme=2)
        #psfMaker_e1[ipsf] = e1
        #psfMaker_e2[ipsf] = e2
        ang = np.arctan2(e2, e1)/2
        ell = np.sqrt(e1**2 + e2**2)
        print(e1,e2, ell, ang)

    #######
    fig = plt.figure(figsize=(10, 15))
    cpp = 128
    dcpp= 10
    ax = plt.subplot(3, 2, 1)
    plt.imshow(imx[cpp-dcpp:cpp+dcpp, cpp-dcpp:cpp+dcpp])

    ax = plt.subplot(3, 2, 2)
    plt.imshow(imy[cpp-dcpp:cpp+dcpp, cpp-dcpp:cpp+dcpp])

    ax = plt.subplot(3, 2, 3)
    imx= psfSetB[neigh[0]]['psfMat']
    plt.imshow(imx[cpp-dcpp:cpp+dcpp, cpp-dcpp:cpp+dcpp])
    
    ax = plt.subplot(3, 2, 4)
    imx= psfSetB[neigh[1]]['psfMat']
    plt.imshow(imx[cpp-dcpp:cpp+dcpp, cpp-dcpp:cpp+dcpp])

    ax = plt.subplot(3, 2, 5)
    imx= psfSetB[neigh[2]]['psfMat']
    plt.imshow(imx[cpp-dcpp:cpp+dcpp, cpp-dcpp:cpp+dcpp])

    ax = plt.subplot(3, 2, 6)
    imx= psfSetB[neigh[3]]['psfMat']
    plt.imshow(imx[cpp-dcpp:cpp+dcpp, cpp-dcpp:cpp+dcpp])

    '''
    ax = plt.subplot(3, 2, 5)
    imx = psfSetB[790]['psfMat']*0.7565569887550834 + psfSetB[789]['psfMat']*0.24344301124491663
    imx = imx/np.nansum(imx)
    plt.imshow(imx[cpp-dcpp:cpp+dcpp, cpp-dcpp:cpp+dcpp])
    cenX, cenY, sz, e1, e2, REE80 = myUtil.psfSizeCalculator(imx, CalcPSFcenter=True, SigRange=True, TailorScheme=2)
    ang = np.arctan2(e2, e1)/2
    ell = np.sqrt(e1**2 + e2**2)
    print(e1,e2, ell, ang)
    '''

    plt.savefig('figs/psfCrossMakerIDW_imCheck.pdf'.format(iccd, iwave))



if __name__ == '__main__':
    iccd = 29  #[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'
    
    test_Check(iccd, iwave, psfPathA, psfPathB)




