from django.core.management.base import BaseCommand, CommandError
from django.db.models import Q

from datetime import date
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

from heasarc.models import HeasarcBase
from artsurvey.models import ArtSurveySource, ArtSurvey, EnergyBand, Catalog, MetaSource
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


def get_survey(version):
    try:
        survey = ArtSurvey.objects.get(version=version)
        print("use ArtSurvey {}".format(survey))
        return survey
    except ArtSurvey.DoesNotExist:
        print("ArtSurvey {} not found".format(version))
        return None

def get_band(slug):
    try:
        band = EnergyBand.objects.get(slug=slug)
        print("use EnergyBand {}".format(band))
        return band
    except EnergyBand.DoesNotExist:
        print("EnergyBand not found")
        return None

def select(catalog,survey,band,threshold):
    srcs = ArtSurveySource.objects.all().filter(survey=survey).filter(band=band).filter(Q(sig__gte=threshold) | Q(ext__gte=1))

    for src in srcs:
        src.catalog.add(catalog)
        ms = src.metasource
        ms.catalog=catalog
        ms.save()
        src.save()
    count=srcs.count()
    print("{} {}, sign. threshold {}, selected {}".format(survey,band.slug,threshold,count))

def select_catalog():
    minrad=40
    maxdist=120
        
    s1s2=get_survey(12.5)
    s1=get_survey(1.7)
    s2=get_survey(2.7)

    e0=get_band('E0')
    e1=get_band('E1')
    e2=get_band('E2')
    e3=get_band('E3')
    
    cat=Catalog.objects.get(year=2021)

    """
    RUN METASOURCE FIRST
    """

    """ clean links in MetaSource """
    msrcs = MetaSource.objects.all()
    for ms in msrcs:
        ms.catalog = None
        ms.save()

    srcs = ArtSurveySource.objects.all()
    for src in srcs:
        src.catalog.clear()
    
    select(cat,s1s2,e0,4.818)

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):
        select_catalog()
        self.stdout.write(self.style.SUCCESS('Done'))