Require Import UniMath.Foundations.All.
Require Import UniMath.MoreFoundations.All.
Require Import UniMath.Combinatorics.Tuples.
Require Import UniMath.Combinatorics.StandardFiniteSets.
Require Import UniMath.Combinatorics.Vectors.

Require Import UniMath.AlgebraicTheories.LambdaTheories.
Require Import UniMath.AlgebraicTheories.AlgebraicTheories.

Require Import Ltac2.Ltac2.

Local Open Scope vec.
Local Open Scope stn.
Local Open Scope algebraic_theories.
Local Open Scope lambda_calculus.

Ltac2 Type state := {
  left : string list;
  right : string list;
  side : int;
}.

Ltac2 print_refine (s : state) (t : string) :=
  let (short: bool) := match s.(left), s.(right) with
    | [], [] => true
    | _, _ => false
  end in
  let (beginning, ending) := match (s.(side)), short with
    | 0, true => ("refine '(", " @ _).")
    | 0, false => ("refine '(maponpaths (λ x, ", ") @ _).")
    | _, true => ("refine '(_ @ !", ").")
    | _, false => ("refine '(_ @ !maponpaths (λ x, ", ")).")
  end in
  let navigation := match short with
    | true => []
    | false => [
        String.concat "" (List.rev (s.(left))) ;
        "x" ;
        String.concat "" (s.(right)) ;
        ") ("
      ]
  end in
  Message.print (Message.of_string (String.concat "" [
    beginning ;
    String.concat "" navigation ;
    t ;
    ending
  ])).

Ltac2 mutable traversals () : ((unit -> unit) * string * string) list := [].

Ltac2 rec traverse_left (s : state) (t : state -> unit) () :=
  t s;
  List.iter (
    fun (m, l, r) =>
    try (m () ; Control.focus 2 2 (
      traverse_left {
        s with
        left := (l :: "(" :: s.(left));
        right := (r :: ")" :: s.(right))
      } t
    ))) (List.rev (traversals ()));
  apply idpath.

Ltac2 traverse (t : state -> unit) :=
  refine '(_ @ !_);
  Control.focus 2 2 (traverse_left {
    side := 0;
    left := [];
    right := [];
  } t);
  refine '(_ @ !_);
  Control.focus 2 2 (traverse_left {
    side := 1;
    left := [];
    right := [];
  } t).

Ltac2 generate_refine' (p : pattern) (t : string) := traverse
  (fun s => match! goal with
  | [ |- $pattern:p = _] => print_refine s t
  | [ |- _ = _ ] => ()
  end).

Ltac2 Notation "generate_refine" p(pattern) := generate_refine' p.

Ltac2 mutable rewrites () : ((unit -> unit) * string) list := [].

Ltac2 propagate_subst () := repeat (progress (traverse (fun s =>
  List.iter (fun (r, t) => try (r (); print_refine s t)) (List.rev (rewrites ()))
))).

