#include "tensor.h"
#include <assert.h>
#include <stdexcept>
#include <algorithm>
#include <numeric>
#include <math.h>

tensor::tensor() {}

tensor::tensor(double v) {
  shape_.push_back(1);
  data_.push_back(v);
}

tensor::tensor(int dim, std::initializer_list<size_t> shape, std::initializer_list<double> data) {
    shape_.assign(shape.begin(), shape.end());
    data_.assign(data.begin(), data.end());
    dim_ = dim;
}

size_t tensor::getTotalSize() const {
  size_t totalSize = 1;
  for (size_t dim : shape_) {
    if (totalSize > std::numeric_limits<size_t>::max() / dim)
      throw std::overflow_error("Size overflow detected");
    totalSize *= dim;
  }
  return totalSize;
}

tensor::tensor(int dim, size_t shape[], double data[]):
    shape_(shape, shape+dim) {
    dim_ = dim;
    // calc N has shape[0]*shape[1]*...*shape[dim-1]
    // calc N as the product of elements in the shape array
    size_t N = 1;
    for (int i=0; i < dim; ++i) {
        N = shape[i] * N;
    }

    data_.assign(data, data+N);
}

int tensor::get_dim() const {
    return shape_.size();
}

double tensor::item() const {
    if (shape_.empty()) {
        return data_[0];
    } else {
        throw std::runtime_error("item() called on a non-scalar tensor.");
    }
}

double &tensor::item() {
    if (shape_.empty()) {
        return data_[0];
    } else {
        throw std::runtime_error("item() called on a non-scalar tensor.");
    }
}

double tensor::at(size_t i) const {
    assert(get_dim() == 1);
    assert(i < shape_[0]);
    return data_[i];
}

double tensor::at(size_t i, size_t j) const {
    assert(get_dim() == 2);
    if (shape_.size() != 2)
        throw std::runtime_error("Tensor is not 2-dimensional");

    assert((i < shape_[0]) && (j < shape_[1]));
    return data_[i * shape_[1] + j];
}

double tensor::at(size_t i, size_t j, size_t k, size_t l) const {
    assert(get_dim() == 4); // Ensure the tensor is 4D
    assert(i < shape_[0] && j < shape_[1] && k < shape_[2] && l < shape_[3]);

    size_t index = i * shape_[1] * shape_[2] * shape_[3] + j * shape_[2] * shape_[3] + k * shape_[3] + l;
    return data_[index];
}

double& tensor::at(size_t i, size_t j, size_t k, size_t l) {
    assert(get_dim() == 4); // Ensure the tensor is 4D
    assert(i < shape_[0] && j < shape_[1] && k < shape_[2] && l < shape_[3]);

    size_t index = i * shape_[1] * shape_[2] * shape_[3] + j * shape_[2] * shape_[3] + k * shape_[3] + l;
    return data_[index];
}

size_t* tensor::get_shape_array() {
    return shape_.empty()? nullptr: &shape_[0];
}

double* tensor::get_data_array() {
    return &data_[0];
}

// This can be used for debugging across the entire program.
// prints the values of the tensor. Be careful though because
// it will clutter the screen output.
void tensor::print() const {
    std::cout << "Tensor - Dim: " << dim_ << ", Shape: {";
    for (const auto& s : shape_) std::cout << s << " ";
    std::cout << "}, Data: {";
    for (const auto& d : data_) std::cout << d << " ";
    std::cout << "}\n";
}

// helper template
template <class T> tensor helper(const tensor &lhs, const tensor &rhs) {
  tensor result;
  result.shape_ = lhs.shape_;
  size_t totalSize = lhs.getTotalSize();
  result.data_.reserve(totalSize);
  std::transform(lhs.data_.begin(), lhs.data_.begin() + totalSize,
                 rhs.data_.begin(), std::back_inserter(result.data_), T());
  return result;
}

