#include "tensor.h"

namespace tensor {

std::shared_ptr<Tensor> Tensor::transpose() {
    // 放心，下面的代码暂时不会被触发，我们假定所有的tensor都是2维的
    // if (shape.size() != 2) {
    //     throw std::runtime_error("Transpose is only supported for 2D tensors.");
    // }

    // 这里能够获得矩阵的行数和列数，但是我们是使用一个一维的vector来存储数据的。该如何实现“转置”呢？
    std::size_t rows = shape[0];
    std::size_t cols = shape[1];
    std::vector<size_t> new_shape = {cols, rows};

    // 创建一个空的 Tensor 对象，不进行随机初始化
    auto transposed_tensor = std::make_shared<Tensor>(new_shape, false);

    // 计算转置后的数据
    for(std::size_t i = 0; i < cols; ++i) {
        for(std::size_t j = 0; j < rows; ++j) {
            // 计算转置后的索引
            std::size_t transposed_index = i * rows + j;
            // 计算原数据的索引
            std::size_t original_index = j * cols + i;
            transposed_tensor->data[transposed_index] = data[original_index];
        }
    }

    return transposed_tensor;
}


std::shared_ptr<Tensor> pyarray_to_tensor(py::array_t<float> array) {
    py::buffer_info info = array.request();
    float* dataPtr = static_cast<float*>(info.ptr);
    std::vector<std::size_t> shape = {};
    for (auto &it: info.shape) {
        shape.push_back(it);
    }
    auto tensor = std::make_shared<Tensor>(shape);
    std::vector<float> result(dataPtr, dataPtr + info.size);
    tensor->data = result;
    return tensor;
}

std::shared_ptr<Tensor> argmax(const std::shared_ptr<Tensor>& tensor, int axis) {
    // you only need to handle the two dimensional tensor, and the axis can be either 0 or 1
    // the tensor's shape is (batch_size, features)
    // if the axis is 0, it outputs a tensor (1, features)
    // if the axis is 1, it outputs a tensor (batch_size, 1)

    // compute the output's shape
    std::vector<std::size_t> output_shape = tensor->shape;
    output_shape.erase(output_shape.begin() + axis);

    auto result = std::make_shared<Tensor>(output_shape);
    // 这个问题似乎有点难，所以我们决定给你送点分。一个简单的办法是分axis为0还是为1来进行讨论，反正我们已经把问题简化为了，在一个二维的tensor里面，找到每一行或者每一列的最大值，并输出一个一维的tensor。
    // 补全这里的代码。
       size_t rows=tensor->shape[0];
       size_t cols=tensor->shape[1];
       if(axis==0)
       {
        output_shape={1,cols};
       }
       else if(axis==1)
       {
        output_shape={rows,1};
       }
       if(axis==0)//qiu lie xiang liang 
       {
            for(size_t j=0;j<cols;j++)
            {
                float maxx=0;
                size_t maxx_id=0;//ji lu ID
                for(size_t i=0;i<rows;i++)
                {
                    if(tensor->data[i*cols+j]>maxx) 
                    {
                        maxx=tensor->data[i*cols+j];
                        maxx_id=i;
                    }
                }
                result->data[j] = static_cast<float>(maxx_id);
            }
       }
       else
       {
            for(size_t i=0;i<rows;i++)
            {
               float maxx=0;
                size_t maxx_id=0;
                for (size_t j=0;j<cols;j++)
                {
                    if(maxx<tensor->data[i*cols+j])
                    {
                        maxx=tensor->data[i*cols+j];
                        maxx_id=j;
                    }
                }
                result->data[i] = static_cast<float>(maxx_id);
            }
       }

    return result;
}

std::shared_ptr<Tensor> mean(const std::shared_ptr<Tensor>& tensor) {
    std::vector<std::size_t> shape = {1};
    auto result = std::make_shared<Tensor>(shape);
    auto sum = 0.0f;
    for (auto &it: tensor->data) {
        sum += it;
    }
    sum /= tensor->size;
    result->data[0] = sum;
    return result;
}

std::shared_ptr<Tensor> exp(const std::shared_ptr<Tensor>& tensor) {
    auto result = std::make_shared<Tensor>(tensor->shape);
    for (auto i = 0; i < tensor->size; i++) {
        result->data[i] = expf(tensor->data[i]);
    }
    return result;
}

}