with(LinearAlgebra[Modular]):
interface (printbytes=false);
pp:=65521;

to_magma:=proc(fic,p,l)
local fd,i;
    printf("save in %s\n",fic):
    fd:=fopen(fic,WRITE):
    fprintf(fd,cat ("K:=GF(%a);\nA:=Matrix(K,%a);\nA:=Transpose(A);",
                    "time KS:=NullSpace(A);\n"),p,l):
    fclose(fd):
end:

# Returns a nonsquare in F_p for an odd p
p_nonsquare:= proc (p)
local roll, t, halfp;
    if (p <= 2 or not isprime (p)) then return ``;
    end if:
    roll := rand (1..p-1);
    halfp:= (p-1)/2;
    do
        t:= roll ();
        if (t &^ halfp mod p <> 1) # t is not a square
        then return t; 
        end if: 
    end do:
end proc:

# Returns the symmetric part of M
# If p = 2, the diagonal will be zero
Sigma:= proc (p, M)
    if (LinearAlgebra[RowDimension] (M) <> LinearAlgebra[ColumnDimension] (M))
    then error "Not a square matrix";
    end if:
    return M+Transpose (p, M) mod p;
end proc:

# Returns the diagonal of M
Delta:= proc (M)
local n;
    n:= LinearAlgebra[RowDimension] (M);
    if (n <> LinearAlgebra[ColumnDimension] (M))
    then error "Not a square matrix";
    end if:
    return [seq (M(i,i), i=1..n)];
end proc:

# Returns the upper triangular matrix of the quadratic form in char 2
TriangRep:= proc (M)
local N,i,j,n,d;
    n:= LinearAlgebra[RowDimension] (M);
    if (n <> LinearAlgebra[ColumnDimension] (M))
    then error "Not a square matrix";
    end if:
    N:= Copy (2, M);
    for i from 1 to n-1 do;
        for j from i+1 to n do;
            N(i,j):= N(i,j)+N(j,i) mod 2;
            N(j,i):= 0;
        end do:
    end do:
    return N;
end proc:


# Creation of two sets of m matrices of size n over F_p, with odd p
# If are_equiv = 1, they are equivalent over F_p
# Otherwise if are_equiv = 2, they are equivalent over F_{p^2} but not F_p
# Otherwise they should not be equivalent over any extension of F_p
# If are_affine > 0, the last row of the matrices is 0,...,0,1
IP1S_rand_instance:= proc (p, m, n, are_equiv:= 1, are_affine:= 0)
local A, tA, H, K, M, t;
    do 
        A:= Random (p, n, n, integer[]);
        if (are_affine > 0) # Affine case, last row is 0,..,0,1
        then A(n,1..n-1):= 0; A(n,n):= 1;
        end if:
        if (Determinant (p, A) <> 0) then break;
        end if:
    end do:

    tA:= Transpose(p, A);
    H:= [seq (Sigma (p, Random (p, n, n, integer[])), i=1..m)];

    if (are_equiv = 1) # Equivalent sets over F_p
    then K:= map (x->tA.x.A mod p, H);
    elif (are_equiv = 2) # Equivalent sets over F_{p^2}
    then
        t:= p_nonsquare (p);
        K:= map (x->t * tA.x.A mod p, H);
    else # Two random sets, should not be equivalent
        K:= [seq (Sigma (p, Random (p, n, n, integer[])), i=1..m)];
    end if:
    return H, K, A;
end proc:

# Creation of two sets of m matrices of even size n over F_2
# If are_equiv = 1, they are equivalent over F_2
# Otherwise if are_equiv = 2, they are equivalent over F_4 but not F_2
# Otherwise they should not be equivalent over any extension of F_p
# If are_affine > 0, the last row of the matrices is 0,...,0,1
IP1S_rand_instance_char2:= proc (m, n, are_equiv:= 1, are_affine:= 0)
local A, tA, H, K, M, t;
    if (n mod 2 <> 0) then error "Not of even dimension";
    end if:
    do 
        A:= Random (2, n, n, integer[]);
        if (are_affine > 0) # Affine case, last row is 0,..,0,1
        then A(n,1..n-1):= 0; A(n,n):= 1;
        end if:
        if (Determinant (2, A) <> 0) then break;
        end if:
    end do:
    H:= [seq (TriangRep (Random (2, n, n, integer[])), i=1..m)];
    tA:= Transpose(2, A);
    if (are_equiv = 1) # Equivalent sets over F_p
    then 
        K:= map (x->TriangRep(tA.x.A mod 2), H);
    #TODO
    # elif (are_equiv = 2)
    # then
    else # Two random sets, should not be equivalent
        K:= [seq (TriangRep (Random (2, n, n, integer[])), i=1..m)];
    end if:
    return H, K, A;
