/*
  Copyright (C) 2018 David Robillard <d@drobilla.net>
  Copyright (C) 2006-2007 Chris Hamilton <chamilton@cs.dal.ca>

  This program is free software: you can redistribute it and/or modify it under
  the terms of the GNU General Public License as published by the Free Software
  Foundation, either version 2 of the License, or (at your option) any later
  version.

  This program is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
  details.

  You should have received a copy of the GNU General Public License along with
  this program.  If not, see <https://www.gnu.org/licenses/>.
*/

#ifndef CHILBERT_ALGORITHM_HPP
#define CHILBERT_ALGORITHM_HPP

#include "chilbert/BigBitVec.hpp"
#include "chilbert/FixBitVec.hpp"
#include "chilbert/GetBits.hpp"
#include "chilbert/GetLocation.hpp"
#include "chilbert/GrayCodeRank.hpp"
#include "chilbert/SetBits.hpp"
#include "chilbert/SetLocation.hpp"

#include <cassert>

// Templated Hilbert functions.
// P - is the class used to represent each dimension
//     of a multidimensional point.
// H - is the class used to represent the point as a Hilbert
//     index.
// I - is the class used to represent interim variables.
//     Needs to have as many bits as there are dimensions.
//
// In general, each of P,H,I should be a FixBitVec if they
// fit in that fixed precision.  Otherwise, use a BigBitVec.
// Whatever you use, they must all be of consistent underlying
// storage types.

// The dimension across which the Hilbert curve travels
// principally.
// D0 is d + 1, where d is the actual dimension you want to
// walk across.
// MUST HAVE 0 <= D0 < n
#define D0 1

