#include <torch/extension.h>

#include <ATen/ops/bmm_ops.h>
#include <ATen/ops/mm_ops.h>
#include <ATen/ops/mul_ops.h>

#ifdef QIBOTN_USE_MKL
#include <mkl_cblas.h>
#include <mkl_service.h>
#endif

#include <cstdint>
#include <cstdlib>
#include <complex>
#include <algorithm>
#include <stdexcept>
#include <string>
#include <tuple>
#include <unordered_map>
#include <vector>

namespace {

using Tensor = torch::Tensor;

struct ContractOp {
  int64_t out_id;
  int64_t left_id;
  int64_t right_id;
  int64_t kind;
  std::vector<int64_t> left_perm;
  std::vector<int64_t> left_shape;
  std::vector<int64_t> right_perm;
  std::vector<int64_t> right_shape;
  std::vector<int64_t> out_shape;
  std::vector<int64_t> out_perm;
};

bool env_enabled(const char* name) {
  const char* raw = std::getenv(name);
  if (raw == nullptr) {
    return false;
  }
  std::string value(raw);
  return value == "1" || value == "true" || value == "TRUE" ||
      value == "yes" || value == "on";
}

int env_int(const char* name, int default_value) {
  const char* raw = std::getenv(name);
  if (raw == nullptr || raw[0] == '\0') {
    return default_value;
  }
  char* end = nullptr;
  long value = std::strtol(raw, &end, 10);
  if (end == raw || value <= 0) {
    return default_value;
  }
  return static_cast<int>(value);
}

int64_t env_i64(const char* name, int64_t default_value) {
  const char* raw = std::getenv(name);
  if (raw == nullptr || raw[0] == '\0') {
    return default_value;
  }
  char* end = nullptr;
  long long value = std::strtoll(raw, &end, 10);
  if (end == raw || value <= 0) {
    return default_value;
  }
  return static_cast<int64_t>(value);
}

struct WorkspaceKey {
  c10::ScalarType dtype;
  c10::DeviceType device_type;
  c10::DeviceIndex device_index;
  int64_t numel;

  bool operator==(const WorkspaceKey& other) const {
    return dtype == other.dtype && device_type == other.device_type &&
        device_index == other.device_index && numel == other.numel;
  }
};

struct WorkspaceKeyHash {
  size_t operator()(const WorkspaceKey& key) const {
    size_t h = static_cast<size_t>(key.numel);
    h ^= static_cast<size_t>(key.dtype) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2);
    h ^= static_cast<size_t>(key.device_type) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2);
    h ^= static_cast<size_t>(key.device_index) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2);
    return h;
  }
};

class Workspace {
 public:
  Workspace()
      : max_cached_bytes_(read_size_env(
            "QIBOTN_CPP_WORKSPACE_MAX_BYTES",
            1LL * 1024 * 1024 * 1024)),
        max_buffer_bytes_(read_size_env(
            "QIBOTN_CPP_WORKSPACE_MAX_BUFFER_BYTES",
            256LL * 1024 * 1024)) {}

  Tensor acquire(
      const std::vector<int64_t>& shape,
      c10::ScalarType dtype,
      c10::Device device) {
    const auto numel = product(shape);
    WorkspaceKey key{dtype, device.type(), device.index(), numel};
    auto& bucket = pool_[key];
    Tensor out;
    if (!bucket.empty()) {
      out = bucket.back();
      bucket.pop_back();
      cached_bytes_ -= tensor_bytes(out);
      out.resize_(shape);
    } else {
      out = torch::empty(shape, torch::TensorOptions().dtype(dtype).device(device));
    }
    return out;
  }

  void release(const Tensor& x) {
    if (!x.defined() || !x.is_contiguous() || x.storage_offset() != 0) {
      return;
    }
    const auto bytes = tensor_bytes(x);
    if (bytes > max_buffer_bytes_) {
      return;
    }
    if (cached_bytes_ + bytes > max_cached_bytes_) {
      return;
    }
    WorkspaceKey key{
        x.scalar_type(),
        x.device().type(),
        x.device().index(),
        x.numel(),
    };
    pool_[key].push_back(x);
    cached_bytes_ += bytes;
  }

