# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
import inspect
import operator
from itertools import accumulate, zip_longest
from typing import Sequence, Union, Optional

import numpy as np

from ._shaped_value import ShapedValue, _indices_to_indexer, _maybe_compute_size
from .arith import ScalarValue, constant, index_cast, ConstantOp
from .tensor import compute_result_shape_reassoc_list
from numba_cuda_mlir._mlir.extras import types as T
from numba_cuda_mlir._mlir.extras.meta import region_op
from ..util import (
    _get_sym_name,
    get_user_code_loc,
    infer_mlir_type,
)
from numba_cuda_mlir._mlir._mlir_libs._mlir import register_value_caster
from numba_cuda_mlir._mlir.dialects import memref, arith, vector, builtin
from numba_cuda_mlir._mlir.dialects._ods_common import (
    get_op_result_or_op_results,
    MixedValues,
    _dispatch_mixed_values,
)
from numba_cuda_mlir._mlir.dialects.memref import (
    _is_static_int_like,
    _generated_subview,
    _is_constant_int_like,
)
from numba_cuda_mlir._mlir.dialects.memref import *
from numba_cuda_mlir._mlir.ir import (
    DenseElementsAttr,
    MemRefType,
    UnrankedMemRefType,
    ShapedType,
    IndexType,
    Type,
    Value,
    VectorType,
    SymbolTable,
    InsertionPoint,
    StridedLayoutAttr,
)

S = ShapedType.get_dynamic_size()


def __alloc(
    op_ctor,
    sizes: Sequence[Union[int, Value]],
    element_type: Type,
    memory_space=None,
    alignment=None,
    loc=None,
    ip=None,
):
    dynamic_sizes = []
    memref_shape = []
    for s in sizes:
        if isinstance(s, int):
            memref_shape.append(s)
        else:
            memref_shape.append(ShapedType.get_dynamic_size())
            dynamic_sizes.append(s)
    result_type = T.memref(*memref_shape, element_type=element_type, memory_space=memory_space)

    symbol_operands = []
    return get_op_result_or_op_results(
        op_ctor(
            result_type,
            dynamic_sizes,
            symbol_operands,
            alignment=alignment,
            loc=loc,
            ip=ip,
        )
    )


_alloc = alloc


def alloc(
    sizes: Union[int, Value],
    element_type: Type = None,
    *,
    memory_space=None,
    alignment=None,
    loc=None,
    ip=None,
):
    return __alloc(
        AllocOp,
        sizes,
        element_type,
        memory_space=memory_space,
        alignment=alignment,
        loc=loc,
        ip=ip,
    )


_alloca = alloca


def alloca(
    sizes: Union[int, Value],
    element_type: Type = None,
    memory_space=None,
    alignment=None,
    loc=None,
    ip=None,
):
    return __alloc(
        AllocaOp,
        sizes,
        element_type,
        memory_space=memory_space,
        alignment=alignment,
        loc=loc,
        ip=ip,
    )


def load(memref: Value, indices: Sequence[Union[Value, int]], *, loc=None, ip=None):
    indices = list(indices)
    for idx, i in enumerate(indices):
        if isinstance(i, int):
            indices[idx] = constant(i, index=True)
        elif isinstance(i, Value):
            if not isinstance(i.type, IndexType):
                i = index_cast(i, to=IndexType.get(), loc=loc, ip=ip)
                indices[idx] = i
        else:
            raise TypeError(f"expected {i=} to be either int or Value")
    return get_op_result_or_op_results(LoadOp(memref, indices, loc=loc, ip=ip))


def store(
    value: Value,
    memref: Value,
    indices: Sequence[Union[Value, int]],
    *,
    loc=None,
    ip=None,
):
    indices = list(indices)
    for idx, i in enumerate(indices):
        if isinstance(i, int):
            indices[idx] = constant(i, index=True)
        elif isinstance(i, Value):
            if not isinstance(i.type, IndexType):
                i = index_cast(i, to=IndexType.get(), loc=loc, ip=ip)
                indices[idx] = i
        else:
            raise TypeError(f"expected {i=} to be either int or Value")
    return get_op_result_or_op_results(StoreOp(value, memref, indices, loc=loc, ip=ip))


rank_reduce = object()


