import collections
import os
import re
import sys
import time
import warnings

import astropy.units as u

from astromodels.core.my_yaml import my_yaml
from astromodels.functions.function import get_function_class
from astromodels.utils import get_user_data_path
from astromodels.xspec import _xspec


class XSpecNotAvailable(ImportWarning):
    pass


# When running in a Anaconda environment, the package xspec-modelsonly
# will install the models data in a specific place, so we set the HEADAS variable to
# point to the right place for the models to find their data
if (
    os.environ.get("CONDA_PREFIX") is not None
    and os.environ.get("HEADAS", None) is None
):

    # NOTE: the 'headas' directory doesn't actually exists, but models are looked for
    # (according to Xspec documentation) into $HEADAS../spectral, so we put
    # HEADAS=[....]/headas so that $HEADAS/../ will be the right place where the
    # 'spectral' directory is

    os.environ["HEADAS"] = os.path.join(
        os.environ.get("CONDA_PREFIX"), "lib", "Xspec", "headas"
    )

    # os.environ['HEADAS'] = os.path.join(os.environ.get("CONDA_PREFIX"),
    # 'Xspec', 'headas')

    # we need to create the directory otherwise an exception casted:
    if not os.path.exists(os.environ["HEADAS"]):
        os.makedirs(os.environ["HEADAS"])

# This list defines all python protected names (names which variables or attributes
# should not have)


illegal_variable_names = [
    "and",
    "assert",
    "break",
    "class",
    "continue",
    "def",
    "del",
    "elif",
    "else",
    "except",
    "exec",
    "finally",
    "for",
    "if",
    "import",
    "in",
    "is",
    "lambda",
    "not",
    "or",
    "pass",
    "print",
    "raise",
    "return",
    "try",
    "while",
    "data",
    "float",
    "int",
    "numeric",
    "array",
    "close",
    "input",
    "open",
    "range",
    "type",
    "write",
    "zeros",
    "from",
    "global",
]


def find_model_dat():
    """
    Find the file containing the definition of all the models in Xspec
    (model.dat) and return its path
    """

    # model.dat is in either $HEADAS/spectral/manager or
    # $HEADAS/../spectral/manager

    headas_env = os.environ.get("HEADAS")

    assert headas_env is not None, (
        "You need to setup the HEADAS variable before importing this module."
        " See Heasoft documentation."
    )

    # Expand all variables and other things like ~
    headas_env = os.path.expandvars(os.path.expanduser(headas_env))

    # Lazy check that it exists
    assert os.path.exists(
        headas_env
    ), (
        "The HEADAS env. variable point to a non-existent directory: %s"
        % (headas_env)
    )

    # First try: $HEADAS/spectral/manager/model.dat
    final_path = os.path.join(headas_env, "spectral", "manager", "model.dat")

    if os.path.exists(final_path):
        return os.path.abspath(final_path)

    # Second try: $HEADAS/../spectral/manager/model.dat
    parent_dir = os.path.dirname(headas_env)
    final_path = os.path.join(parent_dir, "spectral", "manager", "model.dat")

    if os.path.exists(final_path):
        return os.path.abspath(final_path)

    # If neither exists, raise an error
    raise FileNotFoundError(
        "Cannot find Xspec model definition file. Tried: "
        f"{os.path.join(headas_env, 'spectral', 'manager', 'model.dat')} "
        "and "
        f"{os.path.join(parent_dir, 'spectral', 'manager', 'model.dat')}"
    )