 private:
  static int64_t read_size_env(const char* name, int64_t default_value) {
    const char* raw = std::getenv(name);
    if (raw == nullptr || raw[0] == '\0') {
      return default_value;
    }
    char* end = nullptr;
    const auto value = std::strtod(raw, &end);
    if (end == raw || value < 0) {
      return default_value;
    }
    std::string suffix(end);
    int64_t multiplier = 1;
    if (suffix == "g" || suffix == "G" || suffix == "gb" || suffix == "GB") {
      multiplier = 1024LL * 1024 * 1024;
    } else if (
        suffix == "m" || suffix == "M" || suffix == "mb" || suffix == "MB") {
      multiplier = 1024LL * 1024;
    } else if (
        suffix == "k" || suffix == "K" || suffix == "kb" || suffix == "KB") {
      multiplier = 1024LL;
    }
    return static_cast<int64_t>(value * multiplier);
  }

  static int64_t product(const std::vector<int64_t>& shape) {
    int64_t out = 1;
    for (const auto value : shape) {
      out *= value;
    }
    return out;
  }

  static int64_t tensor_bytes(const Tensor& x) {
    return x.numel() * x.element_size();
  }

  std::unordered_map<WorkspaceKey, std::vector<Tensor>, WorkspaceKeyHash> pool_;
  int64_t cached_bytes_ = 0;
  int64_t max_cached_bytes_;
  int64_t max_buffer_bytes_;
};

struct PreparedTensor {
  Tensor tensor;
  bool workspace_owned;
};

std::vector<int64_t> tensor_sizes(const Tensor& x) {
  return std::vector<int64_t>(x.sizes().begin(), x.sizes().end());
}

PreparedTensor apply_preparation(
    Workspace& workspace,
    Tensor x,
    const std::vector<int64_t>& perm,
    const std::vector<int64_t>& shape) {
  bool workspace_owned = false;
  if (!perm.empty()) {
    x = x.permute(perm);
  }
  if (!shape.empty()) {
    try {
      x = x.view(shape);
    } catch (const c10::Error&) {
      Tensor out = workspace.acquire(shape, x.scalar_type(), x.device());
      out.view(tensor_sizes(x)).copy_(x);
      x = out;
      workspace_owned = true;
    }
  }
  return PreparedTensor{x, workspace_owned};
}

std::vector<int64_t> broadcast_shape(
    c10::IntArrayRef left,
    c10::IntArrayRef right) {
  const auto ndim = std::max(left.size(), right.size());
  std::vector<int64_t> out(ndim, 1);
  for (size_t i = 0; i < ndim; ++i) {
    const auto li = left.size() < ndim - i ? 1 : left[i - (ndim - left.size())];
    const auto ri = right.size() < ndim - i ? 1 : right[i - (ndim - right.size())];
    if (li != ri && li != 1 && ri != 1) {
      throw std::runtime_error("Shapes are not broadcastable.");
    }
    out[i] = std::max<int64_t>(li, ri);
  }
  return out;
}

std::vector<int64_t> matmul_output_shape(const Tensor& left, const Tensor& right) {
  if (left.dim() == 2 && right.dim() == 2) {
    return {left.size(-2), right.size(-1)};
  }
  auto batch = broadcast_shape(
      left.sizes().slice(0, left.dim() - 2),
      right.sizes().slice(0, right.dim() - 2));
  batch.push_back(left.size(-2));
  batch.push_back(right.size(-1));
  return batch;
}

std::vector<int64_t> multiply_output_shape(const Tensor& left, const Tensor& right) {
  return broadcast_shape(left.sizes(), right.sizes());
}

