needs "Library/prime.ml";;
needs "Library/pocklington.ml";;
needs "Library/binomial.ml";;

prioritize_num();;

let FERMAT_PRIME_CONV n =
  let tm = subst [mk_small_numeral n,`x:num`] `prime(2 EXP (2 EXP x) + 1)` in
  (RAND_CONV NUM_REDUCE_CONV THENC PRIME_CONV) tm;;

FERMAT_PRIME_CONV 0;;
FERMAT_PRIME_CONV 1;;
FERMAT_PRIME_CONV 2;;
FERMAT_PRIME_CONV 3;;
FERMAT_PRIME_CONV 4;;
FERMAT_PRIME_CONV 5;;
FERMAT_PRIME_CONV 6;;
FERMAT_PRIME_CONV 7;;
FERMAT_PRIME_CONV 8;;

let CONG_TRIVIAL = prove
 (`!x y. n divides x /\ n divides y ==> (x == y) (mod n)`,
  MESON_TAC[CONG_0; CONG_SYM; CONG_TRANS]);;

let LITTLE_CHECK_CONV tm =
  EQT_ELIM((RATOR_CONV(LAND_CONV NUM_EXP_CONV) THENC CONG_CONV) tm);;

LITTLE_CHECK_CONV `(9 EXP 8 == 9) (mod 3)`;;
LITTLE_CHECK_CONV `(9 EXP 3 == 9) (mod 3)`;;
LITTLE_CHECK_CONV `(10 EXP 7 == 10) (mod 7)`;;
LITTLE_CHECK_CONV `(2 EXP 7 == 2) (mod 7)`;;
LITTLE_CHECK_CONV `(777 EXP 13 == 777) (mod 13)`;;

let DIVIDES_FACT_PRIME = prove
 (`!p. prime p ==> !n. p divides (FACT n) <=> p <= n`,
  GEN_TAC THEN DISCH_TAC THEN INDUCT_TAC THEN REWRITE_TAC[FACT; LE] THENL
   [ASM_MESON_TAC[DIVIDES_ONE; PRIME_0; PRIME_1];
    ASM_MESON_TAC[PRIME_DIVPROD_EQ; DIVIDES_LE; NOT_SUC; DIVIDES_REFL;
                  ARITH_RULE `~(p <= n) /\ p <= SUC n ==> p = SUC n`]]);;

let DIVIDES_BINOM_PRIME = prove
 (`!n p. prime p /\ 0 < n /\ n < p ==> p divides binom(p,n)`,
  REPEAT STRIP_TAC THEN
  MP_TAC(AP_TERM `(divides) p` (SPECL [`p - n`; `n:num`] BINOM_FACT)) THEN
  ASM_SIMP_TAC[DIVIDES_FACT_PRIME; PRIME_DIVPROD_EQ; SUB_ADD; LT_IMP_LE] THEN
  ASM_REWRITE_TAC[GSYM NOT_LT; LT_REFL] THEN
  ASM_SIMP_TAC[ARITH_RULE `0 < n /\ n < p ==> p - n < p`]);;

let DIVIDES_NSUM = prove
 (`!m n. (!i. m <= i /\ i <= n ==> p divides f(i)) ==> p divides nsum(m..n) f`,
  GEN_TAC THEN INDUCT_TAC THEN ASM_REWRITE_TAC[NSUM_CLAUSES_NUMSEG] THEN
  ASM_MESON_TAC[LE; LE_TRANS; DIVIDES_0; DIVIDES_ADD; LE_REFL]);;

let FLT_LEMMA = prove
 (`!p a b. prime p ==> ((a + b) EXP p == a EXP p + b EXP p) (mod p)`,
  REPEAT STRIP_TAC THEN REWRITE_TAC[BINOMIAL_THEOREM] THEN
  SUBGOAL_THEN `1 <= p /\ 0 < p` STRIP_ASSUME_TAC THENL
   [FIRST_ASSUM(MP_TAC o MATCH_MP PRIME_IMP_NZ) THEN ARITH_TAC; ALL_TAC] THEN
  ASM_SIMP_TAC[NSUM_CLAUSES_LEFT; LE_0; ARITH; NSUM_CLAUSES_RIGHT] THEN
  REWRITE_TAC[SUB_0; SUB_REFL; EXP; binom; BINOM_REFL; MULT_CLAUSES] THEN
  GEN_REWRITE_TAC LAND_CONV [ARITH_RULE `a + b = (b + 0) + a`] THEN
  REPEAT(MATCH_MP_TAC CONG_ADD THEN REWRITE_TAC[CONG_REFL]) THEN
  REWRITE_TAC[CONG_0] THEN MATCH_MP_TAC DIVIDES_NSUM THEN
  ASM_MESON_TAC[DIVIDES_RMUL; DIVIDES_BINOM_PRIME; ARITH_RULE
   `0 < p /\ 1 <= i /\ i <= p - 1 ==> 0 < i /\ i < p`]);;

