"""
plot PSF场的分布：iccd-1-30
"""
import sys,os,math
import time
from IPython import display
import numpy as np
import matplotlib.pyplot as plt
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 assignTasks(npsf, NTasks, ThisTask):
    npsfPerTasks = int(npsf/NTasks)
    iStart= 0 + npsfPerTasks*ThisTask
    iEnd  = npsfPerTasks + npsfPerTasks*ThisTask
    if ThisTask == NTasks:
        iEnd = npsf
    return iStart, iEnd


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

iccd = 1  #[1, 30]
iwave= 1  #[1, 4]
ipsf = 1  #[1, 100]
psfPathA = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp'
psfPathB = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp_30X30'

for iccd in range(1,31):
  npsfA = 100
  npsfB = 900
  iStartA, iEndA = assignTasks(npsfA, NTasks, ThisTask)
  iStartB, iEndB = assignTasks(npsfB, NTasks, ThisTask)

  imxA = np.zeros(npsfA)
  imyA = np.zeros(npsfA)
  imxB = np.zeros(npsfB)
  imyB = np.zeros(npsfB)

  psfA_e1 = np.zeros(npsfA)
  psfA_e2 = np.zeros(npsfA)
  psfB_e1 = np.zeros(npsfB)
  psfB_e2 = np.zeros(npsfB)

  for ipsf in range(iStartA+1, iEndA+1):
    print('ipsf-{:}'.format(ipsf), end='\r')
    psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPathA, InputMaxPixelPos=False, PSFCentroidWgt=True)
    imxA[ipsf-1] = psfInfo['image_x']+psfInfo['centroid_x']
    imyA[ipsf-1] = psfInfo['image_y']+psfInfo['centroid_y']
    cenX, cenY, sz, e1, e2, REE80 = myUtil.psfSizeCalculator(psfInfo['psfMat'], CalcPSFcenter=True, SigRange=True, TailorScheme=2)
    psfA_e1[ipsf-1] = e1
    psfA_e2[ipsf-1] = e2

  for ipsf in range(iStartB+1, iEndB+1):
    print('ipsf-{:}'.format(ipsf), end='\r')
    psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPathB, InputMaxPixelPos=True, PSFCentroidWgt=True)
    imxB[ipsf-1] = psfInfo['image_x']+psfInfo['centroid_x']
    imyB[ipsf-1] = psfInfo['image_y']+psfInfo['centroid_y']
    cenX, cenY, sz, e1, e2, REE80 = myUtil.psfSizeCalculator(psfInfo['psfMat'], CalcPSFcenter=True, SigRange=True, TailorScheme=2)
    psfB_e1[ipsf-1] = e1
    psfB_e2[ipsf-1] = e2


  #######
  comm.barrier()
  imxA = comm.allreduce(imxA, op=MPI.SUM)
  imyA = comm.allreduce(imyA, op=MPI.SUM)
  imxB = comm.allreduce(imxB, op=MPI.SUM)
  imyB = comm.allreduce(imyB, op=MPI.SUM)

  psfA_e1 = comm.allreduce(psfA_e1, op=MPI.SUM)
  psfA_e2 = comm.allreduce(psfA_e2, op=MPI.SUM)
  psfB_e1 = comm.allreduce(psfB_e1, op=MPI.SUM)
  psfB_e2 = comm.allreduce(psfB_e2, op=MPI.SUM)

  #######
  if ThisTask == 0:
    plt.close()
    fig = plt.figure(figsize=(12,12))
    plt.plot(imxA, imyA, 'b.')
    plt.plot(imxB, imyB, 'r.')
    plt.savefig('figs/psfCrossPos.pdf')

    #######
    fig = plt.figure(figsize=(12, 12))
    plt.subplots_adjust(wspace=0.1, hspace=0.1)
    ax = plt.subplot(1, 1, 1)
    for ipsf in range(npsfA):
        plt.plot(imxA[ipsf], imyA[ipsf], 'b.')

        ang = np.arctan2(psfA_e2[ipsf], psfA_e1[ipsf])/2
        ell = np.sqrt(psfA_e1[ipsf]**2 + psfA_e2[ipsf]**2)
        ell *= 50
        lcos = ell*np.cos(ang)
        lsin = ell*np.sin(ang)
        plt.plot([imxA[ipsf]-lcos, imxA[ipsf]+lcos],[imyA[ipsf]-lsin, imyA[ipsf]+lsin],'b', lw=2)
        ###########
    for ipsf in range(npsfB):
        plt.plot(imxB[ipsf], imyB[ipsf], 'r.', alpha=0.75)

        ang = np.arctan2(psfB_e2[ipsf], psfB_e1[ipsf])/2
        ell = np.sqrt(psfB_e1[ipsf]**2 + psfB_e2[ipsf]**2)
        ell *= 50
        lcos = ell*np.cos(ang)
        lsin = ell*np.sin(ang)
        plt.plot([imxB[ipsf]-lcos, imxB[ipsf]+lcos],[imyB[ipsf]-lsin, imyB[ipsf]+lsin],'r', lw=2, alpha=0.75)
        ###########
    plt.gca().set_aspect(1)
    #plt.savefig('figs/psfCrossEll_iccd{:}.pdf'.format(iccd))


