from myTensor import Tensor_Float as myTensor_f
from myTensor import Tensor_Int as myTensor_i
class Value:
    op: Optional[Op]
    inputs: List["Value"]
    cached_data: myTensor_f # or myTensor_i
    requires_grad: bool

class Tensor(Value):
    grad: "Tensor"
# should know that the device be always 'GPU'
# and that dtype be merely only float32. dtype value unused
    def __init__(
        self,
        array,
        *,
        device='GPU',
        dtype=None,
        requires_grad=True,
        **kwargs
    ):
        if isinstance(array, Tensor):
            # copy initialize deom a tensor
            # ...
        elif isinstance(array, myTensor_f):
            # initialize from myTensor_f
            # ...
        elif isinstance(array, myTensor_i):
            # initialize from myTensor_f
            # ...
        else:
            # simply raise error from that
            raise ValueError("error in Tensor init: from something wrong")

        self._init(
            None,
            [],
            cached_data=cached_data,
            requires_grad=requires_grad,
        )

# ops. The below is all you need.
class EWiseAdd(TensorOp):
    # needed for __add__
    # ...
class Negate(TensorOp):
    # get negative. used to def __sub__
    # ...
class MulScalar(TensorOp):
    # multiply a scalar
    # if need multiplying between Tensors, use FC(use_bias=False) instead
    # used to def __mul__, and used in SGD and Adam
    # ...
class Sqrt(TensorOp):
    # sqrt of a Tensor, element-wise. 
    # used in Adam optimizer
    # ...
class Relu(TensorOp):
    # relu
    # ...
class Sigmoid(TensorOp):
    # sigmoid
    # ...
class FC(TensorOp):
    # fully-connected layers
    # you can set use_bias to control if there need a bias in FC layers
    # ...
class Conv(TensorOp):
    # convolution
    # ...
class Maxpool(TensorOp):
    # maxpooling
    # ...
class SftCrossEn(TensorOp):
    # softmax & cross entropy
    # ...
class Flat(TensorOp):
    # hold the Tensor flat, just as torch.view(-1) 
    # ...

# Step 2: define the CNN
class SimpleCNN():
    def __init__(self):
        # CNN parameters initialized
        self.weights = [self.conv1, self.conv2, self.fcw1, self.fcw2, self.fcb1, self.fcb2]
        # initialize the weights, and upload them in the optimizers.
        # ...

    def forward(self, x, real_labels):
        # forward method defined with TensorOps, which have been build upon myLayers
        # that's what makes a CNN
        x = maxpool(relu(conv(x, self.conv1)))
        x = maxpool(relu(conv(x, self.conv2)))
        x = flat(x)
        x = fc(fc(x, self.fcw1, self.fcb1), self.fcw2, self.fcb2)
        x = sftcrossen(x, real_labels)
        return x
    
    def SGD_epoch(self, X, y, lr=0.1, batch=100):
        # SGD inplemented
        # ...
        # when you need autograd, just do the below:
        tr_loss = self.forward(X_batch, y_batch)
        tr_loss.backward()

        # and then update
        for i in range(len(self.weights)):
            self.weights[i] = self.weights[i] - lr * self.weights[i].grad

    def Adam_epoch(self, X, y, lr=0.1, batch=100, beta1=0.9, beta2=0.999, epsilon=1e-8):
        # Adam optimizer inplemented
        # ...

    def train_nn(self, X_tr, y_tr, X_te, y_te, 
                 epochs=10, lr=0.5, batch=100, 
                 beta1=0.9, beta2=0.999, use_Adam=False):
        # almost the same as that in H06
        # ...