from dataclasses import dataclass
import numpy as np

from ..._dtype import bool as bool_dtype
from ..._dtype import int64 as int64_dtype


@dataclass(frozen=True)
class TensorSpec:
    shape: tuple
    stride: tuple
    dtype: object
    offset: int = 0


def _contiguous_stride(shape):
    stride = []
    acc = 1
    for d in reversed(shape):
        stride.append(acc)
        acc *= d
    return tuple(reversed(stride))


def _broadcast_shape(a_shape, b_shape):
    try:
        return np.broadcast_shapes(a_shape, b_shape)
    except AttributeError:
        return np.broadcast(np.empty(a_shape), np.empty(b_shape)).shape


def infer_binary(a, b):
    shape = _broadcast_shape(a.shape, b.shape)
    return TensorSpec(shape=tuple(shape), stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_binary_bool(a, b):
    shape = _broadcast_shape(a.shape, b.shape)
    return TensorSpec(shape=tuple(shape), stride=_contiguous_stride(shape), dtype=bool_dtype)


def infer_unary(a, *args, **kwargs):
    return TensorSpec(shape=tuple(a.shape), stride=_contiguous_stride(a.shape), dtype=a.dtype)


def infer_unary_bool(a, *args, **kwargs):
    return TensorSpec(shape=tuple(a.shape), stride=_contiguous_stride(a.shape), dtype=bool_dtype)


def infer_sum(a, dim=None, keepdim=False):
    shape = list(a.shape)
    if isinstance(dim, (list, tuple)) and len(dim) == 0:
        dim = None

    ndim = len(shape)

    def _check_dim_range(d):
        if d < -ndim or d >= ndim:
            raise IndexError(
                f"Dimension out of range (expected to be in range of [{-ndim}, {ndim - 1}], but got {d})"
            )

    if isinstance(dim, int):
        _check_dim_range(dim)
    elif isinstance(dim, (list, tuple)):
        for d in dim:
            _check_dim_range(d)

    if dim is None:
        dims = list(range(len(shape)))
    elif isinstance(dim, int):
        dims = [dim]
    else:
        dims = list(dim)
    for d in sorted(dims):
        shape[d] = 1
    if not keepdim:
        shape = [s for i, s in enumerate(shape) if i not in dims]
    shape = tuple(shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_reduce_bool(a, dim=None, keepdim=False):
    spec = infer_sum(a, dim=dim, keepdim=keepdim)
    return TensorSpec(shape=spec.shape, stride=spec.stride, dtype=bool_dtype)


def infer_argmax(a, dim=None, keepdim=False):
    spec = infer_sum(a, dim=dim, keepdim=keepdim)
    return TensorSpec(shape=spec.shape, stride=spec.stride, dtype=int64_dtype)


def infer_cummax(a, dim=0):
    spec = infer_unary(a)
    indices = TensorSpec(shape=spec.shape, stride=spec.stride, dtype=int64_dtype)
    return (spec, indices)


def infer_sort(a, dim=-1, descending=False, stable=False):
    values = infer_unary(a)
    indices = TensorSpec(shape=values.shape, stride=values.stride, dtype=int64_dtype)
    return (values, indices)


def infer_topk(a, k, dim=-1, largest=True, sorted=True):
    shape = list(a.shape)
    shape[dim] = k
    shape = tuple(shape)
    values = TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)
    indices = TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=int64_dtype)
    return (values, indices)


def infer_argsort(a, dim=-1, descending=False, stable=False):
    spec = infer_unary(a)
    return TensorSpec(shape=spec.shape, stride=spec.stride, dtype=int64_dtype)


def infer_stack(tensors, dim=0):
    shape = list(tensors[0].shape)
    shape.insert(dim, len(tensors))
    shape = tuple(shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=tensors[0].dtype)


def infer_cat(tensors, dim=0):
    shape = list(tensors[0].shape)
    shape[dim] = sum(t.shape[dim] for t in tensors)
    shape = tuple(shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=tensors[0].dtype)


def infer_hstack(tensors):
    if len(tensors[0].shape) == 1:
        shape = (sum(t.shape[0] for t in tensors),)
    else:
        shape = list(tensors[0].shape)
        shape[1] = sum(t.shape[1] for t in tensors)
        shape = tuple(shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=tensors[0].dtype)


def infer_vstack(tensors):
    if len(tensors[0].shape) == 1:
        shape = (len(tensors), tensors[0].shape[0])
    else:
        shape = list(tensors[0].shape)
        shape[0] = sum(t.shape[0] for t in tensors)
        shape = tuple(shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=tensors[0].dtype)


def infer_column_stack(tensors):
    if len(tensors[0].shape) == 1:
        shape = (tensors[0].shape[0], len(tensors))
    else:
        shape = list(tensors[0].shape)
        shape[1] = sum(t.shape[1] for t in tensors)
        shape = tuple(shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=tensors[0].dtype)


