import numpy as np

class Tensor:
    __slots__ = ['data', 'grad', 'device', '_prev', '_backward']

    def __init__(self, data, device='cpu'):
        self.device = 'cpu'
        self.data = np.array(data, dtype=np.float16)
        self.grad = None
        self._prev = set()
        self._backward = lambda: None

    def backward(self):
        topo = []
        visited = set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev: build_topo(child)
                topo.append(v)
        build_topo(self)

        if self.grad is None:
            self.grad = np.ones_like(self.data, dtype=np.float16)
            
        for node in reversed(topo):
            node._backward()

    def _match_shape(self, grad, shape):
        if grad.shape == shape:
            return grad
        # 1. Убираем лишние размерности (например, из (10, 1) в (1,))
        if grad.ndim > len(shape):
            grad = grad.sum(axis=tuple(range(grad.ndim - len(shape))))
        # 2. Суммируем по осям, где у данных 1, а у градиента > 1 (broadcasting)
        for i, dim in enumerate(shape):
            if dim == 1 and grad.shape[i] > 1:
                grad = grad.sum(axis=i, keepdims=True)
        return grad.reshape(shape)

    def __add__(self, other):
        other = other if isinstance(other, Tensor) else Tensor(other)
        out = Tensor(self.data + other.data)
        out._prev = {self, other}
        def _backward():
            self.grad = (self.grad if self.grad is not None else 0) + self._match_shape(out.grad, self.data.shape)
            other.grad = (other.grad if other.grad is not None else 0) + self._match_shape(out.grad, other.data.shape)
        out._backward = _backward
        return out

    def __mul__(self, other):
        other = other if isinstance(other, Tensor) else Tensor(other)
        out = Tensor(self.data * other.data)
        out._prev = {self, other}
        def _backward():
            self.grad = (self.grad if self.grad is not None else 0) + self._match_shape(other.data * out.grad, self.data.shape)
            other.grad = (other.grad if other.grad is not None else 0) + self._match_shape(self.data * out.grad, other.data.shape)
        out._backward = _backward
        return out

    def __matmul__(self, other):
        other = other if isinstance(other, Tensor) else Tensor(other)
        out = Tensor(self.data @ other.data)
        out._prev = {self, other}
        def _backward():
            og = out.grad
            if og.ndim == 0: og = og.reshape(1, 1)
            self.grad = (self.grad if self.grad is not None else 0) + og @ other.data.T
            other.grad = (other.grad if other.grad is not None else 0) + self.data.T @ og
        out._backward = _backward
        return out

    def mean(self):
        out = Tensor(self.data.mean())
        out._prev = {self}
        def _backward():
            self.grad = (self.grad if self.grad is not None else 0) + (np.ones_like(self.data) / self.data.size) * out.grad
        out._backward = _backward
        return out

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