Require Import Coq.ZArith.ZArith Coq.ZArith.Znumtheory.
Require Import Coq.Numbers.Natural.Peano.NPeano.
Require Import Coq.Program.Equality.
Require Import Crypto.CompleteEdwardsCurve.Pre.
Require Import Crypto.CompleteEdwardsCurve.CompleteEdwardsCurveTheorems.
Require Import Bedrock.Word.
Require Import Crypto.Encoding.ModularWordEncodingTheorems.
Require Import Crypto.Tactics.VerdiTactics.
Require Import Crypto.Util.ZUtil.
Require Import Crypto.Algebra.

Require Import Crypto.Spec.Encoding Crypto.Spec.ModularWordEncoding Crypto.Spec.ModularArithmetic.

Require Import Crypto.Util.Notations.
Require Export Crypto.Util.FixCoqMistakes.

Generalizable All Variables.
Section PointEncodingPre.
  Context {F eq zero one opp add sub mul inv div} `{field F eq zero one opp add sub mul inv div}.
  Local Infix "==" := eq : type_scope.
  Local Notation "a !== b" := (not (a == b)): type_scope.
  Local Notation "0" := zero.  Local Notation "1" := one.
  Local Infix "+" := add. Local Infix "*" := mul.
  Local Infix "-" := sub. Local Infix "/" := div.
  Local Notation "x ^ 2" := (x*x).

  Add Field EdwardsCurveField : (Field.field_theory_for_stdlib_tactic (T:=F)).

  Context {eq_dec:forall x y : F, {x==y}+{x==y->False} }.
  Definition F_eqb x y := if eq_dec x y then true else false.
  Lemma F_eqb_iff : forall x y, F_eqb x y = true <-> x == y.
  Proof.
    unfold F_eqb; intros; destruct (eq_dec x y); split; auto; discriminate.
  Qed.

  Context {a d:F} {prm:@E.twisted_edwards_params F eq zero one add mul a d}.
  Local Notation point := (@E.point F eq one add mul a d).
  Local Notation onCurve := (@onCurve F eq one add mul a d).
  Local Notation solve_for_x2 := (@E.solve_for_x2 F one sub mul div a d).

  Context {sz : nat} (sz_nonzero : (0 < sz)%nat).
  Context {sqrt : F -> F} (sqrt_square : forall x root, x == (root ^2) -> sqrt x == root)
    (sqrt_subst : forall x y, x == y -> sqrt x == sqrt y).
  Context (FEncoding : canonical encoding of F as (word sz)).
  Context {sign_bit : F -> bool} (sign_bit_zero : forall x, x == 0 -> Logic.eq (sign_bit x) false)
    (sign_bit_opp : forall x, x !== 0 -> Logic.eq (negb (sign_bit x)) (sign_bit (opp x)))
    (sign_bit_subst : forall x y, x == y -> sign_bit x = sign_bit y).

  Definition sqrt_ok (a : F) := (sqrt a) ^ 2 == a.

  Lemma square_sqrt : forall y root, y == (root ^2) ->
    sqrt_ok y.
  Proof.
    unfold sqrt_ok; intros ? ? root2_y.
    pose proof root2_y.
    apply sqrt_square in root2_y.
    rewrite root2_y.
    symmetry; assumption.
  Qed.

  Lemma solve_onCurve: forall x y : F, onCurve (x,y) ->
    onCurve (sqrt (solve_for_x2 y), y).
  Proof.
    intros.
    apply E.solve_correct.
    eapply square_sqrt.
    symmetry.
    apply E.solve_correct; eassumption.
  Qed.

  (* TODO : move? *)
  Lemma square_opp : forall x : F, (opp x ^2) == (x ^2).
  Proof.
    intros. ring.
  Qed.

  Lemma solve_opp_onCurve: forall x y : F, onCurve (x,y) ->
    onCurve (opp (sqrt (solve_for_x2 y)), y).
  Proof.
    intros.
    apply E.solve_correct.
    etransitivity; [ apply square_opp | ].
    eapply square_sqrt.
    symmetry.
    apply E.solve_correct; eassumption.
  Qed.

  Definition point_enc_coordinates (p : (F * F)) : Word.word (S sz) := let '(x,y) := p in
    Word.WS (sign_bit x) (enc y).

  Let point_enc (p : point) : Word.word (S sz) := point_enc_coordinates (E.coordinates p).

  Definition point_dec_coordinates (w : Word.word (S sz)) : option (F * F) :=
    match dec (Word.wtl w) with
    | None => None
    | Some y => let x2 := solve_for_x2 y in
        let x := sqrt x2 in
        if eq_dec (x ^ 2) x2
        then
          let p := (if Bool.eqb (whd w) (sign_bit x) then x else opp x, y) in
          if (andb (F_eqb x 0) (whd w))
          then None (* special case for 0, since its opposite has the same sign; if the sign bit of 0 is 1, produce None.*)
          else Some p
        else None
    end.

  (* Definition of product equality parameterized over equality of underlying types *)
  Definition prod_eq {A B} eqA eqB (x y : (A * B)) := let (xA,xB) := x in let (yA,yB) := y in
    (eqA xA yA) /\ (eqB xB yB).

  Lemma prod_eq_dec : forall {A eq} (A_eq_dec : forall a a' : A, {eq a a'} + {not (eq a a')})
   (x y : (A * A)), {prod_eq eq eq x y} + {not (prod_eq eq eq x y)}.
  Proof.
    intros.
    destruct x as [x1 x2].
    destruct y as [y1 y2].
    match goal with
    | |- {prod_eq _ _ (?x1, ?x2) (?y1,?y2)} + {not (prod_eq _ _ (?x1, ?x2) (?y1,?y2))} =>
      destruct (A_eq_dec x1 y1); destruct (A_eq_dec x2 y2) end;
      unfold prod_eq; intuition.
  Qed.

  Definition option_eq {A} eq (x y : option A) :=
    match x with
    | None    => y = None
    | Some ax => match y with
                 | None => False
                 | Some ay => eq ax ay
                 end
    end.

  Lemma option_eq_dec : forall {A eq} (A_eq_dec : forall a a' : A, {eq a a'} + {not (eq a a')})
    (x y : option A), {option_eq eq x y} + {not (option_eq eq x y)}.
  Proof.
    unfold option_eq; intros; destruct x as [ax|], y as [ay|]; try tauto; auto.
    right; congruence.
  Qed.
  Definition option_coordinates_eq := option_eq (prod_eq eq eq).

  Lemma option_coordinates_eq_NS : forall x, option_coordinates_eq None (Some x) -> False.
  Proof.
    unfold option_coordinates_eq, option_eq.
    intros; discriminate.
  Qed.

  Lemma inversion_option_coordinates_eq : forall x y,
    option_coordinates_eq (Some x) (Some y) -> prod_eq eq eq x y.
  Proof.
    unfold option_coordinates_eq, option_eq; intros; assumption.
  Qed.

  Lemma prod_eq_onCurve : forall p q : F * F, prod_eq eq eq p q ->
    onCurve p -> onCurve q.
  Proof.
    unfold prod_eq; intros.
    destruct p; destruct q.
    eauto using onCurve_subst.
  Qed.

  Lemma option_coordinates_eq_iff : forall x1 x2 y1 y2,
    option_coordinates_eq (Some (x1,y1)) (Some (x2,y2)) <-> and (eq x1 x2) (eq y1 y2).
  Proof.
    unfold option_coordinates_eq, option_eq, prod_eq; tauto.
  Qed.

  Definition point_eq (p q : point) : Prop := prod_eq eq eq (proj1_sig p) (proj1_sig q).
  Definition option_point_eq := option_eq (point_eq).

  Lemma option_point_eq_iff : forall p q,
    option_point_eq (Some p) (Some q) <->
    option_coordinates_eq (Some (proj1_sig p)) (Some (proj1_sig q)).
  Proof.
    unfold option_point_eq, option_coordinates_eq, option_eq, point_eq; intros.
    reflexivity.
  Qed.

  Lemma option_coordinates_eq_dec : forall p q,
     {option_coordinates_eq p q} + {~ option_coordinates_eq p q}.
  Proof.
    intros.
    apply option_eq_dec.
    apply prod_eq_dec.
    apply eq_dec.
  Qed.

  Lemma point_eq_dec : forall p q, {point_eq p q} + {~ point_eq p q}.
  Proof.
    intros.
    apply prod_eq_dec.
    apply eq_dec.
  Qed.

  Lemma option_point_eq_dec : forall p q,
     {option_point_eq p q} + {~ option_point_eq p q}.
  Proof.
    intros.
    apply option_eq_dec.
    apply point_eq_dec.
  Qed.

  Lemma prod_eq_trans : forall p q r, prod_eq eq eq p q -> prod_eq eq eq q r ->
    prod_eq eq eq p r.
  Proof.
    unfold prod_eq; intros.
    repeat break_let.
    intuition; etransitivity; eauto.
  Qed.

  Lemma option_coordinates_eq_trans : forall p q r, option_coordinates_eq p q ->
    option_coordinates_eq q r -> option_coordinates_eq p r.
  Proof.
    unfold option_coordinates_eq, option_eq; intros.
    repeat break_match; subst; congruence || eauto using prod_eq_trans.
  Qed.

  Lemma prod_eq_sym : forall p q, prod_eq eq eq p q -> prod_eq eq eq q p.
  Proof.
    unfold prod_eq; intros.
    repeat break_let.
    intuition auto with relations; etransitivity; eauto.
  Qed.

  Lemma option_coordinates_eq_sym : forall p q, option_coordinates_eq p q ->
    option_coordinates_eq q p.
  Proof.
    unfold option_coordinates_eq, option_eq; intros.
    repeat break_match; subst; congruence || eauto using prod_eq_sym; intuition.
  Qed.

  Opaque option_coordinates_eq option_point_eq point_eq option_eq prod_eq.

  Ltac inversion_Some_eq := match goal with [H: Some ?x = Some ?y |- _] => inversion H; subst end.

  Ltac congruence_option_coord := exfalso; eauto using option_coordinates_eq_NS.

  Lemma point_dec_coordinates_onCurve : forall w p, option_coordinates_eq  (point_dec_coordinates w) (Some p) -> onCurve p.
  Proof.
    unfold point_dec_coordinates; intros.
    edestruct dec; [ | congruence_option_coord ].
    break_if; [ | congruence_option_coord].
    break_if; [ congruence_option_coord | ].
    apply E.solve_correct in e.
    break_if; eapply prod_eq_onCurve;
      eauto using inversion_option_coordinates_eq, solve_onCurve, solve_opp_onCurve.
  Qed.

  Definition point_dec' w p : option point :=
    match (option_coordinates_eq_dec  (point_dec_coordinates w) (Some p)) with
      | left EQ => Some (exist _ p (point_dec_coordinates_onCurve w p EQ))
      | right _ => None (* this case is never reached *)
    end.

  Definition point_dec (w : word (S sz)) : option point :=
    match point_dec_coordinates w with
    | Some p => point_dec' w p
    | None => None
    end.

  Lemma point_coordinates_encoding_canonical : forall w p,
    point_dec_coordinates w = Some p -> point_enc_coordinates p = w.
  Proof.
    unfold point_dec_coordinates, point_enc_coordinates; intros ? ? coord_dec_Some.
    case_eq (dec (wtl w)); [ intros ? dec_Some | intros dec_None; rewrite dec_None in *; congruence ].
    destruct p.
    rewrite (shatter_word w).
    f_equal; rewrite dec_Some in *;
      do 2 (break_if; try congruence); inversion coord_dec_Some; subst.
    + destruct (eq_dec (sqrt (solve_for_x2 f1)) 0) as [sqrt_0 | ?].
      - break_if; rewrite sign_bit_zero in * by (assumption || (rewrite sqrt_0; ring));
          auto using Bool.eqb_prop.
        apply F_eqb_iff in sqrt_0.
        rewrite sqrt_0 in *.
        destruct (whd w); inversion Heqb0; auto.
      - break_if.
        symmetry; auto using Bool.eqb_prop.
        rewrite <- sign_bit_opp by assumption.
        destruct (whd w); inversion Heqb0; break_if; auto.
    + inversion coord_dec_Some; subst.
      auto using encoding_canonical.
  Qed.

  Lemma inversion_point_dec : forall w x, point_dec w = Some x ->
    point_dec_coordinates w = Some (E.coordinates x).
  Proof.
    unfold point_dec, E.coordinates; intros.
    break_match; [ | congruence].
    unfold point_dec' in *; break_match; [ | congruence].
    match goal with [ H : Some _ = Some _ |- _ ] => inversion H end.
    reflexivity.
  Qed.

  Lemma point_encoding_canonical : forall w x, point_dec w = Some x -> point_enc x = w.
  Proof.
    unfold point_enc; intros.
    apply point_coordinates_encoding_canonical.
    auto using inversion_point_dec.
  Qed.

  Lemma y_decode : forall p, dec (wtl (point_enc_coordinates p)) = Some (snd p).
  Proof.
    intros; destruct p. cbv [point_enc_coordinates wtl snd].
    exact (encoding_valid _).
  Qed.

  Lemma F_eqb_false : forall x y, x !== y -> F_eqb x y = false.
  Proof.
    intros; unfold F_eqb; destruct (eq_dec x y); congruence.
  Qed.

  Lemma eqb_sign_opp_r : forall x y, (y !== 0) ->
    Bool.eqb (sign_bit x) (sign_bit y) = false ->
    Bool.eqb (sign_bit x) (sign_bit (opp y)) = true.
  Proof.
    intros x y y_nonzero ?.
    specialize (sign_bit_opp y y_nonzero).
    destruct (sign_bit x), (sign_bit y); try discriminate;
      rewrite <-sign_bit_opp; auto.
  Qed.

  Lemma sign_match : forall x y sqrt_y, sqrt_y !== 0 -> (x ^2) == y -> sqrt_y ^2 == y ->
    Bool.eqb (sign_bit x) (sign_bit sqrt_y) = true ->
    sqrt_y == x.
  Proof.
    intros.
    pose proof (only_two_square_roots_choice x sqrt_y y) as Hchoice.
    destruct Hchoice; try assumption; symmetry; try assumption.
    rewrite (sign_bit_subst x (opp sqrt_y)) in * by assumption.
    rewrite <-sign_bit_opp in * by assumption.
    rewrite Bool.eqb_negb1 in *; discriminate.
  Qed.

  Lemma point_encoding_coordinates_valid : forall p, onCurve p ->
    option_coordinates_eq (point_dec_coordinates (point_enc_coordinates p)) (Some p).
  Proof.
    intros [x y] onCurve_p.
    unfold point_dec_coordinates.
    rewrite y_decode.
    cbv [whd point_enc_coordinates snd].
    pose proof (square_sqrt (solve_for_x2 y) x) as solve_sqrt_ok.
    forward solve_sqrt_ok. {
      symmetry.
      apply E.solve_correct.
      assumption.
    }
    match goal with [ H1 : ?P, H2 : ?P -> _ |- _ ] => specialize (H2 H1); clear H1 end.
    unfold sqrt_ok in solve_sqrt_ok.
    break_if; [ |  congruence].
    assert (solve_for_x2 y == (x ^2)) as solve_correct by (symmetry; apply E.solve_correct; assumption).
    destruct (eq_dec x 0) as [eq_x_0 | neq_x_0].
    + rewrite !sign_bit_zero by
        (eauto || (rewrite eq_x_0 in *; rewrite sqrt_square; [ | eauto]; reflexivity)).
      rewrite Bool.andb_false_r, Bool.eqb_reflx.
      apply option_coordinates_eq_iff; split; try reflexivity.
      transitivity (sqrt (x ^2)); auto.
      apply (sqrt_square); reflexivity.
    + rewrite F_eqb_false, Bool.andb_false_l by (rewrite sqrt_square; [ | eauto]; assumption).
      break_if; [ | apply eqb_sign_opp_r in Heqb];
        try (apply option_coordinates_eq_iff; split; try reflexivity);
        try eapply sign_match with (y := solve_for_x2 y); eauto;
        try solve [symmetry; auto]; rewrite ?square_opp; auto;
        (rewrite sqrt_square; [ | eauto]); try apply Ring.opp_nonzero_nonzero;
        assumption.
Qed.

Lemma point_dec'_valid : forall p q, option_coordinates_eq (Some q) (Some (proj1_sig p)) ->
  option_point_eq (point_dec' (point_enc_coordinates (proj1_sig p)) q) (Some p).
Proof.
  unfold point_dec'; intros.
  break_match.
  + f_equal.
    apply option_point_eq_iff.
    destruct p as [ [ ? ? ] ? ]; simpl in *.
    assumption.
  + exfalso; apply  n.
    eapply option_coordinates_eq_trans; [ | eauto using option_coordinates_eq_sym ].
    apply point_encoding_coordinates_valid.
    apply (proj2_sig p).
Qed.

Lemma point_encoding_valid : forall p,
  option_point_eq (point_dec (point_enc p)) (Some p).
Proof.
  intros.
  unfold point_dec.
  replace (point_enc p) with (point_enc_coordinates (proj1_sig p)) by reflexivity.
  break_match.
  + eapply (point_dec'_valid p).
    rewrite <-Heqo.
    apply point_encoding_coordinates_valid.
    apply (proj2_sig p).
  + exfalso.
    eapply option_coordinates_eq_NS.
    pose proof (point_encoding_coordinates_valid _ (proj2_sig p)).
    rewrite Heqo in *.
    eassumption.
Qed.

End PointEncodingPre.