/*******************************************************************************
 *
 * Copyright (c) 2019 Konduit K.K.
 *
* ******************************************************************************
 *
 *
 * This program and the accompanying materials are made available under the
 * terms of the Apache License, Version 2.0 which is available at
 * https://www.apache.org/licenses/LICENSE-2.0.
 *
 *  See the NOTICE file distributed with this work for additional
 *  information regarding copyright ownership.
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations
 * under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 ******************************************************************************/
//
// @author AbdelRauf
//
#ifndef LIBND4J_LOOPCOORDSHELPER_H
#define LIBND4J_LOOPCOORDSHELPER_H
#include <helpers/shape.h>
#include <system/op_boilerplate.h>

#include <cstddef>
#include <type_traits>
#include <utility>
#include <vector>
namespace sd {

#if defined(__GNUC__)
#define likely(x) __builtin_expect((x), 1)
#define unlikely(x) __builtin_expect((x), 0)
#else
#define likely(x) (x)
#define unlikely(x) (x)
#endif

struct zip_size_t {
  LongType first;
  LongType second;
};

template <size_t Index>
struct CoordsState : CoordsState<Index - 1> {
  LongType coord;
  LongType last_num;
  LongType stride;
  LongType adjust;
  CoordsState() : CoordsState<Index - 1>() {}
};

template <>
struct CoordsState<0> {
  LongType coord;
  LongType last_num;
  LongType stride;
  LongType adjust;
  CoordsState() {}
};

template <size_t Index>
struct ZipCoordsState : ZipCoordsState<Index - 1> {
  LongType coord;
  LongType last_num;
  LongType stride1;
  LongType stride2;
  LongType adjust1;
  LongType adjust2;
  ZipCoordsState() : ZipCoordsState<Index - 1>() {}
};

template <>
struct ZipCoordsState<0> {
  LongType coord;
  LongType last_num;
  LongType stride1;
  LongType stride2;
  LongType adjust1;
  LongType adjust2;
  ZipCoordsState() {}
};

#define COORDS(x, index) ((x).::sd::CoordsState<(index)>::coord)
#define STRIDE(x, index) ((x).::sd::CoordsState<(index)>::stride)
#define LAST_NUM(x, index) ((x).::sd::CoordsState<(index)>::last_num)
#define OF_ADJUST(x, index) ((x).::sd::CoordsState<(index)>::adjust)
#define ZIP_LAST_NUM(x, index) ((x).::sd::ZipCoordsState<(index)>::last_num)
#define ZIP_COORDS(x, index) ((x).::sd::ZipCoordsState<(index)>::coord)
#define ZIP_STRIDE1(x, index) ((x).::sd::ZipCoordsState<(index)>::stride1)
#define ZIP_STRIDE2(x, index) ((x).::sd::ZipCoordsState<(index)>::stride2)
#define ZIP_OF_ADJUST1(x, index) ((x).::sd::ZipCoordsState<(index)>::adjust1)
#define ZIP_OF_ADJUST2(x, index) ((x).::sd::ZipCoordsState<(index)>::adjust2)



SD_INLINE SD_HOST_DEVICE size_t offset_from_coords(const LongType* strides, const LongType* coords,
                                                   const LongType& rank) {
  size_t offset = 0;
  size_t rank_4 = rank & -4;
  for (size_t i = 0; i < rank_4; i += 4) {
    offset = offset + coords[i] * strides[i] + coords[i + 1] * strides[i + 1] + coords[i + 2] * strides[i + 2] +
             coords[i + 3] * strides[i + 3];
  }
  for (int i = rank_4; i < rank; i++) {
    offset += coords[i] * strides[i];
  }
  return offset;
}

SD_INLINE SD_HOST_DEVICE zip_size_t offset_from_coords(const LongType* x_strides, const LongType* z_strides,
                                                       const LongType* coords, const LongType& rank) {
  zip_size_t offset = {0, 0};
  size_t rank_4 = rank & -4;
  for (size_t i = 0; i < rank_4; i += 4) {
    offset.first = offset.first + coords[i] * x_strides[i] + coords[i + 1] * x_strides[i + 1] +
                   coords[i + 2] * x_strides[i + 2] + coords[i + 3] * x_strides[i + 3];
    offset.second = offset.second + coords[i] * z_strides[i] + coords[i + 1] * z_strides[i + 1] +
                    coords[i + 2] * z_strides[i + 2] + coords[i + 3] * z_strides[i + 3];
  }
  for (int i = rank_4; i < rank; i++) {
    offset.first += coords[i] * x_strides[i];
    offset.second += coords[i] * z_strides[i];
  }
  return offset;
}

template <size_t Rank, size_t Index, bool Last_Index_Faster = true>
constexpr size_t StridesOrderInd() {
  return Last_Index_Faster ? Rank - Index - 1 : Index;
}

template <size_t Rank, size_t Index, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE typename std::enable_if<(Rank - 1 == Index), size_t>::type coord_inc_n(
    CoordsState<Rank - 1>& cbs, size_t last_offset) {
  constexpr size_t Ind = StridesOrderInd<Rank, Index, Last_Index_Faster>();

  if (likely(COORDS(cbs, Ind) < LAST_NUM(cbs, Ind))) {
    last_offset += cbs.CoordsState<Ind>::stride;
    COORDS(cbs, Ind) = COORDS(cbs, Ind) + 1;
    return last_offset;
  }
  // overflow case should not happen
  COORDS(cbs, Ind) = 0;
  // last_offset = 0;// last_offset + strides[Ind] - adjust_stride;
  return 0;
}

template <size_t Rank, size_t Index, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE typename std::enable_if<(Rank - 1 != Index), size_t>::type coord_inc_n(
    CoordsState<Rank - 1>& cbs, size_t last_offset) {
  constexpr size_t Ind = StridesOrderInd<Rank, Index, Last_Index_Faster>();

  if (likely(COORDS(cbs, Ind) < LAST_NUM(cbs, Ind))) {
    last_offset = last_offset + cbs.CoordsState<Ind>::stride;
    COORDS(cbs, Ind) = COORDS(cbs, Ind) + 1;
  } else {
    // lets adjust offset
    last_offset -= OF_ADJUST(cbs, Ind);
    COORDS(cbs, Ind) = 0;
    last_offset = coord_inc_n<Rank, Index + 1, Last_Index_Faster>(cbs, last_offset);
  }

  return last_offset;
}

template <size_t Rank, size_t Index = 0, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE size_t inc_coords(CoordsState<Rank - 1>& cbs, size_t last_offset) {
  return coord_inc_n<Rank, Index, Last_Index_Faster>(cbs, /* 1,*/ last_offset /*, 0*/);
}

template <size_t Rank, size_t rankIndex = 0, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE size_t inc_coords_ews(CoordsState<Rank - 1>& cbs, size_t last_offset, size_t ews) {
  if (ews == 1) {
    constexpr size_t Ind = StridesOrderInd<Rank, rankIndex, Last_Index_Faster>();
    return last_offset + STRIDE(cbs, Ind);
  }
  return coord_inc_n<Rank, rankIndex, Last_Index_Faster>(cbs, /* 1,*/ last_offset /*, 0*/);
}

template <size_t Rank, size_t rankIndex, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE typename std::enable_if<(Rank - 1 == rankIndex), zip_size_t>::type coord_inc_n(
    ZipCoordsState<Rank - 1>& cbs, zip_size_t last_offset) {
  constexpr size_t Ind = StridesOrderInd<Rank, rankIndex, Last_Index_Faster>();

  if (likely(ZIP_COORDS(cbs, Ind) < ZIP_LAST_NUM(cbs, Ind))) {
    last_offset.first += ZIP_STRIDE1(cbs, Ind);
    last_offset.second += ZIP_STRIDE2(cbs, Ind);
    ZIP_COORDS(cbs, Ind) = ZIP_COORDS(cbs, Ind) + 1;
    return last_offset;
  }
  // overflow case should not happen
  ZIP_COORDS(cbs, Ind) = 0;
  // last_offset = 0;// last_offset + strides[Ind] - adjust_stride;
  return {0, 0};
}

template <size_t Rank, size_t rankIndex, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE typename std::enable_if<(Rank - 1 != rankIndex), zip_size_t>::type coord_inc_n(
    ZipCoordsState<Rank - 1>& cbs, zip_size_t last_offset) {
  constexpr size_t Ind = StridesOrderInd<Rank, rankIndex, Last_Index_Faster>();

  if (likely(ZIP_COORDS(cbs, Ind) < ZIP_LAST_NUM(cbs, Ind))) {
    last_offset.first += ZIP_STRIDE1(cbs, Ind);
    last_offset.second += ZIP_STRIDE2(cbs, Ind);
    ZIP_COORDS(cbs, Ind) = ZIP_COORDS(cbs, Ind) + 1;
  } else {
    // lets adjust offset
    last_offset.first -= ZIP_OF_ADJUST1(cbs, Ind);
    last_offset.second -= ZIP_OF_ADJUST2(cbs, Ind);
    ZIP_COORDS(cbs, Ind) = 0;
    last_offset = coord_inc_n<Rank, rankIndex + 1, Last_Index_Faster>(cbs, last_offset);
  }

  return last_offset;
}

template <size_t Rank, size_t rankIndex = 0, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE zip_size_t inc_coords(ZipCoordsState<Rank - 1>& cbs, zip_size_t last_offset) {
  return coord_inc_n<Rank, rankIndex, Last_Index_Faster>(cbs, last_offset);
}

template <size_t Rank, size_t rankIndex = 0, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE typename std::enable_if<(Rank - 1 == rankIndex), size_t>::type init_coords(
    CoordsState<Rank - 1>& cbs, const LongType index, const LongType* bases, const LongType* strides,
    size_t offset = 0) {
  constexpr size_t Ind = StridesOrderInd<Rank, rankIndex, Last_Index_Faster>();
  COORDS(cbs, Ind) = index % bases[Ind];
  LAST_NUM(cbs, Ind) = bases[Ind] - 1;
  STRIDE(cbs, Ind) = strides[Ind];
  OF_ADJUST(cbs, Ind) = bases[Ind] * strides[Ind] - strides[Ind];
  offset += COORDS(cbs, Ind) * strides[Ind];
  return offset;
}

template <size_t Rank, size_t rankIndex = 0, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE typename std::enable_if<(Rank - 1 != rankIndex), size_t>::type init_coords(
    CoordsState<Rank - 1>& cbs, const LongType index, const LongType* bases, const LongType* strides,
    size_t offset = 0) {
  constexpr size_t Ind = StridesOrderInd<Rank, rankIndex, Last_Index_Faster>();
  COORDS(cbs, Ind) = index % bases[Ind];
  LAST_NUM(cbs, Ind) = bases[Ind] - 1;
  STRIDE(cbs, Ind) = strides[Ind];
  OF_ADJUST(cbs, Ind) = bases[Ind] * strides[Ind] - strides[Ind];
  offset += COORDS(cbs, Ind) * strides[Ind];
  return init_coords<Rank, rankIndex + 1, Last_Index_Faster>(cbs, index / bases[Ind], bases, strides, offset);
}

template <size_t Rank, size_t rankIndex = 0, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE typename std::enable_if<(Rank - 1 == rankIndex), bool>::type eq_coords(
    CoordsState<Rank - 1>& cbs, const LongType* coords) {
  return COORDS(cbs, rankIndex) == coords[rankIndex];
}

template <size_t Rank, size_t rankIndex = 0>
SD_INLINE SD_HOST_DEVICE typename std::enable_if<(Rank - 1 != rankIndex), bool>::type eq_coords(
    CoordsState<Rank - 1>& cbs, const LongType* coords) {
  return COORDS(cbs, rankIndex) == coords[rankIndex] && eq_coords<Rank, rankIndex + 1>(cbs, coords);
}

template <size_t Rank, size_t rankIndex = 0, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE typename std::enable_if<(Rank - 1 == rankIndex), bool>::type eq_zip_coords(
    ZipCoordsState<Rank - 1>& cbs, const LongType* coords) {
  return ZIP_COORDS(cbs, rankIndex) == coords[rankIndex];
}

template <size_t Rank, size_t rankIndex = 0>
SD_INLINE SD_HOST_DEVICE typename std::enable_if<(Rank - 1 != rankIndex), bool>::type eq_zip_coords(
    ZipCoordsState<Rank - 1>& cbs, const LongType* coords) {
  return ZIP_COORDS(cbs, rankIndex) == coords[rankIndex] && eq_zip_coords<Rank, rankIndex + 1>(cbs, coords);
}

template <size_t Rank, size_t rankIndex = 0, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE typename std::enable_if<(Rank - 1 == rankIndex), zip_size_t>::type init_coords(
    ZipCoordsState<Rank - 1>& cbs, const LongType index, const LongType* bases, const LongType* x_strides,
    const LongType* z_strides, zip_size_t offset = {}) {
  constexpr size_t Ind = StridesOrderInd<Rank, rankIndex, Last_Index_Faster>();
  ZIP_COORDS(cbs, Ind) = index % bases[Ind];
  ZIP_LAST_NUM(cbs, Ind) = bases[Ind] - 1;
  ZIP_STRIDE1(cbs, Ind) = x_strides[Ind];
  ZIP_STRIDE2(cbs, Ind) = z_strides[Ind];
  ZIP_OF_ADJUST1(cbs, Ind) = ZIP_LAST_NUM(cbs, Ind) * ZIP_STRIDE1(cbs, Ind);
  ZIP_OF_ADJUST2(cbs, Ind) = ZIP_LAST_NUM(cbs, Ind) * ZIP_STRIDE2(cbs, Ind);
  offset.first += ZIP_COORDS(cbs, Ind) * ZIP_STRIDE1(cbs, Ind);
  offset.second += ZIP_COORDS(cbs, Ind) * ZIP_STRIDE2(cbs, Ind);
  return offset;
}

template <size_t Rank, size_t rankIndex = 0, bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE typename std::enable_if<(Rank - 1 != rankIndex), zip_size_t>::type init_coords(
    ZipCoordsState<Rank - 1>& cbs, const LongType index, const LongType* bases, const LongType* x_strides,
    const LongType* z_strides, zip_size_t offset = {}) {
  constexpr size_t Ind = StridesOrderInd<Rank, rankIndex, Last_Index_Faster>();
  ZIP_COORDS(cbs, Ind) = index % bases[Ind];
  ZIP_LAST_NUM(cbs, Ind) = bases[Ind] - 1;
  ZIP_STRIDE1(cbs, Ind) = x_strides[Ind];
  ZIP_STRIDE2(cbs, Ind) = z_strides[Ind];
  ZIP_OF_ADJUST1(cbs, Ind) = ZIP_LAST_NUM(cbs, Ind) * ZIP_STRIDE1(cbs, Ind);
  ZIP_OF_ADJUST2(cbs, Ind) = ZIP_LAST_NUM(cbs, Ind) * ZIP_STRIDE2(cbs, Ind);
  offset.first += ZIP_COORDS(cbs, Ind) * ZIP_STRIDE1(cbs, Ind);
  offset.second += ZIP_COORDS(cbs, Ind) * ZIP_STRIDE2(cbs, Ind);
  return init_coords<Rank, rankIndex + 1, Last_Index_Faster>(cbs, index / bases[Ind], bases, x_strides, z_strides,
                                                             offset);
}

// inc coords for non constant Ranks
template <bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE size_t inc_coords(const LongType* bases, const LongType* strides, LongType* coords,
                                           size_t last_offset, const size_t rank, const size_t skip = 0) {
  LongType val;
  for (int i = rank - skip - 1; i >= 0; i--) {
    val = coords[i] + 1;
    if (likely(val < bases[i])) {
      coords[i] = val;
      last_offset += strides[i];
      break;
    } else {
      last_offset -= coords[i] * strides[i];
      coords[i] = 0;
    }
  }
  return last_offset;
}

template <>
SD_INLINE SD_HOST_DEVICE size_t inc_coords<false>(const LongType* bases, const LongType* strides, LongType* coords, size_t last_offset, const size_t rank,
                                                  const size_t skip) {
  LongType val;
  for (size_t i = skip; i < rank; i++) {
    val = coords[i] + 1;
    if (likely(val < bases[i])) {
      coords[i] = val;
      last_offset += strides[i];
      break;
    } else {
      last_offset -= coords[i] * strides[i];
      coords[i] = 0;
    }
  }
  return last_offset;
}

template <bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE zip_size_t inc_coords(const LongType* bases, const LongType* x_strides,
                                               const LongType* z_strides, LongType* coords,
                                               zip_size_t last_offset, const size_t rank, const size_t skip = 0) {
  LongType val = 0;
  for (int i = rank - skip - 1; i >= 0; i--) {
    val = coords[i] + 1;
    if (likely(val < bases[i])) {
      coords[i] = val;
      last_offset.first += x_strides[i];
      last_offset.second += z_strides[i];
      break;
    } else {
      last_offset.first -= coords[i] * x_strides[i];
      last_offset.second -= coords[i] * z_strides[i];
      coords[i] = 0;
    }
  }
  return last_offset;
}

template <>
SD_INLINE SD_HOST_DEVICE zip_size_t inc_coords<false>(const LongType* bases, const LongType* x_strides,
                                                      const LongType* z_strides, LongType* coords,
                                                      zip_size_t last_offset, const size_t rank, const size_t skip) {
  LongType val = 0;
  for (size_t i = skip; i < rank; i++) {
    val = coords[i] + 1;
    if (likely(val < bases[i])) {
      coords[i] = val;

      last_offset.first += x_strides[i];
      last_offset.second += z_strides[i];
      break;
    } else {
      last_offset.first -= coords[i] * x_strides[i];
      last_offset.second -= coords[i] * z_strides[i];
      coords[i] = 0;
    }
  }
  return last_offset;
}

struct triple_size_t {
  size_t first;
  size_t second;
  size_t third;
};

template <bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE triple_size_t inc_coords(const LongType* bases, const LongType* x_strides,
                                                  const LongType* y_strides, const LongType* z_strides,
                                                  LongType* coords, triple_size_t last_offset, const size_t rank,
                                                  const size_t skip = 0) {
  LongType val = 0;
  for (int i = rank - skip - 1; i >= 0; i--) {
    val = coords[i] + 1;
    if (likely(val < bases[i])) {
      coords[i] = val;
      last_offset.first += x_strides[i];
      last_offset.second += y_strides[i];
      last_offset.third += z_strides[i];
      break;
    } else {
      last_offset.first -= coords[i] * x_strides[i];
      last_offset.second -= coords[i] * y_strides[i];
      last_offset.third -= coords[i] * z_strides[i];
      coords[i] = 0;
    }
  }
  return last_offset;
}

template <>
SD_INLINE SD_HOST_DEVICE triple_size_t inc_coords<false>(const LongType* bases, const LongType* x_strides,
                                                         const LongType* y_strides, const LongType* z_strides,
                                                         LongType* coords, triple_size_t last_offset,
                                                         const size_t rank, const size_t skip) {
  LongType val = 0;
  for (size_t i = skip; i < rank; i++) {
    val = coords[i] + 1;
    if (likely(val < bases[i])) {
      coords[i] = val;

      last_offset.first += x_strides[i];
      last_offset.second += y_strides[i];
      last_offset.third += z_strides[i];
      break;
    } else {
      last_offset.first -= coords[i] * x_strides[i];
      last_offset.second -= coords[i] * y_strides[i];
      last_offset.third -= coords[i] * z_strides[i];
      coords[i] = 0;
    }
  }
  return last_offset;
}

SD_INLINE SD_HOST_DEVICE triple_size_t offset_from_coords(const LongType* x_strides, const LongType* y_strides,
                                                          const LongType* z_strides, const LongType* coords,
                                                          const LongType& rank) {
  triple_size_t offset = {0, 0, 0};
  size_t rank_4 = rank & -4;
  for (size_t i = 0; i < rank_4; i += 4) {
    offset.first = offset.first + coords[i] * x_strides[i] + coords[i + 1] * x_strides[i + 1] +
                   coords[i + 2] * x_strides[i + 2] + coords[i + 3] * x_strides[i + 3];
    offset.second = offset.second + coords[i] * y_strides[i] + coords[i + 1] * y_strides[i + 1] +
                    coords[i + 2] * y_strides[i + 2] + coords[i + 3] * y_strides[i + 3];
    offset.third = offset.third + coords[i] * z_strides[i] + coords[i + 1] * z_strides[i + 1] +
                   coords[i + 2] * z_strides[i + 2] + coords[i + 3] * z_strides[i + 3];
  }
  for (int i = rank_4; i < rank; i++) {
    offset.first += coords[i] * x_strides[i];
    offset.second += coords[i] * y_strides[i];
    offset.third += coords[i] * z_strides[i];
  }
  return offset;
}

template <bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE LongType getLength(const LongType* bases, int rank, int skip = 0) {
  if (skip < 0 || skip >= rank) skip = 0;
  LongType total = 1;
  for (int i = 0; i < rank - skip; i++) {
    total *= bases[i];
  }
  return total;
}

template <>
SD_INLINE SD_HOST_DEVICE LongType getLength<false>(const LongType* bases, int rank, int skip) {
  if (skip < 0 || skip >= rank) skip = 0;
  LongType total = 1;
  for (int i = skip; i < rank; i++) {
    total *= bases[i];
  }

  return total;
}

template <bool Last_Index_Faster = true>
SD_INLINE SD_HOST_DEVICE LongType getLength(const LongType* bases, int rank, int skip, LongType& outSkippedLength) {
  if (skip < 0 || skip >= rank) skip = 0;
  LongType total = 1;
  for (int i = 0; i < rank - skip; i++) {
    total *= bases[i];
  }
  if (skip > 0) {
    outSkippedLength = 1;
    for (int i = rank - skip; i < rank; i++) {
      outSkippedLength *= bases[i];
    }
  } else {
    outSkippedLength = 0;
  }
  return total;
}

template <>
SD_INLINE SD_HOST_DEVICE LongType getLength<false>(const LongType* bases, int rank, int skip,
                                                   LongType& outSkippedLength) {
  if (skip < 0 || skip >= rank) skip = 0;
  if (skip > 0) {
    outSkippedLength = 1;
    for (int i = 0; i < skip; i++) {
      outSkippedLength *= bases[i];
    }
  } else {
    outSkippedLength = 0;
  }
  LongType total = 1;
  for (int i = skip; i < rank; i++) {
    total *= bases[i];
  }

  return total;
}

/*
for ODR rule it willbe declared as inline
rePartition for reductions and et cet
Indices mentioned in the dimension list will be moved to the tail
This way it will be splitted into two parts
the first part will contain output part,the second tail part will be used for reductions and other purposes
if squash is True then  it will attempt to minimize the output ( for both orders) and the tail
*/

SD_INLINE SD_HOST_DEVICE void rePartition(char order, const std::vector<LongType> dimensions, const size_t rank,
                                          const LongType* bases, const LongType* strides,
                                          LongType (&new_bases)[SD_MAX_RANK], LongType (&new_strides)[SD_MAX_RANK], LongType& first_begin,
                                          LongType& first_end, LongType& second_begin, LongType& second_end, bool first_squash = false,
                                          bool second_squash = true) {
  bool indices[SD_MAX_RANK] = {};
  int ind = 0;
  size_t second_rank;
  if (dimensions.size() == 0 || (dimensions.size() == 1 && dimensions.at(0) == DataTypeUtils::max<int>())) {
    first_end = 0;
    first_begin = 0;
    // treat it as the whole
    for (size_t i = 0; i < rank; i++) {
      new_bases[i] = bases[i];
      new_strides[i] = strides[i];
    }
    second_rank = rank;
    second_end = rank;
    second_begin = 0;

  } else {
    for (size_t index : dimensions) {
      if (index < rank) {
        indices[index] = true;
      }
    }

    // move output ones and
    for (size_t i = 0; i < rank; i++) {
      if (!indices[i]) {
        new_bases[ind] = bases[i];
        new_strides[ind] = strides[i];
        ind++;
      }
    }

    int first_rank = ind;

    first_end = ind;
    first_begin = 0;
    // squash output rank
    if (first_squash && first_rank > 1) {
      if (order == 'c') {
        int uniq_ind = first_end - 1;
        for (int i = first_end - 2; i >= first_begin; i--) {
          if (new_strides[i] == new_bases[uniq_ind] * new_strides[uniq_ind]) {
            new_bases[uniq_ind] = new_bases[i] * new_bases[uniq_ind];
            new_strides[uniq_ind] = new_strides[uniq_ind];
            --first_rank;
          } else {
            --uniq_ind;
            new_bases[uniq_ind] = new_bases[i];
            new_strides[uniq_ind] = new_strides[i];
          }
        }
        first_begin = first_end - first_rank;
      } else {
        // squash fortran
        int uniq_ind = 0;
        for (int i = 1; i < first_end; i++) {
          if (new_strides[i] == new_bases[uniq_ind] * new_strides[uniq_ind]) {
            new_bases[uniq_ind] = new_bases[i] * new_bases[uniq_ind];
            new_strides[uniq_ind] = new_strides[uniq_ind];
            --first_rank;
          } else {
            uniq_ind++;
            new_bases[uniq_ind] = new_bases[i];
            new_strides[uniq_ind] = new_strides[i];
          }
        }
        first_end = first_begin + first_rank;
      }
      ind = first_end;
    }

    // move process indices
    for (size_t i = 0; i < rank; i++) {
      if (indices[i]) {
        new_bases[ind] = bases[i];
        new_strides[ind] = strides[i];
        ind++;
      }
    }

    second_rank = ind - first_end;
    second_end = ind;
    second_begin = first_end;
  }

  if (second_squash && second_rank > 1) {
    if (order == 'c') {
      int uniq_ind = second_end - 1;
      for (int i = second_end - 2; i >= second_begin; i--) {
        if (new_strides[i] == new_bases[uniq_ind] * new_strides[uniq_ind]) {
          new_bases[uniq_ind] = new_bases[i] * new_bases[uniq_ind];
          new_strides[uniq_ind] = new_strides[uniq_ind];
          --second_rank;
        } else {
          --uniq_ind;
          new_bases[uniq_ind] = new_bases[i];
          new_strides[uniq_ind] = new_strides[i];
        }
      }
      second_begin = second_end - second_rank;
    } else {
      int uniq_ind = second_begin;
      for (int i = second_begin + 1; i < second_end; i++) {
        if (new_strides[i] == new_bases[uniq_ind] * new_strides[uniq_ind]) {
          new_bases[uniq_ind] = new_bases[i] * new_bases[uniq_ind];
          new_strides[uniq_ind] = new_strides[uniq_ind];
          --second_rank;
        } else {
          uniq_ind++;
          new_bases[uniq_ind] = new_bases[i];
          new_strides[uniq_ind] = new_strides[i];
        }
      }
      second_end = second_begin + second_rank;
    }
  }

  return;
}

// basic CRTP static polymorphism classes for offset increments

template <typename Derived>
struct CoordsBaseMovement {
  void init(const LongType* bases, const LongType* strides1, const LongType* strides2, int rank,
            int start = 0) {
    static_cast<Derived*>(this)->initImpl(bases, strides1, strides2, rank, start);
  }

