"""
kernels.py —— 原课 HW1 的四类内核（map、zip、reduce、矩阵乘法）在 SIMT 模拟器上的实现，以及主机端的封装（第 3 讲）
=====================================================================
（第 4 讲用它做对照，第 5 讲把它接成迷你框架的一个"后端"）

【这个文件要回答的问题】
    原课 HW1 要在 combine.cu 里写四个 CUDA 内核，再在 cuda_kernel_ops.py 里把它们接到 miniTorch 的张量上。
    这里用 SIMT 模拟器把同样的东西写一遍：
        内核（设备端）：每个线程算出自己负责的输出位置 → 用步长找到输入的位置 → 计算 → 写回
        封装（主机端）：分配显存 → 拷贝输入 → 算网格大小 → 发射内核 → 拷回结果 → 释放显存
    张量 = 一维存储 + 形状 + 步长（P1 §5）。内核只拿到"指针 + 几个整数元组"，
    所以同一个内核能处理连续的张量、转置之后的视图、广播之后的张量。

【本文件提供】
    HostTensor          : 主机端的张量（NumPy 一维存储 + shape + strides），from_numpy、permute、to_numpy
    内核（设备端）       : map_kernel、zip_kernel、reduce_kernel_thread、reduce_kernel_block、matmul_kernel
    封装（主机端）       : tensor_map、tensor_zip、tensor_reduce、tensor_matmul（每个都返回 (结果, LaunchStats)）
    参考实现            : reduce_reference、matmul_reference（按和内核完全相同的顺序用 float32 计算，用来逐位对比）

【约定】
    - 和 miniTorch 一样：reduce 之后被归约的那一维保留、长度变成 1；matmul 支持三维（第 0 维是批），批维可以广播。
    - 下标换算用 P1 的 tensor_layout.py（to_index、index_to_position、broadcast_index），它们在内核里相当于 __device__ 函数。
    - 形状和步长以 Python 元组的形式直接传给内核（相当于 CUDA 的内核参数 / 常量内存）。
      原课 HW1 的做法是把它们也 cudaMemcpy 到显存里再传指针，原理相同。
"""

import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture02"))
sys.path.insert(0, os.path.join(HERE, "..", "prereq1"))
from simt import ceil_div, launch  # noqa: E402
from tensor_layout import broadcast_index, index_to_position, shape_broadcast, strides_for, to_index  # noqa: E402

THREADS_PER_BLOCK = 32          # 一维内核每块的线程数（本讲义取 32，模拟器跑得快；真 GPU 上常用 128–256）


# =====================================================================
# 0. 主机端的张量
# =====================================================================

class HostTensor:
    """
    主机内存里的张量：一维存储 storage + 形状 shape + 步长 strides（单位是元素）。和 miniTorch 的 TensorData 一样。

    【方法】
        HostTensor.from_numpy(arr) : 从 NumPy 数组建一个连续的张量（复制一份数据）
        t.permute(*order)          : 交换维度，返回共享存储的新视图（零拷贝）
        t.to_numpy()               : 按 shape / strides 读出一个普通的 NumPy 数组（用来检查结果）
    【边界】没有 offset（起点总是存储的第 0 个元素），和 miniTorch 一致。
    """

    def __init__(self, storage, shape, strides=None):
        self.storage = storage
        self.shape = tuple(int(s) for s in shape)
        self.strides = tuple(strides) if strides is not None else strides_for(self.shape)
        self.size = int(np.prod(self.shape)) if self.shape else 1

    @classmethod
    def from_numpy(cls, arr, dtype=np.float32):
        arr = np.ascontiguousarray(arr, dtype=dtype)
        return cls(arr.reshape(-1).copy(), arr.shape)

    def permute(self, *order):
        return HostTensor(self.storage, tuple(self.shape[k] for k in order), tuple(self.strides[k] for k in order))

    def to_numpy(self):
        itemsize = self.storage.itemsize
        return np.lib.stride_tricks.as_strided(self.storage, self.shape, tuple(s * itemsize for s in self.strides)).copy()

    def __repr__(self):
        return f"HostTensor(shape={self.shape}, strides={self.strides})"