def infer_dstack(tensors):
    expanded = []
    for t in tensors:
        shape = t.shape
        if len(shape) == 1:
            expanded.append((1, shape[0], 1))
        elif len(shape) == 2:
            expanded.append((shape[0], shape[1], 1))
        else:
            expanded.append(tuple(shape))
    base = expanded[0]
    if len(base) < 3:
        raise ValueError("dstack expects 1D, 2D, or 3D+ tensors")
    for shape in expanded[1:]:
        if len(shape) != len(base):
            raise ValueError("dstack expects same rank after expansion")
        if shape[0] != base[0] or shape[1] != base[1] or shape[3:] != base[3:]:
            raise ValueError("dstack expects same shape except dim=2")
    out_shape = list(base)
    out_shape[2] = sum(shape[2] for shape in expanded)
    out_shape = tuple(out_shape)
    return TensorSpec(shape=out_shape, stride=_contiguous_stride(out_shape), dtype=tensors[0].dtype)


def _tril_count(row, col, offset):
    if row <= 0 or col <= 0:
        return 0
    start = max(0, -offset)
    if start >= row:
        return 0
    cap = col - offset - 1
    total = 0
    if cap > start:
        end_linear = min(row - 1, cap - 1)
        if end_linear >= start:
            n = end_linear - start + 1
            total += (start + end_linear) * n // 2 + (offset + 1) * n
    start_full = max(start, cap)
    if start_full <= row - 1:
        total += (row - start_full) * col
    return total


def _triu_count(row, col, offset):
    return row * col - _tril_count(row, col, offset - 1)


def infer_tril_indices(row, col, offset=0, dtype=None, device=None, layout=None):
    if row < 0 or col < 0:
        raise ValueError("row and col must be non-negative")
    if dtype is None:
        dtype = int64_dtype
    n = _tril_count(row, col, offset)
    shape = (2, n)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=dtype)


def _normalize_tensor_sequence_args(tensors):
    if len(tensors) == 1 and isinstance(tensors[0], (list, tuple)):
        return tuple(tensors[0])
    return tuple(tensors)


def infer_triu_indices(row, col, offset=0, dtype=None, device=None, layout=None):
    if row < 0 or col < 0:
        raise ValueError("row and col must be non-negative")
    if dtype is None:
        dtype = int64_dtype
    n = _triu_count(row, col, offset)
    shape = (2, n)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=dtype)


def infer_pad_sequence(seqs, batch_first=False, padding_value=0.0, padding_side="right"):
    max_len = max(t.shape[0] for t in seqs)
    batch = len(seqs)
    trailing = seqs[0].shape[1:]
    shape = (batch, max_len, *trailing) if batch_first else (max_len, batch, *trailing)
    shape = tuple(shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=seqs[0].dtype)


def infer_block_diag(*tensors):
    tensors = _normalize_tensor_sequence_args(tensors)
    rows = sum(t.shape[0] for t in tensors)
    cols = sum(t.shape[1] for t in tensors)
    shape = (rows, cols)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=tensors[0].dtype)


def infer_diag(a, diagonal=0):
    if len(a.shape) == 1:
        n = a.shape[0]
        size = n + abs(diagonal)
        shape = (size, size)
    elif len(a.shape) == 2:
        m, n = a.shape
        if diagonal >= 0:
            length = max(0, min(m, n - diagonal))
        else:
            length = max(0, min(m + diagonal, n))
        shape = (length,)
    else:
        raise ValueError("diag expects 1D or 2D tensor")
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_cartesian_prod(*tensors):
    tensors = _normalize_tensor_sequence_args(tensors)
    rows = 1
    for t in tensors:
        rows *= t.shape[0]
    cols = len(tensors)
    shape = (rows, cols)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=tensors[0].dtype)


def infer_chunk(a, chunks, dim=0):
    dim_size = a.shape[dim]
    actual_chunks = min(chunks, dim_size) if dim_size > 0 else chunks
    if actual_chunks == 0:
        return tuple()
    chunk_size = (dim_size + actual_chunks - 1) // actual_chunks
    specs = []
    for i in range(actual_chunks):
        start = i * chunk_size
        end = min(start + chunk_size, dim_size)
        if start >= end:
            break
        shape = list(a.shape)
        shape[dim] = end - start
        shape = tuple(shape)
        specs.append(TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype))
    return tuple(specs)


def infer_split(a, split_size_or_sections, dim=0):
    dim_size = a.shape[dim]
    specs = []
    if isinstance(split_size_or_sections, int):
        step = split_size_or_sections
        for start in range(0, dim_size, step):
            end = min(start + step, dim_size)
            shape = list(a.shape)
            shape[dim] = end - start
            shape = tuple(shape)
            specs.append(TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype))
    else:
        sizes = list(split_size_or_sections)
        if sum(sizes) != dim_size:
            raise ValueError("split sections must sum to dim size")
        for size in sizes:
            shape = list(a.shape)
            shape[dim] = size
            shape = tuple(shape)
            specs.append(TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype))
    return tuple(specs)


