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, iwave, ipsf, psfPath, r50):
    #psf质心位置检查,图像切片程序检查
    plt.cla()
    plt.close("all")

    fig = plt.figure(figsize=(4,8))
    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)

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

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

    npix = psfMat.shape[0]
    pixCutEdge= int(npix/2-15)
    plt.imshow(np.log10(1e-6+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 = psfInfo['image_x']  #in mm
    imgY = psfInfo['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-1 + deltY/pixsize
    #plt.plot([cenPix_X-pixCutEdge],[cenPix_Y-pixCutEdge], 'rx', ms = 20)

    deltX= psfInfoX['centroid_x'] #in mm
    deltY= psfInfoX['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-1 + 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)

    ax = plt.subplot(2,1,2)
    cenPix = [cenPix_X, cenPix_Y]    
    img = myUtil.psfTailor(psfMat, apSizeInArcsec=2*r50, psfSampleSizeInMicrons=2.5, cenPix=cenPix)
    plt.imshow(np.log10(1e-6+img[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--')


    plt.savefig('figs/psf_{:}_{:}_2r50_log.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'

    REE50 = np.loadtxt('REE50_w{:}.txt'.format(iwave))
    iipsf = [1, 15, 30, 450, 465, 480, 870, 885, 900]
    for iccd in range(1, 31):
        r50   = REE50[iccd-1]
        print(r50)
        for ipsf in iipsf:
            test_psfPlot(iccd, iwave, ipsf, psfPath, r50)