bool direct_cblas_mm(const Tensor& left, const Tensor& right, Tensor& out) {
#ifndef QIBOTN_USE_MKL
  return false;
#else
  static const bool direct_mkl_enabled = env_enabled("QIBOTN_CPP_DIRECT_MKL");
  if (!direct_mkl_enabled) {
    return false;
  }
  if (left.device().type() != c10::DeviceType::CPU ||
      right.device().type() != c10::DeviceType::CPU ||
      out.device().type() != c10::DeviceType::CPU) {
    return false;
  }
  if (left.dim() != 2 || right.dim() != 2 || out.dim() != 2) {
    return false;
  }
  if (!left.is_contiguous() || !right.is_contiguous() || !out.is_contiguous()) {
    return false;
  }
  if (left.scalar_type() != right.scalar_type() ||
      left.scalar_type() != out.scalar_type()) {
    return false;
  }
  if (left.size(1) != right.size(0) ||
      out.size(0) != left.size(0) ||
      out.size(1) != right.size(1)) {
    return false;
  }

  const int m = static_cast<int>(left.size(0));
  const int k = static_cast<int>(left.size(1));
  const int n = static_cast<int>(right.size(1));
  const int lda = static_cast<int>(left.stride(0));
  const int ldb = static_cast<int>(right.stride(0));
  const int ldc = static_cast<int>(out.stride(0));
  const int mkl_threads = env_int("QIBOTN_CPP_MKL_THREADS", 0);
  const int previous_threads =
      mkl_threads > 0 ? mkl_set_num_threads_local(mkl_threads) : 0;

  if (left.scalar_type() == c10::ScalarType::ComplexFloat) {
    const std::complex<float> alpha(1.0f, 0.0f);
    const std::complex<float> beta(0.0f, 0.0f);
    cblas_cgemm(
        CblasRowMajor,
        CblasNoTrans,
        CblasNoTrans,
        m,
        n,
        k,
        &alpha,
        left.data_ptr(),
        lda,
        right.data_ptr(),
        ldb,
        &beta,
        out.data_ptr(),
        ldc);
    if (mkl_threads > 0) {
      mkl_set_num_threads_local(previous_threads);
    }
    return true;
  }
  if (left.scalar_type() == c10::ScalarType::ComplexDouble) {
    const std::complex<double> alpha(1.0, 0.0);
    const std::complex<double> beta(0.0, 0.0);
    cblas_zgemm(
        CblasRowMajor,
        CblasNoTrans,
        CblasNoTrans,
        m,
        n,
        k,
        &alpha,
        left.data_ptr(),
        lda,
        right.data_ptr(),
        ldb,
        &beta,
        out.data_ptr(),
        ldc);
    if (mkl_threads > 0) {
      mkl_set_num_threads_local(previous_threads);
    }
    return true;
  }
  if (mkl_threads > 0) {
    mkl_set_num_threads_local(previous_threads);
  }
  return false;
#endif
}

Tensor apply_output_preparation(
    Tensor x,
    const std::vector<int64_t>& shape,
    const std::vector<int64_t>& perm) {
  if (!shape.empty()) {
    x = x.reshape(shape);
  }
  if (!perm.empty()) {
    x = x.permute(perm);
  }
  return x;
}

struct ExecutedOp {
  Tensor left;
  Tensor right;
  bool left_workspace_owned = false;
  bool right_workspace_owned = false;
};

void release_inputs(
    const ContractOp& op,
    const ExecutedOp& executed,
    size_t ninputs,
    std::vector<Tensor>& temps,
    Workspace& workspace) {
  if (executed.left_workspace_owned) {
    workspace.release(executed.left);
  }
  if (executed.right_workspace_owned) {
    workspace.release(executed.right);
  }
  if (op.left_id >= static_cast<int64_t>(ninputs)) {
    workspace.release(temps.at(static_cast<size_t>(op.left_id)));
    temps.at(static_cast<size_t>(op.left_id)) = Tensor();
  }
  if (op.right_id >= static_cast<int64_t>(ninputs)) {
    workspace.release(temps.at(static_cast<size_t>(op.right_id)));
    temps.at(static_cast<size_t>(op.right_id)) = Tensor();
  }
}

ExecutedOp execute_single_op(
    const ContractOp& op,
    std::vector<Tensor>& temps,
    Workspace& workspace) {
  Tensor left = temps.at(static_cast<size_t>(op.left_id));
  Tensor right = temps.at(static_cast<size_t>(op.right_id));
  if (!left.defined() || !right.defined()) {
    throw std::runtime_error("Contraction plan referenced a released tensor.");
  }

  auto prepared_left = apply_preparation(workspace, left, op.left_perm, op.left_shape);
  auto prepared_right = apply_preparation(workspace, right, op.right_perm, op.right_shape);
  left = prepared_left.tensor;
  right = prepared_right.tensor;

  Tensor out;
  if (op.kind == 0) {
    auto out_shape = matmul_output_shape(left, right);
    out = workspace.acquire(out_shape, left.scalar_type(), left.device());
    if (left.dim() == 2 && right.dim() == 2) {
      if (!direct_cblas_mm(left, right, out)) {
        at::_ops::mm_out::call(left, right, out);
      }
    } else if (left.dim() == 3 && right.dim() == 3) {
      at::_ops::bmm_out::call(left, right, out);
    } else {
      out = torch::matmul(left, right);
    }
  } else if (op.kind == 1) {
    auto out_shape = multiply_output_shape(left, right);
    out = workspace.acquire(out_shape, left.scalar_type(), left.device());
    at::_ops::mul_out::call(left, right, out);
  } else {
    throw std::runtime_error("Unknown contraction op kind.");
  }

  out = apply_output_preparation(out, op.out_shape, op.out_perm);
  temps.at(static_cast<size_t>(op.out_id)) = out;
  return ExecutedOp{left, right, prepared_left.workspace_owned, prepared_right.workspace_owned};
}

