from __future__ import annotations
from typing import Optional, Tuple, Union, Any, Dict, Callable, Type, List, ClassVar
from enum import Enum, auto
from abc import ABC

# first, 2+3 as a Tensor, the highest level
from tinygrad.tensor import Tensor
a = Tensor([2])
b = Tensor([3])
result = a + b
print(f"{a.numpy()} + {b.numpy()} = {result.numpy()}")
assert result.numpy()[0] == 5.

class Tensor:
  # these two are pretty straightforward
  grad: Optional[Tensor]
  requires_grad: Optional[bool]
 
  # this is the graph for the autograd engine
  _ctx: Optional[Function]
 
  # this is where the data (and other tensor properties) actually live
  lazydata: LazyBuffer
 
  # high level ops (hlops) are defined on this class. Ex: relu
  def relu(self): return self.maximum(0)
 
  # log is an mlop, this is the wrapper function in Tensor
  def log(self): return mlops.Log.apply(self)

class Function:
  # Ex: types of forward and backward
  def forward(self, x:LazyBuffer) -> LazyBuffer: pass
  def backward(self, x:LazyBuffer) -> LazyBuffer: pass

class LazyBuffer:
  # these three define the "type" of the buffer, and they are returned as Tensor properties
  device: str
  shape: Tuple[int, ...]
  dtype: DType
 
  # a ShapeTracker is used to track things like reshapes and permutes
  # all MovementOps are zero copy in tinygrad!
  # the ShapeTracker specifies how the data in the RawBuffer matches to the shape
  # we'll come back to this later
  st: ShapeTracker
 
  # if the LazyBuffer is realized, it has a RawBuffer
  # we will come back to RawBuffers later
  realized: Optional[RawBuffer]
 
  # if the lazybuffer is unrealized, it has a LazyOp
  # this LazyOp describes the computation needed to realize this LazyBuffer
  op: Optional[LazyOp]

class LazyOp:
  op: Op                                       # the type of the compute
  src: Tuple[Union[LazyOp, LazyBuffer], ...]   # the sources
  arg: Optional[Any] = None                    # and an optional static argument
 
# there's currently 28 Ops you have to implement for an accelerator.
class UnaryOps(Enum):    NOOP = auto(); EXP2 = auto(); LOG2 = auto(); CAST = auto(); SIN = auto();   SQRT = auto()
class BinaryOps(Enum):   ADD = auto();  SUB = auto();  MUL = auto();  DIV = auto();  CMPLT = auto(); MAX = auto()
class ReduceOps(Enum):   SUM = auto();  MAX = auto()
class MovementOps(Enum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); STRIDE = auto()
class TernaryOps(Enum):  MULACC = auto(); WHERE = auto()
class LoadOps(Enum):     EMPTY = auto(); RAND = auto(); CONST = auto(); FROM = auto(); CONTIGUOUS = auto(); CUSTOM = auto()

from tinygrad.tensor import Tensor
from tinygrad.ops import LazyOp, BinaryOps, LoadOps
 
