# 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.
# ruff: noqa: F403
"""JSON Database validation script"""

import argparse
import itertools
import logging
import warnings
from collections.abc import Callable
from statistics import mean
from typing import Any

import numpy as np  # type: ignore
import tvm_ffi
from tvm_ffi import get_global_func, register_global_func

import tvm
from tvm.ir import IRModule
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.tune_utils import generate_input_data
from tvm.s_tir.meta_schedule.utils import remove_build_dir
from tvm.s_tir.schedule import Trace
from tvm.s_tir.tensor_intrin import *  # type: ignore # pylint: disable=wildcard-import,unused-wildcard-import
from tvm.support import describe
from tvm.target import Target
from tvm.testing.utils import strtobool

DELIMITOR = "\n" + "-" * 30 + "\n"


def _parse_args():
    args = argparse.ArgumentParser()
    args.add_argument(
        "--work-dir",
        type=str,
        required=True,
        help="The path to the work directory containing database files.",
    )
    args.add_argument(
        "--target",
        type=Target,
        required=True,
    )
    args.add_argument(
        "--baseline-target",
        type=Target,
        default='{"kind": "llvm", "num-cores": 1}',
        required=False,
        help="The baseline target to compile the original module.",
    )
    args.add_argument(
        "--top-k",
        type=int,
        default=10**9,
        required=False,
        help="The number of top-k tuning records to validate for each unique original workload.",
    )
    args.add_argument(
        "--rpc-host",
        type=str,
    )
    args.add_argument(
        "--rpc-port",
        type=int,
    )
    args.add_argument(
        "--rpc-key",
        type=str,
    )
    args.add_argument(
        "--number",
        type=int,
        default=3,
    )
    args.add_argument(
        "--repeat",
        type=int,
        default=1,
    )
    args.add_argument(
        "--min-repeat-ms",
        type=int,
        default=100,
    )
    args.add_argument(
        "--cpu-flush",
        type=lambda x: bool(strtobool(x)),
        help="example: True / False",
        required=True,
    )
    args.add_argument(
        "--input-generator-func",
        type=str,
        default="tvm.s_tir.meta_schedule.testing.default_input_generator",
    )
    args.add_argument(
        "--check-metric-func",
        type=str,
        default="tvm.s_tir.meta_schedule.testing.default_check_metric",
    )
    parsed = args.parse_args()
    parsed.target = tvm.target.Target(parsed.target)
    if parsed.rpc_host is not None and parsed.rpc_port is not None and parsed.rpc_key is not None:
        parsed.rpc_config = ms.runner.RPCConfig(
            tracker_host=parsed.rpc_host,
            tracker_port=parsed.rpc_port,
            tracker_key=parsed.rpc_key,
            session_timeout_sec=600,
        )
    else:
        parsed.rpc_config = None
        warnings.warn("RPC config is not provided, will use local runner.")
    if parsed.cpu_flush and parsed.target.kind.name != "llvm":
        warnings.warn("cpu_flush is only supported on llvm target")
    return parsed


# arg parser
ARGS = _parse_args()

# logging
logging.basicConfig(
    format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
)
logging.getLogger("tvm.s_tir.meta_schedule").setLevel(logging.DEBUG)
logging.getLogger("tvm.s_tir.meta_schedule.runner").setLevel(logging.WARN)


def get_device_type(target: Target) -> str:
    """Get the device type string from a target.

    Parameters
    ----------
    target : Target
        The target to get the device type from.

    Returns
    -------
    device_type : str
        The device type string.
    """
    if target.kind.name == "llvm":
        return "cpu"
    elif target.kind.name == "cuda":
        return "cuda"
    else:
        raise RuntimeError(f"Unsupported target kind for device type: {target.kind.name}")


def get_runtime_device(target: Target) -> tvm.runtime.Device:
    """Get the runtime device from a target.

    Parameters
    ----------
    target : Target
        The target to get the runtime device from.

    Returns
    -------
    device : tvm.runtime.Device
        The runtime device.
    """
    if target.kind.name == "llvm":
        return tvm.cpu()
    elif target.kind.name == "cuda":
        return tvm.cuda()
    else:
        raise RuntimeError(f"Unsupported target kind for runtime device: {target.kind.name}")


def check_and_run(func: str | Callable, *args, **kwargs) -> Any:
    """Check if the function is a string or a callable, and run it."""
    if isinstance(func, str):
        func = get_global_func(func)
    return func(*args, **kwargs)  # type: ignore


