(***************************************************************)
(*
** File threesquare.ari
**
** File author:
**	 (C) 2004/09 Otto Forster <forster@mathematik.uni-muenchen.de>
** Date of last change:
**    2009-04-08
**
** ARIBAS code for calculating the representation of
** a natural number as a sum of 2,3 or 4 squares
** Requires ARIBAS version 1.52 or higher
**
** This code is placed under the GNU General Public Licence
*)
(***************************************************************)
(*
** Functions:
--------------------------------------------------------------
p2squaresum(p: integer): array[2];
    For a prime p = 1 mod 4, returns a vector (x,y)
    such that p = x**2 + y**2

n2square_test(n: integer): boolean;
    Tests whether n can be written as a sum of 2 squares

n2squaresum(n: integer): array[2];
    Represents n as a sum of two squares (if possible)

p4squaresum(p: integer): array[4];
    For a prime p = 3 mod 4, returns a vector (x1,x2,x3,x4)
    with p = x1**2 + x2**2 + x3**2 + x4**2

n4squaresum(n: integer): array[4];
    Represents n as a sum of four squares

n3square_test(n: integer): boolean;
    Tests whether n can be written as a sum of 3 squares

n3squaresum(n: integer): array[3];
    Represents n as a sum of three squares (if possible)

norm2(X: array): integer;
    Returns the squaresum of the coefficients of the vector X

--------------------------------------------------------------*)
(* Example session:

==> n := random(10**40).
-: 38688_68566_02366_57267_07035_96402_88984_51601

==> n3square_test(n).
-: true

==> n3squaresum(n).
-: (12870_78185_88116_67098, 39803_80022_27470_91674, 46031_17452_54690_27939
)

==> norm2(_).
-: 38688_68566_02366_57267_07035_96402_88984_51601

==> _ = n.
-: true

==> x := 23**19.
-: 7_46154_70927_59071_05619_08487

==> n3square_test(x).
-: false

==> n4squaresum(x).
-: (540_34579_84389, 180_11526_61463, 540_34579_84389, 360_23053_22926)

==> norm2(_).
-: 7_46154_70927_59071_05619_08487

==> _ = x.
-: true

****************************************************************)
(*
** The function n3squaresum() makes essential use of
** Legendre's equation  a*x**2 + b*y**2 = z**2
**
** The following functions are available:

hilbert(a,b,p: integer): integer;
    Calculates the Hilbert symbol of (a,b) with respect
    to the prime p. The value is +1, if Legendre's equation 
    a*x**2 + b*y**2 = z**2 has a non-trival solution
    in the p-adic field Q_p, and -1 otherwise

leg_test(a,b: integer): integer;
    Returns +1, if Legendre's equation has a
    non-trivial solution, and -1 otherwise

leg_construct(a: integer): integer;
    If the argument a is a postive interger /= 7 mod 8,
    constructs a number q such that the equation
    a*x**2 - q*y**2 = z**2 has a non-trivial solution.
    q is either a prime = 1 mod 4, or q = 2*p with
    p a prime = 1 mod 4

leg_solve(a,b: integer): array[3];
    Returns a non-trivial solution (x,y,z) of the equation
    a*x**2 + b*y**2 = z**2, if such a solution exists

--------------------------------------------------------------*)
(* Examples:

==> leg_test(17,19).
-: 1

==> leg_solve(17,19).
-: (1, 1, 6)

==> leg_test(17,20).
-: -1

==> hilbert(17,20,2).
-: 1

==> hilbert(17,20,5).
-: -1

==> a := random(10**30).
-: 22578_80381_15374_44713_13381_25929

==> a mod 8.
-: 1

==> q := leg_construct(a).
-: 11_28940_19057_68722_35656_69062_96449

==> (x,y,z) := leg_solve(a,-q).
-: (73328_22327_30564_68123_64718_70484_11829, 905_78192_15681_86265_79523_
45841_46255, 34710_32109_06443_05878_74113_15978_46857_42287_20979_76392)

==> a*x**2 - q*y**2.
-: 12048_06390_21562_69159_62973_69556_77987_29364_95866_23331_07004_14952_
87137_58993_31388_88460_46843_27502_13893_37664

==> z**2.
-: 12048_06390_21562_69159_62973_69556_77987_29364_95866_23331_07004_14952_
87137_58993_31388_88460_46843_27502_13893_37664

*)
(***************************************************************)
(*
** Factorization routines
*)
(***************************************************************)
(*
** Erstellt ein Array der Primfaktoren von x,
** wobei jeder Primfaktor so oft aufgezaehlt wird,
** wie es seiner Vielfachheit entspricht
** Verwendet rho_factorize und qs_factorize
**
** Wird das optionale Argument verbose mit einem
** Wert /= 0 angegeben, erfolgen Meldungen
*)
function factorlist(x: integer; verbose := 0): array;
var
    st, st1: stack;
    q, y, bound: integer;
    vec: array;
    count: integer;
