import math
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional

import astropy.units as astropy_units
import numba as nb
import numpy as np
from astropy.io import fits
from interpolation import interp

from astromodels.functions.function import Function1D, FunctionMeta
from astromodels.utils import _get_data_file_path
from astromodels.utils.configuration import astromodels_config
from astromodels.utils.logging import setup_logger

log = setup_logger(__name__)

try:

    # ebltable is a Python packages to read in and interpolate tables for the photon density of
    # the Extragalactic Background Light (EBL) and the resulting opacity for high energy gamma
    # rays.

    import ebltable.tau_from_model as ebltau

    has_ebltable = True

except ImportError:

    if astromodels_config.logging.startup_warnings:

        log.warning(
            "The ebltable package is not available. Models that depend on it will not be available"
        )

    has_ebltable = False


class EBLTableNotAvailable(ImportWarning):
    pass


class InvalidUsageForFunction(Exception):
    pass


@dataclass(frozen=False)
class AbundanceTable:
    name: str
    tables: List[str]
    _current_table: str

    def set_table(self, table: str):

        """
        set the current table from AG89, WILM or ASPL

        :param table: 
        :type table: str
        :returns: 

        """
        
        
        
        old_table = self._current_table

        self._current_table = table
        
        
        if self.current_table not in self.tables:

            log.error(f"{self.name} does not contain {table} choose {','.join(self.table)}")

            self._current_table = old_table
            
            raise AssertionError()
            

    @property
    def current_table(self) -> str:

        convert= {
        "AG89": "angr",
        "ASPL": "aspl",
        "WILM": "wilm"
        }

        return convert[self._current_table]
        
    @property
    def info(self) -> str:

        _abund_info = {}
        _abund_info["WILM"] = "wilms\nfrom Wilms, Allen & McCray (2000), ApJ 542, 914 \n except for elements not listed which are given zero abundance)\n https://heasarc.nasa.gov/xanadu/xspec/manual/XSabund.html "
        _abund_info["AG89"] = "angr\nfrom Anders E. & Grevesse N. (1989, Geochimica et Cosmochimica Acta 53, 197)\n https://heasarc.nasa.gov/xanadu/xspec/manual/XSabund.html"
        _abund_info["ASPL"] = "aspl\nfrom Asplund M., Grevesse N., Sauval A.J. & Scott P. (2009, ARAA, 47, 481)\nhttps://heasarc.nasa.gov/xanadu/xspec/manual/XSabund.html"

        return _abund_info[self._current_table]

    @property
    def xsect_table(self) -> np.ndarray:

        """
        returns the XSECT table for the current model

        :returns: 

        """
        _path: Path = Path(
        "xsect") / f"xsect_{self.name}_{self.current_table}.fits"

        path_to_xsect: Path = _get_data_file_path(_path)

        with fits.open(path_to_xsect) as fxs:

            dxs = fxs[1].data
            xsect_ene = dxs["ENERGY"]
            xsect_val = dxs["SIGMA"]

        return np.array(xsect_ene, dtype=np.float64), np.array(xsect_val,
                                                           dtype=np.float64)


        
phabs = AbundanceTable("phabs", ["angr","aspl"], "AG89")
tbabs = AbundanceTable("tbabs", ["angr","aspl","wilm"], "WILM")
wabs = AbundanceTable("wabs", ["angr"], "AG89")
    


# PhAbs class


