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())
        print(f'[{current_time()}] Sources ready. Forming ra & dec arrays...')

        if catalog_file.status == 'INDEX_IN_PROGRESS':
            sources.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])
        print(f'[{current_time()}] ra & dec arrays ready. Creating SkyCoord objects...')

        skycoord = SkyCoord(ra=ra_list, dec=dec_list, unit='deg', frame='fk5')
        print(f'[{current_time()}] SkyCoord objects ready. Calculating healpix indices...')



        ring_healpix = ah.HEALPix(nside=nside, order=ring, frame='fk5')
        ring_healpix_indices = healpix.skycoord_to_healpix(skycoord)
        print(f'[{current_time()}] Ring indices ready. Assigning indices to source instances...')

        for source, healpix_index in zip(sources, ring_healpix_indices):
            source.healpix_ring_index = healpix_index
        print(f'[{current_time()}] Instances ready.')



        nested_healpix = ah.HEALPix(nside=nside, order=nested, frame='fk5')
        nested_healpix_indices = healpix.skycoord_to_healpix(skycoord)
        print(f'[{current_time()}] Nested indices ready. Assigning indices to source instances...')

        for source, healpix_index in zip(sources, healpix_indices):
            source.healpix_nested_index = healpix_index
        print(f'[{current_time()}] Instances ready. Bulk updating the database...')


        update_catalog_file_status(catalog_file, 'INDEX_IN_PROGRESS')


        GaiaSource.objects.bulk_update(sources, ['healpix_ring_index', 'healpix_nested_index'])
        print(f'[{current_time()}] Database updated, sources nested-indexed successfully.')
        update_catalog_file_status(catalog_file, 'INDEXED')


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)