#include "binary_cross_entropy.h"
#include "../tensor.h"
#include <torch/torch.h>

namespace TensorOps {

Napi::Value BinaryCrossEntropy(Tensor* self, const Napi::CallbackInfo& info) {
  Napi::Env env = info.Env();

  try {
    // Binary cross entropy requires a target tensor
    if (info.Length() < 1 || !info[0].IsObject()) {
      Napi::TypeError::New(env, "binary_cross_entropy requires a target tensor").ThrowAsJavaScriptException();
      return env.Undefined();
    }

    // Get target tensor
    Tensor* target = Napi::ObjectWrap<Tensor>::Unwrap(info[0].As<Napi::Object>());

    // Parse optional weight parameter
    torch::Tensor weight;
    bool has_weight = false;
    if (info.Length() > 1 && !info[1].IsNull() && !info[1].IsUndefined() && info[1].IsObject()) {
      Tensor* weight_tensor = Napi::ObjectWrap<Tensor>::Unwrap(info[1].As<Napi::Object>());
      weight = weight_tensor->tensor;
      has_weight = true;
    }

    // Parse optional reduction parameter (default: "mean")
    std::string reduction = "mean";
    if (info.Length() > 2 && info[2].IsString()) {
      reduction = info[2].As<Napi::String>().Utf8Value();
    }

    // Convert reduction string to torch reduction enum
    int64_t reduction_enum;
    if (reduction == "none") {
      reduction_enum = 0;  // at::Reduction::None
    } else if (reduction == "mean") {
      reduction_enum = 1;  // at::Reduction::Mean
    } else if (reduction == "sum") {
      reduction_enum = 2;  // at::Reduction::Sum
    } else {
      Napi::TypeError::New(env, "reduction must be 'none', 'mean', or 'sum'").ThrowAsJavaScriptException();
      return env.Undefined();
    }

    // Compute binary cross entropy loss
    torch::Tensor result;
    if (has_weight) {
      result = torch::binary_cross_entropy(
        self->tensor,
        target->tensor,
        weight,
        reduction_enum
      );
    } else {
      result = torch::binary_cross_entropy(
        self->tensor,
        target->tensor,
        {},
        reduction_enum
      );
    }

    return Tensor::NewInstance(env, result);
  } catch (const std::exception& e) {
    Napi::Error::New(env, e.what()).ThrowAsJavaScriptException();
    return env.Undefined();
  }
}

}  // namespace TensorOps