#include <iostream>
#include <vector>

// Возведение a в степень b по модулю p
int powmod (int a, int b, int p)
{
	int res = 1;
	while (b)
		if (b & 1)
			res = int ((long long) res * a % p),  --b;
		else
			a = int ((long long) a * a % p),  b >>= 1;
	return res;
}

// Вычисление обратного к a элемента по модулю n
int inverse(int a, int n)
{
	int b0 = n, t, q;
	int x0 = 0, x1 = 1;
	if (n == 1) return 1;
	while (a > 1) {
		q = a / n;
		t = n, n = a % n, a = t;
		t = x0, x0 = x1 - q * x0, x1 = t;
	}
	if (x1 < 0) x1 += b0;
	return x1;
}

// Вычисление примитивного корня числового поля
// (Стандартный DFT использует комплексные числа, поэтому изменим алгоритм для работы в Zp)
// Ограничение: p - простое
int generator (int p)
{
	std::vector<int> fact;
	int phi = p - 1,  n = phi;
	for (int i = 2; i * i <= n; ++i)
		if (n % i == 0)
		{
			fact.push_back (i);
			while (n % i == 0)
				n /= i;
		}
	if (n > 1)
		fact.push_back (n);
 
	for (int res = 2; res <= p; ++res) {
		bool ok = true;
		for (size_t i = 0; i < fact.size() && ok; ++i)
			ok &= powmod (res, phi / fact[i], p) != 1;
		if (ok)  return res;
	}
	return -1;
}

// Реализация прямого DFT над Zp
// Принимает вектор длины n, эта длина используется при вычислении
// примитивного корня поля
// Вычисления приводятся по mod (n + 1), т.к. длина вектора на 1 меньше p,
// где p: Zp
std::vector<int> forward_dft (std::vector<int>& a)
{
	std::vector<int> result;
	int prim_root = generator (a.size());	// Корень степени n из единицы
	for (int i = 0; i < a.size(); i++)
	{
		int sum   = 0;
		int power = 0;
		for (int j = 0; j < a.size(); j++)
		{
			int power_of_root = powmod (prim_root, power, a.size() + 1);
			sum	  += power_of_root * a[j];	// Накапливаем сумму произведений элементов строки матрицы на вектор
			power += i;						// Увеличиваем степень, в которую возводится примитивный корень из единицы
		}
		result.push_back(sum % (a.size() + 1)); // Набрали сумму, приводим по mod p, p = длина вектора + 1
	}
	return result;
}


std::vector<int> inversed_dft (std::vector<int>& a)
{
	std::vector<int> result;
	int prim_root = generator (a.size());	// Корень степени n из единицы
	int inv_prim_root = inverse (prim_root, a.size() + 1);	// Обратный к найденному примитивному корню
	int inv_n = inverse(a.size(), a.size() + 1);
	for (int i = 0; i < a.size(); i++)
	{
		int sum   = 0;
		int power = 0;
		for (int j = 0; j < a.size(); j++)
		{
			int power_of_inv_root = powmod (inv_prim_root, power, a.size() + 1);
			sum	  += inv_n * power_of_inv_root * a[j];	// Накапливаем сумму произведений элементов строки матрицы на вектор
			power += i;							// Увеличиваем степень, в которую возводится примитивный корень из единицы
		}
		result.push_back(sum % (a.size() + 1)); // Набрали сумму, приводим по mod p, p = длина вектора + 1
	}
	return result;
}

// Вывод сообщений
void debug(const std::string& str, std::vector<int>& v)
{
	std::cout << "[DEBUG]:" << str << std::endl;
	for(int i = 0; i < v.size(); i++)
		std::cout << v[i] << " ";
	std::cout << std::endl;
}

int main()
{
	int a[] = {4, 3, 2, 1};
	std::vector<int> v(std::begin(a), std::end(a));
 
	debug("Source vector", v);
 
	std::vector<int> result = forward_dft(v);
 
	debug("Forward DFT", result);
 
	result = inversed_dft(result);
 
	debug("Inversed DFT", result);
}