(*** XXX TODO MOVE ALL THINGS IN THIS FILE TO BETTER PLACES *)
Require Import Coq.ZArith.BinInt.
Require Import Coq.Classes.Morphisms.
Require Import Crypto.Util.Tuple.
Require Import Crypto.Util.BoundedWord.
Require Import Crypto.Util.Sigma.Lift.
Require Import Crypto.Util.Sigma.Associativity.
Require Import Crypto.Util.Sigma.MapProjections.
Require Import Crypto.Util.ZRange.
Require Import Crypto.Util.ZUtil.Tactics.LtbToLt.
Require Import Crypto.Util.Tactics.MoveLetIn.
Require Import Crypto.Util.Tactics.DestructHead.

Definition adjust_tuple2_tuple2_sig {A P Q}
           (v : { a : { a : tuple (tuple A 2) 2 | (P (fst (fst a)) /\ P (snd (fst a))) /\ (P (fst (snd a)) /\ P (snd (snd a))) }
                | Q (exist _ _ (proj1 (proj1 (proj2_sig a))),
                     exist _ _ (proj2 (proj1 (proj2_sig a))),
                     (exist _ _ (proj1 (proj2 (proj2_sig a))),
                      exist _ _ (proj2 (proj2 (proj2_sig a))))) })
  : { a : tuple (tuple (@sig A P) 2) 2 | Q a }.
Proof.
  eexists.
  exact (proj2_sig v).
Defined.