def infer_vsplit(a, split_size_or_sections):
    if isinstance(split_size_or_sections, int):
        sizes = _split_sections_from_count(a.shape[0], split_size_or_sections)
        return infer_split(a, sizes, dim=0)
    return infer_split(a, split_size_or_sections, dim=0)


def infer_hsplit(a, split_size_or_sections):
    dim = 0 if len(a.shape) == 1 else 1
    if isinstance(split_size_or_sections, int):
        sizes = _split_sections_from_count(a.shape[dim], split_size_or_sections)
        return infer_split(a, sizes, dim=dim)
    return infer_split(a, split_size_or_sections, dim=dim)


def infer_dsplit(a, split_size_or_sections):
    if len(a.shape) < 3:
        raise ValueError("dsplit expects input with at least 3 dimensions")
    if isinstance(split_size_or_sections, int):
        sizes = _split_sections_from_count(a.shape[2], split_size_or_sections)
        return infer_split(a, sizes, dim=2)
    return infer_split(a, split_size_or_sections, dim=2)


def _split_sections_from_count(dim_size, sections):
    if sections <= 0:
        raise ValueError("sections must be > 0")
    size, extra = divmod(dim_size, sections)
    return [size + 1] * extra + [size] * (sections - extra)


def infer_unbind(a, dim=0):
    dim_size = a.shape[dim]
    specs = []
    for _ in range(dim_size):
        shape = list(a.shape)
        shape.pop(dim)
        shape = tuple(shape)
        specs.append(TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype))
    return tuple(specs)


