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

namespace TensorOps {

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

  try {
    // Parse optional dim parameter (default: -1)
    int64_t dim = -1;
    if (info.Length() > 0 && info[0].IsNumber()) {
      dim = info[0].As<Napi::Number>().Int64Value();
    }

    // Apply softmax activation: exp(x_i) / sum(exp(x_j)) along dimension
    // PyTorch handles numerical stability internally (subtracts max before exp)
    torch::Tensor result = torch::softmax(self->tensor, dim);
    return Tensor::NewInstance(env, result);
  } catch (const std::exception& e) {
    Napi::Error::New(env, e.what()).ThrowAsJavaScriptException();
    return env.Undefined();
  }
}

}  // namespace TensorOps