from django.core.management.base import BaseCommand, CommandError
from django.core.paginator import Paginator

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 multiprocessing import Pool, cpu_count

from datetime import date
import datetime
import time
from django.utils import timezone
import astropy
from astropy.io import ascii
import pandas as pd
import pymysql
from sqlalchemy import create_engine

from heasarc.models import HeasarcBase
from artsurvey.models import ArtSurveySource, ArtSurvey
from srglib.utils import load_vizier_nvss
from srglib.utils import load_vizier_first
from srglib.utils import load_vizier_sumss

from astrobasis.models import GAIADR3
from astrobasis.models import GLIMPSE

from srglib.utils import find_counterparts
from srglib.utils import load_vizier_allwise
from srglib.utils import load_simbad_for_skymap_sources

from heasarc.models import NSIDE_SOURCES, ORDER

NPARTS=10

def do_match():

    minrad=5
    maxdist=30
    hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5())
    
    """
    g=GLIMPSE.objects.all()
    g2=g.filter(healpix__exact=0)
    print("Total {} healpix=NULL {}".format(g.count(),g2.count()))
    print(g2[0].ra,g2[0].dec,g2[0].healpix)
    return
    """

    srcs = HeasarcBase.objects.all().filter(error_radius__lte=minrad)#.order_by("-id")


    pages = Paginator(srcs, 100000)
    pages_count = pages.count
    pages_num = pages.num_pages
    pages_range = pages.page_range

    gaia3_all = GAIADR3.objects.all()

    for i in pages_range:
        page = pages.page(i)
        start_time0 = time.time()
        for src in page.object_list:
            crd = SkyCoord(src.ra, src.dec, frame="fk5", unit="deg")      
            heal = hp.cone_search_skycoord(crd, radius=maxdist*u.arcsecond)

            src.gaia3.clear()
            try:
                gaia3_all_cone = gaia3_all.filter(healpix__in=heal)
            except:
                continue
            gaia_list=[]
            gaia_crd_cone = SkyCoord(gaia3_all_cone.values_list('ra', flat=True), 
                                     gaia3_all_cone.values_list('dec', flat=True), 
                                     frame="fk5", unit="deg")
            sep_cone = crd.separation(gaia_crd_cone).arcsecond
            isel = (sep_cone <= minrad)
        
            index=0
            for gaia in gaia3_all_cone:
                #gaia_crd = SkyCoord(gaia.ra, gaia.dec, frame="fk5", unit="deg")
                #sep=crd.separation(gaia_crd).arcsecond
                if(isel[index]):
                    gaia_list.append(gaia)
                index=index+1
            
            if(gaia_list):
                src.gaia3.add(*gaia_list)
        
        elapsed = time.time() - start_time0
        estimated = (elapsed*(pages_num-i))/60/60
        print("Page {}/{} for {:.2f} sec, est. {:.2f} hours".format(i,pages_num,elapsed,estimated))
        

    
    

    #load_simbad_for_skymap_sources(srcs,minrad=minrad,maxdist=maxdist)
    #load_vizier_allwise(srcs, minrad=minrad, maxdist=maxdist)
    #load_vizier_nvss(srcs, minrad=minrad, maxdist=maxdist)
    #load_vizier_first(srcs, minrad=minrad, maxdist=maxdist)
    #load_vizier_sumss(srcs, minrad=minrad, maxdist=maxdist)        
    #find_counterparts(srcs, GAIADR3.objects.all(), "gaia3", maxdist=maxdist, minrad=minrad)
    #for src in srcs:
    #    src.heasarc.clear()
    
    #find_counterparts(srcs, GAIADR3.objects.all(),'gaia3', maxdist=maxdist, minrad=minrad)
    #find_counterparts(srcs, GLIMPSE.objects.all(),'glimpse', maxdist=maxdist, minrad=minrad)

class Command(BaseCommand):
    help = 'Initiates data dase'

    def handle(self, *args, **options):

        start_time = time.time()
        
        do_match()
        hours = (time.time() - start_time)/60/60
        print("--- {:.2f} hours ---".format(hours))

        self.stdout.write(self.style.SUCCESS('Done'))