Miller Rabin Primality Test in Python (with latex trace)

Miller_Rabin_Primality_Test.py

import gmpy2
from gmpy2 import powmod
import random

def sr(n):
    s = 0
    n = n - 1
    while n%2 == 0:
        n = n // 2
        s += 1
    return (s, n)

def miller_rabin(n):
    a = random.randint(2, n - 2)
    s,r = sr(n)
    y = powmod(a,r, n)
    # print(f"s={s} r={r} a={a} y = {y}")
    print(f"""
\\begin{{split}}
n &= {n} \\\\
n-1 &= {n-1} = 2^{s} \\times {r} = 2^sr \\\\ & \\Rightarrow r={r}, \\ s={s} \\\\
&2 \\le a \\le 49:  a\\leftarrow {a} \\quad \\text{{random}} \\\\
y &= a^r \\bmod n \\\\ &=  {a}^{{{r}}} \\bmod {n} \\\\ &= {y} \\\\
    """)
    print(f"& y = {y} \\ne 1 ({y != 1}) \\quad y = {y} \\ne {n-1} ({y != n-1}) \\\\")
    if y != 1 and y != n-1:
        j = 1
        print("&j \\leftarrow 1 \\\\")
        print(f"&j = {j} \\le {s - 1} = s - 1 \\text{{({j <= s-1})}} \\quad  y = {y} \\ne {n-1} = n-1 \\text{{({y != n-1})}} \\\\")
        while j <= s-1 and y != n-1:
            y_prev=y
            y = powmod(y,2,n)
            print(f"y & \leftarrow y^2 \\bmod n = {y_prev}^2 \\bmod {n} = {y} \\\\")
            if y == 1:
                print("& composite \\\\")
                print("\end{split}")
                return
            j += 1
            print(f"&j \\leftarrow {j} \\\\")
            print(f"&j = {j} \\le {s - 1} = s - 1 \\text{{({j <= s-1})}} \\quad  y = {y} \\ne {n-1} = n-1 \\text{{({y != n-1})}} \\\\")

        if y != n - 1:
            print("& composite \\\\")
            print("\end{split}")
            return
    
    print("& prime \\\\")
    print("\end{split}")


miller_rabin(51)
miller_rabin(53)