Ltac2 Set traversals as traversals0 := fun _ => (
    (fun () => match! goal with | [ |- abs _ = _ ] => refine '(maponpaths abs _ @ _) end),
    "abs ",
    ""
  ) :: traversals0 ().

Ltac2 Set traversals as traversals0 := fun _ => (
    (fun () => match! goal with | [ |- app _ _ = _ ] => refine '(maponpaths (λ x, app x _) _ @ _) end),
    "app ",
    " _"
  ) :: traversals0 ().

Ltac2 Set traversals as traversals0 := fun _ => (
    (fun () => match! goal with | [ |- app _ _ = _ ] => refine '(maponpaths (λ x, app _ x) _ @ _) end),
    "app _ ",
    ""
  ) :: traversals0 ().

Ltac2 Set traversals as traversals0 := fun _ => (
    (fun () => match! goal with | [ |- _ • _ = _ ] => refine '(maponpaths (λ x, x • _) _ @ _) end),
    "",
    " • _"
  ) :: traversals0 ().

Ltac2 Set traversals as traversals0 := fun _ => (
    (fun () => match! goal with | [ |- inflate _ = _ ] => refine '(maponpaths inflate _ @ _) end),
    "inflate ",
    ""
  ) :: traversals0 ().

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- (var ?a) • ?b = _] => refine '(var_subst _ $a $b)
  end), "var_subst _ _ _") :: rewrites0 ().

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- (abs ?a) • ?b = _] => refine '(subst_abs _ $a $b)
  end), "subst_abs _ _ _") :: rewrites0 ().

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- (app ?a ?b) • ?c = _] => refine '(subst_app _ $a $b $c)
  end), "subst_app _ _ _ _") :: rewrites0 ().

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- (?a • ?b) • ?c = _] => refine '(subst_subst _ $a $b $c)
  end), "subst_subst _ _ _ _") :: rewrites0 ().

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- app (abs ?a) ?b = _] => refine '(beta_equality _ _ $a $b); assumption
  end), "beta_equality _ Lβ _ _") :: rewrites0 ().

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- inflate ?a • ?b = _] => refine '(subst_inflate _ $a $b)
  end), "subst_inflate _ _ _") :: rewrites0 ().

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- inflate (var ?a) = _] => refine '(inflate_var _ $a)
  end), "inflate_var _ _") :: rewrites0 ().

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- inflate (abs ?a) = _] => refine '(inflate_abs _ $a)
  end), "inflate_abs _ _") :: rewrites0 ().

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- inflate (app ?a ?b) = _] => refine '(inflate_app _ $a $b)
  end), "inflate_app _ _ _") :: rewrites0 ().

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- inflate (?a • ?b) = _] => refine '(inflate_subst _ $a $b)
  end), "inflate_subst _ _ _") :: rewrites0 ().

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- extend_tuple ?a ?b (_ (inl ?c)) = _] => refine '(extend_tuple_inl $a $b $c)
  end), "extend_tuple_inl _ _ _") :: rewrites0 ().

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- extend_tuple ?a ?b (_ (inr ?c)) = _] => refine '(extend_tuple_inr $a $b $c)
  end), "extend_tuple_inr _ _ _") :: rewrites0 ().

Definition compose
  {L : lambda_theory}
  {n : nat}
  (a b : L n)
  : L n
  := abs
    (app
      (inflate a)
      (app
        (inflate b)
        (var (stnweq (inr tt))))).

Notation "a ∘ b" :=
  (compose a b)
  (at level 40, left associativity)
  : lambda_calculus.

Ltac2 Set traversals as traversals0 := fun _ => (
    (fun () => match! goal with | [ |- _ ∘ _ = _ ] => refine '(maponpaths (λ x, x ∘ _) _ @ _) end),
    "",
    " ∘ _"
  ) :: traversals0 ().

Ltac2 Set traversals as traversals0 := fun _ => (
    (fun () => match! goal with | [ |- _ ∘ _ = _ ] => refine '(maponpaths (λ x, _ ∘ x) _ @ _) end),
    "_ ∘ ",
    ""
  ) :: traversals0 ().

Lemma subst_compose
  (L : lambda_theory)
  {m n : nat}
  (a b : L m)
  (c : stn m → L n)
  : subst (a ∘ b) c = compose (subst a c) (subst b c).
Show proof.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- (?a ∘ ?b) • ?c = _] => refine '(subst_compose _ $a $b $c)
  end), "subst_compose _ _ _ _") :: rewrites0 ().

Definition inflate_compose
  (L : lambda_theory)
  {n : nat}
  (a b : L n)
  : inflate (a ∘ b) = inflate a ∘ inflate b
  := subst_compose L _ _ _.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- inflate (?a ∘ ?b) = _] => refine '(inflate_compose _ $a $b)
  end), "inflate_compose _ _ _") :: rewrites0 ().

Lemma app_compose
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a b c : L n)
  : app (a ∘ b) c = app a (app b c).
Show proof.

Lemma compose_abs
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a : L n)
  (b : L (S n))
  : a ∘ abs b = abs (app (inflate a) b).
Show proof.

Lemma abs_compose
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a : L (S n))
  (b : L n)
  : abs a ∘ b
  = abs
    (subst
      a
      (extend_tuple
        (λ i, var (stnweq (inl i)))
        (app
          (inflate b)
          (var (stnweq (inr tt)))))).
