#ifndef TENSOR_H
#define TENSOR_H

#include <stdio.h>
#include <vector>
#include <iostream>

class tensor {
public:
    tensor(); // scalar 0
    explicit tensor(double v); // scalar v
    tensor(int dim, size_t shape[], double data[]); // From C
    tensor(int dim, std::initializer_list<size_t> shape, std::initializer_list<double> data); // Conv2d

    int get_dim() const;

    // scalar only
    double item() const;
    double &item();

    // more accessors
    double at(size_t i) const;
    double at(size_t i, size_t j) const;
    double at(size_t i, size_t j, size_t k, size_t l) const;
    double& at(size_t i, size_t j, size_t k, size_t l);

    // passing tensor back to C code
    size_t *get_shape_array();
    double *get_data_array();

    // easy way to print out the tensor for logging
    void print() const;

    size_t getTotalSize() const;

  // friend is better than member, because it allows
  // double + tensor (implicit conversion)
  // operator overloads
  friend tensor operator+(const tensor &lhs, const tensor &rhs);
  friend tensor operator-(const tensor &lhs, const tensor &rhs);
  friend tensor operator*(const tensor &lhs, const tensor &rhs);

  template <class T>
  friend tensor helper(const tensor &lhs, const tensor &rhs);

  friend tensor ReLU(const tensor &x);
  friend tensor Flatten(const tensor &x);
  friend tensor Input2d(const tensor &x);
  friend tensor Linear(const tensor &w, const tensor &x, const tensor &b);
  friend tensor MaxPool2d(const tensor &x, const int kernel_size, const int stride);
  friend tensor Conv2D(const tensor &input_tensor, const tensor &kernel_tensor, const tensor &bias_tensor, int in_channels, int out_channels, int kernel_size);

private:
    int dim_;
    std::vector<size_t> shape_;
    std::vector<double> data_;
}; // class tensor

#endif // TENSOR_H;