import srglib.correlate_utils as correlate_utils
from heasarc.models import HeasarcTable, TableColumn, HeasarcObjectClass, NSIDE_SOURCES, ORDER, \
    HeasarcXrayMaster, HeasarcXMMSSC, HeasarcBase, HeasarcCSC, HeasarcRASS2RXS, Heasarc2SXPS
from django.contrib.contenttypes.models import ContentType
import numpy as np

polctypes = {cat: ContentType.objects.get_for_model(model) for cat, model in
             zip(["rxs2", "sxps2", "3xmm", "csc2"], [HeasarcRASS2RXS, Heasarc2SXPS, HeasarcXMMSSC, HeasarcCSC])}

def check_rass_transient(erositasource):
    """
    for provided erotranssource checks wether it has rosat counterparts in the RXS2 catalogue,
    it it has, checks, whether the source is transient relative to all counterparts, if it does not,
    check wether the source is expected to be 10 times brighter then detection limit for rass exposure

    as a result of the check switches roast_transient field of the source

    Parameters:
    --------
        erositasource: django.queryset
            a queryset of erosita sources
    """
    correlate_utils.load_rassemap()
    rasssources = erositasource.heasarc.filter(polymorphic_ctype=polctypes["rxs2"])
    if rasssources.count():
        rassfluxes = np.array([rass.count_rate for rass in rasssources])
        if np.any(erositasource.ML_FLUX_0 < rassfluxes*correlate_utils.FSCALES["2rxs"]/10.) or \
                np.any(erositasource.ML_FLUX_0 > rassfluxes*correlate_utils.FSCALES["2rxs"]*10.):
            erositasource.rosat_transient = True
            erositasource.has_rosat=True
    else:
        erositasource.rosat_transient = correlate_utils.check_rss_cts_individual(erositasource)

def check_sxps_transients(erositasource):
    """
    for provided erotranssource checks wether it has swift counterparts in the sxps2 catalog,
    it it has, checks, whether the source is transient relative to all counterparts

    as a result of the check switches transient field of the source to False, if the erosita not a transient to at least one counterpart

    Parameters:
    --------
        erositasource: django.queryset
            a queryset of erosita sources
    """
    sxpscounterparts = erositasource.heasarc.filter(polymorphic_ctype=polctypes["sxps2"])
    if sxpscounterparts.exists():
        flux = np.array([s.PeakRate_band1 + s.PeakRate_band2 for s in sxpscounterparts])
        transient = np.all(flux*10*correlate_utils.FSCALES["sxps2"] < erositasource.ML_FLUX_0)
        print("sxps flux", flux, "eflux", erositasource.ML_FLUX_0*0.9e-12, "sxps transient", transient)
        erositasource.transient = erositasource.transient & transient

def check_xmm3_transients(erositasource):
    """
    for provided erotranssource checks wether it has XMM-Newton counterparts in the 3xmm catalog,
    it it has, checks, whether the source is transient relative to all counterparts

    as a result of the check switches transient field of the source to False, if the erosita not a transient to at least one counterpart

    Parameters:
    --------
        erositasource: django.queryset
            a queryset of erosita sources
    """
    xmmcounterparts = erositasource.heasarc.filter(polymorphic_ctype=polctypes["3xmm"])
    if xmmcounterparts.exists():
        flux = np.array([s.ep_1_flux + s.ep_3_flux + s.ep_2_flux  for s in xmmcounterparts])
        transient = np.all(flux*10*correlate_utils.FSCALES["3XMM"] < erositasource.ML_FLUX_0)
        print("3xmm flux", flux, "eflux", erositasource.ML_FLUX_0*0.9e-12,  "3xmm transient", transient)
        erositasource.transient = erositasource.transient & transient

def check_csc_transients(erositasource):
    """
    for provided erotranssource checks wether it has chandra counterparts in the csc2 catalog,
    it it has, checks, whether the source is transient relative to all counterparts

    as a result of the check switches transient field of the source to False, if the erosita not a transient to at least one counterpart

    Parameters:
    --------
        erositasource: django.queryset
            a queryset of erosita sources
    """
    chancounterparts = erositasource.heasarc.filter(polymorphic_ctype=polctypes["csc2"])
    if chancounterparts.exists():
        flux = np.array([s.s_photflux_ap + s.m_photflux_ap  for s in chancounterparts])
        transient = np.all(flux*10*correlate_utils.FSCALES["csc2"] < erositasource.ML_FLUX_0)
        print("csc2 flux", flux, "eflux", erositasource.ML_FLUX_0*0.9e-12, "csc2 transient", transient)
        erositasource.transient = erositasource.transient & transient


search_transients = {"sxps2": check_sxps_transients, "xmm3": check_xmm3_transients, "csc2": check_csc_transients}