"""
WRITEME

"""

import logging
import warnings

import theano
import theano.gof.vm
from theano import config, gof
from theano.compile.function.types import Supervisor
from theano.sandbox.jax_linker import JAXLinker


_logger = logging.getLogger("theano.compile.mode")


# If a string is passed as the linker argument in the constructor for
# Mode, it will be used as the key to retrieve the real linker in this
# dictionary
predefined_linkers = {
    "py": gof.PerformLinker(),  # Use allow_gc Theano flag
    "c": gof.CLinker(),  # Don't support gc. so don't check allow_gc
    "c|py": gof.OpWiseCLinker(),  # Use allow_gc Theano flag
    "c|py_nogc": gof.OpWiseCLinker(allow_gc=False),
    "vm": gof.vm.VM_Linker(use_cloop=False),  # Use allow_gc Theano flag
    "cvm": gof.vm.VM_Linker(use_cloop=True),  # Use allow_gc Theano flag
    "vm_nogc": gof.vm.VM_Linker(allow_gc=False, use_cloop=False),
    "cvm_nogc": gof.vm.VM_Linker(allow_gc=False, use_cloop=True),
    "jax": JAXLinker(),
}


def register_linker(name, linker):
    """Add a `Linker` which can be referred to by `name` in `Mode`."""
    if name in predefined_linkers:
        raise ValueError(f"Linker name already taken: {name}")
    predefined_linkers[name] = linker


# If a string is passed as the optimizer argument in the constructor
# for Mode, it will be used as the key to retrieve the real optimizer
# in this dictionary
exclude = []
if not theano.config.cxx:
    exclude = ["cxx_only"]
OPT_NONE = gof.Query(include=[], exclude=exclude)
# Even if multiple merge optimizer call will be there, this shouldn't
# impact performance.
OPT_MERGE = gof.Query(include=["merge"], exclude=exclude)
OPT_FAST_RUN = gof.Query(include=["fast_run"], exclude=exclude)
OPT_FAST_RUN_STABLE = OPT_FAST_RUN.requiring("stable")
# We need fast_compile_gpu here.  As on the GPU, we don't have all
# operation that exist in fast_compile, but have some that get
# introduced in fast_run, we want those optimization to also run in
# fast_compile+gpu. We can't tag them just as 'gpu', as this would
# exclude them if we exclude 'gpu'.
OPT_FAST_COMPILE = gof.Query(
    include=["fast_compile", "fast_compile_gpu"], exclude=exclude
)
OPT_STABILIZE = gof.Query(include=["fast_run"], exclude=exclude)
OPT_STABILIZE.position_cutoff = 1.5000001
OPT_NONE.name = "OPT_NONE"
OPT_MERGE.name = "OPT_MERGE"
OPT_FAST_RUN.name = "OPT_FAST_RUN"
OPT_FAST_RUN_STABLE.name = "OPT_FAST_RUN_STABLE"
OPT_FAST_COMPILE.name = "OPT_FAST_COMPILE"
OPT_STABILIZE.name = "OPT_STABILIZE"

OPT_O2 = OPT_FAST_COMPILE.including("fusion")
OPT_O3 = OPT_FAST_RUN.excluding("inplace")
OPT_UNSAFE = OPT_O3.including("unsafe")

OPT_O2.name = "OPT_O2"
OPT_O3.name = "OPT_O3"
OPT_UNSAFE.name = "OPT_UNSAFE"

predefined_optimizers = {
    None: OPT_NONE,
    "None": OPT_NONE,
    "merge": OPT_MERGE,
    "o4": OPT_FAST_RUN,
    "o3": OPT_O3,
    "o2": OPT_O2,
    "o1": OPT_FAST_COMPILE,
    "unsafe": OPT_UNSAFE,
    "fast_compile": OPT_FAST_COMPILE,
    "fast_run": OPT_FAST_RUN,
    "fast_run_stable": OPT_FAST_RUN_STABLE,
    "stabilize": OPT_STABILIZE,
}


def register_optimizer(name, opt):
    """Add a `GlobalOptimizer` which can be referred to by `name` in `Mode`."""
    if name in predefined_optimizers:
        raise ValueError(f"Optimizer name already taken: {name}")
    predefined_optimizers[name] = opt