begin
    x := abs(x);
    if x < 2 then
        return ();
    end;
    q := 2;
    while q := factor16(x,q) do
        stack_push(st,q);
        x := x div q;
        if verbose then
            writeln(q);
        end;
    end;
    if x < 2**32 then
        stack_push(st,x);
        if verbose then
            writeln(x);
        end;
    else
        stack_push(st1,x);
    end;
    while not stack_empty(st1) do
        x := stack_pop(st1);
        if rab_primetest(x) then
            stack_push(st,x);
            if verbose then
                writeln(x);
            end;
        else
            bound := bit_length(x);
            bound := 4*bound**2;
            if verbose then
                writeln("trying to factor ",x," with Pollard rho")
            end;
            y := rho_factorize(x,bound,0);
            count := 0;
            while (y <= 1 or y >= x) do
                if inc(count) > 2 then
                    writeln("unable to factorize ",x);
                    halt(-1);
                end;
                if verbose then
                    writeln("trying to factorize ",x,
                            " using quadratic sieve");
                end;
                y := qs_factorize(x,0);
            end;
            if verbose then
                writeln("found factor ",y);
            end;
            stack_push(st1,x div y);
            stack_push(st1,y);
        end;
    end;
    vec := stack2array(st);
    return sort(vec);
end;
(*---------------------------------------------------------*)
(*
** Erstellt ein Array aus den Primfaktoren von x, wobei jeder
** Primfaktor nur einmal aufgezaehlt wird.
*)
function primefactors(x: integer): array;
var
    i,p,p0: integer;
    vec: array;
    st: stack;
begin
    vec := factorlist(x);
    p0 := 0;
    for i := 0 to length(vec)-1 do
        p := vec[i];
        if p /= p0 then
            stack_push(st,p);
            p0 := p;
        end;
    end;
    return stack2array(st);
end;
(*---------------------------------------------------------*)
(*
** Returns the union of the primefactors of x and y
*)
function primefactors2(x,y: integer): array;
var
    st: stack;
    vec: array;
    i,p,p0,z: integer;
begin
    z := gcd(x,y);
    vec := factorlist(z);
    stack_arraypush(st,vec);
    vec := factorlist(x div z);
    stack_arraypush(st,vec);
    vec := factorlist(y div z);
    stack_arraypush(st,vec);
    vec := stack2array(st);
    sort(vec);
    p0 := 0;
    for i := 0 to length(vec)-1 do
        p := vec[i];
        if p /= p0 then
            stack_push(st,p);
            p0 := p;
        end;
    end;
    return stack2array(st);
end;
(*---------------------------------------------------------*)
(*
** An integer x is decomposed as
**      x = z**2 * y
** where y is squarefree.
** Return value z
*)
function squareext(x: integer): integer;
var
    vec: array;
    z: integer;
begin
    vec := factorlist(x);
    z := square_ext(vec);
    return z;
