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

import sys
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 django.db.models import Q

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


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,log_nfalse,threshold):
    srcs = ArtSurveySource.objects.all().filter(survey=survey).filter(band=band).filter(Q(log_nfalse__lte=log_nfalse) | 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("{} {}, log_nfalse {}, selected {}, {:.1f}% false count: {:.1f}".format(survey,band.slug,log_nfalse,count,threshold*100,count*threshold))

def select_ml(catalog,survey,band,log_ml_nfalse, threshold):
    srcs = ArtSurveySource.objects.all().filter(survey=survey).filter(band=band).filter(log_ml_nfalse__lte=log_ml_nfalse) # no extended sources in 2023

    bad_category = Category.objects.get(slug='false')

    count=0
    for src in srcs:
        if(src.category == bad_category):
            print(src," --bad")
            continue
        print(src)
        src.catalog.add(catalog)
        if not (src.metasource):
            raise ValueError("{} has no MetaSource".format(src))
        ms = src.metasource
        ms.catalog=catalog
        ms.save()
        src.save()
        count=count+1
    
    #count=srcs.count()
    print("{} {}, log_ml_nfalse {}, selected {}, {:.1f}% false count: {:.1f}".format(survey,band.slug,log_ml_nfalse,count,threshold*100,count*threshold))

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

    #s15=get_survey(12345.12)
    s15=get_survey(12345.14)

    #e0=get_band('E0')
    #e1=get_band('E1')
    #e2=get_band('E2')
    #e3=get_band('E3')
    #e4=get_band('E4')
    e5=get_band('E5')

    cat1=Catalog.objects.get(year=2021)
    cat2=Catalog.objects.get(year=2023)

    """
    RUN METASOURCE FIRST
    """

    """ clean links in MetaSource """
    msrcs = MetaSource.objects.all()
    for ms in msrcs:
        """ clean Pape II only, if exists """
        if(ms.catalog == cat2):
            print("{} Clean catalog {}".format(ms,cat2))
            ms.catalog = None
            ms.save()

    

    srcs = ArtSurveySource.objects.all()
    
    for src in srcs:
        for cat in src.catalog.all():
            if(cat == cat1):
                print("{} {} -- keep".format(src,cat))
                continue

            if(cat == cat2):
                print("{} {} -- clean".format(src,cat))
                #continue
                src.catalog.remove(cat2)
                src.save()
                continue
    

    #sys.exit()
    
    
    #select_ml(cat2,s15,e5,-3.1222,0.02)# ML appeared in second catalog
    
    select_ml(cat2,s15,e5,-3.1264,0.02)# ML appeared in second catalog

    sys.exit()

    #select(cat,s1s2,e0,-2.6751,0.10)
    #select(cat,s1s2,e1,-3.2786)
    #select(cat,s1s2,e2,-3.3014)
    #select(cat,s1s2,e3,-3.9955)

    # false 1%
    #select(cat,s1s2,e1,-4.0384,0.01)
    #select(cat,s1s2,e2,-4.0799,0.01)
    #select(cat,s1s2,e3,-4.5918,0.01)

    # false 5%
    #select(cat,s1,e0,-3.2900)
    #select(cat,s1,e1,-3.4620 )
    #select(cat,s1,e2,-3.5190 )
    #select(cat,s1,e3,-4.0747 )

    # false 5%
    #select(cat,s2,e0,-3.3065)
    #select(cat,s2,e1,-3.5003 )
    #select(cat,s2,e2,-3.5262 )
    #select(cat,s2,e3,-4.1576 )

    # false 1%
    #select(cat,s1,e0,-4.0519,0.01)

    # false 1%
    #select(cat,s2,e0,-4.0747,0.01)

    
    
    

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