# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""MetaSchedule-TIR integration"""

from collections.abc import Mapping

# isort: off
from typing import Literal
from tvm_ffi import register_global_func

# isort: on
from tvm import ir, tirx
from tvm.s_tir.schedule import Schedule as _Schedule
from tvm.target import Target
from tvm.tirx.expr import IntImm

from .builder import Builder
from .cost_model import CostModel
from .database import Database
from .logging import get_loggers_from_work_dir
from .measure_callback import MeasureCallback
from .runner import Runner
from .search_strategy import SearchStrategy
from .space_generator import SpaceGenerator
from .task_scheduler import TaskScheduler
from .tune import tune_tasks
from .tune_context import TuneContext, _normalize_mod
from .utils import fork_seed


def tune_tir(  # pylint: disable=too-many-locals
    mod: ir.IRModule | tirx.PrimFunc,
    target: str | Target,
    work_dir: str,
    max_trials_global: int,
    *,
    max_trials_per_task: int | None = None,
    num_trials_per_iter: int = 64,
    builder: Builder.BuilderType = "local",
    runner: Runner.RunnerType = "local",
    database: Database.DatabaseType = "json",
    cost_model: CostModel.CostModelType = "xgb",
    measure_callbacks: MeasureCallback.CallbackListType = "default",
    task_scheduler: TaskScheduler.TaskSchedulerType = "gradient",
    space: SpaceGenerator.SpaceGeneratorType = "post-order-apply",
    strategy: SearchStrategy.SearchStrategyType = "evolutionary",
    num_tuning_cores: Literal["physical", "logical"] | int = "physical",
    seed: int | None = None,
    module_equality: str = "structural",
    special_space: Mapping[str, SpaceGenerator.SpaceGeneratorType] | None = None,
    post_optimization: bool | None = False,
) -> Database:
    """Tune a TIR function or an IRModule of TIR functions.

    Parameters
    ----------
    mod : Union[ir.IRModule, tirx.PrimFunc]
        The TIR IRModule to tune.
    target : Union[str, Target]
        The target to tune for.
    work_dir : str
        The working directory.
    max_trials_global : int
        The maximum number of trials to run globally.
    max_trials_per_task : Optional[int]
        The maximum number of trials to run per task.
    num_trials_per_iter : int
        The number of trials to run per iteration
    builder : Builder.BuilderType
        The builder.
    runner : Runner.RunnerType
        The runner.
    database : Database.DatabaseType
        The database.
    cost_model : CostModel.CostModelType
        The cost model.
    measure_callbacks : MeasureCallback.CallbackListType
        The measure callbacks.
    task_scheduler : TaskScheduler.TaskSchedulerType
        The task scheduler.
    space : SpaceGenerator.SpaceGeneratorType
        The space generator.
    strategy : SearchStrategy.SearchStrategyType
        The search strategy.
    num_tuning_cores : Union[Literal["physical", "logical"], int]
        The number of CPU cores to use during tuning.
    seed : Optional[int]
        The seed for the random number generator.
    module_equality : Optional[str]
        A string to specify the module equality testing and hashing method.
    special_space : Optional[Mapping[str, SpaceGenerator.SpaceGeneratorType]]
        A mapping from task name to a special space generator for that task.

    Returns
    -------
    database : Database
        The database with all tuning records
    """
    if isinstance(mod, tirx.PrimFunc):
        mod = _normalize_mod(mod)

    named_tasks: list[tuple[str, tirx.PrimFunc]] = []
    for gv, func in mod.functions_items():  # pylint: disable=invalid-name
        if isinstance(func, tirx.PrimFunc):
            named_tasks.append((gv.name_hint, func))
    named_tasks.sort(key=lambda x: x[0])

    task_names = [x for x, _ in named_tasks]
    tasks: list[TuneContext] = []
    for task_name, task_func, logger, rand_state in zip(
        task_names,
        [x for _, x in named_tasks],
        get_loggers_from_work_dir(work_dir, task_names),
        fork_seed(seed, n=len(named_tasks)),
    ):
        if special_space and task_name in special_space:
            task_space = special_space[task_name]
        else:
            task_space = space
        if task_space is None:
            continue
        tasks.append(
            TuneContext(
                mod=task_func,
                target=target,
                space_generator=task_space,
                search_strategy=strategy,
                task_name=task_name,
                rand_state=rand_state,
                num_threads=num_tuning_cores,
                logger=logger,
            ).clone()
        )
    return tune_tasks(
        tasks=tasks,
        task_weights=[1.0] * len(tasks),
        work_dir=work_dir,
        max_trials_global=max_trials_global,
        max_trials_per_task=max_trials_per_task,
        num_trials_per_iter=num_trials_per_iter,
        builder=builder,
        runner=runner,
        database=database,
        cost_model=cost_model,
        measure_callbacks=measure_callbacks,
        task_scheduler=task_scheduler,
        module_equality=module_equality,
        post_optimization=post_optimization,
    )