# =====================================================================
# 1. 内核（设备端）
# =====================================================================

def map_kernel(t, out, out_shape, out_strides, out_size, a, a_shape, a_strides, fn):
    """
    map：out[i] = fn(a[对应位置])，a 可以是任意步长的视图，也可以比 out 小（广播）。

    对应原课 HW1 的 mapKernel。每个线程负责一个输出元素：
        1. i = blockIdx.x * blockDim.x + threadIdx.x；i ≥ out_size 的线程什么也不做
        2. out_index = to_index(i, out_shape)              —— 第 i 个输出元素的多维下标
        3. a_index = broadcast_index(out_index, ...)      —— 广播：a 里大小为 1 的维度下标取 0
        4. 按各自的步长换算成一维位置，读 a、算 fn、写 out
    【边界】fn 是一个普通的 Python 函数（CUDA 里是模板参数 / 函数指针）。
    """
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if i < out_size:
        out_index = to_index(i, out_shape)
        a_index = broadcast_index(out_index, out_shape, a_shape)
        out[index_to_position(out_index, out_strides)] = fn(a[index_to_position(a_index, a_strides)])


def zip_kernel(t, out, out_shape, out_strides, out_size, a, a_shape, a_strides, b, b_shape, b_strides, fn):
    """
    zip：out[i] = fn(a[…], b[…])，a、b 各自按广播规则对应到 out 的形状。对应原课 HW1 的 zipKernel。
    """
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if i < out_size:
        out_index = to_index(i, out_shape)
        ai = broadcast_index(out_index, out_shape, a_shape)
        bi = broadcast_index(out_index, out_shape, b_shape)
        out[index_to_position(out_index, out_strides)] = fn(a[index_to_position(ai, a_strides)],
                                                             b[index_to_position(bi, b_strides)])


def reduce_kernel_thread(t, out, out_shape, out_strides, out_size, a, a_shape, a_strides, dim, start, fn):
    """
    reduce（写法一：一个线程算一个输出）：线程 i 负责第 i 个输出，自己沿 dim 从头加到尾。

    【优点】简单，不需要共享内存和同步。
    【缺点】输出很少时（例如把一个长向量求和成 1 个数）只有一个线程在干活，其余几万个线程闲着。
    """
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if i < out_size:
        out_index = to_index(i, out_shape)
        a_index = list(out_index)
        acc = start
        for k in range(a_shape[dim]):
            a_index[dim] = k
            acc = fn(acc, a[index_to_position(a_index, a_strides)])
        out[index_to_position(out_index, out_strides)] = acc


