#include "transpose.h"
#include "../tensor.h"
#include "../utils.h"
#include <torch/torch.h>

namespace TensorOps {

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

  if (info.Length() < 2) {
    Napi::TypeError::New(env, "Expected two dimension arguments").ThrowAsJavaScriptException();
    return env.Null();
  }

  if (!info[0].IsNumber() || !info[1].IsNumber()) {
    Napi::TypeError::New(env, "Both dimensions must be numbers").ThrowAsJavaScriptException();
    return env.Null();
  }

  int64_t dim0 = info[0].As<Napi::Number>().Int64Value();
  int64_t dim1 = info[1].As<Napi::Number>().Int64Value();

  // Use PyTorch's transpose method
  torch::Tensor result = self->tensor.transpose(dim0, dim1);
  return Tensor::NewInstance(env, result);
}

}  // namespace TensorOps