"""Tensor + 反向模式自动求导（只依赖 numpy）。

每个运算产出一个新 Tensor，并记录:
  _parents  : 参与运算的输入 Tensor
  _backward : 函数 g -> (对每个 parent 的梯度, ...)
backward() 先拓扑排序，再从输出往回逐个调用 _backward，把梯度累加到 .grad。
"""
import numpy as np

_grad_enabled = True
_default_dtype = np.dtype(np.float32)


def set_default_dtype(dtype):
    """所有浮点数据进入框架时统一转成这个类型（默认 float32）。
    只有测试里的数值梯度检查需要切到 float64。整数数组（下标）不受影响。"""
    global _default_dtype
    _default_dtype = np.dtype(dtype)


def get_default_dtype():
    return _default_dtype


_flops = None   # count_flops() 打开时是一个 [计数]，否则 None


class count_flops:
    """with count_flops() as c: ...  之后 c.flops = 期间所有矩阵乘的浮点运算次数（乘加算 2 次）。
    只统计矩阵乘（模型里绝大部分计算量），不统计逐元素运算；不受机器负载影响，比计时可靠。"""

    def __enter__(self):
        global _flops
        self._prev, _flops = _flops, [0]
        return self

    def __exit__(self, *exc):
        global _flops
        self.flops = _flops[0]
        _flops = self._prev


class no_grad:
    """with no_grad(): 内部的运算不建计算图（推理/评估时用，省内存）。"""

    def __enter__(self):
        global _grad_enabled
        self._prev, _grad_enabled = _grad_enabled, False

    def __exit__(self, *exc):
        global _grad_enabled
        _grad_enabled = self._prev


def _unbroadcast(grad, shape):
    """把广播后的梯度求和还原成原始形状。"""
    while grad.ndim > len(shape):
        grad = grad.sum(axis=0)
    for i, s in enumerate(shape):
        if s == 1 and grad.shape[i] != 1:
            grad = grad.sum(axis=i, keepdims=True)
    return grad


def _to_array(x):
    if isinstance(x, Tensor):
        return x.data
    arr = np.asarray(x)
    # 浮点一律转成默认类型，避免 float32 和 float64 混在一起被 numpy 悄悄升成 float64
    if arr.dtype.kind == "f" and arr.dtype != _default_dtype:
        arr = arr.astype(_default_dtype)
    return arr