@register_value_caster(MemRefType.static_typeid)
@ShapedValue
class MemRefValue(Value):
    def __str__(self):
        return f"{self.__class__.__name__}({self.get_name()}, {self.type})"

    def __repr__(self):
        return str(self)

    def __getitem__(self, idx: tuple) -> "MemRefValue":
        loc = get_user_code_loc()

        if not self.has_rank():
            raise ValueError("only ranked memref slicing/indexing supported")

        if idx == Ellipsis or idx == slice(None):
            return self
        if isinstance(idx, tuple) and all(i == slice(None) for i in idx):
            return self
        if idx is None:
            return expand_shape(self, (0,), loc=loc)

        idx = list((idx,) if isinstance(idx, (int, ScalarValue, slice)) else idx)
        should_rank_reduce = rank_reduce in idx
        if should_rank_reduce:
            idx.remove(rank_reduce)

        for i, d in enumerate(idx):
            if isinstance(d, int):
                idx[i] = constant(d, index=True, loc=loc)

        if all(isinstance(d, ScalarValue) for d in idx) and len(idx) == len(self.shape):
            return load(self, idx, loc=loc)
        else:
            return _subview(self, tuple(idx), rank_reduce=should_rank_reduce, loc=loc)

    def __setitem__(self, idx, val):
        loc = get_user_code_loc()

        if not self.has_rank():
            raise ValueError("only ranked memref slicing/indexing supported")

        idx = list((idx,) if isinstance(idx, (ScalarValue, int, Value)) else idx)
        for i, d in enumerate(idx):
            if isinstance(d, int):
                idx[i] = constant(d, index=True, loc=loc)

        if all(isinstance(d, ScalarValue) for d in idx) and len(idx) == len(self.shape):
            if isinstance(val, (int, float)):
                # TODO: this is an unchecked conversion
                val = ScalarValue(val, dtype=self.dtype)
            if isinstance(val, ScalarValue):
                store(val, self, idx, loc=loc)
            elif isinstance(val.type, VectorType):
                return vector.StoreOp(
                    valueToStore=val,
                    base=self,
                    indices=idx,
                    loc=loc,
                )
            else:
                raise ValueError(f"unsupported value type for memref store: {val=}")
        else:
            _copy_to_subview(self, val, tuple(idx), loc=loc)

    def dim(self, idx, *, loc=None, ip=None):
        if self.has_rank():
            if isinstance(idx, Value) and isinstance(idx.owner.opview, ConstantOp):
                idx = idx.owner.opview.value.value
            if not isinstance(idx, int):
                raise TypeError(f"expected {idx=} to be an int")
            d = self.shape[idx]
            if d != S:
                return d
        if isinstance(idx, int):
            idx = constant(idx, IndexType.get())
        if not isinstance(idx.type, IndexType):
            idx = index_cast(idx, to=IndexType.get(), loc=loc, ip=ip)
        return DimOp(self, idx, loc=loc, ip=ip).result

    def dims(self, *, rank=None, loc=None, ip=None):
        if self.has_rank():
            return tuple(self.dim(i, loc=loc, ip=ip) for i in range(self.rank))
        if rank is None:
            raise ValueError("unranked dims requires rank for unranked memref")
        return tuple(self.dim(i, loc=loc, ip=ip) for i in range(rank))


def expand_shape(
    inp,
    newaxis_dims,
    *,
    loc=None,
    ip=None,
) -> MemRefValue:
    """Expand the shape of a memref.

    Insert a new axis that will appear at the `axis` position in the expanded
    memref shape.

    Args:
      inp: Input memref-like.
      axis: Position in the expanded axes where the new axis (or axes) is placed.

    Returns:
       View of `a` with the number of dimensions increased.

    """

    if len(newaxis_dims) == 0:
        return inp

    result_shape, reassoc_list = compute_result_shape_reassoc_list(inp.shape, newaxis_dims)

    return MemRefValue(
        memref.expand_shape(
            T.memref(*result_shape, inp.dtype),
            inp,
            reassoc_list,
            output_shape=[],
            static_output_shape=result_shape,
            loc=loc,
            ip=ip,
        )
    )


def _infer_memref_subview_result_type(source_memref_type, offsets, sizes, strides):
    source_strides, source_offset = source_memref_type.get_strides_and_offset()
    # "canonicalize" from tuple|list -> list
    offsets, sizes, strides, source_strides = map(list, (offsets, sizes, strides, source_strides))

    if any(not _is_static_int_like(i) for i in offsets + [source_offset]):
        target_offset = ShapedType.get_dynamic_size()
    else:
        target_offset = source_offset
        for offset, target_stride in zip(offsets, source_strides):
            target_offset += offset * target_stride

    target_strides = []
    for source_stride, static_stride in zip(source_strides, strides):
        target_strides.append(source_stride * static_stride)

    if all(isinstance(s, int) for s in sizes):
        # If default striding then no need to complicate things for downstream ops (e.g., expand_shape).
        default_strides = list(accumulate(sizes[1:][::-1], operator.mul))[::-1] + [1]
    else:
        default_strides = None
        sizes = [s if isinstance(s, int) else ShapedType.get_dynamic_size() for s in sizes]

    if target_strides == default_strides and target_offset == 0:
        layout = None
    else:
        layout = StridedLayoutAttr.get(target_offset, target_strides)

    return MemRefType.get(
        sizes,
        source_memref_type.element_type,
        layout,
        source_memref_type.memory_space,
    )