Show proof.

Lemma compose_assoc
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a b c : L n)
  : a ∘ (b ∘ c) = a ∘ b ∘ c.
Show proof.

Lemma subst_is_compose
  (L : lambda_theory)
  (Lβ : has_β L)
  (A B : L 0)
  : appx A • (λ _, appx B) = appx (A ∘ B).
Show proof.

Lemma compose_is_subst
  (L : lambda_theory)
  (Lβ : has_β L)
  (A B : L 1)
  : abs A ∘ abs B = abs (A • (λ _, B)).
Show proof.

Definition pair
  {L : lambda_theory}
  {n : nat}
  (a b : L n)
  : L n
  := abs
    (app
      (app
        (var (stnweq (inr tt)))
        (inflate a))
      (inflate b)).

Notation "⟨ a , b ⟩" :=
  (pair a b)
  : lambda_calculus.

Ltac2 Set traversals as traversals0 := fun _ => (
    (fun () => match! goal with | [ |- ⟨_, _⟩ = _ ] => refine '(maponpaths (λ x, ⟨x, _⟩) _ @ _) end),
    "⟨",
    ", _⟩"
  ) :: traversals0 ().

Ltac2 Set traversals as traversals0 := fun _ => (
    (fun () => match! goal with | [ |- ⟨_, _⟩ = _ ] => refine '(maponpaths (λ x, ⟨_, x⟩) _ @ _) end),
    "⟨_, ",
    "⟩"
  ) :: traversals0 ().

Lemma subst_pair
  (L : lambda_theory)
  {m n : nat}
  (a b : L m)
  (c : stn m → L n)
  : subst (⟨a, b⟩) c = (⟨subst a c, subst b c⟩).
Show proof.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- (⟨?a , ?b⟩) • ?c = _] => refine '(subst_pair _ $a $b $c)
  end), "subst_pair _ _ _ _") :: rewrites0 ().

Definition inflate_pair
  (L : lambda_theory)
  {n : nat}
  (a b : L n)
  : inflate (⟨a, b⟩) = ⟨inflate a, inflate b⟩
  := subst_pair _ _ _ _.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- inflate (⟨?a , ?b⟩) = _] => refine '(inflate_pair _ $a $b)
  end), "inflate_pair _ _ _") :: rewrites0 ().

Definition pair_arrow
  {L : lambda_theory}
  {n : nat}
  (a b : L n)
  : L n
  := abs
    (pair
      (app
        (inflate a)
        (var (stnweq (inr tt))))
      (app
        (inflate b)
        (var (stnweq (inr tt))))).

Ltac2 Set traversals as traversals0 := fun _ => (
    (fun () => match! goal with | [ |- pair_arrow _ _ = _ ] => refine '(maponpaths (λ x, pair_arrow x _) _ @ _) end),
    "pair_arrow ",
    " _"
  ) :: traversals0 ().

Ltac2 Set traversals as traversals0 := fun _ => (
    (fun () => match! goal with | [ |- pair_arrow _ _ = _ ] => refine '(maponpaths (λ x, pair_arrow _ x) _ @ _) end),
    "pair_arrow _ ",
    ""
  ) :: traversals0 ().

Lemma subst_pair_arrow
  (L : lambda_theory)
  {m n : nat}
  (a b : L m)
  (c : stn m → L n)
  : (pair_arrow a b) • c = pair_arrow (a • c) (b • c).
Show proof.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- (pair_arrow ?a ?b) • ?c = _] => refine '(subst_pair_arrow _ $a $b $c)
  end), "subst_pair_arrow _ _ _ _") :: rewrites0 ().

Definition inflate_pair_arrow
  (L : lambda_theory)
  {n : nat}
  (a b : L n)
  : inflate (pair_arrow a b) = pair_arrow (inflate a) (inflate b)
  := subst_pair_arrow _ _ _ _.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- inflate (pair_arrow ?a ?b) = _] => refine '(inflate_pair_arrow _ $a $b)
  end), "inflate_pair_arrow _ _ _") :: rewrites0 ().

