import .logic .maps

@[derive decidable_eq]
inductive ty : Type
| bool : ty
| arrow : ty -> ty -> ty
| nat : ty
| prod : ty -> ty -> ty
| unit : ty
| sum : ty -> ty -> ty
| list : ty -> ty

open ty

inductive tm : Type
| var : string -> tm
| abs : string -> ty -> tm -> tm
| app : tm -> tm -> tm
| const : ℕ -> tm
| prd : tm -> tm
| scc : tm -> tm
| mlt : tm -> tm -> tm
| iszro : tm -> tm
| tru : tm
| fls : tm
| tst : tm -> tm -> tm -> tm
| let_ : string -> tm -> tm -> tm
| pair : tm -> tm -> tm
| fst : tm -> tm
| snd : tm -> tm
| unit : tm
| inl : ty -> tm -> tm
| inr : ty -> tm -> tm
| scase : tm -> string -> tm -> string -> tm -> tm
| nil : ty -> tm
| cons : tm -> tm -> tm
| lcase : tm -> tm -> string -> string -> tm -> tm
| fix : tm -> tm

open tm

def idB := abs "x" bool (var "x")
def idBB := abs "x" (arrow bool bool) (var "x")
def idBBBB := abs "x" (arrow (arrow bool bool) (arrow bool bool)) (var "x")
def k := abs "x" bool (abs "y" bool (var "x"))
def notB := abs "x" bool (tst (var "x") fls tru)

inductive value : tm -> Prop
| v_abs {x T t} : value (abs x T t)
| v_const {n} : value (const n)
| v_tru : value tru
| v_fls : value fls
| v_pair {t₁ t₂} : value t₁ -> value t₂ -> value (pair t₁ t₂)
| v_unit : value unit
| v_inl {t T} : value t -> value (inl T t)
| v_inr {t T} : value t -> value (inr T t)
| v_nil {T} : value (nil T)
| v_cons {t₁ t₂} : value t₁ -> value t₂ -> value (cons t₁ t₂)

open value

def subst (x : string) (s : tm) : tm -> tm
| (var y) := if x = y then s else var y
| (abs y T t1) := abs y T (if x = y then t1 else subst t1)
| (app t1 t2) := app (subst t1) (subst t2)
| (const n) := const n
| (prd t) := prd (subst t)
| (scc t) := scc (subst t)
| (mlt t₁ t₂) := mlt (subst t₁) (subst t₂)
| (iszro t) := iszro (subst t)
| tru := tru
| fls := fls
| (tst t1 t2 t3) := tst (subst t1) (subst t2) (subst t3)
| (let_ y t₁ t₂) := let_ y (subst t₁) (if x = y then t₂ else (subst t₂))
| (pair t₁ t₂) := pair (subst t₁) (subst t₂)
| (fst t) := fst (subst t)
| (snd t) := snd (subst t)
| unit := unit
| (inl T t) := inl T (subst t)
| (inr T t) := inr T (subst t)
| (scase t y t₁ z t₂) := scase (subst t)
                               y (if x = y then t₁ else (subst t₁))
                               z (if x = z then t₂ else (subst t₂))
| (nil T) := nil T
| (cons t₁ t₂) := cons (subst t₁) (subst t₂)
| (lcase t t₁ y z t₂) := lcase (subst t)
                               (subst t₁)
                               y z (if x = y ∨ x = z then t₂ else subst t₂)
| (fix t) := fix (subst t)

notation `[` x `:=` s `]` t := subst x s t

inductive substi (s : tm) (x : string) : tm -> tm -> Prop
| s_var1 : substi (var x) s
| s_var2 {y : string} : y ≠ x -> substi (var y) (var y)
| s_abs1 {T : ty} {t : tm} : substi (abs x T t) (abs x T t)
| s_abs2 {y : string} {T : ty} {t t' : tm} :
    y ≠ x -> substi t t' -> substi (abs y T t) (abs y T t')
| s_app {t1 t1' t2 t2' : tm} :
    substi t1 t1' -> substi t2 t2' -> substi (app t1 t2) (app t1' t2')
