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_20210108")
import PSFConfig as myConfig
import PSFUtil as myUtil

NPSF = 900
##############################
##############################
##############################
def test_psfPlot(iccd, ipsf, psfPath):
    plt.cla()
    plt.close('all')
    #psf质心位置检查,图像切片程序检查
    fig = plt.figure(figsize=(4,16))
    plt.subplots_adjust(wspace=0.1, hspace=0.1)

    psfSet = []
    psfSetX= []
    for iwave in range(1, 5):
        psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=True, PSFCentroidWgt=False)
        psfSet.append(psfInfo)

        psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=True, PSFCentroidWgt=True)
        psfSetX.append(psfInfo)

        ax = plt.subplot(4,1,iwave)
        psfMat = psfSet[iwave-1]['psfMat']
        pixsize = psfSet[iwave-1]['pixsize']*1e-3  #microns -> mm
        print('psfMat-shape:', psfMat.shape)

        npix = psfMat.shape[0]
        pixCutEdge= int(npix/2-15)
        plt.imshow((psfMat[pixCutEdge:npix-pixCutEdge, pixCutEdge:npix-pixCutEdge]), origin='lower')
        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--')

        imgX = psfSet[iwave-1]['image_x']  #in mm
        imgY = psfSet[iwave-1]['image_y']  #in mm
        #cenX = imgX +deltX
        #cenY = imgY -deltY

        deltX= psfSet[iwave-1]['centroid_x'] #in mm
        deltY= psfSet[iwave-1]['centroid_y'] #in mm
        cenPix_X = npix/2 + deltX/pixsize
        cenPix_Y = npix/2 + deltY/pixsize
        plt.plot([cenPix_X-pixCutEdge],[cenPix_Y-pixCutEdge], 'rx', ms = 20)

        deltX= psfSetX[iwave-1]['centroid_x'] #in mm
        deltY= psfSetX[iwave-1]['centroid_y'] #in mm
        cenPix_X = npix/2 + deltX/pixsize
        cenPix_Y = npix/2 + deltY/pixsize
        plt.plot([cenPix_X-pixCutEdge],[cenPix_Y-pixCutEdge], 'bx', ms = 10, mew=2)

        maxX = psfSet[iwave-1]['max_x']
        maxY = psfSet[iwave-1]['max_y']
        maxPix_X = npix/2 + maxX/pixsize
        maxPix_Y = npix/2 + maxY/pixsize
        plt.plot([maxPix_X-pixCutEdge],[maxPix_Y-pixCutEdge], 'r+', ms = 20)


        img = psfMat/np.sum(psfMat)
        y, x = ndimage.center_of_mass(img)
        print(x, y)
        print(cenPix_X, cenPix_Y)
        plt.plot([x-pixCutEdge],[y-pixCutEdge], 'wx', ms = 5, mew=1)
        plt.annotate('iccd={:}, iwave={:}, ipsf={:}'.format(iccd, iwave, ipsf), [0,(npix/2-pixCutEdge)*2-5], c='w', size=15)
    plt.savefig('figs/psf_{:}_{:}.pdf'.format(iccd, ipsf))

    print('psfSet has been loaded.')
    print('Usage: psfSet[i][keys]')
    print('psfSet.keys:', psfSet[0].keys())
    

##############################
##############################
##############################
if __name__=='__main__':
    iccd = 2  #[1, 30]
    iwave= 2  #[1, 4]
    ipsf = 46  #[1, 100]
    psfPath = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_20210108/CSST_psf_ciomp_2p5um_cycle3_ccr90'

    iipsf = [1, 15, 30, 450, 465, 480, 870, 885, 900]
    for iccd in range(1, 31):
        for ipsf in iipsf:
            test_psfPlot(iccd, ipsf, psfPath)