struct GemmBatchKey {
  c10::ScalarType dtype;
  int64_t m;
  int64_t k;
  int64_t n;

  bool operator==(const GemmBatchKey& other) const {
    return dtype == other.dtype && m == other.m && k == other.k && n == other.n;
  }
};

struct GemmBatchKeyHash {
  size_t operator()(const GemmBatchKey& key) const {
    size_t h = static_cast<size_t>(key.m);
    h ^= static_cast<size_t>(key.k) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2);
    h ^= static_cast<size_t>(key.n) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2);
    h ^= static_cast<size_t>(key.dtype) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2);
    return h;
  }
};

bool static_small_2d_gemm_key(
    const ContractOp& op,
    int64_t max_flops,
    GemmBatchKey& key) {
  if (op.kind != 0 || op.left_shape.size() != 2 || op.right_shape.size() != 2) {
    return false;
  }
  const int64_t m = op.left_shape[0];
  const int64_t k = op.left_shape[1];
  const int64_t n = op.right_shape[1];
  if (op.right_shape[0] != k || m * k * n > max_flops) {
    return false;
  }
  key = GemmBatchKey{c10::ScalarType::Undefined, m, k, n};
  return true;
}

#ifdef QIBOTN_USE_MKL
bool execute_mkl_batch(
    const std::vector<int>& op_indices,
    const std::vector<ContractOp>& ops,
    std::vector<Tensor>& temps,
    Workspace& workspace,
    std::vector<ExecutedOp>& executed) {
  if (op_indices.empty()) {
    return false;
  }

  const auto& first_op = ops[op_indices[0]];
  std::vector<Tensor> lefts;
  std::vector<Tensor> rights;
  std::vector<Tensor> outs;
  lefts.reserve(op_indices.size());
  rights.reserve(op_indices.size());
  outs.reserve(op_indices.size());
  executed.clear();
  executed.reserve(op_indices.size());

  c10::ScalarType dtype = c10::ScalarType::Undefined;
  int m = 0;
  int k = 0;
  int n = 0;
  for (const int op_index : op_indices) {
    const auto& op = ops[op_index];
    Tensor left = temps.at(static_cast<size_t>(op.left_id));
    Tensor right = temps.at(static_cast<size_t>(op.right_id));
    if (!left.defined() || !right.defined()) {
      throw std::runtime_error("Batch contraction referenced a released tensor.");
    }
    auto prepared_left = apply_preparation(workspace, left, op.left_perm, op.left_shape);
    auto prepared_right = apply_preparation(workspace, right, op.right_perm, op.right_shape);
    left = prepared_left.tensor;
    right = prepared_right.tensor;
    if (left.dim() != 2 || right.dim() != 2 || !left.is_contiguous() ||
        !right.is_contiguous() || left.device().type() != c10::DeviceType::CPU ||
        right.device().type() != c10::DeviceType::CPU ||
        left.scalar_type() != right.scalar_type()) {
      return false;
    }
    if (dtype == c10::ScalarType::Undefined) {
      dtype = left.scalar_type();
      m = static_cast<int>(left.size(0));
      k = static_cast<int>(left.size(1));
      n = static_cast<int>(right.size(1));
    }
    if (left.scalar_type() != dtype || left.size(0) != m || left.size(1) != k ||
        right.size(0) != k || right.size(1) != n) {
      return false;
    }
    Tensor out = workspace.acquire({m, n}, dtype, left.device());
    if (!out.is_contiguous()) {
      return false;
    }
    lefts.push_back(left);
    rights.push_back(right);
    outs.push_back(out);
    executed.push_back(
        ExecutedOp{left, right, prepared_left.workspace_owned, prepared_right.workspace_owned});
  }

  if (dtype != c10::ScalarType::ComplexFloat && dtype != c10::ScalarType::ComplexDouble) {
    return false;
  }

  const int group_count = 1;
  const int group_size[] = {static_cast<int>(op_indices.size())};
  const CBLAS_TRANSPOSE transa[] = {CblasNoTrans};
  const CBLAS_TRANSPOSE transb[] = {CblasNoTrans};
  const int m_array[] = {m};
  const int n_array[] = {n};
  const int k_array[] = {k};
  const int lda[] = {k};
  const int ldb[] = {n};
  const int ldc[] = {n};
  std::vector<const void*> a_array(op_indices.size());
  std::vector<const void*> b_array(op_indices.size());
  std::vector<void*> c_array(op_indices.size());
  for (size_t i = 0; i < op_indices.size(); ++i) {
    a_array[i] = lefts[i].data_ptr();
    b_array[i] = rights[i].data_ptr();
    c_array[i] = outs[i].data_ptr();
  }

  const int mkl_threads = env_int("QIBOTN_CPP_BATCH_MKL_THREADS", 0);
  const int previous_threads =
      mkl_threads > 0 ? mkl_set_num_threads_local(mkl_threads) : 0;
  if (dtype == c10::ScalarType::ComplexFloat) {
    const std::complex<float> alpha(1.0f, 0.0f);
    const std::complex<float> beta(0.0f, 0.0f);
    const void* alpha_array[] = {&alpha};
    const void* beta_array[] = {&beta};
    cblas_cgemm_batch(
        CblasRowMajor,
        transa,
        transb,
        m_array,
        n_array,
        k_array,
        alpha_array,
        a_array.data(),
        lda,
        b_array.data(),
        ldb,
        beta_array,
        c_array.data(),
        ldc,
        group_count,
        group_size);
  } else {
    const std::complex<double> alpha(1.0, 0.0);
    const std::complex<double> beta(0.0, 0.0);
    const void* alpha_array[] = {&alpha};
    const void* beta_array[] = {&beta};
    cblas_zgemm_batch(
        CblasRowMajor,
        transa,
        transb,
        m_array,
        n_array,
        k_array,
        alpha_array,
        a_array.data(),
        lda,
        b_array.data(),
        ldb,
        beta_array,
        c_array.data(),
        ldc,
        group_count,
        group_size);
  }
  if (mkl_threads > 0) {
    mkl_set_num_threads_local(previous_threads);
  }

  for (size_t i = 0; i < op_indices.size(); ++i) {
    const auto& op = ops[op_indices[i]];
    Tensor out = apply_output_preparation(outs[i], op.out_shape, op.out_perm);
    temps.at(static_cast<size_t>(op.out_id)) = out;
  }
  return true;
}
#endif

