using System.Runtime.CompilerServices;

namespace KyberDotnet
{
	public static class Utils
	{
		/// <summary>
		/// Positive modulo defined in Kyber. C# modulo outputs negative values.
		/// </summary>
		/// <param name="x">Number to reduce.</param>
		/// <returns><c>(x mod 3329) >= 0</c></returns>
		[MethodImpl(MethodImplOptions.AggressiveInlining)]
		public static int PosMod(int x)
		{
			x %= Kyber.q;
			return x < 0 ? x + Kyber.q : x;
		}

		//	-------------------------------------------------------------

		/// <summary>
		/// Byte to bit array conversion.
		/// </summary>
		/// <param name="bytes">Bytes to split.</param>
		/// <returns>Bit array.</returns>
		public static byte[] BytesToBits(byte[] bytes)
		{
			var ret = new byte[bytes.Length * 8];
			for (int i = 0; i < ret.Length; i++)
			{
				//For each bit it selects a corresponding byte and bit.
				ret[i] = (byte)((bytes[i >> 3] >> (i & 0b111)) & 0b1);	//div 8, mod 8, select LSB
			}
			return ret;
		}

		/// <summary>
		/// Bit to byte array conversion.
		/// </summary>
		/// <param name="bits">Bits to combine.</param>
		/// <returns>Byte array.</returns>
		public static byte[] BitsToBytes(byte[] bits)
		{
			var ret = new byte[bits.Length / 8];
			for (int i = 0; i < ret.Length; i++)
			{
				for (int j = 0; j < 8; j++)
				{
					ret[i] += (byte)(bits[i * 8 + j] << j);
				}
			}
			return ret;
		}

		//	-------------------------------------------------------------

		/// <summary>
		/// Centered binomial distribution sampling.
		/// <c>bytes.Length = 64*eta</c>
		/// </summary>
		/// <param name="bytes">Bytes to sample from.</param>
		/// <param name="eta">Input length parameter.</param>
		/// <returns>Sampled polynomial.</returns>
		public static Poly CBD(byte[] bytes, int eta)
		{
			var bits = BytesToBits(bytes);
			var ret = new Poly();
			for (int i = 0; i < Kyber.n; i++)
			{
				var A = (from j in Enumerable.Range(0, eta) select (int)bits[2 * i * eta + j]).Sum();
				var B = (from j in Enumerable.Range(0, eta) select (int)bits[2 * i * eta + eta + j]).Sum();
				ret.Data[i] = A - B;
			}
			return ret;
		}

		//	-------------------------------------------------------------

		/// <summary>
		/// Random number sampling.
		/// </summary>
		/// <param name="bytes">Bytes to sample from.</param>
		/// <returns>Sampled polynomial.</returns>
		public static Poly Parse(byte[] bytes)
		{
			int i = 0, j = 0;
			var ret = new Poly();
			while (j < Kyber.n)
			{
				var d1 = bytes[i] + 256 * (bytes[i + 1] & 0b1111); //mod 16
				var d2 = (bytes[i + 1] >> 4) + 16 * bytes[i + 2]; //div 16
				if (d1 < Kyber.q)
				{
					ret.Data[j++] = d1 % Kyber.q;
				}
				if (d2 < Kyber.q && j < Kyber.n)
				{
					ret.Data[j++] = d2 % Kyber.q;
				}
				i += 3;
			}
			return ret;
		}

		//	-------------------------------------------------------------

		/// <summary>
		/// Polynomial to byte array encoding.
		/// <para/><c>bytes.Length = 32*l</c>
		/// </summary>
		/// <param name="poly">Polynomial to encode.</param>
		/// <param name="l">Output length parameter.</param>
		/// <returns>Encoded bytes.</returns>
		public static byte[] Encode(Poly poly, int l)
		{
			var ret = new byte[l * 32 * 8];	//Number of bits for the given length.
			for (int i = 0; i < Kyber.n; i++)
			{
				for (int j = 0; j < l; j++)
				{
					//Inverse of Decode specified in Kyber.
					ret[i * l + j] = (byte)((poly.Data[i] >> j) & 0b1);
				}
			}
			return BitsToBytes(ret);
		}

