# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict
# pyre-ignore-all-errors[6]
# pyre-ignore-all-errors[16]
from __future__ import annotations

import copy

import math
import typing
from typing import Dict, List, Optional, Tuple, Union

import executorch.exir.schema as schema
import torch
from executorch.exir.error import internal_assert
from executorch.exir.schema import ScalarType, TensorShapeDynamism
from executorch.exir.sym_util import eval_shape


class AddressSpaceOverflowException(Exception):
    pass


def num_bytes_from_shape_and_dtype(shape: torch.Size, dtype: torch.dtype) -> int:
    """
    Assume the tensor is a contiguous one.
    """

    return math.prod(shape) * torch._utils._element_size(dtype)


def contiguous_stride_from_shape(shape: torch.Size) -> Tuple[int]:
    strides = []
    accum = 1
    for sz in reversed(shape):
        strides.append(accum)
        # For sizes[i] == 0, treat it as 1 to be consistent with core Pytorch
        # This preserves the PT equivalent behavior for dims with 0 elements
        if isinstance(sz, int):
            if sz != 0:
                accum *= sz
        else:
            # Unbacked symints may error on the != 0 check
            accum *= sz
    return tuple(reversed(strides))


def dim_order_from_stride(stride: Tuple[int]) -> Tuple[bytes]:
    """
    Dimension order represents how dimensions are laid out in memory,
    starting from the outer-most to the inner-most dimension.
    Thus, the conversion from strides is done by sorting the strides
    from larger to smaller since the dimension with the largest stride
    is the outer-most and the dimension with the smallest stride is the inner-most.
    For example, tensor with sizes = (3, 5, 2) and strides = (5, 1, 15), implies
    dimension order of (2, 0, 1). Dimension order of (2, 0, 1) can be obtained
    by sorting strides from large to smaller.

    When strides do not convey dimension order unambiguously, dimension order
    returned is dependent on stability of sort. In python same key elements are kept
    in original order. Thus when strides = (4, 3, 1, 1) returned value is (0, 1, 2, 3)
    Another example is: sizes = (1, 3, 1, 1) with strides = (3, 1, 3, 3), returned
    value is (0, 2, 3, 1)
    """
    for _, s in enumerate(stride):
        if s == 0:
            raise ValueError("0 in strides is not supported for ExecuTorch.")
    sorted_dims = [
        i[0] for i in sorted(enumerate(stride), key=lambda x: x[1], reverse=True)
    ]
    return tuple(typing.cast(Tuple[bytes], sorted_dims))


def stride_from_dim_order(sizes: List[int], dim_order: List[bytes]) -> List[int]:
    """
    Converts dim order to stride using sizes
    e.g. if sizes = (2, 3, 4) and dim_order = (0, 1, 2) then strides = (12, 4, 1)
    while for the same size if dim_order = (0, 2, 1) then strides = (12, 1, 3)
    See executorch/runtime/core/exec_aten/util/dim_order_util.h for details
    Args:
        sizes (Tuple[int]): sizes of the tensor
        dim_order (Tuple[bytes]): dim order of the tensor
    Returns:
        Tuple[int]: stride
    """
    if len(sizes) == 0:
        return []
    strides = copy.deepcopy(sizes)
    ndim = len(sizes)
    strides[dim_order[ndim - 1]] = 1
    for i in range(ndim - 2, -1, -1):
        if sizes[dim_order[i + 1]] == 0:
            strides[dim_order[i]] = strides[dim_order[i + 1]]
        else:
            strides[dim_order[i]] = sizes[dim_order[i + 1]] * strides[dim_order[i + 1]]
    return strides


def calculate_aligned_num_bytes(num: int, alignment: int) -> int:
    return math.ceil(num / alignment) * alignment


def determine_tensor_dynanism(shape: torch.Size) -> TensorShapeDynamism:
    if all(isinstance(s, int) for s in shape):
        return TensorShapeDynamism.STATIC
    else:
        try:
            _ = eval_shape(shape)
            return TensorShapeDynamism.DYNAMIC_BOUND
        except torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode:
            return TensorShapeDynamism.DYNAMIC_UNBOUND


ALIGNMENT = 16