end;
(*---------------------------------------------------------*)
(*
** The integer x = product(pvec) is decomposed as
**      x = z**2 * y
** where y is squarefree.
** Return value z,
** the prime decomposition of y is stored
** in the variable parameter pvec
*)
function square_ext(var pvec: array): integer;
var
    vec: array;
    st: stack;
    k, len, q, z, m: integer;
begin
    len := length(pvec);
    z := 1;
    while k < len do
        q := pvec[k];
        m := 1;
        while (k+1 < len) and (pvec[k+1] = q) do
            inc(m); inc(k);
        end;
        if odd(m) then
            stack_push(st,q);
        end;
        if m > 1 then
            z := z * q**(m div 2);
        end;
        inc(k);
    end;
    pvec := stack2array(st);
    return z;
end;
(***********************************************************)
(*
** Square roots modulo a composite number
** Uses ARIBAS builtin function gfp_sqrt(p,a)
** for the square root modulo an odd prime p
** and the Chinese Remainder Theorem
*)
(***********************************************************)
(*
** Inverse map for the Chinese Remainder Theorem:
** Constructs an integer z such that
**  z = X[i] mod M[i] for all i
** The coefficients of the array M are supposed to be
** pairwise coprime
*)
function chin_inv(X, M: array): integer;
var
    m, m1, i, z: integer;
begin
    m := product(M);
    z := 0;
    for i := 0 to length(M)-1 do
        m1 := m div M[i];
        z := z + X[i] * m1 * mod_inverse(m1,M[i]);
    end;
    return z mod m;
end;
(*---------------------------------------------------------*)
(*
** Calculates a square root of x modulo m
** where m is supposed to be squarefree
** x must be a quadratic residue modulo all odd
** prime factors of m
*)
function sqrtmodm(x,m: integer): integer;
var
    M: array;
begin
    M := primefactors(m);
    return sqrt_modM(x,M);
end;
(*---------------------------------------------------------*)
(*
** Calculates a square root of x mod prod(M)
** It is supposed that x /= 0 and that
** the coefficients of M are pairwise distinct primes
** and that the square root exists
*)
function sqrt_modM(x: integer; M: array): integer;
var
    X: array;
    i,len,y,z,p: integer;
begin
    len := length(M);
    X := alloc(array,len);
    for i := 0 to len-1 do
        p := M[i];
        z := x mod p;
        if p = 2 then
            y := z;
        else
            if jacobi(z,p) = -1 then
                writeln("error in sqrt_modM: not solvable");
                halt(-2);
            end;
            y := gfp_sqrt(p,z);
        end;
        X[i] := y;
    end;
    return chin_inv(X,M);
end;
(***********************************************************)
(*
** Arithmetic in the ring of Gaussian integers
** and the representation of integers as the
** sum of two squares
*)
(***********************************************************)
(*
** Multiplikation im Ring der ganzen Gauss'schen Zahlen
*)
function gauss_mult(x,y: array[2]): array[2];
begin
    return (x[0]*y[0] - x[1]*y[1], x[1]*y[0] + x[0]*y[1]);
end;
(*-----------------------------------------------------------*)
(*
** Seien x,y Elemente des Rings der ganzen Gauss'schen Zahlen,
** y /= 0
** Es wird eine ganze Gauss'sche Zahl r mit Norm(r) < Norm(y)
** bestimmt, so dass
**      x = q*y + r
** mit einer ganzen Gauss'schen Zahl q
*)
function gauss_mod(x,y: array[2]): array[2];
var
    z0,z1,q0,q1,u0,u1,N,Nhalf: integer;
begin
    N := y[0]*y[0] + y[1]*y[1];
    Nhalf := N div 2;
    z0 := x[0]*y[0] + x[1]*y[1];
    z1 := x[1]*y[0] - x[0]*y[1];
    q0 := z0 div N; q1 := z1 div N;
    if (z0 mod N) > Nhalf then inc(q0) end;
    if (z1 mod N) > Nhalf then inc(q1) end;
    u0 := y[0]*q0 - y[1]*q1;
    u1 := y[0]*q1 + y[1]*q0;
    return (x[0]-u0, x[1]-u1);