def subview(
    source: Value,
    offsets: MixedValues,
    sizes: MixedValues,
    strides: MixedValues,
    *,
    rank_reduce=False,
    result_type: Optional[MemRefType] = None,
    loc=None,
    ip=None,
):
    if offsets is None:
        offsets = []
    if sizes is None:
        sizes = []
    if strides is None:
        strides = []

    for s in [offsets, sizes, strides]:
        for idx, i in enumerate(s):
            if _is_constant_int_like(i):
                s[idx] = i.owner.opview.literal_value

    if result_type is None:
        # If any are arith.constant results then this will canonicalize to python int
        # (which can then be used to fully specify the subview).
        result_type = _infer_memref_subview_result_type(source.type, offsets, sizes, strides)

    offsets, _packed_offsets, static_offsets = _dispatch_mixed_values(offsets)
    sizes, _packed_sizes, static_sizes = _dispatch_mixed_values(sizes)
    strides, _packed_strides, static_strides = _dispatch_mixed_values(strides)

    if rank_reduce:
        result_shape = list(result_type.shape)
        layout_strides = None
        if result_type.layout:
            layout_strides = result_type.layout.strides
        for i, (s, ss) in reversed(list(enumerate(list(zip_longest(sizes, static_sizes))))):
            if (s is not None and _is_static_int_like(s) and s.literal_value == 1) or ss == 1:
                del result_shape[i]
                if layout_strides is not None:
                    del layout_strides[i]
        reduced_layout = None
        if layout_strides is not None:
            reduced_layout = StridedLayoutAttr.get(result_type.layout.offset, layout_strides)
        result_type = MemRefType.get(
            result_shape,
            result_type.element_type,
            reduced_layout,
            result_type.memory_space,
        )

    return _generated_subview(
        result_type,
        source,
        offsets,
        sizes,
        strides,
        static_offsets,
        static_sizes,
        static_strides,
        loc=loc,
        ip=ip,
    )


def _subview(
    mem: MemRefValue,
    idx,
    *,
    rank_reduce=False,
    loc=None,
    ip=None,
) -> MemRefValue:
    indexer = _indices_to_indexer(idx, mem.shape)
    out = mem

    if indexer.is_constant():
        offsets = indexer.static_offsets()
        sizes = indexer.static_sizes()
        strides = indexer.static_strides()
    else:
        # special tile case
        offsets = [None] * len(indexer.in_shape)
        sizes = [None] * len(indexer.in_shape)
        strides = [None] * len(indexer.in_shape)
        for i, ind in enumerate(indexer.indices):
            if isinstance(ind, slice):
                maybe_size = _maybe_compute_size(ind.start, ind.stop, ind.step)
                if maybe_size is None:
                    raise RuntimeError(f"failed to canonicalize start, stop, step: {ind=}")
                offsets[i] = ind.start
                sizes[i] = maybe_size
                strides[i] = (
                    ind.step.literal_value if isinstance(ind.step, ScalarValue) else ind.step
                )
            elif isinstance(ind, Value):
                offsets[i] = ind
                sizes[i] = 1
                strides[i] = 1
            else:
                raise RuntimeError(f"indexing of {mem=} not supported by {ind=}")
        assert all(map(lambda x: x is not None, offsets + sizes + strides)), (
            f"not each slice is statically known: {indexer.indices}"
        )

    out = subview(
        out,
        offsets=offsets,
        sizes=sizes,
        strides=strides,
        rank_reduce=rank_reduce,
        loc=loc,
        ip=ip,
    )

    # This adds newaxis/None dimensions.
    return expand_shape(out, indexer.newaxis_dims, loc=loc, ip=ip)


def _copy_to_subview(
    dest: MemRefValue,
    source: MemRefValue,
    idx,
    *,
    loc=None,
    ip=None,
):
    if isinstance(source, ScalarValue):
        source = expand_shape(source, (0,), loc=loc, ip=ip)

    dest_subview = _subview(dest, idx, loc=loc, ip=ip)
    assert dest_subview.shape == source.shape, (
        f"Expected matching shape for dest subview {dest_subview.shape} and source {source.shape=}"
    )

    return memref.copy(source, dest_subview, loc=loc, ip=ip)


alloca_scope = region_op(AllocaScopeOp)

_dim = dim


def dim(source, index, *, loc=None, ip=None):
    if isinstance(index, int):
        index = constant(index, index=True)
    return _dim(source=source, index=index, loc=loc, ip=ip)