class Tensor:
    __array_priority__ = 100  # 让 ndarray + Tensor 走 Tensor.__radd__

    def __init__(self, data, requires_grad=False):
        self.data = _to_array(data)
        self.requires_grad = requires_grad
        self.grad = None
        self._parents = ()
        self._backward = None

    # ---------- 基础属性 ----------
    @property
    def shape(self):
        return self.data.shape

    @property
    def ndim(self):
        return self.data.ndim

    @property
    def dtype(self):
        return self.data.dtype

    def __len__(self):
        return len(self.data)

    def __repr__(self):
        return f"Tensor({self.data}, requires_grad={self.requires_grad})"

    def numpy(self):
        return self.data

    def item(self):
        return self.data.item()

    def detach(self):
        return Tensor(self.data)

    # ---------- 建图 ----------
    @staticmethod
    def _op(data, parents, backward):
        """创建运算结果。backward(g) 返回与 parents 一一对应的梯度（可为 None）。"""
        out = Tensor(data)
        if _grad_enabled and any(p.requires_grad for p in parents):
            out.requires_grad = True
            out._parents = parents
            out._backward = backward
        return out

    def backward(self, grad=None):
        if grad is None:
            if self.data.size != 1:
                raise RuntimeError("非标量输出调用 backward() 需要显式传入 grad")
            grad = np.ones_like(self.data)

        # 迭代式拓扑排序（避免深图递归爆栈）
        topo, visited, stack = [], set(), [(self, False)]
        while stack:
            node, expanded = stack.pop()
            if expanded:
                topo.append(node)
                continue
            if id(node) in visited:
                continue
            visited.add(id(node))
            stack.append((node, True))
            for p in node._parents:
                if p.requires_grad and id(p) not in visited:
                    stack.append((p, False))

        self.grad = grad if self.grad is None else self.grad + grad
        for node in reversed(topo):
            if node._backward is None or node.grad is None:
                continue
            for p, g in zip(node._parents, node._backward(node.grad)):
                if g is not None and p.requires_grad:
                    p.grad = g if p.grad is None else p.grad + g
            # 释放计算图，中间节点的梯度也不再需要
            node._parents, node._backward, node.grad = (), None, None

    # ---------- 逐元素运算 ----------
    def __add__(self, other):
        other = other if isinstance(other, Tensor) else Tensor(other)
        a, b = self, other
        return Tensor._op(a.data + b.data, (a, b),
                          lambda g: (_unbroadcast(g, a.shape), _unbroadcast(g, b.shape)))

    def __mul__(self, other):
        other = other if isinstance(other, Tensor) else Tensor(other)
        a, b = self, other
        return Tensor._op(a.data * b.data, (a, b),
                          lambda g: (_unbroadcast(g * b.data, a.shape), _unbroadcast(g * a.data, b.shape)))

    def __truediv__(self, other):
        other = other if isinstance(other, Tensor) else Tensor(other)
        a, b = self, other
        return Tensor._op(a.data / b.data, (a, b),
                          lambda g: (_unbroadcast(g / b.data, a.shape),
                                     _unbroadcast(-g * a.data / (b.data ** 2), b.shape)))

    def __pow__(self, p):
        assert isinstance(p, (int, float)), "只支持标量指数"
        a = self
        return Tensor._op(a.data ** p, (a,), lambda g: (g * p * a.data ** (p - 1),))

    def __neg__(self):
        return Tensor._op(-self.data, (self,), lambda g: (-g,))

    def __sub__(self, other):
        return self + (-other if isinstance(other, Tensor) else Tensor(-_to_array(other)))

    def __radd__(self, other):
        return self + other

    def __rmul__(self, other):
        return self * other

    def __rsub__(self, other):
        return (-self) + other

    def __rtruediv__(self, other):
        return Tensor(other) / self

    def exp(self):
        out = np.exp(self.data)
        return Tensor._op(out, (self,), lambda g: (g * out,))

    def log(self):
        a = self
        return Tensor._op(np.log(a.data), (a,), lambda g: (g / a.data,))

    def sqrt(self):
        return self ** 0.5

    def tanh(self):
        out = np.tanh(self.data)
        return Tensor._op(out, (self,), lambda g: (g * (1 - out ** 2),))

    def relu(self):
        mask = self.data > 0
        return Tensor._op(self.data * mask, (self,), lambda g: (g * mask,))

    def masked_fill(self, mask, value):
        mask = _to_array(mask).astype(bool)
        return Tensor._op(np.where(mask, value, self.data).astype(self.dtype), (self,),
                          lambda g: (np.where(mask, 0, g),))

    # ---------- 归约 ----------
    def sum(self, axis=None, keepdims=False):
        a = self

        def backward(g):
            if axis is not None and not keepdims:
                g = np.expand_dims(g, axis)
            return (np.broadcast_to(g, a.shape).copy(),)

        return Tensor._op(a.data.sum(axis=axis, keepdims=keepdims), (a,), backward)

    def mean(self, axis=None, keepdims=False):
        n = self.data.size if axis is None else np.prod([self.shape[i] for i in np.atleast_1d(axis)])
        return self.sum(axis, keepdims) * (1.0 / n)

    def max(self, axis=None, keepdims=False):
        a = self
        out = a.data.max(axis=axis, keepdims=True)

        def backward(g):
            if axis is not None and not keepdims:
                g = np.expand_dims(g, axis)
            mask = a.data == out
            return (mask * g / mask.sum(axis=axis, keepdims=True),)  # 并列最大值平分梯度

        return Tensor._op(out if keepdims else out.squeeze(axis), (a,), backward)

    # ---------- 形状 ----------
    def reshape(self, *shape):
        a = self
        return Tensor._op(a.data.reshape(*shape), (a,), lambda g: (g.reshape(a.shape),))

    def transpose(self, *axes):
        axes = axes or tuple(reversed(range(self.ndim)))
        inv = np.argsort(axes)
        return Tensor._op(self.data.transpose(axes), (self,), lambda g: (g.transpose(inv),))

    @property
    def T(self):
        return self.transpose()

    def __getitem__(self, idx):
        a = self
        if isinstance(idx, Tensor):
            idx = idx.data

        parts = idx if isinstance(idx, tuple) else (idx,)
        fancy = any(isinstance(i, (np.ndarray, list)) for i in parts)

        def backward(g):
            grad = np.zeros_like(a.data)
            if fancy:
                np.add.at(grad, idx, g)  # 数组下标可能重复（Embedding 查表），必须累加
            else:
                grad[idx] = g            # 普通切片不会重复，直接写，比 add.at 快得多
            return (grad,)

        return Tensor._op(a.data[idx], (a,), backward)

    # ---------- 矩阵乘 ----------
    def __matmul__(self, other):
        other = other if isinstance(other, Tensor) else Tensor(other)
        a, b = self, other
        assert a.ndim >= 2 and b.ndim >= 2, "matmul 需要至少 2 维"
        out = a.data @ b.data
        if _flops is not None:
            _flops[0] += 2 * out.size * a.shape[-1]

        def backward(g):
            ga = g @ np.swapaxes(b.data, -1, -2)
            gb = np.swapaxes(a.data, -1, -2) @ g
            return _unbroadcast(ga, a.shape), _unbroadcast(gb, b.shape)

        return Tensor._op(out, (a, b), backward)


def concat(tensors, axis=0):
    """沿 axis 拼接。反向时把梯度按原来的大小切回去。"""
    tensors = tuple(t if isinstance(t, Tensor) else Tensor(t) for t in tensors)
    splits = np.cumsum([t.shape[axis] for t in tensors])[:-1]
    return Tensor._op(np.concatenate([t.data for t in tensors], axis=axis), tensors,
                      lambda g: tuple(np.split(g, splits, axis=axis)))


# ---------- 构造函数 ----------
def tensor(data, requires_grad=False):
    return Tensor(data, requires_grad)


def zeros(*shape, requires_grad=False):
    return Tensor(np.zeros(shape, dtype=_default_dtype), requires_grad)


def ones(*shape, requires_grad=False):
    return Tensor(np.ones(shape, dtype=_default_dtype), requires_grad)


def randn(*shape, requires_grad=False):
    return Tensor(np.random.randn(*shape).astype(_default_dtype), requires_grad)