import sys
import ctypes
import numpy as np
from scipy import ndimage
import scipy.spatial as spatial


###定义PSF像素的全局坐标###
def psfPixelLayout(nrows, ncols, cenPosRow, cenPosCol, pixSizeInMicrons=5.0):
    """
    convert psf pixels to physical position

    Parameters:
        nrows, ncols (int, int): psf sampling with [nrows, ncols].
        cenPosRow, cenPosCol (float, float): A physical position of the chief ray for a given psf.
        pixSizeInMicrons (float-optional): The pixel size in microns from the psf sampling.

    Returns:
        psfPixelPos (numpy.array-float): [posx, posy] in mm for [irow, icol]

    Notes:
        1. show positions on ccd, but not position on image only [+/- dy]
    """
    psfPixelPos = np.zeros([2, nrows, ncols])
    if nrows % 2 != 0:
        sys.exit()
    if ncols % 2 != 0:
        sys.exit()

    cenPix_row = nrows/2 + 1 #中心主光线对应pixle [由长光定义]
    cenPix_col = ncols/2 + 1

    for irow in range(nrows):
        for icol in range(ncols):
            delta_row = ((irow + 1) - cenPix_row)*pixSizeInMicrons*1e-3
            delta_col = ((icol + 1) - cenPix_col)*pixSizeInMicrons*1e-3
            psfPixelPos[0, irow, icol] = cenPosCol + delta_col
            psfPixelPos[1, irow, icol] = cenPosRow - delta_row  #note-1:in CCD全局坐标

    return psfPixelPos


###查找最大pixel位置###
def findMaxPix(img):
    """
    get the pixel position of the maximum-value

    Parameters:
        img (numpy.array-float): image

    Returns:
        imgMaxPix_x, imgMaxPix_y (int, int): pixel position in columns & rows
    """
    maxIndx = np.argmax(img)
    maxIndx = np.unravel_index(maxIndx, np.array(img).shape)
    imgMaxPix_x = maxIndx[1]
    imgMaxPix_y = maxIndx[0]

    return imgMaxPix_x, imgMaxPix_y


###查找neighbors位置###
def findNeighbors(tx, ty, px, py, dr=0.1, dn=1, OnlyDistance=True):
    """
    find nearest neighbors by 2D-KDTree

    Parameters:
        tx, ty (float, float): a given position
        px, py (numpy.array, numpy.array): position data for tree
        dr (float-optional): distance
        dn (int-optional): nearest-N
        OnlyDistance (bool-optional): only use distance to find neighbors. Default: True

    Returns:
        dataq (numpy.array): index
    """
    datax = px
    datay = py
    tree = spatial.KDTree(list(zip(datax.ravel(), datay.ravel())))

    dataq=[]
    rr = dr
    if OnlyDistance == True:
        dataq = tree.query_ball_point([tx, ty], rr)
    if OnlyDistance == False:
        while len(dataq) < dn:
            dataq = tree.query_ball_point([tx, ty], rr)
            rr += dr
        dd = np.hypot(datax[dataq]-tx, datay[dataq]-ty)
        ddSortindx = np.argsort(dd)
        dataq = np.array(dataq)[ddSortindx[0:dn]]
    return dataq

###查找neighbors位置-hoclist###
def hocBuild(partx, party, nhocx, nhocy, dhocx, dhocy):
    if np.max(partx) > nhocx*dhocx:
        print('ERROR')
        sys.exit()
    if np.max(party) > nhocy*dhocy:
        print('ERROR')
        sys.exit()

    npart  = partx.size
    hoclist= np.zeros(npart, dtype=np.int32)-1
    hoc = np.zeros([nhocy, nhocx], dtype=np.int32)-1
    for ipart in range(npart):
        ix = int(partx[ipart]/dhocx)
        iy = int(party[ipart]/dhocy)
        hoclist[ipart] = hoc[iy, ix]
        hoc[iy, ix] = ipart
    return hoc, hoclist

def hocFind(px, py, dhocx, dhocy, hoc, hoclist):
    ix = int(px/dhocx)
    iy = int(py/dhocy)

    neigh=[]
    it = hoc[iy, ix]
    while it != -1:
        neigh.append(it)
        it = hoclist[it]
    return neigh

