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
from datetime import datetime, timedelta




def current_time():
    return (datetime.now() + timedelta(hours=3)).strftime("%H:%M:%S")



def update_catalog_file_status(catalog_file, status):              #catalog file status updater
    catalog_file.status = status                                   
    catalog_file.save()



def healpix(nside): #NSIDE 2048, ORDER ring or nested
    
    for catalog_file in CatalogFile.objects.filter(status='INGESTED'):


        print(f'[{current_time()}] Loading sources from {catalog_file.name}...')
        sources = list(catalog_file.sources.all())


        if catalog_file.status == 'INDEX_IN_PROGRESS':
            #sources.update(healpix_ring_index=None, healpix_nested_index=None)
            GaiaSource.objects.filter(catalog_file=catalog_file).update(
                healpix_ring_index=None, 
                healpix_nested_index=None
            )


        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'
        )


        ring_healpix = ah.HEALPix(
            nside=nside, 
            order=ring, 
            frame='fk5'
        )
        ring_healpix_indices = healpix.skycoord_to_healpix(skycoord)
        
        
        nested_healpix = ah.HEALPix(
            nside=nside, 
            order=nested, 
            frame='fk5'
        )
        nested_healpix_indices = healpix.skycoord_to_healpix(skycoord)


        for source, healpix_index in zip(sources, ring_healpix_indices):
            source.healpix_ring_index = healpix_index

        for source, healpix_index in zip(sources, nested_healpix_indices):
            source.healpix_nested_index = healpix_index


        update_catalog_file_status(catalog_file, 'INDEX_IN_PROGRESS')

        GaiaSource.objects.bulk_update(sources, ['healpix_ring_index', 'healpix_nested_index'])

        update_catalog_file_status(catalog_file, 'INDEXED')


        print(f'[{current_time()}] Database updated, sources indexed successfully.')




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


    def add_arguments(self, parser):
        parser.add_argument(
            'nside', 
            type=int, 
            help='NSIDE parameter for HEALPix'
        )


    def handle(self, *args, **options):
        nside = options['nside']

        healpix(nside)