// operator overloads
tensor operator+(const tensor &lhs, const tensor &rhs) {
    tensor tensor_a = tensor(lhs);
    tensor tensor_b = tensor(rhs);

    size_t size_a = lhs.getTotalSize();
    std::vector<double> result;
    for (size_t i =0; i< size_a; ++i) {
        double sum_array = tensor_a.get_data_array()[i] + tensor_b.get_data_array()[i];
        result.push_back(sum_array);
    }
    tensor c(tensor_a.get_dim(), tensor_a.get_shape_array(), &result[0]);
    return c;
}

tensor operator-(const tensor &lhs, const tensor &rhs) {
    tensor tensor_a = tensor(lhs);
    tensor tensor_b = tensor(rhs);

    size_t size_a = lhs.getTotalSize();
    std::vector<double> result;
    for (size_t i =0; i< size_a; ++i) {
        double sum_array = tensor_a.get_data_array()[i] - tensor_b.get_data_array()[i];
        result.push_back(sum_array);
    }
    tensor c(tensor_a.get_dim(), tensor_a.get_shape_array(), &result[0]);
    return c;
}

tensor operator*(const tensor &lhs, const tensor &rhs) 
{   
    try {

        // scalar * scalar multiplication
        if (lhs.getTotalSize() == 1 && rhs.getTotalSize() == 1) {
            return tensor(lhs.data_[0] * rhs.data_[0]);
        }
        if (lhs.shape_.empty() || rhs.shape_.empty()) {
            return tensor(); // for p3-q6, passing in None 
        }

        tensor result;
        size_t lhsSize = lhs.getTotalSize();
        size_t rhsSize = rhs.getTotalSize();
             
        // scalar * tensor multiplication
        if (lhsSize == 1 || rhsSize == 1) {
            double scalar = (lhsSize == 1 ? lhs : rhs).data_[0];
            const tensor &t = lhsSize == 1 ? rhs : lhs;
            result.shape_ = t.shape_;
            result.data_.reserve(lhsSize * rhsSize);
            std::transform(t.data_.begin(), t.data_.end(),
                std::back_inserter(result.data_),
                [&scalar](double x) { return x * scalar; });
        } else {
        // matrix multiplication
            size_t I = lhs.shape_[0];
            size_t J = rhs.shape_[1];
            size_t K = lhs.shape_[1];
            result.shape_ = {I, J};
            size_t totalSize = result.getTotalSize();
            result.data_.resize(totalSize);

            for (size_t i = 0; i < I; ++i) {
                for (size_t j = 0; j < J; ++j) {
                    for (size_t k = 0; k < K; ++k) {
                        result.data_[J * i + j] +=
                            lhs.data_[K * i + k] * rhs.data_[J * k + j];
                    }
                }
            }
        }
        return result;
    } catch (const std::bad_alloc&) {
        throw std::runtime_error("Memory allocation failed in tensor multiplication.");
    }
}

tensor ReLU(const tensor &x){
  tensor result;
  size_t xSize = x.getTotalSize();
  result.shape_ = x.shape_;
  for(size_t i = 0; i < xSize; i++ ){
      result.data_.push_back(x.data_[i]*(x.data_[i]>0));
  }
  return result;
  }

tensor Flatten(const tensor &x){
  tensor result;
  size_t xSize = x.getTotalSize();
  size_t shape_1 = xSize/x.shape_[0];
  std::vector<size_t> shape_;
  shape_.push_back(x.shape_[0]);
  shape_.push_back(shape_1);
  result.shape_ = shape_;
  result.data_ = x.data_;
  return result;
  }

tensor Input2d(const tensor &x){
     tensor result;
     std::vector<size_t> shape_;
     shape_.push_back(x.shape_[0]);
     shape_.push_back(x.shape_[3]);
     shape_.push_back(x.shape_[1]);
     shape_.push_back(x.shape_[2]);
     result.shape_ = shape_;
     for(size_t i = 0; i < x.shape_[0]; i++ ){
      for(size_t j =0; j < x.shape_[3];j++){
        for(size_t k = 0; k <x.shape_[1]*x.shape_[2]; k++){
            result.data_.push_back(x.data_[j + k*x.shape_[3] + i*x.shape_[3]*x.shape_[1]*x.shape_[2]]);
        }
      }
    }
    return result;
}