def findNeighbors_hoclist(px, py, tx=None,ty=None, dn=4, hoc=None, hoclist=None):
    nhocy = nhocx = 20

    pxMin = np.min(px)
    pxMax = np.max(px)
    pyMin = np.min(py)
    pyMax = np.max(py)

    dhocx = (pxMax - pxMin)/(nhocx-1)
    dhocy = (pyMax - pyMin)/(nhocy-1)
    partx = px - pxMin +dhocx/2
    party = py - pyMin +dhocy/2

    if hoc is None:
        hoc, hoclist = hocBuild(partx, party, nhocx, nhocy, dhocx, dhocy)
        return hoc, hoclist

    if hoc is not None:
        tx = tx - pxMin +dhocx/2
        ty = ty - pyMin +dhocy/2
        itx = int(tx/dhocx)
        ity = int(ty/dhocy)

        ps = [-1, 0, 1]
        neigh=[]
        for ii in range(3):
            for jj in range(3):
                ix = itx + ps[ii]
                iy = ity + ps[jj]
                if ix < 0:
                    continue
                if iy < 0:
                    continue
                if ix > nhocx-1:
                    continue
                if iy > nhocy-1:
                    continue

                #neightt = myUtil.hocFind(ppx, ppy, dhocx, dhocy, hoc, hoclist)
                it = hoc[iy, ix]
                while it != -1:
                    neigh.append(it)
                    it = hoclist[it]
                #neigh.append(neightt)
        #ll = [i for k in neigh for i in k]
        if dn != -1:
            ptx = np.array(partx[neigh])
            pty = np.array(party[neigh])
            dd  = np.hypot(ptx-tx, pty-ty)
            idx = np.argsort(dd)
            neigh= np.array(neigh)[idx[0:dn]]
        return neigh



###PSF中心对齐###
def psfCentering(img, apSizeInArcsec=0.5, psfSampleSizeInMicrons=5, focalLengthInMeters=28, CenteringMode=1):
    """
    centering psf within an aperture

    Parameters:
        img (numpy.array): image
        apSizeInArcsec (float-optional): aperture size in arcseconds.
        psfSampleSizeInMicrons (float-optional): psf pixel size in microns.
        focalLengthInMeters (float-optional): csst focal length im meters.
        CenteringMode (int-optional): how to center psf images

    Returns:
        imgT (numpy.array)
    """
    if CenteringMode == 1:
        imgMaxPix_x, imgMaxPix_y = findMaxPix(img)
    if CenteringMode == 2:
        y,x = ndimage.center_of_mass(img)  #y-rows, x-cols
        imgMaxPix_x = int(x)
        imgMaxPix_y = int(y)
    apSizeInMicrons = np.deg2rad(apSizeInArcsec/3600.)*focalLengthInMeters*1e6
    apSizeInPix = apSizeInMicrons/psfSampleSizeInMicrons
    apSizeInPix = np.int(np.ceil(apSizeInPix))
    imgT = np.zeros_like(img)
    ngy, ngx = img.shape
    cy = int(ngy/2)
    cx = int(ngx/2)
    imgT[cy-apSizeInPix:cy+apSizeInPix+1,
         cx-apSizeInPix:cx+apSizeInPix+1] = \
    img[imgMaxPix_y-apSizeInPix:imgMaxPix_y+apSizeInPix+1,
        imgMaxPix_x-apSizeInPix:imgMaxPix_x+apSizeInPix+1]
    return imgT


###插值对齐-fft###
def psfCentering_FFT(image):
    """
    centering psf within an aperture by FFT
    """
    ny, nx = image.shape
    py, px = ndimage.center_of_mass(image)
    dx = (px - nx/2)
    dy = (py - ny/2)
    k=np.zeros((nx,ny,2),dtype=float)
    fg=np.fft.fft2(image)
    ge =np.zeros_like(fg)

    inx = int(nx/2)
    jny = int(ny/2)
    #prepare for the phase multiply matrix
    #left bottom
    for i in range(inx+1):
        for j in range(jny+1):
            k[i][j][0]=i;
            k[i][j][1]=j;
    #top right
    for i in range(inx-1):
        for j in range(jny-1):
            k[i+inx+1][j+jny+1][0]=(-(nx/2-1)+i)
            k[i+inx+1][j+jny+1][1]=(-(ny/2-1)+j)
    #bottom right
    for i in range(inx+1):
        for j in range(jny-1):
            k[i][j+jny+1][0]=i
            k[i][j+jny+1][1]=(-(ny/2-1)+j)
    #top left
    for i in range(inx-1):
        for j in range(int(jny+1)):
            k[i+inx+1][j][0]=(-(nx/2-1)+i)
            k[i+inx+1][j][1]=j
    for i in range(nx):
        for j in range(ny):
            ge[i][j]=fg[i][j]*np.exp(2.*np.pi*(dx*k[i][j][0]/nx+dy*k[i][j][1]/ny)*1j)
    get=np.fft.ifft2(ge).real
    return(get)


