(* Polynomials over a prime field F_p, and factorization in it.

   Part 3 carried a small Pmod module inside Upoly, just enough to take
   a GCD modulo a machine prime. Factorization needs far more from the
   same place, so it becomes its own module: division with remainder,
   modular exponentiation, the Frobenius map, and the three algorithms
   that between them split any polynomial over F_p into irreducibles.

   The prime is passed explicitly to every function rather than fixed in
   a functor, which keeps the call sites honest about which field they
   are working in. Primes stay below 2^30 so a product of two residues
   fits in an OCaml int.

   Representation is dense: coefficient of x^i at index i, last entry
   nonzero, empty array for the zero polynomial. *)

(* ------------------------------------------------------------------ *)
(* The coefficient field                                               *)
(* ------------------------------------------------------------------ *)

module Zp = struct
  let norm p a =
    let r = a mod p in
    if r < 0 then r + p else r

  let add p a b =
    let s = a + b in
    if s >= p then s - p else s

  let sub p a b =
    let s = a - b in
    if s < 0 then s + p else s

  let mul p a b = a * b mod p
  let neg p a = if a = 0 then 0 else p - a

  let pow p a n =
    let rec go acc b n =
      if n = 0 then acc else if n land 1 = 1 then go (mul p acc b) (mul p b b) (n lsr 1)
      else go acc (mul p b b) (n lsr 1)
    in
    go 1 (norm p a) n

  (* Extended Euclid, iterative, tracking only the coefficient needed. *)
  let inv p a =
    let a = norm p a in
    if a = 0 then invalid_arg "Fp.Zp.inv: zero has no inverse";
    let r0 = ref p and r1 = ref a in
    let t0 = ref 0 and t1 = ref 1 in
    while !r1 <> 0 do
      let q = !r0 / !r1 in
      let r = !r0 - (q * !r1) in
      r0 := !r1;
      r1 := r;
      let t = !t0 - (q * !t1) in
      t0 := !t1;
      t1 := t
    done;
    norm p !t0
end

let is_prime (n : int) : bool =
  if n < 2 then false
  else if n mod 2 = 0 then n = 2
  else begin
    let i = ref 3 and ok = ref true in
    while !ok && !i * !i <= n do
      if n mod !i = 0 then ok := false;
      i := !i + 2
    done;
    !ok
  end

let next_prime (n : int) : int =
  let i = ref (if n < 2 then 2 else n) in
  while not (is_prime !i) do
    incr i
  done;
  !i

(* ------------------------------------------------------------------ *)
(* Polynomials                                                         *)
(* ------------------------------------------------------------------ *)

type t = int array

let zero : t = [||]

let normalize (a : t) : t =
  let n = ref (Array.length a) in
  while !n > 0 && a.(!n - 1) = 0 do
    decr n
  done;
  if !n = Array.length a then a else Array.sub a 0 !n

let degree (a : t) : int = Array.length a - 1
let is_zero (a : t) : bool = Array.length a = 0
let lc (a : t) : int = if is_zero a then 0 else a.(Array.length a - 1)
let coeff (a : t) (i : int) : int = if i >= 0 && i < Array.length a then a.(i) else 0
let const p (c : int) : t = normalize [| Zp.norm p c |]
let one : t = [| 1 |]
let x : t = [| 0; 1 |]
let of_list p (l : int list) : t = normalize (Array.of_list (List.map (Zp.norm p) l))
let is_one (a : t) : bool = Array.length a = 1 && a.(0) = 1

let equal (a : t) (b : t) : bool =
  Array.length a = Array.length b
  && (let ok = ref true in
      Array.iteri (fun i v -> if v <> b.(i) then ok := false) a;
      !ok)

let add p (a : t) (b : t) : t =
  let n = max (Array.length a) (Array.length b) in
  normalize (Array.init n (fun i -> Zp.add p (coeff a i) (coeff b i)))

let sub p (a : t) (b : t) : t =
  let n = max (Array.length a) (Array.length b) in
  normalize (Array.init n (fun i -> Zp.sub p (coeff a i) (coeff b i)))

let neg p (a : t) : t = Array.map (Zp.neg p) a

let scale p (a : t) (c : int) : t =
  let c = Zp.norm p c in
  if c = 0 then zero else Array.map (fun v -> Zp.mul p v c) a

