import numpy as np
import matplotlib.pyplot as plt

import scipy.io
import galsim

vc_A = 2.99792458e+18  # speed of light: A/s
vc_m = 2.99792458e+8   # speed of light: m/s
h_Plank = 6.626196e-27 # Plank constant: erg s

def photonEnergy(lambd):
    nu = vc_A / lambd
    eph = h_Plank * nu
    return eph

pixSize = 0.037
mag_star = 15.
Nx = 256
Ny = 256

###加载PSF信息###
def LoadPSF(iccd, iwave, ipsf, psfPath, psfSampleSize=5, PSFCentroidWgt=False):
    psfInfo = {}
    fpath = psfPath +'/' +'ccd{:}'.format(iccd) +'/' + 'wave_{:}'.format(iwave)

    #获取ipsf矩阵
    if not PSFCentroidWgt:
        ##读取PSF原数据
        fpathMat = fpath +'/' +'5_psf_array' +'/' +'psf_{:}.mat'.format(ipsf)
        data = scipy.io.loadmat(fpathMat)
        psfInfo['psfMat'] = data['psf']
    if PSFCentroidWgt:
        ##读取PSFCentroidWgt
        ffpath = psfPath +'_proc/' +'ccd{:}'.format(iccd) +'/' + 'wave_{:}'.format(iwave)
        ffpathMat = ffpath +'/' +'5_psf_array' +'/' +'psf_{:}_centroidWgt.mat'.format(ipsf)
        data = scipy.io.loadmat(ffpathMat)
        psfInfo['psfMat'] = data['psf']
    return psfInfo

def psfCentering(img, apSizeInArcsec=0.5, psfSampleSizeInMicrons=5, focalLengthInMeters=28, CenteringMode=1):
    imgMaxPix_x, imgMaxPix_y = findMaxPix(img)
    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

def findMaxPix(img):
    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

def magToFlux(mag):
    flux = 10**(-0.4*(mag+48.6))
    return flux

def getElectronFluxFilt(mag, exptime=150.):
    pE = photonEnergy(lambd=6199.8)
    flux = magToFlux(mag)
    factor = 1.0e4 * flux/pE * vc_A * (1.0/5370.0 - 1.0/7030.0)
    return factor * 0.5040 * np.pi * (0.5 * 2.0)**2 * exptime


def radial_profile(img, cx, cy, nbins=100, Rmin=16, Rmax=128):
    nx = img.shape[0]
    ny = img.shape[1]
    x = np.arange(0, img.shape[0], 1)
    y = np.arange(0, img.shape[1], 1)
    xx, yy = np.meshgrid(x, y)
    dist = np.sqrt((xx - cx)**2 + (yy - cy)**2)
    img_flat = img.flatten()
    dist_flat = dist.flatten()
    a, bins = np.histogram(dist_flat, range=(Rmin, Rmax), bins=nbins, weights=img_flat)
    b, bins = np.histogram(dist_flat, range=(Rmin, Rmax), bins=nbins)
    b[b==0] = 1.
    mid = (bins[0:-1]+bins[1:])/2.
    return a / b, mid