def global_(
    initial_value=None,
    sym_name=None,
    type=None,
    sym_visibility="private",
    constant=None,
    alignment=None,
    loc=None,
    ip=None,
):
    if sym_name is None:
        previous_frame = inspect.currentframe().f_back
        sym_name = _get_sym_name(previous_frame, check_func_call="memref\\.global_|global_")
        assert sym_name is not None, "couldn't automatically find sym_name in previous frame"
    if initial_value is None:
        assert type is not None
    else:
        assert isinstance(initial_value, np.ndarray)
        if type is None:
            type = infer_mlir_type(initial_value, memref=True)
        initial_value = DenseElementsAttr.get(
            initial_value,
            type=type.element_type,
            context=None,
        )
        constant = True

    return memref.global_(
        sym_name,
        type,
        sym_visibility=sym_visibility,
        initial_value=initial_value,
        constant=constant,
        alignment=alignment,
        loc=loc,
        ip=ip,
    ).opview


def view(source, shape, dtype=None, shift=0, memory_space=None, loc=None, ip=None):
    if dtype is None:
        dtype = source.type.element_type
    byte_width_dtype = dtype.width // 8
    byte_shift = shift * byte_width_dtype
    if isinstance(byte_shift, int):
        byte_shift = constant(byte_shift, index=True)
    elif isinstance(byte_shift, Value):
        if not isinstance(byte_shift.type, IndexType):
            byte_shift = index_cast(byte_shift, to=IndexType.get(), loc=loc, ip=ip)
    else:
        raise TypeError(f"expected {byte_shift=} to be either int or Value")
    assert isinstance(byte_shift.type, IndexType), "expected index type for byte-shift"
    if memory_space is None and source:
        memory_space = source.type.memory_space

    dynamic_sizes = []
    memref_shape = []
    for s in shape:
        if isinstance(s, int):
            memref_shape.append(s)
        else:
            memref_shape.append(ShapedType.get_dynamic_size())
            if not isinstance(s.type, IndexType):
                s = index_cast(s, to=IndexType.get(), loc=loc, ip=ip)
            dynamic_sizes.append(s)

    return memref.view(
        T.memref(*memref_shape, element_type=dtype, memory_space=memory_space),
        source,
        byte_shift,
        dynamic_sizes,
        loc=loc,
        ip=ip,
    )


_get_global = get_global


def get_global(name_or_global, *, name=None, global_=None, result=None, loc=None, ip=None):
    if isinstance(name_or_global, GlobalOp):
        global_ = name_or_global
    elif isinstance(name_or_global, str):
        name = name_or_global
    elif name_or_global is not None:
        raise ValueError(f"only string or GlobalOp can be provided; got {name_or_global}")

    if global_ is None:
        assert name is not None, "name must be provided"

        if result is None:

            def callback(symbol_table_op, _uses_visible):
                nonlocal global_
                sym_table = SymbolTable(symbol_table_op)
                if name in sym_table:
                    global_ = sym_table[name]

            current_owner = InsertionPoint.current.block.owner
            while not isinstance(current_owner.opview, builtin.ModuleOp):
                current_owner = current_owner.parent
            SymbolTable.walk_symbol_tables(current_owner, True, callback)
            if global_ is None:
                raise RuntimeError(f"couldn't find symbol for {name}")

    if not isinstance(global_, GlobalOp):
        raise RuntimeError(f"expected memref.global, got {global_}")
    result = global_.type_.value
    name = global_.sym_name.value
    return GetGlobalOp(result=result, name=name, loc=loc, ip=ip).result


def reinterpret_cast(
    source: Value,
    offsets: MixedValues = None,
    sizes: MixedValues = None,
    strides: MixedValues = None,
    *,
    loc=None,
    ip=None,
) -> Value:
    if offsets is None:
        offsets = []
    if sizes is None:
        sizes = []

    offsets_, _packed_offsets, static_offsets = _dispatch_mixed_values(offsets)
    sizes_, _packed_sizes, static_sizes = _dispatch_mixed_values(sizes)
    strides_, _packed_strides, static_strides = _dispatch_mixed_values(strides)

    if offsets_ or sizes_ or strides_:
        raise NotImplementedError("only static offsets and sizes and strides supported")

    default_strides = None
    if not static_strides and all(_is_static_int_like(s) for s in static_sizes):
        default_strides = list(accumulate(list(static_sizes)[1:][::-1], operator.mul))[::-1] + [1]
        static_strides = default_strides

    target_offset = 0
    for offset, target_stride in zip(static_offsets, static_strides):
        target_offset += offset * target_stride

    if static_strides == default_strides and target_offset == 0:
        layout = None
    else:
        layout = StridedLayoutAttr.get(target_offset, static_strides)

    result = MemRefType.get(
        static_sizes, source.type.element_type, layout, source.type.memory_space
    )
    return ReinterpretCastOp(
        result=result,
        source=source,
        offsets=offsets_,
        sizes=sizes_,
        strides=strides_,
        static_offsets=static_offsets,
        static_sizes=static_sizes,
        static_strides=static_strides,
        loc=loc,
        ip=ip,
    ).result