end;
(*-----------------------------------------------------------*)
(*
** Groesster gemeinsamer Teiler
** im Ring der ganzen Gauss'schen Zahlen
*)
function gauss_gcd(x,y: array[2]): array[2];
var
    temp: array[2];
begin
    while y /= (0,0) do
        temp := y;
        y := gauss_mod(x,y)
        x := temp;
    end;
    return x;
end;
(*-----------------------------------------------------------*)
(*
** The argument must be a prime p = 1 mod 4
** The function returns a pair (x,y) with p = x**2 + y**2
*)
function p2squaresum(p: integer): array[2];
var
    x: integer;
    xi: array[2];
begin
    x := gfp_sqrt(p,-1);
    xi := gauss_gcd((x,1),(p,0));
    return (abs(xi[0]),abs(xi[1]));
end;
(*-----------------------------------------------------------*)
(*
** Returns a pair (a,b) such that x = a**2 + b**2
** It is supposed that such a representation exists
*)
function n2squaresum(x: integer): array[2];
var
    pvec: array;
begin
    if x = 0 then
        return (0,0);
    end;
    pvec := factorlist(x);
    return n2square_sum(pvec);
end;
(*-----------------------------------------------------------*)
(*
** Returns a pair (a,b) such that x = a**2 + b**2,
** where x = prod(pvec)
** It is supposed that such a representation exists
*)
function n2square_sum(pvec: array): array[2];
var
    xy, uv: array;
    t, k, len, q: integer;
begin
    t := square_ext(pvec);
    len := length(pvec);
    xy := (t,0);
    for k := 0 to len-1 do
        q := pvec[k];
        if q = 2 then
            uv := (1,1);
        elsif q mod 4 = 3 then
            writeln("error n2square_sum: not a sum of two squares");
            return ();
        else
            uv := p2squaresum(q);
        end;
        xy := gauss_mult(xy,uv);
    end;
    return xy;
end;
(*-------------------------------------------------------*)
(*
** Tests whether n can be written as a sum of 2 squares
*)
function n2square_test(n: integer): boolean;
var
    pvec: array;
    t, k, p: integer;
begin
    pvec := factorlist(n);
    t := square_ext(pvec);
    for k := 0 to length(pvec)-1 do
        p := pvec[k];
        if (p > 2) and (p mod 4 = 3) then
            return false;
        end;
    end;
    return true;
end;
(*************************************************************)
(*
** Arithmetic in the ring of integer quaternions
** and representation of a natural number as 
** a sum of four squares (Theorem of Lagrange)
*)
(*************************************************************)
(*
** Conjugation of integer quaternions
*)
function h_conj(x: array[4]): array[4];
begin
    return (x[0],-x[1],-x[2],-x[3]);
end;
(*-------------------------------------------------------*)
(*
** Product of two integer quaternions x,y
*)
function h_mult(x,y: array[4]): array[4];
var
    z0,z1,z2,z3: integer;
begin
    z0 := x[0]*y[0] - x[1]*y[1] - x[2]*y[2] - x[3]*y[3];
    z1 := x[0]*y[1] + x[1]*y[0] + x[2]*y[3] - x[3]*y[2];
    z2 := x[0]*y[2] - x[1]*y[3] + x[2]*y[0] + x[3]*y[1];
    z3 := x[0]*y[3] + x[1]*y[2] - x[2]*y[1] + x[3]*y[0];
    return (z0,z1,z2,z3);
end;
(*-------------------------------------------------------*)
(*
** Die Argumente x,y sind ganzzahlige Quaternionen, y /= 0.
** Es wird Division mit Rest durchgefuehrt
**      x = q*y + r
** Dabei ist q eine Hurwitz-Quaternion, der Rest r eine ganzzahlige
** Quaternion mit Norm(r) < Norm(y).
** Der Rest r wird zurueckgegeben.
*)
function h_mod(x,y: array[4]): array[4];
var
    N, Nhalf, i: integer;
    z, q, r, u: array[4];
