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
import time

from django.core.paginator import Paginator

from heasarc.models import NSIDE_SOURCES, ORDER

def astrobasis_assign_healpix(catalogname,perpage=50000,batch_size=200):
        start_time = time.time()

        hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5)

        srcs = catalogname.objects.all()#.filter(healpix__exact=0)

        paginator = Paginator(srcs, perpage)
        for page_idx in range(1, paginator.num_pages):
            start_time0 = time.time()
            ra=paginator.page(page_idx).object_list.values_list('ra', flat=True)
            dec=paginator.page(page_idx).object_list.values_list('dec', flat=True)
            crd = SkyCoord(ra, dec, frame="fk5", unit="deg")
            lon = crd.galactic.l.value
            lat = crd.galactic.b.value
            heal = hp.skycoord_to_healpix(crd)
            index=0
            updates=[]
            for src in paginator.page(page_idx).object_list:
                # here you can do what you want with the row
                #crd = SkyCoord(src.ra, src.dec, frame="fk5", unit="deg")
                #lon = crd.galactic.l.value
                #lat = crd.galactic.b.value
                #heal = hp.skycoord_to_healpix(crd)
                src.healpix=heal[index]
                src.glon=lon[index]
                src.glat=lat[index]
                #src.save()
                updates.append(src)
                #print(src.ra,src.dec,heal)
                index=index+1
            catalogname.objects.bulk_update(updates, ["healpix","glon","glat"], batch_size=batch_size)
            elapsed = time.time() - start_time0
            estimated = (elapsed*(paginator.num_pages-page_idx))/60/60
            print("Page {}/{} for {:.2f} sec, est. {:.2f} hours".format(page_idx,paginator.num_pages,elapsed,estimated))

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