class AddDestroyHandler(gof.GlobalOptimizer):
    """
    This optimizer performs two important functions:

    1) It has a 'requirement' of the destroyhandler. This means that the fgraph
    will include it as a feature for this optimization, and keep this feature
    enabled for subsequent optimizations. All optimizations that work inplace
    on any of their inputs must run *after* this optimization to ensure that
    the DestroyHandler has been included in the fgraph.

    2) It tries to replace each output with an Op that purports to destroy it
    (but it won't I promise). If this replacement succeeds it means that
    there is a bug in theano. It should not be possible to destroy outputs.

    """

    def apply(self, fgraph):
        supervisor_added = False
        for feature in fgraph._features:
            if isinstance(feature, Supervisor):
                supervisor_added = True
                break
        if not supervisor_added:
            warnings.warn(
                "WARNING: Supervisor is not added. Please build a FunctionGraph"
                "via theano.compile.function.types.std_graph()"
                "or add the Supervisor class manually.",
                stacklevel=3,
            )

    def add_requirements(self, fgraph):
        super().add_requirements(fgraph)
        fgraph.attach_feature(gof.DestroyHandler())


class AddFeatureOptimizer(gof.GlobalOptimizer):
    """
    This optimizer adds a provided feature to the function graph.
    """

    def __init__(self, feature):
        self.feature = feature

    def add_requirements(self, fgraph):
        super().add_requirements(fgraph)
        fgraph.attach_feature(self.feature)

    def apply(self, fgraph):
        pass


class PrintCurrentFunctionGraph(gof.GlobalOptimizer):
    """
    This optimizer is for debugging.

    Toss it into the optimization pipeline to see the state of things at any
    given point.

    """

    def __init__(self, header):
        self.header = header

    def apply(self, fgraph):
        import theano.printing

        print("PrintCurrentFunctionGraph:", self.header)
        theano.printing.debugprint(fgraph.outputs)


optdb = gof.SequenceDB()
optdb.register("merge1", gof.MergeOptimizer(), 0, "fast_run", "fast_compile", "merge")

# After scan1 opt at 0.5 and before ShapeOpt at 1
# This should only remove nodes.
# The opt should not do anything that need shape inference.
# New nodes that don't have infer_shape need that the original node
# also don't have infer_shape
local_useless = gof.optdb.LocalGroupDB(apply_all_opts=True, profile=True)
optdb.register(
    "useless",
    gof.optdb.TopoDB(
        local_useless, failure_callback=gof.opt.NavigatorOptimizer.warn_inplace
    ),
    0.6,
    "fast_run",
    "fast_compile",
)

optdb.register(
    "merge1.1", gof.MergeOptimizer(), 0.65, "fast_run", "fast_compile", "merge"
)

# rearranges elemwise expressions
optdb.register(
    "canonicalize",
    gof.EquilibriumDB(ignore_newtrees=False),
    1,
    "fast_run",
    "fast_compile",
    "canonicalize_db",
)
# Register in the canonizer Equilibrium as a clean up opt the merge opt.
# Without this, as the equilibrium have ignore_newtrees=False, we
# won't merge all nodes if it is set as a global optimizer with
# final_opt=True.

# We need a new instance of MergeOptimizer to don't have its name
# changed by other usage of it.
optdb["canonicalize"].register(
    "merge", gof.opt.MergeOptimizer(), "fast_run", "fast_compile", cleanup=True
)

optdb.register(
    "merge1.2", gof.MergeOptimizer(), 1.2, "fast_run", "fast_compile", "merge"
)

optdb.register(
    "Print1.21",
    PrintCurrentFunctionGraph("Post-canonicalize"),
    1.21,
)  # 'fast_run', 'fast_compile')

# replace unstable subgraphs
optdb.register("stabilize", gof.EquilibriumDB(), 1.5, "fast_run")

optdb.register(
    "Print1.51",
    PrintCurrentFunctionGraph("Post-stabilize"),
    1.51,
)  # 'fast_run', 'fast_compile')