begin
    N := norm2(y);
    Nhalf := N div 2;
    z := h_mult(x,h_conj(y));
    q := z div N;
    r := z mod N;
    if even(N) and r = (Nhalf,Nhalf,Nhalf,Nhalf) then
        return (0,0,0,0);
    end;
    for i := 0 to 3 do
        if r[i] > Nhalf then inc(q[i]) end;
    end;
    u := h_mult(q,y);
    return x-u;
end;
(*-------------------------------------------------------*)
(*
** Groesster gemeinsamer Teiler fuer ganzzahlige Quaternionen.
** Dabei wird im Ring der Hurwitz-Quaternionen gearbeitet,
** das Resultat ist aber immer eine ganzzahlige Quaternion.
*)
function h_gcd(x,y: array[4]): array[4];
var
    temp: array[4];
begin
    while y /= (0,0,0,0) do
        temp := y;
        y := h_mod(x,y);
        x := temp;
    end;
    return x;
end;
(*-------------------------------------------------------*)
(*
** Das Argument p muss eine Primzahl = 3 mod 4 sein
** Die Funktion berechnet (x,y), so dass
**      x*x + y*y = - 1 mod p
** Wird von der Funktion p4squaresum benuetzt.
*)
function fp_m1sqsum(p: integer): array[2];
var
    x, y: integer;
begin
    for x := 1 to p do
        if jacobi(-1-x*x,p) = 1 then break end;
    end;
    y := (-1-x*x)**((p+1) div 4) mod p;
    return (x,y);
end;
(*-------------------------------------------------------*)
(*
** Das Argument p muss eine Primzahl = 3 mod 4 sein.
** Das Resultat ist ein Vektor aus 4 natuerlichen Zahlen,
** deren Quadratsumme gleich p ist.
*)
function p4squaresum(p: integer): array[4];
var
    x: array[2];
    z: array[4];
begin
    if p mod 4 /= 3 then
        writeln("p4squaresum: argument must be a prime = 3 mod 4");
        halt(0);
    end;
    x := fp_m1sqsum(p);
    z := h_gcd((p,0,0,0),(1,x[0],x[1],0));
    return (abs(z[0]),abs(z[1]),abs(z[2]),abs(z[3]));
end;
(*---------------------------------------------------------*)
(*
** The argument n must be a positive integer.
** The function returns a vector of 4 integers
** whose square sum equals n
*)
function n4squaresum(n: integer): array[4];
var
    pvec: array;
begin
    if n = 0 then
        return (0,0,0,0);
    end;
    pvec := factorlist(n);
    return n4square_sum(pvec);
end;
(*---------------------------------------------------------*)
function n4square_sum(pvec: array): array[4];
var
    X, U: array[4];
    t, u, v, q, k, len: integer;
begin
    t := square_ext(pvec);
    len := length(pvec);
    X := (t,0,0,0);
    for k := 0 to len-1 do
        q := pvec[k];
        if q = 2 then
            U := (1,1,0,0);
        elsif q mod 4 = 1 then
            (u,v) := p2squaresum(q);
            U := (u,v,0,0);
        else  (* q mod 4 = 3 *)
            U := p4squaresum(q);
        end;
        X := h_mult(X,U);
    end;
    return (abs(X[0]), abs(X[1]), abs(X[2]), abs(X[3]));
end;
(*---------------------------------------------------------*)
(*
** Hypothesis: The integer a is a square modulo b,
** where b is a squarefree integer.
** Constructs a non-trivial solution (x,z) of the congruence
**    x**2 - a*z**2 = 0 mod b
** such that
**    x**2 + abs(a)*z**2
** is minimal
*)
function minsqrt(a,b: integer): array[2];
var
    bvec: array;
begin
    bvec := primefactors(b);
    return min_sqrt(a,bvec);
end;
(*---------------------------------------------------------*)
(*
** Like minsqrt(a,b), with b = product(bvec)
** bvec must be a vector of pairwise different primes
*)
function min_sqrt(a: integer; bvec: array): array[2];
var
    len, w, i, p, y, z, a0, b: integer;
    wvec: array;
    X: array[2];