def reduce_kernel_block(t, out, out_shape, out_strides, out_size, a, a_shape, a_strides, dim, start, fn, mode):
    """
    reduce（写法二：一个线程块算一个输出）：对应原课 HW1 的提示——"每个输出元素由一个块单独计算"，
    以及可选的"共享内存 + 树形归约"优化。

    【步骤】
        1. 块号 blockIdx.x = 第几个输出；块内每个线程沿 dim 以 blockDim.x 为步长各加一部分（跨步循环）
        2. 每个线程把自己的部分和写进共享数组 cache[tid]，然后 __syncthreads()
        3. 树形归约：每一轮把"活着的"部分和两两合并，直到只剩 cache[0]；每一轮之后都要 __syncthreads()
             mode = "interleaved" : 第 s 轮，tid % (2s) == 0 的线程把 cache[tid + s] 加到 cache[tid]（HW1 给的伪代码 reduce0）
             mode = "strided"     : 第 s 轮，前 blockDim/(2s) 个线程各算 index = 2s·tid，把 cache[index + s] 加到 cache[index]
                                    （合并的是和 interleaved 完全相同的两两配对，只是换了"由谁来做"）
             mode = "sequential"  : 步长从 blockDim/2 开始减半，tid < s 的线程把 cache[tid + s] 加到 cache[tid]
        4. 线程 0 把 cache[0] 写到输出
    【边界】blockDim.x 必须是 2 的幂（树形归约每轮减半）。
        注意 __syncthreads() 不能放在只有部分线程会执行的分支里——所以 if 只包住"加"那一行，栅栏在 if 外面。
    """
    tid = t.threadIdx.x
    nthreads = t.blockDim.x
    cache = t.shared("cache", nthreads, np.float32)
    out_index = to_index(t.blockIdx.x, out_shape)
    a_index = list(out_index)
    acc = start
    k = tid
    while k < a_shape[dim]:
        a_index[dim] = k
        acc = fn(acc, a[index_to_position(a_index, a_strides)])
        k += nthreads
    cache[tid] = acc
    yield t.syncthreads()
    if mode == "interleaved":
        s = 1
        while s < nthreads:
            if tid % (2 * s) == 0:
                cache[tid] = fn(cache[tid], cache[tid + s])
            yield t.syncthreads()
            s *= 2
    elif mode == "strided":
        s = 1
        while s < nthreads:
            index = 2 * s * tid
            if index < nthreads:
                cache[index] = fn(cache[index], cache[index + s])
            yield t.syncthreads()
            s *= 2
    else:
        s = nthreads // 2
        while s > 0:
            if tid < s:
                cache[tid] = fn(cache[tid], cache[tid + s])
            yield t.syncthreads()
            s //= 2
    if tid == 0:
        out[index_to_position(out_index, out_strides)] = cache[0]


def matmul_kernel(t, out, out_shape, out_strides, a, a_shape, a_strides, b, b_shape, b_strides):
    """
    批量矩阵乘法 out[z] = a[z] @ b[z]：一个线程算一个输出元素（原课 HW1 的"简单并行"提示）。

    【线程分工】三维网格：blockIdx.z = 第几个批；x 方向对应输出的列 j、y 方向对应行 i：
        j = blockIdx.x * blockDim.x + threadIdx.x，i = blockIdx.y * blockDim.y + threadIdx.y
        这样一个线程束里相邻的线程算同一行相邻的列——读 b 和写 out 都是合并的（第 2 部分会看到反过来有多糟）。
    【广播】a 或 b 的批维长度为 1 时，它的批步长当作 0：所有批共用同一个矩阵。
    【边界】越界检查：i、j 超出输出形状的线程什么也不做（网格是按块大小向上取整的）。
    【数值】acc 用 float32 累加，每步先乘、再加（两次舍入）。真 GPU 上 nvcc 默认把 a*b+c 编译成 FMA（只舍入一次），
        所以真 GPU 的结果可能和这里差最后一位。
    """
    z = t.blockIdx.z
    i = t.blockIdx.y * t.blockDim.y + t.threadIdx.y
    j = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    n, p = out_shape[1], out_shape[2]
    if i < n and j < p:
        m = a_shape[2]
        a_base = (z if a_shape[0] > 1 else 0) * a_strides[0] + i * a_strides[1]
        b_base = (z if b_shape[0] > 1 else 0) * b_strides[0] + j * b_strides[2]
        acc = np.float32(0)
        for k in range(m):
            acc += a[a_base + k * a_strides[2]] * b[b_base + k * b_strides[1]]
        out[z * out_strides[0] + i * out_strides[1] + j * out_strides[2]] = acc


# =====================================================================
# 2. 主机端封装（对应原课 HW1 的 cuda_kernel_ops.py）
# =====================================================================

def _upload(dev, x, name):
    """把一个 HostTensor 的整段存储拷到显存（视图也拷整段存储，步长照用）。"""
    return dev.to_device(x.storage, name)