def get_models(model_dat_path):
    """Parse the model.dat file from Xspec and returns a dictionary containing
    the definition of all the models.

    :param model_dat_path: the path to the model.dat file
    :return: dictionary containing the definition of all XSpec models
    """

    # Check first if we already have a model data file in the data directory

    with open(model_dat_path) as f:

        # model.dat is a text file, no size issues here (will fit in memory)

        model_dat = f.read()

    # Replace \r or \r\n with \n (the former two might be in the file if
    # it originates from Windows)
    if "\n" in model_dat:

        model_dat = model_dat.replace("\r", "")

    else:

        model_dat = model_dat.replace("\r", "\n")

    # Loop through the file and build the model definition
    # dictionary

    lines = model_dat.split("\n")

    model_definitions = collections.OrderedDict()

    for line in lines:

        match = re.match("""(.+(add|mul|con|acn).+)""", line)

        if match is not None:

            # This is a model definition

            tokens = line.split()

            if len(tokens) == 7:

                (
                    model_name,
                    n_parameters,
                    min_energy,
                    max_energy,
                    library_function,
                    model_type,
                    flag,
                ) = line.split()

            else:

                (
                    model_name,
                    n_parameters,
                    min_energy,
                    max_energy,
                    library_function,
                    model_type,
                    flag,
                    flag_2,
                ) = line.split()

            this_model = collections.OrderedDict()

            this_model["description"] = (
                "The %s model from XSpec (https://heasarc.gsfc.nasa.gov/xanadu/"
                "xspec/manual/XspecModels.html)" % model_name
            )

            this_model["parameters"] = collections.OrderedDict()

            model_definitions[(model_name, library_function, model_type)] = this_model

        else:

            # This is a parameter definition

            if len(line.split()) == 0:
                # Empty line
                continue

            # Parameters are free by default, unless they have delta < 0
            # they are switch parameters or scale parameters

            free = True

            if line[0] == "$":

                # Probably a switch parameter

                free = False

                tokens = line.split()

                if len(tokens) == 2:

                    par_name = tokens[0][1:]

                    default_value = tokens[1]

                    par_unit = ""

                    hard_minimum, soft_minimum, soft_maximum, hard_maximum = (
                        0,
                        0,
                        1e9,
                        1e9,
                    )

                elif len(tokens) == 3:

                    par_name = tokens[0][1:]

                    default_value = tokens[2]

                    par_unit = ""

                    hard_minimum, soft_minimum, soft_maximum, hard_maximum = (
                        0,
                        0,
                        1e9,
                        1e9,
                    )

                else:

                    match = re.match(r'(\S+)\s+(".+"|[a-zA-Z]+)?(.+)*', line[1:])

                    par_name, par_unit, par_spec = match.groups()

                    tokens = par_spec.split()

                    if len(tokens) == 1:

                        default_value = tokens[0]
                        par_unit = ""
                        (
                            hard_minimum,
                            soft_minimum,
                            soft_maximum,
                            hard_maximum,
                        ) = (0, 0, 1e9, 1e9)

                    else:

                        par_unit = ""

                        (
                            default_value,
                            hard_minimum,
                            soft_minimum,
                            soft_maximum,
                            hard_maximum,
                            delta,
                        ) = par_spec.split()

            else:

                # A normal parameter

                match = re.match(r'(\S+)\s+(".+"|\S+)(.+)', line)

                if match is None:

                    raise RuntimeError("Cannot parse parameter %s" % line)

                par_name, par_unit, par_spec = match.groups()

                if par_name[0] == "*":

                    # Scale parameter (always frozen)

                    par_name = par_name[1:]

                    free = False

                    tokens = par_spec.split()

                    if len(tokens) == 1:

                        default_value = tokens[0]
                        (
                            hard_minimum,
                            soft_minimum,
                            soft_maximum,
                            hard_maximum,
                            delta,
                        ) = (None, None, None, None, 0.1)

                    else:

                        (
                            default_value,
                            hard_minimum,
                            soft_minimum,
                            soft_maximum,
                            hard_maximum,
                            delta,
                        ) = par_spec.split()

                        delta = abs(float(delta))

                else:

                    (
                        default_value,
                        hard_minimum,
                        soft_minimum,
                        soft_maximum,
                        hard_maximum,
                        delta,
                    ) = par_spec.split()

                    delta = float(delta)

                    if delta <= 0:
                        free = False

                        delta = abs(delta)

            # Now fix the parameter name removing illegal characters
            # (for example 'log(A)' is not a legal name)

            par_name = re.sub(r"[\(,\)]", "_", par_name)
            par_name = re.sub(r"<", "_less_", par_name)
            par_name = re.sub(r">", "_more_", par_name)
            par_name = re.sub(r"/", "_div_", par_name)
            par_name = re.sub(r"\-", "_minus_", par_name)
            par_name = re.sub(r"\+", "_plus_", par_name)
            par_name = re.sub(r"\.", "_dot_", par_name)
            par_name = re.sub(r"@", "_at_", par_name)

            # Parameter names must be lower case
            par_name = par_name.lower()

            # Some parameters are enclosed between ", like "z"
            par_name = par_name.replace('"', "")

            # "z" is a protected name in astromodels.
            if par_name == "z":

                par_name = "redshift"

                # this is illegal because of the 2D functions
            if par_name == "y":

                par_name = "y1"

            # Check that the parameter name is not an illegal Python name
            if par_name in illegal_variable_names:

                par_name = "xs_%s" % par_name

            # Sometimes the unit is " " which is not recognized by astropy
            if par_unit:

                par_unit = par_unit.replace('"', "")

            # There are some errors in model.dat , like KeV instead of keV, and
            # ergcm/s instead of erg cm /s
            # Let's correct them
            if par_unit == "KeV":

                par_unit = "keV"

            elif par_unit == "ergcm/s":

                par_unit = "erg cm / s"

            elif par_unit == "days":

                par_unit = "day"

            elif par_unit == "z" and par_name.lower() == "redshift":

                par_unit = ""

            # There are funny units in model.dat, like "Rs" (which means
            # Schwarzschild radius) or other things so let's try to convert the
            # par_unit into an astropy.Unit instance. If that fails, use a
            # unitless unit
            try:

                _ = u.Unit(par_unit)

            except ValueError:

                # Unit not recognized

                # warnings.warn("Unit %s is not recognized by astropy." % par_unit)

                par_unit = ""

            # Make sure that this is a valid python identifier
            # by matching it with the relative regexp
            if re.match("([a-zA-Z_][a-zA-Z0-9_]*)$", par_name) is None:

                raise ValueError("Illegal identifier name %s" % (par_name))

            if hard_maximum is not None and hard_minimum is not None:

                if float(hard_maximum) < float(hard_minimum):

                    raise ValueError(
                        "Hard maximum (%s) < hard minimum (%s)"
                        % (hard_maximum, hard_minimum)
                    )

                if float(default_value) > float(hard_maximum):

                    default_value = hard_maximum

                elif float(default_value) < float(hard_minimum):

                    default_value = hard_minimum

            this_model["parameters"][par_name] = {
                "initial value": float(default_value),
                "desc": "(see https://heasarc.gsfc.nasa.gov/xanadu/xspec/manual/"
                "XspecModels.html)",
                "min": hard_minimum,
                "max": hard_maximum,
                "delta": float(delta),
                "unit": par_unit,
                "free": free,
            }

    return model_definitions


