"""
WRITEME

"""

import logging
import sys
import warnings
from typing import Any

from pytensor.compile.aliasing import Supervisor
from pytensor.configdefaults import config
from pytensor.graph.destroyhandler import DestroyHandler
from pytensor.graph.rewriting.basic import (
    CheckStackTraceRewriter,
    GraphRewriter,
    MergeOptimizer,
    NodeProcessingGraphRewriter,
)
from pytensor.graph.rewriting.db import (
    EquilibriumDB,
    LocalGroupDB,
    RewriteDatabase,
    RewriteDatabaseQuery,
    SequenceDB,
    TopoDB,
)
from pytensor.link.basic import Linker, PerformLinker
from pytensor.link.c.basic import CLinker, OpWiseCLinker
from pytensor.link.jax.linker import JAXLinker
from pytensor.link.mlx.linker import MLXLinker
from pytensor.link.numba.linker import NumbaLinker
from pytensor.link.pytorch.linker import PytorchLinker
from pytensor.link.vm import VMLinker


_logger = logging.getLogger("pytensor.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": PerformLinker(),  # Use allow_gc PyTensor flag
    "c": CLinker(),  # Don't support gc. so don't check allow_gc
    "c|py": OpWiseCLinker(),  # Use allow_gc PyTensor flag
    "c|py_nogc": OpWiseCLinker(allow_gc=False),
    "vm": VMLinker(use_cloop=False),  # Use allow_gc PyTensor flag
    "cvm": VMLinker(use_cloop=True),  # Use allow_gc PyTensor flag
    "vm_nogc": VMLinker(allow_gc=False, use_cloop=False),
    "cvm_nogc": VMLinker(allow_gc=False, use_cloop=True),
    "jax": JAXLinker(),
    "pytorch": PytorchLinker(),
    "numba": NumbaLinker(),
    "mlx": MLXLinker(),
}


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


OPT_NONE = RewriteDatabaseQuery(include=[])
# Minimum set of rewrites needed to evaluate a function. This is needed for graphs with "dummy" Operations
OPT_MINIMUM = RewriteDatabaseQuery(include=["minimum_compile"])
# Even if multiple merge optimizer call will be there, this shouldn't
# impact performance.
OPT_MERGE = RewriteDatabaseQuery(include=["merge"])
OPT_FAST_RUN = RewriteDatabaseQuery(include=["fast_run"])
OPT_FAST_RUN_STABLE = OPT_FAST_RUN.requiring("stable")

