import os, sys
import numpy as np
import scipy.io
import mpi4py.MPI as MPI



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

def mkdir(path):
    isExists = os.path.exists(path)
    if not isExists:
        os.mkdir(path)
    

############################################
comm = MPI.COMM_WORLD
ThisTask = comm.Get_rank()
NTasks   = comm.Get_size()

npsf = 900
psfPath = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_20210108/CSST_psf_ciomp_2p5um_cycle3_ccr90'

npsfPerTasks = int(npsf/NTasks)
iStart= 0 + npsfPerTasks*ThisTask
iEnd  = npsfPerTasks + npsfPerTasks*ThisTask
if ThisTask == NTasks:
    iEnd = npsf


for iccd in range(1, 31):
    iccdPath = psfPath + '_proc/ccd{:}'.format(iccd)
    if ThisTask == 0:
        mkdir(iccdPath)
    comm.barrier()

    for iwave in range(1, 5):
        iwavePath = iccdPath + '/wave_{:}'.format(iwave)
        if ThisTask == 0:
            mkdir(iwavePath)
        comm.barrier()

        psfMatPath = iwavePath + '/5_psf_array'
        if ThisTask == 0:
            mkdir(psfMatPath)
        comm.barrier()

        for ii in range(iStart, iEnd):
            ipsf = ii+1
            if ThisTask ==0:
                print('iccd-iwave-ipsf: {:4}{:4}{:4}'.format(iccd, iwave, ipsf), end='\r',flush=True)
            #if iccd != 1 or iwave !=1 or ipsf != 1:
            #    continue

            ipsfOutput = psfMatPath + '/psf_{:}_centroidWgt.mat'.format(ipsf)
            psfInfo = PSFConfig.LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=True)
            ipsfMat  = psfInfo['psfMat']
            npix_y, npix_x = ipsfMat.shape

            #if npsf == 100:
            #    ncut = 160
            #if npsf == 900:
            #    ncut = 200
            ncut = int(npix_y*0.90)
            ncut = ncut + ncut%2  #even pixs

            img, cx, cy = PSFUtil.centroidWgt(ipsfMat, nt=ncut)
            dcx = cx - npix_x/2   #pixel coords -> global coords
            dcy = cy - npix_y/2  #pixel coords -> global coords
            dcx*= psfInfo['pixsize']*1e-3 #5e-3  #pixels -> mm
            dcy*= psfInfo['pixsize']*1e-3 #5e-3  #pixels -> mm
            
            nn = npix_y
            dn = int((nn - ncut)/2)
            imgt = np.zeros([nn, nn], dtype=np.float32)
            imgt[dn:-dn, dn:-dn] = img

            scipy.io.savemat(ipsfOutput, {'cx':dcx, 'cy':dcy, 'psf':imgt})
            if iccd != 1 or iwave !=1 or ipsf != 1:
                if ThisTask == 0:
                    print('CHECK::', dcx, dcy, psfInfo['centroid_x'],psfInfo['centroid_y'])