Lemma app_pair_arrow
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a b c : L n)
  : app (pair_arrow a b) c = (⟨app a c, app b c⟩).
Show proof.

Lemma pair_arrow_compose
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a b c : L n)
  : (pair_arrow a b) ∘ c = pair_arrow (a ∘ c) (b ∘ c).
Show proof.

Definition π1
  {L : lambda_theory}
  {n : nat}
  : L n
  := abs
    (app
      (var (stnweq (inr tt)))
      (abs
        (abs
          (var (stnweq (inl (stnweq (inr tt)))))))).

Lemma subst_π1
  (L : lambda_theory)
  {m n : nat}
  (t : stn m → L n)
  : subst π1 t = π1.
Show proof.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- π1 • ?a = _] => refine '(subst_π1 _ $a)
  end), "subst_π1 _ _") :: rewrites0 ().

Definition inflate_π1
  (L : lambda_theory)
  {n : nat}
  : inflate (π1 : L n) = π1
  := subst_π1 _ _.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- inflate π1 = _] => refine '(inflate_π1 _)
  end), "inflate_π1 _") :: rewrites0 ().

Lemma π1_pair
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a b : L n)
  : app π1 (⟨a, b⟩) = a.
Show proof.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- app π1 (⟨?a , ?b⟩) = _] => refine '(π1_pair _ _ $a $b); assumption
  end), "π1_pair _ Lβ _ _") :: rewrites0 ().

Lemma π1_pair_arrow
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a b : L n)
  : π1 ∘ pair_arrow a b = abs (app (inflate a) (var (stnweq (inr tt)))).
Show proof.

Lemma abs_app_abs_var
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a : L (S n))
  : abs
    (app
      (inflate (abs a))
      (var (stnweq (inr tt))))
  = abs a.
Show proof.

Lemma π1_pair_arrow_alt
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a : L (S n))
  (b : L n)
  : π1 ∘ (pair_arrow (abs a) b) = abs a.
Show proof.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- π1 ∘ (pair_arrow ?a ?b) = _] => refine '(π1_pair_arrow_alt _ _ _ $b); assumption
  end), "π1_pair_arrow_alt _ Lβ _ _") :: rewrites0 ().

Definition π2
  {L : lambda_theory}
  {n : nat}
  : L n
  := abs
    (app
      (var (stnweq (inr tt)))
      (abs
        (abs
          (var (stnweq (inr tt)))))).

Lemma subst_π2
  (L : lambda_theory)
  {m n : nat}
  (t : stn m → L n)
  : subst π2 t = π2.
Show proof.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- π2 • ?c = _] => refine '(subst_π2 _ $c)
  end), "subst_π2 _ _") :: rewrites0 ().

Definition inflate_π2
  (L : lambda_theory)
  {n : nat}
  : inflate (π2 : L n) = π2
  := subst_π2 _ _.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- inflate π2 = _] => refine '(inflate_π2 _)
  end), "inflate_π2 _") :: rewrites0 ().

Lemma π2_pair
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a b : L n)
  : app π2 (⟨a, b⟩) = b.
Show proof.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- app π2 (⟨?a , ?b⟩) = _] => refine '(π2_pair _ _ $a $b); assumption
  end), "π2_pair _ Lβ _ _") :: rewrites0 ().

Lemma π2_pair_arrow
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a b : L n)
  : π2 ∘ pair_arrow a b = abs (app (inflate b) (var (stnweq (inr tt)))).
Show proof.

Lemma π2_pair_arrow_alt
  (L : lambda_theory)
  (Lβ : has_β L)
  {n : nat}
  (a : L n)
  (b : L (S n))
  : π2 ∘ pair_arrow a (abs b) = abs b.
Show proof.

Ltac2 Set rewrites as rewrites0 := fun () =>
  ((fun () => match! goal with
  | [ |- π2 ∘ (pair_arrow ?a ?b) = _] => refine '(π2_pair_arrow_alt _ _ $a _); assumption
  end), "π2_pair_arrow_alt _ Lβ _ _") :: rewrites0 ().