| s_const {n} : substi (const n) (const n)
| s_prd {t t'} : substi t t' -> substi (prd t) (prd t')
| s_scc {t t'} : substi t t' -> substi (scc t) (scc t')
| s_mlt {t₁ t₁' t₂ t₂'} :
    substi t₁ t₁' -> substi t₂ t₂' -> substi (mlt t₁ t₂) (mlt t₁' t₂')
| s_iszro {t t'} : substi t t' -> substi (iszro t) (iszro t')
| s_tru : substi tru tru
| s_fls : substi fls fls
| s_tst {t1 t1' t2 t2' t3 t3' : tm} :
    substi t1 t1' -> substi t2 t2' -> substi t3 t3' ->
    substi (tst t1 t2 t3) (tst t1' t2' t3')
| s_let1 {t₁ t₁' t₂} : substi t₁ t₁' -> substi (let_ x t₁ t₂) (let_ x t₁' t₂)
| s_let2 {y t₁ t₁' t₂ t₂'} :
    y ≠ x ->
    substi t₁ t₁' ->
    substi t₂ t₂' ->
    substi (let_ y t₁ t₂) (let_ y t₁' t₂')
| s_pair {t₁ t₁' t₂ t₂' } :
    substi t₁ t₁' -> substi t₂ t₂' -> substi (pair t₁ t₂) (pair t₁' t₂')
| s_fst {t t'} : substi t t' -> substi (fst t) (fst t')
| s_snd {t t'} : substi t t' -> substi (snd t) (snd t')
| s_unit : substi unit unit
| s_inl {t t' T} : substi t t' -> substi (inl T t) (inl T t')
| s_inr {t t' T} : substi t t' -> substi (inr T t) (inr T t')
| s_scase1 {t t' t₁ t₂} :
    substi t t' -> substi (scase t x t₁ x t₂) (scase t' x t₁ x t₂)
| s_scase2 {y t t' t₁ t₁' t₂} :
    y ≠ x ->
    substi t t' ->
    substi t₁ t₁' ->
    substi (scase t y t₁ x t₂) (scase t' y t₁' x t₂)
| s_scase3 {z t t' t₁ t₂ t₂'} :
    z ≠ x ->
    substi t t' ->
    substi t₂ t₂' ->
    substi (scase t x t₁ z t₂) (scase t' x t₁ z t₂')
| s_scase4 {y z t t' t₁ t₁' t₂ t₂'} :
    y ≠ x ->
    z ≠ x ->
    substi t t' ->
    substi t₁ t₁' ->
    substi t₂ t₂' ->
    substi (scase t y t₁ z t₂) (scase t' y t₁' z t₂')
| s_nil {T} : substi (nil T) (nil T)
| s_cons {t₁ t₁' t₂ t₂'} :
    substi t₁ t₁' -> substi t₂ t₂' -> substi (cons t₁ t₂) (cons t₁' t₂')
| s_lcase1 {t t' t₁ t₁' t₂} :
    substi t t' ->
    substi t₁ t₁' ->
    substi (lcase t t₁ x x t₂) (lcase t' t₁' x x t₂)
| s_lcase2 {t t' t₁ t₁' y t₂} :
    y ≠ x ->
    substi t t' ->
    substi t₁ t₁' ->
    substi (lcase t t₁ y x t₂) (lcase t' t₁' y x t₂)
| s_lcase3 {t t' t₁ t₁' z t₂} :
    z ≠ x ->
    substi t t' ->
    substi t₁ t₁' ->
    substi (lcase t t₁ x z t₂) (lcase t' t₁' x z t₂)
| s_lcase4 {t t' t₁ t₁' y z t₂ t₂'} :
    y ≠ x ->
    z ≠ x ->
    substi t t' ->
    substi t₁ t₁' ->
    substi t₂ t₂' ->
    substi (lcase t t₁ y z t₂) (lcase t' t₁' y z t₂')
| s_fix {t t'} : substi t t' -> substi (fix t) (fix t')

theorem substi_correct : ∀ x s t t', ([x:=s]t) = t' ↔ substi s x t t' :=
begin
  intros,
  apply iff.intro,
  { intro prf,
    induction prf,
    induction t,
      case tm.var: y {
        unfold subst,
        by_cases h : x = y,
          { simp [*], apply substi.s_var1 },
          { simp [*], apply substi.s_var2, exact ne.symm h },
      },
      case tm.app: { apply substi.s_app; assumption },
      case tm.abs: y {
        unfold subst,
        by_cases h : x = y,
          { simp [*], apply substi.s_abs1 },
          { simp [*], apply substi.s_abs2, exact ne.symm h, assumption, },
      },
      case tm.const: { apply substi.s_const },
      case tm.prd: { apply substi.s_prd, assumption },
      case tm.scc: { apply substi.s_scc, assumption },
      case tm.iszro: { apply substi.s_iszro, assumption },
      case tm.mlt: { apply substi.s_mlt; assumption },
      case tm.tru: { apply substi.s_tru },
      case tm.fls: { apply substi.s_fls },
      case tm.tst: { apply substi.s_tst; assumption },
      case tm.let_: y _ _ ih₁ ih₂ {
        unfold subst,
        by_cases h : x = y,
          { simp [h], rewrite h at ih₁, exact substi.s_let1 ih₁ },
          { simp [h], exact substi.s_let2 (ne.symm h) ih₁ ih₂, },
      },
      case tm.pair: { apply substi.s_pair; assumption },
      case tm.fst: { apply substi.s_fst, assumption },
      case tm.snd: { apply substi.s_snd, assumption },
      case tm.unit: { apply substi.s_unit },
      case tm.inl: { apply substi.s_inl, assumption },
      case tm.inr: { apply substi.s_inr, assumption },
      case tm.scase: _ y _ z _ ih ih₁ ih₂ {
        unfold subst,
        by_cases hxy : x = y,
          { by_cases hxz : x = z,
              { simp [hxy, hxz, eq.trans (symm hxz) hxy],
                rewrite hxy at ih,
                exact substi.s_scase1 ih },
              { simp [hxy, hxz],
                rewrite hxy at hxz ih ih₂,
                exact substi.s_scase3 (ne.symm hxz) ih ih₂ } },
          { by_cases hxz : x = z,
              { simp [hxy, hxz],
                rewrite hxz at hxy ih ih₁,
                exact substi.s_scase2 (ne.symm hxy) ih ih₁ },
              { simp [hxy, hxz],
                exact
                  substi.s_scase4 (ne.symm hxy) (ne.symm hxz) ih ih₁ ih₂ } },
      },
      case tm.nil: { apply substi.s_nil },
      case tm.cons: { apply substi.s_cons; assumption },
      case tm.lcase: _ _ y z _ ih ih₁ ih₂ {
        unfold subst,
        by_cases hxy : x = y,
          { by_cases hxz : x = z,
              { simp [*],
                rewrite <-hxy,
                rewrite <-hxz,
                exact substi.s_lcase1 ih ih₁ },
              { simp [*],
                rewrite <-hxy,
                exact substi.s_lcase3 (ne.symm hxz) ih ih₁ } },
          { by_cases hxz : x = z,
              { simp [*],
                rewrite <-hxz,
                exact substi.s_lcase2 (ne.symm hxy) ih ih₁ },
              { simp [*],
                exact substi.s_lcase4 (ne.symm hxy) (ne.symm hxz) ih ih₁ ih₂ } }
      },
      case tm.fix: _ ih { exact substi.s_fix ih } },
  { intro sub,
    induction sub,
    repeat { simp [*, subst]; simp [ne.symm sub_a]; simp [ne.symm sub_a_1] } },
end

inductive step : tm -> tm -> Prop
| st_appabs {x T t12 v2} : value v2 -> step (app (abs x T t12) v2) ([x:=v2]t12)
| st_app1 {t1 t1' t2} : step t1 t1' -> step (app t1 t2) (app t1' t2)
| st_app2 {v1 t2 t2'} : value v1 -> step t2 t2' -> step (app v1 t2) (app v1 t2')
| st_prdzro : step (prd (const 0)) (const 0)
| st_prdnzr {n} : step (prd (const (nat.succ n))) (const n)
| st_prd {t t'} : step t t' -> step (prd t) (prd t')
| st_sccn {n} : step (scc (const n)) (const (nat.succ n))
| st_scc {t t'} : step t t' -> step (scc t) (scc t')
| st_mlt1 {t₁ t₁' t₂} : step t₁ t₁' -> step (mlt t₁ t₂) (mlt t₁' t₂)
| st_mlt2 {t₁ t₂ t₂'} :
    value t₁ -> step t₂ t₂' -> step (mlt t₁ t₂) (mlt t₁ t₂')
| st_mltnm {n m} : step (mlt (const n) (const m)) (const (n * m))
| st_iszrozro : step (iszro (const 0)) tru
| st_iszronzr {n} : step (iszro (const (nat.succ n))) fls
| st_iszro {t t'} : step t t' -> step (iszro t) (iszro t')
| st_tsttru {t2 t3} : step (tst tru t2 t3) t2
| st_tstfls {t2 t3} : step (tst fls t2 t3) t3
| st_tst {t₁ t₁' t₂ t₃} : step t₁ t₁' -> step (tst t₁ t₂ t₃) (tst t₁' t₂ t₃)
| st_let {t₁ t₁' x t₂} : step t₁ t₁' -> step (let_ x t₁ t₂) (let_ x t₁' t₂)
| st_letvalue {x t₁ t₂} : value t₁ -> step (let_ x t₁ t₂) ([x:=t₁]t₂)
| st_pair1 {t₁ t₁' t₂} : step t₁ t₁' -> step (pair t₁ t₂) (pair t₁' t₂)
| st_pair2 {t₁ t₂ t₂'} :
    value t₁ -> step t₂ t₂' -> step (pair t₁ t₂) (pair t₁ t₂')
| st_fstpair {t₁ t₂} : value t₁ -> value t₂ -> step (fst (pair t₁ t₂)) t₁
| st_fst {t t'} : step t t' -> step (fst t) (fst t')
| st_sndpair {t₁ t₂} : value t₁ -> value t₂ -> step (snd (pair t₁ t₂)) t₂
| st_snd {t t'} : step t t' -> step (snd t) (snd t')
| st_inl {t t' T} : step t t' -> step (inl T t) (inl T t')
| st_inr {t t' T} : step t t' -> step (inr T t) (inr T t')
| st_scase {t t' y t₁ z t₂} :
    step t t' -> step (scase t y t₁ z t₂) (scase t' y t₁ z t₂)
| st_scaseinl {T t y t₁ z t₂} :
    value t -> step (scase (inl T t) y t₁ z t₂) ([y:=t]t₁)
| st_scaseinr {T t y t₁ z t₂} :
    value t -> step (scase (inr T t) y t₁ z t₂) ([z:=t]t₂)
| st_cons1 {t₁ t₁' t₂} : step t₁ t₁' -> step (cons t₁ t₂) (cons t₁' t₂)
| st_cons2 {t₁ t₂ t₂'} :
    value t₁ -> step t₂ t₂' -> step (cons t₁ t₂) (cons t₁ t₂')
| st_lcase {t t' t₁ y z t₂} :
    step t t' -> step (lcase t t₁ y z t₂) (lcase t' t₁ y z t₂)
| st_lcasenil {T t₁ y z t₂} : step (lcase (nil T) t₁ y z t₂) t₁
| st_lcasecons {t_h t_t t₁ y z t₂} :
    value t_h ->
    value t_t ->
    step (lcase (cons t_h t_t) t₁ y z t₂) ([y:=t_h][z:=t_t]t₂)
| st_fix {t t'} : step t t' -> step (fix t) (fix t')
| st_fixabs {x T t} : step (fix (abs x T t)) ([x:=fix (abs x T t)]t)

open step

notation t ` -+> ` t' := step t t'

def multistep := @multi tm step

notation t ` -+>* ` t' := multistep t t'

open multi

example : app idBB idB -+>* idB :=
multi_step (st_appabs v_abs) multi_refl

example : app idBB (app idBB idB) -+>* idB :=
multi_step (st_app2 v_abs (st_appabs v_abs)) $
multi_step (st_appabs v_abs) $
multi_refl

example : app (app idBB notB) tru -+>* fls :=
multi_step (st_app1 (st_appabs v_abs)) $
multi_step (st_appabs v_tru) $
multi_step st_tsttru $
multi_refl

example : app idBB (app notB tru) -+>* fls :=
multi_step (st_app2 v_abs (st_appabs v_tru)) $
multi_step (st_app2 v_abs st_tsttru) $
multi_step (st_appabs v_fls) $
multi_refl

example : app (app idBBBB idBB) idB -+>* idB :=
multi_step (st_app1 (st_appabs v_abs)) $
multi_step (st_appabs v_abs) $
multi_refl

def context := partial_map ty

def context.empty : context := partial_map.empty

inductive has_type : context -> tm -> ty -> Prop
| t_var {gamma : context} {x T} :
    gamma x = some T -> has_type gamma (var x) T
| t_abs {gamma x T₁₁ T₁₂ t₁₂} :
    has_type (partial_map.update x T₁₁ gamma) t₁₂ T₁₂ ->
    has_type gamma (abs x T₁₁ t₁₂) (arrow T₁₁ T₁₂)
| t_app {gamma T₁₁ T₁₂ t₁ t₂} :
    has_type gamma t₁ (arrow T₁₁ T₁₂) ->
    has_type gamma t₂ T₁₁ ->
    has_type gamma (app t₁ t₂) T₁₂
| t_const {gamma n} : has_type gamma (const n) nat
| t_prd {gamma t} : has_type gamma t nat -> has_type gamma (prd t) nat
| t_scc {gamma t} : has_type gamma t nat -> has_type gamma (scc t) nat
| t_mlt {gamma t₁ t₂} :
    has_type gamma t₁ nat ->
    has_type gamma t₂ nat ->
    has_type gamma (mlt t₁ t₂) nat
| t_iszro {gamma t} : has_type gamma t nat -> has_type gamma (iszro t) bool
| t_tru {gamma} : has_type gamma tru bool
| t_fls {gamma} : has_type gamma fls bool
| t_tst {gamma T t₁ t₂ t₃} :
    has_type gamma t₁ bool ->
    has_type gamma t₂ T ->
    has_type gamma t₃ T ->
    has_type gamma (tst t₁ t₂ t₃) T
| t_let {gamma x T₁ T₂ t₁ t₂} :
    has_type gamma t₁ T₁ ->
    has_type (partial_map.update x T₁ gamma) t₂ T₂ ->
    has_type gamma (let_ x t₁ t₂) T₂
| t_pair {gamma t₁ T₁ t₂ T₂} :
    has_type gamma t₁ T₁ ->
    has_type gamma t₂ T₂ ->
    has_type gamma (pair t₁ t₂) (prod T₁ T₂)
| t_fst {gamma t T₁ T₂} :
    has_type gamma t (prod T₁ T₂) -> has_type gamma (fst t) T₁
| t_snd {gamma t T₁ T₂} :
    has_type gamma t (prod T₁ T₂) -> has_type gamma (snd t) T₂
| t_unit {gamma} : has_type gamma unit unit
| t_inl {gamma t T₁ T₂} :
    has_type gamma t T₁ -> has_type gamma (inl T₂ t) (sum T₁ T₂)
| t_inr {gamma t T₂ T₁} :
    has_type gamma t T₂ -> has_type gamma (inr T₁ t) (sum T₁ T₂)
| t_scase {gamma t T₁ T₂ y t₁ T z t₂} :
    has_type gamma t (sum T₁ T₂) ->
    has_type (partial_map.update y T₁ gamma) t₁ T ->
    has_type (partial_map.update z T₂ gamma) t₂ T ->
    has_type gamma (scase t y t₁ z t₂) T
| t_nil {gamma T} : has_type gamma (nil T) (list T)
| t_cons {gamma t₁ t₂ T} :
    has_type gamma t₁ T ->
    has_type gamma t₂ (list T) ->
    has_type gamma (cons t₁ t₂) (list T)
| t_lcase {gamma t T t₁ T' y z t₂} :
    has_type gamma t (list T) ->
    has_type gamma t₁ T' ->
    has_type (partial_map.update z (list T) $ partial_map.update y T gamma)
             t₂
             T' ->
    has_type gamma (lcase t t₁ y z t₂) T'
| t_fix {gamma t T} : has_type gamma t (arrow T T) -> has_type gamma (fix t) T

open has_type

example : has_type context.empty (abs "x" bool (var "x")) (arrow bool bool) :=
t_abs (t_var rfl)

meta def auto_typing : tactic unit :=
tactic.repeat (   tactic.applyc ``t_tru
              <|> tactic.applyc ``t_fls
              <|> tactic.applyc ``t_const
              <|> (tactic.applyc ``t_var >> tactic.reflexivity)
              <|> tactic.applyc ``t_abs
              <|> tactic.applyc ``t_prd
              <|> tactic.applyc ``t_scc
              <|> tactic.applyc ``t_iszro
              <|> tactic.applyc ``t_tst
              <|> tactic.applyc ``t_let )

example : has_type context.empty (abs "x" bool (var "x")) (arrow bool bool) :=
by auto_typing

example {t : ty} : has_type context.empty
                            (abs "x" t
                              (abs "y" (arrow t t)
                                (app (var "y") (app (var "y") (var "x")))))
                            (arrow t (arrow (arrow t t) t)) :=
t_abs (t_abs (t_app (t_var rfl) (t_app (t_var rfl) (t_var rfl))))

example : ∃t, has_type context.empty
                       (abs "x" (arrow bool bool)
                         (abs "y" (arrow bool bool)
                           (abs "z" bool
                             (app (var "y") (app (var "x") (var "z"))))))
                       t :=
⟨ arrow (arrow bool bool)
        (arrow (arrow bool bool)
               (arrow bool bool))
, t_abs (t_abs (t_abs (t_app (t_var rfl) (t_app (t_var rfl) (t_var rfl))))) ⟩

lemma arrow_no_confusion {t₁ t₂} : arrow t₁ t₂ ≠ t₁ :=
assume h : arrow t₁ t₂ = t₁,
begin
  induction t₁,
    case ty.arrow: _ _ ih₁ ih₂ {
      apply ih₁,
      injection h with h' h'',
      rewrite h'',
      exact h',
    },
    repeat { cases h },
end

example : ¬∃s t, has_type context.empty (abs "x" s (app (var "x") (var "x"))) t :=
assume ⟨ _, _, t_abs (t_app (t_var h₁) (t_var h₂)) ⟩,
have h : some (arrow _ _) = some _, from eq.trans (eq.symm h₁) h₂,
option.no_confusion h (assume h', arrow_no_confusion h')

inductive appears_free_in (x : string) : tm -> Prop
| afi_var : appears_free_in (var x)
| afi_abs {y t T} : y ≠ x -> appears_free_in t -> appears_free_in (abs y T t)
| afi_app1 {t₁ t₂} : appears_free_in t₁ -> appears_free_in (app t₁ t₂)
| afi_app2 {t₁ t₂} : appears_free_in t₂ -> appears_free_in (app t₁ t₂)
| afi_prd {t} : appears_free_in t -> appears_free_in (prd t)
| afi_scc {t} : appears_free_in t -> appears_free_in (scc t)
| afi_mlt1 {t₁ t₂} : appears_free_in t₁ -> appears_free_in (mlt t₁ t₂)
| afi_mlt2 {t₁ t₂} : appears_free_in t₂ -> appears_free_in (mlt t₁ t₂)
| afi_iszro {t} : appears_free_in t -> appears_free_in (iszro t)
| afi_tst1 {t₁ t₂ t₃} : appears_free_in t₁ -> appears_free_in (tst t₁ t₂ t₃)
| afi_tst2 {t₁ t₂ t₃} : appears_free_in t₂ -> appears_free_in (tst t₁ t₂ t₃)
| afi_tst3 {t₁ t₂ t₃} : appears_free_in t₃ -> appears_free_in (tst t₁ t₂ t₃)
| afi_let1 {y t₁ t₂} : appears_free_in t₁ -> appears_free_in (let_ y t₁ t₂)
| afi_let2 {y t₁ t₂} :
    y ≠ x -> appears_free_in t₂ -> appears_free_in (let_ y t₁ t₂)
| afi_pair1 {t₁ t₂} : appears_free_in t₁ -> appears_free_in (pair t₁ t₂)
| afi_pair2 {t₁ t₂} : appears_free_in t₂ -> appears_free_in (pair t₁ t₂)
| afi_fst {t} : appears_free_in t -> appears_free_in (fst t)
| afi_snd {t} : appears_free_in t -> appears_free_in (snd t)
| afi_inl {t T} : appears_free_in t -> appears_free_in (inl T t)
| afi_inr {t T} : appears_free_in t -> appears_free_in (inr T t)
| afi_scase1 {t y t₁ z t₂} :
    appears_free_in t -> appears_free_in (scase t y t₁ z t₂)
| afi_scase2 {t y t₁ z t₂} :
    y ≠ x -> appears_free_in t₁ -> appears_free_in (scase t y t₁ z t₂)
| afi_scase3 {t y t₁ z t₂} :
    z ≠ x -> appears_free_in t₂ -> appears_free_in (scase t y t₁ z t₂)
| afi_cons1 {t₁ t₂} : appears_free_in t₁ -> appears_free_in (cons t₁ t₂)
| afi_cons2 {t₁ t₂} : appears_free_in t₂ -> appears_free_in (cons t₁ t₂)
| afi_lcase1 {t t₁ y z t₂} :
    appears_free_in t -> appears_free_in (lcase t t₁ y z t₂)
| afi_lcase2 {t t₁ y z t₂} :
    appears_free_in t₁ -> appears_free_in (lcase t t₁ y z t₂)
| afi_lcase3 {t t₁ y z t₂} :
    y ≠ x -> z ≠ x -> appears_free_in t₂ -> appears_free_in (lcase t t₁ y z t₂)
| afi_fix {t} : appears_free_in t -> appears_free_in (fix t)

def closed (t : tm) := ∀x, ¬appears_free_in x t

def stuck (t : tm) := normal_form step t ∧ ¬value t

def num_test : tm :=
tst (iszro (prd (scc (prd (mlt (const 2) (const 0))))))
  (const 5)
  (const 6)

example : has_type context.empty num_test nat :=
t_tst (t_iszro (t_prd (t_scc (t_prd (t_mlt t_const t_const))))) t_const t_const

example : num_test -+>* const 5 :=
multi_step (st_tst (st_iszro (st_prd (st_scc (st_prd st_mltnm))))) $
multi_step (st_tst (st_iszro (st_prd (st_scc st_prdzro)))) $
multi_step (st_tst (st_iszro (st_prd st_sccn))) $
multi_step (st_tst (st_iszro st_prdnzr)) $
multi_step (st_tst st_iszrozro) $
multi_step st_tsttru $
multi_refl

def prod_test : tm :=
snd (fst (pair (pair (const 5) (const 6)) (const 7)))

example : has_type context.empty prod_test nat :=
t_snd (t_fst (t_pair (t_pair t_const t_const) t_const))

example : prod_test -+>* const 6 :=
multi_step (st_snd (st_fstpair (v_pair v_const v_const) v_const)) $
multi_step (st_sndpair v_const v_const) $
multi_refl

def let_test : tm :=
let_ "x" (prd (const 6))
  (scc (var "x"))

example : has_type context.empty let_test nat :=
t_let (t_prd t_const) (t_scc (t_var rfl))

example : let_test -+>* const 6 :=
multi_step (st_let st_prdnzr) $
multi_step (st_letvalue v_const) $
multi_step st_sccn $
multi_refl

def sum_test_1 : tm :=
scase (inl nat (const 5))
  "x" (var "x")
  "y" (var "y")

example : has_type context.empty sum_test_1 nat :=
t_scase (t_inl t_const) (t_var rfl) (t_var rfl)

example : sum_test_1 -+>* const 5 :=
multi_step (st_scaseinl v_const) $
multi_refl

def sum_test_2 : tm :=
let_ "processSum" (abs "x" (sum nat nat)
                    (scase (var "x")
                      "n" (var "n")
                      "n" (tst (iszro (var "n")) (const 1) (const 0))))
  (pair (app (var "processSum") (inl nat (const 5)))
        (app (var "processSum") (inr nat (const 5))))

example : has_type context.empty sum_test_2 (prod nat nat) :=
t_let (t_abs (t_scase (t_var rfl)
                      (t_var rfl)
                      (t_tst (t_iszro (t_var rfl)) t_const t_const)))
  (t_pair (t_app (t_var rfl) (t_inl t_const))
          (t_app (t_var rfl) (t_inr t_const)))

example : sum_test_2 -+>* pair (const 5) (const 0) :=
multi_step (st_letvalue v_abs) $
multi_step (st_pair1 (st_appabs (v_inl v_const))) $
multi_step (st_pair1 (st_scaseinl v_const)) $
multi_step (st_pair2 v_const (st_appabs (v_inr v_const))) $
multi_step (st_pair2 v_const (st_scaseinr v_const)) $
multi_step (st_pair2 v_const (st_tst st_iszronzr)) $
multi_step (st_pair2 v_const st_tstfls) $
multi_refl

def list_test : tm :=
let_ "l" (cons (const 5) (cons (const 6) (nil nat)))
  (lcase (var "l")
    (const 0)
    "x" "y" (mlt (var "x") (var "x")))

example : has_type context.empty list_test nat :=
t_let (t_cons t_const (t_cons t_const t_nil))
      (t_lcase (t_var rfl) t_const (t_mlt (t_var rfl) (t_var rfl)))

example : list_test -+>* const 25 :=
multi_step (st_letvalue (v_cons v_const (v_cons v_const v_nil))) $
multi_step (st_lcasecons v_const (v_cons v_const v_nil)) $
multi_step st_mltnm $
multi_refl

def fix_fact : tm :=
fix (abs "f" (arrow nat nat)
      (abs "a" nat
        (tst (iszro (var "a"))
          (const 1)
          (mlt (var "a") (app (var "f") (prd (var "a")))))))

example : has_type context.empty fix_fact (arrow nat nat) :=
t_fix (t_abs (t_abs (t_tst (t_iszro (t_var rfl))
                           t_const
                           (t_mlt (t_var rfl)
                                  (t_app (t_var rfl) (t_prd (t_var rfl)))))))

example : app fix_fact (const 4) -+>* const 24 :=
multi_step (st_app1 st_fixabs) $
multi_step (st_appabs v_const) $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_mlt2 v_const (st_app1 st_fixabs)) $
multi_step (st_mlt2 v_const (st_app2 v_abs st_prdnzr)) $
multi_step (st_mlt2 v_const (st_appabs v_const)) $
multi_step (st_mlt2 v_const (st_tst st_iszronzr)) $
multi_step (st_mlt2 v_const st_tstfls) $
multi_step (st_mlt2 v_const (st_mlt2 v_const (st_app1 st_fixabs))) $
multi_step (st_mlt2 v_const (st_mlt2 v_const (st_app2 v_abs st_prdnzr))) $
multi_step (st_mlt2 v_const (st_mlt2 v_const (st_appabs v_const))) $
multi_step (st_mlt2 v_const (st_mlt2 v_const (st_tst st_iszronzr))) $
multi_step (st_mlt2 v_const (st_mlt2 v_const st_tstfls)) $
multi_step (st_mlt2
             v_const
             (st_mlt2 v_const (st_mlt2 v_const (st_app1 st_fixabs)))) $
multi_step (st_mlt2
             v_const
             (st_mlt2 v_const (st_mlt2 v_const (st_app2 v_abs st_prdnzr)))) $
multi_step (st_mlt2
             v_const
             (st_mlt2 v_const (st_mlt2 v_const (st_appabs v_const)))) $
multi_step (st_mlt2
             v_const
             (st_mlt2 v_const (st_mlt2 v_const (st_tst st_iszronzr)))) $
multi_step (st_mlt2
             v_const
             (st_mlt2 v_const (st_mlt2 v_const st_tstfls))) $
multi_step (st_mlt2
             v_const
             (st_mlt2
               v_const
               (st_mlt2 v_const (st_mlt2 v_const (st_app1 st_fixabs))))) $
multi_step (st_mlt2
             v_const
             (st_mlt2
               v_const
               (st_mlt2 v_const (st_mlt2 v_const (st_app2 v_abs st_prdnzr))))) $
multi_step (st_mlt2
             v_const
             (st_mlt2
               v_const
               (st_mlt2 v_const (st_mlt2 v_const (st_appabs v_const))))) $
multi_step (st_mlt2
             v_const
             (st_mlt2
               v_const
               (st_mlt2 v_const (st_mlt2 v_const (st_tst st_iszrozro))))) $
multi_step (st_mlt2
             v_const
             (st_mlt2 v_const (st_mlt2 v_const (st_mlt2 v_const st_tsttru)))) $
multi_step (st_mlt2 v_const (st_mlt2 v_const (st_mlt2 v_const st_mltnm))) $
multi_step (st_mlt2 v_const (st_mlt2 v_const st_mltnm)) $
multi_step (st_mlt2 v_const st_mltnm) $
multi_step st_mltnm $
multi_refl

def fix_map : tm :=
abs "g" (arrow nat nat)
  (fix (abs "f" (arrow (list nat) (list nat))
         (abs "l" (list nat)
           (lcase (var "l")
             (nil nat)
             "a" "l" (cons (app (var "g") (var "a"))
                           (app (var "f") (var "l")))))))

example : has_type context.empty
                   fix_map
                   (arrow (arrow nat nat) (arrow (list nat) (list nat))) :=
t_abs (t_fix (t_abs (t_abs (t_lcase (t_var rfl)
                                    t_nil
                                    (t_cons (t_app (t_var rfl) (t_var rfl))
                                            (t_app (t_var rfl) (t_var rfl)))))))

example :    app (app fix_map (abs "a" nat (scc (var "a"))))
                 (cons (const 1) (cons (const 2) (nil nat)))
        -+>* cons (const 2) (cons (const 3) (nil nat)) :=
multi_step (st_app1 (st_appabs v_abs)) $
multi_step (st_app1 st_fixabs) $
multi_step (st_appabs (v_cons v_const (v_cons v_const v_nil))) $
multi_step (st_lcasecons v_const (v_cons v_const v_nil)) $
multi_step (st_cons1 (st_appabs v_const)) $
multi_step (st_cons1 st_sccn) $
multi_step (st_cons2 v_const (st_app1 st_fixabs)) $
multi_step (st_cons2 v_const (st_appabs (v_cons v_const v_nil))) $
multi_step (st_cons2 v_const (st_lcasecons v_const v_nil)) $
multi_step (st_cons2 v_const (st_cons1 (st_appabs v_const))) $
multi_step (st_cons2 v_const (st_cons1 st_sccn)) $
multi_step (st_cons2 v_const (st_cons2 v_const (st_app1 st_fixabs))) $
multi_step (st_cons2 v_const (st_cons2 v_const (st_appabs v_nil))) $
multi_step (st_cons2 v_const (st_cons2 v_const st_lcasenil)) $
multi_refl

def fix_equal : tm :=
fix (abs "eq" (arrow nat (arrow nat nat))
      (abs "m" nat
        (abs "n" nat
          (tst (iszro (var "m"))
            (tst (iszro (var "n")) (const 1) (const 0))
            (tst (iszro (var "n"))
              (const 0)
              (app (app (var "eq") (prd (var "m")))
                   (prd (var "n"))))))))

example : has_type context.empty fix_equal (arrow nat (arrow nat nat)) :=
t_fix (t_abs (t_abs (t_abs (t_tst (t_iszro (t_var rfl))
                             (t_tst (t_iszro (t_var rfl)) t_const t_const)
                             (t_tst (t_iszro (t_var rfl))
                               t_const
                               (t_app (t_app (t_var rfl) (t_prd (t_var rfl)))
                                      (t_prd (t_var rfl))))))))

example : app (app fix_equal (const 4)) (const 4) -+>* const 1 :=
multi_step (st_app1 (st_app1 st_fixabs)) $
multi_step (st_app1 (st_appabs v_const)) $
multi_step (st_appabs v_const) $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_app1 (st_app1 st_fixabs)) $
multi_step (st_app1 (st_app2 v_abs st_prdnzr)) $
multi_step (st_app1 (st_appabs v_const)) $
multi_step (st_app2 v_abs st_prdnzr) $
multi_step (st_appabs v_const) $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_app1 (st_app1 st_fixabs)) $
multi_step (st_app1 (st_app2 v_abs st_prdnzr)) $
multi_step (st_app1 (st_appabs v_const)) $
multi_step (st_app2 v_abs st_prdnzr) $
multi_step (st_appabs v_const) $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_app1 (st_app1 st_fixabs)) $
multi_step (st_app1 (st_app2 v_abs st_prdnzr)) $
multi_step (st_app1 (st_appabs v_const)) $
multi_step (st_app2 v_abs st_prdnzr) $
multi_step (st_appabs v_const) $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_app1 (st_app1 st_fixabs)) $
multi_step (st_app1 (st_app2 v_abs st_prdnzr)) $
multi_step (st_app1 (st_appabs v_const)) $
multi_step (st_app2 v_abs st_prdnzr) $
multi_step (st_appabs v_const) $
multi_step (st_tst st_iszrozro) $
multi_step st_tsttru $
multi_step (st_tst st_iszrozro) $
multi_step st_tsttru $
multi_refl

example : app (app fix_equal (const 4)) (const 5) -+>* const 0 :=
multi_step (st_app1 (st_app1 st_fixabs)) $
multi_step (st_app1 (st_appabs v_const)) $
multi_step (st_appabs v_const) $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_app1 (st_app1 st_fixabs)) $
multi_step (st_app1 (st_app2 v_abs st_prdnzr)) $
multi_step (st_app1 (st_appabs v_const)) $
multi_step (st_app2 v_abs st_prdnzr) $
multi_step (st_appabs v_const) $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_app1 (st_app1 st_fixabs)) $
multi_step (st_app1 (st_app2 v_abs st_prdnzr)) $
multi_step (st_app1 (st_appabs v_const)) $
multi_step (st_app2 v_abs st_prdnzr) $
multi_step (st_appabs v_const) $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_app1 (st_app1 st_fixabs)) $
multi_step (st_app1 (st_app2 v_abs st_prdnzr)) $
multi_step (st_app1 (st_appabs v_const)) $
multi_step (st_app2 v_abs st_prdnzr) $
multi_step (st_appabs v_const) $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_step (st_app1 (st_app1 st_fixabs)) $
multi_step (st_app1 (st_app2 v_abs st_prdnzr)) $
multi_step (st_app1 (st_appabs v_const)) $
multi_step (st_app2 v_abs st_prdnzr) $
multi_step (st_appabs v_const) $
multi_step (st_tst st_iszrozro) $
multi_step st_tsttru $
multi_step (st_tst st_iszronzr) $
multi_step st_tstfls $
multi_refl

def evenodd : tm :=
let_ "evenodd" (fix (abs "eo" (prod (arrow nat nat) (arrow nat nat))
                      (pair (abs "n" nat
                              (tst (iszro (var "n"))
                                (const 1)
                                (app (snd (var "eo")) (prd (var "n")))))
                            (abs "n" nat
                              (tst (iszro (var "n"))
                                (const 0)
                                (app (fst (var "eo")) (prd (var "n"))))))))
  (let_ "even" (fst (var "evenodd"))
    (let_ "odd" (snd (var "evenodd"))
      (pair (app (var "even") (const 3))
            (app (var "even") (const 4)))))

example : has_type context.empty evenodd (prod nat nat) :=
t_let (t_fix
        (t_abs
          (t_pair (t_abs (t_tst (t_iszro (t_var rfl))
                           t_const
                           (t_app (t_snd (t_var rfl)) (t_prd (t_var rfl)))))
                  (t_abs (t_tst (t_iszro (t_var rfl))
                           t_const
                           (t_app (t_fst (t_var rfl)) (t_prd (t_var rfl))))))))
  (t_let (t_fst (t_var rfl))
    (t_let (t_snd (t_var rfl))
      (t_pair (t_app (t_var rfl) t_const)
              (t_app (t_var rfl) t_const))))

example : evenodd -+>* pair (const 0) (const 1) :=
multi_step (st_let st_fixabs) $
multi_step (st_letvalue (v_pair v_abs v_abs)) $
multi_step (st_let (st_fstpair v_abs v_abs)) $
multi_step (st_letvalue v_abs) $
multi_step (st_let (st_sndpair v_abs v_abs)) $
multi_step (st_letvalue v_abs) $
multi_step (st_pair1 (st_appabs v_const)) $
multi_step (st_pair1 (st_tst st_iszronzr)) $
multi_step (st_pair1 st_tstfls) $
multi_step (st_pair1 (st_app1 (st_snd st_fixabs))) $
multi_step (st_pair1 (st_app1 (st_sndpair v_abs v_abs))) $
multi_step (st_pair1 (st_app2 v_abs st_prdnzr)) $
multi_step (st_pair1 (st_appabs v_const)) $
multi_step (st_pair1 (st_tst st_iszronzr)) $
multi_step (st_pair1 st_tstfls) $
multi_step (st_pair1 (st_app1 (st_fst st_fixabs))) $
multi_step (st_pair1 (st_app1 (st_fstpair v_abs v_abs))) $
multi_step (st_pair1 (st_app2 v_abs st_prdnzr)) $
multi_step (st_pair1 (st_appabs v_const)) $
multi_step (st_pair1 (st_tst st_iszronzr)) $
multi_step (st_pair1 st_tstfls) $
multi_step (st_pair1 (st_app1 (st_snd st_fixabs))) $
multi_step (st_pair1 (st_app1 (st_sndpair v_abs v_abs))) $
multi_step (st_pair1 (st_app2 v_abs st_prdnzr)) $
multi_step (st_pair1 (st_appabs v_const)) $
multi_step (st_pair1 (st_tst st_iszrozro)) $
multi_step (st_pair1 st_tsttru) $
multi_step (st_pair2 v_const (st_appabs v_const)) $
multi_step (st_pair2 v_const (st_tst st_iszronzr)) $
multi_step (st_pair2 v_const st_tstfls) $
multi_step (st_pair2 v_const (st_app1 (st_snd st_fixabs))) $
multi_step (st_pair2 v_const (st_app1 (st_sndpair v_abs v_abs))) $
multi_step (st_pair2 v_const (st_app2 v_abs st_prdnzr)) $
multi_step (st_pair2 v_const (st_appabs v_const)) $
multi_step (st_pair2 v_const (st_tst st_iszronzr)) $
multi_step (st_pair2 v_const st_tstfls) $
multi_step (st_pair2 v_const (st_app1 (st_fst st_fixabs))) $
multi_step (st_pair2 v_const (st_app1 (st_fstpair v_abs v_abs))) $
multi_step (st_pair2 v_const (st_app2 v_abs st_prdnzr)) $
multi_step (st_pair2 v_const (st_appabs v_const)) $
multi_step (st_pair2 v_const (st_tst st_iszronzr)) $
multi_step (st_pair2 v_const st_tstfls) $
multi_step (st_pair2 v_const (st_app1 (st_snd st_fixabs))) $
multi_step (st_pair2 v_const (st_app1 (st_sndpair v_abs v_abs))) $
multi_step (st_pair2 v_const (st_app2 v_abs st_prdnzr)) $
multi_step (st_pair2 v_const (st_appabs v_const)) $
multi_step (st_pair2 v_const (st_tst st_iszronzr)) $
multi_step (st_pair2 v_const st_tstfls) $
multi_step (st_pair2 v_const (st_app1 (st_fst st_fixabs))) $
multi_step (st_pair2 v_const (st_app1 (st_fstpair v_abs v_abs))) $
multi_step (st_pair2 v_const (st_app2 v_abs st_prdnzr)) $
multi_step (st_pair2 v_const (st_appabs v_const)) $
multi_step (st_pair2 v_const (st_tst st_iszrozro)) $
multi_step (st_pair2 v_const st_tsttru) $
multi_refl


-- Vaguely following an approach taken in this paper:
-- https://arxiv.org/pdf/1401.7886.pdf
-- author: Ben Sherman
import data.list.basic
       galois.list.preds
       galois.tactic
       galois.nat.lemmas
       galois.list

universes u v

def depair {A} : list (A × A) -> list A
| [] := []
| ((x, y) :: zs) := y :: x :: depair zs

namespace binary_tree

/-- An `lptree A` represents a binary tree (without internal nodes)
    whose left subtrees are all perfect, where the sizes of these
    perfect left subtrees increase as the roots of those subtrees
    move towards the root of the entire tree.
    (the "lp" in "lptree" is for left-perfect)

    Effectively, an `lptree unit` is a binary natural number,
    and for any `lptree A`, this corresponding binary natural number
    indicates the number of items of type A in the tree.
-/
inductive lptree : Type u -> Type (u + 1)
| nil : ∀ {A}, lptree A
| cons : ∀ {A}, option A -> lptree (A × A) -> lptree A

def double {A : Type u} {B : Type v} (f : A → B)
  (x : A × A) : B × B :=
  let (a, b) := x in (f a, f b).

namespace lptree

def height : ∀ {A : Type u}, lptree A → ℕ
| _ nil := 0
| _ (cons mx t) := height t + 1

def size : ∀ {A : Type u}, lptree A → ℕ
| _ nil := 0
| _ (cons mx t) := 2 * size t + (match mx with
  | some _ := 1
  | none := 0
  end)

/-- Indicates when an `lptree A` has at least one value of `A` in it
    (i.e., its corresponding binary natural number is nonzero.)
-/
inductive nonzero : ∀  {A : Type u}, lptree A -> Prop
| cons_some : ∀ {A} (a : A) t, nonzero (lptree.cons (some a) t)
| extend : ∀ {A} ma (t : lptree (A × A)), @nonzero (A × A) t ->
    nonzero (lptree.cons ma t)

/-- Add a single element to a left-perfect tree.
    (A generalization of adding one to a binary natural number.)
-/
def insert : ∀ {A : Type u}, A -> lptree A -> lptree A
| A x lptree.nil := lptree.cons (some x) lptree.nil
| A x (lptree.cons none t') := lptree.cons (some x) t'
| A x (lptree.cons (some y) t') := lptree.cons none (insert (y, x) t')

def map : ∀ {A : Type u} {B : Type v}, (A → B) → lptree A → lptree B
| A B f lptree.nil := lptree.nil
| A B f (lptree.cons x t') := lptree.cons (option.map f x) (map (double f) t')

/-- Enumerate the leaves of an lptree from right to left
    (i.e., going up towards larger left-perfect subtrees
-/
def leaves : ∀ {A}, lptree A -> list A
| A nil := []
| A (cons mx t) := (match mx with
  | none := λ zs, zs
  | (some x) := list.cons x
  end) (depair (leaves t))

/-- Decision procedure to determine whether
    an lptree is nonempty
-/
def nonzero_bool
  : ∀ {A : Type}, lptree A -> bool
| _ nil := false
| _ (cons (some x) t) := true
| _ (cons none t) := nonzero_bool t

def nonzero_bool_equiv {A : Type}
  (t : lptree A)
  : nonzero_bool t = true ↔ nonzero t
:=
begin
split; intro H,
{ induction t,
  { contradiction },
  { cases a, simp [nonzero_bool] at H,
    { constructor,
      apply ih_1, assumption },
    { constructor }
  }
},
{ induction H,
  { simp [nonzero_bool] },
  { cases ma,
    { simp [nonzero_bool], assumption },
    { reflexivity }
  }
}
end

def nonzero_dec {A : Type}
  (t : lptree A)
  : decidable (nonzero t)
:=
begin
rw <- nonzero_bool_equiv,
apply bool.decidable_eq
end

end lptree

lemma map_insert {A B} (f : A → B) (t : lptree A) (x : A)
  : (t.insert x).map f = (t.map f).insert (f x)
:= begin
revert B,
induction t; intros, dsimp [lptree.insert, lptree.map], reflexivity,
induction a; simp [lptree.insert, lptree.map, option.map, option.bind],
rw ih_1, reflexivity,
end

lemma insert_size {A : Type u} (x : A) (t : lptree A)
  : (t.insert x).size = t.size + 1
:= begin
induction t; dsimp [lptree.size, lptree.insert],
{ simp },
{ induction a; dsimp [lptree.size, lptree.insert],
  { reflexivity },
  { rw ih_1,
    rw mul_add, repeat { rw add_assoc },
    f_equal, }
}
end

lemma map_height {A B} (f : A → B) (t : lptree A)
  : (t.map f).height = t.height
:= begin
revert B,
induction t; intros; simp [lptree.map, lptree.height],
f_equal, apply ih_1,
end

lemma map_size {A B} (f : A → B) (t : lptree A)
  : (t.map f).size = t.size
:= begin
revert B,
induction t; intros; simp [lptree.map, lptree.size],
induction a; simp [lptree.map, option.map, option.bind, lptree.size],
rw mul_comm, f_equal, apply ih_1,
f_equal, rw mul_comm, f_equal, apply ih_1,
end

lemma size_le_height {A : Type u} (t : lptree A)
  : t.size + 1 ≤ 2 ^ t.height
:= begin
induction t; dsimp [lptree.size, lptree.height, nat.pow],
{ rw add_comm, },
{ induction a; dsimp [lptree.size],
  { rw add_assoc, rw (add_comm 0 1), rw add_zero, rw mul_comm,
    rw nat.mul_2_add, rw nat.mul_2_add, rw add_assoc,
    apply nat.le_add_compat,
    apply le_trans, tactic.swap, assumption,
    constructor, constructor, assumption,
   },
  { rw add_assoc, rw ← nat.mul_2_add, rw ← mul_comm,
    rw ← add_mul, apply mul_le_mul, assumption,
    apply le_refl, apply nat.zero_le, apply nat.zero_le,
  }
}
end

inductive nz_minimal : ∀ {A}, lptree A → Prop
| last : ∀ {A} x : A, nz_minimal (lptree.cons (some x) lptree.nil)
| cons : ∀ {A} (mx : option A) t, nz_minimal t → nz_minimal (lptree.cons mx t)

inductive minimal {A} : lptree A → Prop
| zero : minimal (lptree.nil)
| nz_minimal : ∀ t, nz_minimal t → minimal t

lemma insert_nz_minimal {A : Type u} (t : lptree A)
  (mint : minimal t) (x : A)
  : nz_minimal (t.insert x)
:= begin
induction mint,
{ constructor },
{ induction a,
  { dsimp [lptree.insert], constructor, constructor },
  { induction mx; dsimp [lptree.insert],
    { constructor, assumption },
    { constructor, apply ih_1, assumption }
  }
}
end

/-- If you add one to an lptree, it is nonzero
-/
def insert_nonzero {A : Type u} (t : lptree A)
  : ∀ x : A, lptree.nonzero (t.insert x) :=
begin
induction t with A A ma t IHt; intros x,
{ simp [lptree.insert], constructor },
{ cases ma,
  { simp [lptree.insert], constructor },
  { simp [lptree.insert], constructor, apply IHt }
}
end

def list_to_lptree {A : Type u} : list A -> lptree A :=
  list.foldr lptree.insert lptree.nil

lemma map_list_to_lptree {A : Type u} {B : Type v}
  (f : A → B) (xs : list A)
  : (list_to_lptree xs).map f = list_to_lptree (xs.map f)
:= begin
revert B,
induction xs; intros, reflexivity,
dsimp [list_to_lptree, lptree.map, list.map],
rw map_insert, f_equal, unfold list_to_lptree at ih_1,
apply ih_1,
end

lemma list_to_lptree_length {A : Type u} (xs : list A)
  : (list_to_lptree xs).size = xs.length
:= begin
unfold list_to_lptree,
induction xs; dsimp [list.foldr, lptree.size],
{ reflexivity },
{ rw insert_size, rw ih_1, }
end

/-- This only is correct if the lptree is minimal. -/
def log2_of_succ_lptree {A} (t : lptree A)
  := t.height

def log2_of_succ (n : ℕ) :=
  log2_of_succ_lptree (list_to_lptree (list.repeat unit.star n))

def log2 (n : ℕ) : ℕ := log2_of_succ n.pred

lemma list_to_lptree_nz_minimal {A} (x : A) (xs : list A)
  : nz_minimal (list_to_lptree (x :: xs))
:= begin
revert x,
induction xs; intros,
constructor,
dsimp [list_to_lptree],
apply insert_nz_minimal, apply minimal.nz_minimal,
apply ih_1,
end

lemma list_to_lptree_minimal {A} (xs : list A)
  : minimal (list_to_lptree xs)
:= begin
cases xs, apply minimal.zero,
apply minimal.nz_minimal, apply list_to_lptree_nz_minimal
end

lemma same_length_map {A : Type u} {B : Type v}
  (xs : list A) (ys : list B) (Hlen : xs.length = ys.length)
  : xs.map (λ _, unit.star) = ys.map (λ _, unit.star)
:= begin
apply list.pair_induction_same_length _ _ _ xs ys Hlen,
simp, reflexivity, intros,
simp [list.map], simp [list.map], f_equal,
assumption
end

lemma lptree_height_skeleton {A} (t : lptree A)
  : t.height = (t.map (λ _, unit.star)).height
:= begin
symmetry, apply map_height,
end

lemma lptree_height_length {A B} (xs : list A) (ys : list B)
  (Hlen : xs.length = ys.length)
  : (list_to_lptree xs).height = (list_to_lptree ys).height
:= begin
rw (lptree_height_skeleton (list_to_lptree xs)),
rw (lptree_height_skeleton (list_to_lptree ys)),
repeat { rw map_list_to_lptree },
rw same_length_map, assumption,
end

lemma list_to_lptree_height {A} (xs : list A)
  : (list_to_lptree xs).height = log2_of_succ xs.length
:= begin
rw ← list_to_lptree_length,
unfold log2_of_succ log2_of_succ_lptree,
rw list_to_lptree_length,
apply lptree_height_length,
rw list.repeat_length,
end

lemma list_to_lptree_height' {A} (xs : list A)
  : (list_to_lptree xs).height = log2 xs.length.succ
:= begin
unfold log2, dsimp [nat.pred],
apply list_to_lptree_height,
end

lemma lptree_leaves_cons {A} (t : lptree A)
  (x : A)
  : (t.insert x).leaves =
     x :: t.leaves
  :=
begin
induction t,
{ reflexivity },
{ cases a; simp only [lptree.insert, lptree.leaves],
  { rw ih_1,
    simp only [depair]
   }
}
end

/-- If we produce an lptree from a list of elements,
    then enumerate the leaves, we get the same list back
-/
lemma lptree_leaves_list {A} (xs : list A)
  : (list_to_lptree xs).leaves = xs
  :=
begin
induction xs,
{ reflexivity },
{ simp [list_to_lptree], rw lptree_leaves_cons,
  f_equal, assumption
}
end

lemma depair_nil {A} (xs : list (A × A))
  : xs = [] -> depair xs = []
:= begin
intros H, rw H, reflexivity
end

lemma lptree_not_nonzero_no_leaves {A : Type}
  (t : lptree A)
  (H : ¬ lptree.nonzero t)
  : t.leaves = []
:= begin
rw <- lptree.nonzero_bool_equiv at H,
induction t,
{ reflexivity },
{ cases a,
  { simp [lptree.leaves],
    apply depair_nil, apply ih_1,
    simp [lptree.nonzero_bool] at H,
    assumption
     },
  { simp [lptree.nonzero_bool] at H,
    contradiction }
}
end

lemma list_to_lptree.nonzero {A} (xs : list A)
  (H : xs ≠ []) : lptree.nonzero (list_to_lptree xs) :=
begin
cases xs,
{ contradiction },
{ simp [list_to_lptree], apply insert_nonzero }
end

end binary_tree


import IIT.PropInversion
import IIT.ClarifyIndices

mutual
inductive Conₑ : Type
| nilₑ : Conₑ
| extₑ : Conₑ → Tyₑ → Conₑ

inductive Tyₑ : Type
| baseₑ : Conₑ → Tyₑ
| piₑ   : Conₑ → Tyₑ → Tyₑ → Tyₑ
end

open Conₑ Tyₑ

mutual
inductive Con_w : Conₑ → Prop
| nil_w : Con_w nilₑ
| ext_w : ∀ {Γ}, Con_w Γ → ∀ {A}, Ty_w Γ A → Con_w (extₑ Γ A)

inductive Ty_w : Conₑ → Tyₑ → Prop
| base_w : ∀ {Γ}, Con_w Γ → Ty_w Γ (baseₑ Γ)
| pi_w : ∀ {Γ}, Con_w Γ → ∀ {A}, Ty_w Γ A → ∀ {B}, Ty_w (extₑ Γ A) B → Ty_w Γ (piₑ Γ A B)
end

open Con_w Ty_w

def Con := PSigma Con_w
def Ty := fun (Γ : Con) => PSigma (Ty_w Γ.1)

def nil : Con                                         := ⟨nilₑ,            nil_w⟩
def ext (Γ : Con) (A : Ty Γ) : Con                    := ⟨extₑ Γ.1 A.1,    ext_w Γ.2 A.2⟩
def base (Γ : Con) : Ty Γ                             := ⟨baseₑ Γ.1,       base_w Γ.2⟩
def pi (Γ : Con) (A : Ty Γ) (B : Ty (ext Γ A)) : Ty Γ := ⟨piₑ Γ.1 A.1 B.1, pi_w Γ.2 A.2 B.2⟩

section
variable
  (Conₘ  : Con → Sort _)
  (Tyₘ   : ∀ {Γ}, Conₘ Γ → Ty Γ → Sort)
  (nilₘ  : Conₘ nil)
  (extₘ  : ∀ {Γ} (Γₘ : Conₘ Γ) {A}, Tyₘ Γₘ A → Conₘ (ext Γ A))
  (baseₘ : ∀ {Γ} (Γₘ : Conₘ Γ), Tyₘ Γₘ (base Γ))
  (piₘ   : ∀ {Γ} (Γₘ : Conₘ Γ) {A} (Aₘ : Tyₘ Γₘ A) {B} (Bₘ : Tyₘ (extₘ Γₘ Aₘ) B), Tyₘ Γₘ (pi Γ A B))

mutual
inductive Conᵣ : (Γ : Con) → Conₘ Γ → Prop
| nilᵣ : Conᵣ nil nilₘ
| extᵣ : ∀ {Γ} {Γₘ : Conₘ Γ}, Conᵣ Γ Γₘ →
           ∀ {A} {Aₘ : Tyₘ Γₘ A}, Tyᵣ Γₘ A Aₘ → Conᵣ (ext Γ A) (extₘ Γₘ Aₘ)

inductive Tyᵣ : {Γ : Con} → (Γₘ : Conₘ Γ) → (A : Ty Γ) → Tyₘ Γₘ A → Prop
| baseᵣ: ∀ {Γ} {Γₘ : Conₘ Γ}, Conᵣ Γ Γₘ → Tyᵣ Γₘ (base Γ) (baseₘ Γₘ)
| piᵣ: ∀ {Γ} {Γₘ : Conₘ Γ}, Conᵣ Γ Γₘ →
         ∀ {A} {Aₘ : Tyₘ Γₘ A}, Tyᵣ Γₘ A Aₘ →
           ∀ {B} {Bₘ : Tyₘ (extₘ Γₘ Aₘ) B}, Tyᵣ (extₘ Γₘ Aₘ) B Bₘ →
             Tyᵣ Γₘ (pi Γ A B) (piₘ Γₘ Aₘ Bₘ)
end

open Conᵣ Tyᵣ

noncomputable def Con_tot (Γ : Con) : PSigma (Conᵣ Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ) := by
  cases Γ with | mk Γₑ Γ_w => ?_
  apply Conₑ.recOn Γₑ
    (motive_1 := fun Γₑ => ∀ Γ_w, PSigma (Conᵣ Conₘ Tyₘ nilₘ extₘ baseₘ piₘ ⟨Γₑ, Γ_w⟩))
    (motive_2 := fun Aₑ => ∀ {Γ Γₘ} (Γᵣ : Conᵣ Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ Γₘ)
                   A_w, PSigma (Tyᵣ Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γₘ ⟨Aₑ, A_w⟩))
  · intro Γ_w
    exact ⟨nilₘ, nilᵣ⟩
  · intro Δₑ Aₑ Δ_ih A_ih ctor_w
    inversion ctor_w with Δ_w A_w
    cases Δ_ih Δ_w with | mk Δₘ Δᵣ => ?_
    cases A_ih Δᵣ A_w with | mk Aₘ Aᵣ => ?_
    exact ⟨extₘ Δₘ Aₘ, extᵣ Δᵣ Aᵣ⟩
  · intro Γₑ Γ_ih ⟨Δₑ, Δ_w⟩ Δₘ Δᵣ ctor_w
    simp only at ctor_w
    clarifyIndices ctor_w
    exact ⟨baseₘ Δₘ, baseᵣ Δᵣ⟩
  · intro Δₑ Aₑ Bₑ Δ_ih A_ih B_ih ⟨Δ'ₑ, Δ_w⟩ Δ'ₘ Δ'ᵣ ctor_w
    simp only at ctor_w
    clarifyIndices ctor_w
    inversion ctor_w with Δ_w A_w B_w
    cases A_ih Δ'ᵣ A_w with | mk Aₘ Aᵣ => ?_
    cases B_ih (extᵣ Δ'ᵣ Aᵣ) B_w with | mk Bₘ Bᵣ => ?_
    exact ⟨piₘ Δ'ₘ Aₘ Bₘ, piᵣ Δ'ᵣ Aᵣ Bᵣ⟩

noncomputable def Ty_tot (Γ : Con) (A : Ty Γ) :
  PSigma (Tyᵣ Conₘ Tyₘ nilₘ extₘ baseₘ piₘ (Con_tot Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ).1 A) := by
  cases Γ with | mk Γₑ Γ_w => ?_
  cases A with | mk Aₑ A_w => ?_
  apply Tyₑ.recOn Aₑ
    (motive_1 := fun Γₑ => ∀ Γ_w, PSigma (Conᵣ Conₘ Tyₘ nilₘ extₘ baseₘ piₘ ⟨Γₑ, Γ_w⟩))
    (motive_2 := fun Aₑ => ∀ {Γ Γₘ} (Γᵣ : Conᵣ Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ Γₘ)
                   A_w, PSigma (Tyᵣ Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γₘ ⟨Aₑ, A_w⟩))
  · intro Γ_w
    exact ⟨nilₘ, nilᵣ⟩
  · intro Δₑ Aₑ Δ_ih A_ih ctor_w
    inversion ctor_w with Δ_w A_w
    cases Δ_ih Δ_w with | mk Δₘ Δᵣ => ?_
    cases A_ih Δᵣ A_w with | mk Aₘ Aᵣ => ?_
    exact ⟨extₘ Δₘ Aₘ, extᵣ Δᵣ Aᵣ⟩
  · intro Γₑ Γ_ih ⟨Δₑ, Δ_w⟩ Δₘ Δᵣ ctor_w
    simp only at ctor_w
    clarifyIndices ctor_w
    exact ⟨baseₘ Δₘ, baseᵣ Δᵣ⟩
  · intro Δₑ Aₑ Bₑ Δ_ih A_ih B_ih ⟨Δ'ₑ, Δ_w⟩ Δ'ₘ Δ'ᵣ ctor_w
    simp only at ctor_w
    clarifyIndices ctor_w
    inversion ctor_w with Δ_w A_w B_w
    cases A_ih Δ'ᵣ A_w with | mk Aₘ Aᵣ => ?_
    cases B_ih (extᵣ Δ'ᵣ Aᵣ) B_w with | mk Bₘ Bᵣ => ?_
    exact ⟨piₘ Δ'ₘ Aₘ Bₘ, piᵣ Δ'ᵣ Aᵣ Bᵣ⟩
  · exact (Con_tot Conₘ Tyₘ nilₘ extₘ baseₘ piₘ ⟨Γₑ, Γ_w⟩).2

noncomputable def Con.rec (Γ : Con) : Conₘ Γ :=
(Con_tot Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ).1

noncomputable def Ty.rec (Γ : Con) (A : Ty Γ) : Tyₘ (Con.rec Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ) A :=
(Ty_tot Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ A).1

theorem nil_beta : Con.rec Conₘ Tyₘ nilₘ extₘ baseₘ piₘ nil = nilₘ :=
rfl

theorem ext_beta (Γ : Con) (A : Ty Γ) :
  Con.rec Conₘ Tyₘ nilₘ extₘ baseₘ piₘ (ext Γ A)
  = extₘ (Con.rec Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ)
    (Ty.rec Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ A) :=
rfl

theorem base_beta (Γ : Con) :
  Ty.rec Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ (base Γ)
  = baseₘ (Con.rec Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ) :=
rfl

theorem pi_beta (Γ : Con) (A : Ty Γ) (B : Ty (ext Γ A)) :
  Ty.rec Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ (pi Γ A B)
  = piₘ (Con.rec Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ)
      (Ty.rec Conₘ Tyₘ nilₘ extₘ baseₘ piₘ Γ A)
      (Ty.rec Conₘ Tyₘ nilₘ extₘ baseₘ piₘ (ext Γ A) B) :=
rfl

end


import data.equiv.basic
import data.set.function
import for_mathlib

import data.is_equiv

open set equiv

section Bij_on

variables {α : Type*} {β : Type*} {γ : Type*} {δ : Type*}

-- A constructive version of bij_on: f : a → b is an equivalence. Note
-- that f is defined on all of α, but its inverse need only be defined
-- on b ⊆ β.
structure Bij_on (f : α → β) (a : set α) (b : set β) :=
(e : a ≃ b)
(he : ∀ (x : a), f x = e x)

variables {f : α → β} {a : set α} {b : set β}

-- Alternate constructor for the case when a = univ.
def Bij_on.mk_univ (e : α ≃ b) (he : ∀ x, f x = e x) : Bij_on f univ b :=
{ e := (equiv.set.univ _).trans e,
  he := assume x, by rw [he x]; refl }

-- Construct a bijection from an equivalence.
def Bij_on.of_equiv (e : α ≃ β) : Bij_on e.to_fun univ univ :=
Bij_on.mk_univ (e.trans (equiv.set.univ _).symm) (by intro x; refl)

def Bij_on.of_Is_equiv (h : Is_equiv f) : Bij_on f univ univ :=
by convert Bij_on.of_equiv h.e; funext a; rw [h.h] { occs := occurrences.pos [1] }; refl

def Bij_on.Is_equiv (h : Bij_on f univ univ) : Is_equiv f :=
{ e := (equiv.set.univ _).symm.trans (h.e.trans (equiv.set.univ _)),
  h := funext $ λ a, h.he ⟨a, trivial⟩ }

lemma Bij_on.he' (h : Bij_on f a b) {x : α} (hx : x ∈ a) : f x = h.e ⟨x, hx⟩ :=
h.he ⟨x, hx⟩

instance (f : α → β) (a : set α) (b : set β) :
  subsingleton (Bij_on f a b) :=
⟨assume ⟨e, he⟩ ⟨e', he'⟩,
 have e = e', from coe_fn_injective
   (by funext x; apply subtype.eq; exact (he x).symm.trans (he' x)),
 by cc⟩

lemma Bij_on.maps_to (h : Bij_on f a b) : maps_to f a b :=
assume x hx, show f x ∈ b, by rw h.he' hx; exact (h.e _).property

lemma Bij_on.inj_on (h : Bij_on f a b) : inj_on f a :=
assume x hx x' hx' hh, begin
  rw [h.he' hx, h.he' hx'] at hh,
  have := subtype.eq hh,
  simpa using this
end

lemma Bij_on.injective (h : Bij_on f univ b) : function.injective f :=
injective_iff_inj_on_univ.mpr h.inj_on

lemma Bij_on.right_inv (h : Bij_on f a b) (y) : f (h.e.symm y) = y :=
by rw h.he; simp

@[refl] def Bij_on.refl (a : set α) : Bij_on id a a :=
{ e := equiv.refl _, he := assume x, rfl }

@[trans] def Bij_on.trans {f : α → β} {g : β → γ}
  {a : set α} {b : set β} {c : set γ} (hf : Bij_on f a b) (hg : Bij_on g b c) :
  Bij_on (g ∘ f) a c :=
{ e := hf.e.trans hg.e,
  he := assume x, show g (f x) = hg.e (hf.e x), by rw [hf.he, hg.he] }

def Bij_on.trans_symm {f : α → β} {g : β → γ}
  {a : set α} {b : set β} {c : set γ} (hf : ∀ x, x ∈ a → f x ∈ b)
  (hgf : Bij_on (g ∘ f) a c) (hg : Bij_on g b c) :
  Bij_on f a b :=
{ e := hgf.e.trans hg.e.symm,
  he := assume x, show f x = hg.e.symm (hgf.e x), begin
    apply hg.inj_on,
    exact hf x.val x.property,
    exact (hg.e.symm (hgf.e x)).property,
    change (g ∘ f) x = _,
    rw [hg.he, hgf.he], simp
  end }

-- Product of two bijections.
def Bij_on.prod {f : α → β} {a : set α} {b : set β}
  {g : γ → δ} {c : set γ} {d : set δ} (hf : Bij_on f a b) (hg : Bij_on g c d) :
  Bij_on (λ (p : α × γ), (f p.1, g p.2)) (a.prod c) (b.prod d) :=
{ e :=
  -- This is a bit ugly. But building e out of equiv.prod_congr made
  -- the proof of `he` more difficult. Part of the problem is that
  -- `equiv.set.prod` is too "strict": its to_fun pattern matches on
  -- its argument.
  { to_fun := λ p,
      ⟨(hf.e ⟨p.1.1, p.2.1⟩, hg.e ⟨p.1.2, p.2.2⟩),
       ⟨(hf.e ⟨p.1.1, p.2.1⟩).property, (hg.e ⟨p.1.2, p.2.2⟩).property⟩⟩,
    inv_fun := λ p,
      ⟨(hf.e.symm ⟨p.1.1, p.2.1⟩, hg.e.symm ⟨p.1.2, p.2.2⟩),
       ⟨(hf.e.symm ⟨p.1.1, p.2.1⟩).property, (hg.e.symm ⟨p.1.2, p.2.2⟩).property⟩⟩,
    left_inv := λ ⟨⟨x, y⟩, ⟨hx, hy⟩⟩, by simp,
    right_inv := λ ⟨⟨x, y⟩, ⟨hx, hy⟩⟩, by simp },
  he := λ ⟨⟨x, y⟩, ⟨hx, hy⟩⟩, show (f (subtype.mk x hx), g (subtype.mk y hy)) = _,
    by rw [hf.he, hg.he]; simp }

-- Product of two total bijections.
-- Again, we need to use this instead of `equiv.prod_congr` because
-- the latter is too strict.
def Bij_on.prod' {f : α → β} {g : γ → δ}
  (hf : Bij_on f univ univ) (hg : Bij_on g univ univ) :
  Bij_on (λ (p : α × γ), (f p.1, g p.2)) univ univ :=
begin convert Bij_on.prod hf hg; ext p; simp end

-- Product of a total bijection by a type.
def Bij_on.prod_right' {f : α → β} {b : set β} (hf : Bij_on f univ b) :
  Bij_on (λ (p : α × γ), (f p.1, p.2)) univ (b.prod univ) :=
by convert Bij_on.prod hf (Bij_on.refl univ); simp

-- Restriction of a bijection to a subset.
def Bij_on.restrict (h : Bij_on f a b) (r : set β) : Bij_on f (f ⁻¹' r ∩ a) (r ∩ b) :=
{ e :=
  { to_fun := λ p,
      ⟨h.e ⟨p.val, p.property.right⟩,
       by rw ←h.he; exact p.property.left,
       (h.e _).property⟩,
    inv_fun := λ p,
      ⟨h.e.symm ⟨p.val, p.property.right⟩,
       by rw [mem_preimage, h.he]; simpa using p.property.left,
       (h.e.symm _).property⟩,
    left_inv := λ ⟨p, hp⟩, by simp,
    right_inv := λ ⟨p, hp⟩, by simp },
  he := λ ⟨p, hp⟩, by change f p = _; rw h.he' hp.right; simp }

-- Restriction of a total bijection to a subset.
def Bij_on.restrict' (h : Bij_on f univ b) (r : set β) : Bij_on f (f ⁻¹' r) (r ∩ b) :=
by convert Bij_on.restrict h r; simp

def Bij_on.restrict'' (h : Bij_on f univ univ) (r : set β) : Bij_on f (f ⁻¹' r) r :=
by convert Bij_on.restrict' h r; simp

def Bij_on.restrict_equiv (e : α ≃ β) (r : set β) : Bij_on e.to_fun (e.to_fun ⁻¹' r) r :=
(Bij_on.of_equiv e).restrict'' r

-- Restriction of a bijection to a subtype on both sides.
-- TODO: Reduce duplicated ugliness with Bij_on.restrict?
def Bij_on.restrict_to_subtype (h : Bij_on f a b) (r : β → Prop) :
  Bij_on (λ (x : subtype (f ⁻¹' r)), (⟨f x.val, x.property⟩ : subtype r))
    {x | x.val ∈ a} {y | y.val ∈ b} :=
{ e :=
  { to_fun := λ p,
      ⟨⟨h.e ⟨p.val.val, p.property⟩,
       by rw ←h.he; exact p.val.property⟩,
       (h.e _).property⟩,
    inv_fun := λ p,
      ⟨⟨h.e.symm ⟨p.val.val, p.property⟩,
       show ↑(h.e.symm ⟨p.val.val, p.property⟩) ∈ f ⁻¹' r,
       by rw [mem_preimage, h.he]; simp⟩,
       (h.e.symm _).property⟩,
    left_inv := λ ⟨p, hp⟩, by simp,
    right_inv := λ ⟨p, hp⟩, by simp },
  he := λ ⟨p, hp⟩, by apply subtype.eq; change f p = _; rw h.he'; refl }

-- Bijection between a subtype and a propositionally equal one.
def Bij_on.congr_subtype {r r' : set α} (h : r = r') :
  Bij_on (λ (x : subtype r), (⟨x, _⟩ : subtype r')) univ univ :=
Bij_on.of_equiv $ equiv.set_congr h

-- TODO: Use this to simplify other colimit lemmas?
def Bij_on.congr_subset {α : Type*} {r r' : set α} (h : r = r') : Bij_on id r r' :=
{ e :=
  { to_fun := λ p, ⟨p.val, h ▸ p.property⟩,
    inv_fun := λ p, ⟨p.val, h.symm ▸ p.property⟩,
    left_inv := λ p, by cases p; refl,
    right_inv := λ p, by cases p; refl },
  he := λ p, by cases p; refl }

end Bij_on


import Logic.Predicate.Language

variable (L : Language.{u}) (L₁ L₂ L₃ : Language)

inductive SubTerm (μ : Type v) (n : ℕ)
  | bvar : Fin n → SubTerm μ n
  | fvar : μ → SubTerm μ n
  | func : ∀ {arity}, L.func arity → (Fin arity → SubTerm μ n) → SubTerm μ n

prefix:max "&" => SubTerm.fvar
prefix:max "#" => SubTerm.bvar

variable (μ : Type v) (μ₁ : Type v₁) (μ₂ : Type v₂) (μ₃ : Type v₃)

abbrev Term := SubTerm L μ 0

abbrev SyntacticSubTerm (n : ℕ) := SubTerm L ℕ n

abbrev SyntacticTerm := SyntacticSubTerm L 0

namespace SubTerm
variable {μ μ₁ μ₂ μ₃}

instance [Inhabited μ] : Inhabited (SubTerm L μ n) := ⟨&default⟩

abbrev func! (k) (f : L.func k) (v : Fin k → SubTerm L μ n) := func f v

variable {L}

variable [∀ k, ToString (L.func k)] [∀ k, ToString (L.rel k)] [ToString μ]

def toStr : SubTerm L μ n → String
  | #x                        => "x_{" ++ toString (n - 1 - (x : ℕ)) ++ "}"
  | &x                        => "z_{" ++ toString x ++ "}"
  | func (arity := 0) c _     => toString c
  | func (arity := _ + 1) f v => "{" ++ toString f ++ "} \\left(" ++ String.vecToStr (fun i => toStr (v i)) ++ "\\right)"

instance : Repr (SubTerm L μ n) := ⟨fun t _ => toStr t⟩

instance : ToString (SubTerm L μ n) := ⟨toStr⟩

variable {n n₁ n₂ n₃ m m₁ m₂ m₃ : ℕ}

def bind (bound : Fin n₁ → SubTerm L μ₂ n₂) (free : μ₁ → SubTerm L μ₂ n₂) :
    SubTerm L μ₁ n₁ → SubTerm L μ₂ n₂
  | (#x)       => bound x
  | (&x)       => free x
  | (func f v) => func f (fun i => (v i).bind bound free)

def map (bound : Fin n₁ → Fin n₂) (free : μ₁ → μ₂) : SubTerm L μ₁ n₁ → SubTerm L μ₂ n₂ :=
  bind (fun n => #(bound n)) (fun m => &(free m))

def subst (t : SubTerm L μ n) : SubTerm L μ (n + 1) → SubTerm L μ n :=
  bind (bvar <: t) fvar

def bShift : SubTerm L μ n → SubTerm L μ (n + 1) :=
  map Fin.succ id

def castLe {n n' : ℕ} (h : n ≤ n') : SubTerm L μ n → SubTerm L μ n' :=
  map (Fin.castLe h) id

section bind
variable (bound : Fin n₁ → SubTerm L μ₂ n₂) (free : μ₁ → SubTerm L μ₂ n₂)

@[simp] lemma bind_fvar (m : μ₁) : (&m : SubTerm L μ₁ n₁).bind bound free = free m := rfl

@[simp] lemma bind_bvar (n : Fin n₁) : (#n : SubTerm L μ₁ n₁).bind bound free = bound n := rfl

lemma bind_func {k} (f : L.func k) (v : Fin k → SubTerm L μ₁ n₁) :
    (func f v).bind bound free = func f (fun i => (v i).bind bound free) := rfl

end bind

lemma bind_bind
  (bound₁ : Fin n₁ → SubTerm L μ₂ n₂) (free₁ : μ₁ → SubTerm L μ₂ n₂)
  (bound₂ : Fin n₂ → SubTerm L μ₃ n₃) (free₂ : μ₂ → SubTerm L μ₃ n₃) (t : SubTerm L μ₁ n₁) :
    (t.bind bound₁ free₁).bind bound₂ free₂ = t.bind (fun n => (bound₁ n).bind bound₂ free₂) (fun m => (free₁ m).bind bound₂ free₂) :=
  by induction t <;> simp[*, bind_func]

@[simp] lemma bind_id (t) : @bind L μ μ n n bvar fvar t = t :=
  by induction t <;> simp[*, bind_func]

@[simp] lemma bind_id₀ (t) : @bind L μ μ 0 0 finZeroElim fvar t = t :=
  by simpa[eq_finZeroElim] using bind_id t

lemma bind_id_of_eq (hbound : ∀ x, bound x = #x) (hfree : ∀ x, free x = &x) (t) : @bind L μ μ n n bound free t = t :=
  by
  have e₁ : bvar = bound := by funext x; simp[hbound]
  have e₂ : fvar = free := by funext x; simp[hfree]
  exact e₁ ▸ e₂ ▸ bind_id t

section map
variable (bound : Fin n₁ → Fin n₂) (free : μ₁ → μ₂)

@[simp] lemma map_fvar (m : μ₁) : (&m : SubTerm L μ₁ n₁).map bound free = &(free m) := rfl

@[simp] lemma map_bvar (n : Fin n₁) : (#n : SubTerm L μ₁ n₁).map bound free = #(bound n) := rfl

lemma map_func {k} (f : L.func k) (v : Fin k → SubTerm L μ₁ n₁) :
    (func f v).map bound free = func f (fun i => (v i).map bound free) := rfl

end map

lemma map_map
  (bound₁ : Fin n₁ → Fin n₂) (free₁ : μ₁ → μ₂)
  (bound₂ : Fin n₂ → Fin n₃) (free₂ : μ₂ → μ₃) (t : SubTerm L μ₁ n₁) :
    (t.map bound₁ free₁).map bound₂ free₂ = t.map (bound₂ ∘ bound₁) (free₂ ∘ free₁) :=
  bind_bind _ _ _ _ _

@[simp] lemma map_id (t) : @map L μ μ n n id id t = t :=
  by induction t <;> simp[*, map_func]

lemma map_inj {bound : Fin n₁ → Fin n₂} {free : μ₁ → μ₂} (hb : Function.Injective bound) (hf : Function.Injective free) :
    Function.Injective $ map (L := L) bound free
  | #x,                    t => by cases t <;> simp[map_func]; intro h; exact hb h
  | &x,                    t => by cases t <;> simp[map_func]; intro h; exact hf h
  | func (arity := k) f v, t => by
    cases t <;> simp[*, map_func]
    case func =>
      rintro rfl; simp; rintro rfl h; simp
      funext i; exact map_inj hb hf (congr_fun h i)

@[simp] lemma bShift_bvar (x : Fin n) : bShift (#x : SubTerm L μ n) = #(Fin.succ x) := rfl

@[simp] lemma bShift_fvar (x : μ) : bShift (&x : SubTerm L μ n) = &x := rfl

lemma bShift_func {k} (f : L.func k) (v : Fin k → SubTerm L μ n) :
  bShift (func f v) = func f (fun i => bShift (v i)) := rfl

@[simp] lemma leftConcat_bShift_comp_bvar :
    (#0 :> bShift ∘ bvar : Fin (n + 1) → SubTerm L μ (n + 1)) = bvar :=
  funext (Fin.cases (by simp) (by simp))

@[simp] lemma bShift_comp_fvar :
    (bShift ∘ fvar : μ → SubTerm L μ (n + 1)) = fvar :=
  funext (by simp)

@[simp] lemma subst_bvar_last (s : SubTerm L μ n) : subst s #(Fin.last n) = s :=
  by simp[subst]

@[simp] lemma subst_bvar_castSucc (s : SubTerm L μ n) (x : Fin n) : subst s #(Fin.castSucc x) = #x :=
  by simp[subst]

@[simp] lemma subst_fvar (s : SubTerm L μ n) (x : μ) : subst s &x = &x :=
  by simp[subst]

lemma subst_func (s : SubTerm L μ n) {k} (f : L.func k) (v : Fin k → SubTerm L μ (n + 1)) :
    subst s (func f v) = func f (fun i => subst s (v i)) :=
  by simp[subst, bind_func]

@[simp] lemma castLe_bvar {n'} (h : n ≤ n') (x : Fin n) : castLe h (#x : SubTerm L μ n) = #(Fin.castLe h x) := rfl

@[simp] lemma castLe_fvar {n'} (h : n ≤ n') (x : μ) : castLe h (&x : SubTerm L μ n) = &x := rfl

lemma castLe_func {n'} (h : n ≤ n') {k} (f : L.func k) (v : Fin k → SubTerm L μ n) :
    castLe h (func f v) = func f (fun i => castLe h (v i)) := rfl

section Syntactic

/-
  #0 #1 ... #(n - 1) &0 &1 ...
   ↓shift
  #0 #1 ... #(n - 1) &1 &2 &3 ...
-/

def shift : SyntacticSubTerm L n → SyntacticSubTerm L n :=
  map id Nat.succ

def shift_le (s : ℕ) : SyntacticSubTerm L n → SyntacticSubTerm L n :=
  map id (fun m => m + s)

/-
  #0 #1 ... #(n - 1) #n &0 &1 ...
   ↓free           ↑fix
  #0 #1 ... #(n - 1) &0 &1 &2 ...
 -/

def free : SyntacticSubTerm L (n + 1) → SyntacticSubTerm L n :=
  bind (bvar <: &0) (fun m => &(Nat.succ m))

def fix : SyntacticSubTerm L n → SyntacticSubTerm L (n + 1) :=
  bind (fun x => #(Fin.castSucc x)) (#(Fin.last n) :>ₙ fvar)

@[simp] lemma shift_bvar (x : Fin n) : shift (#x : SyntacticSubTerm L n) = #x := rfl

@[simp] lemma shift_fvar (x : ℕ) : shift (&x : SyntacticSubTerm L n) = &(x + 1) := rfl

@[simp] lemma shift_func {k} (f : L.func k) (v : Fin k → SyntacticSubTerm L n) :
    shift (func f v) = func f (fun i => shift (v i)) := rfl

lemma shift_Injective : Function.Injective (@shift L n) :=
  Function.LeftInverse.injective (g := map id Nat.pred)
    (by intros p; simp[shift, map_map, Function.comp]; exact map_id _)

@[simp] lemma free_bvar_castSucc (x : Fin n) : free (#(Fin.castSucc x) : SyntacticSubTerm L (n + 1)) = #x := by simp[free]

@[simp] lemma free_bvar_last : free (#(Fin.last n) : SyntacticSubTerm L (n + 1)) = &0 := by simp[free]

@[simp] lemma free_fvar (x : ℕ) : free (&x : SyntacticSubTerm L (n + 1)) = &(x + 1) := by simp[free]

lemma free_func {k} (f : L.func k) (v : Fin k → SyntacticSubTerm L (n + 1)) :
    free (func f v) = func f (fun i => free $ v i) := by simp[free, bind_func]

@[simp] lemma fix_bvar (x : Fin n) : fix (#x : SyntacticSubTerm L n) = #(Fin.castSucc x) := by simp[fix]

@[simp] lemma fix_fvar_zero : fix (&0 : SyntacticSubTerm L n) = #(Fin.last n) := by simp[fix]

@[simp] lemma fix_fvar_succ (x : ℕ) : fix (&(x + 1) : SyntacticSubTerm L n) = &x := by simp[fix]

lemma fix_func {k} (f : L.func k) (v : Fin k → SyntacticSubTerm L n) :
    fix (func f v) = func f (fun i => fix $ v i) := by simp[fix, bind_func]

@[simp] lemma free_fix (t : SyntacticSubTerm L n) : free (fix t) = t :=
  by simp[free, fix, bind_bind]; exact bind_id_of_eq (by simp) (by intro x; cases x <;> simp) t

@[simp] lemma fix_free (t : SyntacticSubTerm L (n + 1)) : fix (free t) = t :=
  by simp[free, fix, bind_bind]; exact bind_id_of_eq (by intro x; cases x using Fin.lastCases <;> simp) (by simp) t

lemma bShift_free_eq_shift (t : SyntacticTerm L) : free (bShift t) = shift t :=
  by simp[free, bShift, shift, map, bind_bind, eq_finZeroElim]

end Syntactic

def fvarList : SubTerm L μ n → List μ
  | #_       => []
  | &x       => [x]
  | func _ v => List.join $ Matrix.toList (fun i => fvarList (v i))

abbrev fvar? (t : SubTerm L μ n) (x : μ) : Prop := x ∈ t.fvarList

@[simp] lemma fvarList_bvar : fvarList (#x : SubTerm L μ n) = [] := rfl

@[simp] lemma fvarList_fvar : fvarList (&x : SubTerm L μ n) = [x] := rfl

@[simp] lemma mem_fvarList_func {k} {f : L.func k} {v : Fin k → SubTerm L μ n} :
    x ∈ (func f v).fvarList ↔ ∃ i, x ∈ (v i).fvarList :=
  by simp[fvarList]

lemma bind_eq_of_funEqOn (bound : Fin n₁ → SubTerm L μ₂ n₂) (free₁ free₂ : μ₁ → SubTerm L μ₂ n₂) (t : SubTerm L μ₁ n₁)
  (h : Function.funEqOn t.fvar? free₁ free₂) :
    t.bind bound free₁ = t.bind bound free₂ := by
  induction t <;> simp[bind_func]
  case fvar => simpa[fvar?, Function.funEqOn] using h
  case func k f v ih =>
    funext i
    exact ih i (h.of_subset $ by simp[fvar?]; intro x hx; exact ⟨i, hx⟩)

variable [∀ k, DecidableEq (L.func k)] [∀ k, DecidableEq (L.rel k)]

def languageFunc : SubTerm L μ n → Finset (Σ k, L.func k)
  | #_       => ∅
  | &_       => ∅
  | func f v => insert ⟨_, f⟩ $ Finset.bunionᵢ Finset.univ (fun i => languageFunc (v i))

@[simp] lemma languageFunc_func {k} (f : L.func k) (v : Fin k → SubTerm L μ n) :
    ⟨k, f⟩ ∈ (func f v).languageFunc := by simp[languageFunc]

lemma languageFunc_func_ss {k} (f : L.func k) (v : Fin k → SubTerm L μ n) (i) :
    (v i).languageFunc ⊆ (func f v).languageFunc :=
  by intros x; simp[languageFunc]; intros h; exact Or.inr ⟨i, h⟩

variable [DecidableEq μ]

def hasDecEq : (t u : SubTerm L μ n) → Decidable (Eq t u)
  | #x,                   #y                   => by simp; exact decEq x y
  | #_,                   &_                   => isFalse SubTerm.noConfusion
  | #_,                   func _ _             => isFalse SubTerm.noConfusion
  | &_,                   #_                   => isFalse SubTerm.noConfusion
  | &x,                   &y                   => by simp; exact decEq x y
  | &_,                   func _ _             => isFalse SubTerm.noConfusion
  | func _ _,             #_                   => isFalse SubTerm.noConfusion
  | func _ _,             &_                   => isFalse SubTerm.noConfusion
  | @func L μ _ k₁ r₁ v₁, @func L μ _ k₂ r₂ v₂ => by
      by_cases e : k₁ = k₂
      · rcases e with rfl
        exact match decEq r₁ r₂ with
        | isTrue h => by simp[h]; exact Matrix.decVec _ _ (fun i => hasDecEq (v₁ i) (v₂ i))
        | isFalse h => isFalse (by simp[h])
      · exact isFalse (by simp[e])

instance : DecidableEq (SubTerm L μ n) := hasDecEq

end SubTerm

namespace Language

namespace Hom
variable {L L₁ L₂ L₃} {μ} (Φ : Hom L₁ L₂)

def onSubTerm (Φ : Hom L₁ L₂) : SubTerm L₁ μ n → SubTerm L₂ μ n
  | #x               => #x
  | &x               => &x
  | SubTerm.func f v => SubTerm.func (Φ.onFunc f) (fun i => onSubTerm Φ (v i))

@[simp] lemma onSubTerm_bvar (x : Fin n) : Φ.onSubTerm (#x : SubTerm L₁ μ n) = #x := rfl

@[simp] lemma onSubTerm_fvar (x : μ) : Φ.onSubTerm (&x : SubTerm L₁ μ n) = &x := rfl

lemma onSubTerm_func {k} (f : L₁.func k) (v : Fin k → SubTerm L₁ μ n) :
    Φ.onSubTerm (SubTerm.func f v) = SubTerm.func (Φ.onFunc f) (fun i => onSubTerm Φ (v i)) := rfl

end Hom

end Language

namespace SubTerm
open Language.Hom
section
variable {L₁ L₂ : Language} (Φ : L₁ →ᵥ L₂) {μ₁ μ₂ : Type _} {n₁ n₂ : ℕ}

lemma onSubTerm_bind (bound : Fin n₁ → SubTerm L₁ μ₂ n₂) (free : μ₁ → SubTerm L₁ μ₂ n₂) (t) :
    Φ.onSubTerm (bind bound free t) = bind (fun x => Φ.onSubTerm (bound x)) (fun x => Φ.onSubTerm (free x)) (Φ.onSubTerm t) :=
  by induction t <;> simp[*, onSubTerm_func, bind_func]

lemma onSubTerm_map (bound : Fin n₁ → Fin n₂) (free : μ₁ → μ₂) (t) :
    Φ.onSubTerm (map bound free t) = map bound free (Φ.onSubTerm t) :=
  by simp[map, onSubTerm_bind]

lemma onSubTerm_subst (u) (t : SubTerm L₁ μ (n + 1)) :
    Φ.onSubTerm (subst u t) = subst (Φ.onSubTerm u) (Φ.onSubTerm t) :=
  by simp[subst, onSubTerm_bind, Matrix.comp_vecConsLast, Function.comp]

lemma onSubTerm_bShift (t : SubTerm L₁ μ₁ n) : Φ.onSubTerm (bShift t) = bShift (Φ.onSubTerm t) :=
  by simp[bShift, onSubTerm_map]

lemma onSubTerm_shift (t : SyntacticSubTerm L₁ n) : Φ.onSubTerm (shift t) = shift (Φ.onSubTerm t) :=
  by simp[shift, onSubTerm_map]

lemma onSubTerm_free (t : SyntacticSubTerm L₁ (n + 1)) : Φ.onSubTerm (free t) = free (Φ.onSubTerm t) :=
  by simp[free, onSubTerm_bind]; congr; exact funext $ Fin.lastCases (by simp) (by simp)

lemma onSubTerm_fix (t : SyntacticSubTerm L₁ n) : Φ.onSubTerm (fix t) = fix (Φ.onSubTerm t) :=
  by simp[fix, onSubTerm_bind]; congr; funext x; cases x <;> simp

end

section
open Language
variable {L : Language} [∀ k, DecidableEq (L.func k)] {μ n}

def toSubLanguage' (pf : ∀ k, L.func k → Prop) (pr : ∀ k, L.rel k → Prop) : ∀ t : SubTerm L μ n,
    (∀ k f, ⟨k, f⟩ ∈ t.languageFunc → pf k f) → SubTerm (subLanguage L pf pr) μ n
  | #x,                _ => #x
  | &x,                _ => &x
  | func (arity := k) f v, h => func ⟨f, h k f (by simp)⟩
      (fun i => toSubLanguage' pf pr (v i) (fun k' f' h' => h k' f' (languageFunc_func_ss f v i h')))

@[simp] lemma onSubTerm_toSubLanguage' (pf : ∀ k, L.func k → Prop) (pr : ∀ k, L.rel k → Prop)
  (t : SubTerm L μ n) (h : ∀ k f, ⟨k, f⟩ ∈ t.languageFunc → pf k f) :
    L.ofSubLanguage.onSubTerm (t.toSubLanguage' pf pr h) = t :=
  by induction t <;> simp[*, toSubLanguage', onSubTerm_func]

end

section
open Language
variable {L : Language} [hz : L.HasZero] [ho : L.HasOne] [ha : L.HasAdd] {μ : Type v} {μ₁ μ₂} {n : ℕ} {n₁ n₂}

def natLit : ℕ → SubTerm L μ n
  | 0     => func Language.HasZero.zero ![]
  | n + 1 => func Language.HasAdd.add ![natLit n, func Language.HasOne.one ![]]

lemma natLit_zero : (natLit 0 : SubTerm L μ n) = func Language.HasZero.zero ![] := by rfl

lemma natLit_succ (z : ℕ) :
  (natLit (z + 1) : SubTerm L μ n) = func Language.HasAdd.add ![natLit z, func Language.HasOne.one ![]] := by rfl

@[simp] lemma bind_natLit (z : ℕ) (bound : Fin n₁ → SubTerm L μ₂ n₂) (free : μ₁ → SubTerm L μ₂ n₂) :
    bind bound free (natLit z) = natLit z := by
  induction' z with z ih <;> simp[natLit_zero, natLit_succ, bind_func]
  funext i; cases i using Fin.cases <;> simp[ih, bind_func]

end

declare_syntax_cat subterm
syntax:max "#" num : subterm
syntax:max "&" term:max : subterm
syntax:max "!" term:max : subterm
syntax num : subterm
syntax:70 "const" term:max : subterm
syntax:70 "func¹" term "/[" subterm:0 "]" : subterm
syntax:70 "func²" term "/[" subterm:0 "," subterm:0 "]" : subterm
syntax:50 subterm:50 "+" subterm:51 : subterm
syntax:60 subterm:60 "*" subterm:61 : subterm
syntax "(" subterm ")" : subterm

syntax "T“" subterm "”" : term

macro_rules
  | `(T“ # $n:num ”)                                     => `(#$n)
  | `(T“ & $n:term ”)                                    => `(&$n)
  | `(T“ ! $t:term ”)                                    => `($t)
  | `(T“ $n:num ”)                                       => `(natLit $n)
  | `(T“ const $d:term ”)                                => `(func $d ![])
  | `(T“ func¹ $d:term /[ $t:subterm ] ”)                => `(func $d ![T“$t”])
  | `(T“ func² $d:term /[ $t₁:subterm , $t₂:subterm ] ”) => `(func $d ![T“$t₁”, T“$t₂”])
  | `(T“ $t:subterm + $u:subterm ”)                      => `(func Language.HasAdd.add ![T“$t”, T“$u”])
  | `(T“ $t:subterm * $u:subterm ”)                      => `(func Language.HasMul.mul ![T“$t”, T“$u”])
  | `(T“ ( $x ) ”)                                       => `(T“$x”)

#reduce (T“ func² Language.ORingFunc.mul /[&2 + &0, const Language.ORingFunc.zero]” : SubTerm Language.oring ℕ 8)
#reduce (T“(&2 + &0) * #2” : SubTerm Language.oring ℕ 8)

end SubTerm


/-
Copyright (c) 2018 Sander Dahmen, Johannes Hölzl, Robert Y. Lewis. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Sander Dahmen, Johannes Hölzl, Robert Y. Lewis

"On large subsets of 𝔽ⁿ_q with no three-term arithmetic progression"
by J. S. Ellenberg and D. Gijswijt

This file develops just enough elementary calculus to prove a fact needed in section 13.
It includes a proof of the product rule for functions ℝ → ℝ.
-/
import analysis.calculus.deriv

open filter

lemma is_bounded_linear_map_mul_const (r : ℝ) : is_bounded_linear_map ℝ ((*) r) :=
show is_bounded_linear_map ℝ (λx:ℝ, r • x), from
  is_bounded_linear_map.smul _ is_bounded_linear_map.id

noncomputable def mul_const_bounded_linear_map (r : ℝ) : ℝ →L[ℝ] ℝ :=
(is_bounded_linear_map_mul_const r).to_continuous_linear_map

noncomputable def continuous_linear_map.to_fun' (f : ℝ →L[ℝ] ℝ) : ℝ → ℝ :=
f.to_fun

@[simp] lemma mul_const_bounded_linear_map_apply (r x : ℝ) :
  (mul_const_bounded_linear_map r).to_fun' x = r * x :=
rfl

open asymptotics

section
variables {α : Type*} {β : Type*} {γ : Type*}
variables [normed_field β] [normed_field γ]

theorem is_o_mul_right_one {f₁ f₂ : α → β} {g : α → γ} {l : filter α}
    (h₁ : is_o f₁ g l) (h₂ : is_O f₂ (λx, (1:γ)) l):
  is_o (λ x, f₁ x * f₂ x) (λ x, g x) l :=
have is_o (λ x, f₁ x * f₂ x) (λ x, g x * 1) l := is_o_mul_right h₁ h₂,
by convert this; funext; rw mul_one

end

def has_deriv (f : ℝ → ℝ) (f' : ℝ) (x : ℝ) : Prop :=
has_fderiv_at f (mul_const_bounded_linear_map f') x --((*) f') x

lemma has_deriv.congr₂ {f : ℝ → ℝ} {f' f'' x : ℝ} (eq : f' = f'') (h : has_deriv f f' x) :
  has_deriv f f'' x :=
by rwa ← eq

lemma has_deriv_const (r x : ℝ) : has_deriv (λx, r) 0 x :=
(has_fderiv_at_const r _).congr (assume x, rfl) (assume x, (zero_mul _).symm)

lemma has_deriv_id (x : ℝ) : has_deriv (λx, x) 1 x :=
(has_fderiv_at_id x).congr (assume x, rfl) (assume x, (one_mul _).symm)

lemma has_deriv.add {x : ℝ} {f g : ℝ → ℝ} {f' g' : ℝ}
  (hf : has_deriv f f' x) (hg : has_deriv g g' x) : has_deriv (λx, f x + g x) (f' + g') x :=
(hf.add hg).congr (assume x, rfl) (assume x, (add_mul _ _ _).symm)

lemma has_deriv.sub {x : ℝ} {f g : ℝ → ℝ} {f' g' : ℝ}
  (hf : has_deriv f f' x) (hg : has_deriv g g' x) : has_deriv (λx, f x - g x) (f' - g') x :=
(hf.sub hg).congr (assume x, rfl) $ λ x,
  show f' * x - g' * x = (f' - g') * x, by simp [right_distrib]

lemma has_deriv.neg {x : ℝ} {f : ℝ → ℝ} {f' : ℝ} (hf : has_deriv f f' x) :
  has_deriv (λx, - f x) (- f') x :=
(hf.neg).congr (assume x, rfl) $ λ x, show -(f' * x) = (-f') * x, by simp

lemma has_deriv_finset_sum {α : Type*} {x : ℝ} {f : α → ℝ → ℝ} {f' : α → ℝ}
  (s : finset α) (hf : ∀a, has_deriv (f a) (f' a) x) :
  has_deriv (λx, s.sum (λa, f a x)) (s.sum f') x :=
begin
  letI := classical.dec_eq α,
  refine s.induction_on _ _,
  { simp only [finset.sum_empty],
    exact has_deriv_const 0 x },
  { assume a s has ih,
    simp only [finset.sum_insert, has, not_false_iff],
    exact (hf a).add ih }
end

lemma has_deriv.mul {f g : ℝ → ℝ} {f' g' : ℝ} {x : ℝ}
  (hf : has_deriv f f' x) (hg : has_deriv g g' x) :
  has_deriv (λx, f x * g x) (f x * g' + f' * g x) x :=
begin
  unfold has_deriv,
  convert has_fderiv_at.mul hf hg using 1,
  ext, dsimp,
  change continuous_linear_map.to_fun' _ _ = _ * continuous_linear_map.to_fun' _ _ + _ * continuous_linear_map.to_fun' _ _,
  simp, ring
end

lemma has_deriv.mul_left {f : ℝ → ℝ} {f' : ℝ} {x : ℝ} (c : ℝ) (hf : has_deriv f f' x) :
  has_deriv (λx, c * f x) (c * f') x :=
have _ := (has_deriv_const c x).mul hf,
by simpa

lemma has_deriv.pow {f : ℝ → ℝ} {f' : ℝ} {x : ℝ} (hf : has_deriv f f' x) :
  ∀n:ℕ, has_deriv (λx, (f x) ^ n) (n * (f x)^(n - 1) * f') x
| 0       := by simp only [has_deriv_const 1 x, nat.zero_sub, nat.cast_zero, zero_mul, pow_zero]
| 1       := by simp only [hf, mul_one, one_mul, nat.sub_self, pow_one, nat.cast_one, pow_zero]
| (n + 1 + 1) :=
  begin
    refine (hf.mul (has_deriv.pow (n + 1))).congr (assume x, rfl) (assume x, _),
    change continuous_linear_map.to_fun' _ _ = continuous_linear_map.to_fun' _ _,
    simp only [mul_const_bounded_linear_map_apply, nat.add_sub_cancel, nat.cast_add, nat.cast_one],
    simp only [add_mul, mul_add, pow_add, pow_one, one_mul, add_comm, pow_one],
    ac_refl
  end

lemma increasing_of_deriv_zero_pos (f : ℝ → ℝ) (f' : ℝ) (hf : has_deriv f f' 0) (hf' : f' > 0) :
  ∃ε>0, ∀x, 0 < x → x < ε → f 0 < f x :=
begin
  have := (has_fderiv_at_filter_iff_tendsto.1 hf),
  simp only [sub_zero, (norm_inv _).symm, (normed_field.norm_mul _ _).symm] at this,
  rw [← @tendsto_zero_iff_norm_tendsto_zero ℝ ℝ, metric.tendsto_nhds_nhds] at this,
  specialize this f' hf',
  rcases this with ⟨ε, hε, h⟩,
  refine ⟨ε, hε, assume x hx0 hxε, _⟩,
  have : dist x 0 < ε,
  { rwa [dist_zero_right, real.norm_eq_abs, abs_of_pos hx0] },
  specialize @h x this,
  change dist (_*(_ - continuous_linear_map.to_fun' _ _)) _ < _ at h,
  rw [mul_const_bounded_linear_map_apply, dist_zero_right, mul_comm f', mul_sub, ← mul_assoc, inv_mul_cancel (ne_of_gt hx0), one_mul,
    norm_sub_rev, real.norm_eq_abs, abs_sub_lt_iff, sub_lt_self_iff] at h,
  exact (sub_pos.1 $ pos_of_mul_pos_left h.1 $ inv_nonneg.2 $ le_of_lt $ hx0)
end

lemma decreasing_of_fderiv_pos (f : ℝ → ℝ) (f' : ℝ) (x : ℝ) (hf : has_deriv f f' x) (hf' : 0 < f') :
  ∃ε>0, ∀y, x - ε < y → y < x → f y < f x :=
begin
  have : mul_const_bounded_linear_map (-f') = continuous_linear_map.comp (mul_const_bounded_linear_map (f')) (mul_const_bounded_linear_map (-1)),
  { ext x, show -f' * x = f' * (-1 * x), simp },
  have : has_deriv (λx':ℝ, - (f ∘ (λy, x - y)) x') (f') 0,
  { convert @has_deriv.neg _ _ (-f') _ using 1,
    { rw neg_neg },
    { unfold has_deriv at hf ⊢, dsimp, rw ←sub_zero x at hf,
      convert has_fderiv_at.comp _ hf _ using 2,
      convert @has_deriv.sub _ _ _ 0 1 _ _,
      { norm_num },
      { apply has_deriv_const },
      { apply has_deriv_id } } },
  rcases increasing_of_deriv_zero_pos _ _ this hf' with ⟨ε, hε, h⟩,
  refine ⟨ε, hε, assume y hyε hyx, _⟩,
  specialize h (x - y),
  simp [-sub_eq_add_neg, sub_sub_cancel, sub_zero] at h,
  refine h hyx (sub_lt.2 hyε)
end