class TensorSpec:
    """
    Captures the metadata for a given Tensor (ex. scalar type, storage, etc.).
    """

    def __init__(
        self,
        dtype: torch.dtype,
        shape: torch.Size,
        layout: torch.layout = torch.strided,
        is_sparse: bool = False,
        const: bool = False,
        requires_grad: bool = False,
    ) -> None:
        self.scalar_type = dtype
        self.const = const
        self.alignment: int = ALIGNMENT
        self.storage: Optional[torch.UntypedStorage] = None
        # convert to list making it easier to handle type checking
        self.shape: List[int] = list(shape)
        self.stride: Tuple[int] = contiguous_stride_from_shape(shape)
        self.dim_order: Tuple[bytes] = dim_order_from_stride(self.stride)
        self.requires_grad = requires_grad
        self.layout = layout
        self.is_sparse = is_sparse
        self.init_mem_planning_fields()
        self.shape_dynamism: TensorShapeDynamism = determine_tensor_dynanism(self.shape)

    @property
    def allocated_memory(self) -> int:
        nbytes = num_bytes_from_shape_and_dtype(self.shape, self.dtype)
        return calculate_aligned_num_bytes(nbytes, self.alignment)

    def realign(self, new_alignment: int) -> int:
        self.alignment = new_alignment
        return self.allocated_memory

    def nbytes(self) -> int:
        return num_bytes_from_shape_and_dtype(self.shape, self.dtype)

    @classmethod
    def from_tensor(cls, tensor: torch.Tensor, const: bool = False) -> TensorSpec:
        if const:
            # for non-contigous tensors, convert to a contiguous one
            tensor = tensor.contiguous()
            # Weights cannot be views during emission or serialization
            if tensor.nbytes != tensor.untyped_storage().nbytes():
                tensor = tensor.clone()

        spec = cls(
            dtype=tensor.dtype,
            shape=tensor.shape,
            layout=tensor.layout,
            const=const,
            is_sparse=tensor.is_sparse,
        )
        spec.stride = tensor.stride()
        spec.dim_order = dim_order_from_stride(spec.stride)
        spec.requires_grad = tensor.requires_grad
        spec.storage = tensor.untyped_storage() if const else None

        return spec

    def init_mem_planning_fields(self) -> None:
        self.lifetime = [None, None]
        self.mem_id = None
        self.mem_obj_id = None
        self.mem_offset = None

    @property
    def dtype(self) -> torch.dtype:
        return self.scalar_type

    @property
    def is_dynamic_shape_tensor(self) -> bool:
        return self.shape_dynamism != schema.TensorShapeDynamism.STATIC

    @property
    def is_static_shape_tensor(self) -> bool:
        return self.shape_dynamism == TensorShapeDynamism.STATIC

    @property
    def is_upper_bound_tensor(self) -> bool:
        return self.shape_dynamism == TensorShapeDynamism.DYNAMIC_BOUND

    @property
    def is_dynamic_unbound_tensor(self) -> bool:
        return self.shape_dynamism == TensorShapeDynamism.DYNAMIC_UNBOUND

    def debug(self) -> str:
        return (
            f"TensorSpec(id={id(self)}, const={self.const}, scalar_type={self.scalar_type}"
            + f", allocated_memory={self.allocated_memory}, mem_id={self.mem_id}"
            + f", mem_offset={self.mem_offset}, lifetime={self.lifetime}"
            + f", shape_dynamism={self.shape_dynamism}"
            + (f", shape={self.shape}")
            + ")"
        )

    def __repr__(self) -> str:
        """
        Round-trippable printing function
        """
        return (
            f"TensorSpec(dtype={self.scalar_type}, shape={self.shape}"
            + f", layout={self.layout}"
            + f", is_sparse={self.is_sparse}"
            + f", shape_dynamism={self.shape_dynamism}"
            + f", const={self.const}, requires_grad={self.requires_grad}"
            + ")"
        )


def memory_format_enum(memory_format: torch.memory_format) -> int:
    internal_assert(
        isinstance(memory_format, torch.memory_format),
        "We only support torch.memory_format",
    )
    table = {
        torch.contiguous_format: 0,
        torch.preserve_format: 1,
    }
    return table[memory_format]


