import unittest

import sys,os,math
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
from scipy import ndimage

#sys.path.append("/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF_v4_20210326")
sys.path.append("../")
from ObservationSim.PSF.PSFInterp import PSFConfig as myConfig
#import PSFUtil as myUtil

NPSF = 400
##############################

###计算PSF椭率###
def psfSecondMoments(psfMat, cenX, cenY, pixSize=1):
    apr = 0.5 #arcsec, 0.5角秒内测量
    fl  = 28. #meters
    pxs = 2.5 #microns
    apr = np.deg2rad(apr/3600.)*fl*1e6
    apr = apr/pxs
    apr = np.int(np.ceil(apr))

    I = psfMat
    ncol = I.shape[1]
    nrow = I.shape[0]
    w   = 0.0
    w11 = 0.0
    w12 = 0.0
    w22 = 0.0
    for icol in range(ncol):
        for jrow in range(nrow):
            x = icol*pixSize - cenX
            y = jrow*pixSize - cenY
            rr = np.sqrt(x*x + y*y)
            wgt= 0.0
            if rr <= apr:
                wgt = 1.0
            w   += I[jrow, icol]*wgt
            w11 += x*x*I[jrow, icol]*wgt
            w12 += x*y*I[jrow, icol]*wgt
            w22 += y*y*I[jrow, icol]*wgt
    w11 /= w
    w12 /= w
    w22 /= w
    sz = w11 + w22
    e1 = (w11 - w22)/sz
    e2 = 2.0*w12/sz

    return sz, e1, e2

##############################
##############################
'''
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
'''

#def test_psfEll(iccd, iwave, psfPath, ThisTask, NTasks):
def test_psfEll(iccd, iwave, psfPath):
    nccd = 30
    npsf = NPSF

    #iStart, iEnd = assignTasks(npsf, NTasks, ThisTask)

    imx = np.zeros(npsf)
    imy = np.zeros(npsf)
    psf_e1 = np.zeros(npsf)
    psf_e2 = np.zeros(npsf)
    psf_sz = np.zeros(npsf)

    #for ipsf in range(iStart+1, iEnd+1):
    for ipsf in range(1, 401):
        if ipsf != 1:
            continue
        print('ipsf-{:}'.format(ipsf), end='\r')
        psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=True, PSFCentroidWgt=True, VPSF=False)
        imx[ipsf-1] = psfInfo['image_x']+psfInfo['centroid_x']
        imy[ipsf-1] = psfInfo['image_y']+psfInfo['centroid_y']
        psfMat = psfInfo['psfMat']
        cenX = 256
        cenY = 256
        sz, e1, e2  = psfSecondMoments(psfMat, cenX, cenY, pixSize=1)
        psf_e1[ipsf-1] = e1
        psf_e2[ipsf-1] = e2
        psf_sz[ipsf-1] = sz
        print('test:' ,sz, e1, e2)

    #######
    #comm.barrier()
    #imx = comm.allreduce(imx, op=MPI.SUM)
    #imy = comm.allreduce(imy, op=MPI.SUM)

    #psf_e1 = comm.allreduce(psf_e1, op=MPI.SUM)
    #psf_e2 = comm.allreduce(psf_e2, op=MPI.SUM)
    #psf_sz = comm.allreduce(psf_sz, op=MPI.SUM)

    #comm.barrier()
    #if ThisTask == 0:
    #    arr = [imx, imy, psf_e1, psf_e2, psf_sz]
    #    np.save('/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF_v4_20210326/test4report/data/psfEll20_{:}_{:}'.format(iccd, iwave), arr)


def test_psfEllPlot(OVERPLOT=False):
    #if ThisTask == 0:
    if True:
        prefix = 'psfEll30'
        iccd = 1
        iwave= 1
        data = np.load('data/'+prefix+'_1_1.npy')
        imx= data[0]
        imy= data[1]
        psf_e1 = data[2]
        psf_e2 = data[3]
        print(np.shape(imx))
        npsf = np.shape(imx)[0]

        plt.cla()
        plt.close("all")
        fig = plt.figure(figsize=(12,12))
        plt.plot(imx, imy, 'r.')
        plt.savefig('figs/psfPos.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(npsf):
            plt.plot(imx[ipsf], imy[ipsf], 'r.')

            ang = np.arctan2(psf_e2[ipsf], psf_e1[ipsf])/2
            ell = np.sqrt(psf_e1[ipsf]**2 + psf_e2[ipsf]**2)
            ell *= 15
            lcos = ell*np.cos(ang)
            lsin = ell*np.sin(ang)
            plt.plot([imx[ipsf]-lcos, imx[ipsf]+lcos],[imy[ipsf]-lsin, imy[ipsf]+lsin],'r', lw=2)
            ###########
        ang = 0.
        ell = 0.05
        ell*= 15
        lcos = ell*np.cos(ang)
        lsin = ell*np.sin(ang)
        plt.plot([imx[898]-lcos, imx[898]+lcos],[imy[898]+5.-lsin, imy[898]+5.+lsin],'k', lw=2)
        plt.annotate('{:}'.format(ell/15), (imx[898]-2., imy[898]+6.), xycoords='data', fontsize=10)
        plt.xlabel('CCD X (mm)')
        plt.ylabel('CCD Y (mm)')

        if OVERPLOT == True:
            prefix = 'psfEll20'
            data = np.load('data/'+prefix+'_1_1.npy')
            imx= data[0]
            imy= data[1]
            psf_e1 = data[2]
            psf_e2 = data[3]

            npsf = np.shape(imx)[0]
            for ipsf in range(npsf):
                plt.plot(imx[ipsf], imy[ipsf], 'b.')

                ang = np.arctan2(psf_e2[ipsf], psf_e1[ipsf])/2
                ell = np.sqrt(psf_e1[ipsf]**2 + psf_e2[ipsf]**2)
                ell *= 15
                lcos = ell*np.cos(ang)
                lsin = ell*np.sin(ang)
                plt.plot([imx[ipsf]-lcos, imx[ipsf]+lcos],[imy[ipsf]-lsin, imy[ipsf]+lsin],'b', lw=2)

        plt.gca().set_aspect(1)
        if OVERPLOT == True:
            prefix = 'psfEllOP'
        plt.savefig('figs/'+prefix+'_iccd{:}.pdf'.format(iccd))


class PSFMatsEll_coverage(unittest.TestCase):
    def test_psfEll_(self):
        #comm = MPI.COMM_WORLD
        #ThisTask = comm.Get_rank()
        #NTasks   = comm.Get_size()
        print('#####haha#####')

        iccd = 1
        iwave= 1
        psfPath = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_20210326/CSST_psf_ciomp_20x20field'
        #test_psfEll(iccd, iwave, psfPath, ThisTask, NTasks)
        test_psfEll(iccd, iwave, psfPath)

        test_psfEllPlot(OVERPLOT=True)


##############################
##############################
##############################
if __name__=='__main__':
    unittest.main()