		/// <summary>
		/// Polynomial vector to byte array encoding.
		/// <para/><c>bytes_n.Length = 32*l</c>
		/// </summary>
		/// <param name="poly">Polynomial vector to encode.</param>
		/// <param name="l">Output length parameter.</param>
		/// <returns>Encoded bytes.</returns>
		public static byte[] Encode(Poly[] poly, int l)
		{
			List<byte> ret = new();
			foreach (Poly p in poly) { ret.AddRange(Encode(p, l)); }
			return ret.ToArray();
		}

		//	-------------------------------------------------------------

		/// <summary>
		/// Byte array to polynomial decoding.
		/// <para/><c>bytes.Length = 32*l</c>
		/// </summary>
		/// <param name="bytes">Bytes do decode.</param>
		/// <param name="l">Input length parameter.</param>
		/// <returns>Decoded polynomial.</returns>
		public static Poly Decode(byte[] bytes, int l)
		{
			var bits = BytesToBits(bytes);
			var ret = new Poly();
			for (int i = 0; i < Kyber.n; i++)
			{
				for (int j = 0; j < l; j++)
				{
					ret.Data[i] += bits[i * l + j] << j;
				}
			}
			return ret;
		}

		/// <summary>
		/// Byte array to polynomial vector decoding.
		/// <para/><c>bytes_n.Length = 32*l</c>
		/// </summary>
		/// <param name="bytes">Bytes do decode.</param>
		/// <param name="l">Input length parameter.</param>
		/// <param name="k">Vector length.</param>
		/// <returns>Decoded vector.</returns>
		public static Poly[] DecodeVector(byte[] bytes, int l, int k)
		{
			var ret = new Poly[k];
			var step = l * 32;
			for (int i = 0; i < k; i++)
			{
				ret[i] = Decode(bytes.Skip(i*step).Take(step).ToArray(), l);
			}
			return ret;
		}

		//	-------------------------------------------------------------

		/// <summary>
		/// Compression specified in Kyber.
		/// </summary>
		/// <param name="x">Number to compress.</param>
		/// <param name="d">Compression parameter.</param>
		/// <returns>Compressed number.</returns>
		public static int Compress(double x, int d)
		{
			return (int)Math.Round(x * (2 << d) / Kyber.q) % (2 << d);
		}

		/// <summary>
		/// In-place polynomial compression.
		/// </summary>
		/// <param name="poly">Polynomial to compress</param>
		/// <param name="d">Compression parameter.</param>
		public static void Compress(Poly poly, int d)
		{
			for (int i = 0; i < Kyber.n; i++)
			{
				poly.Data[i] = Compress(poly.Data[i], d);
			}
		}

		/// <summary>
		/// In-place polynomial vector compression.
		/// </summary>
		/// <param name="vec">Polynomial vector to compress</param>
		/// <param name="d">Compression parameter.</param>
		public static void Compress(Poly[] vec, int d)
		{
			foreach (Poly p in vec) Compress(p, d);
		}

		/// <summary>
		/// Decompression specified in Kyber.
		/// </summary>
		/// <param name="x">Number to decompress.</param>
		/// <param name="d">Decompression parameter.</param>
		/// <returns>Decompressed number.</returns>
		public static int Decompress(double x, int d)
		{
			return (int)Math.Round(x * Kyber.q / (2 << d)) % Kyber.q;
		}

		/// <summary>
		/// In-place polynomial decompression.
		/// </summary>
		/// <param name="poly">Polynomial to decompress.</param>
		/// <param name="d">Decompression parameter.</param>
		public static void Decompress(Poly poly, int d)
		{
			for (int i = 0; i < Kyber.n; i++)
			{
				poly.Data[i] = Decompress(poly.Data[i], d);
			}
		}

		/// <summary>
		/// In-place polynomial vector decompression.
		/// </summary>
		/// <param name="vec">Polynomial vector to decompress.</param>
		/// <param name="d">Decompression parameter.</param>
		public static void Decompress(Poly[] vec, int d)
		{
			foreach (Poly p in vec) Decompress(p, d);
		}
	}
}
