# NAME:
#     indexFits_hdf5
# PURPOSE:
#     Return a healpix indexed catalog from a set of Fits
# CALLING:
#     write_StampsIndex(dir_cat=dir_temp)
# INPUTS:
#     dir_cat - the directory of Fits catalog, "<dir_cat>/stampCats/*.fits"
# OUTPUTS:
#     <dir_cat>/stampCatsIndex.hdf5
# OPTIONAL:
#     test_fits(nfits=100, dir_cat=None) - generate a set of Fits by galsim gaussian
# HISTORY:
#     Written by Chengliang Wei, 13 Apr. 2023
#     Included by csst-simulation, C.W. 25 Apr. 2023
#
#

import os
import numpy as np
import astropy.io.fits as fitsio
import h5py
import healpy
import galsim


def test_fits(nfits=100, dir_cat=None):
    for ifits in range(nfits):
        gal = galsim.Gaussian(sigma=np.random.uniform(0.2, 0.3)).shear(
            g1=np.random.uniform(-0.5, 0.5), g2=np.random.uniform(-0.5, 0.5))
        arr = gal.drawImage(nx=64, ny=64, scale=0.074).array

        hdu = fitsio.PrimaryHDU()

        hdu.data = arr
        hdu.header.set('index', ifits)
        hdu.header.set('ra',  60.+np.random.uniform(-0.2, 0.2))
        hdu.header.set('dec', -40.+np.random.uniform(-0.2, 0.2))
        hdu.header.set('mag_g', 22+np.random.uniform(-1, 1))
        hdu.header.set('pixScale', 0.074)

        fout = dir_cat+"stampCats/testStamp_{:}.fits".format(ifits)
        if os.path.exists(fout):
            os.remove(fout)
        hdu.writeto(fout)


def write_StampsIndex(dir_cat=None, DEBUG=False):
    MAXNUMBERINDEX = 10000
    NSIDE = 128
    fp = h5py.File(dir_cat+'stampCatsIndex.hdf5', 'w')
    grp1 = fp.create_group('Stamps')

    dataSet_Size = np.zeros(healpy.nside2npix(NSIDE), dtype=np.int64)
    fitsList = os.listdir(dir_cat+'stampCats/')  # 获取fits文件列表
    for istamp in range(len(fitsList)):
        print(istamp, ': ', fitsList[istamp], end='\r')

        hdu = fitsio.open(dir_cat+"stampCats/"+fitsList[istamp])
        tra = hdu[0].header['RA']
        tdec = hdu[0].header['DEC']

        healpixID = healpy.ang2pix(NSIDE, tra, tdec, nest=False, lonlat=True)

        if not (str(healpixID) in grp1):
            grp2 = grp1.create_group(str(healpixID))
        else:
            grp2 = grp1[str(healpixID)]

        if not ('ra' in grp2):
            dset_ra = grp2.create_dataset(
                'ra', (0,), dtype='f16',  maxshape=(MAXNUMBERINDEX, ))
            dset_dec = grp2.create_dataset(
                'dec', (0,), dtype='f16',  maxshape=(MAXNUMBERINDEX, ))
            dt = h5py.special_dtype(vlen=str)
            dset_fn = grp2.create_dataset(
                'filename', (0,), dtype=dt, maxshape=(MAXNUMBERINDEX, ))
        else:
            dset_ra = grp2['ra']
            dset_dec = grp2['dec']
            dset_fn = grp2['filename']

        dataSet_Size[healpixID] = dataSet_Size[healpixID]+1
        grp2['ra'].resize((dataSet_Size[healpixID],))
        grp2['dec'].resize((dataSet_Size[healpixID],))
        grp2['filename'].resize((dataSet_Size[healpixID],))

        dset_ra[dataSet_Size[healpixID]-1] = tra
        dset_dec[dataSet_Size[healpixID]-1] = tdec
        dset_fn[dataSet_Size[healpixID]-1] = fitsList[istamp]
    fp.close()

    if DEBUG:
        print('\n')
        ff = h5py.File(dir_cat+"stampCatsIndex.hdf5", "r")
        ss = 0
        for kk in ff['Stamps'].keys():
            print(kk, ff['Stamps'][kk]['ra'].size)
            ss = ss+ff['Stamps'][kk]['ra'].size
        print(ss)


if __name__ == '__main__':
    dir_temp = "./Catalog_test/"
    # test_fits(dir_cat=dir_temp)
    write_StampsIndex(dir_cat=dir_temp)