"""
hoclist 查好最近的采样点
"""
import sys,os,math
import time
from IPython import display
import numpy as np
import matplotlib.pyplot as plt

sys.path.append("/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF")
import PSFConfig as myConfig
import PSFUtil as myUtil

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 = myUtil.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




#####################################
iccd = 13  #[1, 30]
iwave= 1  #[1, 4]
ipsf = 1  #[1, 100]
psfPathA = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp'
psfPathB = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp_30X30'

for iccd in range(1,31):
  if iccd != 1:
    continue
  npsfA = 100
  npsfB = 900

  imxA = np.zeros(npsfA)
  imyA = np.zeros(npsfA)
  imxB = np.zeros(npsfB)
  imyB = np.zeros(npsfB)

  for ipsf in range(npsfA):
    print('ipsf-{:}'.format(ipsf), end='\r')
    psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf+1, psfPathA, InputMaxPixelPos=False, PSFCentroidWgt=True)
    imxA[ipsf] = psfInfo['image_x']+psfInfo['centroid_x']
    imyA[ipsf] = psfInfo['image_y']+psfInfo['centroid_y']

  for ipsf in range(npsfB):
    print('ipsf-{:}'.format(ipsf), end='\r')
    psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf+1, psfPathB, InputMaxPixelPos=True, PSFCentroidWgt=True)
    imxB[ipsf] = psfInfo['image_x']+psfInfo['centroid_x']
    imyB[ipsf] = psfInfo['image_y']+psfInfo['centroid_y']

  #######
  px = imxB
  py = imyB
  hoc,hoclist = findNeighbors_hoclist(px, py)

  #if True:
  for ipsfA in range(npsfA):
    #plt.close()
    plt.clf()
    #fig = plt.figure(figsize=(12,12))
    plt.plot(imxA, imyA, 'b.')
    plt.plot(imxB, imyB, 'r.')
    

    #ipsfA = 19
    tx = imxA[ipsfA]
    ty = imyA[ipsfA]
    neigh = findNeighbors_hoclist(px, py, tx=tx,ty=ty, dn=4, hoc=hoc, hoclist=hoclist)

    plt.plot(imxB[neigh],imyB[neigh],'x', ms=10)
    plt.plot(tx, ty,'+', ms=10)

    #plt.savefig('figs/psfCrossPos.pdf')
    display.clear_output(wait=True)
    plt.pause(0.05)