scalar_type_table: Dict[torch.dtype, ScalarType] = {
    torch.uint8: ScalarType.BYTE,
    torch.int8: ScalarType.CHAR,
    torch.int16: ScalarType.SHORT,
    torch.int32: ScalarType.INT,
    torch.int64: ScalarType.LONG,
    torch.half: ScalarType.HALF,
    torch.float: ScalarType.FLOAT,
    torch.double: ScalarType.DOUBLE,
    torch.complex32: ScalarType.COMPLEX32,
    torch.complex64: ScalarType.COMPLEX64,
    torch.complex128: ScalarType.COMPLEX128,
    torch.bool: ScalarType.BOOL,
    torch.qint8: ScalarType.QINT8,
    torch.quint8: ScalarType.QUINT8,
    torch.qint32: ScalarType.QINT32,
    torch.bfloat16: ScalarType.BFLOAT16,
    torch.quint4x2: ScalarType.QUINT4x2,
    torch.uint16: ScalarType.UINT16,
}


enum_to_scalar_map: Dict[ScalarType, torch.dtype] = {
    scalar_type_table[key]: key for key in scalar_type_table
}


def scalar_type_enum(dtype: torch.dtype) -> ScalarType:
    # TODO (zhengxu) single source of truth from c10/core/ScalarType.h.
    internal_assert(
        isinstance(dtype, torch.dtype), "We only support dtypes defined in Pytorch Core"
    )
    return scalar_type_table[dtype]


def get_scalar_type(enum: ScalarType) -> torch.dtype:
    return enum_to_scalar_map[enum]


def layout_enum(layout: torch.layout) -> int:
    # TODO single source of truth.
    table = {
        torch.strided: 0,
        torch.sparse_coo: 1,
    }
    return table[layout]


def make_allocation_info(mem_id: int, mem_offset: int) -> schema.AllocationDetails:
    """
    Creates the allocation_details object for creating tensors
    """
    if mem_offset < 0:
        raise ValueError(f"mem_offset {mem_offset} must not be negative")
    memory_offset_low = mem_offset & ((1 << 32) - 1)
    memory_offset_high = mem_offset >> 32
    if memory_offset_high >= 1 << 32:
        raise AddressSpaceOverflowException(
            f"mem_offset {mem_offset} does not fit in 64 bits"
        )

    allocation_info = schema.AllocationDetails(
        memory_id=mem_id,
        memory_offset_low=memory_offset_low,
        memory_offset_high=memory_offset_high,
    )
    return allocation_info


def make_tensor_value(
    data_buffer_idx: int,
    allocation_info: Optional[schema.AllocationDetails],
    spec: TensorSpec,
) -> schema.Tensor:
    """
    Converts the normal torch tensor to a flatbuffer tensor.
    """

    def to_list(
        x: Union[torch.Size, int, List[int], Tuple[int]]
    ) -> Union[List[int], List[torch.Size]]:
        if isinstance(x, torch.Size) or isinstance(x, tuple):
            return list(x)
        elif isinstance(x, int):
            return [x]
        else:
            return x

    tensor_size = to_list(spec.shape)
    tensor_dim_order = to_list(spec.dim_order)

    flatbuffer_tensor = schema.Tensor(
        scalar_type=scalar_type_enum(spec.scalar_type),
        # The runtime currently only supports tensors with offsets of zero.
        storage_offset=0,
        sizes=tensor_size,
        dim_order=tensor_dim_order,
        requires_grad=spec.requires_grad,
        data_buffer_idx=data_buffer_idx,
        allocation_info=allocation_info,
        layout=layout_enum(spec.layout),
        shape_dynamism=spec.shape_dynamism,
    )
    return flatbuffer_tensor


def check_spec(tensor: torch.Tensor, spec: TensorSpec) -> None:
    internal_assert(
        tensor.is_sparse == spec.is_sparse,
        f"Tensor attribute 'is_sparse' is expected to be equal to '{spec.is_sparse}', actually got: '{tensor.is_sparse}'",
    )
    internal_assert(
        tensor.shape == spec.shape,
        f"Tensor attribute 'shape' is expected to be equal to '{spec.shape}', actually got: '{tensor.shape}'",
    )
    internal_assert(
        tensor.dtype == spec.dtype,
        f"Tensor attribute 'dtype' is expected to be equal to '{spec.dtype}', actually got: '{tensor.dtype}'",
    )