def tensor_map(dev, fn, a, out_shape=None, block=THREADS_PER_BLOCK, **launch_kw):
    """
    out = map(fn, a)。out_shape 默认等于 a 的形状；给出更大的形状时 a 会被广播。

    【步骤】和原课 HW1 Part B 一样：建输出 → cudaMalloc / cudaMemcpy → 算块数 → 发射 → 拷回 → cudaFree。
    【返回】(输出 HostTensor（连续）, LaunchStats)
    """
    out_shape = tuple(out_shape or a.shape)
    out = HostTensor(np.zeros(int(np.prod(out_shape)), np.float32), out_shape)
    d_a, d_out = _upload(dev, a, "a"), dev.malloc(out.size, np.float32, "out")
    stats = launch(map_kernel, ceil_div(out.size, block), block,
                   d_out, out.shape, out.strides, out.size, d_a, a.shape, a.strides, fn, **launch_kw)
    out.storage = dev.to_host(d_out)
    dev.free(d_a)
    dev.free(d_out)
    return out, stats


def tensor_zip(dev, fn, a, b, block=THREADS_PER_BLOCK, **launch_kw):
    """out = zip(fn, a, b)，输出形状是两者广播之后的形状。【返回】(输出, LaunchStats)"""
    out_shape = shape_broadcast(a.shape, b.shape)
    out = HostTensor(np.zeros(int(np.prod(out_shape)), np.float32), out_shape)
    d_a, d_b, d_out = _upload(dev, a, "a"), _upload(dev, b, "b"), dev.malloc(out.size, np.float32, "out")
    stats = launch(zip_kernel, ceil_div(out.size, block), block,
                   d_out, out.shape, out.strides, out.size, d_a, a.shape, a.strides, d_b, b.shape, b.strides, fn,
                   **launch_kw)
    out.storage = dev.to_host(d_out)
    for d in (d_a, d_b, d_out):
        dev.free(d)
    return out, stats


def tensor_reduce(dev, fn, a, dim, start, method="block", mode="sequential", block=None, **launch_kw):
    """
    沿 dim 归约，输出在 dim 上的长度变成 1。

    【参数】fn：二元函数（例如加法）；start：初始值（求和 0、求积 1、求最大值 -inf）；
        method："thread"（一个线程一个输出）或 "block"（一个块一个输出 + 树形归约）；mode：树形归约的两种写法；
        block：每块线程数（"block" 方法默认 32，必须是 2 的幂）
    【返回】(输出, LaunchStats)
    """
    out_shape = tuple(1 if k == dim else s for k, s in enumerate(a.shape))
    out = HostTensor(np.zeros(int(np.prod(out_shape)), np.float32), out_shape)
    d_a, d_out = _upload(dev, a, "a"), dev.malloc(out.size, np.float32, "out")
    if method == "thread":
        block = block or THREADS_PER_BLOCK
        stats = launch(reduce_kernel_thread, ceil_div(out.size, block), block,
                       d_out, out.shape, out.strides, out.size, d_a, a.shape, a.strides, dim, np.float32(start), fn,
                       **launch_kw)
    else:
        block = block or 32
        if block & (block - 1):
            raise ValueError("树形归约要求每块线程数是 2 的幂")
        stats = launch(reduce_kernel_block, out.size, block,
                       d_out, out.shape, out.strides, out.size, d_a, a.shape, a.strides, dim, np.float32(start), fn, mode,
                       **launch_kw)
    out.storage = dev.to_host(d_out)
    dev.free(d_a)
    dev.free(d_out)
    return out, stats


