#################################################################################################
#
# Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#################################################################################################

"""
Layout algebras
"""

from pycute import Layout, composition, make_layout, flatten, product


def _infer_split(old_shape, new_shape):
    old_shape = _tuple_to_list(old_shape)
    new_shape = _tuple_to_list(new_shape)
    if len(old_shape) == 0 and len(new_shape) == 0:
        return []
    if len(old_shape) == 0:
        if product(tuple(new_shape)) != 1:
            raise ValueError("Invalid reshape size")
        else:
            return new_shape
    if len(new_shape) == 0:
        if product(tuple(old_shape)) != 1:
            raise ValueError("Invalid reshape size")
        else:
            return old_shape
    # This is done recursively by only process the last dimension at each time
    old_dim = old_shape[-1]
    new_dim = new_shape[-1]
    # Exact match
    if old_dim == new_dim:
        return _infer_split(old_shape[:-1], new_shape[:-1]) + [new_dim,]
    # Needs split
    if old_dim > new_dim and old_dim % new_dim == 0:
        residual = old_dim // new_dim
        return _infer_split(old_shape[:-1] + [residual,], new_shape[:-1]) + [new_dim,]
    # Needs merge
    if old_dim < new_dim and new_dim % old_dim == 0:
        residual = new_dim // old_dim
        return _infer_split(old_shape[:-1], new_shape[:-1] + [residual,]) + [old_dim,]

    raise NotImplementedError(f"Unsupported split: {old_shape} -> {new_shape}")

def _infer_merge(flatten_shape, shape):
    flatten_shape = _tuple_to_list(flatten_shape)
    shape = _tuple_to_list(shape)
    idx_flat = 0
    merged_shape = []
    for dim in shape:
        # Exact match
        if dim == flatten_shape[idx_flat]:
            merged_shape.append(dim)
            idx_flat += 1
        # Need group
        elif dim > flatten_shape[idx_flat] and dim % flatten_shape[idx_flat] == 0:
            residual = dim
            group = []
            while(residual > 1):
                group.append(flatten_shape[idx_flat])
                residual = residual // flatten_shape[idx_flat]
                idx_flat += 1
            merged_shape.append(group)
        else:
            raise NotImplementedError(f"Unsupported merge: {flatten_shape} -> {shape}")

    return merged_shape

def _list_to_tuple(nested_list):
    if isinstance(nested_list, list) or isinstance(nested_list, tuple):
        return tuple(_list_to_tuple(item) for item in nested_list)
    return nested_list

def _tuple_to_list(nested_tuple):
    if isinstance(nested_tuple, list) or isinstance(nested_tuple, tuple):
        return list(_tuple_to_list(item) for item in nested_tuple)
    return nested_tuple

def _reverse_tuple(nested_tuple: tuple):
    if isinstance(nested_tuple, tuple):
        return tuple([_reverse_tuple(item) for item in nested_tuple][::-1])
    return nested_tuple

def _get_first_lhs_nonzero_stride(stride_list, idx):
    for i in reversed(range(idx)):
        if stride_list[i] != 0:
            return i
    else:
        return None

def _get_first_rhs_nonzero_stride(stride_list, idx):
    for i in range(idx+1, len(stride_list)):
        if stride_list[i] != 0:
            return i
        else:
            return None