class OriginalModule:
    """Original module class for deduplication."""

    def __init__(self, mod: IRModule):
        self.mod = mod

    def __eq__(self, __o: "OriginalModule") -> bool:  # type: ignore
        return tvm_ffi.structural_equal(self.mod, __o.mod)

    def __hash__(self) -> int:
        return tvm_ffi.structural_hash(self.mod)


def initializer() -> None:
    """Initializer function to register the functions on PopenWorker."""

    @register_global_func("tvm.s_tir.meta_schedule.testing.default_check_metric")
    def default_check_metric(  # pylint: disable=unused-variable,unreachable-code
        lhs: list[tvm.runtime.Tensor], rhs: list[tvm.runtime.Tensor]
    ) -> bool:
        """Check if the outputs are equal

        Parameters
        ----------
        lhs : List[tvm.runtime.Tensor]
            The first list of Tensors to compare.

        rhs : List[tvm.runtime.Tensor]
            The second list of Tensors to compare.

        Returns
        -------
        is_equal : bool
            Whether the two lists of Tensors are equal.
        """
        assert len(lhs) == len(rhs), "Different number of outputs from two modules"
        for i in range(len(lhs)):  # pylint: disable=consider-using-enumerate
            if not np.allclose(lhs[i].numpy(), rhs[i].numpy(), rtol=1e-3, atol=2e-3):
                return False
        return True


@register_global_func("tvm.s_tir.meta_schedule.testing.default_input_generator")
def default_input_generator(  # pylint: disable=unused-variable
    mod: IRModule,
) -> list[tvm.runtime.Tensor]:
    """Default input generator function

    Parameters
    ----------
    mod : IRModule
        The IRModule to generate the input data for.

    Returns
    -------
    inputs : List[tvm.runtime.Tensor]
        The generated input data.
    """

    args_info = ms.arg_info.TensorInfo.from_prim_func(mod["main"])
    inputs = [
        tvm.runtime.tensor(
            generate_input_data(input_shape=arg_info.shape, input_dtype=arg_info.dtype)
        )
        for arg_info in args_info
    ]
    return inputs


def to_numpy(a: list[tvm.runtime.Tensor]) -> list[np.ndarray]:
    """Convert a list of TVM Tensor to a list of numpy array

    Parameters
    ----------
    a : List[tvm.runtime.Tensor]
        The list of TVM Tensor to be converted

    Returns
    -------
    b : List[np.ndarray]
        The list of numpy array
    """
    assert a is not None, "Empty result cannot be converted to numpy"
    return [x.numpy() for x in a]


def to_tvm_tensor(a: list[np.ndarray]) -> list[tvm.runtime.Tensor]:
    """Convert a list of numpy array to a list of TVM Tensor

    Parameters
    ----------
    a : List[np.ndarray]
        The list of numpy array to be converted.

    Returns
    -------
    b : List[tvm.runtime.Tensor]
        The list of TVM Tensor.
    """
    assert a is not None, "Empty result cannot be converted to TVM Tensor"
    return [tvm.runtime.tensor(x) for x in a]


def is_failed_record(record: ms.database.TuningRecord) -> bool:
    """Check if a tuning record is failed.

    Parameters
    ----------
    record : TuningRecord
        The tuning record to check.

    Returns
    -------
    is_failed : bool
    """
    return len(record.run_secs) == 1 and record.run_secs[0] == 1e9


def print_with_counter_func(counter: int, total: int) -> Callable:
    """Print with counter

    Parameters
    ----------
    counter : int
        The counter to print with.
    total : int
        The total number of items to print with.

    Returns
    -------
    print_result : Callable
        The print result function.
    """

    def print_result(
        result: str,
        *,
        original_mod: IRModule = None,
        scheduled_mod: IRModule = None,
        inputs: list[np.ndarray] | None = None,
        original_res: list[np.ndarray] | None = None,
        scheduled_res: list[np.ndarray] | None = None,
        original_run_secs: list[float] | None = None,
        scheduled_run_secs: list[float] | None = None,
        exception: Exception | None = None,
        trace: str | None = None,
    ) -> None:
        """Print the validation result."""
        status = f"Progress {counter: 6d} / {total: 6d} (estimated) checked, result: {result:>10}, "

        if result in ["pass", "wrong answer"]:
            status += (
                f"original: {mean(original_run_secs) * 1e3: 10.3f} ms, "
                f"scheduled: {mean(scheduled_run_secs) * 1e3: 10.3f} ms"
            )

        output = [status]
        if result not in ["pass", "skip"]:
            output.extend(
                [
                    "Original IRModule:" + DELIMITOR + original_mod.script(),
                    "Scheduled IRModule:" + DELIMITOR + scheduled_mod.script(),
                    "Trace" + DELIMITOR + str(trace),
                ]
            )
            if result == "wrong answer":
                output.extend(
                    [
                        "Input:" + DELIMITOR + str(inputs),
                        "Original Result:" + DELIMITOR + str(original_res),
                        "Scheduled Result:" + DELIMITOR + str(scheduled_res),
                        "Max Diff:"
                        + DELIMITOR
                        + str(
                            [
                                np.max(np.abs(original_res[i] - scheduled_res[i]))
                                for i in range(len(original_res))
                            ]
                        )
                        + "\n",
                    ]
                )
            elif result == "exception":
                output.extend(["Exception:" + DELIMITOR + str(exception) + "\n"])
            else:
                raise ValueError(f"Unknown result: {result}")
        print("\n\n".join(output))

    return print_result