begin
    len := length(bvec);
    wvec := alloc(array,len);
    for i := 0 to len-1 do
        p := bvec[i];
        z := a mod p;
        if p = 2 then
            y := z;
        elsif jacobi(z,p) = -1 then
            writeln("error in minsqrt: not solvable",(z,p));
            halt(0);
        else
            y := gfp_sqrt(p,z);
        end;
        wvec[i] := y;
    end;
    w := chin_inv(wvec,bvec);
    a0 := abs(a);
    b := product(bvec);
    X := gauss_redA((b,0),(w,1),a0);
    return X;
end;
(*---------------------------------------------------------*)
(*
** Gauss lattice reduction
** Constucts a non-zero vector of minimal length in the
** 2-dimensional lattice generated by the vectors X,Y
** with respect to the norm X[0]**2 + a0*X[1]**2
** a0 must be positive
*)
function gauss_redA(X,Y: array[2]; a0: integer): array[2];
var
    T: array[2];
    Tnorm, Ynorm, XYscal, lambda: integer;
begin
    Tnorm := X[0]**2 + a0*X[1]**2;
    Ynorm := Y[0]**2 + a0*Y[1]**2;
    if Tnorm < Ynorm then
        T := X;
        X := Y;
        Y := T;
        Ynorm := Tnorm;
    end;
    while true do
        XYscal := X[0]*Y[0] + a0*X[1]*Y[1];
        lambda := (XYscal + (Ynorm div 2)) div Ynorm;
        T := X - lambda*Y;
        Tnorm := T[0]**2 + a0*T[1]**2;
        if Tnorm >= Ynorm then
            return Y;
        end;
        X := Y;
        Y := T;
        Ynorm := Tnorm;
    end;
end;
(*---------------------------------------------------------*)
(*
** Calculates the Hilbert symbol (a,b/p)
** where a and b are non-zero integers
** This symbol has value 1 if the Legendre equation
**   a*x**2 + b*y**2 = z**2
** has a non-trivial solution in the p-adic field Q_p
** and value -1 otherwise
*)
function hilbert(a,b,p: integer): integer;
var
    m1,m2,u,v,u8,v8,res: integer;
begin
    if p < 2 then
        writeln("hilbert: prime number expected: ",p);
        halt(0);
    elsif a = 0 or b = 0 then
        writeln("hilbert: a and b must be nonzero");
        halt(0);
    end;
    m1 := 0;
    u := a;
    while u mod p = 0 do
        u := u div p;
        inc(m1);
    end;
    m1 := m1 mod 2;
    m2 := 0;
    v := b;
    while v mod p = 0 do
        v := v div p;
        inc(m2);
    end;
    m2 := m2 mod 2;
    if p = 2 then
        if (u mod 4 = 3) and (v mod 4 = 3) then
            res := -1;
        else
            res := 1;
        end;
        if m2 then
            u8 := u mod 8;
            if u8 = 3 or u8 = 5 then
                res := -res;
            end;
        end;
        if m1 then
            v8 := v mod 8;
            if v8 = 3 or v8 = 5 then
                res := -res;
            end;
        end;
    else (* p >= 3 *)
        if (m1 = 1) and (m2 = 1) and (p mod 4 = 3) then
            res := -1;
        else
            res := 1;
        end;
        if m2 then
            res := res*jacobi(u,p);
        end;
        if m1 then
            res := res*jacobi(v,p);
        end;
    end;
    return res;
end;
(*---------------------------------------------------------*)
(*
** The arguments a,b must be non-zero integers.
** Decides whether the Legendre equation
**  a*x**2 + b*y**2 = z**2
** has a non-trivial solution in the ring Z
** If so, returns 1, if not, -1
*)
function leg_test(a,b: integer): integer;
var
    vec,vec1,vec2: array;
    k,k0,p,len,c: integer;