def reshape(layout, new_shape):
    """
    General reshape of input layout.
    It takes two steps:
    1. split the dimensions of the old layout
    2. merge the splitted dimensions according to the new shape
    """
    #
    # Step 1: Split the dimensions of the old layout
    #
    # 1.1 Flat old and new shape
    old_flatten_shape = list(flatten(layout.shape))
    new_flatten_shape = list(flatten(new_shape))

    # 1.2 Infer the flatten splitted shape
    splitted_flatten_shape = _infer_split(old_flatten_shape, new_flatten_shape)

    # 1.3 Unflat the splitted shape based on the old shape
    splited_shape = _infer_merge(splitted_flatten_shape, old_flatten_shape)

    # 1.4 Infer the type of each split
    # If the split type is in row-major (R), the dimension list is reversed because
    # the cute::composition only support column-major split
    split_type = []  # the type of each split (ColumnMajor or RowMajor)
    permuted_splitted_shape = []
    old_flatten_stride = list(flatten(layout.stride))
    for idx, dim in enumerate(splited_shape):
        if not isinstance(dim, list):
            permuted_splitted_shape.append(dim)
            split_type.append("C")
        else:
            lhs_stride = _get_first_lhs_nonzero_stride(old_flatten_stride, idx)
            rhs_stride = _get_first_rhs_nonzero_stride(old_flatten_stride, idx)
            # Special case for single tuple
            # Use column-major by default
            if lhs_stride is None and rhs_stride is None:
                permuted_splitted_shape.append(dim)
                split_type.append("C")
            else:
                if lhs_stride is not None and rhs_stride is not None:
                    # We consider shape[idx]:stride[idx]
                    # Case 1: stride[idx - 1] <= stride[idx] <= stride[idx + 1]: column major
                    if lhs_stride <= old_flatten_stride[idx] and old_flatten_stride[idx] <= rhs_stride:
                        permuted_splitted_shape.append(dim)
                        split_type.append("C")
                    # Case 2: stride[idx - 1] > stride[idx] > stride[idx + 1]: row major
                    elif lhs_stride > old_flatten_stride[idx] and old_flatten_stride[idx] > rhs_stride:
                        permuted_splitted_shape.append([d for d in reversed(dim)])
                        split_type.append("R")
                    # Case 3: stride[idx - 1] <= stride[idx] > stride[idx + 1]: concave
                    elif lhs_stride <= old_flatten_stride[idx] and old_flatten_stride[idx] > rhs_stride:
                        if lhs_stride >= rhs_stride:
                            permuted_splitted_shape.append(dim)
                            split_type.append("C")
                        else:
                            permuted_splitted_shape.append([d for d in reversed(dim)])
                            split_type.append("R")
                    # Case 4: stride[idx - 1] > stride[idx] <= stride[idx + 1]: concave
                    elif lhs_stride > old_flatten_stride[idx] and old_flatten_stride[idx] <= rhs_stride:
                        if lhs_stride >= rhs_stride:
                            permuted_splitted_shape.append(dim)
                            split_type.append("C")
                        else:
                            permuted_splitted_shape.append([d for d in reversed(dim)])
                            split_type.append("R")
                    else:
                        raise NotImplementedError()
                elif lhs_stride is None:
                    # Case 1: dim's stride < dim+1's stride, expand in column major
                    if old_flatten_stride[idx] > rhs_stride:
                        permuted_splitted_shape.append([d for d in reversed(dim)])
                        split_type.append("R")
                    else:
                        permuted_splitted_shape.append(dim)
                        split_type.append("C")
                else:
                    # Case 1: dim's stride > dim-1's stride
                    if old_flatten_stride[idx] < lhs_stride:
                        permuted_splitted_shape.append([d for d in reversed(dim)])
                        split_type.append("R")
                    else:
                        permuted_splitted_shape.append(dim)
                        split_type.append("C")

    # 1.4 Generate the splitted layout
    permuted_splitted_layout = composition(layout, Layout(_list_to_tuple(permuted_splitted_shape)))

    # 1.5 Reverse the permutation in 1.4 before merge
    splitted_shape = []
    splitted_stride = []
    for shape_dim, stride_dim, type in zip(
            permuted_splitted_layout.shape,
            permuted_splitted_layout.stride,
            split_type):
        if type == "C":
            splitted_shape.append(shape_dim)
            splitted_stride.append(stride_dim)
        else:
            splitted_shape.append(tuple([d for d in reversed(shape_dim)]))
            splitted_stride.append(tuple([d for d in reversed(stride_dim)]))
    splitted_layout = Layout(tuple(splitted_shape), tuple(splitted_stride))


    #
    # Step 2: Merge the splitted dimensions according to the new shape
    #
    # 2.1 Merge layout
    merged_layout = composition(splitted_layout, Layout(new_shape))

    # 2.2 Cleaning up
    output_layout = composition(merged_layout, Layout(new_shape))
    return output_layout


def permutation(layout, permutation):
    """
    Permute the layout
    """
    new_shape = tuple([layout.shape[idx] for idx in permutation])
    new_stride = tuple([layout.stride[idx] for idx in permutation])
    return Layout(new_shape, new_stride)


def _broadcast(layout, new_shape):
    if len(layout) == 1 and isinstance(new_shape, int):
        old_dim = layout.shape
        old_stride = layout.stride
        new_dim = new_shape
        if old_dim == new_dim:
            return Layout(old_dim, old_stride)
        elif old_dim == 1:
            return Layout(new_dim, 0)
        else:
            raise NotImplementedError(f"Invalid Broadcast: {old_dim} -> {new_dim}")

    # Align the dimensions
    old_shape = layout.shape
    if isinstance(old_shape, int):
        old_shape = (old_shape,)
        sub_layouts = [layout,]
    else:
        sub_layouts = [sub_layout for sub_layout in layout]
    rhs_broadcast_layouts = [Layout(1, 0)] * (len(new_shape) - len(old_shape))
    # Get the broadcasted layout
    broadcast_layouts = []
    try:
        layout = make_layout(*sub_layouts, *rhs_broadcast_layouts)
        broadcast_layouts = []
        for idx, sub_layout in enumerate(layout):
            broadcast_layouts.append(_broadcast(sub_layout, new_shape[idx]))
    except NotImplementedError:
        layout = make_layout(*rhs_broadcast_layouts, *sub_layouts)
        for idx, sub_layout in enumerate(layout):
            broadcast_layouts.append(_broadcast(sub_layout, new_shape[idx]))
    return make_layout(*broadcast_layouts)


def broadcast(layout, new_shape):
    """
    Broadcast the new layout based on the input shape
    The broadcasted shape equals to the new shape
    The stride of broadcasted dimensions are 0
    """
    return _broadcast(layout, new_shape)


def debroadcast(layout, dims):
    """
    Squeeze the 0-stride
    """
    for dim in dims:
        if layout.stride[dim] != 0:
            raise ValueError(f"Dim{dim} cannot be debroadcasted as it has stride {layout.stride[dim]}")
    new_shape = tuple([s for idx, s in enumerate(layout.shape) if idx not in dims])
    new_stride = tuple([s for idx, s in enumerate(layout.stride) if idx not in dims])
    return Layout(new_shape, new_stride)


def canonicalization_(shapes, strides):
    if isinstance(shapes, tuple):
        c_shapes = []
        c_strides = []
        for shape, stride in zip(shapes, strides):
            c_shape, c_stride = canonicalization_(shape, stride)
            c_shapes.append(c_shape)
            c_strides.append(c_stride)
        return tuple(c_shapes), tuple(c_strides)
    else:
        if shapes == 1:
            return 1, 0
        else:
            return shapes, strides

def canonicalization(layout):
    """
    Canonicalize the input layout
    1. set the stride of shape "1" to 0
    """
    new_shape, new_stride = canonicalization_(layout.shape, layout.stride)
    return Layout(new_shape, new_stride)