'''
centroidWgt: 测试stack后质心与波长平均质心位置的相对偏差
'''

import sys
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
#import xlrd
from scipy import ndimage

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


def test_psfCentroidWgt(iccd, iwave, ipsf, psfPath):
    psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=False)
    ipsfMat  = psfInfo['psfMat']
    print('centroid_x/y', psfInfo['centroid_x']/0.005, psfInfo['centroid_y']/0.005)
    img, cx, cy = myUtil.centroidWgt(ipsfMat, nt=160)
    return img, cx, cy


def test_wavePSFCentroidWgt(iccd, ipsf, psfPath):
    '''
    读取四个波长采样以及计算centroidWgt质心位置偏差
    '''
    psf4iwave = []
    for iwave in range(1, 5):
        psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=True)
        psf4iwave.append(psfInfo['psfMat'])
    tt = myUtil.psfStack(psf4iwave[0], psf4iwave[1], psf4iwave[2], psf4iwave[3])
 
    imgt = []
    cxt = []
    cyt = []
    for iwave in range(4):
        ipsfMat = psf4iwave[iwave]
        img, cx, cy = myUtil.centroidWgt(ipsfMat, nt=160)
        #print('iwave-{:}:'.format(iwave), cx, cy, flush=True)
        imgt.append(img)
        cxt.append(cx)
        cyt.append(cy)
    cxm = np.mean(np.array(cxt))
    cym = np.mean(np.array(cyt))

    img, cx, cy = myUtil.centroidWgt(tt, nt=160)
    imgt.append(img)
    return psf4iwave, imgt, cx - cxm, cy - cym, cxt, cyt


def test_wavePSFCentroidWgt_mpi(iccd, psfPath, ThisTask, NTasks):
    npsf = 900
    npsfPerTasks = int(npsf/NTasks)
    iStart= 0 + npsfPerTasks*ThisTask
    iEnd  = npsfPerTasks + npsfPerTasks*ThisTask
    if ThisTask == NTasks:
        iEnd = npsf
    #print(ThisTask, ':', iStart, iEnd)

    psf_dcx = np.zeros(npsf)
    psf_dcy = np.zeros(npsf)
    for ipsf in range(iStart, iEnd):
        if ThisTask == 0:
            print('ThisTask-ipsf:', ThisTask, ipsf, flush=True)
        psf4iwave, imgt, dcx, dcy, cxt, cyt = test_wavePSFCentroidWgt(iccd, ipsf+1, psfPath)
        psf_dcx[ipsf] = dcx
        psf_dcy[ipsf] = dcy
        '''
        if dcx>=1.0 or dcy>=1.0:
            fig = plt.figure(figsize=(12,18))
            for iax in range(5):
                ax = plt.subplot(3,2,iax+1)
                plt.imshow(imgt[iax][80-10:80+10, 80-10:80+10], origin='lower')
            plt.savefig('figs/psfCentroidWgt_imgt{:}.pdf'.format(ipsf))
            fig = plt.figure(figsize=(12,18))
            for iax in range(4):
                ax = plt.subplot(3,2,iax+1)
                plt.imshow(psf4iwave[iax][90-10:90+10, 90-10:90+10], origin='lower')
                plt.plot(cxt[iax]-90+10, cyt[iax]-90+10, 'b+', ms=20)
            ax = plt.subplot(3,2,5)
            tt = myUtil.psfStack(psf4iwave[0], psf4iwave[1], psf4iwave[2], psf4iwave[3])
            img, cx, cy = myUtil.centroidWgt(tt, nt=160)
            plt.imshow(tt[90-10:90+10, 90-10:90+10], origin='lower')
            plt.plot(cx-90+10, cy-90+10, 'b+', ms=20)
            plt.savefig('figs/psfCentroidWgt_psfMat{:}.pdf'.format(ipsf))
            print('ok?:', ipsf, dcx, dcy, cxt, cyt, cx, cy)
        '''

    comm.barrier()
    psf_dcx = comm.allreduce(psf_dcx, op=MPI.SUM)
    psf_dcy = comm.allreduce(psf_dcy, op=MPI.SUM)

    if ThisTask == 0:
        fig = plt.figure(figsize=(15,4))
        ax = plt.subplot(1,3,1)
        hist = plt.hist(psf_dcx)
        plt.xlabel('dcx=cx-cxm')
        plt.ylabel('PDF')
        plt.annotate('psfCentroid_dcx', (0.65, 0.85), xycoords='axes fraction')
    
        ax = plt.subplot(1,3,2)
        hist = plt.hist(psf_dcy)
        plt.xlabel('dcy=cy-cym')
        plt.ylabel('PDF')
        plt.annotate('psfCentroid_dcy', (0.65, 0.85), xycoords='axes fraction')

        ax = plt.subplot(1,3,3)
        hist = plt.hist(np.sqrt(psf_dcx*psf_dcx + psf_dcy*psf_dcy))
        plt.xlabel('|dc|')
        plt.ylabel('PDF')
        plt.annotate('psfCentroid_|dc|', (0.65, 0.85), xycoords='axes fraction')

        plt.savefig('figs/psfCentroidWgt_ccd{:}.pdf'.format(iccd))
        plt.close()






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]
    #psfPath = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp'
    psfPath = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp_30X30'

    '''
    img, cx, cy = test_psfCentroidWgt(iccd, iwave, ipsf, psfPath)
    print(cx, cy, flush=True)
    plt.imshow(img[80-10:80+10, 80-10:80+10], origin='lower')
    plt.show()
    '''

    '''
    imgt = test_wavePSFCentroidWgt(iccd, ipsf, psfPath)
    fig = plt.figure(figsize=(12,18))
    for iax in range(5):
        ax = plt.subplot(3,2,iax+1)
        plt.imshow(imgt[iax][80-10:80+10, 80-10:80+10], origin='lower')
    plt.savefig('figs/psfCentroidWgt.pdf')
    '''

    for iccd in range(1, 31):
        if ThisTask ==0:
            print('iccd:', iccd, flush=True)
        test_wavePSFCentroidWgt_mpi(iccd, psfPath, ThisTask, NTasks)