# TODO: gabs is returning all 1.0

class_definition_code = '''

import numpy as np
import astropy.units as u
from astromodels.functions import Function1D,FunctionMeta
from astromodels.utils.file_utils import copy_if_needed
from astromodels.xspec import _xspec

# These are multiplicative functions which need numerical differentiation
_force_differentiation = [
    "XS_gabs",
    "XS_expfac",
    "XS_plabs",
    "XS_pwab",
    "XS_spline",
    "XS_swind1",
    "XS_xion",
    "XS_zxipcf",
    "XS_cabs",
    "XS_wabs",
    "XS_zwabs",
]


class XS_$MODEL_NAME$(Function1D, metaclass=FunctionMeta):
    """
    $DOCSTRING$
    """

    def _setup(self):

        # Link to the Xspec function
        self._model = _xspec.$XSPEC_FUNCTION$

        self._model_type = "$MODEL_TYPE$"

        # Decide whether we need numerical differentiation

        if self._model_type == "add" or self._name in _force_differentiation:

            self._scale = 1e5

            self._differentiate = True

        else:

            # For multiplicative models, usually there is no differentiation to be made
            # (but there are exceptions, such as the gabs model)

            self._scale = 1

            self._differentiate = False

        # Now set the units as appropriate

        if self._model_type == "add":

            self._fixed_units = (u.keV, 1 / (u.keV * u.cm**2 * u.s))

        else:

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

    def evaluate(self, x, $PARAMETERS_NAMES$):

        quantity = False

        if isinstance(x, u.Quantity):

            x = np.array(x.to("keV").value, ndmin=1, copy=copy_if_needed, dtype=float)

            quantity = True

            parameters_tuple = tuple(map(lambda x:x.value, ($PARAMETERS_NAMES$,)))

        else:

            # Create a tuple of the current values of the parameters

            parameters_tuple = ($PARAMETERS_NAMES$,)

        # We need to make sure that the energy array is sorted because otherwise some
        # Xspec models will give incorrect values
        idx = np.argsort(x)

        # This is needed to be able to "reverse" the sort operation
        rev_idx = np.argsort(idx)

        # Ordered input vector

        xx = x[idx]

        if self._differentiate:

            # Finite difference differentiation of the Xspec
            # function. Indeed, the xspec function return the integral
            # of the function on energy ranges, while we need the
            # differential flux.

            # Adapt the epsilon to the value to reduce the error

            epsilon = xx / self._scale

            # In the normal case, when x is an array of more than one
            # element, the first call will succeed. If however there
            # is only one element in the array xspec will complain,
            # so handle that as a special case

            try:

                val = self._model(parameters_tuple, xx - epsilon, xx + epsilon)

            except TypeError:
                msg = "This is a bug, xspec call failed and x is not only one element"
                assert xx.shape[0] == 1, msg

                val = np.array(
                    self._model(
                        parameters_tuple, ((xx - epsilon)[0], (xx + epsilon)[0])
                    )[0],
                    ndmin=1,
                )

            if self._model_type == "add":

                # val is now F(x-epsilon,x+epsilon) ~ f(x) * ( 2 * epsilon )

                # In a additive model the function returns the integral over the bins

                final_value = val / (2 * epsilon)

            else:

                final_value = val

        else:

            # In a multiplicative model usually the function returns the average factor
            # over the bins (but there are exceptions, handled through the list
            # _force_differentiation)

            try:

                final_value = self._model(parameters_tuple, xx, xx)

            except TypeError:

                assert (
                    xx.shape[0] == 1
                ), "This is a bug, xspec call failed and x is not only one element"

                final_value = self._model(parameters_tuple, ((xx)[0], (xx)[0]))

        if quantity:

            if self._model_type == "add":

                return final_value[rev_idx] * u.Unit("1 / (keV cm^2 s)")

            else:

                return final_value[rev_idx] * u.dimensionless_unscaled

        else:

            return final_value[rev_idx]

    def has_fixed_units(self):

        return True

    def _set_units(self, x_unit, y_unit):

        # Make sure this is an energy
        # assert str(x_unit.physical_type) == 'energy', "Trying to set x-unit with
        # (%s), which is not energy" % (x_unit)
        assert "energy" in str(
            x_unit.physical_type
        ), "Trying to set x-unit with (%s), which is not energy" % (x_unit)

        # Make sure the y_unit is the correct one
        try:

            y_unit.to(self._fixed_units[1])

        except Exception as e:

            raise RuntimeError(
                f"{e} - Xspec model {self.name} cannot have units of {y_unit}."
            )

    def _integral(self, low_bounds, hi_bounds, $PARAMETERS_NAMES$):

        # Create a tuple of the current values of the parameters
        parameters_tuple = ($PARAMETERS_NAMES$,)

        return self._model(parameters_tuple, low_bounds, hi_bounds)
'''