namespace chilbert {

// 'Transforms' a point.
template <class I>
inline void
transform(const I& e, int d, int n, I& a)
{
	a ^= e;
	a.rotr(d, n); //#D d+1, n );
}

// Inverse 'transforms' a point.
template <class I>
inline void
transformInv(const I& e, int d, int n, I& a)
{
	a.rotl(d, n); //#D d+1, n );
	a ^= e;
}

// Update for method 1 (GrayCodeInv in the loop)
template <class I>
inline void
update1(const I& l, const I& t, const I& w, int n, I& e, int& d)
{
	assert(0 <= d && d < n);
	e = l;
	e.flip(d); //#D d == n-1 ? 0 : d+1 );

	// Update direction
	d += 1 + t.find_first();
	if (d >= n) {
		d -= n;
	}
	if (d >= n) {
		d -= n;
	}
	assert(0 <= d && d < n);

	if (!(w.rack() & 1)) {
		e.flip(d == 0 ? n - 1 : d - 1); //#D d );
	}
}

// Update for method 2 (GrayCodeInv out of loop)
template <class I>
inline void
update2(const I& l, const I& t, const I& w, int n, I& e, int& d)
{
	assert(0 <= d && d < n);
	e = l;
	e.flip(d); //#D d == n-1 ? 0 : d+1 );

	// Update direction
	d += 1 + t.find_first();
	if (d >= n) {
		d -= n;
	}
	if (d >= n) {
		d -= n;
	}
	assert(0 <= d && d < n);
}

template <class P, class H, class I>
inline void
_coordsToIndex(const P* p,
               int      m,
               int      n,
               H&       h,
               I&&      scratch,
               int*     ds = nullptr // #HACK
)
{
	I e{std::move(scratch)};
	I l{e};
	I t{e};
	I w{e};

	// Initialize
	e.reset();
	l.reset();
	h = 0U;

	// Work from MSB to LSB
	int d  = D0;
	int ho = m * n;
	for (int i = m - 1; i >= 0; i--) {
		// #HACK
		if (ds) {
			ds[i] = d;
		}

		// Get corner of sub-hypercube where point lies.
		getLocation<P, I>(p, n, i, l);

		// Mirror and reflect the location.
		// t = T_{(e,d)}(l)
		t = l;
		transform<I>(e, d, n, t);

		w = t;
		if (i < m - 1) {
			w.flip(n - 1);
		}

		// Concatenate to the index.
		ho -= n;
		setBits<H, I>(h, n, ho, w);

		// Update the entry point and direction.
		update2<I>(l, t, w, n, e, d);
	}

	grayCodeInv(h);
}

// This is wrapper to the basic Hilbert curve index
// calculation function.  It will support fixed or
// arbitrary precision, templated.  Depending on the
// number of dimensions, it will use the most efficient
// representation for interim variables.
// Assumes h is big enough for the output (n*m bits!)
template <class P, class H>
inline void
coordsToIndex(const P* p, // [in ] point
              int      m, // [in ] precision of each dimension in bits
              int      n, // [in ] number of dimensions
              H&       h  // [out] Hilbert index
)
{
	if (n <= FBV_BITS) {
		// Intermediate variables will fit in fixed width
		_coordsToIndex<P, H, CFixBitVec>(p, m, n, h, CFixBitVec{});
	} else {
		// Otherwise, they must be BigBitVecs
		_coordsToIndex<P, H, CBigBitVec>(p, m, n, h, CBigBitVec(n));
	}
}

template <class P, class H, class I>
inline void
_indexToCoords(P* p, int m, int n, const H& h, I&& scratch)
{
	I e{std::move(scratch)};
	I l{e};
	I t{e};
	I w{e};

	// Initialize
	e.reset();
	l.reset();
	for (int j = 0; j < n; j++) {
		p[j] = 0U;
	}

	// Work from MSB to LSB
	int d  = D0;
	int ho = m * n;
	for (int i = m - 1; i >= 0; i--) {
		// Get the Hilbert index bits
		ho -= n;
		getBits<H, I>(h, n, ho, w);

		// t = GrayCode(w)
		t = w;
		grayCode(t);

		// Reverse the transform
		// l = T^{-1}_{(e,d)}(t)
		l = t;
		transformInv<I>(e, d, n, l);

		// Distribute these bits
		// to the coordinates.
		setLocation<P, I>(p, n, i, l);

		// Update the entry point and direction.
		update1<I>(l, t, w, n, e, d);
	}
}

// This is wrapper to the basic Hilbert curve inverse
// index function.  It will support fixed or
// arbitrary precision, templated.  Depending on the
// number of dimensions, it will use the most efficient
// representation for interim variables.
// Assumes each entry of p is big enough to hold the
// appropriate variable.
template <class P, class H>
inline void
indexToCoords(P*       p, // [out] point
              int      m, // [in ] precision of each dimension in bits
              int      n, // [in ] number of dimensions
              const H& h  // [out] Hilbert index
)
{
	if (n <= FBV_BITS) {
		// Intermediate variables will fit in fixed width
		_indexToCoords<P, H, CFixBitVec>(p, m, n, h, CFixBitVec{});
	} else {
		// Otherwise, they must be BigBitVecs
		_indexToCoords<P, H, CBigBitVec>(p, m, n, h, CBigBitVec(n));
	}
}

template <class P, class HC, class I>
inline void
_coordsToCompactIndex(const P*   p,
                      const int* ms,
                      int        n,
                      HC&        hc,
                      I&&        scratch,
                      int        M = 0,
                      int        m = 0)
{
	// Get total precision and max precision if not supplied
	if (M == 0 || m == 0) {
		M = m = 0;
		for (int i = 0; i < n; i++) {
			if (ms[i] > m) {
				m = ms[i];
			}
			M += ms[i];
		}
	}

	const int mn = m * n;

	// If we could avoid allocation altogether (ie: have a
	// fixed buffer allocated on the stack) then this increases
	// speed by a bit (4% when n=4, m=20)
	int* const ds = new int[m];

	if (mn > FBV_BITS) {
		CBigBitVec h(mn);
		_coordsToIndex<P, CBigBitVec, I>(p, m, n, h, std::move(scratch), ds);
		compactIndex<CBigBitVec, HC>(ms, ds, n, m, h, hc);
	} else {
		CFixBitVec h;
		_coordsToIndex<P, CFixBitVec, I>(p, m, n, h, std::move(scratch), ds);
		compactIndex<CFixBitVec, HC>(ms, ds, n, m, h, hc);
	}

	delete[] ds;
}

// This is wrapper to the basic Hilbert curve index
// calculation function.  It will support fixed or
// arbitrary precision, templated.  Depending on the
// number of dimensions, it will use the most efficient
// representation for interim variables.
// Assumes h is big enough for the output (n*m bits!)
template <class P, class HC>
inline void
coordsToCompactIndex(const P*   p,  // [in ] point
                     const int* ms, // [in ] precision of each dimension in bits
                     int        n,  // [in ] number of dimensions
                     HC&        hc, // [out] Hilbert index
                     int        M = 0,
                     int        m = 0)
{
	if (n <= FBV_BITS) {
		// Intermediate variables will fit in fixed width?
		_coordsToCompactIndex<P, HC, CFixBitVec>(
		  p, ms, n, hc, CFixBitVec{}, M, m);
	} else {
		// Otherwise, they must be BigBitVecs.
		_coordsToCompactIndex<P, HC, CBigBitVec>(
		  p, ms, n, hc, CBigBitVec(n), M, m);
	}
}

template <class P, class HC, class I>
inline void
_compactIndexToCoords(P*         p,
                      const int* ms,
                      int        n,
                      const HC&  hc,
                      I&&        scratch,
                      int        M = 0,
                      int        m = 0)
{
	I e{std::move(scratch)};
	I l{e};
	I t{e};
	I w{e};
	I r{e};
	I mask{e};
	I ptrn{e};

	// Get total precision and max precision
	// if not supplied
	if (M == 0 || m == 0) {
		M = m = 0;
		for (int i = 0; i < n; i++) {
			if (ms[i] > m) {
				m = ms[i];
			}
			M += ms[i];
		}
	}

	// Initialize
	e.reset();
	l.reset();
	for (int j = 0; j < n; j++) {
		p[j] = 0;
	}

	// Work from MSB to LSB
	int d = D0;

	for (int i = m - 1; i >= 0; i--) {
		// Get the mask and ptrn
		int b = 0;
		extractMask<I>(ms, n, d, i, mask, b);
		ptrn = e;
		ptrn.rotr(d, n); //#D ptrn.Rotr(d+1,n);

		// Get the Hilbert index bits
		M -= b;
		r.reset(); // GetBits doesn't do this
		getBits<HC, I>(hc, b, M, r);

		// w = GrayCodeRankInv(r)
		// t = GrayCode(w)
		grayCodeRankInv<I>(mask, ptrn, r, n, b, t, w);

		// Reverse the transform
		// l = T^{-1}_{(e,d)}(t)
		l = t;
		transformInv<I>(e, d, n, l);

		// Distribute these bits
		// to the coordinates.
		setLocation<P, I>(p, n, i, l);

		// Update the entry point and direction.
		update1<I>(l, t, w, n, e, d);
	}
}

// This is wrapper to the basic Hilbert curve inverse
// index function.  It will support fixed or
// arbitrary precision, templated.  Depending on the
// number of dimensions, it will use the most efficient
// representation for interim variables.
// Assumes each entry of p is big enough to hold the
// appropriate variable.
template <class P, class HC>
inline void
compactIndexToCoords(P*         p,  // [out] point
                     const int* ms, // [in ] precision of each dimension in bits
                     int        n,  // [in ] number of dimensions
                     const HC&  hc, // [out] Hilbert index
                     int        M = 0,
                     int        m = 0)
{
	if (n <= FBV_BITS) {
		// Intermediate variables will fit in fixed width
		CFixBitVec scratch;
		_compactIndexToCoords<P, HC, CFixBitVec>(
		  p, ms, n, hc, CFixBitVec{}, M, m);
	} else {
		// Otherwise, they must be BigBitVecs
		CBigBitVec scratch(n);
		_compactIndexToCoords<P, HC, CBigBitVec>(
		  p, ms, n, hc, std::move(scratch), M, m);
	}
}

} // namespace chilbert

#endif