/*!
 *  Copyright (c) 2019 by Contributors
 * \file array.cc
 * \brief DGL array utilities implementation
 */
#include <dgl/array.h>

namespace dgl {

// TODO(minjie): currently these operators are only on CPU.

IdArray NewIdArray(int64_t length) {
  return IdArray::Empty({length}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
}

BoolArray NewBoolArray(int64_t length) {
  return BoolArray::Empty({length}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
}

IdArray VecToIdArray(const std::vector<dgl_id_t>& vec) {
  IdArray ret = NewIdArray(vec.size());
  std::copy(vec.begin(), vec.end(), static_cast<dgl_id_t*>(ret->data));
  return ret;
}

IdArray Clone(IdArray arr) {
  IdArray ret = NewIdArray(arr->shape[0]);
  ret.CopyFrom(arr);
  return ret;
}

IdArray AsNumBits(IdArray arr, uint8_t bits) {
  if (arr->dtype.bits == bits) {
    return arr;
  } else {
    const int64_t len = arr->shape[0];
    IdArray ret = IdArray::Empty({len},
        DLDataType{kDLInt, bits, 1}, DLContext{kDLCPU, 0});
    if (arr->dtype.bits == 32 && bits == 64) {
      const int32_t* arr_data = static_cast<int32_t*>(arr->data);
      int64_t* ret_data = static_cast<int64_t*>(ret->data);
      for (int64_t i = 0; i < len; ++i) {
        ret_data[i] = arr_data[i];
      }
    } else if (arr->dtype.bits == 64 && bits == 32) {
      const int64_t* arr_data = static_cast<int64_t*>(arr->data);
      int32_t* ret_data = static_cast<int32_t*>(ret->data);
      for (int64_t i = 0; i < len; ++i) {
        ret_data[i] = arr_data[i];
      }
    } else {
      LOG(FATAL) << "Invalid type conversion.";
    }
    return ret;
  }
}

IdArray Add(IdArray lhs, IdArray rhs) {
  IdArray ret = NewIdArray(lhs->shape[0]);
  const dgl_id_t* lhs_data = static_cast<dgl_id_t*>(lhs->data);
  const dgl_id_t* rhs_data = static_cast<dgl_id_t*>(rhs->data);
  dgl_id_t* ret_data = static_cast<dgl_id_t*>(ret->data);
  for (int64_t i = 0; i < lhs->shape[0]; ++i) {
    ret_data[i] = lhs_data[i] + rhs_data[i];
  }
  return ret;
}

IdArray Sub(IdArray lhs, IdArray rhs) {
  IdArray ret = NewIdArray(lhs->shape[0]);
  const dgl_id_t* lhs_data = static_cast<dgl_id_t*>(lhs->data);
  const dgl_id_t* rhs_data = static_cast<dgl_id_t*>(rhs->data);
  dgl_id_t* ret_data = static_cast<dgl_id_t*>(ret->data);
  for (int64_t i = 0; i < lhs->shape[0]; ++i) {
    ret_data[i] = lhs_data[i] - rhs_data[i];
  }
  return ret;
}

IdArray Mul(IdArray lhs, IdArray rhs) {
  IdArray ret = NewIdArray(lhs->shape[0]);
  const dgl_id_t* lhs_data = static_cast<dgl_id_t*>(lhs->data);
  const dgl_id_t* rhs_data = static_cast<dgl_id_t*>(rhs->data);
  dgl_id_t* ret_data = static_cast<dgl_id_t*>(ret->data);
  for (int64_t i = 0; i < lhs->shape[0]; ++i) {
    ret_data[i] = lhs_data[i] * rhs_data[i];
  }
  return ret;
}

IdArray Div(IdArray lhs, IdArray rhs) {
  IdArray ret = NewIdArray(lhs->shape[0]);
  const dgl_id_t* lhs_data = static_cast<dgl_id_t*>(lhs->data);
  const dgl_id_t* rhs_data = static_cast<dgl_id_t*>(rhs->data);
  dgl_id_t* ret_data = static_cast<dgl_id_t*>(ret->data);
  for (int64_t i = 0; i < lhs->shape[0]; ++i) {
    ret_data[i] = lhs_data[i] / rhs_data[i];
  }
  return ret;
}

IdArray Add(IdArray lhs, dgl_id_t rhs) {
  IdArray ret = NewIdArray(lhs->shape[0]);
  const dgl_id_t* lhs_data = static_cast<dgl_id_t*>(lhs->data);
  dgl_id_t* ret_data = static_cast<dgl_id_t*>(ret->data);
  for (int64_t i = 0; i < lhs->shape[0]; ++i) {
    ret_data[i] = lhs_data[i] + rhs;
  }
  return ret;
}

IdArray Sub(IdArray lhs, dgl_id_t rhs) {
  IdArray ret = NewIdArray(lhs->shape[0]);
  const dgl_id_t* lhs_data = static_cast<dgl_id_t*>(lhs->data);
  dgl_id_t* ret_data = static_cast<dgl_id_t*>(ret->data);
  for (int64_t i = 0; i < lhs->shape[0]; ++i) {
    ret_data[i] = lhs_data[i] - rhs;
  }
  return ret;
}

IdArray Mul(IdArray lhs, dgl_id_t rhs) {
  IdArray ret = NewIdArray(lhs->shape[0]);
  const dgl_id_t* lhs_data = static_cast<dgl_id_t*>(lhs->data);
  dgl_id_t* ret_data = static_cast<dgl_id_t*>(ret->data);
  for (int64_t i = 0; i < lhs->shape[0]; ++i) {
    ret_data[i] = lhs_data[i] * rhs;
  }
  return ret;
}

IdArray Div(IdArray lhs, dgl_id_t rhs) {
  IdArray ret = NewIdArray(lhs->shape[0]);
  const dgl_id_t* lhs_data = static_cast<dgl_id_t*>(lhs->data);
  dgl_id_t* ret_data = static_cast<dgl_id_t*>(ret->data);
  for (int64_t i = 0; i < lhs->shape[0]; ++i) {
    ret_data[i] = lhs_data[i] / rhs;
  }
  return ret;
}

IdArray Add(dgl_id_t lhs, IdArray rhs) {
  return Add(rhs, lhs);
}

IdArray Sub(dgl_id_t lhs, IdArray rhs) {
  IdArray ret = NewIdArray(rhs->shape[0]);
  const dgl_id_t* rhs_data = static_cast<dgl_id_t*>(rhs->data);
  dgl_id_t* ret_data = static_cast<dgl_id_t*>(ret->data);
  for (int64_t i = 0; i < rhs->shape[0]; ++i) {
    ret_data[i] = lhs - rhs_data[i];
  }
  return ret;
}

IdArray Mul(dgl_id_t lhs, IdArray rhs) {
  return Mul(rhs, lhs);
}

IdArray Div(dgl_id_t lhs, IdArray rhs) {
  IdArray ret = NewIdArray(rhs->shape[0]);
  const dgl_id_t* rhs_data = static_cast<dgl_id_t*>(rhs->data);
  dgl_id_t* ret_data = static_cast<dgl_id_t*>(ret->data);
  for (int64_t i = 0; i < rhs->shape[0]; ++i) {
    ret_data[i] = lhs / rhs_data[i];
  }
  return ret;
}

IdArray HStack(IdArray arr1, IdArray arr2) {
  CHECK_EQ(arr1->shape[0], arr2->shape[0]);
  const int64_t L = arr1->shape[0];
  IdArray ret = NewIdArray(2 * L);
  const dgl_id_t* arr1_data = static_cast<dgl_id_t*>(arr1->data);
  const dgl_id_t* arr2_data = static_cast<dgl_id_t*>(arr2->data);
  dgl_id_t* ret_data = static_cast<dgl_id_t*>(ret->data);
  for (int64_t i = 0; i < L; ++i) {
    ret_data[i] = arr1_data[i];
    ret_data[i + L] = arr2_data[i];
  }
  return ret;
}

CSRMatrix SliceRows(const CSRMatrix& csr, int64_t start, int64_t end) {
  const dgl_id_t* indptr = static_cast<dgl_id_t*>(csr.indptr->data);
  const dgl_id_t* indices = static_cast<dgl_id_t*>(csr.indices->data);
  const dgl_id_t* data = static_cast<dgl_id_t*>(csr.data->data);
  const int64_t num_rows = end - start;
  const int64_t nnz = indptr[end] - indptr[start];
  CSRMatrix ret;
  ret.indptr = NewIdArray(num_rows + 1);
  ret.indices = NewIdArray(nnz);
  ret.data = NewIdArray(nnz);
  dgl_id_t* r_indptr = static_cast<dgl_id_t*>(ret.indptr->data);
  dgl_id_t* r_indices = static_cast<dgl_id_t*>(ret.indices->data);
  dgl_id_t* r_data = static_cast<dgl_id_t*>(ret.data->data);
  for (int64_t i = start; i < end + 1; ++i) {
    r_indptr[i - start] = indptr[i] - indptr[start];
  }
  std::copy(indices + indptr[start], indices + indptr[end], r_indices);
  std::copy(data + indptr[start], data + indptr[end], r_data);
  return ret;
}

}  // namespace dgl