let FERMAT_LITTLE = prove
 (`!p a. prime p ==> (a EXP p == a) (mod p)`,
  GEN_TAC THEN REWRITE_TAC[RIGHT_FORALL_IMP_THM] THEN DISCH_TAC THEN
  INDUCT_TAC THENL
   [ASM_MESON_TAC[EXP_EQ_0; CONG_REFL; PRIME_0];
    ASM_MESON_TAC[ADD1; FLT_LEMMA; EXP_ONE; CONG_ADD; CONG_TRANS; CONG_REFL]]);;

let FERMAT_LITTLE_COPRIME = prove
 (`!p a. prime p /\ coprime(a,p) ==> (a EXP (p - 1) == 1) (mod p)`,
  REPEAT STRIP_TAC THEN MATCH_MP_TAC CONG_MULT_LCANCEL THEN
  EXISTS_TAC `a:num` THEN ASM_REWRITE_TAC[GSYM(CONJUNCT2 EXP)] THEN
  ASM_SIMP_TAC[PRIME_IMP_NZ; ARITH_RULE `~(p = 0) ==> SUC(p - 1) = p`] THEN
  ASM_SIMP_TAC[FERMAT_LITTLE; MULT_CLAUSES]);;

let FERMAT_LITTLE_VARIANT = prove
 (`!p a. prime p ==> (a EXP (1 + m * (p - 1)) == a) (mod p)`,
  REPEAT STRIP_TAC THEN
  FIRST_ASSUM(DISJ_CASES_TAC o SPEC `a:num` o MATCH_MP PRIME_COPRIME_STRONG)
  THENL [ASM_MESON_TAC[CONG_TRIVIAL; ADD_AC; ADD1; DIVIDES_REXP_SUC]; 
         ALL_TAC] THEN
  GEN_REWRITE_TAC LAND_CONV [ARITH_RULE `a = a * 1`] THEN
  REWRITE_TAC[EXP_ADD; EXP_1] THEN MATCH_MP_TAC CONG_MULT THEN
  REWRITE_TAC[GSYM EXP_EXP; CONG_REFL] THEN
  ASM_MESON_TAC[COPRIME_SYM; COPRIME_EXP; PHI_PRIME; FERMAT_LITTLE_COPRIME]);;

let RSA = prove
 (`prime p /\ prime q /\ ~(p = q) /\
   (d * e == 1) (mod ((p - 1) * (q - 1))) /\
   plaintext < p * q /\ (ciphertext = (plaintext EXP e) MOD (p * q))
   ==> (plaintext = (ciphertext EXP d) MOD (p * q))`,
  REPEAT STRIP_TAC THEN ASM_REWRITE_TAC[] THEN
  ASM_SIMP_TAC[MOD_EXP_MOD; MULT_EQ_0; PRIME_IMP_NZ; EXP_EXP] THEN
  SUBGOAL_THEN `(plaintext == plaintext EXP (e * d)) (mod (p * q))` MP_TAC THENL
   [ALL_TAC; ASM_SIMP_TAC[CONG; MULT_EQ_0; PRIME_IMP_NZ; MOD_LT]] THEN
  GEN_REWRITE_TAC (LAND_CONV o ONCE_DEPTH_CONV) [MULT_SYM] THEN
  FIRST_X_ASSUM(DISJ_CASES_TAC o GEN_REWRITE_RULE I [CONG_TO_1]) THENL
   [ASM_MESON_TAC[MULT_EQ_1; ARITH_RULE `p - 1 = 1 <=> p = 2`]; ALL_TAC] THEN
  MATCH_MP_TAC CONG_CHINESE THEN ASM_SIMP_TAC[DISTINCT_PRIME_COPRIME] THEN
  ASM_MESON_TAC[FERMAT_LITTLE_VARIANT; MULT_AC; CONG_SYM]);;