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 TwoMASS

from srglib.utils import notnan
from srglib.utils import boolean_notnan

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 = 500

    print("Delete all 2MASS")
    if(reset):
        srcs = TwoMASS.objects.all()
        srcs.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)
    """

 
    objs = (TwoMASS(name=data[i]['2MASS'],
                    ra=float(data[i]['RAJ2000']),
                    dec=float(data[i]['DEJ2000']),
                    jmag = notnan(data[i]['Jmag']),
                    e_jmag = notnan(data[i]['e_Jmag']),
                    hmag = notnan(data[i]['Hmag']),
                    e_hmag = notnan(data[i]['e_Hmag']),
                    kmag = notnan(data[i]['Kmag']),
                    e_kmag = notnan(data[i]['e_Kmag']),
                    qkfg = data[i]['Qflg'],
                    rkfg = data[i]['Rflg'],
                    bkfg = data[i]['Bflg'],
                    ckfg = data[i]['Cflg'],
                    prox = notnan(data[i]['prox']),
                    xflg = boolean_notnan(data[i]['Xflg']),
                    aflg = boolean_notnan(data[i]['Aflg']),)
            for i in range(len(data)))
    TwoMASS.objects.bulk_create(objs,batch_size)

    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/2MASS/l_0_180.tsv',reset=True)
        print("--- %s seconds ---" % (time.time() - start_time))

        start_time = time.time()
        bulk_load_heasarc_table('/data/2MASS/l_180_360.tsv')
        print("--- %s seconds ---" % (time.time() - start_time))

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