end proc:


# Returns  1, ``, `` if they are regular instances
# returns  0, ``, `` if they do not have the same ranks
# returns -1, P, Q if both are singular 
IP1S_is_regular:= proc (p, m, n, H, K)
local HH, KK, r, s;
    HH:= <seq (H[i], i=1..m)>;
    KK:= <seq (K[i], i=1..m)>;
    r := Rank (p, HH);
    s := Rank (p, KK);
    if (r = n and s = n) then return 1, ``, ``;
    elif (r <> s) then return 0, ``, ``;
    else #TODO should be P and Q
        return -1, ``, ``
    end if
end proc:

# Normalization of the sets in odd characteristic p
# If normalization is ok,
#     returns 1, [Id, H1^(-1)H2,H1^(-1)H3], [Id, K1^(-1)K2,K1^(-1)K3]
# else if a matrix is invertible and not the other,
#     returns 0, ``, ``,
# else it is the irregular case
#     returns -1, ``, ``
IP1S_normalize:= proc (p, m, n, H, K)
local H1, K1, det1, det2, invH1, invK1, roll, count, c, mm;
    H1 := H[1];
    K1 := K[1];
    roll:=rand (0..1);
    count:= 0;
    do
        det1 := Determinant (p, H1);
        det2 := Determinant (p, K1);
        if (det1 <> 0)
        then 
            if (det2 <> 0) then break;
            else return 0, ``, `` # Not equivalent
            end if:
        elif (det2 <> 0) then return 0, ``, `` # Not equivalent
        end if:
        if (count > n^3) then return -1, ``, `` # Too many tries, irregular?
        end if:
        # change H1 to be regular
        c :=[1,seq (roll (), i=2..m)];
        H1:= add (c[i].H[i], i=1..m) mod p;
        K1:= add (c[i].K[i], i=1..m) mod p;
        count:= count+1;
    end do:
    invH1:= Inverse (p, H1);
    invK1:= Inverse (p, K1);
    mm:= min (m, 3);
    return 1,
    [Identity (p, n, integer[]), seq (invH1.H[i] mod p, i=2..mm)],
    [Identity (p, n, integer[]), seq (invK1.K[i] mod p, i=2..mm)];
end proc:

# Normalization of the sets in characteristic 2
# If normalization is ok,
#     returns 1, [Id, Sigma(H1)^(-1)Sigma(H2),Sigma(H1)^(-1)Sigma(H3)],
#                [Id, Sigma(K1)^(-1)Sigma(K2),Sigma(K1)^(-1)K3]
# else if a matrix is invertible and not the other,
#     returns 0, ``, ``,
# else it is the irregular case
#     returns -1, ``, ``
IP1S_normalize_char2:= proc (m, n, H, K)
local SH, SK, SH1, SK1, det1, det2, invSH1, invSK1, roll, count, c, mm;
    SH := map (x->Sigma(2,x), H);
    SK := map (x->Sigma(2,x), K);
    SH1 := SH[1];
    SK1 := SK[1];
    roll:=rand (0..1);
    count:= 0;
    do
        det1 := Determinant (2, SH1);
        det2 := Determinant (2, SK1);
        if (det1 <> 0)
        then 
            if (det2 <> 0) then break;
            else return 0, ``, `` # Not equivalent
            end if:
        elif (det2 <> 0) then return 0, ``, `` # Not equivalent
        end if:
        if (count > n^3) then return -1, ``, `` # Too many tries, irregular?
        end if:
        # change H1 to be regular
        c :=[1,seq (roll (), i=2..m)];
        SH1:= add (c[i].SH[i], i=1..m) mod 2;
        SK1:= add (c[i].SK[i], i=1..m) mod 2;
        count:= count+1;
    end do:
    invSH1:= Inverse (2, SH1);
    invSK1:= Inverse (2, SK1);
    if (n >= 8) then mm:= 3;
    elif (n = 6) then mm:= 4;
    elif (n = 4) then mm:= 6;
    end if:
    mm:= min (m, mm);
    return 1,
    [Identity (2, n, integer[]), seq (invSH1.SH[i] mod 2, i=2..mm)],
    [Identity (2, n, integer[]), seq (invSK1.SK[i] mod 2, i=2..mm)];