class PhAbs(Function1D, metaclass=FunctionMeta):
    r"""
    description :
        Photometric absorption (phabs implementation), f(E) = exp(- NH * sigma(E))
        contributed by Dominique Eckert
    parameters :
        NH :
            desc : absorbing column density in units of 1e22 particles per cm^2
            initial value : 1.0
            is_normalization : False
            transformation : log10
            min : 1e-4
            max : 1e4
            delta : 0.1

        redshift :
            desc : the redshift of the source
            initial value : 0.
            is_normalization : False
            min : 0
            max : 15
            delta : 0.1
            fix: True

    """
    def _setup(self):
        self._fixed_units = (astropy_units.keV,
                             astropy_units.dimensionless_unscaled)
        self.init_xsect(astromodels_config.absorption_models.phabs_table.value)

    def _set_units(self, x_unit, y_unit):
        self.NH.unit = astropy_units.cm**(-2)
        self.redshift.unit = astropy_units.dimensionless_unscaled

    def init_xsect(self, abund_table="AG89"):
        """
        Set the abundance table

        :param abund_table: "ASPL", "AG89" 
        :returns: 
        :rtype: 

        """

        # load cross section data

        try:
            phabs.set_table(abund_table)

            log.debug(f"phabs model set to {abund_table}")
            
        except:

            log.info(f"defaulting to {astromodels_config.absorption_models.phabs_table.value}")

            phabs.set_table(astromodels_config.absorption_models.phabs_table.value)
            

        self.xsect_ene, self.xsect_val = phabs.xsect_table

    @property
    def abundance_table(self):
        print(phabs.info)

        
    def evaluate(self, x, NH, redshift):

        if isinstance(x, astropy_units.Quantity):

            _unit = astropy_units.cm**2
            _y_unit = astropy_units.dimensionless_unscaled
            _x = x.value
            _redshift = redshift.value
        else:

            _unit = 1.0
            _y_unit = 1.0
            _redshift = redshift
            _x = x

        xsect_interp = interp(self.xsect_ene, self.xsect_val,
                              _x * (1 + _redshift))

        # evaluate the exponential with numba
        
        spec = _numba_eval(NH, xsect_interp) * _y_unit
        
        return spec


# TbAbs class


class TbAbs(Function1D, metaclass=FunctionMeta):
    r"""
    description :
        Photometric absorption (Tbabs implementation), f(E) = exp(- NH * sigma(E))
        contributed by Dominique Eckert
    parameters :
        NH :
            desc : absorbing column density in units of 1e22 particles per cm^2
            initial value : 1.0
            is_normalization : True
            transformation : log10
            min : 1e-4
            max : 1e4
            delta : 0.1

        redshift :
            desc : the redshift of the source
            initial value : 0.
            is_normalization : False
            min : 0
            max : 15
            delta : 0.1
            fix: True

    """
    def _setup(self):

        self.init_xsect(astromodels_config.absorption_models.tbabs_table.value)

        self._fixed_units = (astropy_units.keV,
                             astropy_units.dimensionless_unscaled)

    def _set_units(self, x_unit, y_unit):

        self.NH.unit = astropy_units.cm**(-2)
        self.redshift.unit = astropy_units.dimensionless_unscaled

    def init_xsect(self, abund_table="WILM"):
        """
        Set the abundance table

        :param abund_table: "WILM", "ASPL", "AG89" 
        :returns: 
        :rtype: 

        """
        
        try:
            tbabs.set_table(abund_table)
            
            log.debug(f"tbabs model set to {abund_table}")
            
        except:

            log.info(f"defaulting to {astromodels_config.absorption_models.tbabs_table.value}")

            tbabs.set_table(astromodels_config.absorption_models.tbabs_table.value)
            

        self.xsect_ene, self.xsect_val = tbabs.xsect_table

    @property
    def abundance_table(self):
        print(tbabs.info)

    def evaluate(self, x, NH, redshift):

        if isinstance(x, astropy_units.Quantity):

            _unit = astropy_units.cm**2
            _y_unit = astropy_units.dimensionless_unscaled
            _x = x.value
            _redshift = redshift.value

        else:

            _unit = 1.0
            _y_unit = 1.0
            _redshift = redshift
            _x = x

        xsect_interp = interp(self.xsect_ene, self.xsect_val,
                              _x * (1 + _redshift))

        spec = _numba_eval(NH, xsect_interp) * _y_unit

        return spec


# WAbs class


