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

    ingested_files = CatalogFile.objects.filter(status='INGESTED')
    total_files = ingested_files.count()
    print(f'[{current_time()}] There are {total_files} ingested files.')
    current_file_index = 0

    for catalog_file in ingested_files:
        
        current_file_index += 1
        
        print(f'[{current_time()}] Processing sources from {catalog_file.name}')
        sources = list(catalog_file.sources.all())
        print(f'[{current_time()}] Sources loaded, calculating indices.')

        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 = ring_healpix.skycoord_to_healpix(skycoord)
        

        nested_healpix = ah.HEALPix(
            nside=nside, 
            order='nested', 
            frame='fk5'
        )
        nested_healpix_indices = nested_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

        print(f'[{current_time()}] Calculations done, updating database.')

        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.')
        print(f'[{current_time()}] Progress: {current_file_index}/{total_files} files 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)