# misc special cases for speed
optdb.register("specialize", gof.EquilibriumDB(), 2, "fast_run", "fast_compile_gpu")

# misc special cases for speed that break canonicalization
optdb.register("uncanonicalize", gof.EquilibriumDB(), 3, "fast_run")

# misc special cases for speed that are dependent on the device.
optdb.register(
    "specialize_device", gof.EquilibriumDB(), 48.6, "fast_compile", "fast_run"
)  # must be after gpu stuff at 48.5

# especially constant merge
optdb.register("merge2", gof.MergeOptimizer(), 49, "fast_run", "merge")

optdb.register("add_destroy_handler", AddDestroyHandler(), 49.5, "fast_run", "inplace")

# final pass just to make sure
optdb.register("merge3", gof.MergeOptimizer(), 100, "fast_run", "merge")

if theano.config.check_stack_trace in ["raise", "warn", "log"]:
    _tags = ("fast_run", "fast_compile")

if theano.config.check_stack_trace == "off":
    _tags = ()

optdb.register("CheckStackTrace", gof.CheckStackTraceOptimization(), -1, *_tags)
del _tags


class Mode:
    """
    The Mode represents a way to optimize and then link a computation graph.

    Parameters
    ----------
    optimizer : a structure of type Optimizer
        An Optimizer may simplify the math, put similar computations together,
        improve numerical stability and various other improvements.
    linker : a structure of type Linker
        A Linker decides which implementations to use (C or Python, for example)
        and how to string them together to perform the computation.

    See Also
    --------
    predefined_linkers
    predefined_optimizers
    predefined_modes

    """

    def __init__(self, linker=None, optimizer="default"):
        if linker is None:
            linker = config.linker
        if type(optimizer) == str and optimizer == "default":
            optimizer = config.optimizer
        Mode.__setstate__(self, (linker, optimizer))

        # self.provided_optimizer - typically the `optimizer` arg.
        # But if the `optimizer` arg is keyword corresponding to a predefined
        # Query, then this stores the query
        # self._optimizer - typically same as provided_optimizer??

        # self.__get_optimizer - returns self._optimizer (possibly querying
        # optdb with self._optimizer)
        # self.optimizer - property that returns __get_optimizer()

    def __getstate__(self):
        return (self.provided_linker, self.provided_optimizer)

    def __setstate__(self, state):
        linker, optimizer = state
        self.provided_linker = linker
        self.provided_optimizer = optimizer
        if isinstance(linker, str) or linker is None:
            linker = predefined_linkers[linker]
        self.linker = linker
        if isinstance(optimizer, str) or optimizer is None:
            optimizer = predefined_optimizers[optimizer]
        if isinstance(optimizer, gof.Query):
            self.provided_optimizer = optimizer
        self._optimizer = optimizer
        self.call_time = 0
        self.fn_time = 0

    def __str__(self):
        return "{}(linker = {}, optimizer = {})".format(
            self.__class__.__name__,
            self.provided_linker,
            self.provided_optimizer,
        )

    def __get_optimizer(self):
        if isinstance(self._optimizer, gof.Query):
            return optdb.query(self._optimizer)
        else:
            return self._optimizer

    optimizer = property(__get_optimizer)

    def get_linker_optimizer(self, linker, optimizer):
        if isinstance(linker, str) or linker is None:
            linker = predefined_linkers[linker]
        if isinstance(optimizer, str) or optimizer is None:
            optimizer = predefined_optimizers[optimizer]
        return (linker, optimizer)

    def including(self, *tags):
        link, opt = self.get_linker_optimizer(
            self.provided_linker, self.provided_optimizer
        )
        # N.B. opt might be a Query instance, not sure what else it might be...
        #     string? Optimizer? OptDB? who knows???
        return self.clone(optimizer=opt.including(*tags), linker=link)

    def register(self, *optimizations):
        """Adds new optimization instances to a mode.

        This method adds new optimization instances to a compilation mode. It
        works like the `including()` method but takes as inputs optimization
        instances to add instead of tags.

        Parameters
        ----------
        optimizations :
            Every element of `optimizations` is a tuple containing an
            optimization instance and a floating point value indicating the
            position at which to insert the optimization in the mode.

        Returns
        -------
        Mode
            Copy of the current Mode which includes the provided
            optimizations.
        """

        link, opt = self.get_linker_optimizer(
            self.provided_linker, self.provided_optimizer
        )
        return self.clone(optimizer=opt.register(*optimizations))

    def excluding(self, *tags):
        link, opt = self.get_linker_optimizer(
            self.provided_linker, self.provided_optimizer
        )
        return self.clone(optimizer=opt.excluding(*tags), linker=link)

    def requiring(self, *tags):
        link, opt = self.get_linker_optimizer(
            self.provided_linker, self.provided_optimizer
        )
        return self.clone(optimizer=opt.requiring(*tags), linker=link)

    def clone(self, link_kwargs=None, optimizer="", **kwargs):
        """
        Create a new instance of this Mode.

        Keyword arguments can be provided for the linker,
        in which case its `clone` method will be called with these
        arguments.

        """
        if link_kwargs is None:
            link_kwargs = {}
        new_linker = self.linker.clone(**link_kwargs)

        if optimizer == "":
            optimizer = self.provided_optimizer
        new_mode = type(self)(linker=new_linker, optimizer=optimizer)
        return new_mode