class WAbs(Function1D, metaclass=FunctionMeta):
    r"""
    description :
        Photometric absorption (Wabs implementation), f(E) = exp(- NH * sigma(E))
        contributed by Dominique Eckert
    parameters :
        NH :
            desc : absorbing column density in units of 1e22 particles per cm^2
            initial value : 1.0
            is_normalization : True
            transformation : log10
            min : 1e-4
            max : 1e4
            delta : 0.1
        redshift :
            desc : the redshift of the source
            initial value : 0.
            is_normalization : False
            min : 0
            max : 15
            delta : 0.1
            fix: True

    """

    def _setup(self):
        self._fixed_units = (astropy_units.keV,
                             astropy_units.dimensionless_unscaled)
        self.init_xsect()

    def _set_units(self, x_unit, y_unit):
        self.NH.unit = astropy_units.cm**(-2)
        self.redshift.unit = astropy_units.dimensionless_unscaled

    def init_xsect(self):
        """
        Set the abundance table

        :returns:
        :rtype:

        """            

        self.xsect_ene, self.xsect_val = wabs.xsect_table


    @property
    def abundance_table(self):
        print(wabs.info)


    def evaluate(self, x, NH, redshift):

        if isinstance(x, astropy_units.Quantity):

            _unit = astropy_units.cm**2
            _y_unit = astropy_units.dimensionless_unscaled
            _x = x.value
            _redshift = redshift.value
        else:

            _unit = 1.0
            _y_unit = 1.0
            _redshift = redshift
            _x = x

        xsect_interp = interp(self.xsect_ene, self.xsect_val,
                              _x * (1 + _redshift))

        xsect_interp = interp( self.xsect_ene, self.xsect_val, _x * (1+ _redshift))

        spec = _numba_eval(NH,xsect_interp) * _y_unit

        return spec



if has_ebltable:

    class EBLattenuation(Function1D, metaclass=FunctionMeta):
        r"""
        description :
            Attenuation factor for absorption in the extragalactic background light (EBL) ,
            to be used for extragalactic source spectra. Based on package "ebltable" by
            Manuel Meyer, https://github.com/me-manu/ebltable .

        latex: not available

        parameters :

          redshift :
                desc : redshift of the source
                initial value : 1.0
                fix : yes

          attenuation :
                desc : scaling factor for the strength of attenuation
                initial value : 1.0
                min : 0.0
                max : 10.0
                fix : yes

        """

        def _setup(self):

            # define EBL model, use dominguez as default
            self._tau = ebltau.OptDepth.readmodel(model=astromodels_config.absorption_models.ebl_table.value)

        def set_ebl_model(self, modelname):

            # passing modelname to ebltable, which will check if defined
            self._tau = ebltau.OptDepth.readmodel(model=modelname)

        def _set_units(self, x_unit, y_unit):

            if not hasattr(x_unit, "physical_type") or x_unit.physical_type != "energy":

                # x should be energy
                raise InvalidUsageForFunction(
                    "Unit for x is not an energy. The function "
                    "EBLOptDepth calculates energy-dependent "
                    "absorption."
                )

            # y should be dimensionless
            if (
                not hasattr(y_unit, "physical_type")
                or y_unit.physical_type != "dimensionless"
            ):
                raise InvalidUsageForFunction(
                    "Unit for y is not dimensionless.")

            self.redshift.unit = astropy_units.dimensionless_unscaled
            self.attenuation.unit = astropy_units.dimensionless_unscaled

        def evaluate(self, x, redshift, attenuation):

            if isinstance(x, astropy_units.Quantity):

                # ebltable expects TeV
                eTeV = x.to(astropy_units.TeV).value
                _unit = astropy_units.dimensionless_unscaled
                _redshift = redshift.value
                _attenuation = attenuation.value

            else:

                # otherwise it's in keV
                eTeV = x/1.e9

                _unit = 1.
                _redshift = redshift
                _attenuation = attenuation
                
                
            return _numba_eval(self._tau.opt_depth(_redshift, eTeV), _attenuation) * _unit



@nb.vectorize
def _exp(x):
    return math.exp(x)

@nb.njit(fastmath=True)
def _numba_eval(nh, xsect_interp):

    return _exp(-nh * xsect_interp )
