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

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 GaiaSourceFileDR3, GAIADR3

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

from multiprocessing import Pool, cpu_count


USE_NCORES=1

def myfunc(n):
    print(n)
    
    batch_size=500
    hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5)
    
    plates=GaiaSourceFileDR3.objects.filter(status__exact="ready")
    
    total_ready = len(plates)
    print('Total <ready> Gaia files: '+str(total_ready))
    
    step=int(total_ready/USE_NCORES)+1
    idx1 = step*n
    idx2 = step*(n+1)
    if idx2 > total_ready:
        idx2=total_ready
    for i in range(idx1,idx2):
        print(str(n)+": "+plates[i].filename)
        phots=plates[i].gaiadr3_set.exclude(healpix__gt=0.0)
        ra=phots.values_list('ra', flat=True)
        dec=phots.values_list('dec', flat=True)
        crd = SkyCoord(ra, dec, frame="fk5", unit="deg")
        heal = hp.skycoord_to_healpix(crd)
        for j, phot in enumerate(phots):
            phot.healpix=heal[j]
        GAIADR3.objects.bulk_update(phots, ['healpix'], batch_size=batch_size)

        plates[i].status='healpix'
        plates[i].save()

    pass

def run_healpix():
    
    start_time = time.time()

    batch_size=500
    hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5)
    
    plates=GaiaSourceFileDR3.objects.filter(status__exact="ready")
    
    total_ready = len(plates)
    
    print('Total <ready> Gaia files: '+str(total_ready))
    
    for plate in plates:
        print(plate.filename)
        phots=plate.gaiadr3_set.exclude(healpix__gt=0.0)
        ra=phots.values_list('ra', flat=True)
        dec=phots.values_list('dec', flat=True)
        crd = SkyCoord(ra, dec, frame="fk5", unit="deg")
        heal = hp.skycoord_to_healpix(crd)
        for j, phot in enumerate(phots):
            phot.healpix=heal[j]
        GAIADR3.objects.bulk_update(phots, ['healpix'], batch_size=batch_size)

        plate.status='healpix'
        plate.save()

    print("--- %s seconds ---" % (time.time() - start_time))
    pass

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

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

        run_healpix()
        
        """
        does not work for multitheading for some reason

        with Pool(10) as p:
            p.map(myfunc, range(USE_NCORES))
        """