std::vector<int64_t> list_to_i64_vector(const pybind11::handle& object) {
  std::vector<int64_t> values;
  for (const auto item : object) {
    values.push_back(pybind11::cast<int64_t>(item));
  }
  return values;
}

ContractOp parse_op(const pybind11::handle& object) {
  auto tuple = pybind11::cast<pybind11::tuple>(object);
  if (tuple.size() != 10) {
    throw std::runtime_error("Compiled contraction op must have 10 fields.");
  }

  ContractOp op;
  op.out_id = pybind11::cast<int64_t>(tuple[0]);
  op.left_id = pybind11::cast<int64_t>(tuple[1]);
  op.right_id = pybind11::cast<int64_t>(tuple[2]);
  op.kind = pybind11::cast<int64_t>(tuple[3]);
  op.left_perm = list_to_i64_vector(tuple[4]);
  op.left_shape = list_to_i64_vector(tuple[5]);
  op.right_perm = list_to_i64_vector(tuple[6]);
  op.right_shape = list_to_i64_vector(tuple[7]);
  op.out_shape = list_to_i64_vector(tuple[8]);
  op.out_perm = list_to_i64_vector(tuple[9]);
  return op;
}

std::vector<ContractOp> parse_plan(const pybind11::list& plan) {
  std::vector<ContractOp> ops;
  ops.reserve(plan.size());
  for (const auto object : plan) {
    ops.push_back(parse_op(object));
  }
  return ops;
}