def make_alloc_arg_and_check(
    inputs: list[np.ndarray],
    original_mod: IRModule,
    scheduled_mod: IRModule,
    trace: str,
    original_res: list[np.ndarray],
    original_run_secs: list[float],
    print_result: Callable,
) -> tuple[Callable, Callable]:
    """Make alloc_arg and check functions for the given inputs and collect results.

    Parameters
    ----------
    inputs : List[np.ndarray]
        The inputs to the two modules.
    original_mod : IRModule
        The original IRModule.
    scheduled_mod : IRModule
        The scheduled IRModule.
    trace : str
        The trace of the scheduled IRModule.
    original_res : List[np.ndarray]
        The original results.
    original_run_secs : List[float]
        The original run times.
    print_result : Callable
        The print result function.

    Returns
    -------
    f_with_args_alloc_argument : Callable
        The function to allocate arguments.

    f_with_args_run_evaluator : Callable
        The function to run evaluator.
    """

    def f_with_args_alloc_argument_common(
        device: tvm.runtime.Device,
        args_info: ms.runner.rpc_runner.T_ARG_INFO_JSON_OBJ_LIST,  # pylint: disable=unused-argument
        alloc_repeat: int,
    ) -> list[ms.runner.rpc_runner.T_ARGUMENT_LIST]:
        """Allocate arguments using the given inputs.

        Parameters
        ----------
        session : RPCSession
            The RPC session.
        device : Device
            The device.
        args_info : T_ARG_INFO_JSON_OBJ_LIST
            argument information.
        alloc_repeat : int
            The number of times to repeat the allocation.

        Returns
        -------
        args_list : List[T_ARGUMENT_LIST]
            The list of argument lists.
        """
        return [
            [tvm.runtime.tensor(arg, device=device) for arg in inputs] for _ in range(alloc_repeat)
        ]

    def f_with_args_run_evaluator_common(
        rt_mod: tvm.runtime.Module,
        device: tvm.runtime.Device,
        evaluator_config: ms.runner.EvaluatorConfig,
        repeated_args: list[ms.runner.rpc_runner.T_ARGUMENT_LIST],
    ) -> list[float]:
        """With args function to run the evaluator

        Parameters
        ----------
        session : tvm.rpc.RPCSession
            The RPC session
        rt_mod: Module
            The runtime module
        device: Device
            The device to run the evaluator
        evaluator_config: EvaluatorConfig
            The evaluator config
        repeated_args: List[T_ARGUMENT_LIST]
            The repeated arguments

        Returns
        -------
        costs: List[float]
            The evaluator results
        """
        evaluator = rt_mod.time_evaluator(
            func_name=rt_mod.entry_name,
            dev=device,
            number=evaluator_config.number,
            repeat=evaluator_config.repeat,
            min_repeat_ms=evaluator_config.min_repeat_ms,
            f_preproc=(
                "cache_flush_cpu_non_first_arg" if evaluator_config.enable_cpu_cache_flush else ""
            ),
        )

        repeated_costs: list[list[float]] = []
        for args in repeated_args:
            device.sync()
            profile_result = evaluator(*args)
            repeated_costs.append(profile_result.results)
        costs = [float(cost) for cost in itertools.chain.from_iterable(repeated_costs)]

        assert len(repeated_args) == 1, "Only support one set of arguments"
        scheduled_res = [arg.numpy() for arg in repeated_args[0]]  # type: ignore
        # fetch comparison function
        passed = check_and_run(
            ARGS.check_metric_func,
            to_tvm_tensor(original_res),
            to_tvm_tensor(scheduled_res),
        )

        print_result(
            result="pass" if passed else "wrong answer",
            original_mod=original_mod,
            scheduled_mod=scheduled_mod,
            trace=trace,
            inputs=inputs,
            original_res=original_res,
            scheduled_res=scheduled_res,
            original_run_secs=original_run_secs,
            scheduled_run_secs=costs,
        )

        return costs

    def f_with_args_alloc_argument_rpc(
        rpc_session: ms.runner.rpc_runner.RPCSession,  # pylint: disable=unused-argument
        device: tvm.runtime.Device,
        args_info: ms.runner.rpc_runner.T_ARG_INFO_JSON_OBJ_LIST,
        alloc_repeat: int,
    ) -> list[ms.runner.rpc_runner.T_ARGUMENT_LIST]:
        return f_with_args_alloc_argument_common(device, args_info, alloc_repeat)

    def f_with_args_run_evaluator_rpc(
        rpc_session: ms.runner.rpc_runner.RPCSession,  # pylint: disable=unused-argument
        rt_mod: tvm.runtime.Module,
        device: tvm.runtime.Device,
        evaluator_config: ms.runner.EvaluatorConfig,
        repeated_args: list[ms.runner.rpc_runner.T_ARGUMENT_LIST],
    ) -> list[float]:
        return f_with_args_run_evaluator_common(rt_mod, device, evaluator_config, repeated_args)

    if ARGS.rpc_config is None:
        return f_with_args_alloc_argument_common, f_with_args_run_evaluator_common
    else:
        return f_with_args_alloc_argument_rpc, f_with_args_run_evaluator_rpc


