#include "utils.h"

namespace TyTorchUtils {

// Helper: Parse dtype string to torch::ScalarType
torch::ScalarType ParseDtype(const std::string& dtype_str) {
  if (dtype_str == "float32" || dtype_str == "f32" || dtype_str == "float") return torch::kFloat32;
  if (dtype_str == "float64" || dtype_str == "f64" || dtype_str == "double") return torch::kFloat64;
  if (dtype_str == "int32" || dtype_str == "i32" || dtype_str == "int") return torch::kInt32;
  if (dtype_str == "int64" || dtype_str == "i64" || dtype_str == "long") return torch::kInt64;
  if (dtype_str == "int16" || dtype_str == "i16" || dtype_str == "short") return torch::kInt16;
  if (dtype_str == "int8" || dtype_str == "i8") return torch::kInt8;
  if (dtype_str == "uint8" || dtype_str == "u8") return torch::kUInt8;
  if (dtype_str == "bool") return torch::kBool;
  return torch::kFloat32; // default
}

// Helper: Parse device string to torch::Device
torch::Device ParseDevice(const std::string& device_str) {
  if (device_str == "cpu") return torch::kCPU;
  if (device_str == "cuda" || device_str == "cuda:0") return torch::Device(torch::kCUDA, 0);
  if (device_str.rfind("cuda:", 0) == 0) {
    int device_id = std::stoi(device_str.substr(5));
    return torch::Device(torch::kCUDA, device_id);
  }
  if (device_str == "mps" || device_str == "mps:0") return torch::Device(torch::kMPS, 0);
  return torch::kCPU; // default
}

}