#ifndef _RECURSIVESHUFFLE_HPP
#define _RECURSIVESHUFFLE_HPP

#include "UnrolledShuffle.hpp"
#include "../Tensor/Tensor.hpp"

template <typename T, unsigned char LOG_WIDTH, unsigned char LOG_SUB_WIDTH, unsigned long R, unsigned long C>
class LogSquareTranspose {
public:
  static void apply(T* __restrict x) {
    constexpr unsigned long SUB_WIDTH = 1ul<<LOG_SUB_WIDTH;
    
    LogSquareTranspose<T, LOG_WIDTH, LOG_SUB_WIDTH-1, R, C>::apply(x);
    LogSquareTranspose<T, LOG_WIDTH, LOG_SUB_WIDTH-1, R, C+(SUB_WIDTH>>1)>::apply(x);
    LogSquareTranspose<T, LOG_WIDTH, LOG_SUB_WIDTH-1, R+(SUB_WIDTH>>1), C>::apply(x);
    LogSquareTranspose<T, LOG_WIDTH, LOG_SUB_WIDTH-1, R+(SUB_WIDTH>>1), C+(SUB_WIDTH>>1)>::apply(x);
  }
};

const unsigned char LOG_BLOCK_WIDTH = 4;

template <typename T, unsigned char LOG_WIDTH, unsigned long R, unsigned long C>
class LogSquareTranspose<T, LOG_WIDTH, LOG_BLOCK_WIDTH, R, C> {
public:
  static void apply(T* __restrict x) {
    constexpr unsigned long WIDTH = 1ul<<LOG_WIDTH;
    constexpr unsigned long BLOCK_WIDTH = 1ul<<LOG_BLOCK_WIDTH;
    for (unsigned int i=R; i<R+BLOCK_WIDTH; ++i)
      for (unsigned int j=C; j<C+BLOCK_WIDTH; ++j)
	if (i < j) {
	  std::swap(x[i*WIDTH+j], x[j*WIDTH+i]);
	}
  }
};

// Uses the idea that rev(AB) = rev(B)rev(A). The word reversal (i.e.,
// changing AB to BA) is a transposition; if A and B have the same
// number of bits (i.e., if the total number of bits is divisible by
// 2), then it is a square matrix transposition. The matrix
// transposition can be performed with improved cache
// performance. Likewise, the rev(A) operations are row operations,
// and so they have good cache locality. Thus the method is more
// efficient than TableShuffle or XORShuffle.

// When the number of bits is odd, first peel off the single least
// significant bit (LSB) and then perform a transposition with a
// buffer (this is essentially ordering of even and odd indices in
// FFT).

template <typename T, unsigned char NUM_BITS>
// Buffer must be at least N/2 in length
inline void lsb_to_msb_with_existing_buffer(T* __restrict x, T* __restrict buffer) {
  constexpr unsigned long N = 1ul<<NUM_BITS;
    
  // Copy odd indices to buffer:
  for (unsigned long k=1; k<N; k+=2)
    buffer[k>>1] = x[k];

  // Pack even indices into first half of x (start at x[1] = x[2]
  // since x[0] already = x[0]):
  for (unsigned long k=2; k<N; k+=2)
    x[k>>1] = x[k];

  // Copy odd indices into second half of x (memcpy equivalent to
  // loopy below):
  memcpy(x+(N>>1), buffer, (N>>1)*sizeof(T));
  //    for (unsigned long k=0; k<N>>1; ++k)
  //      x[k+(N>>1)] = odds[k];
}

template <typename T, unsigned char NUM_BITS>
inline void lsb_to_msb(T* __restrict x) {
  constexpr unsigned long N = 1ul<<NUM_BITS;
    
  // Copy odd indices to buffer:
  T* __restrict odds = aligned_malloc<T>(N>>1);
  lsb_to_msb_with_existing_buffer<T, NUM_BITS>(x, odds);
  free(odds);
}

template <typename T, unsigned char NUM_BITS>
class RecursiveShuffle {
private:
public:
  inline static void apply(T* __restrict x) {
    // & 1 is the same as % 2:
    if ((NUM_BITS & 1) == 1) {
      // allocate buffer and perform single LSB --> MSB:
      lsb_to_msb<T,NUM_BITS>(x);

      RecursiveShuffle<T, NUM_BITS-1>::apply(x);
      RecursiveShuffle<T, NUM_BITS-1>::apply(x+(1ul<<(NUM_BITS-1)));
    }
    else {
      constexpr unsigned char HALF_NUM_BITS = NUM_BITS>>1;
      constexpr unsigned long SQRT_N = 1ul<<HALF_NUM_BITS;

      for (unsigned long k=0; k<SQRT_N; ++k)
      	RecursiveShuffle< T, HALF_NUM_BITS >::apply(x+(k<<HALF_NUM_BITS));
      
      MatrixTranspose<T>::apply_square(x, SQRT_N);
      //      LogSquareTranspose<T, NUM_BITS/2, NUM_BITS/2, 0, 0>::apply(x);
      
      for (unsigned long k=0; k<SQRT_N; ++k)
      	RecursiveShuffle< T, HALF_NUM_BITS >::apply(x+(k<<HALF_NUM_BITS));
    }
  }
};

// For small problems (NUM_BITS <= 9), simply use the closed form
// generated by UnrolledShuffle:
template <typename T>
class RecursiveShuffle<T, 9> {
public:
  inline static void apply(T* __restrict x) {
    UnrolledShuffle<T, 9>::apply(x);
  }
};

template <typename T>
class RecursiveShuffle<T, 8> {
public:
  inline static void apply(T* __restrict x) {
    UnrolledShuffle<T, 8>::apply(x);
  }
};

template <typename T>
class RecursiveShuffle<T, 7> {
public:
  inline static void apply(T* __restrict x) {
    UnrolledShuffle<T, 7>::apply(x);
  }
};

template <typename T>
class RecursiveShuffle<T, 6> {
public:
  inline static void apply(T* __restrict x) {
    UnrolledShuffle<T, 6>::apply(x);
  }
};

template <typename T>
class RecursiveShuffle<T, 5> {
public:
  inline static void apply(T* __restrict x) {
    UnrolledShuffle<T, 5>::apply(x);
  }
};

template <typename T>
class RecursiveShuffle<T, 4> {
public:
  inline static void apply(T* __restrict x) {
    UnrolledShuffle<T, 4>::apply(x);
  }
};

template <typename T>
class RecursiveShuffle<T, 3> {
public:
  inline static void apply(T* __restrict x) {
    UnrolledShuffle<T, 3>::apply(x);
  }
};

template <typename T>
class RecursiveShuffle<T, 2> {
public:
  inline static void apply(T* __restrict x) {
    UnrolledShuffle<T, 2>::apply(x);
  }
};

template <typename T>
class RecursiveShuffle<T, 1> {
public:
  inline static void apply(T* __restrict x) {
    UnrolledShuffle<T, 1>::apply(x);
  }
};

template <typename T>
class RecursiveShuffle<T, 0> {
public:
  inline static void apply(T* __restrict x) {
  }
};

#endif
