#include "conv2d.hpp"

#include <format>
#include <stdexcept>

#include "layer.hpp"
#include "tensor.hpp"

using namespace CUDANet::Layers;

Conv2d::Conv2d(
    CUDANet::Shape    input_shape,
    CUDANet::Shape    kernel_shape,
    CUDANet::Shape    stride_shape,
    CUDANet::Shape    padding_shape,
    CUDANet::Backend* backend
)
    : in_shape(input_shape),
      kernel_shape(kernel_shape),
      stride_shape(stride_shape),
      padding_shape(padding_shape),
      backend(backend) {
    if (in_shape.size() != 3) {
        throw std::runtime_error(
            std::format(
                "Invalid input shape. Expected 3 dims, got {}", in_shape
            )
        );
    }

    if (kernel_shape.size() != 3) {
        throw std::runtime_error(
            std::format(
                "Invalid kernel shape. Expected 3 dims, got {}", kernel_shape
            )
        );
    }

    if (stride_shape.size() != 2) {
        throw std::runtime_error(
            std::format(
                "Invalid stride shape. Expected 2 dims, got {}", stride_shape
            )
        );
    }

    if (padding_shape.size() != 2) {
        throw std::runtime_error(
            std::format(
                "Invalid padding shape. Expected 2 dims, got {}", padding_shape
            )
        );
    }

    size_t out_h = (in_shape[0] - kernel_shape[0] + 2 * padding_shape[0]) /
                       stride_shape[0] +
                   1;
    size_t out_w = (in_shape[1] - kernel_shape[1] + 2 * padding_shape[1]) /
                       stride_shape[1] +
                   1;
    out_shape.resize(3);
    out_shape[0] = out_h;
    out_shape[1] = out_w;
    out_shape[2] = kernel_shape[2];
    output = CUDANet::Tensor(
        Shape{out_shape[0] * out_shape[1] * out_shape[3]},
        CUDANet::DType::FLOAT32, backend
    );

    weights = CUDANet::Tensor(
        Shape{
            kernel_shape[0] * kernel_shape[1] * kernel_shape[2] * in_shape[2]
        },
        CUDANet::DType::FLOAT32, backend
    );
    biases = CUDANet::Tensor(
        Shape{kernel_shape[2]}, CUDANet::DType::FLOAT32, backend
    );

    weights.zero();
    biases.zero();
}

Conv2d::~Conv2d() {}

CUDANet::Tensor& Conv2d::forward(const CUDANet::Tensor& input) {
    output.zero();
    backend->conv2d(
        weights,
        biases,
        input,
        output,
        in_shape,
        padding_shape,
        kernel_shape,
        stride_shape,
        out_shape
    );
    return output;
}

CUDANet::Shape Conv2d::input_shape() {
    return in_shape;
}

CUDANet::Shape Conv2d::output_shape() {
    return out_shape;
}

size_t Conv2d::input_size() {
    return sizeof(float) * in_shape[0] * in_shape[1] * in_shape[2];
}

size_t Conv2d::output_size() {
    return sizeof(float) * out_shape[0] * out_shape[1] * out_shape[2];
}

void Conv2d::set_weights(void* input) {
    weights.set_data<float>(static_cast<float*>(input));
}

CUDANet::Tensor& Conv2d::get_weights() {
    return weights;
}

void Conv2d::set_biases(void* input) {
    biases.set_data<float>(static_cast<float*>(input));
}

CUDANet::Tensor& Conv2d::get_biases() {
    return biases;
}

CUDANet::Shape Conv2d::get_padding_shape() {
    return padding_shape;
}