modular_gaussian_integers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9782aa60",
   "metadata": {},
   "source": [
    "# Modular Gaussian Integers\n",
    "\n",
    "*Dmytro Fedoriaka, July 2024*\n",
    "\n",
    "## 1. Introduction\n",
    "\n",
    "In competitive programming it is common to do computations modulo a prime number $M$. It is convenient because one can define and efficiently compute an inverse modulo prime number (e.g. using formula $x^{-1} \\equiv x^{M-2} (\\text{mod} ~M))$. In other words, $\\mathbb{Z}_M$ is a field. So, one can do arithmetic pretty much like if these numbers are real numbers and get correct answers.\n",
    "\n",
    "However, there are other operations that are not well-defined in a field $\\mathbb{Z}_M$. In particular, a square root is not defined for all elements in $\\mathbb{Z}_M$. For example, there is no solutions to equation $x^2 = 3$ in $\\mathbb{Z}_7$, which means square roots of 3 is not defined.\n",
    "\n",
    "What to do in this case when we still need to get square root for any number? Well, what we do for real numbers $\\mathbb{R}$ when we want to extract square root of a negative number? We define a new symbol $i$ and say that by definition $i^2=1$, then add it to $\\mathbb{R}$ and also close new structure under all arithmetic operations (this construction is called a [field extension](https://en.wikipedia.org/wiki/Field_extension)). As a result we get new field $\\mathbb{R}[i]=\\mathbb{C}$, which we call \"complex numbers\".\n",
    "\n",
    "Here I propose to do the same thing to a field of residuals modulo M. It turns out that resulting construction has some interesting properties. It is a field and it allows to extract square roots of any number from $\\mathbb{Z}_M$, and much more.\n",
    "\n",
    "Below I will describe some properties of $\\mathbb{Z}_M$ and describe how to efficiently (in $O(poly(\\log M))$) compute roots of any degree in it, and order of any non-zero element. Finally, I will show a practical application of this constructions: solving cubic equations in $\\mathbb{Z}_M$.\n",
    "\n",
    "**Note.** *I couldn't find literature explicitly describing properties $\\mathbb{Z}_M[i]$, but it's hard to believe it doesn't exist. So I didn't bother to give rigorous mathematical proofs (so far). Instead, I focused on describing interesting facts and providing explicit algorithms in Python. Also, my notation might be different from conventional.*\n",
    "\n",
    "## 2. Definition\n",
    "\n",
    "Let $M$ be a prime number such that $M \\equiv 3 ~ (\\text{mod} ~ 4)$.\n",
    "\n",
    "Define $\\mathbb{Z}_M$ - a field of residuals modulo $M$.\n",
    "\n",
    "Define $\\mathbb{Z}_M[i] = \\big \\{ a+b \\cdot i ~|~ a,b \\in \\mathbb{Z}_M \\big \\}$, where $i$ is such a symbol that by definition $i^2 = 1$. \n",
    "\n",
    "I call these \"Modular Gaussian Integers\", because these are like [Gaussian Integers](https://en.wikipedia.org/wiki/Gaussian_integer) $\\mathbb{Z}[i]$, except their coefficients are from $\\mathbb{Z}_M$ instead of $\\mathbb{Z}$.\n",
    "\n",
    "## 3. Elementary operations\n",
    "\n",
    "\n",
    "**Arithmetic operations** are defined the same way as for regular complex numbers $\\mathbb{C} = \\mathbb{R}[i]$. Note that coefficients are from $\\mathbb{Z}_M$, so it's assumed that all operations with coefficients are modulo M (e.g. when we write $ab$, we imply $(ab)\\%M$) .\n",
    "\n",
    "* Addition: $(a+b i) + (c+di) = (a+c) + (b+d)i$\n",
    "* Subtraction: $(a+b i) - (c+di) = (a+c) - (b+d)i$\n",
    "* Multiplication: $(a+b i) \\cdot (c+di) = (ac-bdc) - (ad+bc)i$\n",
    "* Division: $\\frac{a+bi}{c+di} = \\frac{(a+bi)(c-di)}{(c+di)(c-di)} = \\frac{(ac+bd) + (bc-ad)i}{c^2+d^2} = \n",
    "(ac+bd)\\cdot (c^2+d^2)^{-1} + (bc-ad) \\cdot (c^2+d^2)^{-1} \\cdot i $\n",
    "\n",
    "Division involves finding inverse of $(c^2+d^2)$ modulo $M$. So, $\\mathbb{Z}_M[i]$ is a field if and only if this inverse exists for any non-zero $c+d i$. This holds if and only if $M$ is prime and $M \\equiv 3 ~(\\text{mod}~ 4)$. See proof [here](https://math.stackexchange.com/questions/2160330/about-ring-of-gaussian-integers-modulo-n).\n",
    "\n",
    "Formally, we can claim that:\n",
    "* $\\mathbb{Z}_M$ is a ring for any integer $M \\ge 2$.\n",
    "* $\\mathbb{Z}_M$ is a field if and only if $M$ is prime and $M \\equiv 3 ~(\\text{mod}~ 4)$.\n",
    "\n",
    "**Exponentiation.** We can define $a^b$ as multiplying $a$ by itself $b$ times. This can be computed in $O(\\log M)$ using binary exponentiation.\n",
    "\n",
    "**Complex conjugate.** Define $\\overline{a+b i}= a-bi$.\n",
    "\n",
    "**Modulus.** Define $|a+b i| = a^2+b^2 \\in Z_M$. Note that unlike in $\\mathbb{C}$, we **are not taking square root**, because it's not always possible in $Z_M$. \n",
    "\n",
    "It is easy to check that:\n",
    "  * $|x|=x \\cdot \\overline{x}$\n",
    "  * $|x| \\cdot |y| = |x \\cdot y|$\n",
    "  * $|x|=0 \\Leftrightarrow x=0$\n",
    "  * $|x/\\overline{x}|=|\\overline{x}/x|=1$ for any non-zero $x$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c54144d9",
   "metadata": {},
   "source": [
    "## 4. Structure of multiplicative group of $\\mathbb{Z}_M[i]$\n",
    "\n",
    "*This section is not very useful for algorithms and can be skipped.*\n",
    "\n",
    "We know that in $\\mathbb{C}$ any number can be written as $z = A \\cdot e^{i \\phi}$, where $A \\in \\mathbb{R}, A\\ge 0$ and $\\phi \\in [0, 2\\cdot \\pi)$. Let's do something similar in $Z_M[i]$. In $\\mathbb{C}$ we could just use $A=|z|$, $e^{i \\phi} = z/|z|$. here it's a bit more complicated.\n",
    "\n",
    "**Lemma.** Any non-zero $x = \\mathbb{Z}_M[i]$ can be written as $x \\in A \\cdot B$,\n",
    "where are $A \\in Z_M$ (i.e. it's \"real\") and $|B|=\\pm 1$.\n",
    "\n",
    "**Proof.**\n",
    "Consider $|x|$. There are two possible cases\n",
    "   \n",
    "   1. $|x|$ has square roots $\\pm r_1$ with $r_1 \\in \\mathbb{Z}_M, r_1 \\ne 0$, Then $r1_2=|x|$. Take $A=r_1, B=x/r_1$. Then $A \\in \\mathbb{Z}_M, |B| = \\frac{|x|}{r_1^2}=1$. \n",
    "   \n",
    "   2. $|x|$ has square roots $\\pm r_2 \\cdot i$ with $r_2 \\in \\mathbb{Z}_M, r_2 \\ne 0$. Then $r2_2=-|x|$. Take $A=r_2, B = x/r_2$. Then $A \\in \\mathbb{Z}_M, |B| = \\frac{|x|}{r_2^2}=-1$. \n",
    "\n",
    "\n",
    "Next, note that set of all numbers with modulus $\\pm 1$ forms a group under multiplication of order $2(M+1)$. Denote $\\omega_2$ - a primitive root in it (an element with order $2(M+1))$. Then we can write our representation as\n",
    "\n",
    "$$ x = A \\cdot \\omega_2 ^{n_2},$$\n",
    "\n",
    "where $n_2 \\in [0, 2(M+1))$.\n",
    "\n",
    "All real numbers also form a subgroup, which has a primitive root $\\omega_1$. So, we can write \n",
    "\n",
    "$$ x = \\omega_1 ^{n_1} \\cdot \\omega_2 ^{n_2},$$\n",
    "\n",
    "where $n_1 \\in [0, M-1),  n_2 \\in [0, 2(M+1))$.\n",
    "\n",
    "This representation is non-unique. In fact, every $x$ has exactly two representations of this form, because \n",
    "\n",
    "$$  \\omega_1 ^{n_1} \\cdot \\omega_2 ^{n_2} =  (\\omega_1 ^{n_1} \\cdot (-1)) \\cdot (\\omega_2 ^{n_2} \\cdot (-1)) = \n",
    "\\omega_1 ^{(n_1+(M-1)/2)\\%(M-1)} \\cdot \\omega_2 ^{(n_2+(M+1))\\%(2M+2)} $$\n",
    "\n",
    "We can break ambiguity by restricting $n_1$ to $[0, (M-1)/2)$. Define $K=(M-1)/2$. Then $2(M+1)=4(K+1)$. Recall that $M\\%4=3$, so $K$ is odd, then $gcd(K,4(K+1))=1$. We can write $x = \\omega_1 ^{n_1} \\cdot \\omega_2 ^{n_2}$, where $n_1 \\in [0, K), n_2 \\in [0,4(K+1)$. Using Chinese Remainder Theorem, now we can claim that:\n",
    "\n",
    "**Lemma.** Any non-zero $x \\in \\mathbb{Z}_M[i]$ can be written in unique way:\n",
    "\n",
    "* as $x = \\omega_1 ^{n_1} \\cdot \\omega_2 ^{n_2}$, where $n_1 \\in [0, K), n_2 \\in [0,4(K+1)$,\n",
    "* as $x = (\\omega_1 \\omega_2) ^{n}$, where $n \\in [0, M^2-1)$.\n",
    "\n",
    "So, we have proven that multiplicative group of $Z_M[i]$ is cyclic with generator $\\omega = \\omega_1 \\cdot \\omega_2$.\n",
    "\n",
    "\n",
    "How to efficiently find $\\omega_1$ and $\\omega_2$?\n",
    "* To find $\\omega_1$, randomly generate an integer $a \\in [2, M-1]$, until $ord(a + 0 \\cdot i)=M-1$.\n",
    "* To find $\\omega_2$, randomly generate an integer $a \\in [1, M-1]$. Consider $b = \\sqrt{-1-a^2}$. If it's not real. Otherwise, check if $ord(a + b \\cdot i)=2M+1$.\n",
    "\n",
    "In both cases we will need to do only $O(1)$ attempts.\n",
    "\n",
    "Of course, we didn't need to go to all this trouble, as it is a known fact that that multiplicative group of a finite field is a cyclic group. However, this representation will be useful for taking discrete logarithms.\n",
    "\n",
    "\n",
    "\n",
    "Some other interesting facts about the multiplicative group:\n",
    "\n",
    "\n",
    "* Of $M-1$ non-zero elements in $Z_M$, exactly half have roots in $Z_M$.\n",
    "\n",
    "* For any two non-zero numbers $x$ and $-x$ from $Z_M$, one of them has two real roots $\\pm y$, and the other has two purely imaginary roots $\\pm y \\cdot i$. \n",
    "    * **Proof sketch.** Consider mapping $f: Z_M \\to Z_M$, $f(x)=x^2$. Obviously, it maps numbers $x$ and $-x$\n",
    "       to the same number. But also if it maps some number $x$ to $y$, it doesn't map any number to $-y$ (in\n",
    "       other words, there are no such $a,b \\ne 0$ that $a^2=-b^2$ because that would mean $a^2+b^2=0$).\n",
    "* Of $M^2-1$ non-zero elements in $Z_M[i]$, exactly half have roots in $Z_M$ (again, each of them having two roots summing to 0), and the other half don't.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60abdf0f",
   "metadata": {},
   "source": [
    "## 5. Computing order in $\\mathbb{Z}_M[i]$ in $O(\\log(M)^2)$\n",
    "\n",
    "**Definition.** Order of non-zero $x \\in \\mathbb{Z}_M[i]$ is a smallest positive integer $n$ such that $x^n=1$.\n",
    "\n",
    "By the [Lagrange theorem](https://en.wikipedia.org/wiki/Lagrange%27s_theorem_(group_theory)), the order of $x$ is divisor of the size of group, i.e. divisor of $(M-1)(M+1)$, so we can just check all divisors $d$ in order and find the smallest of them such that $x^d=1$. However, we can do much faster.\n",
    "\n",
    "Let's say that if $x^d=1$, then $d$ is *candidate order*. Consider known prime factorization $M^2-1 = p_1^{a_1}p_2^{a_2}...p_l^{a_l}$. We know that the answer will be \n",
    "$p_1^{b_1}p_2^{b_2}...p_l^{b_l}$ where $0 \\le b_i \\le a_i$. Also we know that if $X$ is a candidate order, any its multiple is candidate order, and if $X$ is not candidate order, any divisor of $X$ is not a candidate order.\n",
    "\n",
    "Then we can propose the following algorithm:\n",
    "* Start with $A = p_1^{a_1}p_2^{a_2}...p_l^{a_l}$.\n",
    "* For all $i=1 \\dots l$:\n",
    "   * While $a_i > 0$ and $x^A = 1$ decrement $a_i$.\n",
    "   * If $x^A \\ne 1$, increment $a_i$.\n",
    "* Resulting X is the order of $x$.\n",
    "\n",
    "This algorithms needs to compute power at most $a_1+a_2 +\\cdot + a_l = O(\\log(M))$ times, and computing power is done in $O(\\log(M))$, so the total complexity is $O(\\log(M^2)).$\n",
    "\n",
    "\n",
    "**Note about complexity.** Here and below we assume that we already have factorization of $M^2-1$, and don't include cost of computing this factorization into complexity of operations. In practice, for $M$ up to $10^9$ we can use naive factorization of $M^2-1=(M-1)(M+1)$ by checking all divisors up to root, having compliexity $O(\\sqrt{M})$. For larger $M$ we could use more advanced factoring algorithms. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "131c944f",
   "metadata": {},
   "source": [
    "## 6. Computing $r^{th}$ roots in $\\mathbb{Z}_M[i]$ in $O(\\log^4(M) + r \\log^3(M))$\n",
    "\n",
    "There is a classic [Tonelli-Shanks](https://en.wikipedia.org/wiki/Tonelli%E2%80%93Shanks_algorithm) algorithm that computes square roots in $Z_M$. Actually, it can be generalized to compute any root in any finite field. This is called Adleman-Manders-Miller algorithm, and it is well-described in [this paper](https://arxiv.org/pdf/1111.4877) by Zhengjun Cao, Qian Sha and Xiao Fan (table 4 on page 8), so I won't repeat it here.\n",
    "\n",
    "Some implementation details:\n",
    "* $q = M^2$\n",
    "* When it says, \"Compute the least nonnegative integer α such that s|rα − 1.\", this just means $\\alpha = r^{-1} (\\text{mod} ~s)$.\n",
    "* When it says to compute $\\log_a d$, just find smallest $j$ such that $a^j=d$ by checking $j=0,1,\\dots,r-1$, and it's guaranteed that you'll find the answer.\n",
    "* It covers only cases when $r$ is prime and divisor of $M^2-1$.\n",
    "  * If $r$ is composite, just compose roots for all prime factors, e.g. $\\sqrt[6]{x} = \\sqrt[2]{\\sqrt[3]{x}}.$\n",
    "  * If $r$ is prime but not a divisor of $M^2-1$, then $\\sqrt[r]{x}=x^{r^{-1}(\\text{mod}~M^2-1)}$.\n",
    "  \n",
    "The complexity $O(\\log^4(M) + r \\log^3(M))$ is from the same paper.\n",
    "\n",
    "Now when we found one root, how to find all roots? Just multiply the found root by all $r$-th roots of unity. The $r$-th roots of unity are easy to compute using the fact that multiplicative group is cyclic group. There will be $g=GCD(r, M^2-1)$ of them, and they will be of form $\\omega^{kj}$ where $k=\\frac{M^2-1}{g}$ and $j<=0,1,\\dots,g-1$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5df1277",
   "metadata": {},
   "source": [
    "## 7. Computing discrete logarithm in $\\mathbb{Z}_M[i]$ in $O(\\sqrt{M})$ (sketch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df84a44d",
   "metadata": {},
   "source": [
    "Let's say that $\\log_A{x}=y$ if $A^y=x$. How to efficiently find $\\log_A{x}$.\n",
    "\n",
    "Decompose $x=\\omega_1^{x_1} \\omega_2^{x_2}$, $A=\\omega_1^{A} \\omega_2^{A}$. Then we are looking for such an integer $y$ that\n",
    "\n",
    "$$\\omega_1^{A_1 y} \\omega_2^{A_2 y} = \\omega_1^{x_1} \\omega_2^{x_2}\n",
    "\\Rightarrow\n",
    "\\begin{cases}\n",
    "    A_1 y \\equiv x_1 (\\text{mod} ~M-1) \\\\\n",
    "    A_2 y \\equiv x_2 (\\text{mod}~ 2(M+1))\n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "This can be solved using Chinese Remainder Theorem and the [Extended Euclidean Algorithm](https://cp-algorithms.com/algebra/extended-euclid-algorithm.html). Note that answer doesn't always exist.\n",
    "\n",
    "But how to compute the decompositions? Above we have shown how to compute decomposition $x=AB$ where $A=\\omega_1^{n_1}$, $B=\\omega_2^{n_2}$ but not how to find $n_1,n_2$. Now we have the problem: given $x$ in multiplicative subgroup of order $s$ and primitive root $\\omega$ of that subgroup, find $n$ such that $x=\\omega^x$. This can be solved using the [Baby-step giant-step algorithm](https://cp-algorithms.com/algebra/discrete-log.html) in $O(s)$. So the overall complexity for the original task is $O(\\sqrt{M})$.\n",
    "\n",
    "Note that we could apply the BSGS algorithm directly in $\\mathbb{Z}_M[i]$ and that would have complexity $O(\\sqrt{M^2-1}) = O(M)$. We reduced this complexity by factoring the group. This makes me think that we can continue this and factor the multiplicative group into as many \"orthogonal\" subgroups as many there are factors in prime decompositions of $M^2-1$ and then complexity of discrete log would be $O(\\sqrt{p})$, where $p$ is the largest prime factor in $M^2-1$. However, in the worst case it's still $O(\\sqrt{M})$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07bb8b87",
   "metadata": {},
   "source": [
    "## 8. Implementation\n",
    "\n",
    "Below is the full implementation of the Modular Gaussian Integers in Python.\n",
    "\n",
    "Notes:\n",
    "* It's deliberately written with no dependencies on external libraries, so it could be used in competitive programming.\n",
    "* Prime factoring is done using naive algorithm, which is fine for $M \\le 10^9$. The factoring needs to be only computed once for given $M$ and the is cached. To work with larger $M$, use `sympy.primefactors` from SymPy instead of `primefactors` (although this code will still work for any $M$, just be slower).\n",
    "* If rewriting in C++, for $M>2^{31}$ you will need long arithmetic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "915d1417",
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools\n",
    "import math\n",
    "import random\n",
    "\n",
    "\n",
    "@functools.cache\n",
    "def primefactors(x: int) -> set[int]:\n",
    "    ans = set()\n",
    "    for d in range(2, x):\n",
    "        if d*d > x:\n",
    "            break\n",
    "        while x % d == 0:\n",
    "            ans.add(d)\n",
    "            x //= d\n",
    "    if x > 1:\n",
    "        ans.add(x)\n",
    "    return ans\n",
    "  \n",
    "is_prime = lambda x: primefactors(x) == {x}\n",
    "  \n",
    "def pow_mod(x: int, y: int, M: int) -> int:\n",
    "    \"\"\"Returns (x**y)%M.\"\"\"\n",
    "    ans = 1\n",
    "    while y != 0:\n",
    "        if (y & 1):\n",
    "            ans = (ans*x) % M\n",
    "        x = (x*x) % M\n",
    "        y >>= 1\n",
    "    return ans\n",
    "\n",
    "\n",
    "class ModGaussianPool:\n",
    "    \"\"\"Manages all Modular Gaussian numbers with the same modulo M.\"\"\"\n",
    "\n",
    "    def __init__(self, M: int):\n",
    "        assert is_prime(M), \"M is not prime.\"\n",
    "        assert M % 4 == 3\n",
    "        self.M = M\n",
    "\n",
    "        # Prime factors of M^2-1 (order or multiplicative group).\n",
    "        self.M2m1_prime_factors = primefactors(M-1).union(primefactors(M+1))\n",
    "\n",
    "        # Find primitive root in subgroup of real numbers.\n",
    "        for i in range(2, M):\n",
    "            if self.num(i).get_order() == M-1:\n",
    "                self.w1 = self.num(i)\n",
    "                break\n",
    "        assert (self.w1.pow(self.M - 1).is_one())\n",
    "        assert (self.w1.pow((self.M - 1)//2) == self.num(-1))\n",
    "\n",
    "        # Find primitive root in subgroup of elements having modulus +1/-1.\n",
    "        while True:\n",
    "            a = random.randint(1, self.M-1)\n",
    "            b2 = self.num(-1-a**2)\n",
    "            b = b2._discrete_root_amm(2)\n",
    "            if b.b != 0:\n",
    "                continue\n",
    "            x = self.num(a, b.a)\n",
    "            assert (x.a**2+x.b**2) % M == M - 1\n",
    "            if x.get_order() == 2*(self.M + 1):\n",
    "                self.w2 = x\n",
    "                break\n",
    "        assert self.w2.pow(2*(self.M + 1)).is_one()\n",
    "        assert self.w2.pow(self.M + 1) == self.num(-1)\n",
    "        assert (self.w1*self.w2).get_order() == self.M**2 - 1\n",
    "\n",
    "    def __repr__(self):\n",
    "        return \"ModGaussianPool(%d)\" % self.M\n",
    "\n",
    "    @functools.cache\n",
    "    @staticmethod\n",
    "    def make_pool(M: int):\n",
    "        return ModGaussianPool(M)\n",
    "\n",
    "    @functools.cache\n",
    "    def unity_roots(self, r: int):\n",
    "        n = self.M**2-1\n",
    "        gen = self.w1*self.w2\n",
    "        return [gen.pow(i) for i in range(0, n, n // math.gcd(r, n))]\n",
    "      \n",
    "    @functools.cache\n",
    "    def totient(self, x: int) -> int:\n",
    "        assert (self.M**2 - 1) % x == 0\n",
    "        for p in self.M2m1_prime_factors:\n",
    "            if x % p == 0:\n",
    "                x = (p-1)*(x//p)\n",
    "        return x\n",
    "    \n",
    "    def inv_mod(self, x: int, s: int) -> int:\n",
    "        \"\"\"Returns x^-1 (mod s), where s is divisor of M^2-1.\"\"\"\n",
    "        assert math.gcd(x, s) == 1, \"Not invertible\"\n",
    "        return pow_mod(x, self.totient(s)-1, s)  \n",
    "\n",
    "    def all_numbers(self):\n",
    "        return (ModGaussian(i, j, self) for i in range(self.M) for j in range(self.M))\n",
    "\n",
    "    def num(self, a: int, b: int = 0) -> \"ModGaussian\":\n",
    "        return ModGaussian(a, b, self)\n",
    "\n",
    "    def random(self) -> \"ModGaussian\":\n",
    "        return ModGaussian(random.randint(0, self.M-1), random.randint(0, self.M-1), self)\n",
    "\n",
    "    @functools.cache\n",
    "    def get_rho(self, r: int) -> \"ModGaussian\":\n",
    "        \"\"\"Returns non-zero element that is not an r-th root.\"\"\"\n",
    "        p = (self.M**2-1)//r\n",
    "        while True:\n",
    "            rho = self.random()\n",
    "            if not (rho.is_zero() or rho.pow(p).is_one()):\n",
    "                return rho\n",
    "\n",
    "\n",
    "class ModGaussian:\n",
    "    def __init__(self, a: int, b: int, pool: ModGaussianPool):\n",
    "        self.a = a % pool.M\n",
    "        self.b = b % pool.M\n",
    "        self.pool = pool\n",
    "\n",
    "    def __add__(self, other):\n",
    "        assert self.pool is other.pool\n",
    "        return ModGaussian(self.a+other.a, self.b+other.b, self.pool)\n",
    "\n",
    "    def __sub__(self, other):\n",
    "        assert self.pool is other.pool\n",
    "        return ModGaussian(self.a-other.a, self.b-other.b, self.pool)\n",
    "\n",
    "    def __mul__(self, other):\n",
    "        if type(other) is int:\n",
    "            return ModGaussian(self.a*other, self.b*other, self.pool)\n",
    "        assert self.pool is other.pool\n",
    "        return ModGaussian(self.a*other.a - self.b*other.b, self.a*other.b + self.b*other.a, self.pool)\n",
    "\n",
    "    def __truediv__(self, other):\n",
    "        if type(other) is int:\n",
    "            assert other != 0, \"Division by zero\"\n",
    "            inv = pow_mod(other, self.pool.M-2, self.pool.M)\n",
    "            return ModGaussian(self.a*inv, self.b*inv, self.pool)\n",
    "        assert self.pool is other.pool\n",
    "        a = self.a*other.a + self.b*other.b\n",
    "        b = self.b*other.a - self.a*other.b\n",
    "        denom = (other.a**2 + other.b**2) % self.pool.M\n",
    "        assert denom != 0, \"Division by zero\"\n",
    "        denom_inv = pow_mod(denom, self.pool.M-2, self.pool.M)\n",
    "        return ModGaussian(a * denom_inv, b * denom_inv, self.pool)\n",
    "\n",
    "    def __pow__(self, other):\n",
    "        return self.pow(other)\n",
    "\n",
    "    def is_zero(self):\n",
    "        return self.a == 0 and self.b == 0\n",
    "\n",
    "    def is_one(self):\n",
    "        return self.a == 1 and self.b == 0\n",
    "\n",
    "    def __repr__(self):\n",
    "        return str(self.a) if self.b == 0 else \"%d+%dj\" % (self.a, self.b)\n",
    "\n",
    "    def __eq__(self, other):\n",
    "        if type(other) is not ModGaussian:\n",
    "            return False\n",
    "        return self.a == other.a and self.b == other.b and self.pool is other.pool\n",
    "\n",
    "    def __hash__(self):\n",
    "        return hash(self.a * self.pool.M + self.b)\n",
    "\n",
    "    def pow(self, y):\n",
    "        ans = ModGaussian(1, 0, self.pool)\n",
    "        x = self\n",
    "        y %= (self.pool.M**2 - 1)\n",
    "        while y != 0:\n",
    "            if (y & 1):\n",
    "                ans = ans*x\n",
    "            x = x*x\n",
    "            y >>= 1\n",
    "        return ans\n",
    "\n",
    "    def conj(self):\n",
    "        return ModGaussian(self.a, -self.b, self.pool)\n",
    "\n",
    "    def get_order(self) -> int:\n",
    "        \"\"\"Returns smallest i, such that self^i=1.\"\"\"\n",
    "        assert not self.is_zero(), \"Order is not defined for the zero element.\"\n",
    "        ans = self.pool.M**2 - 1\n",
    "        assert self.pow(ans).is_one()\n",
    "        for f in self.pool.M2m1_prime_factors:\n",
    "            while ans % f == 0:\n",
    "                ans //= f\n",
    "                if not self.pow(ans).is_one():\n",
    "                    ans *= f\n",
    "                    break\n",
    "        return ans\n",
    "\n",
    "    # Finds disrete root for the case when r is prime.\n",
    "    # Reports None if root doesn't exist.\n",
    "    # Uses Adleman-Manders-Miller algorithm, see https://arxiv.org/pdf/1111.4877\n",
    "    def _discrete_root_amm(self, r) -> \"Optional[ModGaussian]\":\n",
    "        # Handle trivial case when r is not a divisor of M^2-1.\n",
    "        qm1 = self.pool.M**2-1  # \"q-1\" in the paper.\n",
    "        if qm1 % r != 0:\n",
    "            return self.pow(self.pool.inv_mod(r, qm1))\n",
    "\n",
    "        # Check if it's a r-th residue.\n",
    "        if not self.pow(qm1//r).is_one():\n",
    "            return None\n",
    "\n",
    "        rho = self.pool.get_rho(r)\n",
    "        s, t = qm1, 0\n",
    "        while s % r == 0:\n",
    "            s //= r\n",
    "            t += 1\n",
    "        assert (t >= 1)\n",
    "\n",
    "        al = self.pool.inv_mod(r, s)  # least non-negative such that (r*al-1)%s==0\n",
    "        assert ((r*al-1) % s == 0)\n",
    "\n",
    "        a = rho.pow(r**(t-1)*s)\n",
    "        assert not a.is_zero()\n",
    "        b = self.pow(r*al-1)\n",
    "        c = rho.pow(s)\n",
    "        h = self.pool.num(1)\n",
    "        for i in range(1, t):\n",
    "            d = b.pow(r**(t-1-i))\n",
    "\n",
    "            # Compute the discrete logarithm.\n",
    "            a_j = self.pool.num(1)\n",
    "            j = None\n",
    "            for j1 in range(r):\n",
    "                if d == a_j:\n",
    "                    j = -j1\n",
    "                    break\n",
    "                a_j = a_j * a\n",
    "            assert j is not None\n",
    "\n",
    "            b = b * c.pow(r).pow(j)\n",
    "            h = h * c.pow(j)\n",
    "            c = c.pow(r)\n",
    "        return self.pow(al) * h\n",
    "\n",
    "    def discrete_root(self, r) -> \"Optional[ModGaussian]\":\n",
    "        \"\"\"Finds some r-th root or reports that there is none.\"\"\"\n",
    "        if self.is_zero():\n",
    "            return self\n",
    "        ans = self\n",
    "        for f in primefactors(r):\n",
    "            while r % f == 0:\n",
    "                r /= f\n",
    "                ans = ans._discrete_root_amm(f)\n",
    "                if ans is None:\n",
    "                    return None\n",
    "        return ans\n",
    "\n",
    "    def all_roots(self, r) -> \"List[ModGaussian]\":\n",
    "        \"\"\"Returns all r-th roots.\"\"\"\n",
    "        if self.is_zero():\n",
    "            return [self]\n",
    "        r %= (self.pool.M**2 - 1)\n",
    "        r0 = self.discrete_root(r)\n",
    "        if r0 is None:\n",
    "            return []\n",
    "        return [r0 * ur for ur in self.pool.unity_roots(r)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52eb2984",
   "metadata": {},
   "source": [
    "### Example usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fc70b400",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a=1+14j, b=3+4j, a+b=4+18j, a-b=4+18j, a*b=4+8j, a/b=13\n",
      "a**5=7\n",
      "order(a)=15\n",
      "Square roots of a: [14+10j, 5+9j]\n",
      "Cubic roots of a: [10+1j, 15+11j, 13+7j]\n"
     ]
    }
   ],
   "source": [
    "pool = ModGaussianPool.make_pool(19)\n",
    "a = pool.num(1,14)\n",
    "b = pool.num(3,4)\n",
    "print(f\"a={a}, b={b}, a+b={a+b}, a-b={a+b}, a*b={a*b}, a/b={a/b}\")\n",
    "print(f\"a**5={a**5}\")\n",
    "print(f\"order(a)={a.get_order()}\")\n",
    "print(f\"Square roots of a: {a.all_roots(2)}\")\n",
    "print(f\"Cubic roots of a: {a.all_roots(3)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "172be4fa",
   "metadata": {},
   "source": [
    "### Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bdfdcd5d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exhaustive order check for M=3 OK\n",
      "Exhaustive order check for M=7 OK\n",
      "Exhaustive order check for M=11 OK\n",
      "Exhaustive order check for M=19 OK\n",
      "Exhaustive order check for M=23 OK\n",
      "Exhaustive order check for M=31 OK\n"
     ]
    }
   ],
   "source": [
    "# Exhaustive check for orders.\n",
    "def order_bf(x):\n",
    "  assert not x.is_zero()\n",
    "  ans = 1\n",
    "  a = x\n",
    "  while not a.is_one():\n",
    "    a = a * x\n",
    "    ans += 1\n",
    "  return ans\n",
    "\n",
    "for M in [3, 7, 11, 19, 23, 31]:\n",
    "  pool = ModGaussianPool.make_pool(M)\n",
    "  for x in pool.all_numbers():\n",
    "    if x.is_zero(): continue\n",
    "    assert order_bf(x) == x.get_order()\n",
    "  print(\"Exhaustive order check for M=%d OK\" % pool.M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c6402b68",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exhaustive root check for M=3 OK\n",
      "Exhaustive root check for M=7 OK\n",
      "Exhaustive root check for M=11 OK\n",
      "Exhaustive root check for M=19 OK\n",
      "Exhaustive root check for M=23 OK\n",
      "Exhaustive root check for M=31 OK\n"
     ]
    }
   ],
   "source": [
    "for M in [3,7,11,19,23, 31]:\n",
    "    pool = ModGaussianPool.make_pool(M)\n",
    "    roots_bruteforce = dict()\n",
    "    for i in range(2, 6):\n",
    "        roots_bruteforce[i] = {x:[] for x in pool.all_numbers()}\n",
    "        for x in pool.all_numbers():\n",
    "            roots_bruteforce[i][x.pow(i)].append(x)\n",
    "\n",
    "    for x in pool.all_numbers():\n",
    "        for i in range(2, 6):\n",
    "            assert set(roots_bruteforce[i][x]) == set(x.all_roots(i))\n",
    "\n",
    "    print(\"Exhaustive root check for M=%d OK\" % pool.M)\n",
    "  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "802e148a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Random root check for M=103 OK\n",
      "Random root check for M=271 OK\n",
      "Random root check for M=991 OK\n",
      "Random root check for M=1000000007 OK\n"
     ]
    }
   ],
   "source": [
    "for M in [103, 271, 991, 10**9+7]:\n",
    "  pool = ModGaussianPool.make_pool(M)\n",
    "  for _ in range(1000):\n",
    "      x = pool.random()\n",
    "      r = random.randint(2, 10)\n",
    "      roots = x.all_roots(r)\n",
    "      for root in roots:\n",
    "          assert root**r == x\n",
    "  print(\"Random root check for M=%d OK\" % pool.M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "681971b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exhaustive structure check for M=3 OK\n",
      "Exhaustive structure check for M=7 OK\n",
      "Exhaustive structure check for M=11 OK\n",
      "Exhaustive structure check for M=19 OK\n",
      "Exhaustive structure check for M=23 OK\n"
     ]
    }
   ],
   "source": [
    "for M in [3,7,11,19,23]:\n",
    "    pool = ModGaussianPool.make_pool(M)\n",
    "    non_zeros = set(x for x in pool.all_numbers() if not x.is_zero())\n",
    "    # Verifying that multiplicative group is cyclic group.\n",
    "    gen = pool.w1 * pool.w2\n",
    "    assert set(gen.pow(i) for i in range(M**2-1)) == non_zeros\n",
    "    \n",
    "    # Verifying the unique representation.\n",
    "    s1 = set()\n",
    "    for i1 in range((pool.M-1)//2):\n",
    "        for i2 in range(2*(pool.M+1)):\n",
    "            v1 = pool.w1.pow(i1)\n",
    "            v2 = pool.w2.pow(i2)\n",
    "            s1.add(v1 * v2)\n",
    "    assert s1 == non_zeros\n",
    "            \n",
    "    # Verifying number of primitve roots (aka generators).\n",
    "    pr_count = 0\n",
    "    for x in pool.all_numbers():\n",
    "        if x.is_zero(): continue\n",
    "        if x.get_order() == pool.M**2-1: pr_count += 1\n",
    "    assert pr_count/(pool.M**2-1) >= 0.25    \n",
    "    print(\"Exhaustive structure check for M=%d OK\" % pool.M)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6a897e4",
   "metadata": {},
   "source": [
    "## 9. Application: solving cubic equations in $\\mathbb{Z}_M$\n",
    "\n",
    "Consider the following problem: find all integer $x \\in [0, M)$, such that\n",
    "\n",
    "$$ax^3+bx^2+cx+d \\equiv 0 ~(\\text{mod} ~M).$$\n",
    "\n",
    "After reducing it to \"depressed cubic\" $t^3+pt+q=0$ and considering trivial case $p=0$, we can solve it using the [Cardano's formula](https://en.wikipedia.org/wiki/Cubic_equation#Cardano's_formula):\n",
    "\n",
    "$$t = C-\\frac{p}{3C}, \\text{where} ~C = \\sqrt[\\leftroot{-1}\\uproot{2}\\scriptstyle 3]{-\\frac{q}{2}+\\sqrt{\\frac{q^2}{4}+\\frac{p^3}{27}}}$$\n",
    "\n",
    "When coefficients are from $\\mathbb{R}$, this formula can be used as follows: take any of 2 square roots, then then take all 3 cubic roots, and then keep those that are real. However, it is well-known that even if all solutions are real, intermediary results might be complex, and there is no alternative way to compute this without going into $\\mathbb{R}$.\n",
    "\n",
    "The same works for $\\mathbb{Z}_M$. In order to solve the equation where coefficients and answers are $\\mathbb{Z}_M$, we will have to go into $\\mathbb{Z}_M[i]$ for intermediary calculations.\n",
    "\n",
    "To be concrete, for in $\\mathbb{Z}_{11}$, the equation $x^3+x+2$ has 3 roots: 5,7 and 10. However, when solving it using Cardano's formula, we will need to compute square root from 10, which doesn't exist in $\\mathbb{Z}_{11}$.\n",
    "\n",
    "Below is code to solve cubic equation in $\\mathbb{Z}_M$ for case $a \\ne 0$ and tests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b0550a38",
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_cubic_eqn(a: int, b: int, c: int, d: int, M: int) -> set[int]:\n",
    "    \"\"\"Solves a*x^3+b*x^2+c*x+d=0 in Z[M] using the Cardano's formula.\"\"\"\n",
    "    assert M>3  # Doesn't work for case M=3, but it can be trivially solved with brute force.\n",
    "    pool = ModGaussianPool.make_pool(M)\n",
    "    a, b, c, d = map(lambda x: pool.num(x), [a, b, c, d])\n",
    "    assert not a.is_zero()\n",
    "  \n",
    "    # Reduce this to the \"Depressed cubic\" t^3+pt+q=0.\n",
    "    p = (a*c*3-b*b)/(a*a*3)\n",
    "    q = (b*b*b*2 - a*b*c*9+a*a*d*27)/(a*a*a*27)\n",
    "\n",
    "    ts=[]\n",
    "    if p.is_zero():\n",
    "        ts = (q*-1).all_roots(3)\n",
    "    else:\n",
    "        r = ((q*q/4)+(p*p*p/27)).discrete_root(2)        \n",
    "        ts = [C - p/(C*3) for C in (q/(-2) + r).all_roots(3)]\n",
    "        \n",
    "    xs = set(t-b/(a*3) for t in ts)\n",
    "    return {x.a for x in xs if x.b==0}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "43261bba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exhaustive OK for 7\n",
      "Exhaustive OK for 11\n",
      "Random OK for 19\n",
      "Random OK for 23\n",
      "Random OK for 31\n",
      "Random OK for 43\n",
      "Random OK for 103\n",
      "Random OK for 271\n",
      "Random OK for 991\n"
     ]
    }
   ],
   "source": [
    "import itertools\n",
    "  \n",
    "def solve_cubic_bf(a: int, b: int, c: int, d: int, M: int) -> set[int]:\n",
    "    return {x for x in range(M) if (a*x*x*x+b*x*x+c*x+d)%M==0}  \n",
    "  \n",
    "for M in [7, 11]:\n",
    "    for a,b,c,d in itertools.product(range(1,M), range(M), range(M), range(M)):\n",
    "        bf_sol = solve_cubic_bf(a,b,c,d,M)\n",
    "        sol = solve_cubic_eqn(a,b,c,d,M)\n",
    "        assert sol == bf_sol\n",
    "    print(\"Exhaustive OK for %d\" % M)\n",
    "\n",
    "for M in [19, 23, 31, 43, 103, 271, 991]:\n",
    "    for _ in range(1000):\n",
    "        a=random.randint(1,M-1)\n",
    "        b=random.randint(0,M-1)\n",
    "        c=random.randint(0,M-1)\n",
    "        d=random.randint(0,M-1)\n",
    "        bf_sol = solve_cubic_bf(a,b,c,d,M)\n",
    "        sol = solve_cubic_eqn(a,b,c,d,M)\n",
    "        assert sol == bf_sol\n",
    "    print(\"Random OK for %d\" % M)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a55b43a",
   "metadata": {},
   "source": [
    "## 10. Open questions\n",
    "\n",
    "*I'm pretty sure these are answered in some book or paper, but at the moment I don't know the answers.* \n",
    "\n",
    "**1. Further field extension.**\n",
    "\n",
    "What if we wanted to be able  to compute square root for **any** element in $\\mathbb{Z}_M[i]$? For this we need to go to further extension $\\mathbb{Z}_M[i,j]$, where $j$ is a symbol defined such that $j^2 = \\omega_2$. Will this be a field? How many elements are in this field (I think $M^4$)? Will it be isomorphic to field of quaternions over $\\mathbb{Z}_M$?\n",
    "\n",
    "**2. Infinite field extension.**\n",
    "\n",
    "What if we added ALL roots of 1 of any degree? Formally, for any integer $r$ we add $r$-th root of unity if it's not already there (obviously result will have infinite size)? Is this a field? Is it described anywhere?\n",
    "\n",
    "**3. Fundamental theorem of algebra.**\n",
    "\n",
    "Is it true that any polynomial with coefficients from $\\mathbb{Z}_M$ has at least one root in $\\mathbb{Z}_M[i]$?\n",
    "\n",
    "No, this is not true starting from 3rd degree. Just take polynomial $P(x) = x^r-a$ where a is not $r$-th residue. Example: $x^3+2$ for M=7.\n",
    "\n",
    "**4. Solving cubic equation when M%4=1.** \n",
    "\n",
    "Do we really need the requirement $M \\%3= 4$ to solve cubic equation? We only need to go into $Z_M[i]$ to compute square root, do an addition and compute cubic root, and computing roots doesn't involve divisions in $\\mathbb{Z}_M[i]$. So it seems like exactly same code must work for all prime $M$s larger than 3, even if $Z_M[i]$ is not a field. Is this true?\n",
    "\n",
    "**5. Computing discrete logarithm faster than $O(\\sqrt{M})$.** \n",
    "\n",
    "Is it possible to compute discrete logarithm in an cyclic group of size $q$ in $O(poly(\\log(q))$? If not, what is the lower bound? Is GSBS algorithm the fastest we can do?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}