# the 2+3 from before
result = Tensor([2]) + Tensor([3])
print(type(result.lazydata), result.lazydata)  # let's look at the lazydata of result
``` python
 you'll see it has a LazyOp
 the op type is BinaryOps.ADD
  and it has two sources, the 2 and the 3
lazyop: LazyOp = result.lazydata.op
assert lazyop.op == BinaryOps.ADD
assert len(lazyop.src) == 2
# the first source is the 2, it comes from the CPU
# the source is a LazyBuffer that is a "CPU" Tensor
# again, a LazyOp AST is like a GPU kernel. you have to copy the data on the device first
assert lazyop.src[0].op.op == LoadOps.FROM
assert lazyop.src[0].op.src[0].device == "CPU"
assert lazyop.src[0].op.src[0].op.src[0].realized._buf[0] == 2, "the src of the FROM LazyOP is a LazyBuffer on the CPU holding [2.]"
assert result.lazydata.realized is None, "the LazyBuffer is not realized yet"
 
# now we realize the LazyBuffer
result.realize()
assert result.lazydata.realized is not None, "the LazyBuffer is realized!"
# this brings us nicely to DeviceBuffer, of which the realized ClangBuffer is a subclass
assert 'RawMallocBuffer' in str(type(result.lazydata.realized))
# getting ahead of ourselves, but we can copy the DeviceBuffer toCPU
assert result.lazydata.realized.toCPU()[0] == 5, "when put in numpy with toCPU, it's 5"

class Interpreted:
  # they have a backing RawBuffer
  buffer: Type[RawBuffer]
 
  # and they have a lookup table to functions for the Ops
  fxn_for_op: Dict[Op, Callable] = {
    UnaryOps.EXP2: lambda x: np.exp2(x),
    BinaryOps.ADD: lambda x,y: x+y}
 
# Compiled backends take a little more (Ex: GPU and LLVM)
class Compiled:
  # they also have a backing RawBuffer
  buffer: Type[RawBuffer]
 
  # a code generator, which compiles the AST
  codegen: Type[Linearizer]
 
  # and a runtime, which runs the generated code
  runtime: Type[Runtime]
 
# Runtime is what actually runs the kernels for a compiled backend
class Runtime(ABC):
  # `name` is the name of the function, and `prg` is the code
  # the constructor compiles the code
  def __init__(self, name:str, prg:str): pass
  # call runs the code on the bufs. NOTE: the output is always bufs[0], but this is just a convention
  def __call__(self, global_size:Optional[List[int]], local_size:Optional[List[int]], *bufs:List[RawBuffer]): pass

import numpy as np
 
 RawBuffer is where the data is actually held. it's pretty close to just memory
 
class RawBuffer(ABC):
  # create an empty rawbuffer that holds `size` elements of type `dtype`
  # `buf` is an opaque container class
  def __init__(self, size:int, dtype:DType, buf:Any): raise NotImplementedError("must be implemented")
 
  # fromCPU is classmethod that creates a RawBuffer, it's a classmethod since some runtimes are 0 copy
  @classmethod
  def fromCPU(cls:RawBuffer, x:np.ndarray) -> RawBuffer: raise NotImplementedError("must be implemented")
 
  # toCPU converts the RawBuffer to a numpy array with shape (size,). many backends are 0 copy here
  def toCPU(self) -> np.ndarray: raise NotImplementedError("must be implemented")

class RawNumpyBuffer(RawBuffer):
  # NOTE: the "np.ndarray" is stored in the opaque container
  def __init__(self, buf:np.ndarray):
    super().__init__(buf.size, dtypes.from_np(buf.dtype), buf)
  @classmethod
  def fromCPU(cls, x): return cls(x)
  def toCPU(self): return self._buf

 
 in CLANG, global_size and local_size are ignored
from tinygrad.runtime.ops_clang import ClangProgram, compile_clang
# a concrete Ex: looks like this, this adds two size 1 RawBuffer
# first we create two numpy buffers containing 2 and 3
# then we copy the numpy in to RawMallocBuffers
# last, we create an empty output buffer
from tinygrad.helpers import dtypes
numpy_a, numpy_b = np.array([2], dtype=np.float32), np.array([3], dtype=np.float32)
input_a, input_b = RawMallocBuffer.fromCPU(numpy_a), RawMallocBuffer.fromCPU(numpy_b)
output = RawMallocBuffer(1, dtypes.float32)
 
# compile the program, run it, and 2+3 does indeed equal 5
program = ClangProgram("add", compile_clang(f"void add(float *a, float *b, float *c) {{ *a = *b + *c; }}"))
program(output, input_a, input_b)
print(output.toCPU())
assert output.toCPU()[0] == 5, "it's still 5"
np.testing.assert_allclose(output.toCPU(), numpy_a+numpy_b)

class UOps(Enum): LOOP = auto(); DEFINE_LOCAL = auto(); LOAD = auto(); ALU = auto(); CONST = auto(); ENDLOOP = auto(); STORE = auto();
 
class UOp:
  uop: UOps
  dtype: Optional[DType]
  vin: Tuple[UOp, ...]
  arg: Any
  num: int  # UOps are unique
 
class Linearizer:
  # create the kernel with the AST
  # NOTE: the AST contains the CompiledBuffers themselves as the root nodes. this will change
  def __init__(self, ast:LazyOp): pass
  def linearize(self): pass
 
  # when linearize is run, it fills in this list
  uops: List[UOp]
 
from tinygrad.tensor import Tensor
result = Tensor(2).realize() + Tensor(3).realize()
 
# use the real Linearizer to linearize 2+3
from tinygrad.codegen.linearizer import Linearizer
sched = result.lazydata.schedule()
linearizer = Linearizer(sched[-1].ast)
linearizer.linearize()

from tinygrad.tensor import Tensor
result = Tensor(2) + Tensor(3)
 
# we have a global cache used by the JIT
# from there, we can see the generated clang code
from tinygrad.jit import CacheCollector
CacheCollector.start()       # enables the cache
result.realize()             # create the program and runs it
cache_saved = CacheCollector.finish()  # disable the cache
 
# there's one ASTRunner in the cache
assert len(cache_saved) == 1
 
# print the C Program :
print(cache_saved[0].prg.prg)
 
# after some formatting (the compiler doesn't care)
# NOTE: the 2 and 3 are constant folded
"""
void E_n2(float* restrict data0) {
  data0[0] = (2.0f+3.0f);
}
"""

create a virtual (10, 10) Tensor. this is just a shape, there's no actual tensor
a = ShapeTracker.from_shape((10, 10))
 
# you'll see it has one view. the (10, 1 are the strides)
print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (10, 1), 0)])
 
# we can permute it, and the strides change
a = a.permute((1,0))
print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (1, 10), 0)])
 
# we can then reshape it, and the strides change again
# note how the permute stays applied
a = a.reshape((5,2,5,2))
print(a) # ShapeTracker(shape=(5, 2, 5, 2), views=[View((5, 2, 5, 2), (2, 1, 20, 10), 0)])
 
# now, if we were to reshape it to a (100,) shape tensor, we have to create a second view
a = a.reshape((100,))
print(a) # ShapeTracker(shape=(100,), views=[
         #   View((5, 2, 5, 2), (2, 1, 20, 10), 0),
         #   View((100,), (1,), 0)])
 
# Views stack on top of each other, to allow zero copy for any number of MovementOps
# we can render a Python expression for the index at any time
idx, _ = a.expr_idxs()
print(idx.render())  # (((idx0%10)*10)+(idx0//10))
 
# of course, if we reshape it back, the indexes get simple again
a = a.reshape((10,10))
idx, _ = a.expr_idxs()
print(idx.render())  # ((idx1*10)+idx0)
 
# the ShapeTracker still has two views though...
print(a) # ShapeTracker(shape=(10, 10), views=[
         #   View((5, 2, 5, 2), (2, 1, 20, 10), 0),
         #   View((10, 10), (10, 1), 0)])
 
# ...until we simplify it!
a = a.simplify()
print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (1, 10), 0)])
 
# and now we permute it back
a = a.permute((1,0))
print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (10, 1), 0)])
 
# and it's even contiguous
assert a.contiguous == True
 
# %%
# == Ex: Variable (in tinygrad/shape/symbolic.py, code 6/10) ==
 
# Under the hood, ShapeTracker is powered by a small symbolic algebra library
from tinygrad.shape.symbolic import Variable
 
# Variable is the basic class from symbolic
# it's created with a name and a min and max (inclusive)
a = Variable("a", 0, 10)
b = Variable("b", 0, 10)
 
# some math Ex:s
print((a*10).min, (a*10).max)  # you'll see a*10 has a min of 0 and max of 100
print((a+b).min, (a+b).max)    # 0 20, you get the idea
 
# but complex expressions are where it gets fun
expr = (a + b*10) % 10
print(expr.render())   # (a%10)
# as you can see, b is gone!
 
# one more
expr = (a*40 + b) // 20
print(expr.render())       # (a*2)
print(expr.min, expr.max)  # 0 20
# this is just "(a*2)"
# since b only has a range from 0-10, it can't affect the output