if __name__ == '__main__':
    iccd   = 1
    iwave  = 1
    ipsf   = 1
    psfPath= '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp_30X30'
    psfInfo= LoadPSF(iccd, iwave, ipsf, psfPath, psfSampleSize=5, PSFCentroidWgt=False)    

    ipsfMat = psfInfo['psfMat']
    print(ipsfMat[0:100,0:100])
    cgrid   = 128
    dgrid   = 15

    cx = int(Nx/2)
    cy = int(Ny/2)

    with np.printoptions(precision=5, suppress=True):
        fig = plt.figure(figsize=(18,12))
        ax  = plt.subplot(2,2,1)
        plt.imshow(ipsfMat[cgrid-dgrid:cgrid+dgrid, cgrid-dgrid:cgrid+dgrid], origin='lower')
        plt.annotate('originalPSF', [0.1, 0.85], xycoords='axes fraction', color='w')
        
        ax  = plt.subplot(2,2,2)
        img = galsim.ImageF(ncol=Nx, nrow=Ny, scale=pixSize)
        psf_img = galsim.ImageF(ipsfMat, scale=pixSize)
        psf = galsim.InterpolatedImage(psf_img)
        psf_list = [psf] * 4
        mag_star = 15
        nphotons_tot = getElectronFluxFilt(mag=mag_star)
        print(nphotons_tot)
        for i in range(4):
            nphotons = nphotons_tot / 4.
            # star = galsim.Gaussian(sigma=1.e-8, flux=1.)
            star = galsim.DeltaFunction()
            star = star.withFlux(nphotons)
            star = galsim.Convolve(psf_list[i], star)
            img = star.drawImage(image=img, method='phot', center=galsim.PositionD(cx, cy), add_to_image=True)
        print(img.array.shape)
        plt.imshow(img.array[cx-dgrid:cx+dgrid, cy-dgrid:cy+dgrid], origin='lower')
        plt.annotate('photon shooting 4 times', [0.1, 0.85], xycoords='axes fraction', color='w')
        plt.colorbar()
        ax = plt.subplot(2,2,3)
        val, bins = radial_profile(img=img.array, cx=cx, cy=cy, nbins=100, Rmax=cx)
        # plt.hist(val, bins=bins)
        plt.plot(bins, val, label='mag=%d'%(mag_star))

        img = galsim.ImageF(ncol=Nx, nrow=Ny, scale=pixSize)
        psf_img = galsim.ImageF(ipsfMat, scale=pixSize)
        psf = galsim.InterpolatedImage(psf_img)
        psf_list = [psf] * 4
        mag_star = 13
        nphotons_tot = getElectronFluxFilt(mag=mag_star)
        print(nphotons_tot)
        for i in range(4):
            nphotons = nphotons_tot / 4.
            # star = galsim.Gaussian(sigma=1.e-8, flux=1.)
            star = galsim.DeltaFunction()
            star = star.withFlux(nphotons)
            star = galsim.Convolve(psf_list[i], star)
            img = star.drawImage(image=img, method='phot', center=galsim.PositionD(cx, cy), add_to_image=True)
        print(img.array.shape)
        val, bins = radial_profile(img=img.array, cx=cx, cy=cy, nbins=100, Rmax=cx)
        # plt.hist(val, bins=bins)
        plt.plot(bins, val, label='mag=%d'%(mag_star))
        plt.legend(loc='upper right', fancybox=True)
        plt.xlabel("R [pix]", size='x-large')
        plt.ylabel("photon count", size='x-large')
        plt.ylim([0, 500])
        # print(img.array[0:100,0:100])
        #psf Convolve galsim.DeltaFunction
        #photon shooting ?
        #plot image?

        # ax  = plt.subplot(2,2,2)
        # plt.imshow(ipsfMat[cgrid-dgrid:cgrid+dgrid, cgrid-dgrid:cgrid+dgrid], origin='lower')
        # plt.annotate('originalPSF', [0.1, 0.85], xycoords='axes fraction', color='w')
        
        # ax  = plt.subplot(2,2,4)
        # img = galsim.ImageF(ncol=Nx, nrow=Ny, scale=pixSize)
        # psf_img = galsim.ImageF(ipsfMat, scale=pixSize)
        # psf = galsim.InterpolatedImage(psf_img)
        # psf_list = [psf] * 4
        # nphotons_tot = getElectronFluxFilt(mag=mag_star)
        # print(nphotons_tot)
        # obj_list = []
        # for i in range(4):
        #     nphotons = nphotons_tot / 4.
        #     # star = galsim.Gaussian(sigma=1.e-8, flux=1.)
        #     star = galsim.DeltaFunction()
        #     star = star.withFlux(nphotons)
        #     star = galsim.Convolve(psf_list[i], star)
        #     obj_list.append(star)
        # star = galsim.Sum(obj_list)
        # img = star.drawImage(image=img, method='phot', center=galsim.PositionD(cx, cy), add_to_image=True)
        # plt.annotate('photon shooting once', [0.1, 0.85], xycoords='axes fraction', color='w')
        # print(img.array.shape)
        # plt.imshow(img.array[cx-dgrid:cx+dgrid, cy-dgrid:cy+dgrid], origin='lower')
        # plt.colorbar()
        # print(img.array[0:100,0:100])

        # ax  = plt.subplot(2,3,2)
        # imy = psfCentering(ipsfMat, apSizeInArcsec=2.0)
        # plt.imshow(imy[cgrid-dgrid:cgrid+dgrid, cgrid-dgrid:cgrid+dgrid], origin='lower')
        # plt.annotate('PSFCentroidOld', [0.1, 0.85], xycoords='axes fraction', color='w')
        # ax  = plt.subplot(2,3,5)
        # #psf Convolve galsim.DeltaFunction
        # #photon shooting ?
        # #plot image?

        # ax  = plt.subplot(2,3,3)
        # psfInfo= LoadPSF(iccd, iwave, ipsf, psfPath, psfSampleSize=5, PSFCentroidWgt=True)
        # ipsfMat = psfInfo['psfMat']
        # # print(ipsfMat[0:100,0:100])
        # plt.imshow(ipsfMat[cgrid-dgrid:cgrid+dgrid, cgrid-dgrid:cgrid+dgrid], origin='lower')
        # plt.annotate('PSFCentroidNew', [0.1, 0.85], xycoords='axes fraction', color='w')
        # ax  = plt.subplot(2,3,6)
        # #psf Convolve galsim.DeltaFunction
        # #photon shooting ?
        # #plot image?
        
        plt.savefig('testPlot.pdf')