  void increment(int skipRank = 0) { static_cast<Derived*>(this)->incrementImpl(skipRank); }

  LongType First() { return static_cast<Derived*>(this)->FirstImpl(); };
  LongType Second() { return static_cast<Derived*>(this)->SecondImpl(); };
};

struct ZipGenericCoordsRank1Stride1 : CoordsBaseMovement<ZipGenericCoordsRank1Stride1> {
  size_t offset1;
  size_t offset2;

  void initImpl(const LongType* bases, const LongType* strides1, const LongType* strides2, int rank,
                int start = 0) {
    offset1 = start;
    offset2 = start;
  }

  void incrementImpl(int skipRank = 0) {
    offset1 += 1;
    offset2 += 1;
  }

  LongType FirstImpl() { return offset1; };
  LongType SecondImpl() { return offset2; };
};

struct ZipGenericCoordsRank1BothStrideN : CoordsBaseMovement<ZipGenericCoordsRank1BothStrideN> {
  size_t stride1;
  size_t stride2;
  size_t offset1;
  size_t offset2;

  void initImpl(const LongType* bases, const LongType* strides1, const LongType* strides2, int rank,
                int start = 0) {
    stride1 = strides1[0];
    stride2 = strides2[0];
    offset1 = start * stride1;
    offset2 = start * stride2;
  }