let shift (a : t) (k : int) : t =
  if is_zero a || k = 0 then a
  else Array.init (Array.length a + k) (fun i -> if i < k then 0 else a.(i - k))

let mul p (a : t) (b : t) : t =
  if is_zero a || is_zero b then zero
  else begin
    let n = Array.length a and m = Array.length b in
    let r = Array.make (n + m - 1) 0 in
    for i = 0 to n - 1 do
      if a.(i) <> 0 then
        for j = 0 to m - 1 do
          r.(i + j) <- Zp.add p r.(i + j) (Zp.mul p a.(i) b.(j))
        done
    done;
    normalize r
  end

let monic p (a : t) : t = if is_zero a then a else scale p a (Zp.inv p (lc a))

(* Division with remainder. F_p is a field, so unlike Z there is no
   pseudo-division here: the leading coefficient is always invertible. *)
let divmod p (a : t) (b : t) : t * t =
  if is_zero b then raise Division_by_zero;
  let db = degree b and ib = Zp.inv p (lc b) in
  if degree a < db then (zero, a)
  else begin
    let q = Array.make (degree a - db + 1) 0 in
    let r = ref a in
    while (not (is_zero !r)) && degree !r >= db do
      let d = degree !r - db in
      let c = Zp.mul p (lc !r) ib in
      q.(d) <- c;
      r := sub p !r (shift (scale p b c) d)
    done;
    (normalize q, !r)
  end

let rem p a b = snd (divmod p a b)
let div p a b = fst (divmod p a b)

let gcd p (a : t) (b : t) : t =
  let u = ref a and v = ref b in
  while not (is_zero !v) do
    let r = rem p !u !v in
    u := !v;
    v := r
  done;
  monic p !u

(* Extended Euclid: returns (g, s, t) with s*a + t*b = g, g monic.
   Hensel lifting needs the cofactors, not just the GCD. *)
let ext_gcd p (a : t) (b : t) : t * t * t =
  let r0 = ref a and r1 = ref b in
  let s0 = ref one and s1 = ref zero in
  let t0 = ref zero and t1 = ref one in
  while not (is_zero !r1) do
    let q, r = divmod p !r0 !r1 in
    r0 := !r1;
    r1 := r;
    let s = sub p !s0 (mul p q !s1) in
    s0 := !s1;
    s1 := s;
    let t = sub p !t0 (mul p q !t1) in
    t0 := !t1;
    t1 := t
  done;
  let c = Zp.inv p (lc !r0) in
  (scale p !r0 c, scale p !s0 c, scale p !t0 c)

let diff p (a : t) : t =
  if Array.length a <= 1 then zero
  else normalize (Array.init (Array.length a - 1) (fun i -> Zp.mul p a.(i + 1) (Zp.norm p (i + 1))))

let eval p (a : t) (v : int) : int =
  let acc = ref 0 in
  for i = Array.length a - 1 downto 0 do
    acc := Zp.add p (Zp.mul p !acc v) a.(i)
  done;
  !acc

(* a^n mod m, by repeated squaring in the quotient ring F_p[x]/(m). *)
let pow_mod p (a : t) (n : int) (m : t) : t =
  let rec go acc b n =
    if n = 0 then acc
    else if n land 1 = 1 then go (rem p (mul p acc b) m) (rem p (mul p b b) m) (n lsr 1)
    else go acc (rem p (mul p b b) m) (n lsr 1)
  in
  go one (rem p a m) n

(* The same with an exponent too large for an int: equal-degree
   factorization needs a^((p^d - 1)/2), which overflows immediately. *)
let pow_mod_big p (a : t) (n : Bigint.t) (m : t) : t =
  let two = Bigint.of_int 2 in
  let acc = ref one and b = ref (rem p a m) and e = ref n in
  while not (Bigint.is_zero !e) do
    let q, r = Bigint.divmod !e two in
    if not (Bigint.is_zero r) then acc := rem p (mul p !acc !b) m;
    b := rem p (mul p !b !b) m;
    e := q
  done;
  !acc

let random p (deg : int) : t =
  normalize (Array.init (deg + 1) (fun _ -> Random.int p))

