
import numpy as np
import h5py
import astropy.io.fits
from astropy.table import Table
from astropy import units as u
from astropy.coordinates import Angle
from astropy.coordinates import SkyCoord

catalogFile = '../catalogs/GLADE.hdf5'

with h5py.File(catalogFile, 'r') as f:
    ra, dec = f['ra'][:], f['dec'][:]
    distmpc, z = f['distmpc'][:], f['z'][:]
    magb, magk = f['magb'][:], f['magk'][:]
    GWGC, PGC, _2MASS = f['GWGC'][:], f['PGC'][:], f['2MASS'][:]
    HyperLEDA, SDSS = f['HyperLEDA'][:], f['SDSS'][:]
    # Convert bytestring to unicode
    GWGC = GWGC.astype('U')
    PGC = PGC.astype('U')
    HyperLEDA = HyperLEDA.astype('U')
    _2MASS = _2MASS.astype('U')
    SDSS = SDSS.astype('U')

filename = "../data/GW190425/zadko.dat"
lines = [line.rstrip('\n') for line in open(filename)]
ras, decs = [], []
for line in lines:
    lineSplit = line.split(" ")
    thisra  = Angle(lineSplit[0], unit=u.hour).deg
    thisdec = Angle(lineSplit[1], unit=u.deg).deg

    ras.append(thisra)
    decs.append(thisdec)
ras, decs = np.array(ras), np.array(decs)

dist_mean = 156
dist_std = 41
distance_min = dist_mean - 2*dist_std
distance_max = dist_mean + 2*dist_std

idx = np.where((distmpc >= distance_min) & (distmpc <= distance_max))[0]
ra, dec, distmpc = ra[idx], dec[idx], distmpc[idx]
magb, magk = magb[idx], magk[idx]
GWGC, PGC, HyperLEDA = GWGC[idx], PGC[idx], HyperLEDA[idx]
_2MASS, SDSS = _2MASS[idx], SDSS[idx]

catalog1 = SkyCoord(ra=ra*u.degree,
                    dec=dec*u.degree, frame='icrs')

for jj in range(len(ras)):
    thisra, thisdec = ras[jj], decs[jj]
    print(lines[jj])

    coord = SkyCoord(ra=thisra*u.degree, dec=thisdec*u.degree, frame='icrs')
    sep = coord.separation(catalog1)
    sep = sep.deg
    idx = np.where(sep <  0.3/2.0)[0]
    for ii in range(len(idx)):
        print(ra[idx[ii]], dec[idx[ii]], distmpc[idx[ii]], GWGC[idx[ii]], PGC[idx[ii]], _2MASS[idx[ii]], HyperLEDA[idx[ii]], SDSS[idx[ii]]) 
    print(" ")

