FF2 := GF(2);

// List of precomputed elements a in GF(2^f) for even f <= 30,
// s.t. x^2+x+a is irreducible.
precomputed := [
    [ FF2 | 0, 1 ],
    [ FF2 | 0, 0, 0, 1 ],
    [ FF2 | 0, 0, 0, 1, 0, 0 ],
    [ FF2 | 0, 0, 0, 0, 0, 1, 0, 0 ],
    [ FF2 | 0, 0, 0, 0, 0, 1, 0, 0, 0, 0 ],
    [ FF2 | 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0 ],
    [ FF2 | 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0 ],
    [ FF2 | 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0 ],
    [ FF2 | 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1 ],
    [ FF2 | 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0 ],
    [ FF2 | 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ],
    [ FF2 | 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ],
    [ FF2 | 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0 ],
    [ FF2 | 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ],
    [ FF2 | 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
    0 ]
];

function special_local_square_class( p )
// construct a p-unit u with square defect = 4*O = p^(2e)
// implementation of Remark 3.5

    K := NumberField(Place(p));
    // mKp : OK -> Kp
    Kp, mKp := ResidueClassField(p);
    f := Degree(Kp);
    if IsOdd(f) then 
        u_ := One(Kp);
    elif (f div 2) lt #precomputed then
        F := GF(2, f);
        Embed(F, Kp);
        u_ := F!precomputed[ f div 2 ];
        u_ := Kp!u_;
    else
        Kpy<y> := PolynomialRing(Kp);
        u_ := Random(Kp);
        while not IsIrreducible(y^2 + y + u_) do
            u_ := Random(Kp);
        end while;
    end if;
    u := 1 + 4*K!(u_ @@ mKp);

    assert2 QuadraticDefect(u, p) eq 2*AbsoluteRamificationIndex(p);
    return u;
end function;

function local_square_classes( p )
// compute a basis of the local square-class group
// in a dyadic completion
// implementation of Algorithm 2

    OK := Order(p);
    K := NumberField(Place(p));
    e := AbsoluteRamificationIndex(p); 
    f := AbsoluteInertiaDegree(p); 
    pi := UniformizingElement(p);
    
    // the square-class corresponding to the unique unramified quadratic extension
    Delta := special_local_square_class(p);
    // mKp : OK -> Kp
    Kp, mKp := ResidueClassField(p);
    // lifting of a basis of Kp over FF2
    basis := [ (Kp.1^k) @@ mKp : k in [0..f-1] ];
    // sequence of els of valuations 1,3,5,..,2e-1
    if IsOdd(e) then
        powers1 := [ pi^r : r in [1..e-1 by 2] ];
        powers2 := [ 2*pi^r : r in [0..e by 2] ];
    else
        powers1 := [ pi^r : r in [1..e-1 by 2] ];
        powers2 := [ 2*pi^r : r in [1..e-1 by 2] ];
    end if;
    // reduce modulo 2 and 4, respectively
    powers1 := [ OK![ Integers()!bi mod 2 : bi in Eltseq(OK!b)] : b in powers1 ];
    powers2 := [ OK![ Integers()!bi mod 4 : bi in Eltseq(OK!b)] : b in powers2 ];
    powers := powers1 cat powers2;
    // e*f linearly independent square classes
    U := Flat([[ OK!(1 + u*a) : u in basis ] : a in powers ]);
    // the complete basis of the local group of square-class
    return U cat [ Delta, pi ], basis, powers;
end function;

function reduce( a )
// Helper function. Given an element a of the field K,
// a priori known to be a unit of OK, this function 
// outputs an element in the same dyadic square class
// but with smaller coefficients.

    K := Parent(a);
    n := Numerator(a);
    d := Denominator(a);
    v := Min([ Valuation(Integers()!c, 2) : c in Eltseq(n) ]);
    v := Max(v, Valuation(d, 2));
    n_ := K![ Integers()!c mod 2^(v + 3) : c in Eltseq(n) ];
    d_ := d mod 2^(v + 3);
    
    return n_/d_;
end function;

