from django.core.management.base import BaseCommand
from GaiaDBInterface.models import GaiaSource, CatalogFile
from astropy.coordinates import SkyCoord
import astropy_healpix as ah
import numpy as np


def healpix():

    NSIDE=2048
    ORDER='ring'
    
    for catalog_file in CatalogFile.objects.all():

        sources = list(catalog_file.sources.all())
        
        ra_list = np.array([source.ra for source in sources])
        dec_list = np.array([source.dec for source in sources])
        
        skycoord = SkyCoord(ra=ra_list, dec=dec_list, unit='deg', frame="fk5")
        
        healpix = ah.HEALPix(nside=NSIDE, order=ORDER)
        healpix_indices = healpix.skycoord_to_healpix(skycoord)
        
        for source, healpix_index in zip(sources, healpix_indices):
            source.healpix_ring_index = healpix_index
        
        GaiaSource.objects.bulk_update(sources, ['healpix_ring_index'])


class Command(BaseCommand):
    help = 'Index sources using healpix.'

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

        healpix()