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

from datetime import date
from heasarc.models import INPUT_DATA_DIR
import datetime
from django.utils import timezone
import astropy
from astropy.io import ascii
import pandas as pd
import pymysql
from sqlalchemy import create_engine
import numpy.ma as ma

from astrobasis.models import VLASS

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 NSIDE_SOURCES, ORDER

import os
import glob
import logging
import time

class Command(BaseCommand):
    help = 'Initiates healpix for Gaia'

    def handle(self, *args, **options):
        start_time = time.time()
        hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5)

        
        #srcs = VLASS.objects.all().filter(healpix__gt=0)[:50]
        #for src in srcs:
        #    crd = SkyCoord(src.ra, src.dec, frame="fk5", unit="deg")
        #    heal = hp.skycoord_to_healpix(crd)
        #    print(src.ra,src.dec,src.healpix,heal)
        #return
        

        srcs = VLASS.objects.all()#.filter(healpix__exact=0)
        paginator = Paginator(srcs, 50000)
        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
            VLASS.objects.bulk_update(updates, ["healpix","glon","glat"], batch_size=200)
            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))