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

from datetime import date
import datetime, os, sys
from django.utils import timezone
import astropy
from astropy.io import ascii
import pandas as pd
import pymysql
from sqlalchemy import create_engine
from math import isfinite, sqrt

from heasarc.models import HeasarcBase
from heasarc.models import Heasarc4XMMDR12

from artsurvey.models import ArtSurveySource, ArtSurvey, EnergyBand, Category
from srglib.utils import load_vizier_nvss
from srglib.utils import load_vizier_first
from srglib.utils import load_vizier_sumss
from astrobasis.models import GAIADR3
from srglib.utils import find_counterparts
from srglib.utils import load_vizier_allwise
from srglib.utils import load_simbad_for_skymap_sources
from artsurvey.utils import find_heasarc

from heasarc.models import HeasarcBase
from heasarc.models import HeasarcSwiftBAT105m
from heasarc.models import HeasarcIntRefCat
from heasarc.models import HeasarcXTEASSCAT
from heasarc.models import HeasarcXMMSL2
from heasarc.models import Heasarc4FGL
from heasarc.models import HeasarcIntegral2020  
from heasarc.models import HeasarcMAXI7YR
from heasarc.models import HeasarcRASS2RXS

def do_heasarc(survey_version,slug,log_ml_nfalse=None):
    bad_category = Category.objects.get(slug='false')
    minrad=40
    maxdist=600
    try:
        band = EnergyBand.objects.get(slug=slug)
    except EnergyBand.DoesNotExist:
        print("EnergyBand bot found")
        return


    try:
        survey = ArtSurvey.objects.get(version=survey_version)
        print("use ArtSurvey {}".format(survey))
    except ArtSurvey.DoesNotExist:
        print("ArtSurvey {} not found".format(survey_version))
        return
        
    srcs = survey.artsurveysource_set.all().filter(band=band).filter(log_ml_nfalse__lte=log_ml_nfalse).exclude(category__exact=bad_category)
    print("Selected {}".format(srcs.count()))

    # normal run
    find_heasarc(srcs, maxdist=maxdist, minrad=minrad, offset=False, syserr=10.0)
    return


    # make offset

    """
    # script to read this file:
    import pandas as pd
    df=pd.read_csv('offset.csv')  

    for k in df.keys():
        print("{} {:.1f} +/- {:.1f} std={:.1f}".format(k,df[k].mean(),df[k].sem(),df[k].std()))
    """

    fout='offset.csv'
    try:
        os.remove(fout)
    except OSError:
        pass
    with open(fout, "a") as myfile:
        myfile.write("RASS2RXS,XMMSL2,MAXI7YR,SwiftBAT105m,Integral2020,4FGL\n")

    for x in range(100):
        find_heasarc(srcs, maxdist=maxdist, minrad=minrad, offset=True, syserr=10.0)

        SwiftBAT105m=0
        XMMSL2=0
        FGL=0
        Integral2020=0
        MAXI7YR=0
        RASS2RXS=0
        
        for src in srcs:
            heasarcs = src.heasarc.all()
            if not (heasarcs.count()):
                continue
            if (heasarcs.filter(instance_of=HeasarcSwiftBAT105m).count()):
                SwiftBAT105m+=1
            if (heasarcs.filter(instance_of=HeasarcXMMSL2).count()):
                XMMSL2+=1
            if (heasarcs.filter(instance_of=Heasarc4FGL).count()):
                FGL+=1
            if (heasarcs.filter(instance_of=HeasarcIntegral2020).count()):
                Integral2020+=1
            if (heasarcs.filter(instance_of=HeasarcMAXI7YR).count()):
                MAXI7YR+=1
            if (heasarcs.filter(instance_of=HeasarcRASS2RXS).count()):
                RASS2RXS+=1
        with open(fout, "a") as myfile:
            myfile.write("{}, {}, {}, {}, {}, {}\n".format(RASS2RXS,XMMSL2,MAXI7YR,SwiftBAT105m,Integral2020,FGL))

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):

        #all=HeasarcBase.objects.all()
        #print(all.count())

        #do_heasarc(12345.12,'E5',log_ml_nfalse=-3.1222)
        #do_heasarc(12345.13,'E5',log_ml_nfalse=-3.1264)
        do_heasarc(12345.14,'E5',log_ml_nfalse=-3.1264)

        """
        srcs = ArtSurveySource.objects.all().filter(heasarc__heasarc4xmmdr12__isnull=False)[:10]
        print("Selected {} in total".format(srcs.count()))
        for src in srcs:
            print(src.get_absolute_url())
        """
        
        self.stdout.write(self.style.SUCCESS('Done'))