Library PrimeGapS1.MaynardSpec

From Stdlib Require Import ZArith List.
From mathcomp Require Import all_ssreflect all_algebra.
From PrimeGapS1 Require Import MaynardFactQ MaynardBasis.

Import ListNotations.

Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.

Import GRing.Theory.


Local Open Scope ring_scope.

Fixpoint compositions_aux (r remaining : nat) : seq (seq nat) :=
  match r with
  | 0 ⇒ if (remaining == 0)%N then [:: [::] ] else [::]
  | S r' ⇒
      flatten
        [seq [seq (a :: tl) | tl <- compositions_aux r' (remaining - a)%N]
           | a <- iota 1 remaining]
  end.

Definition compositions (r n : nat) : seq (seq nat) := compositions_aux r n.

Definition cff (a : seq nat) : rat :=
  \prod_(x <- a) (factQ (2 × x) / factQ x).

Definition G_2 (n k : nat) : rat :=
  if (n == 0)%N then 1
  else
    factQ n
    × \sum_(r <- iota 1 n) binQ k r × \sum_(a <- compositions r n) cff a.

Definition K1 : nat := 105.
Definition K2 : nat := 104.

Definition M1_entry (bi ci bj cj : nat) : rat :=
  let b := (bi + bj)%N in
  let c := (ci + cj)%N in
  factQ b / factQ (K1 + b + 2 × c)%N × G_2 c K1.

Definition alpha (b c cp : nat) : rat :=
  binQ c cp × factQ b × factQ (2 × c - 2 × cp)%N
    / factQ (b + 2 × c - 2 × cp + 1)%N.

Definition M2_entry (bi ci bj cj : nat) : rat :=
  \sum_(cp1 <- iota 0 ci.+1)
    \sum_(cp2 <- iota 0 cj.+1)
      let bp1 := (bi + 2 × ci - 2 × cp1 + 1)%N in
      let bp2 := (bj + 2 × cj - 2 × cp2 + 1)%N in
      let bsum := (bp1 + bp2)%N in
      let csum := (cp1 + cp2)%N in
      alpha bi ci cp1 × alpha bj cj cp2
        × factQ bsum / factQ (K2 + bsum + 2 × csum)%N
        × G_2 csum K2.

Definition M1_spec_ij (i j : nat) : rat :=
  let bci := List.nth i maynard_basis (0%N, 0%N) in
  let bcj := List.nth j maynard_basis (0%N, 0%N) in
  M1_entry bci.1 bci.2 bcj.1 bcj.2.

Definition M2_spec_ij (i j : nat) : rat :=
  let bci := List.nth i maynard_basis (0%N, 0%N) in
  let bcj := List.nth j maynard_basis (0%N, 0%N) in
  M2_entry bci.1 bci.2 bcj.1 bcj.2.


Local Close Scope ring_scope.
Local Open Scope Z_scope.

Fixpoint factZ (n : nat) : Z :=
  match n with
  | O ⇒ 1
  | S k ⇒ Z.of_nat (S k) × factZ k
  end.

Definition dblratZ (x : nat) : Z := factZ (2 × x)%nat / factZ x.

Fixpoint prod_dblratZ (xs : list nat) : Z :=
  match xs with
  | nil ⇒ 1
  | x :: r ⇒ dblratZ x × prod_dblratZ r
  end.

Fixpoint compositions_auxZ (r remaining : nat) : list (list nat) :=
  match r with
  | O ⇒ if Nat.eqb remaining 0 then [nil] else nil
  | S r' ⇒
      List.flat_map
        (fun a ⇒
           List.map (fun tl ⇒ a :: tl)
                    (compositions_auxZ r' (remaining - a)%nat))
        (List.seq 1 remaining)
  end.

Definition compositionsZ (r n : nat) : list (list nat) :=
  compositions_auxZ r n.

Definition cffZ (a : list nat) : Z := prod_dblratZ a.

Definition binZ (n k : nat) : Z :=
  if Nat.leb k n then factZ n / (factZ k × factZ (n - k)%nat) else 0.

Definition G2Z (n k : nat) : Z :=
  if Nat.eqb n O then 1
  else
    factZ n
    × List.fold_left Z.add
        (List.map
           (fun r ⇒
              binZ k r
              × List.fold_left Z.add
                  (List.map cffZ (compositionsZ r n)) 0)
           (List.seq 1 n)) 0.


Definition K1n : nat := 105.
Definition K2n : nat := 104.

Definition m1_num_den (bi ci bj cj : nat) : Z × Z :=
  let b := (bi + bj)%nat in
  let c := (ci + cj)%nat in
  (factZ b × G2Z c K1n, factZ (K1n + b + 2 × c)%nat).


Definition qplus (p q : Z × Z) : Z × Z :=
  let '(a, b) := p in
  let '(c, d) := q in
  (a × d + c × b, b × d).

Definition qmul (p q : Z × Z) : Z × Z :=
  let '(a, b) := p in
  let '(c, d) := q in
  (a × c, b × d).

Definition alphaZ (b c cp : nat) : Z × Z :=
  (binZ c cp × factZ b × factZ (2 × c - 2 × cp)%nat,
   factZ (b + 2 × c - 2 × cp + 1)%nat).

Definition m2_term_num_den (bi ci bj cj cp1 cp2 : nat) : Z × Z :=
  let bp1 := (bi + 2 × ci - 2 × cp1 + 1)%nat in
  let bp2 := (bj + 2 × cj - 2 × cp2 + 1)%nat in
  let bsum := (bp1 + bp2)%nat in
  let csum := (cp1 + cp2)%nat in
  qmul (qmul (alphaZ bi ci cp1) (alphaZ bj cj cp2))
       (factZ bsum × G2Z csum K2n, factZ (K2n + bsum + 2 × csum)%nat).

Definition m2_num_den (bi ci bj cj : nat) : Z × Z :=
  List.fold_left
    (fun acc cp1 ⇒
       List.fold_left
         (fun acc2 cp2 ⇒ qplus acc2 (m2_term_num_den bi ci bj cj cp1 cp2))
         (List.seq 0 (S cj)) acc)
    (List.seq 0 (S ci)) (0, 1).


Definition m1_num_den_at (i j : nat) : Z × Z :=
  let bci := List.nth i maynard_basis (0%nat, 0%nat) in
  let bcj := List.nth j maynard_basis (0%nat, 0%nat) in
  m1_num_den bci.1 bci.2 bcj.1 bcj.2.

Definition m2_num_den_at (i j : nat) : Z × Z :=
  let bci := List.nth i maynard_basis (0%nat, 0%nat) in
  let bcj := List.nth j maynard_basis (0%nat, 0%nat) in
  m2_num_den bci.1 bci.2 bcj.1 bcj.2.