/* Legendre symbol */
legendresymbol(a, p) =
{
    lift(Mod(a, p)^((p-1)/2));
}

/* Tonelli-Shanks Algorithm */
tonelli(n, p) =
{
    local(q, s, z, c, r, t, m, t2, i, b);

    /* Check that n is a quadratic residue mod p */
    if(legendresymbol(n, p) != 1,
        error("n is not a quadratic residue mod p"));

    /* Step 1: Factor out powers of 2 from p - 1 */
    q = p - 1;
    s = 0;
    while(q % 2 == 0,
        q = q/2;
        s++
    );

    /* Special case when s == 1 */
    if(s == 1,
        return(lift(Mod(n, p)^((p+1)/4)))
    );

    /* Find a non-residue z */
    z = 2;
    while(z < p,
        if(legendresymbol(z, p) == p - 1, break());
        z++
    );

    /* Initialize variables */
    c = lift(Mod(z, p)^q);
    r = lift(Mod(n, p)^((q+1)/2));
    t = lift(Mod(n, p)^q);
    m = s;

    /* Main loop */
    while((t - 1) % p != 0,
        t2 = t^2 % p;
        i = 1;
        while(i < m,
            if((t2 - 1) % p == 0, break());
            t2 = t2^2 % p;
            i++
        );
        b = lift(Mod(c, p)^(2^(m - i - 1)));
        r = (r * b) % p;
        c = b^2 % p;
        t = (t * c) % p;
        m = i
    );

    return(r);
}

/* Test the functions */
print("Testing Tonelli-Shanks algorithm:");
print();

/* Test case 1 */
n = 10; p = 13;
r = tonelli(n, p);
print("n = ", n, ", p = ", p);
print("roots: ", r, ", ", p - r);
print("verification: ", r, "^2 mod ", p, " = ", r^2 % p);
print();

/* Test case 2 */
n = 56; p = 101;
r = tonelli(n, p);
print("n = ", n, ", p = ", p);
print("roots: ", r, ", ", p - r);
print("verification: ", r, "^2 mod ", p, " = ", r^2 % p);
print();

/* Test case 3 */
n = 1030; p = 10009;
r = tonelli(n, p);
print("n = ", n, ", p = ", p);
print("roots: ", r, ", ", p - r);
print("verification: ", r, "^2 mod ", p, " = ", r^2 % p);
print();


/* Test case 4 */
n = 44402; p = 100049;
r = tonelli(n, p);
print("n = ", n, ", p = ", p);
print("roots: ", r, ", ", p - r);
print("verification: ", r, "^2 mod ", p, " = ", r^2 % p);
print();


/* Test case 5 */
n = 665820697; p = 1000000009;
r = tonelli(n, p);
print("n = ", n, ", p = ", p);
print("roots: ", r, ", ", p - r);
print("verification: ", r, "^2 mod ", p, " = ", r^2 % p);
print();


/* Test case 6 */
n = 881398088036; p = 1000000000039;
r = tonelli(n, p);
print("n = ", n, ", p = ", p);
print("roots: ", r, ", ", p - r);
print("verification: ", r, "^2 mod ", p, " = ", r^2 % p);
print();


/* Test case 7 */
n = 41660815127637347468140745042827704103445750172002; p = 100000000000000000000000000000000000000000000000577;
r = tonelli(n, p);
print("n = ", n, ", p = ", p);
print("roots: ", r, ", ", p - r);
print("verification: ", r, "^2 mod ", p, " = ", r^2 % p);
print();