import sys
import numpy as np
import scipy.io
from itertools import islice

from PSF.PSFInterp.PSFUtil import *



###加载PSF信息###
def LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=True, PSFCentroidWgt=False, PixSizeInMicrons=2.5):
    """
    load psf informations from psf matrix.

    Parameters:
        iccd (int): ccd number [1,30].
        iwave(int): wave-index [1,4].
        ipsf (int): psf number [1, 100].
        psfPath (int): path to psf matrix
        InputMaxPixelPos(bool-optional): only True for 30*30 psf-matrix 
        PSFCentroidWgt(bool-optional): True for using psfMat with libCentroid.lib
    Returns:
        psfInfo (dirctionary)
    """

    if iccd not in np.linspace(1, 30, 30, dtype='int'):
        print('Error - iccd should be in [1, 30].')
        sys.exit()
    if iwave not in np.linspace(1, 4, 4, dtype='int'):
        print('Error - iwave should be in [1, 4].')
        sys.exit()
    if ipsf not in np.linspace(1, 900, 900, dtype='int'):
        print('Error - ipsf should be in [1, 900].')
        sys.exit()

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

    psfInfo['pixsize'] = PixSizeInMicrons  #microns

    #获取ipsf波长
    fpathWave = fpath +'/' +'1_wavelength.txt'
    f = open(fpathWave, 'r')
    wavelength = np.float(f.readline())
    f.close()
    psfInfo['wavelength'] = wavelength

    #获取ipsf位置
    fpathCoordinate = fpath +'/' +'4_PSF_coordinate.txt'
    f = open(fpathCoordinate, 'r')
    header = f.readline()
    for line in islice(f, ipsf-1, ipsf):
        line = line.strip()
        columns = line.split()
    f.close()
    icol = 0
    psfInfo['field_x'] = float(columns[icol])    #deg, 视场采样位置
    icol+= 1
    psfInfo['field_y'] = float(columns[icol])    #deg
    icol+= 1
    psfInfo['centroid_x'] = float(columns[icol]) #mm, psf质心相对主光线的偏移量
    icol+= 1
    psfInfo['centroid_y'] = float(columns[icol]) #mm
    icol+= 1
    if InputMaxPixelPos == True:
        psfInfo['max_x'] = float(columns[icol])  #mm, max pixel postion
        icol+= 1
        psfInfo['max_y'] = float(columns[icol])  #mm
        icol+= 1
    psfInfo['image_x'] = float(columns[icol])    #mm, 主光线位置
    icol+= 1
    psfInfo['image_y'] = float(columns[icol])    #mm
    

    if PSFCentroidWgt:
        psfInfo['centroid_x'] = data['cx'][0][0] #mm, psfCentroidWgt质心相对主光线的偏移量
        psfInfo['centroid_y'] = data['cy'][0][0] #mm
    return psfInfo



###插值PSF图像-IDW方法###
def psfMaker_IDW(px, py, PSFMat, cen_col, cen_row, IDWindex=2, OnlyNeighbors=True, hoc=None, hoclist=None, PSFCentroidWgt=False):
    """
    psf interpolation by IDW

    Parameters:
        px, py (float, float): position of the target
        PSFMat (numpy.array): image
        cen_col, cen_row (numpy.array, numpy.array): potions of the psf centers
        IDWindex (int-optional): the power index of IDW
        OnlyNeighbors (bool-optional): only neighbors are used for psf interpolation

    Returns:
        psfMaker (numpy.array)
    """

    minimum_psf_weight = 1e-8
    ref_col = px
    ref_row = py

    ngy, ngx = PSFMat[0, :, :].shape
    npsf = PSFMat[:, :, :].shape[0]
    psfWeight = np.zeros([npsf])

    if OnlyNeighbors == True:
        if hoc is None:
            neigh = findNeighbors(px, py, cen_col, cen_row, dr=5., dn=4, OnlyDistance=False)
        if hoc is not None:
            neigh = findNeighbors_hoclist(cen_col, cen_row, tx=px,ty=py, dn=4, hoc=hoc, hoclist=hoclist)

        neighFlag = np.zeros(npsf)
        neighFlag[neigh] = 1

    for ipsf in range(npsf):
        if OnlyNeighbors == True:
            if neighFlag[ipsf] != 1:
                continue

        dist = np.sqrt((ref_col - cen_col[ipsf])**2 + (ref_row - cen_row[ipsf])**2)
        if IDWindex == 1:
            psfWeight[ipsf] = dist
        if IDWindex == 2:
            psfWeight[ipsf] = dist**2
        if IDWindex == 3:
            psfWeight[ipsf] = dist**3
        if IDWindex == 4:
            psfWeight[ipsf] = dist**4
        psfWeight[ipsf] = max(psfWeight[ipsf], minimum_psf_weight)
        psfWeight[ipsf] = 1./psfWeight[ipsf]
    psfWeight /= np.sum(psfWeight)

    psfMaker  = np.zeros([ngy, ngx], dtype=np.float32)
    for ipsf in range(npsf):
        if OnlyNeighbors == True:
            if neighFlag[ipsf] != 1:
                continue

        iPSFMat = PSFMat[ipsf, :, :].copy()
        ipsfWeight = psfWeight[ipsf]

        psfMaker += iPSFMat * ipsfWeight
    psfMaker /= np.nansum(psfMaker)

    return psfMaker



###TEST###
if __name__ == '__main__':
    iccd = 1
    iwave= 1
    ipsf = 1 
    psfPath = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_20210108/CSST_psf_ciomp_2p5um_cycle3_ccr90'

    psfInfo = LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=True, PSFCentroidWgt=True)
    print('psfInfo.keys:', psfInfo.keys())
    print(psfInfo)