let to_string ?(var = "x") (a : t) : string =
  if is_zero a then "0"
  else begin
    let buf = Buffer.create 64 in
    let first = ref true in
    for i = Array.length a - 1 downto 0 do
      if a.(i) <> 0 then begin
        if not !first then Buffer.add_string buf " + ";
        first := false;
        let unit = a.(i) = 1 && i > 0 in
        if not unit then Buffer.add_string buf (string_of_int a.(i));
        if i > 0 then begin
          if not unit then Buffer.add_string buf "*";
          Buffer.add_string buf var;
          if i > 1 then Buffer.add_string buf ("^" ^ string_of_int i)
        end
      end
    done;
    Buffer.contents buf
  end

(* ------------------------------------------------------------------ *)
(* Square-free decomposition in characteristic p                       *)
(* ------------------------------------------------------------------ *)

(* Characteristic zero lets you argue that f and f' share exactly the
   repeated factors. In characteristic p that argument has a hole: the
   derivative of x^p is zero, so a p-th power is invisible to it. Yun
   becomes Musser, with an extra branch that takes a p-th root and
   recurses. Over a prime field the root is cheap, because a^p = a for
   every a, so only the exponents move. *)
let pth_root p (a : t) : t =
  let n = degree a in
  normalize (Array.init ((n / p) + 1) (fun i -> a.(i * p)))

let rec square_free p (f : t) : (t * int) list =
  if degree f <= 0 then []
  else begin
    let f = monic p f in
    let fd = diff p f in
    if is_zero fd then
      (* f is a p-th power: f(x) = g(x)^p, so every multiplicity in g
         is multiplied by p. *)
      List.map (fun (g, m) -> (g, m * p)) (square_free p (pth_root p f))
    else begin
      let c = ref (gcd p f fd) in
      let w = ref (div p f !c) in
      let out = ref [] and i = ref 1 in
      while degree !w > 0 do
        let y = gcd p !w !c in
        let z = div p !w y in
        if degree z > 0 then out := (z, !i) :: !out;
        w := y;
        c := div p !c y;
        incr i
      done;
      (* Whatever is left in c is a p-th power. *)
      let tail =
        if degree !c > 0 then
          List.map (fun (g, m) -> (g, m * p)) (square_free p (pth_root p !c))
        else []
      in
      List.rev_append !out tail
    end
  end

(* ------------------------------------------------------------------ *)
(* Distinct-degree factorization                                       *)
(* ------------------------------------------------------------------ *)

(* x^(p^d) - x is the product of every monic irreducible whose degree
   divides d. Taking a GCD with it therefore peels off exactly the
   factors of degree d, once the smaller ones are already gone. The
   whole algorithm is that identity plus the Frobenius map applied
   repeatedly. Input must be monic and square-free. *)
let distinct_degree p (f : t) : (t * int) list =
  let out = ref [] in
  let fstar = ref f in
  let xp = ref x in
  let d = ref 1 in
  while degree !fstar >= 2 * !d do
    xp := pow_mod p !xp p !fstar;
    let g = gcd p !fstar (sub p !xp x) in
    if degree g > 0 then begin
      out := (g, !d) :: !out;
      fstar := div p !fstar g;
      xp := rem p !xp !fstar
    end;
    incr d
  done;
  (* Anything left is irreducible: its degree exceeds half the degree
     of what remained, so it cannot be a product of two pieces. *)
  if degree !fstar > 0 then out := (!fstar, degree !fstar) :: !out;
  List.rev !out

(* ------------------------------------------------------------------ *)
(* Equal-degree factorization (Cantor-Zassenhaus)                      *)
(* ------------------------------------------------------------------ *)

(* Splits a product of r distinct irreducibles all of degree d. The
   trick: for odd p, half the nonzero elements of F_p^d are squares, so
   a random a has a^((p^d-1)/2) equal to 1 in some of the r factor
   fields and to -1 in others. A GCD with that value minus one then
   separates the two groups. Each attempt succeeds with probability
   close to one half, so the expected number of attempts is constant. *)
let rec equal_degree p (f : t) (d : int) : t list =
  if degree f = 0 then []
  else if degree f = d then [ monic p f ]
  else begin
    let n = degree f in
    let e =
      (* (p^d - 1) / 2, which needs a bignum for any interesting d. *)
      Bigint.div (Bigint.sub (Bigint.pow (Bigint.of_int p) d) Bigint.one) (Bigint.of_int 2)
    in
    let result = ref None in
    while !result = None do
      let a = random p (n - 1) in
      if degree a > 0 then begin
        let g = gcd p a f in
        let split =
          if degree g > 0 && degree g < n then Some g
          else begin
            let b = pow_mod_big p a e f in
            let g = gcd p (sub p b one) f in
            if degree g > 0 && degree g < n then Some g else None
          end
        in
        match split with
        | Some g -> result := Some (equal_degree p g d @ equal_degree p (div p f g) d)
        | None -> ()
      end
    done;
    match !result with Some l -> l | None -> assert false
  end