end proc:



# Deterministic version
# Returns 1 if X is indeed a solution of XT.H.X = K
# Returns 0 otherwise
IP1S_verif:= proc (p, m, n, H, K, X)
local i, XT;
    XT:= Transpose (p,X);
    for i from 1 to m do           
        if (not LinearAlgebra[Equal] (XT.H[i].X mod p, K[i] mod p))
        then return 0;
        end if:
    end do:
    return 1;
end proc:

# Probabilistic version
# Returns 1 if X is indeed a solution of XT.H.X = K with good probability
# Returns 0 otherwise
IP1S_verifProba:= proc (p, m, n, H, K, X)
local i, v, Xv, XT;
    v := Random (p, n, 1, integer[]);
    Xv:= X.v mod p;
    XT:= Transpose (p, X);
    for i from 1 to m do           
        if (not LinearAlgebra[Equal] (XT.(H[i].Xv) mod p, K[i].v mod p))
        then return 0;
        end if:
    end do:
    return 1;
end proc:

# Deterministic version
# Returns 1 if X is indeed a solution of XT.Sigma(H).X = Sigma(K) and
# of Delta(XT.H.X)=Delta(K)
# Returns 0 otherwise
IP1S_verif_char2:= proc (m, n, H, K, X)
local i, XT;
    XT:= Transpose (2,X);
    for i from 1 to m do           
        if (not LinearAlgebra[Equal] (XT.Sigma(2,H[i]).X mod 2,
                                      Sigma(2,K[i]) mod 2))
        then return 0;
        elif (Delta(XT.H[i].X mod 2) <> Delta(K[i])) then return 0;
        end if:
    end do:
    return 1;
end proc:

# Probabilistic version
# Returns 1 if X is indeed a solution of XT.Sigma(H).X = K and
# of Delta(XT.H.X)=Delta(K) with good probability
# Returns 0 otherwise
IP1S_verifProba_char2:= proc (m, n, H, K, X)
local i, j, v, Xv, XT;
    v := Random (2, n, 1, integer[]);
    Xv:= X.v mod 2;
    XT:= Transpose (2, X);
    for i from 1 to m do           
        if (not LinearAlgebra[Equal] (XT.(Sigma(2,H[i]).Xv) mod 2,
                                      Sigma(2,K[i]).v mod 2))
        then return 0;
        end if:
        for j from 1 to n do # Diagonal terms
            if (add (X(a,j)*add (H[i](a,b)*X(b,j), b=1..n), a=1..n) mod 2 <>
                K[i](j,j)) then return 0;
            end if:
        end do:
    end do:
    return 1;
end proc:


# Dimension 1 space case: finds a good value for the indeterminate
# Returns 1, A where A is a specialization of A solution of AT.H.A=K
# over F_p
# Returns 0, `` if the only solution is over F_{p^2}
IP1S_orthogonalize:= proc (p, m, n, H, K, X)
local a, b, i, j, k, h, s;
    k:= 1;
    while (k <= n*n) do
        if (K[1](k) <> 0)
        then
            i:= iquo (k, n, 'j');
            if (j = 0) then j:= n; else i:= i+1;
            end if:
            h:= add (X(a,i)*add (H[1](a,b)*X(b,j), b=1..n) mod p, a=1..n) mod p;
            if (nops (indets (h)) <> 0) # found lambda in the expression
            then
                s:= msolve(h-K[1](i,j), p);
                if (s = NULL)
                then return 0, ``;
                else return 1, subs(s[1], X) mod p;
                end if:
            end if:
        end if:
        k:= k+1;
    end do:
end proc:

# Generates the matrix of Eq in xvar
# Uses the structure of Eq!
IP1S_genMatrix:= proc (p, n, Eq, xvar)
local X, nr, nc, i, j, k, l, n_im1;
    nr:= nops (Eq);
    nc:= n*n;
    X:= Mod (p, Matrix (nr*nc,nc), integer[]);
    n_im1:= 0;
    for i from 1 to n do
        for j from 1 to n do
            for k from 1 to n do
                for l from 1 to nr do
                    X((l-1)*nc+n_im1+j,n_im1+k)  := coeff (Eq[l](i,j),
                                                           xvar[n_im1+k]);
                    X((l-1)*nc+n_im1+j,n*(k-1)+j):= coeff (Eq[l](i,j),
                                                           xvar[n*(k-1)+j]);
                end do:
            end do:
        end do:
        n_im1:= n_im1+n;
    end do:
    return X, Vector[column] (nr*nc);
end proc:

# Computation of the space of solutions
# Returns 1, X if there is a nonzero general solution X
# Returns 0, `` otherwise
IP1S_linearSystem:= proc (p, m, n, H, K, are_affine:=0)
local X, xvar, Eq, i, S;
    xvar:= [seq (seq (cat (`x`,i,`_`,j), j=1..n), i=1..n)];
    X := Matrix (n, n, xvar);
    
    if (are_affine > 0) # Affine case, last row is 0,..,0,1
    then X(n,1..n-1):= 0; X(n,n):= 1;
    end if:

    # Solve the linear system
    Eq:= [seq (H[i].X-X.K[i] mod p, i=2..m)];
    Eq:= IP1S_genMatrix (p,n,Eq,xvar);
    # Eq:= {seq (op (convert (H[i].X-X.K[i] mod p, set)), i=2..m)};
    # S := LinearAlgebra[GenerateMatrix] (Eq, xvar);
    S:= Eq[1];
    ##### MAGMA
    # to_magma (cat("/tmp/bench",n,"_char_",p), p, convert (S, listlist)):
    # return 0,"End!";
    ##### MAGMA

    S := Linsolve (Eq) mod p;
    if (LinearAlgebra[Equal] (S,LinearAlgebra[ZeroVector] (n^2)))
    then return 0, ``;
    else return 1, subs ({seq (xvar[i]=S[i], i=1..n^2)}, X) mod p;
    end if:
end proc:

# Specialize a matrix in all but one variable
IP1S_specialization:= proc (p, X)
local v, roll;
    v:= indets (X);
    roll:= rand (0..p-1);
    return (subs ({seq (v[i]= roll(), i=1..nops(v)-1)}, X) mod p);
end proc:


# Main function
# Returns 1, A if AT.H.A = K over F_p
# Returns 0, S with a string explaining the reason otherwise
IP1S:= proc (p, m, n, H, K, are_affine:=0)
local answer, P, Q, HH, KK, X, Y, lambda, count;
    # Regular / Singular ?
    answer, P, Q:= IP1S_is_regular (p, m, n, H, K);
    if (answer = 1) then HH:= H; KK:= K;
    elif (answer = 0)
    then return 0, "Different numbers of essential variables";
    else return 0, "Singular" #TODO
    end if:

    # Normalization
    answer, HH, KK:= IP1S_normalize (p, m, n, HH, KK);
    if (answer = 0) # Normalization failed
    then return 0, "Zero Det vs. Nonzero Det"
    elif (answer = -1)
    then return 0, "Irregular case";
    end if:

    # Linear system
    answer, Y:= IP1S_linearSystem (p, nops (HH), n, HH, KK, are_affine);
    lambda:= indets (Y);
    if (answer = 0) # Zero solution
    then return 0, "Linear system failed";         
    elif (nops (lambda) = 0) # Unique solution
    then return Y;
    elif (nops (lambda) = 1) # Solution space of dimension 1
    then 
        answer, X:= IP1S_orthogonalize (p, m, n, H, K, Y);
        if (answer = 0) # Orthogonalization failed
        then return 0, "Equivalent over F_{p^2} not F_p";
        end if:
        answer   := IP1S_verifProba (p, m, n, H, K, X);
        if (answer = 0) # Verification failed
        then return 0, "Wrong solution!";
        else return 1, X;
        end if:
    else # Solution space of dimension > 1
        count:= 0;
        while (count < n^2) do
            X:= IP1S_specialization (p, Y);      # Pick up at random a line
            count:= count+1;
            if (Determinant (p, X) mod p = 0) then next;
            end if:
            answer, X:= IP1S_orthogonalize (p, m, n, H,
                                            K, X); # Orthogonalization
            if (answer = 0) then next;
            end if:
            answer:= IP1S_verifProba (p, m, n, H, K, X); # Verification
            if (answer = 0) then next;
            end if:
            return 1, X;
        end do:
        return 0, cat ("ProbablyNoSolution - Too many tries: ", count);
    end if:
end proc:

# Main function
# Returns 1, A if AT.H.A = K over F_p
# Returns 0, S with a string explaining the reason otherwise
IP1S_char2:= proc (m, n, H, K, are_affine:=0)
local answer, P, Q, HH, KK, X, Y, lambda;
    # Regular / Singular ?
    # answer, P, Q:= IP1S_is_regular (p, m, n, H, K);
    # if (answer = 1) then HH:= H; KK:= K;
    # elif (answer = 0)
    # then return 0, "Different numbers of essential variables";
    # else return 0, "Singular" #TODO
    # end if:
    HH:= H; KK:= K;
    # Normalization
    answer, HH, KK:= IP1S_normalize_char2 (m, n, HH, KK);
    if (answer = 0) # Normalization failed
    then return 0, "Zero Det vs. Nonzero Det"
    elif (answer = -1)
    then return 0, "Irregular case";
    end if:

    # Linear system
    answer, Y:= IP1S_linearSystem (2, nops (HH), n, HH, KK, are_affine);
    lambda:= indets (Y);
    if (answer = 0) # Zero solution
    then return 0, "Linear system failed";         
    elif (nops (lambda) = 0) # Unique solution
    then return Y;
    elif (nops (lambda) = 1) # Solution space of dimension 1
    then
        # TODO for general field of char 2
        # answer, X:= IP1S_orthogonalize (2, m, n, H, K, Y);
        # if (answer = 0) # Orthogonalization failed
        # then return 0, "Equivalent over F_4 not F_2";
        # end if:
        X:= subs (op(1,lambda) = 1, Y);
        answer   := IP1S_verifProba_char2 (m, n, H, K, X);
        if (answer = 0) # Verification failed
        then return 0, "Wrong solution!";
        else return 1, X;
        end if:
    else # Solution space of dimension > 1
        return 0, "Dimension strictly greater than 2";
    end if:
end proc:


# Create two sets of matrices and check if they are equivalent
# If are_equiv=1, they should be equivalent over F_p
# Otherwise if are_equiv=2, they should be equivalent over F_{p^2} but not F_p
# Otherwise they should not be equivalent over any extension of F_p
# If are_affine=1, the last row will be 0,..,0,1
IP1S_test:= proc (p, m, n, are_equiv:=1, are_affine:=0)
local H, K, A, B, answer;
    if (p < 2) then error "p not a prime";
    elif (p = 2)
    then
        H,K,A:= IP1S_rand_instance_char2 (m, n, are_equiv, are_affine);
        answer, B:= IP1S_char2 (m, n, H, K);
        if (answer = 1)
        then return
            [LinearAlgebra[Equal] (A,B), B];
        else return B;
        end if:
    else
        H,K,A:= IP1S_rand_instance (p, m, n, are_equiv, are_affine);
        answer, B:= IP1S (p, m, n, H, K);
        if (answer = 1)
        then return
            [LinearAlgebra[Equal] (A,B) or LinearAlgebra[Equal] (A,-B mod p),
             B];
        else return B;
        end if:
    end if:
end proc:

IP1S_bench:= proc (p, m, n, are_equiv:=1, are_affine:=0)
local H, K, A, st, et;
    if (p < 2) then error "p not a prime";
    elif (p = 2)
    then
        H,K,A:= IP1S_rand_instance_char2 (m, n, are_equiv, are_affine);
        st   := time ();
        IP1S_char2 (m, n, H, K);
        et:= time () - st;
        return et;
    else
        H,K,A:= IP1S_rand_instance (p, m, n, are_equiv, are_affine);
        st   := time ();
        IP1S (p, m, n, H, K);
        et:= time () - st;
        return et;
    end if:
end proc:

roll:= rand (1..2): #always equivalent but over F_p or F_{p^2}
for n from 10 to 100 by 10 do
    printf ("size %d, char %d, time: %f\n", n, pp,
            IP1S_bench (pp, n, n, roll(), 0));
    printf ("size %d, char %05d, time: %f\n", n, 2,
            IP1S_bench (2, n, n, 1, 0));
end do;