Tensor run_contract(
    const std::vector<Tensor>& arrays,
    const std::vector<ContractOp>& ops,
    int64_t ntemps,
    int64_t root_id,
    Workspace& workspace) {
  c10::InferenceMode inference_guard(true);
  std::vector<Tensor> temps(static_cast<size_t>(ntemps));

  if (arrays.size() > temps.size()) {
    throw std::runtime_error("Temporary tensor table is smaller than input arrays.");
  }
  for (size_t i = 0; i < arrays.size(); ++i) {
    temps[i] = arrays[i];
  }

  for (const auto& op : ops) {
    Tensor left = temps.at(static_cast<size_t>(op.left_id));
    Tensor right = temps.at(static_cast<size_t>(op.right_id));
    if (!left.defined() || !right.defined()) {
      throw std::runtime_error("Contraction plan referenced a released tensor.");
    }

    auto prepared_left = apply_preparation(workspace, left, op.left_perm, op.left_shape);
    auto prepared_right = apply_preparation(workspace, right, op.right_perm, op.right_shape);
    left = prepared_left.tensor;
    right = prepared_right.tensor;

    Tensor out;
    if (op.kind == 0) {
      auto out_shape = matmul_output_shape(left, right);
      out = workspace.acquire(out_shape, left.scalar_type(), left.device());
      if (left.dim() == 2 && right.dim() == 2) {
        if (!direct_cblas_mm(left, right, out)) {
          at::_ops::mm_out::call(left, right, out);
        }
      } else if (left.dim() == 3 && right.dim() == 3) {
        at::_ops::bmm_out::call(left, right, out);
      } else {
        out = torch::matmul(left, right);
      }
    } else if (op.kind == 1) {
      auto out_shape = multiply_output_shape(left, right);
      out = workspace.acquire(out_shape, left.scalar_type(), left.device());
      at::_ops::mul_out::call(left, right, out);
    } else {
      throw std::runtime_error("Unknown contraction op kind.");
    }

    out = apply_output_preparation(out, op.out_shape, op.out_perm);
    temps.at(static_cast<size_t>(op.out_id)) = out;

    if (prepared_left.workspace_owned) {
      workspace.release(left);
    }
    if (prepared_right.workspace_owned) {
      workspace.release(right);
    }
    if (op.left_id >= static_cast<int64_t>(arrays.size())) {
      workspace.release(temps.at(static_cast<size_t>(op.left_id)));
      temps.at(static_cast<size_t>(op.left_id)) = Tensor();
    }
    if (op.right_id >= static_cast<int64_t>(arrays.size())) {
      workspace.release(temps.at(static_cast<size_t>(op.right_id)));
      temps.at(static_cast<size_t>(op.right_id)) = Tensor();
    }
  }

  Tensor result = temps.at(static_cast<size_t>(root_id));
  if (!result.defined()) {
    throw std::runtime_error("Contraction plan did not produce a result tensor.");
  }
  return result;
}

