from aesara.compile import optdb
from aesara.graph.opt import GraphToGPULocalOptGroup, TopoOptimizer, local_optimizer
from aesara.graph.optdb import (
    EquilibriumDB,
    LocalGroupDB,
    OptimizationDatabase,
    SequenceDB,
)


gpu_optimizer = EquilibriumDB()
gpu_cut_copies = EquilibriumDB()

# Not used for an EquilibriumOptimizer. It has the "tracks" that we need for GraphToGPUDB.
gpu_optimizer2 = EquilibriumDB()

gpu_seqopt = SequenceDB()

# do not add 'fast_run' to these two as this would always enable gpuarray mode
optdb.register(
    "gpuarray_opt",
    gpu_seqopt,
    optdb.__position__.get("add_destroy_handler", 49.5) - 1,
    "gpuarray",
)


pool_db = LocalGroupDB()
pool_db2 = LocalGroupDB(local_opt=GraphToGPULocalOptGroup)
pool_db2.__name__ = "pool_db2"

matrix_ops_db = LocalGroupDB()
matrix_ops_db2 = LocalGroupDB(local_opt=GraphToGPULocalOptGroup)
matrix_ops_db2.__name__ = "matrix_ops_db2"

abstract_batch_norm_db = LocalGroupDB()
abstract_batch_norm_db2 = LocalGroupDB(local_opt=GraphToGPULocalOptGroup)
abstract_batch_norm_db2.__name__ = "abstract_batch_norm_db2"

abstract_batch_norm_groupopt = LocalGroupDB()
abstract_batch_norm_groupopt.__name__ = "gpuarray_batchnorm_opts"


def register_opt(*tags, **kwargs):
    def f(local_opt):
        name = (kwargs and kwargs.pop("name")) or local_opt.__name__
        gpu_optimizer.register(name, local_opt, "fast_run", "gpuarray", *tags)
        return local_opt

    return f


def register_opt2(tracks, *tags, **kwargs):
    """
    Decorator for the new GraphToGPU optimizer.
    Takes an extra parameter(Op) compared to register_opt decorator.

    Parameters
    ----------
    tracks : List of Op class Or Op instance or None
        The Node's Op to which optimization is being applied.

    tags : String
        The optimization tag to which the optimizer will be registered.

    """

    def f(local_opt):
        name = (kwargs and kwargs.pop("name")) or local_opt.__name__
        if isinstance(local_opt, OptimizationDatabase):
            opt = local_opt
        else:
            opt = local_optimizer(tracks)(local_opt)
        gpu_optimizer2.register(name, opt, "fast_run", "gpuarray", *tags)
        return local_opt

    return f


def register_inplace(*tags, **kwargs):
    def f(local_opt):
        name = (kwargs and kwargs.pop("name")) or local_opt.__name__
        optdb.register(
            name,
            TopoOptimizer(local_opt, failure_callback=TopoOptimizer.warn_inplace),
            60,
            "fast_run",
            "inplace",
            "gpuarray",
            *tags,
        )
        return local_opt

    return f


# Register GPU convolution implementation
# They are tried in a specific order so we can control
# which ones take precedence over others.
abstractconv_groupopt = LocalGroupDB()
abstractconv_groupopt.__name__ = "gpuarray_abstractconv_opts"
register_opt("fast_compile")(abstractconv_groupopt)


class GraphToGPUDB(OptimizationDatabase):
    """
    Retrieves the list local optimizers based on the optimizer flag's value
    from EquilibriumOptimizer by calling the method query.

    """

    def query(self, *tags, **kwtags):
        from aesara.gpuarray.opt import GraphToGPU

        opt = gpu_optimizer2.query(*tags, **kwtags)
        return GraphToGPU(opt.local_optimizers_all, opt.local_optimizers_map)

[(11, 'aesara.graph.optdb.EquilibriumDB', 'EquilibriumDB', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), (12, 'aesara.graph.optdb.EquilibriumDB', 'EquilibriumDB', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), (15, 'aesara.graph.optdb.EquilibriumDB', 'EquilibriumDB', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), (17, 'aesara.graph.optdb.SequenceDB', 'SequenceDB', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), (28, 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), (29, 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), (32, 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), (33, 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), (36, 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), (37, 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), (40, 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), (100, 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), (23, 'aesara.compile.optdb.__position__.get', 'optdb.__position__.get', False, 'from aesara.compile import optdb\n'), (116, 'aesara.gpuarray.opt.GraphToGPU', 'GraphToGPU', False, 'from aesara.gpuarray.opt import GraphToGPU\n'), (85, 'aesara.graph.opt.TopoOptimizer', 'TopoOptimizer', False, 'from aesara.graph.opt import GraphToGPULocalOptGroup, TopoOptimizer, local_optimizer\n'), (73, 'aesara.graph.opt.local_optimizer', 'local_optimizer', False, 'from aesara.graph.opt import GraphToGPULocalOptGroup, TopoOptimizer, local_optimizer\n')]