tensor Linear(const tensor &w, const tensor &x, const tensor &b){
    tensor result_wT;
    std::vector<size_t> shape_0;
    shape_0.push_back(x.shape_[1]);
    shape_0.push_back(w.shape_[0]);
    size_t totalSize_w = w.getTotalSize();
    result_wT.data_.reserve(totalSize_w);
    result_wT.shape_ = shape_0;
    for(size_t i = 0; i < w.shape_[1]; i++){
      for(size_t j = 0; j < w.shape_[0]; j++){
         result_wT.data_.push_back(w.data_[i + j*w.shape_[1]]);
      }
    }
    tensor result_xwT;
    std::vector<size_t> shape_xwT;
    shape_xwT.push_back(x.shape_[0]);
    shape_xwT.push_back(result_wT.shape_[1]);
    size_t totalSize_xwT = x.shape_[0]*result_wT.shape_[1];
    result_xwT.data_.resize(totalSize_xwT);
    result_xwT.shape_ = shape_xwT;
    result_xwT =  x * result_wT;

    tensor result_bBrodcast;
    std::vector<size_t> shape_1;
    shape_1.push_back(b.shape_[0]);
    shape_1.push_back(1*x.shape_[0]);
    result_bBrodcast.shape_ = shape_1;
    for(size_t i = 0; i < result_bBrodcast.shape_[0]; i++){
       for(size_t j = 0; j < result_bBrodcast.shape_[1]; j++){
          result_bBrodcast.data_.push_back(b.data_[i]);
       }
    }

   tensor result_bBroadcastT;
   std::vector<size_t> shape_2;
   size_t totalSize_b = result_bBrodcast.getTotalSize();
   result_bBrodcast.data_.resize(totalSize_b);
   shape_2.push_back(result_bBrodcast.shape_[1]);
   shape_2.push_back(result_bBrodcast.shape_[0]);
   result_bBroadcastT.shape_ = shape_2;
  for(size_t i = 0; i < result_bBrodcast.shape_[1]; i++){
      for(size_t j = 0; j < result_bBrodcast.shape_[0]; j++){
         result_bBroadcastT.data_.push_back(result_bBrodcast.data_[i + j*result_bBrodcast.shape_[1]]);
      }
    }
    tensor result;
    result.data_.resize(result_xwT.shape_[0]*result_xwT.shape_[1]);
    result = result_xwT + result_bBroadcastT;
    return result;
}


std::vector<double> maxpooling(std::vector<std::vector<double> > map, int k_w, int k_h, int s_w, int s_h)
{
	int row = map.size(); int col = map[0].size();
	int out_row = 1+(row-k_h)/s_h, out_col = 1+(col-k_w)/s_w;
	int mod_row = row%s_h, mod_col = col%s_w;
	if (mod_row != 0) out_row++;
	if (mod_col != 0) out_col++;

	std::vector<std::vector<double> > res(out_row, std::vector<double>(out_col, 0));
	std::vector<std::vector<double> > pad_map(map);

	for (int i=0;i<row;i++)
		for (int j = 0; j < k_h - mod_row; j++)
		{
			pad_map[i].push_back(map[i][col - 1]);
		}
	for (int j = 0; j < k_w - mod_col; j++)
	{
		pad_map.push_back(pad_map[row - 1]);
	}

	for(int i=0;i<out_row;i++)
		for (int j = 0; j < out_col; j++)
		{
			int start_x = j*s_w;
			int start_y = i*s_h;
			std::vector<double> temp;
			for(int ii=0;ii<k_w;ii++)
				for (int jj = 0; jj < k_h; jj++)
				{
					temp.push_back(pad_map[start_y + jj][start_x + ii]);
				}
			sort(temp.begin(), temp.end());
			res[i][j] = temp[temp.size() - 1];
		}
       std::vector<double> res_;
      for (auto l : res)
	     {
		       for (auto ll : l)
		       {
			          res_.push_back(ll);
		       }
	      }
	return res_;
}