# If a string is passed as the mode argument in function or
# FunctionMaker, the Mode will be taken from this dictionary using the
# string as the key
# Use VM_linker to allow lazy evaluation by default.
FAST_COMPILE = Mode(
    theano.gof.vm.VM_Linker(use_cloop=False, c_thunks=False), "fast_compile"
)
if theano.config.cxx:
    FAST_RUN = Mode("cvm", "fast_run")
else:
    FAST_RUN = Mode("vm", "fast_run")

JAX = Mode(
    JAXLinker(), gof.Query(include=["fast_run"], exclude=["cxx_only", "BlasOpt"])
)


predefined_modes = {
    "FAST_COMPILE": FAST_COMPILE,
    "FAST_RUN": FAST_RUN,
    "JAX": JAX,
}

instantiated_default_mode = None


def get_mode(orig_string):
    if orig_string is None:
        string = config.mode
    else:
        string = orig_string
    if not isinstance(string, str):
        return string  # it is hopefully already a mode...

    global instantiated_default_mode
    # The default mode is cached. However, config.mode can change
    # If instantiated_default_mode has the right class, use it.
    if orig_string is None and instantiated_default_mode:
        if string in predefined_modes:
            default_mode_class = predefined_modes[string].__class__.__name__
        else:
            default_mode_class = string
        if instantiated_default_mode.__class__.__name__ == default_mode_class:
            return instantiated_default_mode

    if string in ["Mode", "DebugMode", "NanGuardMode"]:
        if string == "DebugMode":
            # need to import later to break circular dependency.
            from .debugmode import DebugMode

            # DebugMode use its own linker.
            ret = DebugMode(optimizer=config.optimizer)
        elif string == "NanGuardMode":
            # need to import later to break circular dependency.
            from .nanguardmode import NanGuardMode

            # NanGuardMode use its own linker.
            ret = NanGuardMode(True, True, True, optimizer=config.optimizer)
        else:
            # TODO: Can't we look up the name and invoke it rather than using eval here?
            ret = eval(string + "(linker=config.linker, optimizer=config.optimizer)")
    elif string in predefined_modes:
        ret = predefined_modes[string]
    else:
        raise Exception(f"No predefined mode exist for string: {string}")

    if orig_string is None:
        # Build and cache the default mode
        if theano.config.optimizer_excluding:
            ret = ret.excluding(*theano.config.optimizer_excluding.split(":"))
        if theano.config.optimizer_including:
            ret = ret.including(*theano.config.optimizer_including.split(":"))
        if theano.config.optimizer_requiring:
            ret = ret.requiring(*theano.config.optimizer_requiring.split(":"))
        instantiated_default_mode = ret

    return ret


def get_default_mode():
    return get_mode(None)


def register_mode(name, mode):
    """
    Add a `Mode` which can be referred to by `name` in `function`.

    """
    if name in predefined_modes:
        raise ValueError(f"Mode name already taken: {name}")
    predefined_modes[name] = mode