Tensor run_contract_dag(
    const std::vector<Tensor>& arrays,
    const std::vector<ContractOp>& ops,
    int64_t ntemps,
    int64_t root_id,
    Workspace& workspace) {
  c10::InferenceMode inference_guard(true);
  std::vector<Tensor> temps(static_cast<size_t>(ntemps));
  if (arrays.size() > temps.size()) {
    throw std::runtime_error("Temporary tensor table is smaller than input arrays.");
  }
  for (size_t i = 0; i < arrays.size(); ++i) {
    temps[i] = arrays[i];
  }

  std::unordered_map<int64_t, int> producer;
  producer.reserve(ops.size());
  for (size_t i = 0; i < ops.size(); ++i) {
    producer[ops[i].out_id] = static_cast<int>(i);
  }

  std::vector<std::vector<int>> dependents(ops.size());
  std::vector<int> pending_inputs(ops.size(), 0);
  std::vector<int> remaining_consumers(static_cast<size_t>(ntemps), 0);
  for (size_t i = 0; i < ops.size(); ++i) {
    for (const auto tensor_id : {ops[i].left_id, ops[i].right_id}) {
      if (tensor_id >= 0 && tensor_id < ntemps) {
        remaining_consumers[static_cast<size_t>(tensor_id)] += 1;
      }
      auto it = producer.find(tensor_id);
      if (it != producer.end()) {
        pending_inputs[i] += 1;
        dependents[it->second].push_back(static_cast<int>(i));
      }
    }
  }

  std::vector<int> ready;
  ready.reserve(ops.size());
  for (size_t i = 0; i < ops.size(); ++i) {
    if (pending_inputs[i] == 0) {
      ready.push_back(static_cast<int>(i));
    }
  }

  const bool batch_enabled = env_enabled("QIBOTN_CPP_BATCH_SMALL_GEMM");
  const int min_batch = env_int("QIBOTN_CPP_BATCH_MIN_SIZE", 4);
  const int64_t max_flops = env_i64("QIBOTN_CPP_BATCH_MAX_FLOPS", 1000000);
  std::vector<char> done(ops.size(), 0);
  size_t completed = 0;

  auto mark_consumed_and_release = [&](const ContractOp& op, const ExecutedOp& executed) {
    if (executed.left_workspace_owned) {
      workspace.release(executed.left);
    }
    if (executed.right_workspace_owned) {
      workspace.release(executed.right);
    }
    for (const auto tensor_id : {op.left_id, op.right_id}) {
      if (tensor_id < static_cast<int64_t>(arrays.size()) || tensor_id < 0) {
        continue;
      }
      auto& count = remaining_consumers[static_cast<size_t>(tensor_id)];
      count -= 1;
      if (count == 0) {
        workspace.release(temps.at(static_cast<size_t>(tensor_id)));
        temps.at(static_cast<size_t>(tensor_id)) = Tensor();
      }
    }
  };

  while (completed < ops.size()) {
    if (ready.empty()) {
      throw std::runtime_error("DAG contractor has no ready ops before completion.");
    }

    std::vector<int> to_execute;
#ifdef QIBOTN_USE_MKL
    if (batch_enabled) {
      std::unordered_map<GemmBatchKey, std::vector<int>, GemmBatchKeyHash> groups;
      for (const int op_index : ready) {
        if (done[op_index]) {
          continue;
        }
        GemmBatchKey key;
        if (!static_small_2d_gemm_key(ops[op_index], max_flops, key)) {
          continue;
        }
        Tensor left = temps.at(static_cast<size_t>(ops[op_index].left_id));
        if (!left.defined()) {
          continue;
        }
        key.dtype = left.scalar_type();
        groups[key].push_back(op_index);
      }
      for (const auto& item : groups) {
        if (static_cast<int>(item.second.size()) >= min_batch) {
          to_execute = item.second;
          break;
        }
      }
    }
#endif

    if (to_execute.empty()) {
      to_execute.push_back(ready.back());
    }

    std::vector<ExecutedOp> executed_batch;
    bool batch_ok = false;
#ifdef QIBOTN_USE_MKL
    if (to_execute.size() > 1) {
      batch_ok = execute_mkl_batch(to_execute, ops, temps, workspace, executed_batch);
    }
#endif
    if (!batch_ok) {
      if (to_execute.size() > 1) {
        to_execute.resize(1);
      }
      executed_batch.clear();
      executed_batch.push_back(execute_single_op(ops[to_execute[0]], temps, workspace));
    }

    for (size_t j = 0; j < to_execute.size(); ++j) {
      const int op_index = to_execute[j];
      if (done[op_index]) {
        continue;
      }
      done[op_index] = 1;
      completed += 1;
      mark_consumed_and_release(ops[op_index], executed_batch[j]);
      for (const int dep : dependents[op_index]) {
        pending_inputs[dep] -= 1;
        if (pending_inputs[dep] == 0) {
          ready.push_back(dep);
        }
      }
    }

    ready.erase(
        std::remove_if(
            ready.begin(),
            ready.end(),
            [&](int op_index) { return done[op_index] != 0; }),
        ready.end());
  }

  Tensor result = temps.at(static_cast<size_t>(root_id));
  if (!result.defined()) {
    throw std::runtime_error("DAG contractor did not produce a result tensor.");
  }
  return result;
}

Tensor contract_impl(
    const std::vector<Tensor>& arrays,
    const pybind11::list& plan,
    int64_t ntemps,
    int64_t root_id) {
  auto ops = parse_plan(plan);
  Workspace workspace;
  return run_contract(arrays, ops, ntemps, root_id, workspace);
}

class Contractor {
 public:
  Contractor(const pybind11::list& plan, int64_t ntemps, int64_t root_id)
      : ops_(parse_plan(plan)), ntemps_(ntemps), root_id_(root_id) {
    compute_batch_summary();
  }

