#pragma once
#include <stdint.h>
#include <stddef.h>
#include <algorithm>
#include <array>

template<typename word_t>
size_t lg(word_t v);


/**
Sean Eron Anderson
seander@cs.stanford.edu
**/
template<>
size_t lg<uint32_t>(uint32_t v)
{
	static const int MultiplyDeBruijnBitPosition[32] = 
	{
	0, 9, 1, 10, 13, 21, 2, 29, 11, 14, 16, 18, 22, 25, 3, 30,
	8, 12, 20, 28, 15, 17, 24, 7, 19, 27, 23, 6, 26, 5, 4, 31
	};

	v |= v >> 1;
	v |= v >> 2;
	v |= v >> 4;
	v |= v >> 8;
	v |= v >> 16;

	return MultiplyDeBruijnBitPosition[(uint32_t)(v * 0x07C4ACDDU) >> 27];
}
template<>
size_t lg<uint64_t>(uint64_t v)
{
	static const int MultiplyDeBruijnBitPosition[64] = 
	{
		0, 58, 1, 59, 47, 53, 2, 60, 39, 48, 27, 54, 33, 42, 3, 61,
        51, 37, 40, 49, 18, 28, 20, 55, 30, 34, 11, 43, 14, 22, 4, 62,
        57, 46, 52, 38, 26, 32, 41, 50, 36, 17, 19, 29, 10, 13, 21, 56,
        45, 25, 31, 35, 16, 9, 12, 44, 24, 15, 8, 23, 7, 6, 5, 63
	};

	v |= v >> 1;
	v |= v >> 2;
	v |= v >> 4;
	v |= v >> 8;
	v |= v >> 16;
	v |= v >> 32;

	return MultiplyDeBruijnBitPosition[(uint64_t)(v * 0x03f6eaf2cd271461) >> 58];
}



template<typename word_t = uint32_t, size_t r = 20, size_t b = 128, word_t P = 0xb7e15163L, word_t Q = 0x9e3779b9L>
class RC6 {
	static constexpr size_t word_size = 8*sizeof(word_t);

	constexpr static word_t r_l(const word_t& w, size_t v) {
		return (w << v) | ( w >> (word_size-v));
	}

	constexpr static word_t r_r(const word_t& w, size_t v) {
		return (w >> v) | ( w << (word_size-v));
	}


	class RC6_KeySched {
	public:
		static constexpr size_t c = (b+word_size-1)/word_size;
		static constexpr size_t v_3 = std::max(c, 2*r+4);
		static constexpr size_t v = v_3*3;
	private:
		std::array<word_t, 2*r+4> S;
	public:
		RC6_KeySched(std::array<word_t, c> L)
		{
			assert(r_l(r_r(13,13),13) == 13);
			S[0] = P;
			for(auto it = S.begin()+1; it < S.end(); ++it)
			{
				*it = *(it-1) + Q;
			}
			word_t A = 0;
			word_t B = 0;
			word_t i = 0;
			word_t j = 0;

			for(size_t s = 0; s < v; ++s)
			{
				i = s % (2*r+4);
				j = s % c;
				A = S[i] = r_l( S[i] + A + B, 3 );
				B = L[j] = r_l( L[j] + A + B, (A + B)%(word_size));
			}
		}

		const word_t& operator[](const size_t pos) {
			return S[pos];
		}

		auto begin()
		{
			return S.begin();
		}

		auto end()
		{
			return S.end();
		}

		auto rbegin()
		{
			return S.rbegin();
		}

		auto rend()
		{
			return S.rend();
		}
	};

	RC6_KeySched S;

public:

	typedef std::array<word_t, RC6_KeySched::c> key_type;
	typedef std::array<word_t, 4> block_type;

	RC6(const key_type& key)
	: S(key)
	{}

	block_type encrypt(block_type plaintext) {
		auto& A = plaintext[0];
		auto& B = plaintext[1];
		auto& C = plaintext[2];
		auto& D = plaintext[3];

		auto it = S.begin();

		B += *(it++);
		D += *(it++);

		for(size_t i = 0; i < r; ++i)
		{
			auto u = r_l( D * ( 2 * D + 1 ), 5);
			auto t = r_l( B * ( 2 * B + 1 ), 5);
			A = r_l((A ^ t), u % word_size) + *(it++);
			C = r_l((C ^ u), t % word_size) + *(it++);
			std::rotate(plaintext.begin(), plaintext.begin()+1, plaintext.end());
		}

		A += *(it++);
		C += *(it++);
		assert(it == S.end());
		return plaintext;
	}

	block_type decrypt(block_type plaintext) {
		auto& A = plaintext[0];
		auto& B = plaintext[1];
		auto& C = plaintext[2];
		auto& D = plaintext[3];
		auto it = S.rbegin();

		C -= *(it++);
		A -= *(it++);

		for(size_t i = 0; i < r; ++i)
		{
			std::rotate(plaintext.begin(), plaintext.end()-1, plaintext.end());
			auto u = r_l( D * ( 2 * D + 1 ), 5);
			auto t = r_l( B * ( 2 * B + 1 ), 5);
			C = r_r( (C - *(it++)) , t % word_size) ^ u ;
			A = r_r( (A - *(it++)) , u % word_size) ^ t ;
		}

		D -= *(it++);
		B -= *(it++);
		assert(it == S.rend());
		return plaintext;
	}

};