function coordinates( a, p, B, powers, mOKmod8 )
// Given a basis B of the vector space K_p^*/(K_p^*)2
// and a nonzero element a, this function outputs
// the coordinates of a wrt B
// implementation of Algorithm 3
 
    error if IsZero(a), "Element must not be zero.";
    
    OK := Order(p);
    e := AbsoluteRamificationIndex(p); 
    f := AbsoluteInertiaDegree(p); 
    // mKp : OK -> Kp
    Kp, mKp := ResidueClassField(p);
    pi := UniformizingElement(p);
    
    a *:= Denominator(a)^2;
    // initialize the result
    dim := #B;
    E := [ Zero(FF2) : i in [1..dim] ];
    
    // compute the valuation
    v := Valuation(a, p);
    // last coord is the parity of v
    E[dim] := v mod 2;
    // update a
    a := reduce( a*pi^(-v) );

    // square root in the residue field
    s := Sqrt( mKp(a)^-1 ) @@ mKp;
    // update a
    a *:= s^2;
    
    // valuation for computing quadratic defect
    v := Valuation(a - 1, p);
    // main loop
    while v lt 2*e do
        if IsOdd(v) then
            // we hit the quadratic defect => extract coordinates
            k := (v + 1) div 2;
            inx := f*(k-1);
            // coords in K(p) over FF2
            coords := Eltseq( ((a-1)/powers[k]) @ mKp );
            for j in [1..f] do
                E[ inx + j ] := coords[j];
            end for;
            // compute the product of inverses of basis els
            prodB := &* [ B[inx + j] @ mOKmod8 : j in [1..f] | IsOne(coords[j]) ];
            prodB := ((prodB)^(-1)) @@ mOKmod8;
            // update a 
//             a /:= &* [ B[inx + j] : j in [1..f] | IsOne(coords[j]) ];
            a := reduce(a*prodB);
            v := Valuation(a - 1, p);
        else
            // keep computing the quadratic defect
            // square root  in the residue field
            s := Sqrt( mKp((a-1)*pi^(-v)) ) @@ mKp;
            // update a
            b := (((1 + s*pi^(v div 2)) @ mOKmod8)^-2) @@ mOKmod8;
            a := reduce(a*b);
//             a /:= (1 + s*pi^(v div 2))^2;
            v := Valuation(a - 1, p);
        end if;
    end while;
    
    //assert v eq 2*e;
    if v eq 2*e and IsIrreducible( Polynomial([ Kp | mKp((a-1)/4), 1, 1]) ) then
        E[dim - 1] := 1;
    end if;
    
    // output the result
    return E;
end function;

intrinsic LocalGroupOfSquaresClasses( p::RngOrdIdl ) -> ModFld, Map
{
Given a prime ideal over a number field K, this function outputs a vector space V
over FF2 isomorphic to the group G := K_p^*/(K_p^*)^2, together with an isomorphism G -> V.

This function is functionally equivalent to the built-in function LocalMultiplicativeGroupModSquares, only much faster for dyadic primes.
}
    R:= Order(p);
    require IsPrime(p) and IsMaximal(R): "The ideal must be prime ideal of a maximal order";
  
    // the non-dyadic case is handled by the built-in method
    if Minimum(p) ne 2 then
        return LocalMultiplicativeGroupModSquares(p);
    end if;
    
    // the dyadic case is handled by the new algorithm
    K:= NumberField(R);
    OKmod8, mOKmod8 := quo< R | 8 >;
    // construct a basis
    B, basis, powers := local_square_classes( p );
    dim := #B;
    // construct a vector space isomorphic to span(B)
    V := VectorSpace(FF2, dim);
    // construct the two isomorphisms
    m := map< V -> K | 
        x :-> ((&* ( [ One(OKmod8) ] cat [ B[i] @ mOKmod8 : i in [1..dim - 1] | IsOne(x[i]) ] )) @@ mOKmod8) * (IsOne(x[dim]) select B[dim] else One(K)),
        y :-> coordinates(y, p, B, powers, mOKmod8)>;
        
    // output the result
    return V, m;
end intrinsic;