tensor MaxPool2d(const tensor &x, const int kernel_size, const int stride){
    std::cout <<"======1================="<<std::endl;
    std::cout << x.shape_[0] <<"_"<< x.shape_[1] << x.shape_[2] << x.shape_[3] << std::endl;
    std::cout << kernel_size << std::endl;
    std::cout << stride << std::endl;
    std::cout <<"======2================="<<std::endl;

    tensor result;
    size_t xSize = x.getTotalSize();
    std::cout << xSize << std::endl;
    std::vector<size_t> shape_;
    shape_.push_back(x.shape_[0]);
    shape_.push_back(x.shape_[1]);
    shape_.push_back(floor(x.shape_[2]/kernel_size));
    shape_.push_back(floor(x.shape_[3]/kernel_size));
    result.shape_ = shape_;
    for(size_t i = 0; i < x.shape_[0]; i++ ){
      for(size_t j =0; j < x.shape_[1];j++){
        std::vector<std::vector<double> > tmp;
        for(size_t k =0; k < x.shape_[2];k++){
          std::vector<double> tmp_;
          for(size_t l =0; l < x.shape_[3];l++){
            double int_data = x.data_[l + k*x.shape_[3] + j*x.shape_[3]*x.shape_[2] + i*x.shape_[3]*x.shape_[2]*x.shape_[1] ];
            tmp_.push_back(int_data);
          }
          tmp.push_back(tmp_);
        }
       auto res = maxpooling(tmp,kernel_size,kernel_size,stride,stride);
       for(size_t i = 0; i< res.size(); i++){
         result.data_.push_back(res[i]);
       }
      }
    }
    return result;
  }

tensor Conv2D(const tensor &input_tensor,
              const tensor &kernel_tensor,
              const tensor &bias_tensor,
              int in_channels,
              int out_channels,
              int kernel_size)

{
    std::cout << "In Channels: " << in_channels << ", Out Channels: " << out_channels << std::endl;
    std::cout << "Kernel Size: " << kernel_size << std::endl;

    // Calculate output dimensions
    int out_height = input_tensor.shape_[2] - kernel_size + 1;
    int out_width = input_tensor.shape_[3] - kernel_size + 1;

    // Ensure out_height and out_width are non-negative
    out_height = std::max(out_height, 0);
    out_width = std::max(out_width, 0);

    // Initialize shape array for output tensor
    size_t output_shape[] = {input_tensor.shape_[0], static_cast<size_t>(out_channels), static_cast<size_t>(out_height), static_cast<size_t>(out_width)};

    // Initialize data array for output tensor (fill with zeros)
    std::vector<double> output_data(output_shape[0] * output_shape[1] * output_shape[2] * output_shape[3], 0);

    // Create output tensor
    tensor output_tensor(4, output_shape, output_data.data());

    // Convolution operation
    for (size_t n = 0; n < input_tensor.shape_[0]; ++n) { // for each batch
        for (int m = 0; m < out_channels; ++m) { // for each filter
            for (int y = 0; y < out_height; ++y) {
                for (int x = 0; x < out_width; ++x) {
                    double sum = 0.0;
                    for (int c = 0; c < in_channels; ++c) { // for each channel
                        for (int ky = 0; ky < kernel_size; ++ky) {
                            for (int kx = 0; kx < kernel_size; ++kx) {
                                size_t ix = x + kx;
                                size_t iy = y + ky;
                                if (ix < input_tensor.shape_[3] && iy < input_tensor.shape_[2]) {
                                    sum += input_tensor.at(n, c, iy, ix) * kernel_tensor.at(m, c, ky, kx);
                                }
                            }
                        }
                    }
                    sum += bias_tensor.at(m); // add bias
                    output_tensor.at(n, m, y, x) = sum;
                }
            }
        }
    }

    return output_tensor;
}