begin
    if a = 0 or b = 0 then
        return 0;
    end;
    if a < 0 and b < 0 then
        return -1;
    end;
    if hilbert(a,b,2) /= 1 then
        return -1;
    end;
    vec := primefactors2(a,b);
    len := length(vec);
    if len = 0 then
        return 1;
    end;
    if vec[0] = 2 then
        k0 := 1;
    else
        k0 := 0;
    end;
    for k := k0 to len-1 do
        p := vec[k];
        if hilbert(a,b,p) /= 1 then
            return -1;
        end;
    end;
    return 1;
end;
(*-----------------------------------------------------------*)
(*
** The argument a must be a squarefree positive
** integer /= 7 mod 8
** The function returns an integer q such that the
** Legendre equation
**      a*x**2 - q*y**2 = z**2
** has a non-trivial solution.
** If a = 1,2,5,6 mod 8, then q is a prime = 1 mod 4;
** if a = 3 mod 8, then q = 2*p, p = 1 mod 4 prime
*)
function leg_construct(a: integer): integer;
var
    a8, p: integer
begin
    a8 := a mod 8;
    if a < 0 or a8 = 0 or a8 = 4 or a8 = 7 then
        write("leg_construct: argument must be ");
        writeln("a squarefree positive integer /= 7 mod 8");
        halt(0);
    end;
    if a8 = 1 or a8 = 5 then
        p := (2*a - 1);
    elsif a8 = 2 or a8 = 6 then
        p := a - 1;
    else (* a8 = 3 *)
        p := a - 2;
    end;
    while not rab_primetest(p) do
        p := p + 4*a;
    end;
    if a8 /= 3 then
        return p;
    else
        return 2*p;
    end;
end;
(*---------------------------------------------------------*)
(*
** Constructs a solution (x,y,z) of the equation
**    a*x**2 + b*y**2 = z**2
** Supposes that a solution exists, i.e. leg_test(a,b) = 1
*)
function leg_solve(a,b: integer): array[3];
var
    avec, bvec: array;
begin
    avec := factorlist(a);
    bvec := factorlist(b);
    return leg_solve1(a,b,avec,bvec);
end;
(*-----------------------------------------------------------*)
(*
** Constructs a solution (x,y,z) of the equation
**    a*x**2 + b*y**2 = z**2
** Supposes that a solution exists, i.e. leg_test(a,b) = 1
** avec resp. bvec must be the prime decompositions
** of a and b, ignoring signs
*)
function leg_solve1(a,b: integer; avec,bvec: array): array[3];
var
    a2, b2, c2, a1, b1, u, v, w, d: integer;
begin
    a2 := square_ext(avec);
    a1 := a div a2**2;
    b2 := square_ext(bvec);
    b1 := b div b2**2;
    (u,v,w) := leg_solvaux1(a1,b1,avec,bvec);
    d := gcd(a2,b2);
    a2 := a2 div d;
    c2 := a2*b2;
    b2 := b2 div d;
    return (abs(u*b2), abs(v*a2), abs(w*c2));
end;
(*---------------------------------------------------------*)
(*
** Auxiliary function for leg_solve1
*)
function leg_solvaux1(a,b: integer; avec,bvec: array): array[3];
var
    t, s, u, v, w, x, y, z, b1, d: integer;
    swap: boolean;
begin
    if a = 1 then
        return (1,0,1);
    elsif b = 1 then
        return (0,1,1);
    elsif abs(a) = abs(b) then
        if a = -b then
            return (1,-1,0);
        elsif a > 0 and b > 0 then
            (u,v) := n2square_sum(avec);
            return (u,v,a);
        else (* a < 0 and b < 0 *)
            writeln("leg_solve: error, not solvable");
            halt(-1);
        end;
    end;
    if abs(a) > abs(b) then
        (a,b) := (b,a);
        (avec,bvec) := (bvec,avec);
        swap := true;
    else
        swap := false;
    end;
    (t,s) := min_sqrt(a,bvec);
    b1 := (t**2 - a*s**2) div b;
    bvec := factorlist(b1);
    (u,v,w) := leg_solve1(a,b1,avec,bvec);
    x := t*u + s*w;
    y := b1*v;
    z := t*w + a*s*u;
    if swap then
        (x,y) := (y,x);
    end;
    d := gcd(x,y,z);
    return (x,y,z) div d;