def local_build_and_run(
    mod: IRModule,
    target: Target,
    device: tvm.runtime.Device,
    inputs: list[np.ndarray],
) -> tuple[list[np.ndarray], list[float]]:
    """Build and run the module locally.

    Parameters
    ----------
    mod: IRModule
        The module to build and run
    target: Target
        The target to build the module
    device: Device
        The device to run the module
    inputs: List[np.ndarray]
        The inputs to run the module

    Returns
    -------
    res: List[np.ndarray]
        The results of running the module
    run_secs: List[float]
        The running time of running the module
    """
    # potential memory leak https://github.com/apache/tvm/issues/11096
    lib = tvm.compile(mod, target=target)
    tvm_inputs = [tvm.runtime.tensor(inp, device=device) for inp in inputs]
    device.sync()
    func = lib.time_evaluator(lib.entry_name, dev=device, number=ARGS.number, repeat=ARGS.repeat)
    benchmark_res = func(*tvm_inputs)
    device.sync()
    return [arg.numpy() for arg in tvm_inputs], list(benchmark_res.results)


def _check_builder_result(builder_result: ms.builder.BuilderResult) -> None:
    """Check if the builder result is defined.

    Parameters
    ----------
    builder_result: BuilderResult
        The builder result
    """
    assert builder_result.error_msg is None, "Builder failed: " + str(
        builder_result.error_msg if builder_result.error_msg else "Empty error message"
    )


def _apply_trace(mod: IRModule, trace: Trace) -> IRModule:
    """Apply the trace to the module.

    Parameters
    ----------
    mod: IRModule
        The module to apply the trace to
    trace: Trace
        The trace to apply

    Returns
    -------
    mod: IRModule
        The module with the trace applied
    """
    sch = Schedule(mod)
    trace.apply_to_schedule(sch, remove_postproc=False)
    return sch.mod


def _build_all_mods(
    mods: list[IRModule], builder: ms.builder.Builder, target: Target
) -> list[ms.builder.BuilderResult]:
    """Build all the modules.

    Parameters
    ----------
    mods: List[IRModule]
        The modules to build
    builder: Builder
        The builder to build the modules
    target: Target
        The target to build the modules

    Returns
    -------
    builder_results: List[BuilderResult]
        The builder results
    """
    builder_results = builder.build([ms.builder.BuilderInput(mod, target) for mod in mods])
    assert len(builder_results) == len(mods), (
        f"Unexpected number of build results, expected {len(mods)} got {len(builder_results)}"
    )
    return builder_results


def _run_single_mod(
    builder_result: ms.builder.BuilderResult,
    runner: ms.runner.Runner,
    dev_type: str,
) -> None:
    """Run a single module.

    Parameters
    ----------
    builder_result: BuilderResult
        The builder result
    runner: Runner
        The runner to run the module
    dev_type: str
        The device type
    """
    runner_futures = runner.run(
        # arginfo is not used in this case so we can pass an empty list
        [ms.runner.RunnerInput(builder_result.artifact_path, device_type=dev_type, args_info=[])]
    )
    assert len(runner_futures) == 1, (
        f"Unexpected number of runner futures, expected 1 got {len(runner_futures)}"
    )
    (runner_future,) = runner_futures  # pylint: disable=unbalanced-tuple-unpacking
    runner_res = runner_future.result()
    assert runner_res.error_msg is None, "Runner failed: " + (
        runner_res.error_msg if runner_res.error_msg else "Empty error message"
    )


