"""
tensor_layout.py —— 张量 = 一维内存 + 形状 + 步长（前置讲 P1 §5；第 3、5 讲还会用）
=====================================================================

【这个文件要回答的问题】
    内存是一维的：地址 0, 1, 2, ……一字排开。可我们用的是二维矩阵、三维、四维的张量。
    张量库（NumPy、PyTorch、原课的 miniTorch）的做法都一样：数据放在一段一维存储里，
    再用"形状 shape"和"步长 strides"描述怎样把多维下标翻译成一维位置：
        位置 = Σ_k 下标[k] × 步长[k]          （步长的单位是"元素个数"）
    这几个小函数和原课 HW1 / miniTorch 的 tensor_data.py 里的同名函数做的是同一件事
    （原课作业里的实现要你自己写，这里只是讲解用的参考版本）。

【本文件提供】
    strides_for(shape)                  —— 行优先（row-major）连续存放时的步长
    index_to_position(index, strides)   —— 多维下标 → 一维位置
    to_index(ordinal, shape)            —— 第 ordinal 个元素（按行优先顺序数）→ 多维下标
    shape_broadcast(shape_a, shape_b)   —— 两个形状广播之后的形状
    broadcast_index(big_index, big_shape, shape) —— 广播后的大张量的下标 → 原来小张量的下标
    is_contiguous(shape, strides)       —— 这组步长是不是"行优先连续"的
    permute(shape, strides, order)      —— 交换维度（转置的推广）：只改形状和步长，不动数据

【约定】shape、strides、index 都用元组；步长单位是"元素"，不是字节（NumPy 的 .strides 是字节，要除以 itemsize）。

【系统联系】
    GPU 内核拿到的永远只是"一个指针 + 若干整数"。张量的每一种变形（转置、切片、广播、reshape）
    要么只改这几个整数（零拷贝，叫"视图 view"），要么必须真正搬运数据（拷贝）。
    分清这两种情况，是写正确又快速的内核的第一步（第 3 讲）。
"""


def strides_for(shape):
    """
    行优先连续存放时的步长：最后一维步长为 1，往前每一维的步长 = 后面所有维度大小的乘积。

    【例子】shape (2, 3, 4) → strides (12, 4, 1)：第 0 维加 1 要跳过一整个 3×4 的"面"（12 个元素）。
    【边界】shape 为空元组（标量）时返回空元组。
    """
    strides = []
    acc = 1
    for n in reversed(shape):
        strides.append(acc)
        acc *= n
    return tuple(reversed(strides))


def index_to_position(index, strides, offset=0):
    """
    多维下标 → 一维存储里的位置：offset + Σ index[k] × strides[k]。

    【参数】index：下标元组；strides：步长元组（元素为单位）；offset：视图在存储里的起点（切片会产生非零起点）
    【边界】不检查越界（和 C 一样）——调用者要保证 0 ≤ index[k] < shape[k]。
    """
    return offset + sum(i * s for i, s in zip(index, strides))


def to_index(ordinal, shape):
    """
    按行优先顺序数的第 ordinal 个元素 → 多维下标（strides_for 的"逆运算"）。

    【思路】从最后一维开始，反复做"除以这一维的大小，余数就是这一维的下标"。
        GPU 内核里常见的写法：一个线程拿到一维编号 i，先用它算出自己负责的多维下标。
    【例子】shape (2, 3, 4)，ordinal 17 → (1, 1, 1)，因为 17 = 1×12 + 1×4 + 1。
    【边界】要求 0 ≤ ordinal < 元素总数。
    """
    idx = [0] * len(shape)
    for k in range(len(shape) - 1, -1, -1):
        idx[k] = ordinal % shape[k]
        ordinal //= shape[k]
    return tuple(idx)


def shape_broadcast(shape_a, shape_b):
    """
    广播规则（NumPy / PyTorch 通用）：从最后一维往前对齐；每一维要么相等，要么其中一个是 1；
    缺的维度当作 1。结果的每一维取两者中较大的那个。

    【例子】(2, 1, 4) 和 (3, 1) → (2, 3, 4)
    【边界】不兼容时（例如 3 和 4）抛出 ValueError。
    """
    n = max(len(shape_a), len(shape_b))
    a = (1,) * (n - len(shape_a)) + tuple(shape_a)
    b = (1,) * (n - len(shape_b)) + tuple(shape_b)
    out = []
    for x, y in zip(a, b):
        if x != y and x != 1 and y != 1:
            raise ValueError(f"形状 {shape_a} 和 {shape_b} 无法广播")
        out.append(max(x, y))
    return tuple(out)


def broadcast_index(big_index, big_shape, shape):
    """
    广播后的大张量里的下标 → 原来那个小张量里对应的下标。

    【思路】小张量的维度和大张量的最后几维对齐；小张量里大小为 1 的维度，下标一律取 0
        （同一个元素被"复用"到大张量的每一个位置）；大张量多出来的前几维直接丢掉。
    【例子】big_shape (2, 3, 4)，shape (3, 1)，big_index (1, 2, 3) → (2, 0)
    【边界】假设 shape 确实能广播成 big_shape。
    """
    offset = len(big_shape) - len(shape)
    return tuple(0 if shape[k] == 1 else big_index[k + offset] for k in range(len(shape)))


def is_contiguous(shape, strides):
    """步长是否等于行优先连续存放的步长（大小为 1 的维度步长不影响，忽略）。"""
    want = strides_for(shape)
    return all(n == 1 or s == w for n, s, w in zip(shape, strides, want))


def permute(shape, strides, order):
    """
    交换维度：新的第 k 维 = 原来的第 order[k] 维。只重排 shape 和 strides，数据一个字节都不动。

    【例子】二维转置就是 permute(shape, strides, (1, 0))。
    【返回】(新 shape, 新 strides)
    """
    return tuple(shape[k] for k in order), tuple(strides[k] for k in order)