end;
(*************************************************************)
(*
** Davenport-Cassels Reduction
*)
(*************************************************************)
(*
** Calculates the next integer vector to X/t
*)
function nextintvec(X: array; t: integer): array;
var
    x, z, thalf, k: integer;
begin
    thalf := t div 2;
    for k := 0 to length(X)-1 do
        x := X[k];
        z := x div t;
        if (x mod t) > thalf then
            inc(z);
        end;
        X[k] := z;
    end;
    return X;
end;
(*---------------------------------------------------------*)
(*
** Calculates the dot product of the integer vectors X,Y
*)
function dotprod(X,Y: array): integer;
var
    k, t, len: integer;
begin
    len := length(X);
    if len /= length(Y) then
        writeln("error in dotprod: vectors must have same dim");
        halt(0);
    end;
    t := 0;
    for k := 0 to len-1 do
        t := t + X[k]*Y[k];
    end;
    return t;
end;
(*-------------------------------------------------------*)
(*
** Calculates the square sum of the coefficients
** of the integer vector X (of arbitrary length)
*)
function norm2(X: array): integer;
var
    k, t: integer;
begin
    t := 0;
    for k := 0 to length(X)-1 do
        t := t + X[k]**2;
    end;
    return t;
end;
(*---------------------------------------------------------*)
(*
** Reduction theorem of Davenport-Cassels
** Supposes that X is an integer vector such that
**   norm2(X) = n*t**2
** for some positive integer n.
** Returns a vector Z with 
**   norm2(Z) = n
*)
function davp_reduce(X: array; t: integer): array;
var
    Y: array;
    nn: integer;
    alfa, beta, len: integer;
begin
    len := length(X);
    if len > 3 then
        writeln("davenport: dimension must be <= 3");
        halt(0);
    end;
    nn := norm2(X);
    if nn mod t**2 /= 0 then
        writeln("davenport: norm of X must be multiple of t**2");
        halt(-1);
    end;
    t := abs(t);
    nn := nn div t**2;
    while t > 1 do
        Y := nextintvec(X,t);
        alfa := norm2(Y) - nn;
        if alfa = 0 then
            return Y;
        end;
        beta := 2*(nn*t - dotprod(X,Y));
        X := alfa*X + beta*Y;
        t := alfa*t + beta;
    end;
    return X;
end;
(*************************************************************)
(*
** Representation of a natural number as a sum
** of three squares (Theorem of Gauss)
*)
(*************************************************************)
(*
** Tests whether the positive integer n can be
** represented as a sum of three squares
*)
function n3square_test(n: integer): boolean;
begin
    while n mod 4 = 0 do
        n := n div 4;
    end;
    return (n mod 8 /= 7);
end;
(*---------------------------------------------------------*)
(*
** Returns an integer vector (x,y,z) such that 
**   n = x**2 + y**2 + z**2
** Supposes that such a representation exists,
** i.e. n3square_test(n) = true
*)
function n3squaresum(n: integer): array[3];
var
    avec, bvec: array;
    X: array[3];
    t,a,p,u,v,x,y,z: integer;
begin
    if n = 0 then
        return (0,0,0);
    end;
    avec := factorlist(n);
    t := square_ext(avec);
    a := product(avec);
    p := leg_construct(a);
    if odd(p) then
        (u,v) := p2squaresum(p);
        bvec := {p};
    else
        (u,v) := p2squaresum(p div 2);
        (u,v) := (u+v,u-v);
        bvec := (2, p div 2);
    end;
    (x,y,z) := leg_solve1(a,-p,avec,bvec);
    X := (u*y, v*y, z);
    X := davp_reduce(X,x);
    X := (abs(X[0]), abs(X[1]), abs(X[2]));
    return t*X;
end;
(*********************************************************)