import numpy as np
import matplotlib.pyplot as plt
import ctypes
import galsim

libCentroid = ctypes.CDLL('../libCentroid.so')  # CDLL加载库
print('load libCenroid')

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)]


from scipy.io import loadmat
data = loadmat('/Users/chengliangwei/csstPSFdata/CSSOS_psf_ciomp/ccd20/wave_1/5_psf_array/psf_10.mat')
imx = data['psf']
imx = imx/np.sum(imx)
ny,nx = imx.shape
print(nx, ny)

#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_cx = para[3]
imx_cy = para[4]

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

#imy centroid
libCentroid.centroidWgt(yat, nyt, nxt, para)
imy_cx = para[3]
imy_cy = para[4]


#plot check
fig = plt.figure(figsize=(12, 6))

##imx_plot
ax = plt.subplot(1,2,1)
cpix= int(nx/2)
dpix= 10
plt.imshow(imx[cpix-dpix:cpix+dpix, cpix-dpix:cpix+dpix], origin='lower')
plt.plot(imx_cy-cpix+dpix, imx_cx-cpix+dpix, 'bx', ms = 20)

from scipy import ndimage
cx, cy = ndimage.center_of_mass(imx)
plt.plot(cy-cpix+dpix, cx-cpix+dpix, 'r+', ms = 20)

maxIndx = np.argmax(imx)
maxIndx = np.unravel_index(maxIndx, np.array(imx).shape)
imgMaxPix_x = maxIndx[1]
imgMaxPix_y = maxIndx[0]
apSizeInPix = 23
imgT = np.zeros_like(imx)
imgT[imgMaxPix_y-apSizeInPix:imgMaxPix_y+apSizeInPix+1,
    imgMaxPix_x-apSizeInPix:imgMaxPix_x+apSizeInPix+1] = \
    imx[imgMaxPix_y-apSizeInPix:imgMaxPix_y+apSizeInPix+1,
        imgMaxPix_x-apSizeInPix:imgMaxPix_x+apSizeInPix+1]
cx, cy = ndimage.center_of_mass(imgT)
plt.plot(cy - cpix+dpix, cx - cpix+dpix, 'b+', ms = 15)

##imy_plot
ax = plt.subplot(1,2,2)
cpix= int(nxt/2)
dpix= 10
plt.imshow(imy[cpix-dpix:cpix+dpix, cpix-dpix:cpix+dpix], origin='lower')
plt.plot(imy_cy-cpix+dpix, imy_cx-cpix+dpix, 'bx', ms = 20)
plt.plot([dpix, dpix],[0,dpix],'w:')
plt.plot([0, dpix],[dpix,dpix],'w:')

cx, cy = ndimage.center_of_mass(imy)
plt.plot(cy-cpix+dpix, cx-cpix+dpix, 'r+', ms = 20)

maxIndx = np.argmax(imy)
maxIndx = np.unravel_index(maxIndx, np.array(imy).shape)
imgMaxPix_x = maxIndx[1]
imgMaxPix_y = maxIndx[0]
apSizeInPix = 23
imgT = np.zeros_like(imy)
imgT[imgMaxPix_y-apSizeInPix:imgMaxPix_y+apSizeInPix+1,
    imgMaxPix_x-apSizeInPix:imgMaxPix_x+apSizeInPix+1] = \
    imy[imgMaxPix_y-apSizeInPix:imgMaxPix_y+apSizeInPix+1,
        imgMaxPix_x-apSizeInPix:imgMaxPix_x+apSizeInPix+1]
cx, cy = ndimage.center_of_mass(imgT)
plt.plot(cy - cpix+dpix, cx - cpix+dpix, 'b+', ms = 15)





plt.show()