def infer_take(a, index):
    shape = tuple(index.shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_take_along_dim(a, indices, dim):
    if dim < 0:
        dim += len(a.shape)
    if dim < 0 or dim >= len(a.shape):
        raise ValueError("dim out of range")
    if len(indices.shape) != len(a.shape):
        raise ValueError("indices shape mismatch")
    for i, size in enumerate(indices.shape):
        if i != dim and size != a.shape[i]:
            raise ValueError("indices shape mismatch")
    shape = tuple(indices.shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_index_select(a, dim, index):
    if len(index.shape) != 1:
        raise ValueError("index must be 1D")
    if dim < 0:
        dim += len(a.shape)
    if dim < 0 or dim >= len(a.shape):
        raise ValueError("dim out of range")
    shape = list(a.shape)
    shape[dim] = index.shape[0]
    shape = tuple(shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_gather(a, dim, index):
    if dim < 0:
        dim += len(a.shape)
    if dim < 0 or dim >= len(a.shape):
        raise ValueError("dim out of range")
    if len(index.shape) != len(a.shape):
        raise ValueError("index shape mismatch")
    for i, size in enumerate(index.shape):
        if i != dim and size != a.shape[i]:
            raise ValueError("index shape mismatch")
    shape = tuple(index.shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_scatter(a, dim, index, src):
    if dim < 0:
        dim += len(a.shape)
    if dim < 0 or dim >= len(a.shape):
        raise ValueError("dim out of range")
    if len(index.shape) != len(a.shape):
        raise ValueError("index shape mismatch")
    for i, size in enumerate(index.shape):
        if i != dim and size != a.shape[i]:
            raise ValueError("index shape mismatch")
    shape = tuple(a.shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_view(a, shape):
    shape = tuple(shape)
    size = 1
    for d in a.shape:
        size *= d
    new_size = 1
    for d in shape:
        new_size *= d
    if size != new_size:
        raise ValueError("reshape size mismatch")
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype, offset=a.offset)


def infer_transpose(a, dim0, dim1):
    shape = list(a.shape)
    stride = list(a.stride)
    shape[dim0], shape[dim1] = shape[dim1], shape[dim0]
    stride[dim0], stride[dim1] = stride[dim1], stride[dim0]
    return TensorSpec(shape=tuple(shape), stride=tuple(stride), dtype=a.dtype, offset=a.offset)


def infer_matmul(a, b):
    a_shape = a.shape
    b_shape = b.shape
    a_dim = len(a_shape)
    b_dim = len(b_shape)

    if a_dim == 1 and b_dim == 1:
        if a_shape[0] != b_shape[0]:
            raise ValueError("matmul shape mismatch")
        out_shape = ()
    elif a_dim == 1:
        k = a_shape[0]
        if b_dim < 2 or b_shape[-2] != k:
            raise ValueError("matmul shape mismatch")
        batch = b_shape[:-2]
        out_shape = batch + (b_shape[-1],)
    elif b_dim == 1:
        k = b_shape[0]
        if a_shape[-1] != k:
            raise ValueError("matmul shape mismatch")
        batch = a_shape[:-2]
        out_shape = batch + (a_shape[-2],)
    else:
        if a_shape[-1] != b_shape[-2]:
            raise ValueError("matmul shape mismatch")
        batch = _broadcast_shape(a_shape[:-2], b_shape[:-2])
        out_shape = batch + (a_shape[-2], b_shape[-1])
    return TensorSpec(shape=tuple(out_shape), stride=_contiguous_stride(out_shape), dtype=a.dtype)


def infer_masked_select(a, mask):
    return TensorSpec(shape=(0,), stride=_contiguous_stride((0,)), dtype=a.dtype)


def infer_flip(a, dims):
    return TensorSpec(shape=tuple(a.shape), stride=_contiguous_stride(a.shape), dtype=a.dtype)


def infer_roll(a, shifts, dims=None):
    return TensorSpec(shape=tuple(a.shape), stride=_contiguous_stride(a.shape), dtype=a.dtype)


def infer_rot90(a, k=1, dims=(0, 1)):
    dim0, dim1 = dims
    shape = list(a.shape)
    if k % 2 != 0:
        shape[dim0], shape[dim1] = shape[dim1], shape[dim0]
    shape = tuple(shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_repeat(a, repeats):
    if isinstance(repeats, int):
        repeats = (repeats,)
    if len(repeats) < len(a.shape):
        repeats = (1,) * (len(a.shape) - len(repeats)) + tuple(repeats)
    if len(repeats) != len(a.shape):
        raise ValueError("repeats must match input rank")
    shape = tuple(s * r for s, r in zip(a.shape, repeats))
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_repeat_interleave(a, repeats, dim=None):
    shape = list(a.shape)
    if dim is None:
        total = 1
        for s in shape:
            total *= s
        if isinstance(repeats, int):
            total *= repeats
        else:
            total = len(repeats)
        return TensorSpec(shape=(total,), stride=_contiguous_stride((total,)), dtype=a.dtype)
    if dim < 0:
        dim += len(shape)
    if isinstance(repeats, int):
        shape[dim] *= repeats
    else:
        shape[dim] = len(repeats)
    shape = tuple(shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_tile(a, dims):
    if isinstance(dims, int):
        dims = (dims,)
    if len(dims) < len(a.shape):
        dims = (1,) * (len(a.shape) - len(dims)) + tuple(dims)
    if len(dims) != len(a.shape):
        raise ValueError("dims must match input rank")
    shape = tuple(s * d for s, d in zip(a.shape, dims))
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_nonzero(a, as_tuple=False):
    if as_tuple:
        spec = TensorSpec(shape=(0,), stride=_contiguous_stride((0,)), dtype=int64_dtype)
        return tuple(spec for _ in range(len(a.shape)))
    shape = (0, len(a.shape))
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=int64_dtype)


# ---------------------------------------------------------------------------
# Meta infer for newly added ops
# ---------------------------------------------------------------------------

def infer_dot(a, b):
    # dot of two 1D vectors returns a scalar
    shape = ()
    return TensorSpec(shape=shape, stride=(), dtype=a.dtype)


def infer_outer(a, b):
    shape = (a.shape[0], b.shape[0])
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_flatten(a, start_dim=0, end_dim=-1):
    ndim = len(a.shape)
    start = start_dim if start_dim >= 0 else start_dim + ndim
    end = end_dim if end_dim >= 0 else end_dim + ndim
    flat_size = 1
    for i in range(start, end + 1):
        flat_size *= a.shape[i]
    shape = tuple(a.shape[:start]) + (flat_size,) + tuple(a.shape[end + 1:])
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_unflatten(a, dim, sizes):
    ndim = len(a.shape)
    d = dim if dim >= 0 else dim + ndim
    shape = tuple(a.shape[:d]) + tuple(sizes) + tuple(a.shape[d + 1:])
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_broadcast_to(a, shape):
    shape = tuple(shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_movedim(a, source, destination):
    # Shape doesn't change, only strides
    shape = tuple(a.shape)
    return TensorSpec(shape=shape, stride=_contiguous_stride(shape), dtype=a.dtype)


def infer_diagonal(a, offset=0, dim1=0, dim2=1):
    shape = list(a.shape)
    m = shape[dim1]
    n = shape[dim2]
    if offset >= 0:
        diag_len = max(0, min(m, n - offset))
    else:
        diag_len = max(0, min(m + offset, n))
    # Remove dim1 and dim2, append diag_len
    remaining = [s for i, s in enumerate(shape) if i not in (dim1, dim2)]
    out_shape = tuple(remaining) + (diag_len,)
    return TensorSpec(shape=out_shape, stride=_contiguous_stride(out_shape), dtype=a.dtype)