def main():
    """Main function"""
    describe()
    with ms.Profiler() as profiler:
        # initialize
        target = ARGS.target
        dev_type = get_device_type(target)
        builder = ms.builder.LocalBuilder()
        database = ms.database.create(work_dir=ARGS.work_dir)

        # collect records
        with profiler.timeit("collect records"):
            records = database.get_all_tuning_records()
        total = len(records)
        print(
            f"Total {total} records to be validated. "
            f"Collected in {float(profiler.get()['collect records']): 3.3f} sec."
        )

        # collect unique original TIR
        with profiler.timeit("deduplicate records"):
            workloads = set()
            for record in records:
                workloads.add(OriginalModule(record.workload.mod))
        print(
            f"Total {len(workloads)} unique original TIR to validate. "
            f"Deduplicated in {float(profiler.get()['deduplicate records']): 3.3f} sec."
        )
        if ARGS.top_k < 10**9:
            print(f"Top {ARGS.top_k} records for each original TIR will be validated.")
            total = len(workloads) * ARGS.top_k
        print()

        # validate correctness
        counter = 0
        for item in workloads:
            original_mod = item.mod
            records = database.get_top_k(
                workload=database.commit_workload(original_mod), top_k=ARGS.top_k
            )
            if len(records) < ARGS.top_k:
                total -= ARGS.top_k - len(records)
            inputs = to_numpy(check_and_run(ARGS.input_generator_func, original_mod))
            original_res, original_run_secs = local_build_and_run(
                original_mod,
                target=ARGS.baseline_target,
                inputs=inputs,
                device=get_runtime_device(ARGS.baseline_target),
            )
            scheduled_mods = [_apply_trace(original_mod, record.trace) for record in records]
            builder_results = _build_all_mods(scheduled_mods, builder, target)  # type: ignore
            for i, record in enumerate(records):
                counter += 1
                print_result = print_with_counter_func(counter=counter, total=total)
                if is_failed_record(record):
                    # skip failed records where run_secs is 1e9
                    # these records are only negative samples for cost model
                    print_result(result="skip")
                    continue
                try:
                    # prepare scheduled module
                    scheduled_mod = scheduled_mods[i]
                    # check build result
                    builder_result = builder_results[i]
                    _check_builder_result(builder_result)
                    # fetch functions
                    (
                        f_with_args_alloc_argument,
                        f_with_args_run_evaluator,
                    ) = make_alloc_arg_and_check(
                        inputs,
                        original_mod,
                        scheduled_mod,
                        str(record.trace),
                        original_res=original_res,
                        original_run_secs=original_run_secs,
                        print_result=print_result,
                    )
                    # create runner
                    evaluator_config = ms.runner.EvaluatorConfig(
                        number=ARGS.number,
                        repeat=ARGS.repeat,
                        min_repeat_ms=ARGS.min_repeat_ms,
                        enable_cpu_cache_flush=ARGS.cpu_flush,
                    )
                    if ARGS.rpc_config is not None:
                        runner: ms.Runner = ms.runner.RPCRunner(  # type: ignore
                            ARGS.rpc_config,
                            evaluator_config=evaluator_config,
                            alloc_repeat=1,
                            f_alloc_argument=f_with_args_alloc_argument,
                            f_run_evaluator=f_with_args_run_evaluator,
                            initializer=initializer,
                        )
                    else:
                        runner: ms.Runner = ms.runner.LocalRunner(  # type: ignore
                            evaluator_config=evaluator_config,
                            alloc_repeat=1,
                            f_alloc_argument=f_with_args_alloc_argument,
                            f_run_evaluator=f_with_args_run_evaluator,
                            initializer=initializer,
                        )

                    # run and validate
                    _run_single_mod(builder_result, runner, dev_type)  # type: ignore
                except Exception as e:  # pylint: disable=broad-except, invalid-name
                    # validation failed with exception
                    print_result(
                        result="exception",
                        original_mod=original_mod,
                        scheduled_mod=scheduled_mod,
                        trace=str(record.trace),
                        exception=e,
                    )
                # clean up
                remove_build_dir(builder_result.artifact_path)
    print(f"Validation finished! Total time spent: {float(profiler.get()['Total']): 3.3f} sec.")


if __name__ == "__main__":
    main()