def tensor_matmul(dev, a, b, block=(8, 8), **launch_kw):
    """
    批量矩阵乘法：a (B, n, m) @ b (B, m, p) → (B, n, p)；二维输入会先补一个长度为 1 的批维，结果再去掉。

    【边界】a 的最后一维必须等于 b 的倒数第二维；两者的批维要么相等，要么有一个是 1（广播）。
    【返回】(输出, LaunchStats)
    """
    squeeze = len(a.shape) == 2 and len(b.shape) == 2
    if len(a.shape) == 2:
        a = HostTensor(a.storage, (1,) + a.shape, (0,) + a.strides)
    if len(b.shape) == 2:
        b = HostTensor(b.storage, (1,) + b.shape, (0,) + b.strides)
    if a.shape[2] != b.shape[1]:
        raise ValueError(f"矩阵乘法的形状不匹配：{a.shape} @ {b.shape}")
    if a.shape[0] != b.shape[0] and 1 not in (a.shape[0], b.shape[0]):
        raise ValueError(f"批维无法广播：{a.shape[0]} vs {b.shape[0]}")
    B = max(a.shape[0], b.shape[0])
    out_shape = (B, a.shape[1], b.shape[2])
    out = HostTensor(np.zeros(int(np.prod(out_shape)), np.float32), out_shape)
    d_a, d_b, d_out = _upload(dev, a, "a"), _upload(dev, b, "b"), dev.malloc(out.size, np.float32, "out")
    bx, by = block
    grid = (ceil_div(out_shape[2], bx), ceil_div(out_shape[1], by), B)
    stats = launch(matmul_kernel, grid, (bx, by), d_out, out.shape, out.strides,
                   d_a, a.shape, a.strides, d_b, b.shape, b.strides, **launch_kw)
    out.storage = dev.to_host(d_out)
    for d in (d_a, d_b, d_out):
        dev.free(d)
    if squeeze:
        out = HostTensor(out.storage, out_shape[1:])
    return out, stats


# =====================================================================
# 3. 参考实现（按和内核相同的顺序计算，用来逐位对比）
# =====================================================================

def reduce_reference(x, dim, start, fn_np, method="thread", block=32, mode="sequential"):
    """
    用 NumPy 按和内核完全相同的运算顺序做归约（float32），结果应该和内核逐位相同。

    【参数】x：NumPy 数组；fn_np：向量化的二元函数（np.add、np.maximum……）
        method="thread"：从头到尾依次合并；method="block"：先按 tid 跨步累加成 block 份，再按 mode 做树形合并
    """
    x = np.moveaxis(np.asarray(x, np.float32), dim, -1)
    L = x.shape[-1]
    if method == "thread":
        acc = np.full(x.shape[:-1], start, np.float32)
        for k in range(L):
            acc = fn_np(acc, x[..., k]).astype(np.float32)
    else:
        parts = np.full(x.shape[:-1] + (block,), start, np.float32)
        for k in range(L):
            parts[..., k % block] = fn_np(parts[..., k % block], x[..., k])
        if mode in ("interleaved", "strided"):         # 两种写法合并的是同样的配对，结果相同
            s = 1
            while s < block:
                idx = np.arange(0, block, 2 * s)
                parts[..., idx] = fn_np(parts[..., idx], parts[..., idx + s])
                s *= 2
        else:
            s = block // 2
            while s > 0:
                parts[..., :s] = fn_np(parts[..., :s], parts[..., s:2 * s])
                s //= 2
        acc = parts[..., 0]
    return np.expand_dims(acc, dim)


def matmul_reference(a, b):
    """按 k 从 0 到 m−1 依次"先乘后加"的 float32 矩阵乘法（和 matmul_kernel 的运算顺序相同）。"""
    a = np.asarray(a, np.float32)
    b = np.asarray(b, np.float32)
    acc = np.zeros(np.broadcast_shapes(a[..., :1].shape, b[..., :1, :].shape), np.float32)   # 输出形状（批维广播）
    for k in range(a.shape[-1]):
        acc = (acc + (a[..., :, k:k + 1] * b[..., k:k + 1, :]).astype(np.float32)).astype(np.float32)
    return acc