# 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: E722
"""Common schedule strategies for TIR."""

from collections.abc import Callable

from tvm import s_tir

from ..analysis import SBlockInfo


def try_inline(
    sch: s_tir.Schedule,
    blocks: list[SBlockInfo],
) -> list[SBlockInfo]:
    """Try to inline as many blocks as possible, and return the remaining blocks.

    Parameters
    ----------
    sch : s_tir.Schedule
        The TIR schedule used to inline blocks.
    blocks : List[SBlockInfo]
        The blocks to be inlined.

    Returns
    -------
    remaining : List[SBlockInfo]
        The remaining blocks that cannot be inlined.
    """

    def _trial(func: Callable):
        for i, block in enumerate(blocks):
            try:
                func(block.block_rv)
            except:  # pylint: disable=bare-except
                continue
            return i
        return None

    while True:
        i = _trial(sch.compute_inline)
        if i is None:
            i = _trial(sch.reverse_compute_inline)
        if i is None:
            break
        blocks.pop(i)
    return blocks


def try_inline_contiguous_spatial(
    sch: s_tir.Schedule,
    block_infos: list[SBlockInfo],
) -> list[SBlockInfo]:
    """Try to inline contiguous spatial blocks in a schedule

    Parameters
    ----------
    sch : s_tir.Schedule
        The TIR schedule used to inline blocks.
    block_infos : List[SBlockInfo]
        The blocks to be try.

    Returns
    -------
    remaining : List[SBlockInfo]
        The remaining blocks that cannot be inlined.
    """

    if block_infos is None:
        return None
    results = []
    spatial_blocks = []
    block: SBlockInfo
    for block in block_infos:
        if block.is_injective():
            spatial_blocks.append(block)
        elif spatial_blocks:
            results.extend(try_inline(sch, spatial_blocks))
            results.append(block)
            spatial_blocks = []
        else:
            results.append(block)
    if spatial_blocks:
        results.extend(try_inline(sch, spatial_blocks))
    return results