(** TODO MOVE ME *)
(** The [eexists_sig_etransitivity_R R] tactic takes a goal of the form
    [{ a | R (f a) b }], and splits it into two goals, [R ?b' b] and
    [{ a | R (f a) ?b' }], where [?b'] is a fresh evar. *)
Definition sig_R_trans_exist1 {B} (R : B -> B -> Prop) {HT : Transitive R} {A} (f : A -> B)
           (b b' : B)
           (pf : R b' b)
           (y : { a : A | R (f a) b' })
  : { a : A | R (f a) b }
  := let 'exist a p := y in exist _ a (transitivity (R:=R) p pf).
Ltac eexists_sig_etransitivity_R R :=
  lazymatch goal with
  | [ |- @sig ?A ?P ]
    => let RT := type of R in
       let B := lazymatch (eval hnf in RT) with ?B -> _ => B end in
       let lem := constr:(@sig_R_trans_exist1 B R _ A _ _ : forall b' pf y, @sig A P) in
       let lem := open_constr:(lem _) in
       simple refine (lem _ _)
  end.
Tactic Notation "eexists_sig_etransitivity_R" open_constr(R) := eexists_sig_etransitivity_R R.
(** The [eexists_sig_etransitivity] tactic takes a goal of the form
      [{ a | f a = b }], and splits it into two goals, [?b' = b] and
      [{ a | f a = ?b' }], where [?b'] is a fresh evar. *)
Definition sig_eq_trans_exist1 {A B}
  := @sig_R_trans_exist1 B (@eq B) _ A.
Ltac eexists_sig_etransitivity :=
  lazymatch goal with
  | [ |- { a : ?A | @?f a = ?b } ]
    => let lem := open_constr:(@sig_eq_trans_exist1 A _ f b _) in
       simple refine (lem _ (_ : { a : A | _ }))
  end.
Definition sig_R_trans_rewrite_fun_exist1 {B} (R : B -> B -> Prop) {HT : Transitive R}
{A} (f : A -> B) (b : B) (f' : A -> B)
           (pf : forall a, R (f a) (f' a))
           (y : { a : A | R (f' a) b })
  : { a : A | R (f a) b }
  := let 'exist a p := y in exist _ a (transitivity (R:=R) (pf a) p).
Ltac eexists_sig_etransitivity_for_rewrite_fun_R R :=
  lazymatch goal with
  | [ |- @sig ?A ?P ]
    => let RT := type of R in
       let B := lazymatch (eval hnf in RT) with ?B -> _ => B end in
       let lem := constr:(@sig_R_trans_rewrite_fun_exist1 B R _ A _ _ : forall f' pf y, @sig A P) in
       let lem := open_constr:(lem _) in
       simple refine (lem _ _); cbv beta
  end.
Tactic Notation "eexists_sig_etransitivity_for_rewrite_fun_R" open_constr(R)
  := eexists_sig_etransitivity_for_rewrite_fun_R R.
Definition sig_eq_trans_rewrite_fun_exist1 {A B} (f f' : A -> B)
           (b : B)
           (pf : forall a, f' a = f a)
           (y : { a : A | f' a = b })
  : { a : A | f a = b }
  := let 'exist a p := y in exist _ a (eq_trans (eq_sym (pf a)) p).
Ltac eexists_sig_etransitivity_for_rewrite_fun :=
  lazymatch goal with
  | [ |- { a : ?A | @?f a = ?b } ]
    => let lem := open_constr:(@sig_eq_trans_rewrite_fun_exist1 A _ f _ b) in
       simple refine (lem _ _); cbv beta
  end.
Definition sig_conj_by_impl2 {A} {P Q : A -> Prop} (H : forall a : A, Q a -> P a)
           (H' : { a : A | Q a })
  : { a : A | P a /\ Q a }
  := let (a, p) := H' in exist _ a (conj (H a p) p).


(** [apply_lift_sig] picks out which version of the [liftN_sig] lemma
    to apply, and builds the appropriate arguments *)
Ltac make_P_for_apply_lift_sig P :=
  lazymatch P with
  | fun (f : ?F) => forall (a : ?A), @?P f a
    => constr:(fun (a : A)
               => ltac:(lazymatch constr:(fun (f : F)
                                          => ltac:(let v := (eval cbv beta in (P f a)) in
                                                   lazymatch (eval pattern (f a) in v) with
                                                   | ?k _ => exact k
                                                   end))
                        with
                        | fun _ => ?P
                          => let v := make_P_for_apply_lift_sig P in
                             exact v
                        end))
  | _ => P
  end.
Ltac apply_lift_sig :=
  let P := lazymatch goal with |- sig ?P => P end in
  let P := make_P_for_apply_lift_sig P in
  lazymatch goal with
  | [ |- { f | forall a b c d e, _ } ]
    => fail "apply_lift_sig does not yet support ≥ 5 binders"
  | [ |- { f | forall (a : ?A) (b : ?B) (c : ?C) (d : ?D), _ } ]
    => apply (@lift4_sig A B C D _ P)
  | [ |- { f | forall (a : ?A) (b : ?B) (c : ?C), _ } ]
    => apply (@lift3_sig A B C _ P)
  | [ |- { f | forall (a : ?A) (b : ?B), _ } ]
    => apply (@lift2_sig A B _ P)
  | [ |- { f | forall (a : ?A), _ } ]
    => apply (@lift1_sig A _ P)
  | [ |- { f | _ } ]
    => idtac
  end.
Ltac get_proj2_sig_map_arg_helper P :=
  lazymatch P with
  | (fun e => ?A -> @?B e)
    => let B' := get_proj2_sig_map_arg_helper B in
       uconstr:(A -> B')
  | _ => uconstr:(_ : Prop)
  end.
Ltac get_proj2_sig_map_arg _ :=
  lazymatch goal with
  | [ |- { e : ?T | @?E e } ]
    => let P := get_proj2_sig_map_arg_helper E in
       uconstr:(fun e : T => P)
  end.
Ltac get_phi1_for_preglue _ :=
  lazymatch goal with
  | [ |- { e | @?E e } ]
    => lazymatch E with
       | context[Tuple.map (Tuple.map ?phi) _ = _]
         => phi
       | context[?phi _ = _]
         => phi
       end
  end.
Ltac get_phi2_for_preglue _ :=
  lazymatch goal with
  | [ |- { e | @?E e } ]
    => lazymatch E with
       | context[_ = ?f (Tuple.map ?phi _)]
         => phi
       | context[_ = ?f (?phi _)]
         => phi
       | context[_ = ?phi _]
         => phi
       end
  end.
Ltac start_preglue :=
  apply_lift_sig; intros; cbv beta iota zeta;
  let phi := get_phi1_for_preglue () in
  let phi2 := get_phi2_for_preglue () in
  let P' := get_proj2_sig_map_arg () in
  refine (proj2_sig_map (P:=P') _ _);
  [ let FINAL := fresh "FINAL" in
    let a := fresh "a" in
    intros a FINAL;
    repeat (let H := fresh in intro H; specialize (FINAL H));
    lazymatch goal with
    | [ |- ?phi _ = ?RHS ]
      => refine (@eq_trans _ _ _ RHS FINAL _); cbv [phi phi2]; clear a FINAL
    | [ |- _ /\ Tuple.map (Tuple.map ?phi) _ = _ ]
      => split; cbv [phi phi2]; [ refine (proj1 FINAL); shelve | ]
    end
  | cbv [phi phi2] ].
Ltac do_set_sig f_sig :=
  let fZ := fresh f_sig in
  set (fZ := proj1_sig f_sig);
  context_to_dlet_in_rhs fZ;
  try cbv beta iota delta [proj1_sig f_sig] in fZ;
  cbv beta delta [fZ]; clear fZ;
  cbv beta iota delta [fst snd].
Ltac do_set_sig_1arg f_sig :=
  let fZ := fresh f_sig in
  set (fZ := proj1_sig f_sig);
  context_to_dlet_in_rhs (fZ _);
  try cbn beta iota delta [proj1_sig f_sig] in fZ;
  try cbv [f_sig] in fZ;
  cbv beta delta [fZ]; clear fZ;
  cbv beta iota delta [fst snd].
Ltac do_set_sigs _ :=
  lazymatch goal with
  | [ |- context[@proj1_sig ?a ?b ?f_sig] ]
    => let fZ := fresh f_sig in
       set (fZ := proj1_sig f_sig);
       context_to_dlet_in_rhs fZ;
       do_set_sigs (); (* we recurse before unfolding, because that's faster *)
       try cbv beta iota delta [proj1_sig f_sig] in fZ;
       cbv beta delta [fZ];
       cbv beta iota delta [fst snd]
  | _ => idtac
  end.
Ltac trim_after_do_rewrite_with_sig _ :=
  repeat match goal with
         | [ |- Tuple.map ?f _ = Tuple.map ?f _ ]
           => apply f_equal
         end.
Ltac do_rewrite_with_sig_no_set_by f_sig by_tac :=
  let lem := constr:(proj2_sig f_sig) in
  let lemT := type of lem in
  let lemT := (eval cbv beta zeta in lemT) in
  rewrite <- (lem : lemT) by by_tac ();
  trim_after_do_rewrite_with_sig ().
Ltac do_rewrite_with_sig_by f_sig by_tac :=
  do_rewrite_with_sig_no_set_by f_sig by_tac;
  do_set_sig f_sig.
Ltac do_rewrite_with_sig_1arg_by f_sig by_tac :=
  do_rewrite_with_sig_no_set_by f_sig by_tac;
  do_set_sig_1arg f_sig.
Ltac do_rewrite_with_sig f_sig := do_rewrite_with_sig_by f_sig ltac:(fun _ => idtac).
Ltac do_rewrite_with_sig_1arg f_sig := do_rewrite_with_sig_1arg_by f_sig ltac:(fun _ => idtac).
Ltac do_rewrite_with_1sig_add_carry_by f_sig carry_sig by_tac :=
  let fZ := fresh f_sig in
  rewrite <- (proj2_sig f_sig) by by_tac ();
  symmetry; rewrite <- (proj2_sig carry_sig) by by_tac (); symmetry;
  pose (fun a => proj1_sig carry_sig (proj1_sig f_sig a)) as fZ;
  lazymatch goal with
  | [ |- context G[proj1_sig carry_sig (proj1_sig f_sig ?a)] ]
    => let G' := context G[fZ a] in change G'
  end;
  context_to_dlet_in_rhs fZ; cbv beta delta [fZ];
  try cbv beta iota delta [proj1_sig f_sig carry_sig];
  cbv beta iota delta [fst snd].
Ltac do_rewrite_with_1sig_add_carry f_sig carry_sig := do_rewrite_with_1sig_add_carry_by f_sig carry_sig ltac:(fun _ => idtac).
Ltac do_rewrite_with_2sig_add_carry_by f_sig carry_sig by_tac :=
  let fZ := fresh f_sig in
  rewrite <- (proj2_sig f_sig) by by_tac ();
  symmetry; rewrite <- (proj2_sig carry_sig) by by_tac (); symmetry;
  pose (fun a b => proj1_sig carry_sig (proj1_sig f_sig a b)) as fZ;
  lazymatch goal with
  | [ |- context G[proj1_sig carry_sig (proj1_sig f_sig ?a ?b)] ]
    => let G' := context G[fZ a b] in change G'
  end;
  context_to_dlet_in_rhs fZ; cbv beta delta [fZ];
  try cbv beta iota delta [proj1_sig f_sig carry_sig];
  cbv beta iota delta [fst snd].
Ltac do_rewrite_with_2sig_add_carry f_sig carry_sig := do_rewrite_with_2sig_add_carry_by f_sig carry_sig ltac:(fun _ => idtac).
Ltac unmap_map_tuple _ :=
  repeat match goal with
         | [ |- context[Tuple.map (n:=?N) (fun x : ?T => ?f (?g x))] ]
           => rewrite <- (Tuple.map_map (n:=N) f g
                          : pointwise_relation _ eq _ (Tuple.map (n:=N) (fun x : T => f (g x))))
         end.
Ltac get_feW_bounded boundedT :=
  lazymatch boundedT with
  | and ?X ?Y => get_feW_bounded X
  | ?feW_bounded _ => feW_bounded
  end.
Ltac subst_feW _ :=
  let T := lazymatch goal with |- @sig ?T _ => T end in
  let boundedT := lazymatch goal with |- { e | ?A -> _ } => A end in
  let feW_bounded := get_feW_bounded boundedT in
  let feW := lazymatch type of feW_bounded with ?feW -> Prop => feW end in
  cbv [feW feW_bounded];
  try clear feW feW_bounded.
Ltac finish_conjoined_preglue _ :=
  [ > match goal with
      | [ FINAL : _ /\ ?e |- _ ] => is_evar e; refine (proj2 FINAL)
      end
  | try subst_feW () ].
Ltac fin_preglue :=
  [ > reflexivity
  | repeat sig_dlet_in_rhs_to_context;
    apply (fun f => proj2_sig_map (fun THIS_NAME_MUST_NOT_BE_UNDERSCORE_TO_WORK_AROUND_CONSTR_MATCHING_ANAOMLIES___BUT_NOTE_THAT_IF_THIS_NAME_IS_LOWERCASE_A___THEN_REIFICATION_STACK_OVERFLOWS___AND_I_HAVE_NO_IDEA_WHATS_GOING_ON p => f_equal f p)) ].

Ltac factor_out_bounds_and_strip_eval op_bounded op_sig_side_conditions_t :=
  let feBW_small := lazymatch goal with |- { f : ?feBW_small | _ } => feBW_small end in
  Associativity.sig_sig_assoc;
  apply sig_conj_by_impl2;
  [ let H := fresh in
    intros ? H;
    try lazymatch goal with
        | [ |- (?eval _ < _)%Z ]
          => cbv [eval]
        end;
    rewrite H; clear H;
    eapply Z.le_lt_trans;
      [ apply Z.eq_le_incl, f_equal | apply op_bounded ];
      [ repeat match goal with
               | [ |- ?f ?x = ?g ?y ]
                 => is_evar y; unify x y;
                    apply (f_equal (fun fg => fg x))
               end;
        clear; abstract reflexivity
      | .. ];
      op_sig_side_conditions_t ()
  | apply (fun f => proj2_sig_map (fun THIS_NAME_MUST_NOT_BE_UNDERSCORE_TO_WORK_AROUND_CONSTR_MATCHING_ANAOMLIES___BUT_NOTE_THAT_IF_THIS_NAME_IS_LOWERCASE_A___THEN_REIFICATION_STACK_OVERFLOWS___AND_I_HAVE_NO_IDEA_WHATS_GOING_ON p => f_equal f p));
    cbv [proj1_sig];
    repeat match goal with
           | [ H : feBW_small |- _ ] => destruct H as [? _]
           end ].

Ltac op_sig_side_conditions_t _ :=
  try (hnf; rewrite <- (ZRange.is_bounded_by_None_repeat_In_iff_lt _ _ _)); destruct_head_hnf' sig; try assumption.

Local Open Scope Z_scope.
(* XXX TODO: Clean this up *)
Ltac nonzero_preglue op_sig cbv_runtime :=
  let phi := lazymatch goal with
             | [ |- context[Decidable.dec (?phi _ = _)] ] => phi
             end in
  let do_red _ :=
      lazymatch (eval cbv [phi] in phi) with
      | (fun x => ?montgomery_to_F (?meval (?feBW_of_feBW_small _)))
        => cbv [feBW_of_feBW_small phi meval]
      end in
  lazymatch goal with
  | [ |- @sig (?feBW_small -> BoundedWord 1 _ ?bound1) _ ]
    => let b1 := (eval vm_compute in bound1) in
       change bound1 with b1
  end;
  apply_lift_sig; intros; eexists_sig_etransitivity;
  do_red ();
  [ refine (_ : (if Decidable.dec (_ = 0) then true else false) = _);
    lazymatch goal with
    | [ |- (if Decidable.dec ?x then _ else _) = (if Decidable.dec ?y then _ else _) ]
      => cut (x <-> y);
         [ destruct (Decidable.dec x), (Decidable.dec y); try reflexivity; intros [? ?];
           generalize dependent x; generalize dependent y; solve [ intuition congruence ]
         | ]
    end;
    etransitivity; [ | eapply (proj2_sig op_sig) ];
    [ | solve [ op_sig_side_conditions_t () ].. ];
    reflexivity
  | ];
  let decP := lazymatch goal with |- { c | _ = if Decidable.dec (?decP = 0) then _ else _ } => decP end in
  apply (@proj2_sig_map _ (fun c => BoundedWordToZ 1 _ _ c = decP) _);
  [ let a' := fresh "a'" in
    let H' := fresh "H'" in
    intros a' H'; rewrite H';
    let H := fresh in
    lazymatch goal with |- context[Decidable.dec ?x] => destruct (Decidable.dec x) as [H|H]; try rewrite H end;
    [ reflexivity
    | let H := fresh in
      lazymatch goal with |- context[?x =? 0] => destruct (x =? 0) eqn:? end;
      try reflexivity;
      Z.ltb_to_lt; congruence ]
  | ];
  eexists_sig_etransitivity;
  [ do_set_sig op_sig; cbv_runtime (); reflexivity
  | ];
  sig_dlet_in_rhs_to_context;
  cbv [proj1_sig];
  match goal with
  | [ |- context[match ?v with exist _ _ => _ end] ]
    => is_var v; destruct v as [? _]
  end.