###图像叠加###
def psfStack(*psfMat):
    """
    stacked image from the different psfs

    Parameters:
        *psfMat (numpy.array): the different psfs for stacking

    Returns:
        img (numpy.array): image
    """
    nn = len(psfMat)
    img = np.zeros_like(psfMat[0])
    for ii in range(nn):
        img += psfMat[ii]/np.sum(psfMat[ii])
    img /= np.sum(img)
    return img


###计算PSF椭率-接口###
def psfSizeCalculator(psfMat, psfSampleSize=2.5, CalcPSFcenter=True, SigRange=True, TailorScheme=2, cenPix=None):
    """
    calculate psf size & ellipticity

    Parameters:
        psfMat (numpy.array): image
        psfSampleSize (float-optional): psf size in microns.
        CalcPSFcenter (bool-optional): whether calculate psf center. Default: True
        SigRange (bool-optional): whether use psf tailor. Default: False
        TailorScheme (int-optional): which method for psf tailor. Default: 1
    Returns:
        cenX, cenY (float, float): the pixel position of the mass center
        sz (float): psf size
        e1, e2 (float, float): psf ellipticity
        REE80 (float): radius of REE80 in arcseconds
    """
    psfSampleSize = psfSampleSize*1e-3 #mm

    REE80 = -1.0  ##encircling 80% energy
    if SigRange is True:
        if TailorScheme == 1:
            psfMat = imSigRange(psfMat, fraction=0.80)
            psfInfo['psfMat'] = psfMat  #set on/off
        if TailorScheme == 2:
            #img = psfTailor(psfMat, apSizeInArcsec=0.5)
            imgX, REE80 = psfEncircle(psfMat, cenPix=cenPix)
            #psfMat = img
            REE80 = REE80[0]

    if CalcPSFcenter is True:
        img = psfMat/np.sum(psfMat)
        y,x = ndimage.center_of_mass(img)  #y-rows, x-cols
        cenX = x
        cenY = y
    if CalcPSFcenter is False:
        cenPix_X = psfMat.shape[1]/2 #90
        cenPix_Y = psfMat.shape[0]/2 #90
        cenX = cenPix_X + psfInfo['centroid_x']/psfSampleSize
        cenY = cenPix_Y - psfInfo['centroid_y']/psfSampleSize

    pixSize = 1
    sz, e1, e2 = psfSecondMoments(psfMat, cenX, cenY, pixSize=pixSize)

    return cenX, cenY, sz, e1, e2, REE80


###计算PSF椭率###
def psfSecondMoments(psfMat, cenX, cenY, pixSize=1):
    """
    estimate the psf ellipticity by the second moment of surface brightness

    Parameters:
        psfMat (numpy.array-float): image
        cenX, cenY (float, float): pixel position of the psf center
        pixSize (float-optional): pixel size

    Returns:
        sz (float): psf size
        e1, e2 (float, float): psf ellipticity
    """
    apr = 0.5 #arcsec, 0.5角秒内测量
    fl  = 28. #meters
    pxs = 5.0 #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


###计算REE80###
def psfEncircle(img, fraction=0.8, psfSampleSizeInMicrons=2.5, focalLengthInMeters=28, cenPix=None):
    """
    psf tailor within a given percentage.

    Parameters:
        img (numpy.array-float): image
        fraction (float-optional): a percentage for psf tailor.
        psfSampleSizeInMicrons (float-optional): psf pixel size in microns.
        focalLengthInMeters (float-optional): csst focal length im meters.
    Returns:
        img*wgt (numpy.array-float): image
        REE80 (float): radius of REE80 in arcseconds.
    """
    #imgMaxPix_x, imgMaxPix_y = findMaxPix(img)
    y,x = ndimage.center_of_mass(img)  #y-rows, x-cols
    imgMaxPix_x = x #int(x)
    imgMaxPix_y = y #int(y)
    if cenPix != None:
        imgMaxPix_x = cenPix[0]
        imgMaxPix_y = cenPix[1]

    im1 = img.copy()
    im1size = im1.shape

    dis = np.zeros_like(img)
    for irow in range(im1size[0]):
        for icol in range(im1size[1]):
            dx = icol - imgMaxPix_x
            dy = irow - imgMaxPix_y
            dis[irow, icol] = np.hypot(dx, dy)

    nn = im1size[1]*im1size[0]
    disX = dis.reshape(nn)
    disXsortId = np.argsort(disX)

    imgX = img.reshape(nn)
    imgY = imgX[disXsortId]
    psfFrac = np.cumsum(imgY)/np.sum(imgY)
    ind = np.where(psfFrac > fraction)[0][0]

    wgt = np.ones_like(dis)
    #wgt[np.where(dis > dis[np.where(img == imgY[ind])])] = 0

    REE80 = np.rad2deg(dis[np.where(img == imgY[ind])]*psfSampleSizeInMicrons*1e-6/focalLengthInMeters)*3600
    return img*wgt, REE80