  void incrementImpl(int skipRank = 0) {
    offset1 += stride1;
    offset2 += stride2;
  }

  LongType FirstImpl() { return offset1; };
  LongType SecondImpl() { return offset2; };
};

template <int ConstRank, bool LastIndexFaster = true>
struct ZipGenericCoordsConstMovementSecondStride1
    : CoordsBaseMovement<ZipGenericCoordsConstMovementSecondStride1<ConstRank, LastIndexFaster>> {
  CoordsState<ConstRank - 1> cst;
  LongType coords[SD_MAX_RANK];
  size_t offset1;
  size_t offset2;
  int _rank;

  void initImpl(const LongType* bases, const LongType* strides1, const LongType* strides2, int rank,
                int start = 0) {
    offset1 = sd::init_coords<ConstRank, 0, LastIndexFaster>(cst, start, bases, strides1);
    offset2 = start * 1;
  }

  void incrementImpl(int skipRank = 0) {
    offset1 = sd::inc_coords<ConstRank, 0, LastIndexFaster>(cst, offset1);
    offset2 += 1;
  }

  LongType FirstImpl() { return offset1; };
  LongType SecondImpl() { return offset2; };
};

template <int ConstRank, bool LastIndexFaster = true>
struct ZipGenericCoordsConstMovementSecondStrideN
    : CoordsBaseMovement<ZipGenericCoordsConstMovementSecondStrideN<ConstRank, LastIndexFaster>> {
  CoordsState<ConstRank - 1> cst;
  LongType _stride2;
  LongType coords[SD_MAX_RANK];
  size_t offset1;
  size_t offset2;
  int _rank;

  void initImpl(const LongType* bases, const LongType* strides1, const LongType* strides2, int rank,
                int start = 0) {
    _stride2 = strides2[0];
    offset1 = sd::init_coords<ConstRank, 0, LastIndexFaster>(cst, start, bases, strides1);
    offset2 = start * _stride2;
  }

  void incrementImpl(int skipRank = 0) {
    offset1 = sd::inc_coords<ConstRank, 0, LastIndexFaster>(cst, offset1);
    offset2 += _stride2;
  }

  LongType FirstImpl() { return offset1; };
  LongType SecondImpl() { return offset2; };
};
template <bool LastIndexFaster = true>
struct ZipGenericCoordsMovementSecondStrideN
    : CoordsBaseMovement<ZipGenericCoordsMovementSecondStrideN<LastIndexFaster>> {
  const LongType* _bases;
  const LongType* _strides1;
  LongType _stride2;
  LongType coords[SD_MAX_RANK];
  zip_size_t offset;
  int _rank;

  void initImpl(const LongType* bases, const LongType* strides1, const LongType* strides2, int rank,
                int start = 0) {
    _bases = bases;
    _strides1 = strides1;
    _stride2 = strides2[0];
    _rank = rank;
    if (start == 0) {
      for (int i = 0; i < SD_MAX_RANK; i++) {
        coords[i] = 0;
      }
      offset = {0, 0};

    } else {
      INDEX2COORDS(start, rank, bases, coords);
      COORDS2INDEX(rank, strides1, coords, offset.first);
      offset.second = start * _stride2;
    }
  }

  void incrementImpl(int skipRank = 0) {
    offset.first = inc_coords<LastIndexFaster>(_bases, _strides1, coords, offset.first, _rank, skipRank);
    offset.second += _stride2;
  }

  LongType FirstImpl() { return offset.first; };
  LongType SecondImpl() { return offset.second; };
};

template <bool LastIndexFaster = true>
struct ZipGenericCoordsMovement : CoordsBaseMovement<ZipGenericCoordsMovement<LastIndexFaster>> {
  const LongType* _bases;
  const LongType* _strides1;
  const LongType* _strides2;
  LongType coords[SD_MAX_RANK];
  zip_size_t offset;
  int _rank;

  void initImpl(const LongType* bases, const LongType* strides1, const LongType* strides2, int rank,
                int start = 0) {
    _bases = bases;
    _strides1 = strides1;
    _strides2 = strides2;
    _rank = rank;
    if (start == 0) {
      for (int i = 0; i < SD_MAX_RANK; i++) {
        coords[i] = 0;
      }
      offset = {0, 0};

    } else {
      INDEX2COORDS(start, rank, bases, coords);
      COORDS2INDEX(rank, strides1, coords, offset.first);
      COORDS2INDEX(rank, strides2, coords, offset.second);
    }
  }

  void incrementImpl(int skipRank = 0) {
    offset = inc_coords<LastIndexFaster>(_bases, _strides1, _strides2, coords, offset, _rank, skipRank);
  }

  LongType FirstImpl() { return offset.first; };
  LongType SecondImpl() { return offset.second; };
};
}  // namespace sd

#endif