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, 16, 11, 31), 'aesara.graph.optdb.EquilibriumDB', 'EquilibriumDB', ({}, {}), '()', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), ((12, 17, 12, 32), 'aesara.graph.optdb.EquilibriumDB', 'EquilibriumDB', ({}, {}), '()', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), ((15, 17, 15, 32), 'aesara.graph.optdb.EquilibriumDB', 'EquilibriumDB', ({}, {}), '()', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), ((17, 13, 17, 25), 'aesara.graph.optdb.SequenceDB', 'SequenceDB', ({}, {}), '()', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), ((28, 10, 28, 24), 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', ({}, {}), '()', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), ((29, 11, 29, 58), 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', (), '', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), ((32, 16, 32, 30), 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', ({}, {}), '()', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), ((33, 17, 33, 64), 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', (), '', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), ((36, 25, 36, 39), 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', ({}, {}), '()', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), ((37, 26, 37, 73), 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', (), '', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), ((40, 31, 40, 45), 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', ({}, {}), '()', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), ((100, 24, 100, 38), 'aesara.graph.optdb.LocalGroupDB', 'LocalGroupDB', ({}, {}), '()', False, 'from aesara.graph.optdb import EquilibriumDB, LocalGroupDB, OptimizationDatabase, SequenceDB\n'), ((23, 4, 23, 55), 'aesara.compile.optdb.__position__.get', 'optdb.__position__.get', ({(23, 27, 23, 48): '"""add_destroy_handler"""', (23, 50, 23, 54): '(49.5)'}, {}), "('add_destroy_handler', 49.5)", False, 'from aesara.compile import optdb\n'), ((116, 15, 116, 77), 'aesara.gpuarray.opt.GraphToGPU', 'GraphToGPU', ({(116, 26, 116, 50): 'opt.local_optimizers_all', (116, 52, 116, 76): 'opt.local_optimizers_map'}, {}), '(opt.local_optimizers_all, opt.local_optimizers_map)', False, 'from aesara.gpuarray.opt import GraphToGPU\n'), ((85, 12, 85, 81), 'aesara.graph.opt.TopoOptimizer', 'TopoOptimizer', (), '', False, 'from aesara.graph.opt import GraphToGPULocalOptGroup, TopoOptimizer, local_optimizer\n'), ((73, 18, 73, 41), 'aesara.graph.opt.local_optimizer', 'local_optimizer', ({(73, 34, 73, 40): 'tracks'}, {}), '(tracks)', False, 'from aesara.graph.opt import GraphToGPULocalOptGroup, TopoOptimizer, local_optimizer\n')]