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

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_ciomp_30X30'
#npsf = 100
#psfPath = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp'

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)
            #ipsfOutput = './psf_{:}_centroidWgt.mat'.format(ipsf)

            psfInfo = myConfig.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
            img, cx, cy = myUtil.centroidWgt(ipsfMat, nt=ncut)
            dcx =  cx - npix_x/2   #pixel coords -> global coords
            dcy =-(cy - npix_y/2)  #pixel coords -> global coords
            dcx*= 5e-3  #pixels -> mm
            dcy*= 5e-3  #pixels -> mm
            
            nn = 256
            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'])