/// @file   util/dilated_int.hpp
/// @brief  Z-order or Morton style indexing utilities for arbitrary dimension
/// @author will sheffler

// inspired by code from here: http://www.marcusbannerman.co.uk/dynamo
// see "Converting to and from Dilated Integers"(doi:10.1109/TC.2007.70814)
// maybe redo with modern code?

#pragma once

#include "hgeom/util/template_math.hpp"
// #include <iostream>
#ifdef DEBUG
#include <cstddef>
#include <stdexcept>
#endif

namespace hgeom {
namespace util {

namespace impl {

template <uint64_t D> struct MAXBITS {
  static const uint64_t VAL = (8 * sizeof(uint64_t)) / D;
};
template <uint64_t D> struct DILATION_ROUNDS {
  static const uint64_t VAL = CEIL_LOG<MAXBITS<D>::VAL, D - 1>::VAL;
};
template <> struct DILATION_ROUNDS<2> {
  static const uint64_t VAL = CEIL_LOG<MAXBITS<2>::VAL, 2>::VAL;
};
template <uint64_t D> struct UNDILATION_ROUNDS {
  static const uint64_t VAL = CEIL_LOG<MAXBITS<D>::VAL, D>::VAL;
};

template <uint64_t P, uint64_t Q> struct X_P_Q {
  template <uint64_t l, uint64_t MOCK> struct XPQ_IMPL {
    static const uint64_t VAL =
        XPQ_IMPL<l - 1, MOCK>::VAL + SAFE_LSHIFT<uint64_t, 1, l * Q>::VAL;
  };
  template <uint64_t MOCK> struct XPQ_IMPL<0, MOCK> {
    static const uint64_t VAL = 1;
  };
  static const uint64_t VAL = XPQ_IMPL<P - 1, 0>::VAL;
};

template <uint64_t I_PLUS_1, uint64_t D> struct UMULT_I_D {
  static const uint64_t VAL =
      X_P_Q<D, (D - 1) * POW<D, I_PLUS_1 - 1>::VAL>::VAL;
};
template <uint64_t ITER, uint64_t D> struct DMULT_I_D {
  static const uint64_t VAL =
      X_P_Q<D, POW<D - 1, DILATION_ROUNDS<D>::VAL - ITER + 1>::VAL>::VAL;
};
template <uint64_t N> struct N_BITS_LO {
  static const uint64_t VAL = (SAFE_LSHIFT<uint64_t, 1, N>::VAL - 1);
};
template <uint64_t D> struct MAX_DILATABLE {
  static const uint64_t VAL = N_BITS_LO<MAXBITS<D>::VAL>::VAL;
};

template <uint64_t ITER, uint64_t D> struct UMASK_I_D {
  static const uint64_t _di = POW<D, ITER>::VAL;
  static const uint64_t BITCOUNT =
      (_di < MAXBITS<D>::VAL) ? _di : MAXBITS<D>::VAL;
  template <uint64_t ITER_2, uint64_t MOCK> struct ZDI_IMPL {
    static const uint64_t VAL = ZDI_IMPL<ITER_2 - 1, MOCK>::VAL |
                                SAFE_RSHIFT<uint64_t, ZDI_IMPL<0, MOCK>::VAL,
                                            ((_di * D) * ITER_2)>::VAL;
  };
  template <uint64_t MOCK> struct ZDI_IMPL<0, MOCK> {
    static const uint64_t VAL = SAFE_LSHIFT < uint64_t,
                          N_BITS_LO<BITCOUNT>::VAL &
                              N_BITS_LO<MAXBITS<D>::VAL>::VAL,
                          D *(MAXBITS<D>::VAL - 1) + 1 - BITCOUNT > ::VAL;
  };
  static const uint64_t VAL =
      ZDI_IMPL<MAXBITS<D>::VAL / POW<D, ITER>::VAL, 0>::VAL;
};

// bit mask used after each round of the dilation algorithm.
template <uint64_t ITER, uint64_t D> struct DMASK_I_D {
  static const uint64_t BITCOUNT =
      POW<D - 1, DILATION_ROUNDS<D>::VAL - ITER>::VAL;
  static const uint64_t BITSEP =
      POW<D - 1, DILATION_ROUNDS<D>::VAL - ITER + 1>::VAL + BITCOUNT;
  template <uint64_t ITER_2, uint64_t MOCK> struct YDI_IMPL {
    static const uint64_t VAL =
        YDI_IMPL<ITER_2 - 1, MOCK>::VAL |
        SAFE_LSHIFT<uint64_t, N_BITS_LO<BITCOUNT>::VAL, BITSEP * ITER_2>::VAL;
  };
  template <uint64_t MOCK> struct YDI_IMPL<0, MOCK> {
    static const uint64_t VAL = N_BITS_LO<BITCOUNT>::VAL;
  };
  static const uint64_t VAL = YDI_IMPL<MAXBITS<D>::VAL / BITCOUNT, 0>::VAL;
};
template <uint64_t ITER> struct DMASK_I_D<ITER, 2> {
  static const uint64_t BITCOUNT = MAXBITS<2>::VAL / (1 << ITER);
  static const uint64_t BITSEP = 2 * BITCOUNT;
  template <uint64_t ITER_2, uint64_t MOCK> struct YDI_IMPL {
    static const uint64_t VAL =
        YDI_IMPL<ITER_2 - 1, MOCK>::VAL |
        SAFE_LSHIFT<uint64_t, N_BITS_LO<BITCOUNT>::VAL, BITSEP * ITER_2>::VAL;
  };
  template <uint64_t MOCK> struct YDI_IMPL<0, MOCK> {
    static const uint64_t VAL = N_BITS_LO<BITCOUNT>::VAL;
  };
  static const uint64_t VAL = YDI_IMPL<MAXBITS<2>::VAL / BITCOUNT, 0>::VAL;
};

template <uint64_t D> struct MAX_DILATED {
  static const uint64_t VAL = DMASK_I_D<DILATION_ROUNDS<D>::VAL, D>::VAL;
};

template <uint64_t D> struct UNDILATE {
  template <uint64_t ITER, uint64_t MOCK> struct UNDILATE_IMPL {
    static inline uint64_t eval(uint64_t val) {
      return (UNDILATE_IMPL<ITER - 1, MOCK>::eval(val) *
              UMULT_I_D<ITER, D>::VAL) &
             UMASK_I_D<ITER, D>::VAL;
    }
  };
  template <uint64_t MOCK> struct UNDILATE_IMPL<1, MOCK> {
    static inline uint64_t eval(uint64_t val) {
      return (val * UMULT_I_D<1, D>::VAL) & UMASK_I_D<1, D>::VAL;
    }
  };
  static inline uint64_t eval(uint64_t val) {
    static const uint64_t undilate_shift =
        (D * (MAXBITS<D>::VAL - 1) + 1 - MAXBITS<D>::VAL);
    return UNDILATE_IMPL<UNDILATION_ROUNDS<D>::VAL, 0>::eval(val) >>
           undilate_shift;
  };
};

// see "Converting to and from Dilated Integers"(doi:10.1109/TC.2007.70814)
template <uint64_t D> struct DILATE {
  template <uint64_t ITER, uint64_t MOCK> struct DILATE_IMPL {
    static inline uint64_t eval(uint64_t val) {
      return (DILATE_IMPL<ITER - 1, D>::eval(val) * DMULT_I_D<ITER, D>::VAL) &
             DMASK_I_D<ITER, D>::VAL;
    }
  };
  template <uint64_t MOCK> struct DILATE_IMPL<1, MOCK> {
    static inline uint64_t eval(uint64_t val) {
      return (val * DMULT_I_D<1, D>::VAL) & DMASK_I_D<1, D>::VAL;
    }
  };
  static inline uint64_t eval(uint64_t val) {
    return DILATE_IMPL<DILATION_ROUNDS<D>::VAL, 0>::eval(val);
  };
};

// for 2, must use Shift-Or rather than multiply, see paper
// see "Converting to and from Dilated Integers"(doi:10.1109/TC.2007.70814)
template <> struct DILATE<2> {
  template <uint64_t ITER> struct shiftval {
    static const uint64_t VAL =
        SAFE_LSHIFT<uint64_t, 1, DILATION_ROUNDS<2>::VAL - ITER>::VAL;
  };
  template <uint64_t ITER, uint64_t MOCK> struct DILATE_IMPL {
    static inline uint64_t eval(uint64_t val) {
      const uint64_t val2 = DILATE_IMPL<ITER - 1, 2>::eval(val);
      return (val2 | (val2 << shiftval<ITER>::VAL)) & DMASK_I_D<ITER, 2>::VAL;
    }
  };
  template <uint64_t MOCK> struct DILATE_IMPL<1, MOCK> {
    static inline uint64_t eval(uint64_t val) {
      return (val | (val << shiftval<1>::VAL)) & DMASK_I_D<1, 2>::VAL;
    }
  };
  static inline uint64_t eval(uint64_t val) {
    return DILATE_IMPL<DILATION_ROUNDS<2>::VAL, 0>::eval(val);
  };
};
} // namespace impl

template <uint64_t D> inline uint64_t dilate(uint64_t val) {
#ifdef DEBUG
  if (val > impl::MAX_DILATABLE<D>::VAL)
    throw std::out_of_range("too big to dilate");
#endif
  return impl::DILATE<D>::eval(val);
}
template <uint64_t D> inline uint64_t undilate(uint64_t val) {
#ifdef DEBUG
  if (val > impl::MAX_DILATED<D>::VAL)
    throw std::out_of_range("too big to undilate");
#endif
  val = val & impl::MAX_DILATED<D>::VAL;
  return impl::UNDILATE<D>::eval(val);
}
template <> inline uint64_t dilate<1>(uint64_t val) { return val; }
template <> inline uint64_t undilate<1>(uint64_t val) { return val; }

template <class Index> uint64_t undilate(uint64_t dim, Index val) {
  switch (dim) {
  case 1:
    return undilate<1>(val);
  case 2:
    return undilate<2>(val);
  case 3:
    return undilate<3>(val);
  case 4:
    return undilate<4>(val);
  case 5:
    return undilate<5>(val);
  case 6:
    return undilate<6>(val);
  case 7:
    return undilate<7>(val);
  case 8:
    return undilate<8>(val);
  case 9:
    return undilate<9>(val);
  case 10:
    return undilate<10>(val);
  case 11:
    return undilate<11>(val);
  case 12:
    return undilate<12>(val);
  case 13:
    return undilate<13>(val);
  case 14:
    return undilate<14>(val);
  case 15:
    return undilate<15>(val);
  case 16:
    return undilate<16>(val);
  case 17:
    return undilate<17>(val);
  case 18:
    return undilate<18>(val);
  case 19:
    return undilate<19>(val);
  case 20:
    return undilate<20>(val);
  case 21:
    return undilate<21>(val);
  case 22:
    return undilate<22>(val);
  case 23:
    return undilate<23>(val);
  case 24:
    return undilate<24>(val);
  case 25:
    return undilate<25>(val);
  case 26:
    return undilate<26>(val);
  case 27:
    return undilate<27>(val);
  case 28:
    return undilate<28>(val);
  case 29:
    return undilate<29>(val);
  case 30:
    return undilate<30>(val);
  case 31:
    return undilate<31>(val);
  case 32:
    return undilate<32>(val);
  }
  return 0;
}

template <class Index> uint64_t dilate(uint64_t dim, Index val) {
  switch (dim) {
  case 1:
    return dilate<1>(val);
  case 2:
    return dilate<2>(val);
  case 3:
    return dilate<3>(val);
  case 4:
    return dilate<4>(val);
  case 5:
    return dilate<5>(val);
  case 6:
    return dilate<6>(val);
  case 7:
    return dilate<7>(val);
  case 8:
    return dilate<8>(val);
  case 9:
    return dilate<9>(val);
  case 10:
    return dilate<10>(val);
  case 11:
    return dilate<11>(val);
  case 12:
    return dilate<12>(val);
  case 13:
    return dilate<13>(val);
  case 14:
    return dilate<14>(val);
  case 15:
    return dilate<15>(val);
  case 16:
    return dilate<16>(val);
  case 17:
    return dilate<17>(val);
  case 18:
    return dilate<18>(val);
  case 19:
    return dilate<19>(val);
  case 20:
    return dilate<20>(val);
  case 21:
    return dilate<21>(val);
  case 22:
    return dilate<22>(val);
  case 23:
    return dilate<23>(val);
  case 24:
    return dilate<24>(val);
  case 25:
    return dilate<25>(val);
  case 26:
    return dilate<26>(val);
  case 27:
    return dilate<27>(val);
  case 28:
    return dilate<28>(val);
  case 29:
    return dilate<29>(val);
  case 30:
    return dilate<30>(val);
  case 31:
    return dilate<31>(val);
  case 32:
    return dilate<32>(val);
  }
  return 0;
}

} // namespace util
} // namespace hgeom
