from django.core.management.base import BaseCommand, CommandError
import time
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 heasarc.tdat import tDat
from heasarc.models import HeasarcTable, TableColumn, HeasarcObjectClass
from heasarc.models import NSIDE_SOURCES, ORDER 
from astrobasis.models import GLIMPSE

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

def bulk_load_heasarc_table(filename,reset=False):
    batch_size = 300

    print("Delete all GLIMPSE")
    if(reset):
        tables = GLIMPSE.objects.all()
        tables.delete()

    print("Loading {}".format(filename))
    data = astropy.table.Table.read(filename, format='ascii.csv',delimiter=';',comment='#',encoding='latin1')
    print(data.info)
    print("Bulk load to DB")

    """
    cols = data.columns
    for key, value in cols.items():
        print(key, '->', value)
    """

    start_time = time.time()
 
    objs = (GLIMPSE(name=data[i]['GLIMPSE'],
                    ra=float(data[i]['RAJ2000']),
                    dec=float(data[i]['DEJ2000']),
                    twomass = data[i]['2MASS'],
                    jmag = data[i]['Jmag'],
                    hmag = data[i]['Hmag'],
                    kmag = data[i]['Kmag'],
                    b1mag = data[i]['3.6mag'],
                    e_b1mag = data[i]['e_3.6mag'],
                    b2mag = data[i]['4.5mag'],
                    e_b2mag = data[i]['e_4.5mag'],
                    b3mag = data[i]['5.8mag'],
                    e_b3mag = data[i]['e_5.8mag'],
                    b4mag = data[i]['8.0mag'],
                    e_b4mag = data[i]['e_8.0mag'],)
            for i in range(len(data)))
    GLIMPSE.objects.bulk_create(objs,batch_size)
    print("--- %s seconds ---" % (time.time() - start_time))

    print('--> Successfully loaded "%s"' % filename)
    pass

def bulk_healpix_update_tooslow():
    glimpse_all = GLIMPSE.objects.all()
    hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5())
    objs = []
    start_time = time.time()
    for obj in glimpse_all:
        crd = SkyCoord(obj.ra, obj.dec, frame=FK5(), unit="deg")
        lon = crd.galactic.l.value
        lat = crd.galactic.b.value
        healpix = hp.skycoord_to_healpix(crd)
        obj.healpix = healpix
        obj.glon = lon
        obj.glat = lat
        objs.append(obj)
    GLIMPSE.objects.bulk_update(objs, ['healpix','glon','glat'], batch_size=1000)
    print("--- %s seconds ---" % (time.time() - start_time))
        
def load_heasarc_table(filename,reset=False):

    if(reset):
        tables = GLIMPSE.objects.all()
        tables.delete()

    return

    data = astropy.table.Table.read(filename, format='ascii.csv',delimiter=';',comment='#',encoding='latin1')
    print(data.info)
    
    """
    cols = data.columns
    for key, value in cols.items():
        print(key, '->', value)
    """

    hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5())
    for item in data:
        
        healpix = 0
        ra = 0.0
        dec = 0.0
        lon = 0.0
        lat = 0.0
        if not (ma.is_masked(item['RAJ2000']) or ma.is_masked(item['DEJ2000'])):
            ra = float(item['DEJ2000'])
            dec = float(item['DEJ2000'])
            crd = SkyCoord(ra, dec, frame=FK5(), unit="deg")
            lon = crd.galactic.l.value
            lat = crd.galactic.b.value
            healpix = hp.skycoord_to_healpix(crd)
        else:
            continue
        

        obj = GLIMPSE.objects.create(healpix=healpix, 
                                            ra=ra,
                                            dec=dec,
                                            glon = lon,
                                            glat = lat,
                                            name=item['GLIMPSE'])
        twomass = item['2MASS']
        jmag = item['Jmag']
        hmag = item['Hmag']
        kmag = item['Kmag']
        b1mag = item['3.6mag']
        e_b1mag = item['e_3.6mag']
        b2mag = item['4.5mag']
        e_b2mag = item['e_4.5mag']
        b3mag = item['5.8mag']
        e_b3mag = item['e_5.8mag']
        b4mag = item['8.0mag']
        e_b4mag = item['e_8.0mag']
        obj.save()
#C;GLIMPSE;2MASS;RAJ2000;DEJ2000;Jmag;Hmag;Kmag;3.6mag;e_3.6mag;4.5mag;e_4.5mag;5.8mag;e_5.8mag;8.0mag;e_8.0mag

    print('--> Successfully loaded "%s"' % filename)
    pass


class Command(BaseCommand):
    help = 'Initiates data dase'

#    def add_arguments(self, parser):
#        parser.add_argument('poll_id', nargs='+', type=int)

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

        start_time = time.time()
        bulk_load_heasarc_table('/data/GLIMPSE/glimpse_spitzer.tsv',reset=True)
        print("--- %s seconds ---" % (time.time() - start_time))
        # Model.objects.filter(amount__isnull=True).update(amount=F('pre_tax') * 1.2)


        self.stdout.write(self.style.SUCCESS('Done'))