  Tensor contract(const std::vector<Tensor>& arrays) {
    if (env_enabled("QIBOTN_CPP_DAG_SCHEDULER")) {
      return run_contract_dag(arrays, ops_, ntemps_, root_id_, workspace_);
    }
    return run_contract(arrays, ops_, ntemps_, root_id_, workspace_);
  }

  int64_t nsteps() const {
    return static_cast<int64_t>(ops_.size());
  }

  pybind11::dict batch_summary() const {
    pybind11::dict out;
    out["groups"] = batch_groups_;
    out["covered_ops"] = batch_covered_ops_;
    out["calls_saved"] = batch_calls_saved_;
    return out;
  }

 private:
  void compute_batch_summary() {
    const int min_batch = env_int("QIBOTN_CPP_BATCH_MIN_SIZE", 4);
    const int64_t max_flops = env_i64("QIBOTN_CPP_BATCH_MAX_FLOPS", 1000000);
    std::unordered_map<int64_t, int> producer;
    for (size_t i = 0; i < ops_.size(); ++i) {
      producer[ops_[i].out_id] = static_cast<int>(i);
    }
    std::vector<int> pending_inputs(ops_.size(), 0);
    std::vector<std::vector<int>> dependents(ops_.size());
    for (size_t i = 0; i < ops_.size(); ++i) {
      for (const auto tensor_id : {ops_[i].left_id, ops_[i].right_id}) {
        auto it = producer.find(tensor_id);
        if (it != producer.end()) {
          pending_inputs[i] += 1;
          dependents[it->second].push_back(static_cast<int>(i));
        }
      }
    }
    std::vector<int> ready;
    for (size_t i = 0; i < ops_.size(); ++i) {
      if (pending_inputs[i] == 0) {
        ready.push_back(static_cast<int>(i));
      }
    }
    std::vector<char> done(ops_.size(), 0);
    size_t completed = 0;
    while (completed < ops_.size() && !ready.empty()) {
      std::unordered_map<GemmBatchKey, int, GemmBatchKeyHash> counts;
      int best_count = 0;
      for (const int op_index : ready) {
        GemmBatchKey key;
        if (!static_small_2d_gemm_key(ops_[op_index], max_flops, key)) {
          continue;
        }
        key.dtype = c10::ScalarType::ComplexFloat;
        int count = ++counts[key];
        best_count = std::max(best_count, count);
      }
      if (best_count >= min_batch) {
        batch_groups_ += 1;
        batch_covered_ops_ += best_count;
        batch_calls_saved_ += best_count - 1;
      }
      // Simulate one scheduler step: batch best group if present, otherwise one op.
      std::vector<int> executed;
      if (best_count >= min_batch) {
        GemmBatchKey best_key;
        for (const auto& item : counts) {
          if (item.second == best_count) {
            best_key = item.first;
            break;
          }
        }
        for (const int op_index : ready) {
          GemmBatchKey key;
          if (static_small_2d_gemm_key(ops_[op_index], max_flops, key)) {
            key.dtype = c10::ScalarType::ComplexFloat;
            if (key == best_key) {
              executed.push_back(op_index);
            }
          }
        }
      } else {
        executed.push_back(ready.back());
      }
      for (const int op_index : executed) {
        if (done[op_index]) {
          continue;
        }
        done[op_index] = 1;
        completed += 1;
        for (const int dep : dependents[op_index]) {
          pending_inputs[dep] -= 1;
          if (pending_inputs[dep] == 0) {
            ready.push_back(dep);
          }
        }
      }
      ready.erase(
          std::remove_if(
              ready.begin(),
              ready.end(),
              [&](int op_index) { return done[op_index] != 0; }),
          ready.end());
    }
  }

  std::vector<ContractOp> ops_;
  int64_t ntemps_;
  int64_t root_id_;
  Workspace workspace_;
  int64_t batch_groups_ = 0;
  int64_t batch_covered_ops_ = 0;
  int64_t batch_calls_saved_ = 0;
};

}  // namespace

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def(
      "contract",
      &contract_impl,
      "Run a compiled cotengra contraction plan with torch CPU tensors.");
  pybind11::class_<Contractor>(m, "Contractor")
      .def(pybind11::init<const pybind11::list&, int64_t, int64_t>())
      .def("contract", &Contractor::contract)
      .def("batch_summary", &Contractor::batch_summary)
      .def_property_readonly("nsteps", &Contractor::nsteps);
}