OPT_FAST_COMPILE = RewriteDatabaseQuery(include=["fast_compile"])
OPT_STABILIZE = RewriteDatabaseQuery(include=["fast_run"])
OPT_STABILIZE.position_cutoff = 1.5000001
OPT_NONE.name = "OPT_NONE"
OPT_MINIMUM.name = "OPT_MINIMUM"
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,
    "minimum_compile": OPT_MINIMUM,
    "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 `GraphRewriter` 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(GraphRewriter):
    """
    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 pytensor. 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(
                (
                    f"A Supervisor feature is missing from {fgraph}.\n"
                    "This is needed for inplace rewrites. Either exclude inplace rewrites or add a Supervisor feature.\n"
                    "A Supervisor feature can be added via `pytensor.compile.aliasing.add_supervisor_to_fgraph`."
                ),
                stacklevel=3,
            )

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


class AddFeatureOptimizer(GraphRewriter):
    """
    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(GraphRewriter):
    """
    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 pytensor.printing

        print("PrintCurrentFunctionGraph:", self.header)  # noqa: T201
        pytensor.printing.debugprint(fgraph.outputs)


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


# 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 = LocalGroupDB(apply_all_rewrites=True, profile=True)
optdb.register(
    "useless",
    TopoDB(local_useless, failure_callback=NodeProcessingGraphRewriter.warn_inplace),
    "fast_run",
    "fast_compile",
    position=0.6,
)

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

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

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

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

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

# replace unstable subgraphs
# In fast_compile only the rewrites tagged "fast_compile" are included: those needed for the
# scipy/numpy look-alike helpers to be stable. Naive forms get no such promise.
optdb.register("stabilize", EquilibriumDB(), "fast_run", "fast_compile", position=1.5)

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

# misc special cases for speed
optdb.register("specialize", EquilibriumDB(), "fast_run", "fast_compile", position=2)

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

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

optdb.register("py_only", EquilibriumDB(), "fast_compile", position=49.1)

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

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

_tags: tuple[str, str] | tuple

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

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

optdb.register("CheckStackTrace", CheckStackTraceRewriter(), *_tags, position=-1)
del _tags


class _ActiveModeGraphRewriter(GraphRewriter):
    """Wrap a mode's optimizer so applying it records the active compile mode.

    Exposes the compile mode as ``fgraph._compile_mode`` (see ``get_active_mode``)
    so rewrites can make decisions based on the compile target (fusion split rules,
    inner graph specialization, ...). Stamps the mode when the fgraph does not already carry one, and
    restores the previous state after.
    """

    def __init__(self, rewriter: GraphRewriter, mode: "Mode"):
        self._rewriter = rewriter
        self._mode = mode

    def add_requirements(self, fgraph):
        return self._rewriter.add_requirements(fgraph)

    def apply(self, fgraph, *args, **kwargs):
        # Only stamp when unset, so a mode already established by an outer
        # optimizer run (the outermost compilation, or a nested inner-graph
        # rewrite) wins.
        stamped = getattr(fgraph, "_compile_mode", None) is None
        if stamped:
            fgraph._compile_mode = self._mode
        try:
            return self._rewriter.apply(fgraph, *args, **kwargs)
        finally:
            if stamped:
                del fgraph._compile_mode

    def print_summary(self, stream=sys.stdout, level=0, depth=-1):
        return self._rewriter.print_summary(stream=stream, level=level, depth=depth)

    def print_profile(self, stream, prof, level=0):
        # ``GraphRewriter.print_profile`` is a classmethod that raises unless
        # overridden, so delegate to the wrapped rewriter's implementation.
        return self._rewriter.print_profile(stream, prof, level=level)

    def __getattr__(self, name):
        # Delegate everything else (profiling helpers, query metadata, ...) to the
        # wrapped rewriter. ``__getattr__`` only fires for names not found normally,
        # so the overrides above take precedence. Guard ``_rewriter`` itself to avoid
        # infinite recursion before ``__init__`` has run.
        if name == "_rewriter":
            raise AttributeError(name)
        return getattr(self._rewriter, name)


class Mode:
    """A class that specifies the rewrites/optimizations used during function compilation.

    Parameters
    ----------
    optimizer
        An Optimizer may simplify the math, put similar computations together,
        improve numerical stability and various other improvements.
    linker
        A Linker decides which implementations to use (C or Python, for example)
        and how to string them together to perform the computation.
    db
        The `RewriteDatabase` used by this `Mode`.  Note: This value
        is *not* part of a `Mode` instance's pickled state.

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

    """

    def __init__(
        self,
        linker: str | Linker | None = None,
        optimizer: str | RewriteDatabaseQuery | GraphRewriter = "default",
        db: RewriteDatabase = None,
    ):
        if linker is None:
            linker = config.linker
        if isinstance(linker, str) and linker == "auto":
            linker = "numba"
        if isinstance(optimizer, str) and optimizer == "default":
            optimizer = config.optimizer

        self.__setstate__((linker, optimizer))

        if db is None:
            global optdb
            self.optdb = optdb
        else:
            self.optdb = db

        # self.provided_optimizer - typically the `optimizer` arg.
        # But if the `optimizer` arg is keyword corresponding to a predefined
        # RewriteDatabaseQuery, 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):
        global optdb

        linker, optimizer = state
        self.optdb = optdb
        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, RewriteDatabaseQuery):
            # TODO: From the __init__ signature this should always be the case
            # But some tests and internal logic allow passing a GraphRewriter directly as optimizer
            # Cleanup!
            self.provided_optimizer = optimizer
            if r := linker.required_rewrites:
                optimizer = optimizer.including(*r)
            if r := linker.incompatible_rewrites:
                optimizer = optimizer.excluding(*r)
        self._optimizer = optimizer
        self.call_time = 0
        self.fn_time = 0

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

    @property
    def function_maker(self):
        from pytensor.compile.maker import FunctionMaker

        return FunctionMaker

    @property
    def optimizer(self):
        if isinstance(self._optimizer, RewriteDatabaseQuery):
            rewriter = self.optdb.query(self._optimizer)
        else:
            rewriter = self._optimizer
        # Stamp this mode as the fgraph's active compile mode while rewriting, so
        # inner-graph rewrites can recover the right linker even when the optimizer
        # is run directly (outside ``FunctionMaker``, e.g. ``mode.optimizer.rewrite(fgraph)``).
        return _ActiveModeGraphRewriter(rewriter, self)

    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 RewriteDatabaseQuery 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


C = Mode("c", "fast_run")
CVM = Mode("cvm", "fast_run")
VM = Mode("vm", "fast_run")

NUMBA = Mode(
    NumbaLinker(),
    RewriteDatabaseQuery(include=["fast_run", "numba"]),
)

JAX = Mode(
    JAXLinker(),
    # `fusion` is not incompatible with JAX (it can run `Composite`), just not desired:
    # XLA does its own fusion, so building `Composite` only adds graph-translation overhead.
    RewriteDatabaseQuery(include=["fast_run", "jax"], exclude=["fusion"]),
)
PYTORCH = Mode(
    PytorchLinker(),
    RewriteDatabaseQuery(include=["fast_run"]),
)

MLX = Mode(
    MLXLinker(),
    RewriteDatabaseQuery(include=["fast_run", "mlx"], exclude=["fusion"]),
)

FAST_COMPILE = Mode(
    VMLinker(use_cloop=False, c_thunks=False),
    RewriteDatabaseQuery(include=["fast_compile", "py_only"]),
)

fast_run_linkers_to_mode = {
    "cvm": CVM,
    "vm": VM,
    "numba": NUMBA,
}

predefined_modes = {
    "FAST_COMPILE": FAST_COMPILE,
    "C": C,
    "CVM": CVM,
    "JAX": JAX,
    "NUMBA": NUMBA,
    "PYTORCH": PYTORCH,
    "MLX": MLX,
}

_CACHED_RUNTIME_MODES: dict[Any, Mode] = {}


def get_mode(orig_string=None):
    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...

    # Keep the original string for error messages
    upper_string = string.upper()

    if upper_string in predefined_modes:
        return predefined_modes[upper_string]

    if upper_string == "FAST_RUN":
        linker = config.linker
        if linker == "auto":
            return NUMBA
        return fast_run_linkers_to_mode[linker]

    global _CACHED_RUNTIME_MODES

    cache_key = ("MODE", config.linker) if upper_string == "MODE" else upper_string

    try:
        return _CACHED_RUNTIME_MODES[cache_key]
    except KeyError:
        pass

    # Need to define the mode for the first time
    if upper_string == "MODE":
        ret = Mode(linker=config.linker, optimizer=config.optimizer)
    elif upper_string in ("DEBUGMODE", "DEBUG_MODE"):
        from pytensor.compile.debug.debugmode import DebugMode

        # DebugMode use its own linker.
        ret = DebugMode(optimizer=config.optimizer)
    elif upper_string == "NANGUARDMODE":
        from pytensor.compile.debug.nanguardmode import NanGuardMode

        # NanGuardMode use its own linker.
        ret = NanGuardMode(True, True, True, optimizer=config.optimizer)

    else:
        raise ValueError(f"No predefined mode exist for string: {string}")

    if config.optimizer_excluding:
        ret = ret.excluding(*config.optimizer_excluding.split(":"))
    if config.optimizer_including:
        ret = ret.including(*config.optimizer_including.split(":"))
    if config.optimizer_requiring:
        ret = ret.requiring(*config.optimizer_requiring.split(":"))
    # Cache the mode for next time
    _CACHED_RUNTIME_MODES[cache_key] = 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