###图像能量百分比查找###
def imSigRange(img, fraction=0.80):
    """
    extract the image within x-percent (DISCARD)

    Parameters:
        img (numpy.array-float): image
        fraction (float-optional): a percentage

    Returns:
        im1 (numpy.array-float): image
    """
    im1 = img.copy()
    im1size = im1.shape
    im2 = np.sort(im1.reshape(im1size[0]*im1size[1]))
    im2 = im2[::-1]
    im3 = np.cumsum(im2)/np.sum(im2)
    loc = np.where(im3 > fraction)
    #print(im3[loc[0][0]], im2[loc[0][0]])
    im1[np.where(im1 <= im2[loc[0][0]])]=0

    return im1


###孔径内图像裁剪###
def psfTailor(img, apSizeInArcsec=0.5, psfSampleSizeInMicrons=5, focalLengthInMeters=28, cenPix=None):
    """
    psf tailor within a given aperture size

    Parameters:
        img (numpy.array-float): image
        apSizeInArcsec (float-optional): aperture size in arcseconds.
        psfSampleSizeInMicrons (float-optional): psf pixel size in microns.
        focalLengthInMeters (float-optional): csst focal length im meters.
    Returns:
        imgT (numpy.array-float): image
    """
    #imgMaxPix_x, imgMaxPix_y = findMaxPix(img)
    y,x = ndimage.center_of_mass(img)  #y-rows, x-cols
    imgMaxPix_x = int(x)
    imgMaxPix_y = int(y)
    if cenPix != None:
        imgMaxPix_x = int(cenPix[0])
        imgMaxPix_y = int(cenPix[1])

    apSizeInMicrons = np.deg2rad(apSizeInArcsec/3600.)*focalLengthInMeters*1e6
    apSizeInPix = apSizeInMicrons/psfSampleSizeInMicrons
    apSizeInPix = np.int(np.ceil(apSizeInPix))
    print('apSizeInPix=',apSizeInPix)
    imgT = np.zeros_like(img)
    imgT[imgMaxPix_y-apSizeInPix:imgMaxPix_y+apSizeInPix+1,
         imgMaxPix_x-apSizeInPix:imgMaxPix_x+apSizeInPix+1] = \
    img[imgMaxPix_y-apSizeInPix:imgMaxPix_y+apSizeInPix+1,
        imgMaxPix_x-apSizeInPix:imgMaxPix_x+apSizeInPix+1]
    return imgT


###centroid with a window###
def centroidWgt(img, nt=160):
    #libCentroid = ctypes.CDLL('/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF/libCentroid/libCentroid.so')  # CDLL加载库
    libCentroid = ctypes.CDLL('/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF_20210108/libCentroid/libCentroid.so')  # CDLL加载库
    libCentroid.centroidWgt.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_double)]
    libCentroid.imSplint.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_double), ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_float)]

    imx = img/np.sum(img)
    ny, nx = imx.shape

    #imx centroid
    nn = nx*ny
    arr   = (ctypes.c_float*nn)()
    arr[:] = imx.reshape(nn)
    para   = (ctypes.c_double*10)()
    libCentroid.centroidWgt(arr, ny, nx, para)
    imx_cy = para[3] #irow
    imx_cx = para[4] #icol

    #imx -> imy
    nxt=nyt=nt
    nn=nxt*nyt
    yat = (ctypes.c_float*nn)()
    libCentroid.imSplint(arr, ny, nx, para, nxt, nyt, yat)
    imy = np.array(yat[:]).reshape([nxt, nyt])

    return imy, imx_cx, imx_cy

'''
def psfCentering_wgt(ipsfMat, psf_image_x, psf_image_y, psfSampleSizeInMicrons=5.0):
    img, cx, cy = centroidWgt(ipsfMat, nt=160)
    
    nrows, ncols = ipsfMat.shape
    cyt = (cy + nrows/2)*psfSampleSizeInMicrons*1e-3 +psf_image_y
    cxt = (cx + ncols/2)*psfSampleSizeInMicrons*1e-3 +psf_image_x
    return img, cxt, cyt
'''




###binning image###
def binningPSF(img, ngg):
    imgX = img.reshape(ngg, img.shape[0]//ngg, ngg, img.shape[1]//ngg).mean(-1).mean(1)
    return imgX