(* ------------------------------------------------------------------ *)
(* The whole factorization over F_p                                    *)
(* ------------------------------------------------------------------ *)

(* Square-free, then distinct-degree, then equal-degree: three passes,
   each one making the next one's job easier. Returns the leading
   coefficient and the monic irreducible factors with multiplicities. *)
let factor p (f : t) : int * (t * int) list =
  if is_zero f then invalid_arg "Fp.factor: zero";
  let c = lc f in
  let f = monic p f in
  let out = ref [] in
  List.iter
    (fun (g, m) ->
      List.iter
        (fun (h, d) -> List.iter (fun q -> out := (q, m) :: !out) (equal_degree p h d))
        (distinct_degree p g))
    (square_free p f);
  (c, List.sort (fun (a, _) (b, _) -> compare a b) !out)

let is_irreducible p (f : t) : bool =
  degree f > 0 && (match factor p f with _, [ (_, 1) ] -> true | _ -> false)

(* ------------------------------------------------------------------ *)
(* Berlekamp, as an independent check                                  *)
(* ------------------------------------------------------------------ *)

(* Berlekamp attacks the same problem from linear algebra rather than
   from gcds. The map v -> v^p - v is linear on F_p[x]/(f), and its
   kernel has dimension exactly equal to the number of irreducible
   factors of f. That number is worth having on its own: it is a check
   on distinct-degree plus equal-degree that shares none of their code.
   Building the matrix costs p-th powers of x^i, so this is for small
   primes only. *)
let berlekamp_matrix p (f : t) : int array array =
  let n = degree f in
  let q = Array.make_matrix n n 0 in
  for i = 0 to n - 1 do
    let row = pow_mod p (shift one i) p f in
    for j = 0 to n - 1 do
      q.(j).(i) <- coeff row j
    done;
    (* Subtract the identity: the matrix of v -> v^p - v. *)
    q.(i).(i) <- Zp.sub p q.(i).(i) 1
  done;
  q

(* Gaussian elimination over F_p, returning a basis of the nullspace. *)
let nullspace p (m : int array array) : int array list =
  let rows = Array.length m in
  if rows = 0 then []
  else begin
    let cols = Array.length m.(0) in
    let a = Array.map Array.copy m in
    let pivot_of_col = Array.make cols (-1) in
    let r = ref 0 in
    for c = 0 to cols - 1 do
      if !r < rows then begin
        (* Find a row with a nonzero entry in this column. *)
        let piv = ref (-1) in
        for i = rows - 1 downto !r do
          if a.(i).(c) <> 0 then piv := i
        done;
        if !piv >= 0 then begin
          let tmp = a.(!r) in
          a.(!r) <- a.(!piv);
          a.(!piv) <- tmp;
          let inv = Zp.inv p a.(!r).(c) in
          for j = 0 to cols - 1 do
            a.(!r).(j) <- Zp.mul p a.(!r).(j) inv
          done;
          for i = 0 to rows - 1 do
            if i <> !r && a.(i).(c) <> 0 then begin
              let factor = a.(i).(c) in
              for j = 0 to cols - 1 do
                a.(i).(j) <- Zp.sub p a.(i).(j) (Zp.mul p factor a.(!r).(j))
              done
            end
          done;
          pivot_of_col.(c) <- !r;
          incr r
        end
      end
    done;
    (* One basis vector per free column. *)
    let basis = ref [] in
    for c = cols - 1 downto 0 do
      if pivot_of_col.(c) < 0 then begin
        let v = Array.make cols 0 in
        v.(c) <- 1;
        for c2 = 0 to cols - 1 do
          if pivot_of_col.(c2) >= 0 then v.(c2) <- Zp.neg p a.(pivot_of_col.(c2)).(c)
        done;
        basis := v :: !basis
      end
    done;
    !basis
  end

(* The number of irreducible factors of a square-free f, counted
   without finding any of them. *)
let berlekamp_count p (f : t) : int =
  if degree f <= 0 then 0 else List.length (nullspace p (berlekamp_matrix p f))