#include "eval_conv2d.h"

eval_conv2d::eval_conv2d(const expression &expr): eval_op(expr), 
      in_channels(expr.get_op_param_int("in_channels")),
      kernel_size(expr.get_op_param_int("kernel_size")),
      out_channels(expr.get_op_param_int("out_channels")),
      padding(expr.get_op_param_double("padding")),
      weight(expr.get_op_param_tensor("weight")),
      bias(expr.get_op_param_tensor("bias")),
      stride(expr.get_op_param_int("stride"))
    {

    }

void eval_conv2d::eval(vars_type &variables, const kwargs_type &kwargs) {
    if (inputs_.empty()) {
        throw std::runtime_error("No input tensor for eval_conv2d.");
    }

    // Get the input tensor from variables using the expression ID of the input
    int inputExprId = inputs_[0];
    if (variables.find(inputExprId) == variables.end()) {
        throw std::runtime_error("Input tensor not found in variables map.");
    }
    tensor input_tensor = variables.at(inputExprId);

    // Perform the Conv2D operation
    tensor conv2d_tensor = Conv2D(input_tensor, weight, bias, in_channels, out_channels, kernel_size);

    // Store the result in the variables map using the current expression ID
    variables[expr_id_] = conv2d_tensor;
}

std::shared_ptr<eval_op> eval_conv2d::clone(const expression &expr) {
    return std::make_shared<eval_conv2d>(expr);
}