def xspec_model_factory(model_name, xspec_function, model_type, definition):

    class_name = "XS_%s" % model_name

    # Get the path to the user data directory
    user_data_path = str(get_user_data_path())

    # Check if the code for this function already exists

    code_file_name = os.path.join(user_data_path, "%s.py" % class_name)

    if os.path.exists(code_file_name):

        # Code already exists
        pass

    else:

        print("Generating code for Xspec model %s..." % model_name)

        # If this is an additive model (model_type == 'add') we need to add
        # one more parameter (normalization)

        if model_type == "add":

            definition["parameters"]["norm"] = {
                "initial value": 1.0,
                "desc": "(see https://heasarc.gsfc.nasa.gov/xanadu/xspec/manual/"
                "XspecModels.html)",
                "min": 0,
                "max": None,
                "delta": 0.1,
                "unit": "keV / (cm2 s)",
                "free": True,
            }

        assert model_type != "con", "Convolution models are not yet supported"

        # Get a list of the parameter names
        parameters_names = ", ".join(list(definition["parameters"].keys()))

        # Create the docstring
        docstring = my_yaml.dump(definition, default_flow_style=False)

        # Create the class by substituting in the class_definition_code the
        # relevant things for this model

        code = class_definition_code.replace("$MODEL_NAME$", model_name)
        code = code.replace("$DOCSTRING$", docstring)
        code = code.replace("$PARAMETERS_NAMES$", parameters_names)
        code = code.replace("$XSPEC_FUNCTION$", xspec_function)
        code = code.replace("$MODEL_TYPE$", model_type)

        # Write to the file

        with open(code_file_name, "w+") as f:

            f.write("# This code has been automatically generated. Do not edit.\n")
            f.write("\n\n%s\n" % code)

        time.sleep(1)

    # Add the path to sys.path if it doesn't
    if user_data_path not in sys.path:

        sys.path.append(user_data_path)

    namespace = {}
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        exec(f"from {class_name} import {class_name}", globals(), namespace)

    # Return the class we just imported
    return class_name, namespace[class_name]


def setup_xspec_models():

    classes = []

    sys.stdout.write("Loading xspec models...")

    all_models = get_models(find_model_dat())

    for model_name, xspec_function, model_type in all_models:

        if model_type == "con":

            # convolution models are not supported
            continue

        if not hasattr(_xspec, xspec_function):

            # Some function do not exist in the wrapper. Let's ignore them

            continue

        this_model = all_models[(model_name, xspec_function, model_type)]

        # When the class is created it is registered among the known functions in the
        # function module (it happens in the metaclass), so we don't need to do
        # anything special here after the class type is created

        this_class_name, this_class = xspec_model_factory(
            model_name, xspec_function, model_type, this_model
        )

        classes.append(this_class_name)

    sys.stdout.write("done\n")

    return classes


# This will either work or issue a warning if XSpec is not available

new_functions = setup_xspec_models()

# Now import the new classes in the local namespace (if any)
# This is needed to make the classes pickeable

__all__ = []

for function_name in new_functions:

    __all__.append(function_name)

    locals()[function_name] = get_function_class(function_name)