@register_global_func("tvm.s_tir.meta_schedule.tune_tir")
def _tune_tir(
    mod: ir.IRModule | tirx.PrimFunc,
    target: str | Target,
    work_dir: str,
    max_trials_global: int,
    *,
    num_trials_per_iter: int = 64,
    builder: Builder.BuilderType = "local",
    runner: Runner.RunnerType = "local",
    database: Database.DatabaseType = "json",
    cost_model: CostModel.CostModelType = "xgb",
    measure_callbacks: MeasureCallback.CallbackListType = "default",
    task_scheduler: TaskScheduler.TaskSchedulerType = "round-robin",
    space: SpaceGenerator.SpaceGeneratorType = "post-order-apply",
    strategy: SearchStrategy.SearchStrategyType = "evolutionary",
    num_tuning_cores: Literal["physical", "logical"] | int = "physical",
    seed: int | None = None,
) -> Database:
    """Interface with tuning api to tune a TIR program.

    Parameters
    ----------
    mod : Union[ir.IRModule, tirx.PrimFunc]
        The TIR function to tune.
    target : Union[str, Target]
        The target to tune for.
    work_dir : str
        The working directory.
    max_trials_global : int
        The maximum number of trials to run globally.
    num_trials_per_iter : int
        The number of trials to run per iteration
    builder : Builder.BuilderType
        The builder.
    runner : Runner.RunnerType
        The runner.
    database : Database.DatabaseType
        The database.
    cost_model : CostModel.CostModelType
        The cost model.
    measure_callbacks : MeasureCallback.CallbackListType
        The measure callbacks.
    task_scheduler : TaskScheduler.TaskSchedulerType
        The task scheduler.
    space : SpaceGenerator.SpaceGeneratorType
        The space generator.
    strategy : SearchStrategy.SearchStrategyType
        The search strategy.
    num_tuning_cores : Union[Literal["physical", "logical"], int]
        The number of CPU cores to use during tuning.
    seed : Optional[int]
        The seed for the random number generator.

    Returns
    -------
    ret_mod : IRModule
        IRModule
    """
    if isinstance(max_trials_global, IntImm):
        max_trials_global = int(max_trials_global)
    tune_tir(
        mod,
        target,
        work_dir,
        max_trials_global,
        num_trials_per_iter=num_trials_per_iter,
        builder=builder,
        runner=runner,
        database=database,
        cost_model=cost_model,
        measure_callbacks=measure_callbacks,
        task_scheduler=task_scheduler,
        space=space,
        strategy=strategy,
        num_tuning_cores=num_tuning_cores,
        seed=seed,
    )
    # Return original IRModule
    # This pass only makes optimization decision
    return mod


def compile_tir(
    database: Database,
    mod: ir.IRModule | tirx.PrimFunc,
    target: Target | str,
) -> _Schedule:
    """Compile a TIR to s_tir.Schedule, according to the records in the database.

    Parameters
    ----------
    database : Database
        The database of tuning records.
    mod : Union[ir.IRModule, tirx.PrimFunc]
        The TIR function to tune.
    target : Union[str, Target]
        The target to tune for.

    Returns
    -------
    sch : s_tir.Schedule
        The best schedule found in the database.
    """
    mod = _normalize_mod(mod)
    if not isinstance(target, Target):
        target = Target(target)
    return database.query_schedule(mod, target, workload_name="main")
