from django.core.paginator import Paginator
import time

from astropy_healpix import HEALPix
from astropy.coordinates import SkyCoord  # High-level coordinates
from astropy.coordinates import ICRS, Galactic, FK4, FK5  # Low-level frames
from astropy.coordinates import Angle, Latitude, Longitude  # Angles
import astropy.units as u

from heasarc.models import HeasarcBase
from heasarc.models import NSIDE_SOURCES, ORDER

from astrobasis.models import VLASS
from astrobasis.models import GLIMPSE
from astrobasis.models import GAIADR3
from astrobasis.models import AllWise
from astrobasis.models import TwoMASS

def match_src(hp,src,counterpatscat, counterpartname,minrad=5,maxdist=30):

    # testing for selected source
    #src = HeasarcBase.objects.get(pk=15314383)
    #print(src.name)
    
    crd = SkyCoord(src.ra, src.dec, frame="fk5", unit="deg")
    heal = hp.cone_search_skycoord(crd, radius=maxdist*u.arcsecond)

    try:
        cntp_cone = counterpatscat.filter(healpix__in=heal)
        #print("Selected in cone {} within {}".format(cntp_cone.count(),maxdist*u.arcsecond))
    except:
        return

    cntp_list=[]
    for cntp in cntp_cone:
        cntp_crd = SkyCoord(cntp.ra, cntp.dec, frame="fk5", unit="deg")
        cntp_sep=crd.separation(cntp_crd).arcsecond
        if(cntp_sep <= minrad):
            #print("{} {}".format(cntp.name,cntp_sep))
            cntp_list.append(cntp)

    getattr(src, counterpartname).clear()
    if(cntp_list):
        getattr(src, counterpartname).add(*cntp_list)
             
def heasarcbase_match(counterpatscat, counterpartname, perpage=10000, minrad=5,maxdist=30, survey=HeasarcBase):
    hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5())
    srcs = survey.objects.all().filter(error_radius__lte=minrad)#.order_by("-id")
    print("Selected {} from {} objects with error_radius <= {} arcsec".format(srcs.count(), survey, minrad))
    print("Perform cone search for {} within {} srcsec".format(counterpartname,maxdist))
    
    pages = Paginator(srcs, perpage)
    pages_count = pages.count
    pages_num = pages.num_pages
    pages_range = pages.page_range

    for i in pages_range:
        page = pages.page(i)
        start_time0 = time.time()
        for src in page.object_list:
            match_src(hp, src, counterpatscat, counterpartname, minrad=minrad, maxdist=maxdist)
        elapsed = time.time() - start_time0
        estimated = (elapsed*(pages_num-i))/60/60
        print("{} match for {}: Page {}/{} for {:.2f} sec, est. {:.2f} hours".format(survey,
                                                                                     counterpartname,i,
                                                                                     pages_num,
                                                                                     elapsed,
                                                                                     estimated))
def healpix_usercat(src):
    hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5())
    crd = SkyCoord(src.ra, src.dec, frame="fk5", unit="deg")
    src.healpix = hp.skycoord_to_healpix(crd)
    src.lii=crd.galactic.l.value 
    src.bii=crd.galactic.b.value 
    src.save()
        
def match_usercat(src, counterpatscat, counterpartname, minrad=5,maxdist=30):
    hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5())
    match_src(hp, src, counterpatscat, counterpartname, minrad=minrad, maxdist=maxdist)

def match_usercat_all(src, minrad=5,maxdist=30):
    match_usercat(src, VLASS.objects.all(),"vlass", minrad=minrad,maxdist=maxdist)
    match_usercat(src, GLIMPSE.objects.all(),"glimpse", minrad=minrad,maxdist=maxdist)
    match_usercat(src, GAIADR3.objects.all(),"gaia3", minrad=minrad,maxdist=maxdist)
    match_usercat(src, AllWise.objects.all(),"allwise", minrad=minrad,maxdist=maxdist)
    match_usercat(src, TwoMASS.objects.all(),"twomass", minrad=minrad,maxdist=maxdist)