{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The FFT algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_poly(f, n):\n",
    "    '''\n",
    "    Given f as an array of coefficients, pad the polynomial to length n\n",
    "    by adding zeros\n",
    "    '''\n",
    "    if len(f) >= n: return f\n",
    "    \n",
    "    return f + [f[0] - f[0] for _ in range(n - len(f))]\n",
    "\n",
    "def is_power_of_two(n):\n",
    "    if n == 1: return True\n",
    "    if n % 2 == 1: return False\n",
    "    return is_power_of_two(n/2)\n",
    "\n",
    "def FFT(f, omega, n):\n",
    "    '''\n",
    "    Assuming that: \n",
    "     - f is an array of n coefficients\n",
    "     - omega is an n-th principal root of unity\n",
    "     - n is a power of 2\n",
    "    '''\n",
    "    \n",
    "    if n == 1: return [f[0]]\n",
    "    \n",
    "    if not is_power_of_two(n): raise ValueError(f\"{n} is not a power of two\")\n",
    "    \n",
    "    f = pad_poly(f, n)\n",
    "        \n",
    "    \n",
    "    # Precompute all powers of omega\n",
    "    pow_omega = []\n",
    "    omega_i = (omega / omega)\n",
    "    for i in range(n):\n",
    "        pow_omega.append(omega_i)\n",
    "        omega_i = omega_i * omega\n",
    "    \n",
    "    f_even = f[::2]\n",
    "    f_odd  = f[1::2]\n",
    "    \n",
    "    a_even = FFT(f_even, omega*omega, n/2)\n",
    "    a_odd  = FFT(f_odd , omega*omega, n/2)\n",
    "    \n",
    "    output = []\n",
    "    for i in range(n):\n",
    "        output.append(a_even[i % (n/2)] + pow_omega[i] * a_odd[i % (n/2)])\n",
    "    return output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-1"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "R = CC\n",
    "I**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[10, -2*I - 2, -2, 2*I - 2]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "FFT([1,2,3,4], I, 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Polynomial multiplication"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first write down the function where we are provided a suitable root of unity. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def poly_mult_with_rou(f, g, omega, n):\n",
    "    '''\n",
    "    f and g are polynomials given as an array of n coefficients,\n",
    "    and omega is a 2n-th root of unity.\n",
    "    \n",
    "    Assuming that n is a power of 2\n",
    "    '''\n",
    "    \n",
    "    if n == 1: return [f[0]*g[0]]\n",
    "\n",
    "    f_prime = pad_poly(f, 2*n)\n",
    "    g_prime = pad_poly(g, 2*n)\n",
    "        \n",
    "    FFT_f = FFT(f_prime, omega, 2*n)\n",
    "    FFT_g = FFT(g_prime, omega, 2*n)\n",
    "    \n",
    "    FFT_h_times_2n = [FFT_f[i] * FFT_g[i] for i in range(2*n)]\n",
    "    \n",
    "    h_times_2n = FFT(FFT_h_times_2n, omega^-1, 2*n)\n",
    "    \n",
    "    return [h_times_2n[i] / (2*n) for i in range(2*n) ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[3, 10, 8, 0]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f = [1,2] # 1 + 2x\n",
    "g = [3,4] # 3 + 4x\n",
    "poly_mult_with_rou(f, g, omega = I, n = 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "R.<x> = PolynomialRing(QQ)\n",
    "f = R(f)\n",
    "g = R(g)\n",
    "h = R([3,10,8,0])\n",
    "h == f * g"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## When the ring _doesn't_ contain a ROU\n",
    "\n",
    "We'll just look at an example where we wish to multiply polynomials $f$ and $g$ of degree < $n = 4$. This would need an $8$-th principal root of unity but the field $\\mathbb{Q}$ does not have one. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "R.<x,y> = PolynomialRing(QQ)\n",
    "\n",
    "f = 1 + 2*x - x^2 + 4*x^3\n",
    "g = 7 - 2*x + 0   + 3*x^3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's write the above polynomials as a bivariate polynomial ($k = m = 2$)\n",
    "\n",
    "\\begin{align*}\n",
    "f(x) &= 1 + 2x - x^2 + 4x^3 & g(x) &= 7 - 2x + 3x^3\\\\\n",
    "\\tilde{f}(x,y) & = (1 + 2x) + (-1 + 4x)y & \\tilde{g}(x,y) & = (7 - 2x) + (3x)y\\\\\n",
    "f(x) & = \\tilde{f}(x,x^2) & g(x) & = \\tilde{g}(x,x^2)\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_tilde = (1 + 2*x) + (-1 + 4*x)*y\n",
    "g_tilde = (7 - 2*x) + (0  + 3*x)*y\n",
    "\n",
    "f_tilde_coeff = [1 + 2*x , (-1 + 4*x)]\n",
    "g_tilde_coeff = [7 - 2*x , (0  + 3*x)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to multiply $\\tilde{f}$ and $\\tilde{g}$, which are polynomials in $y$ of degree < $2$, we need a $4$-th root of unity. So let's expand the ring to $$R' = \\frac{\\mathbb{Q}[x]}{(x^4 + 1)}.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "R_prime = R.quotient(R.ideal(x^4 + 1))\n",
    "R_prime\n",
    "xbar, ybar = R_prime.gens()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above ring, $\\overline{x}$ is actually an $8$-th root of unity but we only need a $4$-th root of unity so we will choose $\\omega = \\bar{x}^2$ when we need to multiply."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Powers of omega where we evaluate: [1, xbar^2, -1, -xbar^2]\n"
     ]
    }
   ],
   "source": [
    "omega = xbar^2\n",
    "print(f\"Powers of omega where we evaluate: {[omega^i for i in range(4)]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now interpret the two polynomials $\\tilde{f}$ and $\\tilde{g}$ modulo $x^4 + 1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "f_tilde_bar = 4*xbar*ybar + 2*xbar - ybar + 1\n",
      "g_tilde_bar = 3*xbar*ybar - 2*xbar + 7\n"
     ]
    }
   ],
   "source": [
    "f_tilde_bar = R_prime(f_tilde)\n",
    "print(f\"f_tilde_bar = {f_tilde_bar}\")\n",
    "g_tilde_bar = R_prime(g_tilde)\n",
    "print(f\"g_tilde_bar = {g_tilde_bar}\")\n",
    "\n",
    "\n",
    "f_tilde_bar_coeff = [R_prime(a) for a in f_tilde_coeff]\n",
    "g_tilde_bar_coeff = [R_prime(b) for b in g_tilde_coeff]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{align*}\n",
    "\\tilde{f} & = 4 \\bar{x} \\bar{y} + 2\\bar{x} - \\bar{y} + 1\\\\\n",
    "\\tilde{g} & = 3 \\bar{x} \\bar{y} - 2\\bar{x} + 7\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we already have a 4th principal root of unity in our ring ($\\omega = \\overline{x}^2$) we can use FFT to compute the product."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you want to see what FFT does with omega = xbar^2, you can try out the line below:\n",
    "#\n",
    "# print(FFT(f_tilde_bar_coeff, omega = xbar^2, n = 4))\n",
    "\n",
    "# # [\n",
    "# #      6*xbar,\n",
    "# #      4*xbar^3 - xbar^2 + 2*xbar + 1, \n",
    "# #     -2*xbar + 2,\n",
    "# #     -4*xbar^3 + xbar^2 + 2*xbar + 1\n",
    "# # ]\n",
    "# # which is basically just\n",
    "# # [f_tilde(xbar, 1), f_tilde(xbar, xbar^2), f_tilde(xbar, -1), f_tilde(xbar, -xbar^2)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[-4*xbar^2 + 12*xbar + 7, -2*xbar^2 + 33*xbar - 7, 12*xbar^2 - 3*xbar, 0]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h_tilde_bar_coeff = poly_mult_with_rou(\n",
    "                        f_tilde_bar_coeff, \n",
    "                        g_tilde_bar_coeff, \n",
    "                        omega = xbar^2, \n",
    "                        n = 2\n",
    "                    )\n",
    "h_tilde_bar_coeff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h_tilde = (-4*x^2 + 12*x + 7) + (-2*x^2 + 33*x - 7)*y + (12*x^2 - 3*x)*y^2\n",
    "h_tilde == f_tilde * g_tilde"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h = h_tilde.subs(y = x^2)\n",
    "f*g == h"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "SageMath 9.4",
   "language": "sage",
   "name": "sagemath-9.4"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}