import numpy as np
import astropy.units as u
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.io import ascii
from sklearn.neighbors import BallTree

TU_catalog = "injected_bkgsub_img.cat"
source_catalog = "extracted_injected_bkgsub_img.fits"

def convert_catalog(catname):
    text_file = ascii.read(catname)
    text_file.write('test_ascii_to_fits.fits', overwrite=True)

def read_catalog(catname, ext_num=1, ra_name='ra', dec_name='dec', col_list=[]):
    hdu = fits.open(catname)
    hdr = hdu[ext_num].header
    # print(hdr)
    data = hdu[ext_num].data
    ra = data[ra_name]
    dec = data[dec_name]
    col_other = []
    if len(col_list) > 0:
        for col in col_list:
            col_other.append(data[col])
    return ra, dec, col_other

def match_catalogs_sky(ra1, dec1, ra2, dec2, max_dist=0.6, others1=[], others2=[], thresh=[]):
    cat1 = SkyCoord(ra=ra1*u.degree, dec=dec1*u.degree)
    cat2 = SkyCoord(ra=ra2*u.degree, dec=dec2*u.degree)
    idx1, idx2, d2d, d3d = cat2.search_around_sky(cat1, 1*u.deg)
    # print(idx1)
    # print(idx2)
    # print(np.shape(idx1))
    # print(np.shape(idx2))
    # TODO

def match_catalogs_img(x1, y1, x2, y2, max_dist=2, others1=[], others2=[], thresh=[]):
    cat1 = np.array([(x, y) for x,y in zip(x1, y1)])
    cat2 = np.array([(x, y) for x,y in zip(x2, y2)])
    tree = BallTree(cat2)
    idx1 = tree.query_radius(cat1, r = max_dist)
    tree = BallTree(cat1)
    idx2 = tree.query_radius(cat2, r = max_dist)
    tot = 0
    for idx in idx1:
        if len(idx) == 0:
            continue
        if len(idx) > 1: print(len(idx))
        tot += 1
    print("number of matched sources = ", tot)
    return idx1, idx2

if __name__=="__main__":
    convert_catalog(TU_catalog)
    ra_TU, dec_TU, _ = read_catalog('test_ascii_to_fits.fits', ext_num=1, ra_name="ra", dec_name="dec")
    x_TU, y_TU, col_list = read_catalog('test_ascii_to_fits.fits', ext_num=1, ra_name="xImage", dec_name="yImage", col_list=["mag"])
    mag_TU = col_list[0]
    # ra_source, dec_source, _ = read_catalog(source_catalog, ext_num=2, ra_name="ALPHAPEAK_J2000", dec_name="DELTAPEAK_J2000")
    x_source, y_source, _ = read_catalog(source_catalog, ext_num=1, ra_name="X_IMAGE", dec_name="Y_IMAGE")
    # match_catalogs_sky(ra1=ra_TU, dec1=dec_TU, ra2=ra_source, dec2=dec_source)
    idx1, idx2, = match_catalogs_img(x1=x_TU, y1=y_TU, x2=x_source, y2=y_source)
    # print(ra_TU, dec_TU)