import data.vector
import data.fin.vec_notation
import data.fin.tuple.basic
import data.list.of_fn
import data.list.alist
import data.finsupp.basic
import control.bifunctor
import tactic.derive_fintype
import tactic.fin_cases
import finsupp_lemmas
import frames
import verification.vars
import verification.misc
import verification.stream
import data.pfun

section

parameters (R : Type)
open Types (nn rr bb)
open NameSpace (reserved)
open Vars (ind₀ vals len output)

section compiler

parameters [add_comm_monoid R] [has_one R] [has_mul R]

@[reducible]
def ExprVal : Types → Type
| nn := ℕ
| rr := R
| bb := bool

parameter {R}
namespace ExprVal

instance : ∀ b, inhabited (ExprVal b)
| nn := ⟨0⟩
| rr := ⟨0⟩
| bb := ⟨ff⟩

instance [has_to_string R] :
∀ b, has_to_string (ExprVal b)
| nn := infer_instance
| rr := infer_instance
| bb := infer_instance

end ExprVal

@[derive decidable_eq]
inductive Op : Types → Type
| nadd : Op nn | radd : Op rr
| nmul : Op nn | rmul : Op rr
| nsub : Op nn
| and : Op bb
| or : Op bb
| not : Op bb
| nat_eq : Op bb
| lt : Op bb
| le : Op bb
| cast_r : Op rr

namespace Op
instance : ∀ b, has_to_string (Op b)
| rr := ⟨λ v, match v with
| radd := "+"
| rmul := "*"
| cast_r := "cast"
end⟩
| nn := ⟨λ v, match v with
| nadd := "+"
| nmul := "*"
| nsub := "-"
end⟩
| bb := ⟨λ v, match v with
| and := "&&"
| or := "||"
| not := "!"
| nat_eq := "="
| lt := "<"
| le := "<="
end⟩

@[reducible]
def arity : ∀ {b}, Op b → ℕ
| _ nadd := 2
| _ radd := 2
| _ nmul := 2
| _ rmul := 2
| _ nsub := 2
| _ and := 2 | _ or := 2 | _ not := 1 | _ nat_eq := 2 | _ lt := 2 | _ le := 2
| _ cast_r := 1

def is_not_infix : finset (Σ b, Op b) :=
{⟨_, Op.cast_r⟩}

def to_str_with_args {b} (o : Op b) (args : list string) : string :=
if H : (sigma.mk b o) ∈ is_not_infix ∨ 3 ≤ o.arity then
  (to_string o) ++ "(" ++ ", ".intercalate args ++ ")"
else match o.arity, (show ¬(3 ≤ o.arity), by tauto!) with
  0, _ := (to_string o),
  1, _ := (to_string o) ++ args.head,
  2, _ := "(" ++ (args.inth 0) ++ " " ++ (to_string o) ++ " " ++ (args.inth 1) ++ ")",
  (n + 3), h := by { exfalso, simpa using h, }
end

@[reducible]
def signature : ∀ {b} (o : Op b), (fin o.arity → Types)
| _ nadd := ![nn, nn] | _ radd := ![rr, rr]
| _ nmul := ![nn, nn] | _ rmul := ![rr, rr]
| _ nsub := ![nn, nn]
| _ and := ![bb, bb] | _ or := ![bb, bb] | _ not := ![bb]
| _ nat_eq := ![nn, nn] | _ lt := ![nn, nn] | _ le := ![nn, nn]
| _ cast_r := ![nn]

@[simp]
def eval : ∀ {b} (o : Op b), (Π (n : fin o.arity), ExprVal (o.signature n)) → ExprVal b
| _ nadd := λ args, ((+) : ℕ → ℕ → ℕ) (args 0) (args 1)
| _ radd := λ args, ((+) : R → R → R) (args 0) (args 1)
| _ nmul := λ args, ((*) : ℕ → ℕ → ℕ) (args 0) (args 1)
| _ rmul := λ args, ((*) : R → R → R) (args 0) (args 1)
| _ nsub := λ args, nat.sub (args 0) (args 1)
| _ and := λ args, (args 0 : bool) && (args 1 : bool)
| _ or := λ args, (args 0 : bool) || (args 1 : bool)
| _ not := λ args, bnot (args 0)
| _ nat_eq := λ args, args 0 = args 1
| _ lt := λ args, (show ℕ, from args 0) < args 1
| _ le := λ args, (show ℕ, from args 0) ≤ args 1
| _ cast_r := λ args, show ℕ, from args 0

end Op

parameter (R)
inductive Expr : Types → Type
| lit {b} : ExprVal b → Expr b
| ident {b} : Ident b → Expr b
| access {b} : Ident b → Expr nn → Expr b
| call {b} : ∀ o : Op b, (Π (n : fin o.arity), Expr (o.signature n)) → Expr b
| ternary {b} : Expr bb → Expr b → Expr b → Expr b

abbreviation EContext := HeapContext ExprVal
abbreviation Frame := finset (Σ b, Ident b)
instance : inhabited Frame := ⟨(default : finset (Σ b, Ident b))⟩

parameter {R}

def Expr.eval (ctx : EContext) : ∀ {b}, Expr b → ExprVal b
| _ (Expr.lit r) := r
| b (Expr.ident x) := ctx.store.get x
| b (Expr.access x i) := (ctx.heap.get x).inth i.eval
| _ (Expr.call o args) := o.eval (λ i, (args i).eval)
| _ (Expr.ternary c e₁ e₂) := cond c.eval e₁.eval e₂.eval

@[simp] def Expr.frame : ∀ {b}, Expr b → Frame
| _ (Expr.lit r) := ∅
| _ (Expr.ident x) := {sigma.mk _ x}
| _ (Expr.access x i) := insert (sigma.mk _ x) i.frame
| _ (Expr.call o args) := finset.bUnion finset.univ (λ i, (args i).frame)
| _ (Expr.ternary c e₁ e₂) := c.frame ∪ e₁.frame ∪ e₂.frame

-- local notation a ` ⟪<⟫ ` b := Expr.call Op.lt (fin.cons (a : Expr nn) (fin.cons (b : Expr nn) default))

class has_comp (α : Type*) (β : out_param Type*) :=
(eq : α → α → β)
(le : α → α → β)
(lt : α → α → β)
(ge : α → α → β)
(gt : α → α → β)

infix ` ⟪≤⟫ `:50   := has_comp.le
infix ` ⟪<⟫ `:50   := has_comp.lt
infix ` ⟪≥⟫ `:50   := has_comp.ge
infix ` ⟪>⟫ `:50   := has_comp.gt
infix ` ⟪=⟫ `:50   := has_comp.eq

@[simps { attrs := [] }] instance Expr.has_comp : has_comp (Expr nn) (Expr bb) :=
{ eq := λ a b, Expr.call Op.nat_eq $ fin.cons a $ fin.cons b default,
  lt := λ a b, Expr.call Op.lt $ fin.cons a $ fin.cons b default,
  le := λ a b, Expr.call Op.le $ fin.cons a $ fin.cons b default,
  ge := λ a b, Expr.call Op.le $ fin.cons b $ fin.cons a default,
  gt := λ a b, Expr.call Op.lt $ fin.cons b $ fin.cons a default }

section Expr

def expr_repr [has_to_string R] : ∀ {b : Types}, (Expr b) → string
| _ (Expr.lit r) := to_string r
| _ (Expr.ident x) := to_string x
| _ (Expr.access x i) := (to_string x) ++ "[" ++ (expr_repr i) ++ "]"
| _ (Expr.call o args) := o.to_str_with_args (vector.of_fn (λ i, expr_repr $ args i)).to_list
| _ (Expr.ternary c e₁ e₂) := (expr_repr c) ++ " ? " ++ (expr_repr e₁) ++ " : " ++ (expr_repr e₂)

instance {b : Types} [has_to_string R] : has_to_string (Expr b) := ⟨expr_repr⟩

instance Expr.zero_nn : has_zero (Expr nn) := ⟨Expr.lit (0 : ℕ)⟩
instance Expr.one_nn : has_one (Expr nn) := ⟨Expr.lit (1 : ℕ)⟩
instance Expr.zero_rr : has_zero (Expr rr) := ⟨Expr.lit (0 : R)⟩
instance Expr.one_rr : has_one (Expr rr) := ⟨Expr.lit (1 : R)⟩

instance Expr.has_coe_from_nat : has_coe ℕ (Expr nn) := ⟨λ n, Expr.lit n⟩
instance Expr.has_coe_from_R : has_coe R (Expr rr) := ⟨λ r, Expr.lit r⟩

@[simp] lemma Expr_frame_coe_nat (n : ℕ) : (n : Expr nn).frame = ∅ := rfl
@[simp] lemma Expr_frame_coe_R (r : R) : (r : Expr rr).frame = ∅ := rfl
@[simp] lemma Expr_frame_zero_nat : (0 : Expr nn).frame = ∅ := rfl
@[simp] lemma Expr_frame_one_nat : (1 : Expr nn).frame = ∅ := rfl

@[simps { attrs := [] }] instance add_nn : has_add (Expr nn) :=
⟨λ a b, Expr.call Op.nadd (fin.cons a (fin.cons b default))⟩
@[simps { attrs := [] }] instance add_rr : has_add (Expr rr) :=
⟨λ a b, Expr.call Op.radd (fin.cons a (fin.cons b default))⟩
@[simps { attrs := [] }] instance mul_nn : has_mul (Expr nn) :=
⟨λ a b, Expr.call Op.nmul (fin.cons a (fin.cons b default))⟩
@[simps { attrs := [] }] instance mul_rr : has_mul (Expr rr) :=
⟨λ a b, Expr.call Op.rmul (fin.cons a (fin.cons b default))⟩
@[simps { attrs := [] }] instance sub_nn : has_sub (Expr nn) :=
⟨λ a b, Expr.call Op.nsub (fin.cons a (fin.cons b default))⟩

instance inf_bb : has_inf (Expr bb) :=
⟨λ a b, Expr.call Op.and (fin.cons a (fin.cons b default))⟩

instance sup_bb : has_sup (Expr bb) :=
⟨λ a b, Expr.call Op.or (fin.cons a (fin.cons b default))⟩

def Expr.not : Expr bb → Expr bb := λ e, Expr.call Op.not (fin.cons e default)

instance has_coe_to_expr {b : Types} : has_coe (Ident b) (Expr b) := ⟨Expr.ident⟩

@[reducible] def Ident.to_expr {b} : Ident b → Expr b := Expr.ident
@[simp] lemma Expr_frame_coe_ident {b} (i : Ident b) : (i : Expr b).frame = {sigma.mk _ i} := rfl

/- Warning! Lean 3 uses zero, add, one instead of coe from ℕ for numerals -/
example : (3 : Expr nn) = 1 + 1 + 1 := rfl
example : (3 : Expr nn) ≠ Expr.lit 3 := by trivial
example : ((3 : ℕ) : Expr nn) = Expr.lit 3 := rfl

@[simp] lemma Expr.eval_lit {b : Types} (x : ExprVal b) (ctx : EContext) :
  (Expr.lit x).eval ctx = x := rfl
@[simp] lemma Expr.lit_eq_nn (x : ℕ) : @Expr.lit nn x = ↑x := rfl
@[simp] lemma Expr.lit_eq_rr (x : R) : @Expr.lit rr x = ↑x := rfl
@[simp] lemma Expr.eval_lit_nn (x : ℕ) (ctx : EContext) :
  (x : Expr nn).eval ctx = x := rfl
@[simp] lemma Expr.eval_lit_rr (x : R) (ctx : EContext) :
  (x : Expr rr).eval ctx = x := rfl
@[simp] lemma Expr.eval_zero_nn (ctx : EContext) : (0 : Expr nn).eval ctx = 0 := rfl
@[simp] lemma Expr.eval_zero_rr (ctx : EContext) : (0 : Expr rr).eval ctx = 0 := rfl
@[simp] lemma Expr.eval_one_nn (ctx : EContext) : (1 : Expr nn).eval ctx = 1 := rfl
@[simp] lemma Expr.eval_one_rr (ctx : EContext) : (1 : Expr rr).eval ctx = 1 := rfl
@[simp] lemma Expr.eval_ident {b : Types} (x : Ident b) (ctx : EContext) :
  (Expr.ident x).eval ctx = ctx.store.get x := rfl
@[simp] lemma Expr.eval_ident' {b : Types} (x : Ident b) (ctx : EContext) :
  (x : Expr b).eval ctx = ctx.store.get x := rfl
@[simp] lemma Expr.eval_access {b : Types} (x : Ident b) (ind : Expr nn) (ctx : EContext) :
  (Expr.access x ind).eval ctx = (ctx.heap.get x).inth (ind.eval ctx ) := rfl

-- TODO: Derive automatically?
@[simp] lemma Expr.eval_nadd (e₁ e₂ : Expr nn) (ctx : EContext) :
  (e₁ + e₂).eval ctx = (e₁.eval ctx) + (e₂.eval ctx) :=
by simp [add_nn_add, Expr.eval]
@[simp] lemma Expr.eval_radd (e₁ e₂ : Expr rr) (ctx : EContext) :
  (e₁ + e₂).eval ctx = e₁.eval ctx + e₂.eval ctx :=
by simp [add_rr_add, Expr.eval]
@[simp] lemma Expr.frame_nadd (e₁ e₂ : Expr nn) : (e₁ + e₂).frame = e₁.frame ∪ e₂.frame :=
by { simp [add_nn_add], ext, simp [fin.exists_fin_two], }

@[simp] lemma Expr.eval_nmul (e₁ e₂ : Expr nn) (ctx : EContext) :
  (e₁ * e₂).eval ctx = (e₁.eval ctx) * (e₂.eval ctx) :=
by simp [mul_nn_mul, Expr.eval]
@[simp] lemma Expr.eval_rmul (e₁ e₂ : Expr rr) (ctx : EContext) :
  (e₁ * e₂).eval ctx = (e₁.eval ctx) * (e₂.eval ctx) :=
by simp [mul_rr_mul, Expr.eval]

@[simp] lemma Expr.eval_lt (e₁ e₂ : Expr nn) (ctx : EContext) :
  Expr.eval ctx (e₁ ⟪<⟫ e₂) = (e₁.eval ctx < e₂.eval ctx : bool) :=
by simp [(⟪<⟫), Expr.eval]
@[simp] lemma Expr.eval_le (e₁ e₂ : Expr nn) (ctx : EContext) :
  Expr.eval ctx (e₁ ⟪≤⟫ e₂) = (e₁.eval ctx ≤ e₂.eval ctx : bool) :=
by simp [(⟪≤⟫), Expr.eval]
@[simp] lemma Expr.eval_gt (e₁ e₂ : Expr nn) (ctx : EContext) :
  Expr.eval ctx (e₁ ⟪>⟫ e₂) = (e₂.eval ctx < e₁.eval ctx : bool) :=
by simp [(⟪>⟫), Expr.eval]
@[simp] lemma Expr.eval_ge (e₁ e₂ : Expr nn) (ctx : EContext) :
  Expr.eval ctx (e₁ ⟪≥⟫ e₂) = (e₂.eval ctx ≤ e₁.eval ctx : bool) :=
by { simp [(⟪≥⟫), Expr.eval] }

@[simp] lemma Expr.eval_eq (e₁ e₂ : Expr nn) (ctx : EContext) :
  Expr.eval ctx (e₁ ⟪=⟫ e₂) = (e₁.eval ctx = e₂.eval ctx : bool) :=
by simp [(⟪=⟫), Expr.eval]
@[simp] lemma Expr.eval_and (e₁ e₂ : Expr bb) (ctx : EContext) :
  (e₁ ⊓ e₂).eval ctx = e₁.eval ctx && e₂.eval ctx :=
by { simp [has_inf.inf, Expr.eval] }
@[simp] lemma Expr.eval_or  (e₁ e₂ : Expr bb) (ctx : EContext) :
  (e₁ ⊔ e₂).eval ctx = e₁.eval ctx || e₂.eval ctx :=
by { simp [has_sup.sup, Expr.eval] }

@[simp] lemma Expr.eval_not (e : Expr bb) (ctx : EContext) :
  e.not.eval ctx = !(e.eval ctx) :=
by { simp [Expr.not, Expr.eval] }

end Expr

parameter (R)
structure LoopBound :=
(frame : Frame)
(to_fun : EContext → ℕ)
(has_frame : true /- TODO: function.has_frame to_fun frame -/)

section LoopBound

instance : has_coe_to_fun LoopBound (λ _, EContext → ℕ) :=
⟨LoopBound.to_fun⟩
instance has_coe_from_nat : has_coe ℕ LoopBound := ⟨λ n, ⟨finset.empty, (λ _, n), true.intro⟩⟩

@[simp] lemma LoopBound.mk_apply (a f c x) : (LoopBound.mk a f c) x = f x := rfl

end LoopBound

parameter (R)
inductive Prog
| skip : Prog
| store {b : Types} (dst : Ident b) (val : Expr b)
| store_arr {b : Types} (dst : Ident b) (ind : Expr nn) (val : Expr b)
| seq (a : Prog) (b : Prog)
| branch (cond : Expr bb) (a : Prog) (b : Prog)
| loop (n : LoopBound) (cond : Expr bb) (b : Prog)

section Prog

parameter {R}
def prog_repr [has_to_string R] : Prog → list string
| Prog.skip := ["pass"]
| (Prog.store dst val) := [(to_string dst) ++ " := " ++ (to_string val)]
| (Prog.store_arr dst ind val) := [(to_string dst) ++ ("[" ++ to_string ind ++ "]") ++ " := " ++ (to_string val)]
| (Prog.seq a b) := (prog_repr a) ++ (prog_repr b)
| (Prog.branch c a b) := ["if " ++ (to_string c) ++ ":"]
    ++ (prog_repr a).map (λ s, "  " ++ s)
    ++ ["else:"]
    ++ (prog_repr b).map (λ s, "  " ++ s)
| (Prog.loop n cond b) := ["while " ++ (to_string cond) ++ ":"]
    ++ (prog_repr b).map (λ s, "  " ++ s)

instance [has_to_string R] : has_to_string Prog :=
⟨λ p, "\n".intercalate (prog_repr p)⟩

@[simp] def Prog.eval : Prog → EContext → EContext
| Prog.skip ctx := ctx
| (Prog.store dst val) ctx := ctx.update dst (val.eval ctx)
| (Prog.store_arr dst ind val) ctx :=
  let i : ℕ := ind.eval ctx in
  if i < (ctx.heap.get dst).length then ctx.update_arr dst i (val.eval ctx) else ctx
| (Prog.seq a b) ctx := b.eval (a.eval ctx)
| (Prog.branch condition a b) ctx := cond (condition.eval ctx) (a.eval ctx) (b.eval ctx)
| (Prog.loop n c b) ctx :=
(λ ctx, cond (c.eval ctx) (b.eval ctx) ctx)^[(n ctx)] ctx

@[simp] def Prog.frame : Prog → Frame
| Prog.skip := ∅
| (Prog.store dst val) := insert (sigma.mk _ dst) val.frame
| (Prog.store_arr dst ind val) := insert (sigma.mk _ dst) (ind.frame ∪ val.frame)
| (Prog.seq a b) := a.frame ∪ b.frame
| (Prog.branch c a b) := c.frame ∪ a.frame ∪ b.frame
| (Prog.loop n c b) := c.frame ∪ b.frame

end Prog

local infixr ` <;> `:1 := Prog.seq
local notation a ` ::= `:20 c := Prog.store a c
local notation a ` ⟬ `:9000 i ` ⟭ ` ` ::= `:20 c := Prog.store_arr a i c
local notation x ` ⟬ `:9000 i ` ⟭ ` := Expr.access x i

class TRAble (α : Type*) (β : out_param Type*) :=
(tr : EContext → α → β)

open TRAble (tr)

@[simps]
instance tr_expr_nn : TRAble (Expr nn) ℕ :=
{ tr := λ ctx e, e.eval ctx }

@[simps]
instance tr_expr_nn' : TRAble (Expr nn) (ExprVal nn) := tr_expr_nn

@[simps]
instance tr_expr_rr : TRAble (Expr rr) R :=
{ tr := λ ctx e, e.eval ctx }

section stream

parameter (R)
structure BoundedStreamGen (ι α : Type) :=
(current : ι)
(value : α)
(ready : Expr bb)
(next : Prog)
(valid : Expr bb)
(bound : LoopBound)
(initialize : Prog)

parameter {R}
variables {ι α ι' β : Type}

@[ext]
lemma BoundedStreamGen.ext {s₁ s₂ : BoundedStreamGen ι α} (h₁ : s₁.current = s₂.current)
  (h₂ : s₁.value = s₂.value) (h₃ : s₁.ready = s₂.ready) (h₄ : s₁.next = s₂.next) (h₅ : s₁.valid = s₂.valid)
  (h₆ : s₁.bound = s₂.bound) (h₇ : s₁.initialize = s₂.initialize) : s₁ = s₂ :=
by { cases s₁, cases s₂, dsimp only at *, subst_vars, }

section functorality

@[simps]
instance : bifunctor BoundedStreamGen :=
{ bimap := λ _ _ _ _ f g s, { s with current := f s.current, value := g s.value } }

instance : is_lawful_bifunctor BoundedStreamGen :=
{ id_bimap := by { intros, ext; simp, },
  bimap_bimap := by { intros, ext; simp, } }

end functorality

@[simps]
def BoundedStreamGen.to_stream_aux  [TRAble ι ι'] [TRAble α β] (s : BoundedStreamGen ι α) :
  Stream ι' β :=
{ σ := EContext,
  valid := λ ctx, s.valid.eval ctx,
  ready := λ ctx, s.valid.eval ctx && s.ready.eval ctx,
  next := λ ctx h, s.next.eval ctx,
  index := λ ctx h, tr ctx s.current,
  value := λ ctx h, tr ctx s.value }

@[simps]
def BoundedStreamGen.to_stream [TRAble ι ι'] [TRAble α β] (ctx₀ : EContext) (s : BoundedStreamGen ι α) : StreamExec ι' β :=
{ stream := s.to_stream_aux,
  bound := s.bound (s.initialize.eval ctx₀),
  state := s.initialize.eval ctx₀,
  bound_valid := sorry, }

instance eval_stream [TRAble ι ι'] [TRAble α β] : TRAble (BoundedStreamGen ι α) (StreamExec ι' β) :=
{ tr := BoundedStreamGen.to_stream }

section translate
open_locale classical

structure tr_to_stream [TRAble ι ι'] [TRAble α β] (s : BoundedStreamGen ι α)
  (t : Stream ι' β) (f : EContext → t.σ) (ctx : EContext) : Prop :=
(hvalid : s.valid.eval ctx ↔ t.valid (f ctx))
(hready' : s.valid.eval ctx → t.valid (f ctx) → (s.ready.eval ctx ↔ t.ready (f ctx)))
(hnext : s.valid.eval ctx → ∀ h, f (s.next.eval ctx) = t.next (f ctx) h)
(hcurr : s.valid.eval ctx → ∀ h, tr ctx s.current = t.index (f ctx) h)
(hval : s.valid.eval ctx → s.ready.eval ctx → ∀ h, tr ctx s.value = t.value (f ctx) h)

lemma tr_to_stream.hready [TRAble ι ι'] [TRAble α β] {s : BoundedStreamGen ι α}
  {t : Stream ι' β} {f : EContext → t.σ} {ctx : EContext} (h : tr_to_stream s t f ctx)
  (hv : t.valid (f ctx)) : s.ready.eval ctx ↔ t.ready (f ctx) := h.hready' (h.hvalid.mpr hv) hv

@[simp] def EContext.is_length {b : Types} (ctx : EContext) (arr : Ident b) (len : Ident nn) : Prop :=
(ctx.heap.get arr).length = ctx.store.get len

def preserves (next : Prog) (inv : EContext → Prop) : Prop :=
∀ {c}, inv c → inv (next.eval c)

@[mk_iff]
structure EContext.unmodified (inv : NameSpace) (c₀ : EContext) (ctx : EContext) : Prop :=
(h : ∀ {b : Types} (v : Vars), ctx.heap.get (inv∷v : Ident b) = c₀.heap.get inv∷v)
(s : ∀ {b : Types} (v : Vars), ctx.store.get (inv∷v : Ident b) = c₀.store.get inv∷v)

@[refl] lemma EContext.unmodified.rfl (inv : NameSpace) (c₀ : EContext) :
  c₀.unmodified inv c₀ :=
⟨λ _ v, rfl, λ _ v, rfl⟩

section preserves
variables {next : Prog} {ctx : EContext} {p₁ p₂ : EContext → Prop}
lemma preserves.and (h₀ : preserves next p₁) (h₁ : preserves next p₂) : (preserves next (λ c, p₁ c ∧ p₂ c)) :=
by { rw [preserves] at *, tauto, }

--#check
lemma preserves.unmodified (c₀ : EContext) {b} {inv : Ident b}
(h : {⟨_, inv⟩} # next.frame) :
  preserves next (c₀.unmodified inv) := sorry -- FOOTPRINT: nothing in `inv` is modified

lemma preserves.unmodified (c₀ : EContext) {inv : NameSpace} (h : inv ∉ next.frame.image (λ x : Σ b, Ident b, x.2.ns)) :
  preserves next (c₀.unmodified inv) := sorry -- FOOTPRINT: nothing in `inv` is modified

lemma preserves.is_length {b : Types} (v : Ident b) (e : Ident nn)  (h : (sigma.mk _ e) ∉ next.frame) :
  preserves next (λ c, c.is_length v e) := sorry -- FOOTPRINT: If `e` (length variable) is not modified, this is preserved

end preserves

structure tr_to [TRAble ι ι'] [TRAble α β] (s : BoundedStreamGen ι α)
  (t : StreamExec ι' β) (f : EContext → t.stream.σ) (ctx : EContext) : Type :=
(inv : EContext → Prop)
(to_stream : ∀ c, inv c → tr_to_stream s t.stream f c)
(hinit : f (s.initialize.eval ctx) = t.state)
(init_inv : inv (s.initialize.eval ctx))
(hbound : s.bound (s.initialize.eval ctx) = t.bound)
(preserves : preserves s.next inv)

variables [TRAble ι ι'] [TRAble α β] {s : BoundedStreamGen ι α}
  {t : Stream ι' β} {f : EContext → t.σ} {ctx : EContext}
  [add_zero_class β]

lemma tr_to_stream.eval₀ (h : tr_to_stream s t f ctx) (h₀ h₁) :
  s.to_stream_aux.eval₀ ctx h₀ =
  t.eval₀ (f ctx) h₁ :=
begin
  simp [Stream.eval₀, h.hready h₁, h.hvalid, h₁],
  split_ifs with h₂,
  { /- If both are ready -/ rw [h.hcurr, h.hval]; simpa [h.hready h₁, h.hvalid], },
  { /- If both are not ready -/ refl, },
end

lemma tr_to_stream.eval_steps_eq {inv : EContext → Prop}
  (hinv : ∀ {c}, inv c → inv (s.next.eval c)) (hc : inv ctx)
  (h : ∀ {c}, inv c → tr_to_stream s t f c) (n : ℕ) :
  s.to_stream_aux.eval_steps n ctx =
  t.eval_steps n (f ctx) :=
begin
  induction n with n ih generalizing ctx, { refl, },
  specialize h hc,
  simp [StreamExec.valid, h.hvalid],
  split_ifs with hv, swap, { refl, },
  congr' 1, swap,   { /- The first step is the same -/ rw tr_to_stream.eval₀ h, },
  /- The rest of the steps are the same -/
  rw [ih, h.hnext (h.hvalid.mpr hv)],
  exact hinv hc,
end

lemma tr_to.eval_finsupp_eq {t : StreamExec ι' β} {f : EContext → t.stream.σ} (h : tr_to s t f ctx) :
  StreamExec.eval (tr ctx s) = t.eval :=
by { dsimp [StreamExec.eval, BoundedStreamGen.to_stream, tr], simp [h.hbound, tr_to_stream.eval_steps_eq h.preserves h.init_inv h.to_stream, h.hinit], }

end translate

instance eval_unit : TRAble unit unit := ⟨λ _ _, ()⟩

def singleton (x : α) : BoundedStreamGen unit α := sorry

def range_nn (n : Expr nn) : BoundedStreamGen (Expr nn) (Expr nn) := sorry

def range_rr (n : Expr nn) : BoundedStreamGen (Expr nn) (Expr rr) := sorry

def externSparseVec (scratch : NameSpace) : BoundedStreamGen (Expr nn) (Expr rr) :=
let i : Ident nn := scratch∷Vars.i,
    len : Ident nn := reserved∷Vars.len,
    inds : Ident nn := reserved∷ind₀,
    vals : Ident rr := reserved∷vals in
{ current := inds⟬i⟭,
  value := vals⟬i⟭,
  ready := Expr.lit tt,
  next := i ::= i + 1,
  valid := (i : Expr nn) ⟪<⟫ len,
  bound := ⟨default, λ ctx, ctx.store.get len, /- TODO: Frame -/ trivial⟩,
  initialize := i ::= 0, }

def contract (x : BoundedStreamGen ι α) : BoundedStreamGen unit α :=
bifunctor.fst default x

@[simp] lemma contract_spec [TRAble ι ι'] [TRAble α β] (x : BoundedStreamGen ι α)
  (ctx : EContext) :
  tr ctx (contract x) = contract_stream (tr ctx x) := rfl

section sparse_vectors
open NameSpace (reserved) Vars (ind₀ vals len)

@[mk_iff]
structure externSparseVecCond (ctx : EContext) : Prop :=
(inds_len : (ctx.heap.get reserved∷ₙind₀).length = ctx.store.get reserved∷ₙlen)
(vals_len : (ctx.heap.get reserved∷ᵣvals).length = ctx.store.get reserved∷ₙlen)

lemma externSparseVec_tr_to_stream (scratch : NameSpace) (c : EContext) {l : ℕ} (is : vector ℕ l) (vs : vector R l)
  (hc₁ : c.heap.get reserved∷ₙind₀ = is.to_list) (hc₂ : c.heap.get reserved∷ᵣvals = vs.to_list) (hc₃ : c.store.get reserved∷ₙlen = l) :
  tr_to_stream (externSparseVec scratch) (primitives.externSparseVec_stream is vs)
    (λ ctx, ctx.store.get scratch∷ₙVars.i) c :=
{ hvalid := by simp [externSparseVec, primitives.externSparseVec_stream, hc₁, hc₂, hc₃],
  hready' := by simp [externSparseVec, primitives.externSparseVec_stream, hc₁, hc₂, hc₃],
  hnext := by { intros, simp [externSparseVec, primitives.externSparseVec_stream, hc₁, hc₂, hc₃], },
  hcurr := by { simp [externSparseVec, primitives.externSparseVec_stream, hc₁, hc₂, hc₃, vector.nth_eq_nth_le], intros, rw list.nth_le_nth, },
  hval := by { simp [externSparseVec, primitives.externSparseVec_stream, hc₁, hc₂, hc₃, vector.nth_eq_nth_le], intros, rw list.nth_le_nth, } }

def externSparseVec_tr (scratch : NameSpace) (hs : reserved ≠ scratch) (c : EContext)
  (hc : externSparseVecCond c) :
  tr_to (externSparseVec scratch) (primitives.externSparseVec ⟨c.heap.get reserved∷ₙind₀, hc.inds_len⟩ ⟨c.heap.get reserved∷ᵣvals, hc.vals_len⟩)
    (λ ctx, ctx.store.get scratch∷ₙVars.i) c :=
{ inv := c.unmodified reserved,
  to_stream := λ c' hc', by apply externSparseVec_tr_to_stream scratch c'; simp [hc'.h, hc'.s],
  hinit := by simp [primitives.externSparseVec, externSparseVec],
  init_inv := by { apply preserves.unmodified, { simpa [externSparseVec], }, refl, },
  hbound := by simp [primitives.externSparseVec, externSparseVec, hc],
  preserves := by { apply preserves.unmodified, simpa [externSparseVec], } }

open_locale big_operators

@[simp] lemma externSparseVec_spec (scratch : NameSpace) (hs : reserved ≠ scratch) (c : EContext) (hc : externSparseVecCond c) :
  StreamExec.eval (tr c (externSparseVec scratch)) = ∑ i : fin (c.store.get reserved∷ₙlen), finsupp.single ((c.heap.get reserved∷ₙind₀).nth_le i (by rw hc.1; exact i.prop)) ((c.heap.get reserved∷ᵣvals).nth_le i (by rw hc.2; exact i.prop)) :=
by { simp [(externSparseVec_tr scratch hs c hc).eval_finsupp_eq, vector.nth_eq_nth_le], }

end sparse_vectors

def BoundedStreamGen.body (x : BoundedStreamGen unit (Expr rr)) : Prog :=
Prog.branch x.ready
  (reserved∷ᵣoutput ::= reserved∷ᵣoutput + x.value)
/- else -/ Prog.skip <;>
x.next

def compile_scalar (x : BoundedStreamGen unit (Expr rr)) : Prog :=
let out : Ident rr := reserved∷output in
out ::= 0 <;>
x.initialize <;>
Prog.loop x.bound x.valid x.body

section compile_sound

lemma eval_body (x : BoundedStreamGen unit (Expr rr)) (c c' : EContext)
  (hc : c.heap = c'.heap ∧ ∀ v, v ≠ reserved∷ᵣoutput → c.store.get v = c'.store.get v)
  (h : x.to_stream_aux.valid c') (h' : x.valid.eval c') :
  (x.body.eval c).store.get reserved∷ᵣoutput = (x.to_stream_aux.eval₀ c' h ()) + (c.store.get reserved∷ᵣoutput) :=
begin
  have F₁ : x.ready.eval c' = x.ready.eval c := sorry, -- FOOTPRINT: `out ∉ ready.footprint`
  have F₂ : ∀ ctx, (x.next.eval ctx).store.get reserved∷ᵣoutput = ctx.store.get reserved∷ᵣoutput := sorry, -- FOOTPRINT: out ∉ next.footprint
  have F₃ : x.value.eval c' = x.value.eval c := sorry, -- FOOTPRINT: `out ∉ value.footprint`
  simp [BoundedStreamGen.body, Stream.eval₀, h', F₁],
  cases H : x.ready.eval c; simp [H, F₂, F₃, punit_eq_star (tr _ _), add_comm],
end

lemma iterate_body (x : BoundedStreamGen unit (Expr rr)) (c c' : EContext)
  (hc : c.heap = c'.heap ∧ ∀ v, v ≠ reserved∷ᵣoutput → c.store.get v = c'.store.get v)
  (n : ℕ) :
  ((λ ctx, cond (x.valid.eval ctx) (x.body.eval ctx) ctx)^[n] c).store.get reserved∷ᵣoutput =
    (Stream.eval_steps x.to_stream_aux n c' ()) + (c.store.get reserved∷ᵣoutput) :=
begin
  induction n with n ih generalizing c c', { simp, },
  simp,
  have F₁ : x.valid.eval c' = x.valid.eval c := sorry, -- FOOTPRINT: `out ∉ valid.footprint`
  cases H : x.valid.eval c,
  { /- Invalid: both are `0` -/ simp [H, F₁], rw function.iterate_fixed, simp [H], },
  simp [H, F₁],
  rw ih _ (x.next.eval c'), swap,
  { simp [BoundedStreamGen.body], sorry, /- FOOTPRINT: `next` preserves all variables besides `out` -/},
  rw [add_assoc, eval_body x c c' hc],
  simp [F₁, H],
end

lemma compile_scalar_sound (x : BoundedStreamGen unit (Expr rr)) (ctx : EContext) :
  ((compile_scalar x).eval ctx).store.get (reserved∷output : Ident rr) = StreamExec.eval (tr ctx x) () :=
begin
  simp [compile_scalar],
  set ctx' : EContext := ctx.update reserved∷ᵣoutput 0,
  simp [tr, StreamExec.eval],
  rw iterate_body x (x.initialize.eval ctx') (x.initialize.eval ctx),
  have F₁ : (x.initialize.eval ctx').store.get reserved∷ᵣoutput = ctx'.store.get reserved∷ᵣoutput := sorry, -- FOOTPRINT:
  have F₂ : x.bound (x.initialize.eval ctx') = x.bound (x.initialize.eval ctx) := sorry, -- FOOTPRINT:
  { dsimp [BoundedStreamGen.to_stream], simp [F₁, F₂], },
  sorry, /- FOOTPRINT: since `ctx` and `ctx'` only differ in `out`, `init ctx` and `init ctx'` can only differ in `out`
            (since `init` does not read/write from `out`) -/
end

end compile_sound

end stream

end compiler

section examples
open TRAble (tr)

parameters [add_comm_monoid R] [has_one R] [has_mul R]

open_locale big_operators

def sum_vec (scratch : NameSpace) : BoundedStreamGen unit (Expr rr) :=
contract (externSparseVec scratch)

@[simp] lemma sum_vec_spec (scratch : NameSpace) (hs : reserved ≠ scratch) (ctx : EContext) (hctx : externSparseVecCond ctx) :
  StreamExec.eval (tr ctx (sum_vec scratch)) = finsupp.single () (ctx.heap.get reserved∷ᵣvals).sum :=
begin
  simp [sum_vec, *],
  rw [map_sum],
  simp [finset.sum, multiset.map_nth_le hctx.2],
end

lemma sum_vec_compile_spec (scratch : NameSpace) (hs : reserved ≠ scratch) (ctx : EContext) (hctx : externSparseVecCond ctx) :
  ((compile_scalar (sum_vec scratch)).eval ctx).store.get reserved∷ᵣoutput = (ctx.heap.get reserved∷ᵣvals).sum :=
by { rw [compile_scalar_sound, sum_vec_spec _ hs _ hctx], simp, }

end examples

-- Final theorem will be something like:
-- ∀ (x : BoundedStreamGen ι α) [TRAble ι → ι'] [TRAble α → β] [FinsuppEval (StreamExec EContext ι' β)]
--  (hind₁ : ι compiles correctly) (hind₂ : α compiles correctly) : BoundedStreamGen ι α compiles correctly

end

section examples
open Types

notation ` Σ_c ` := contract
@[derive [add_comm_monoid, has_one, has_mul, has_to_string], irreducible]
def R := ℤ
abbreviation compile := @compile_scalar R

def sum_vec' : BoundedStreamGen R unit (Expr R rr) :=
Σ_c (externSparseVec (fresh ∅))

#eval do io.print_ln (compile sum_vec')

end examples


import data.list.big_operators
import tactic.expand_exists
import complexity_class.tactic
import encode

open tencodable function tree
variables {α β γ δ ε : Type*}

class polysize (α : Type*) [tencodable α] :=
(size : α → ℕ)
(upper [] : ∃ p : polynomial ℕ, ∀ x, size x ≤ p.eval (encode x).num_nodes)
(lower [] : ∃ p : polynomial ℕ, ∀ x, (encode x).num_nodes ≤ p.eval (size x))

open polysize
variables [tencodable α] [tencodable β] [tencodable γ] [tencodable δ] [tencodable ε]

@[instance, priority 10]
def default_polysize : polysize α :=
{ size := λ x, (encode x).num_nodes,
  upper := ⟨polynomial.X, by simp⟩,
  lower := ⟨polynomial.X, by simp⟩ }

@[simps]
instance fintype.polysize [fintype α] : polysize α :=
{ size := λ _, 0,
  upper := ⟨0, λ x, zero_le'⟩,
  lower := ⟨polynomial.C ((@finset.univ α _).sup (λ x, (encode x).num_nodes)),
    λ x, by { simp, exact finset.le_sup (finset.mem_univ x), }⟩ }

lemma list.encode_num_nodes_eq : ∀ (x : list α),
  (encode x).num_nodes = x.length + (x.map $ λ e, (encode e).num_nodes).sum
| [] := rfl
| (hd :: tl) := by { simp [encode_cons, tl.encode_num_nodes_eq], abel, }

lemma list.len_le_encode (x : list α) : x.length ≤ (encode x).num_nodes :=
by { rw x.encode_num_nodes_eq, exact le_self_add, }

lemma list.encode_lt_encode_of_mem {x : list α} {y : α} (h : y ∈ x) :
  (encode y).num_nodes < (encode x).num_nodes :=
begin
  rw [x.encode_num_nodes_eq, ← multiset.coe_sum],
  refine lt_add_of_pos_of_le (list.length_pos_of_mem h) (multiset.le_sum_of_mem _),
  simp, exact ⟨_, h, rfl⟩,
end

@[simp] lemma encode_sum_inl_num_nodes (x : α) :
  (encode (sum.inl x : α ⊕ β)).num_nodes = (encode x).num_nodes + 1 := by simp [encode]

@[simp] lemma encode_sum_inr_num_nodes (x : β) :
  (encode (sum.inr x : α ⊕ β)).num_nodes = (encode x).num_nodes + 2 :=
by { simp [encode], ring, }

variables [polysize α] [polysize β] [polysize γ] [polysize δ] [polysize ε]

instance : polysize (α × β) :=
{ size := λ x : α × β, size x.1 + size x.2,
  upper := begin
    cases upper α with p hp, cases upper β with q hq,
    use p + q, rintro ⟨x₁, x₂⟩,
    rw polynomial.eval_add,
    refine add_le_add ((hp _).trans $ p.eval_mono _) ((hq _).trans $ q.eval_mono _);
    simp [encode]; linarith only,
  end,
  lower := begin
    cases lower α with p hp, cases lower β with q hq,
    use p + q + 1, rintro ⟨x₁, x₂⟩,
    simp only [encode, num_nodes, polynomial.eval_add, polynomial.eval_one, add_le_add_iff_right],
    exact add_le_add ((hp _).trans $ p.eval_mono le_self_add)
      ((hq _).trans $ q.eval_mono le_add_self),
  end }

instance : polysize (option α) :=
{ size := λ x, x.elim 0 size,
  upper := begin
    cases upper α with p hp, use p,
    rintro (_|x), { simp, },
    simp only [encode, option.elim, num_nodes, zero_add],
    refine (hp _).trans (p.eval_mono _), simp,
  end,
  lower := begin
    cases lower α with p hp, use p + 1,
    rintro (_|x), { simp [encode], },
    simpa [encode] using hp _,
  end }

@[simp] lemma polysize.none_size : size (@none α) = 0 := rfl
@[simp] lemma polysize.some_size (x : α) : size (some x) = size x := rfl
@[simp] lemma polysize.prod_size (x : α) (y : β) : size (x, y) = size x + size y := rfl

instance : polysize (list α) :=
{ size := λ l, l.length + (l.map size).sum,
  upper := begin
    cases upper α with p hp,
    use polynomial.X + polynomial.X * p,
    intro x, simp,
    refine add_le_add x.len_le_encode ((list.sum_le_card_nsmul _ (p.eval (encode x).num_nodes) _).trans _),
    { simpa using λ a ha, (hp a).trans (p.eval_mono (list.encode_lt_encode_of_mem ha).le), },
    simpa using nat.mul_le_mul_right _ x.len_le_encode,
  end,
  lower := begin
    cases lower α with p hp,
    use polynomial.X + polynomial.X * p,
    intro x, simp [x.encode_num_nodes_eq, add_assoc],
    refine le_add_left ((list.sum_le_card_nsmul _ (p.eval (x.map size).sum) _).trans _),
    { simp, intros a ha, refine (hp _).trans (p.eval_mono _),
      rw ← multiset.coe_sum, apply multiset.le_sum_of_mem,
      simp, exact ⟨_, ha, rfl⟩, },
    simpa using nat.mul_le_mul le_self_add (p.eval_mono le_add_self),
  end }

/-- The same as `list.encode_lt_encode_of_mem` but for `size` -/
lemma list.size_lt_of_mem {x : list α} {e : α} (h : e ∈ x) :
  size e < size x :=
begin
  dsimp only [polysize.size],
  rw [← multiset.coe_sum],
  refine lt_add_of_pos_of_le (list.length_pos_of_mem h) (multiset.le_sum_of_mem _),
  simp, exact ⟨_, h, rfl⟩,
end

lemma list.size_le_of_sublist {x y : list α} (h : y <+ x) :
  size y ≤ size x :=
add_le_add h.length_le ((h.map size).sum_le_sum $ λ _ _, zero_le')

lemma list.length_le_size (l : list α) : l.length ≤ size l := le_self_add

@[simp] lemma size_nil : size ([] : list α) = 0 := rfl

@[simp] lemma size_cons (x : α) (xs : list α) : size (x :: xs) = size x + size xs + 1 :=
by { simp [size], abel, }

@[simp] lemma size_append (x y : list α) : size (x ++ y) = size x + size y :=
by { simp [size], abel, }

@[simp] lemma size_reverse (x : list α) : size x.reverse = size x :=
by simp [size, list.sum_reverse]

@[simp] lemma size_list_fintype {α : Type*} [tencodable α] [fintype α] (x : list α) :
  size x = x.length := by simp [size]

lemma list.size_le_mul_of_le (a b : ℕ) (l : list α)
  (h₁ : l.length ≤ a) (h₂ : ∀ x ∈ l, size x ≤ b) :
  size l ≤ a * (b + 1) :=
begin
  simp only [size, add_comm b 1, mul_add, mul_one],
  refine add_le_add h₁ ((list.sum_le_card_nsmul _ b _).trans _),
  { simpa using h₂, }, { simpa using mul_le_mul_right' h₁ b, }
end

lemma list.perm.size_eq {l₁ l₂ : list α} (h : l₁ ~ l₂) : size l₁ = size l₂ :=
by simp only [size, h.length_eq, (h.map size).sum_eq]

instance : polysize (α ⊕ β) :=
{ size := λ x, x.elim size size,
  upper := begin
    cases upper α with p hp, cases upper β with q hq,
    use p + q,
    rintros (x|x); simp,
    exacts [le_add_right ((hp x).trans (p.eval_mono le_self_add)),
            le_add_left ((hq x).trans (q.eval_mono le_self_add))],
  end,
  lower := begin
    cases lower α with p hp, cases lower β with q hq,
    use p + q + 2,
    rintros (x|x); simp,
    exacts [add_le_add (le_add_right (hp x)) one_le_two,
            le_add_left (hq x)],
  end }

@[simp] lemma size_inl (x : α) : size (sum.inl x : α ⊕ β) = size x := rfl
@[simp] lemma size_inr (x : β) : size (sum.inr x : α ⊕ β) = size x := rfl

instance {n : ℕ} : polysize (vector α n) :=
{ size := λ v, (v.map size).to_list.sum,
  upper := begin
    obtain ⟨p, hp⟩ := polysize.upper (list α),
    refine ⟨p, λ x, trans _ (hp x.to_list)⟩,
    simp only [size, vector.to_list_map],
    exact le_add_self,
  end,
  lower := begin
    obtain ⟨p, hp⟩ := polysize.lower (list α),
    refine ⟨p.comp (polynomial.C n + polynomial.X), λ x, (hp x.to_list).trans _⟩,
    simp [size],
  end }

lemma polysize_vector_def {n} (v : vector α n) : size v = (v.map size).to_list.sum := rfl

lemma _root_.vector.polysize_tail_le_self {n : ℕ} (v : vector α (n + 1)) :
  polysize.size v.tail ≤ polysize.size v :=
by { rcases v.exists_eq_cons with ⟨hd, tl, rfl⟩, simp [polysize_vector_def], }

-- Equal to `default_polysize` but more useful defeq
instance : polysize ℕ :=
{ size := λ n, n,
  upper := ⟨polynomial.X, by simp⟩,
  lower := ⟨polynomial.X, by simp⟩ }

instance : polysize (tree unit) := default_polysize

@[simp] lemma polysize_unary_nat (n : ℕ) : size n = n := rfl

@[simp] lemma polysize_tree_unit (x : tree unit) : size x = x.num_nodes := rfl

def polysize_fun {γ : Type*} [has_uncurry γ α β] (f : γ) : Prop :=
∃ (p : polynomial ℕ), ∀ x : α, size (↿f x) ≤ p.eval (size x)

def polysize_safe (f : α → β → γ) : Prop :=
∃ (p : polynomial ℕ), ∀ x y, size (f x y) ≤ size y + p.eval (size x)

@[expand_exists polysize_fun.poly polysize_fun.spec]
lemma polysize_fun.def {γ : Type} [has_uncurry γ α β] {f : γ} (hf : polysize_fun f) :
  ∃ (p : polynomial ℕ), ∀ x : α, size (↿f x) ≤ p.eval (size x) := hf

theorem polysize_fun.id : polysize_fun (@id α) := ⟨polynomial.X, by simp [has_uncurry.uncurry]⟩

theorem polysize_fun.comp {f : α → β} {g : γ → α} : polysize_fun f → polysize_fun g → polysize_fun (f ∘ g)
| ⟨p₁, h₁⟩ ⟨p₂, h₂⟩ := ⟨p₁.comp p₂, (λ x, by { rw polynomial.eval_comp, exact (h₁ (g x)).trans (p₁.eval_mono (h₂ x)), })⟩

theorem polysize_fun.head' : polysize_fun (@list.head' α) :=
⟨polynomial.X, λ x, by { cases x, { simp [has_uncurry.uncurry], }, simp [has_uncurry.uncurry, add_assoc], }⟩

theorem polysize_fun.tail : polysize_fun (@list.tail α) :=
⟨polynomial.X, λ x, by { cases x, { simp [has_uncurry.uncurry], }, simp [has_uncurry.uncurry], linarith only, }⟩

theorem polysize_fun.fst : polysize_fun (@prod.fst α β) :=
⟨polynomial.X, λ ⟨x, y⟩, by { simp [has_uncurry.uncurry], }⟩

theorem polysize_fun.snd : polysize_fun (@prod.snd α β) :=
⟨polynomial.X, λ ⟨x, y⟩, by { simp [has_uncurry.uncurry], }⟩

theorem polysize_fun.ite {P : α → Prop} [decidable_pred P] {f₁ f₂ : α → β} :
  polysize_fun f₁ → polysize_fun f₂ → polysize_fun (λ x, if P x then f₁ x else f₂ x)
| ⟨p₁, h₁⟩ ⟨p₂, h₂⟩ := ⟨p₁ + p₂, λ x, begin
  dsimp only [has_uncurry.uncurry, id],
  split_ifs, { refine (h₁ x).trans _, simp, }, { refine (h₂ x).trans _, simp, }
end⟩

@[expand_exists polysize_safe.poly polysize_safe.spec]
lemma polysize_safe.def {f : α → β → γ} (hf : polysize_safe f) :
  ∃ (p : polynomial ℕ), ∀ x y, size (f x y) ≤ size y + p.eval (size x) := hf

theorem polysize_safe.size_le {f : α → β → β} (hf : polysize_safe f) (n : ℕ) (x : α) (y : β) :
  size ((f x)^[n] y) ≤ size y + n * (hf.poly.eval $ size x) :=
begin
  induction n with n ih generalizing y, { simp, },
  rw [iterate_succ_apply, nat.succ_mul, nat.add_comm _  (hf.poly.eval (size x)), ← add_assoc],
  refine (ih _).trans _,
  simpa using hf.spec x y,
end

theorem polysize_safe.comp {f : γ → δ → ε} {g : α → γ} {h : α → β → δ} :
  polysize_safe f → polysize_fun g → polysize_safe h → polysize_safe (λ x y, f (g x) (h x y))
| ⟨pf, hf⟩ ⟨pg, hg⟩ ⟨ph, hh⟩ :=
begin
  use ph + (pf.comp pg),
  intros x y,
  refine (hf _ _).trans _,
  simp [← add_assoc], mono*,
end

@[complexity] theorem polysize_safe.id : polysize_safe (λ _ : α, @id β) :=
⟨0, by simp⟩

@[complexity] theorem polysize_safe.id' : polysize_safe (λ (_ : α) (y : β), y) := polysize_safe.id

@[complexity] theorem polysize_safe.of_polysize_fun {f : α → γ} :
  polysize_fun f → polysize_safe (λ x (_ : β), f x)
| ⟨p, hp⟩ := ⟨p, λ x y, le_add_left (hp x)⟩

theorem polysize_safe.comp' {f : γ → δ} {g : α → β → γ}
  (hf : polysize_safe (λ (_ : unit) x, f x)) (hg : polysize_safe g) :
  polysize_safe (λ x y, f (g x y)) :=
hf.comp (show polysize_fun (default : α → unit), from ⟨0, by simp⟩) hg

@[complexity] theorem polysize_safe.cons {f : α → γ} {g : α → β → list γ} (hf : polysize_fun f) (hg : polysize_safe g) :
  polysize_safe (λ (x : α) (y : β), (f x) :: (g x y)) :=
by { apply polysize_safe.comp _ hf hg, use polynomial.X + 1, intros, apply le_of_eq, simp, abel, }

@[complexity] theorem polysize_safe.ordered_insert (r : α → γ → γ → Prop) [∀ x, decidable_rel (r x)] {f : α → γ} {g : α → β → list γ} (hf : polysize_fun f) (hg : polysize_safe g) :
  polysize_safe (λ x y, (g x y).ordered_insert (r x) (f x)) :=
by { cases hf with pf hf, cases hg with pg hg, use pg + pf + 1, intros x y, simp [← add_assoc, ((g x y).perm_ordered_insert (r x) (f x)).size_eq], rw add_comm, mono*, }

@[complexity] theorem polysize_safe.const (C : γ) : polysize_safe (λ (_ : α) (_ : β), C) :=
⟨polynomial.C (size C), λ x y, by simp⟩

@[complexity] theorem polysize_nat_encode {f : α → β → ℕ} :
  polysize_safe f → polysize_safe (λ x y, encode (f x y))
| ⟨p, hp⟩ := ⟨p, by simpa using hp⟩

@[complexity] theorem polysize_num_nodes {f : α → β → tree unit} :
  polysize_safe f → polysize_safe (λ x y, (f x y).num_nodes)
| ⟨p, hp⟩ := ⟨p, by simpa using hp⟩

@[complexity] theorem polysize_safe.fst {f : α → β → γ × δ} (hf : polysize_safe f) :
  polysize_safe (λ x y, (f x y).1) :=
by { refine polysize_safe.comp' _ hf, use 0, simp, }

@[complexity] theorem polysize_safe.left {f : α → β → tree unit} (hf : polysize_safe f) :
  polysize_safe (λ x y, (f x y).left) :=
by { refine polysize_safe.comp' _ hf, use 0, simpa using tree.left_num_nodes_le, }

@[complexity] theorem polysize_safe.right {f : α → β → tree unit} (hf : polysize_safe f) :
  polysize_safe (λ x y, (f x y).right) :=
by { refine polysize_safe.comp' _ hf, use 0, simpa using tree.right_num_nodes_le, }

@[complexity] theorem polysize_safe.snd {f : α → β → γ × δ} (hf : polysize_safe f) :
  polysize_safe (λ x y, (f x y).2) :=
by { refine polysize_safe.comp' _ hf, use 0, simp, }

@[complexity] theorem polysize_safe.some {f : α → β → γ} (hf : polysize_safe f) :
  polysize_safe (λ x y, some (f x y)) :=
by { refine polysize_safe.comp' _ hf, use 0, simp, }

@[complexity] theorem polysize_safe.tail {f : α → β → list γ} (hf : polysize_safe f) :
  polysize_safe (λ x y, (f x y).tail) :=
by { refine polysize_safe.comp' _ hf, use 0, rintros ⟨⟩ (_|⟨hd, tl⟩); simp, exact le_add_right le_add_self, }

@[complexity] theorem polysize_safe.head' {f : α → β → list γ} (hf : polysize_safe f) :
  polysize_safe (λ x y, (f x y).head') :=
by { refine polysize_safe.comp' _ hf, use 0, rintros ⟨⟩ (_|⟨hd, tl⟩); simp [add_assoc], }

@[complexity] theorem polysize_safe.head [inhabited γ] {f : α → β → list γ} (hf : polysize_safe f) :
  polysize_safe (λ x y, (f x y).head) :=
by { refine polysize_safe.comp' _ hf, use polynomial.C (size (default : γ)), rintros ⟨⟩ (_|⟨hd, tl⟩); simp [add_assoc], }

@[complexity] theorem polysize_safe.list_split [inhabited γ] {f : α → β → list γ} : polysize_safe f →
  polysize_safe (λ x y, ((f x y).head, (f x y).tail))
| ⟨p, hp⟩ := ⟨p + size (default : γ), λ x y, begin
  specialize hp x y,
  cases H : f x y, { simp [H, ← add_assoc], },
  simp [H] at hp ⊢, linarith only [hp],
end⟩

@[complexity] theorem polysize_safe.append_left {f : α → list γ} {g : α → β → list γ} :
  polysize_fun f → polysize_safe g → polysize_safe (λ x y, (f x) ++ (g x y))
| ⟨pf, hf⟩ ⟨pg, hg⟩ := ⟨pf + pg, λ x y, by { dsimp [has_uncurry.uncurry] at hf, simp, linarith only [hf x, hg x y], }⟩

@[complexity] theorem polysize_safe.append_right {f : α → list γ} {g : α → β → list γ} :
  polysize_safe g → polysize_fun f → polysize_safe (λ x y, (g x y) ++ (f x))
|  ⟨pg, hg⟩ ⟨pf, hf⟩ := ⟨pf + pg, λ x y, by { dsimp [has_uncurry.uncurry] at hf, simp, linarith only [hf x, hg x y], }⟩

@[complexity] theorem polysize_safe.pair_left {f : α → γ} {g : α → β → δ} :
  polysize_fun f → polysize_safe g → polysize_safe (λ x y, (f x, g x y))
| ⟨pf, hf⟩ ⟨pg, hg⟩ := ⟨pf + pg, λ x y, by { dsimp [has_uncurry.uncurry] at hf, simp, linarith only [hf x, hg x y], }⟩

@[complexity] theorem polysize_safe.pair_right {f : α → γ} {g : α → β → δ} :
  polysize_safe g → polysize_fun f → polysize_safe (λ x y, (g x y, f x))
|  ⟨pg, hg⟩ ⟨pf, hf⟩ := ⟨pf + pg, λ x y, by { dsimp [has_uncurry.uncurry] at hf, simp, linarith only [hf x, hg x y], }⟩

@[complexity] theorem polysize_safe.add_unary_left {f : α → ℕ} {g : α → β → ℕ} :
  polysize_fun f → polysize_safe g → polysize_safe (λ x y, (f x) + (g x y))
| ⟨pf, hf⟩ ⟨pg, hg⟩ := ⟨pf + pg, λ x y, by { dsimp [has_uncurry.uncurry, polysize_unary_nat] at hf hg, simp, linarith only [hf x, hg x y], }⟩

@[complexity] theorem polysize_safe.add_unary_right {f : α → ℕ} {g : α → β → ℕ} :
  polysize_safe g → polysize_fun f → polysize_safe (λ x y, (g x y) + (f x))
| ⟨pg, hg⟩ ⟨pf, hf⟩ := ⟨pf + pg, λ x y, by { dsimp [has_uncurry.uncurry, polysize_unary_nat] at hf hg, simp, linarith only [hf x, hg x y], }⟩

@[complexity] theorem polysize_safe.ite {f g : α → β → γ} {P : α → β → Prop} [∀ x y, decidable (P x y)] (hf : polysize_safe f) (hg : polysize_safe g) :
  polysize_safe (λ x y, if P x y then f x y else g x y) :=
begin
  rcases hf with ⟨pf, hf⟩, rcases hg with ⟨pg, hg⟩, use pf + pg,
  intros x y, dsimp only, split_ifs,
  { refine (hf _ _).trans _, simp, }, { refine (hg _ _).trans _, simp, },
end

theorem foldl_size_le [polysize α] [polysize β] (f : β → α → β) (p : ℕ → ℕ) (hp : monotone p)
  (hf : ∀ x y, size (f x y) ≤ size x + p (size y)) (ls : list α) (x₀ : β) :
  size (ls.foldl f x₀) ≤ size x₀ + ls.length * p (size ls) :=
begin
  induction ls with hd tl ih generalizing x₀, { simp, },
  rw [list.foldl, list.length_cons, nat.succ_mul, ← add_assoc, add_right_comm],
  refine (ih _).trans (add_le_add ((hf _ _).trans $ add_le_add_left (hp _) _) (mul_le_mul_left' (hp _) _));
  simp; linarith only,
end

lemma polysize_safe.foldl [polysize α] [polysize β] [polysize γ] {lst : γ → list α} {acc : γ → β} {f : γ → β → α → β} :
  polysize_fun lst → polysize_fun acc → polysize_safe (λ (usf : γ × α) (sf : β), f usf.1 sf usf.2) →
  polysize_fun (λ x, (lst x).foldl (f x) (acc x))
| ⟨plst, hlst⟩ ⟨pacc, hacc⟩ ⟨pf, hf⟩ := ⟨pacc + plst * pf.comp (polynomial.X + plst), λ x, begin
  refine (foldl_size_le (f x) (λ a, pf.eval (size (x, a))) (λ a b h, pf.eval_mono $ add_le_add_left h _) (λ b a, hf (x, a) b) _ _).trans _,
  simp, mono*,
  exacts [(lst x).length_le_size.trans (hlst _), zero_le', zero_le'],
end⟩

@[complexity] theorem polysize_safe.option_bind₁ {f : α → option γ} {g : α → β → γ → option δ} :
  polysize_fun f → polysize_safe (λ (usf : α × γ) (sf : β), g usf.1 sf usf.2) →
  polysize_safe (λ x y, (f x).bind (g x y))
| ⟨pf, hf⟩ ⟨pg, hg⟩ := ⟨pg.comp (polynomial.X + pf), λ x y, begin
  cases H : f x with v, { simp [H], },
  specialize hf x, simp only [H, has_uncurry.uncurry, id] at hf ⊢,
  refine (hg (x, v) y).trans _,
  simp, mono*,
end⟩

@[complexity] theorem polysize_safe.option_bind₂ {f : α → β → option γ} {g : α → γ → option δ} :
  polysize_safe f → polysize_safe g → polysize_safe (λ x y, (f x y).bind (g x))
| ⟨pf, hf⟩ ⟨pg, hg⟩ := ⟨pf + pg, λ x y, begin
  cases H : f x y with v, { simp [H], },
  specialize hf x y,
  simp only [H, polysize.some_size, option.bind_some, polynomial.eval_add, ← add_assoc] at hf ⊢,
  exact (hg x v).trans (add_le_add_right hf _),
end⟩

@[complexity] theorem polysize_safe.option_map₁ {f : α → option γ} {g : α → β → γ → δ}
  (hf : polysize_fun f) (hg : polysize_safe (λ (usf : α × γ) (sf : β), g usf.1 sf usf.2)) :
  polysize_safe (λ x y, (f x).map (g x y)) :=
by { apply polysize_safe.option_bind₁ hf, exact polysize_safe.some hg, }

@[complexity] theorem polysize_safe.option_map₂ {f : α → β → option γ} {g : α → γ → δ}
  (hf : polysize_safe f) (hg : polysize_safe g) : polysize_safe (λ x y, (f x y).map (g x)) :=
by { apply polysize_safe.option_bind₂ hf, exact polysize_safe.some hg, }

@[complexity] theorem polysize_safe.get_or_else {f : α → β → option γ} {g : α → β → γ} :
  polysize_safe f → polysize_safe g → polysize_safe (λ x y, (f x y).get_or_else (g x y))
| ⟨pf, hpf⟩ ⟨pg, hpg⟩ := ⟨pf + pg, λ x y, by { specialize hpf x y, cases H : f x y, { simp [H], linarith only [hpg x y], }, simp [H] at hpf ⊢, linarith only [hpf], }⟩

section comp
variables {α₀ α₁ α₂ α₃ α₄ α₅ α₆ : Type*}
  [tencodable α₀] [tencodable α₁] [tencodable α₂] [tencodable α₃] [tencodable α₄] [tencodable α₅] [tencodable α₆]
  [polysize α₀]   [polysize α₁]   [polysize α₂]   [polysize α₃]   [polysize α₄]   [polysize α₅]   [polysize α₆]

-- Convention: compₙ_i₁... means composition of `n`-ary function where i₁, i₂, are safe indices
-- TODO automate

theorem polysize_safe.comp₃_1 {f : α₀ → α₁ → α₂ → γ}
  {g₀ : α → α₀} {g₁ : α → β → α₁} {g₂ : α → α₂} :
  polysize_safe (λ (usf : α₀ × α₂) (sf : α₁), f usf.1 sf usf.2) →
  polysize_fun g₀ → polysize_safe g₁ → polysize_fun g₂ →
  polysize_safe (λ x y, f (g₀ x) (g₁ x y) (g₂ x))
| ⟨pf, hf⟩ ⟨p₀, h₀⟩ ⟨p₁, h₁⟩ ⟨p₂, h₂⟩ := ⟨p₁ + pf.comp (p₀ + p₂),
  λ x y, by { refine (hf (g₀ x, g₂ x) (g₁ x y)).trans _, simp [← add_assoc], mono*, }⟩

theorem polysize_safe.comp₃_2 {f : α₀ → α₁ → α₂ → γ}
  {g₀ : α → α₀} {g₁ : α → α₁} {g₂ : α → β → α₂} :
  polysize_safe (λ (usf : α₀ × α₁) (sf : α₂), f usf.1 usf.2 sf) →
  polysize_fun g₀ → polysize_fun g₁ → polysize_safe g₂ →
  polysize_safe (λ x y, f (g₀ x) (g₁ x) (g₂ x y))
| ⟨pf, hf⟩ ⟨p₀, h₀⟩ ⟨p₁, h₁⟩ ⟨p₂, h₂⟩ := ⟨p₂ + pf.comp (p₀ + p₁),
  λ x y, by { refine (hf (g₀ x, g₁ x) (g₂ x y)).trans _, simp [← add_assoc], mono*, }⟩

theorem polysize_safe.comp₄_1 {f : α₀ → α₁ → α₂ → α₃ → γ}
  {g₀ : α → α₀} {g₁ : α → β → α₁} {g₂ : α → α₂} {g₃ : α → α₃} :
  polysize_safe (λ (usf : α₀ × α₂ × α₃) (sf : α₁), f usf.1 sf usf.2.1 usf.2.2) →
  polysize_fun g₀ → polysize_safe g₁ → polysize_fun g₂ → polysize_fun g₃ →
  polysize_safe (λ x y, f (g₀ x) (g₁ x y) (g₂ x) (g₃ x))
| ⟨pf, hf⟩ ⟨p₀, h₀⟩ ⟨p₁, h₁⟩ ⟨p₂, h₂⟩ ⟨p₃, h₃⟩ := ⟨p₁ + pf.comp (p₀ + p₂ + p₃),
  λ x y, by { refine (hf (g₀ x, g₂ x, g₃ x) (g₁ x y)).trans _, simp [← add_assoc], mono*, }⟩

theorem polysize_safe.comp₄_2 {f : α₀ → α₁ → α₂ → α₃ → γ}
  {g₀ : α → α₀} {g₁ : α → α₁} {g₂ : α → β → α₂} {g₃ : α → α₃} :
  polysize_safe (λ (usf : α₀ × α₁ × α₃) (sf : α₂), f usf.1 usf.2.1 sf usf.2.2) →
  polysize_fun g₀ → polysize_fun g₁ → polysize_safe g₂ → polysize_fun g₃ →
  polysize_safe (λ x y, f (g₀ x) (g₁ x) (g₂ x y) (g₃ x))
| ⟨pf, hf⟩ ⟨p₀, h₀⟩ ⟨p₁, h₁⟩ ⟨p₂, h₂⟩ ⟨p₃, h₃⟩ := ⟨p₂ + pf.comp (p₀ + p₁ + p₃),
  λ x y, by { refine (hf (g₀ x, g₁ x, g₃ x) (g₂ x y)).trans _, simp [← add_assoc], mono*, }⟩

theorem polysize_safe.comp₄_3 {f : α₀ → α₁ → α₂ → α₃ → γ}
  {g₀ : α → α₀} {g₁ : α → α₁} {g₂ : α → α₂} {g₃ : α → β → α₃} :
  polysize_safe (λ (usf : α₀ × α₁ × α₂) (sf : α₃), f usf.1 usf.2.1 usf.2.2 sf) →
  polysize_fun g₀ → polysize_fun g₁ → polysize_fun g₂ → polysize_safe g₃ →
  polysize_safe (λ x y, f (g₀ x) (g₁ x) (g₂ x) (g₃ x y))
| ⟨pf, hf⟩ ⟨p₀, h₀⟩ ⟨p₁, h₁⟩ ⟨p₂, h₂⟩ ⟨p₃, h₃⟩ := ⟨p₃ + pf.comp (p₀ + p₁ + p₂),
  λ x y, by { refine (hf (g₀ x, g₁ x, g₂ x) (g₃ x y)).trans _, simp [← add_assoc], mono*, }⟩

theorem polysize_safe.comp₅_0 {f : α₀ → α₁ → α₂ → α₃ → α₄ → γ}
  {g₀ : α → β → α₀} {g₁ : α → α₁} {g₂ : α → α₂} {g₃ : α → α₃} {g₄ : α → α₄} :
  polysize_safe (λ (usf : α₁ × α₂ × α₃ × α₄) (sf : α₀), f sf usf.1 usf.2.1 usf.2.2.1 usf.2.2.2) →
  polysize_safe g₀ → polysize_fun g₁ → polysize_fun g₂ → polysize_fun g₃ → polysize_fun g₄ →
  polysize_safe (λ x y, f (g₀ x y) (g₁ x) (g₂ x) (g₃ x) (g₄ x))
| ⟨pf, hf⟩ ⟨p₀, h₀⟩ ⟨p₁, h₁⟩ ⟨p₂, h₂⟩ ⟨p₃, h₃⟩ ⟨p₄, h₄⟩ := ⟨p₀ + pf.comp (p₁ + p₂ + p₃ + p₄),
  λ x y, by { refine (hf (g₁ x, g₂ x, g₃ x, g₄ x) (g₀ x y)).trans _, simp [← add_assoc], mono*, }⟩

theorem polysize_safe.comp₅_1 {f : α₀ → α₁ → α₂ → α₃ → α₄ → γ}
  {g₀ : α → α₀} {g₁ : α → β → α₁} {g₂ : α → α₂} {g₃ : α → α₃} {g₄ : α → α₄} :
  polysize_safe (λ (usf : α₀ × α₂ × α₃ × α₄) (sf : α₁), f usf.1 sf usf.2.1 usf.2.2.1 usf.2.2.2) →
  polysize_fun g₀ → polysize_safe g₁ → polysize_fun g₂ → polysize_fun g₃ → polysize_fun g₄ →
  polysize_safe (λ x y, f (g₀ x) (g₁ x y) (g₂ x) (g₃ x) (g₄ x))
| ⟨pf, hf⟩ ⟨p₀, h₀⟩ ⟨p₁, h₁⟩ ⟨p₂, h₂⟩ ⟨p₃, h₃⟩ ⟨p₄, h₄⟩ := ⟨p₁ + pf.comp (p₀ + p₂ + p₃ + p₄),
  λ x y, by { refine (hf (g₀ x, g₂ x, g₃ x, g₄ x) (g₁ x y)).trans _, simp [← add_assoc], mono*, }⟩

end comp

@[complexity] theorem polysize_safe.list_cases_on {f : α → list γ} {g : α → β → δ}
  {h : α → β → γ → list γ → δ} (hf : polysize_fun f) (hg : polysize_safe g)
  (hh : polysize_safe (λ (usf : α × γ × list γ) (sf : β), h usf.1 sf usf.2.1 usf.2.2)) :
  @polysize_safe _ _ δ _ _ _ _ _ _ (λ x y, @list.cases_on _ (λ _, δ) (f x) (g x y) (h x y)) :=
begin
  convert_to polysize_safe (λ x y, ((f x).head'.map (λ hd, h x y hd (f x).tail)).get_or_else (g x y)),
  { ext x y, cases f x; simp, },
  refine polysize_safe.get_or_else _ hg,
  refine polysize_safe.option_map₁ (polysize_fun.head'.comp hf) _,
  exact hh.comp₄_1 polysize_fun.fst polysize_safe.id polysize_fun.snd
    (polysize_fun.tail.comp $ hf.comp polysize_fun.fst),
end


/-
Copyright (c) 2020 Bhavik Mehta. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Bhavik Mehta
-/
import category_theory.limits.preserves.shapes.binary_products
import category_theory.limits.preserves.shapes.products
import category_theory.limits.shapes.binary_products
import category_theory.limits.shapes.finite_products
import logic.equiv.fin

/-!
# Constructing finite products from binary products and terminal.

> THIS FILE IS SYNCHRONIZED WITH MATHLIB4.
> Any changes to this file require a corresponding PR to mathlib4.

If a category has binary products and a terminal object then it has finite products.
If a functor preserves binary products and the terminal object then it preserves finite products.

# TODO

Provide the dual results.
Show the analogous results for functors which reflect or create (co)limits.
-/

universes v v' u u'

noncomputable theory
open category_theory category_theory.category category_theory.limits
namespace category_theory

variables {J : Type v} [small_category J]
variables {C : Type u} [category.{v} C]
variables {D : Type u'} [category.{v'} D]

/--
Given `n+1` objects of `C`, a fan for the last `n` with point `c₁.X` and a binary fan on `c₁.X` and
`f 0`, we can build a fan for all `n+1`.

In `extend_fan_is_limit` we show that if the two given fans are limits, then this fan is also a
limit.
-/
@[simps {rhs_md := semireducible}]
def extend_fan {n : ℕ} {f : fin (n+1) → C}
  (c₁ : fan (λ (i : fin n), f i.succ))
  (c₂ : binary_fan (f 0) c₁.X) :
  fan f :=
fan.mk c₂.X
begin
  refine fin.cases _ _,
  { apply c₂.fst },
  { intro i, apply c₂.snd ≫ c₁.π.app ⟨i⟩ },
end

/--
Show that if the two given fans in `extend_fan` are limits, then the constructed fan is also a
limit.
-/
def extend_fan_is_limit {n : ℕ} (f : fin (n+1) → C)
  {c₁ : fan (λ (i : fin n), f i.succ)} {c₂ : binary_fan (f 0) c₁.X}
  (t₁ : is_limit c₁) (t₂ : is_limit c₂) :
  is_limit (extend_fan c₁ c₂) :=
{ lift := λ s,
  begin
    apply (binary_fan.is_limit.lift' t₂ (s.π.app ⟨0⟩) _).1,
    apply t₁.lift ⟨_, discrete.nat_trans (λ ⟨i⟩, s.π.app ⟨i.succ⟩)⟩
  end,
  fac' := λ s ⟨j⟩,
  begin
    apply fin.induction_on j,
    { apply (binary_fan.is_limit.lift' t₂ _ _).2.1 },
    { rintro i -,
      dsimp only [extend_fan_π_app],
      rw [fin.cases_succ, ← assoc, (binary_fan.is_limit.lift' t₂ _ _).2.2, t₁.fac],
      refl }
  end,
  uniq' := λ s m w,
  begin
    apply binary_fan.is_limit.hom_ext t₂,
    { rw (binary_fan.is_limit.lift' t₂ _ _).2.1,
      apply w ⟨0⟩ },
    { rw (binary_fan.is_limit.lift' t₂ _ _).2.2,
      apply t₁.uniq ⟨_, _⟩,
      rintro ⟨j⟩,
      rw assoc,
      dsimp only [discrete.nat_trans_app, extend_fan_is_limit._match_1],
      rw ← w ⟨j.succ⟩,
      dsimp only [extend_fan_π_app],
      rw fin.cases_succ }
  end }

section
variables [has_binary_products C] [has_terminal C]

/--
If `C` has a terminal object and binary products, then it has a product for objects indexed by
`fin n`.
This is a helper lemma for `has_finite_products_of_has_binary_and_terminal`, which is more general
than this.
-/
private lemma has_product_fin :
  Π (n : ℕ) (f : fin n → C), has_product f
| 0 := λ f,
  begin
    letI : has_limits_of_shape (discrete (fin 0)) C :=
      has_limits_of_shape_of_equivalence (discrete.equivalence.{0} fin_zero_equiv'.symm),
    apply_instance,
  end
| (n+1) := λ f,
  begin
    haveI := has_product_fin n,
    apply has_limit.mk ⟨_, extend_fan_is_limit f (limit.is_limit _) (limit.is_limit _)⟩,
  end

/-- If `C` has a terminal object and binary products, then it has finite products. -/
lemma has_finite_products_of_has_binary_and_terminal : has_finite_products C :=
begin
  refine ⟨λ n, ⟨λ K, _⟩⟩,
  letI := has_product_fin n (λ n, K.obj ⟨n⟩),
  let : discrete.functor (λ n, K.obj ⟨n⟩) ≅ K := discrete.nat_iso (λ ⟨i⟩, iso.refl _),
  apply has_limit_of_iso this,
end

end

section preserves
variables (F : C ⥤ D)
variables [preserves_limits_of_shape (discrete walking_pair) F]
variables [preserves_limits_of_shape (discrete.{0} pempty) F]
variables [has_finite_products.{v} C]

/--
If `F` preserves the terminal object and binary products, then it preserves products indexed by
`fin n` for any `n`.
-/
noncomputable def preserves_fin_of_preserves_binary_and_terminal  :
  Π (n : ℕ) (f : fin n → C), preserves_limit (discrete.functor f) F
| 0 := λ f,
  begin
    letI : preserves_limits_of_shape (discrete (fin 0)) F :=
      preserves_limits_of_shape_of_equiv.{0 0}
        (discrete.equivalence fin_zero_equiv'.symm) _,
    apply_instance,
  end
| (n+1) :=
  begin
    haveI := preserves_fin_of_preserves_binary_and_terminal n,
    intro f,
    refine preserves_limit_of_preserves_limit_cone
      (extend_fan_is_limit f (limit.is_limit _) (limit.is_limit _)) _,
    apply (is_limit_map_cone_fan_mk_equiv _ _ _).symm _,
    let := extend_fan_is_limit (λ i, F.obj (f i))
              (is_limit_of_has_product_of_preserves_limit F _)
              (is_limit_of_has_binary_product_of_preserves_limit F _ _),
    refine is_limit.of_iso_limit this _,
    apply cones.ext _ _,
    apply iso.refl _,
    rintro ⟨j⟩,
    apply fin.induction_on j,
    { apply (category.id_comp _).symm },
    { rintro i -,
      dsimp only [extend_fan_π_app, iso.refl_hom, fan.mk_π_app],
      rw [fin.cases_succ, fin.cases_succ],
      change F.map _ ≫ _ = 𝟙 _ ≫ _,
      rw [id_comp, ←F.map_comp],
      refl }
  end

/--
If `F` preserves the terminal object and binary products, then it preserves limits of shape
`discrete (fin n)`.
-/
def preserves_shape_fin_of_preserves_binary_and_terminal (n : ℕ) :
  preserves_limits_of_shape (discrete (fin n)) F :=
{ preserves_limit := λ K,
  begin
    let : discrete.functor (λ n, K.obj ⟨n⟩) ≅ K := discrete.nat_iso (λ ⟨i⟩, iso.refl _),
    haveI := preserves_fin_of_preserves_binary_and_terminal F n (λ n, K.obj ⟨n⟩),
    apply preserves_limit_of_iso_diagram F this,
  end }

/-- If `F` preserves the terminal object and binary products then it preserves finite products. -/
def preserves_finite_products_of_preserves_binary_and_terminal
  (J : Type) [fintype J] :
  preserves_limits_of_shape (discrete J) F :=
begin
  classical,
  let e := fintype.equiv_fin J,
  haveI := preserves_shape_fin_of_preserves_binary_and_terminal F (fintype.card J),
  apply preserves_limits_of_shape_of_equiv.{0 0}
    (discrete.equivalence e).symm,
end

end preserves

/--
Given `n+1` objects of `C`, a cofan for the last `n` with point `c₁.X`
and a binary cofan on `c₁.X` and `f 0`, we can build a cofan for all `n+1`.

In `extend_cofan_is_colimit` we show that if the two given cofans are colimits,
then this cofan is also a colimit.
-/
@[simps {rhs_md := semireducible}]
def extend_cofan {n : ℕ} {f : fin (n+1) → C}
  (c₁ : cofan (λ (i : fin n), f i.succ))
  (c₂ : binary_cofan (f 0) c₁.X) :
  cofan f :=
cofan.mk c₂.X
begin
  refine fin.cases _ _,
  { apply c₂.inl },
  { intro i,
    apply c₁.ι.app ⟨i⟩ ≫ c₂.inr },
end

/--
Show that if the two given cofans in `extend_cofan` are colimits,
then the constructed cofan is also a colimit.
-/
def extend_cofan_is_colimit {n : ℕ} (f : fin (n+1) → C)
  {c₁ : cofan (λ (i : fin n), f i.succ)} {c₂ : binary_cofan (f 0) c₁.X}
  (t₁ : is_colimit c₁) (t₂ : is_colimit c₂) :
  is_colimit (extend_cofan c₁ c₂) :=
{ desc := λ s,
  begin
    apply (binary_cofan.is_colimit.desc' t₂ (s.ι.app ⟨0⟩) _).1,
    apply t₁.desc ⟨_, discrete.nat_trans (λ i, s.ι.app ⟨i.as.succ⟩)⟩
  end,
  fac' := λ s,
  begin
    rintro ⟨j⟩,
    apply fin.induction_on j,
    { apply (binary_cofan.is_colimit.desc' t₂ _ _).2.1 },
    { rintro i -,
      dsimp only [extend_cofan_ι_app],
      rw [fin.cases_succ, assoc, (binary_cofan.is_colimit.desc' t₂ _ _).2.2, t₁.fac],
      refl }
  end,
  uniq' := λ s m w,
  begin
    apply binary_cofan.is_colimit.hom_ext t₂,
    { rw (binary_cofan.is_colimit.desc' t₂ _ _).2.1,
      apply w ⟨0⟩ },
    { rw (binary_cofan.is_colimit.desc' t₂ _ _).2.2,
      apply t₁.uniq ⟨_, _⟩,
      rintro ⟨j⟩,
      dsimp only [discrete.nat_trans_app],
      rw ← w ⟨j.succ⟩,
      dsimp only [extend_cofan_ι_app],
      rw [fin.cases_succ, assoc], }
  end }

section
variables [has_binary_coproducts C] [has_initial C]

/--
If `C` has an initial object and binary coproducts, then it has a coproduct for objects indexed by
`fin n`.
This is a helper lemma for `has_cofinite_products_of_has_binary_and_terminal`, which is more general
than this.
-/
private lemma has_coproduct_fin :
  Π (n : ℕ) (f : fin n → C), has_coproduct f
| 0 := λ f,
  begin
    letI : has_colimits_of_shape (discrete (fin 0)) C :=
      has_colimits_of_shape_of_equivalence (discrete.equivalence.{0} fin_zero_equiv'.symm),
    apply_instance,
  end
| (n+1) := λ f,
  begin
    haveI := has_coproduct_fin n,
    apply has_colimit.mk
      ⟨_, extend_cofan_is_colimit f (colimit.is_colimit _) (colimit.is_colimit _)⟩,
  end

/-- If `C` has an initial object and binary coproducts, then it has finite coproducts. -/
lemma has_finite_coproducts_of_has_binary_and_initial : has_finite_coproducts C :=
begin
  refine ⟨λ n, ⟨λ K, _⟩⟩,
  letI := has_coproduct_fin n (λ n, K.obj ⟨n⟩),
  let : K ≅ discrete.functor (λ n, K.obj ⟨n⟩) := discrete.nat_iso (λ ⟨i⟩, iso.refl _),
  apply has_colimit_of_iso this,
end

end

section preserves
variables (F : C ⥤ D)
variables [preserves_colimits_of_shape (discrete walking_pair) F]
variables [preserves_colimits_of_shape (discrete.{0} pempty) F]
variables [has_finite_coproducts.{v} C]

/--
If `F` preserves the initial object and binary coproducts, then it preserves products indexed by
`fin n` for any `n`.
-/
noncomputable def preserves_fin_of_preserves_binary_and_initial  :
  Π (n : ℕ) (f : fin n → C), preserves_colimit (discrete.functor f) F
| 0 := λ f,
  begin
    letI : preserves_colimits_of_shape (discrete (fin 0)) F :=
      preserves_colimits_of_shape_of_equiv.{0 0}
        (discrete.equivalence fin_zero_equiv'.symm) _,
    apply_instance,
  end
| (n+1) :=
  begin
    haveI := preserves_fin_of_preserves_binary_and_initial n,
    intro f,
    refine preserves_colimit_of_preserves_colimit_cocone
      (extend_cofan_is_colimit f (colimit.is_colimit _) (colimit.is_colimit _)) _,
    apply (is_colimit_map_cocone_cofan_mk_equiv _ _ _).symm _,
    let := extend_cofan_is_colimit (λ i, F.obj (f i))
              (is_colimit_of_has_coproduct_of_preserves_colimit F _)
              (is_colimit_of_has_binary_coproduct_of_preserves_colimit F _ _),
    refine is_colimit.of_iso_colimit this _,
    apply cocones.ext _ _,
    apply iso.refl _,
    rintro ⟨j⟩,
    apply fin.induction_on j,
    { apply category.comp_id },
    { rintro i -,
      dsimp only [extend_cofan_ι_app, iso.refl_hom, cofan.mk_ι_app],
      rw [fin.cases_succ, fin.cases_succ],
      erw [comp_id, ←F.map_comp],
      refl, }
  end

/--
If `F` preserves the initial object and binary coproducts, then it preserves colimits of shape
`discrete (fin n)`.
-/
def preserves_shape_fin_of_preserves_binary_and_initial (n : ℕ) :
  preserves_colimits_of_shape (discrete (fin n)) F :=
{ preserves_colimit := λ K,
  begin
    let : discrete.functor (λ n, K.obj ⟨n⟩) ≅ K := discrete.nat_iso (λ ⟨i⟩, iso.refl _),
    haveI := preserves_fin_of_preserves_binary_and_initial F n (λ n, K.obj ⟨n⟩),
    apply preserves_colimit_of_iso_diagram F this,
  end }

/-- If `F` preserves the initial object and binary coproducts then it preserves finite products. -/
def preserves_finite_coproducts_of_preserves_binary_and_initial
  (J : Type) [fintype J] :
  preserves_colimits_of_shape (discrete J) F :=
begin
  classical,
  let e := fintype.equiv_fin J,
  haveI := preserves_shape_fin_of_preserves_binary_and_initial F (fintype.card J),
  apply preserves_colimits_of_shape_of_equiv.{0 0} (discrete.equivalence e).symm,
end

end preserves

end category_theory


import local.h_principle

/-!
In this file we prove the parametric version of the local h-principle.

We will not use this to prove the global version of the h-principle, but we do use this to conclude
the existence of sphere eversion from the local h-principle, which is proven in `local.h_principle`.

The parametric h-principle states the following: Suppose that `R` is a local relation,
`𝓕₀ : P → J¹(E, F)` is a family of formal solutions of `R` that is holonomic near some set
`C ⊆ P × E`, `K ⊆ P × E` is compact and `ε : ℝ`,
then there exists a homotopy `𝓕 : ℝ × P → J¹(E, F)` between `𝓕` and a solution that is holonomic
near `K`, that agrees with `𝓕₀` near `C` and is everywhere `ε`-close to `𝓕₀`
-/

noncomputable theory

open metric finite_dimensional set function rel_loc linear_map (ker)
open_locale topology pointwise

section parameter_space

variables
{E : Type*} [normed_add_comm_group E] [normed_space ℝ E]
{F : Type*} [normed_add_comm_group F] [normed_space ℝ F]
-- `G` will be `ℝ` in the proof of the parametric h-principle.
-- It indicates the homotopy variable `t`.
{G : Type*} [normed_add_comm_group G] [normed_space ℝ G]
{P : Type*} [normed_add_comm_group P] [normed_space ℝ P]

variables {R : rel_loc E F}

/-- The projection `J¹(P × E, F) → J¹(E, F)`. -/
def one_jet_snd : one_jet (P × E) F → one_jet E F :=
λ p, (p.1.2, p.2.1, p.2.2 ∘L fderiv ℝ (λ y, (p.1.1, y)) p.1.2)

lemma continuous_one_jet_snd :
  continuous (one_jet_snd : one_jet (P × E) F → one_jet E F) :=
continuous_fst.snd.prod_mk $ continuous_snd.fst.prod_mk $ continuous_snd.snd.clm_comp $
  continuous.fderiv (cont_diff_fst.fst.prod_map cont_diff_id) continuous_fst.snd le_top

lemma one_jet_snd_eq (p : one_jet (P × E) F) :
  one_jet_snd p = (p.1.2, p.2.1, p.2.2 ∘L continuous_linear_map.inr ℝ P E) :=
by simp_rw [one_jet_snd, fderiv_prod_right]

variables (P)
/-- The relation `R.relativize P` (`𝓡 ^ P` in the blueprint) is the relation on `J¹(P × E, F)`
induced by `R`. -/
def rel_loc.relativize (R : rel_loc E F) : rel_loc (P × E) F :=
one_jet_snd ⁻¹' R
variables {P}

lemma rel_loc.mem_relativize (R : rel_loc E F) (w : one_jet (P × E) F) :
 w ∈ R.relativize P ↔ (w.1.2, w.2.1, w.2.2 ∘L continuous_linear_map.inr ℝ P E) ∈ R :=
by simp_rw [rel_loc.relativize, mem_preimage, one_jet_snd_eq]

lemma rel_loc.is_open_relativize (R : rel_loc E F) (h2 : is_open R) :
  is_open (R.relativize P) :=
h2.preimage continuous_one_jet_snd

lemma relativize_slice_loc {σ : one_jet (P × E) F}
  {p : dual_pair (P × E)}
  (q : dual_pair E)
  (hpq : p.π.comp (continuous_linear_map.inr ℝ P E) = q.π) :
  (R.relativize P).slice p σ =
  σ.2.2 (p.v - (0, q.v)) +ᵥ R.slice q (one_jet_snd σ) :=
begin
  have h2pq : ∀ x : E, p.π ((0 : P), x) = q.π x := λ x, congr_arg (λ f : E →L[ℝ] ℝ, f x) hpq,
  ext1 w,
  have h1 : (p.update σ.2.2 w).comp (continuous_linear_map.inr ℝ P E) =
    q.update (one_jet_snd σ).2.2 (-σ.2.2 (p.v - (0, q.v)) +ᵥ w),
  { ext1 x,
    simp_rw [continuous_linear_map.comp_apply, continuous_linear_map.inr_apply,
      ← continuous_linear_map.map_neg, neg_sub],
    obtain ⟨u, hu, t, rfl⟩ := q.decomp x,
    have hv : (0, q.v) - p.v ∈ ker p.π,
    { rw [linear_map.mem_ker, map_sub, p.pairing, h2pq, q.pairing, sub_self] },
    have hup : ((0 : P), u) ∈ ker p.π := (h2pq u).trans hu,
    rw [q.update_apply _ hu, ← prod.zero_mk_add_zero_mk, map_add, p.update_ker_pi _ _ hup,
      ← prod.smul_zero_mk, map_smul, vadd_eq_add],
    nth_rewrite 0 [← sub_add_cancel (0, q.v) p.v],
    rw [map_add, p.update_ker_pi _ _ hv, p.update_v, one_jet_snd_eq],
    refl },
  have := preimage_vadd_neg (show F, from σ.2.2 (p.v - (0, q.v))) (R.slice q (one_jet_snd σ)),
  dsimp only at this,
  simp_rw [← this, mem_preimage, mem_slice, R.mem_relativize, h1],
  refl,
end

lemma relativize_slice_eq_univ_loc {σ : one_jet (P × E) F}
  {p : dual_pair (P × E)}
  (hp : p.π.comp (continuous_linear_map.inr ℝ P E) = 0) :
  ((R.relativize P).slice p σ).nonempty ↔
  (R.relativize P).slice p σ = univ :=
begin
  have h2p : ∀ x : E, p.π ((0 : P), x) = 0 := λ x, congr_arg (λ f : E →L[ℝ] ℝ, f x) hp,
  have : ∀ y : F, (p.update σ.2.2 y).comp (continuous_linear_map.inr ℝ P E) =
    σ.2.2.comp (continuous_linear_map.inr ℝ P E),
  { intro y,
    ext1 x,
    simp_rw [continuous_linear_map.comp_apply, continuous_linear_map.inr_apply,
      p.update_ker_pi _ _ (h2p x)] },
  simp_rw [set.nonempty, eq_univ_iff_forall, mem_slice, R.mem_relativize, this, exists_const,
    forall_const]
end

variables (P)

lemma rel_loc.is_ample.relativize (hR : R.is_ample) : (R.relativize P).is_ample :=
begin
  intros p σ,
  let p2 := p.π.comp (continuous_linear_map.inr ℝ P E),
  rcases eq_or_ne p2 0 with h|h,
  { intros w hw,
    rw [(relativize_slice_eq_univ_loc h).mp ⟨w, hw⟩, connected_component_in_univ,
      preconnected_space.connected_component_eq_univ, convex_hull_univ] },
  obtain ⟨u', hu'⟩ := continuous_linear_map.exists_ne_zero h,
  let u := (p2 u')⁻¹ • u',
  let q : dual_pair E :=
  ⟨p2, u, by rw [p2.map_smul, smul_eq_mul, inv_mul_cancel hu']⟩,
  rw [relativize_slice_loc q rfl],
  exact (hR q _).vadd
end

variables {P}

/-- Turn a family of sections of `J¹(E, E')` parametrized by `P` into a section of `J¹(P × E, E')`.
-/
@[simps]
def family_jet_sec.uncurry (S : family_jet_sec E F P) : jet_sec (P × E) F :=
{ f := λ p, S.f p.1 p.2,
  φ := λ p, fderiv ℝ (λ z : P × E, S.f z.1 p.2) p +
    S.φ p.1 p.2 ∘L fderiv ℝ prod.snd p,
  f_diff := S.f_diff,
  φ_diff := begin
    refine (cont_diff.fderiv _ cont_diff_id le_top).add (S.φ_diff.clm_comp _),
    { exact S.f_diff.comp (cont_diff_snd.fst.prod cont_diff_fst.snd) },
    { exact cont_diff.fderiv cont_diff_snd.snd cont_diff_id le_top }
    end }

lemma family_jet_sec.uncurry_φ' (S : family_jet_sec E F P) (p : P × E) :
  (S.uncurry).φ p = fderiv ℝ (λ z, S.f z p.2) p.1 ∘L continuous_linear_map.fst ℝ P E +
  S.φ p.1 p.2 ∘L continuous_linear_map.snd ℝ P E :=
begin
  simp_rw [S.uncurry_φ, fderiv_snd, add_left_inj],
  refine (fderiv_comp p
    ((S.f_diff.comp (cont_diff_id.prod cont_diff_const)).differentiable le_top p.1)
    differentiable_at_fst).trans _,
  rw [fderiv_fst],
  refl,
end

lemma family_jet_sec.uncurry_mem_relativize (S : family_jet_sec E F P) {s : P}
  {x : E} : ((s, x), S.uncurry (s, x)) ∈ R.relativize P ↔ (x, S s x) ∈ R :=
begin
  simp_rw [rel_loc.relativize, mem_preimage, one_jet_snd_eq, jet_sec.coe_apply, S.uncurry_f,
    S.uncurry_φ'],
  congr' 2,
  refine prod.ext rfl (prod.ext rfl _),
  ext v,
  simp_rw [continuous_linear_map.comp_apply, continuous_linear_map.add_apply,
    continuous_linear_map.comp_apply, continuous_linear_map.inr_apply,
    continuous_linear_map.coe_fst', continuous_linear_map.coe_snd',
    continuous_linear_map.map_zero, zero_add],
  refl,
end

lemma family_jet_sec.is_holonomic_at_uncurry (S : family_jet_sec E F P) {p : P × E} :
  S.uncurry.is_holonomic_at p ↔ (S p.1).is_holonomic_at p.2 :=
begin
  simp_rw [jet_sec.is_holonomic_at, S.uncurry_φ],
  rw [show S.uncurry.f = λ x, S.uncurry.f x, from rfl, funext S.uncurry_f,
    show (λ x : P × E, S.f x.1 x.2) = ↿S.f, from rfl],
  simp_rw [fderiv_prod_eq_add (S.f_diff.differentiable le_top _), fderiv_snd],
  refine (add_right_inj _).trans _,
  have := fderiv_comp p ((S p.1).f_diff.cont_diff_at.differentiable_at le_top)
    differentiable_at_snd,
  rw [show D (λ (z : P × E), ↿(S.f) (p.fst, z.snd)) p = _, from this, fderiv_snd,
    (show surjective (continuous_linear_map.snd ℝ P E), from prod.snd_surjective)
      .clm_comp_injective.eq_iff],
  refl
end

/-- Turn a family of formal solutions of `R ⊆ J¹(E, E')` parametrized by `P` into a formal solution
of `R.relativize P`. -/
def rel_loc.family_formal_sol.uncurry (S : R.family_formal_sol P) : formal_sol (R.relativize P) :=
begin
  refine ⟨S.to_family_jet_sec.uncurry, _⟩,
  rintro ⟨s, x⟩,
  exact S.to_family_jet_sec.uncurry_mem_relativize.mpr (S.is_sol s x)
end

lemma rel_loc.family_formal_sol.uncurry_φ' (S : R.family_formal_sol P) (p : P × E) :
  (S.uncurry p).2 = fderiv ℝ (λ z, S.f z p.2) p.1 ∘L continuous_linear_map.fst ℝ P E +
  S.φ p.1 p.2 ∘L continuous_linear_map.snd ℝ P E :=
S.to_family_jet_sec.uncurry_φ' p

/-- Turn a family of sections of `J¹(P × E, F)` parametrized by `G` into a family of sections of
`J¹(E, F)` parametrized by `G × P`. -/
def family_jet_sec.curry (S : family_jet_sec (P × E) F G) :
  family_jet_sec E F (G × P) :=
{ f := λ p x, (S p.1).f (p.2, x),
  φ := λ p x, (S p.1).φ (p.2, x) ∘L fderiv ℝ (λ x, (p.2, x)) x,
  f_diff := S.f_diff.comp (cont_diff_prod_assoc : cont_diff ℝ ⊤ (equiv.prod_assoc G P E)),
  φ_diff := begin
    refine (S.φ_diff.comp (cont_diff_prod_assoc : cont_diff ℝ ⊤ (equiv.prod_assoc G P E))).clm_comp
      _,
    refine cont_diff.fderiv _ cont_diff_snd le_top,
    exact cont_diff_fst.fst.snd.prod cont_diff_snd
  end }

lemma family_jet_sec.curry_f (S : family_jet_sec (P × E) F G) (p : G × P)
  (x : E) : (S.curry p).f x = (S p.1).f (p.2, x) :=
rfl

lemma family_jet_sec.curry_φ (S : family_jet_sec (P × E) F G) (p : G × P)
  (x : E) : (S.curry p).φ x = (S p.1).φ (p.2, x) ∘L fderiv ℝ (λ x, (p.2, x)) x :=
rfl

lemma family_jet_sec.curry_φ' (S : family_jet_sec (P × E) F G) (p : G × P)
  (x : E) : (S.curry p).φ x = (S p.1).φ (p.2, x) ∘L continuous_linear_map.inr ℝ P E :=
begin
  rw [S.curry_φ],
  congr' 1,
  refine ((differentiable_at_const _).fderiv_prod differentiable_at_id).trans _,
  rw [fderiv_id, fderiv_const],
  refl,
end

lemma family_jet_sec.is_holonomic_at_curry
  (S : family_jet_sec (P × E) F G)
  {t : G} {s : P} {x : E} (hS : (S t).is_holonomic_at (s, x)) :
  (S.curry (t, s)).is_holonomic_at x :=
begin
  simp_rw [jet_sec.is_holonomic_at, S.curry_φ] at hS ⊢,
  rw [show (S.curry (t, s)).f = λ x, (S.curry (t, s)).f x, from rfl, funext (S.curry_f _)],
  dsimp only,
  refine (fderiv_comp x ((S t).f_diff.cont_diff_at.differentiable_at le_top)
    ((differentiable_at_const _).prod differentiable_at_id)).trans _,
  rw [id, hS],
  refl,
end

lemma family_jet_sec.curry_mem (S : family_jet_sec (P × E) F G)
  {p : G × P} {x : E} (hR : ((p.2, x), S p.1 (p.2, x)) ∈ R.relativize P) :
  (x, S.curry p x) ∈ R :=
begin
  simp_rw [rel_loc.relativize, mem_preimage, jet_sec.coe_apply, one_jet_snd_eq, S.curry_φ'] at hR ⊢,
  exact hR
end

/-- Turn a family of formal solutions of `R.relativize P` parametrized by `G` into a family of
formal solutions of `R` parametrized by `G × P`. -/
def rel_loc.family_formal_sol.curry (S : family_formal_sol G (R.relativize P)) :
  family_formal_sol (G × P) R :=
⟨S.to_family_jet_sec.curry, λ p x, S.to_family_jet_sec.curry_mem (S.is_sol _ _)⟩

lemma rel_loc.family_formal_sol.curry_φ (S : family_formal_sol G (R.relativize P)) (p : G × P)
  (x : E) : (S.curry p).φ x = (S p.1).φ (p.2, x) ∘L fderiv ℝ (λ x, (p.2, x)) x :=
rfl

lemma rel_loc.family_formal_sol.curry_φ' (S : family_formal_sol G (R.relativize P)) (p : G × P)
  (x : E) : (S.curry p x).2 = (S p.1 (p.2, x)).2 ∘L continuous_linear_map.inr ℝ P E :=
S.to_family_jet_sec.curry_φ' p x

lemma curry_eq_iff_eq_uncurry_loc {𝓕 : family_formal_sol G (R.relativize P)}
  {𝓕₀ : R.family_formal_sol P} {t : G} {x : E} {s : P}
  (h : 𝓕 t (s, x) = 𝓕₀.uncurry (s, x)) :
  (𝓕.curry (t, s)) x = 𝓕₀ s x :=
begin
  simp_rw [prod.ext_iff] at h ⊢,
  refine ⟨h.1, _⟩,
  simp_rw [𝓕.curry_φ', h.2, 𝓕₀.uncurry_φ'],
  ext v,
  simp_rw [continuous_linear_map.comp_apply, continuous_linear_map.add_apply,
    continuous_linear_map.comp_apply,
    continuous_linear_map.inr_apply, continuous_linear_map.coe_fst',
    continuous_linear_map.coe_snd', continuous_linear_map.map_zero, zero_add],
  refl
end

end parameter_space

section parametric_h_principle

variables {E : Type*} [normed_add_comm_group E] [normed_space ℝ E] [finite_dimensional ℝ E]
          {F : Type*} [normed_add_comm_group F] [normed_space ℝ F] [finite_dimensional ℝ F]
          {G : Type*} [normed_add_comm_group G] [normed_space ℝ G]
          {P : Type*} [normed_add_comm_group P] [normed_space ℝ P] [finite_dimensional ℝ P]

variables {R : rel_loc E F} (h_op: is_open R) (h_ample: R.is_ample) (L : landscape E)
include h_op h_ample

/- The local parametric h-principle. -/
lemma rel_loc.family_formal_sol.improve_htpy
  {ε : ℝ} (ε_pos : 0 < ε)
  (C : set (P × E)) (hC : is_closed C) (K : set (P × E)) (hK : is_compact K)
  (𝓕₀ : family_formal_sol P R)
  (h_hol : ∀ᶠ (p : P × E) near C, (𝓕₀ p.1).is_holonomic_at p.2) :
  ∃ 𝓕 : family_formal_sol (ℝ × P) R,
    (∀ s x, 𝓕 (0, s) x = 𝓕₀ s x) ∧
    (∀ᶠ (p : P × E) near C, ∀ t, 𝓕 (t, p.1) p.2 = 𝓕₀ p.1 p.2) ∧
    (∀ s x t, ‖(𝓕 (t, s)).f x - 𝓕₀.f s x‖ ≤ ε)  ∧
    (∀ᶠ (p : P × E) near K, (𝓕 (1, p.1)).is_holonomic_at p.2) :=
begin
  let parametric_landscape : landscape (P × E) :=
  { C := C,
    K₀ := K,
    K₁ := (exists_compact_superset hK).some,
    hC := hC,
    hK₀ := hK,
    hK₁ := (exists_compact_superset hK).some_spec.1,
    h₀₁ := (exists_compact_superset hK).some_spec.2 },
  obtain ⟨𝓕, h₁, -, h₂, -, h₄, h₅⟩ :=
    𝓕₀.uncurry.improve_htpy' (R.is_open_relativize h_op) (h_ample.relativize P)
    parametric_landscape ε_pos (h_hol.mono (λ p hp, 𝓕₀.is_holonomic_at_uncurry.mpr hp)),
  have h₁ : ∀ p, 𝓕 0 p = 𝓕₀.uncurry p,
  { intro p, rw h₁.on_set 0 right_mem_Iic, refl },
  refine ⟨𝓕.curry, _, _, _, _⟩,
  { intros s x, exact curry_eq_iff_eq_uncurry_loc (h₁ (s, x)) },
  { refine h₂.mono _, rintro ⟨s, x⟩ hp t, exact curry_eq_iff_eq_uncurry_loc (hp t) },
  { intros s x t, exact (h₄ (s, x) t).le },
  { refine h₅.mono _, rintros ⟨s, x⟩ hp, exact 𝓕.to_family_jet_sec.is_holonomic_at_curry hp }
end

open filter
open_locale unit_interval

/--
A corollary of the local parametric h-principle, forgetting the homotopy and `ε`-closeness,
and just stating the existence of a solution that is holonomic near `K`.
Furthermore, we assume that `P = ℝ` and `K` is of the form `compact set × I`.
This is sufficient to prove sphere eversion. -/
lemma rel_loc.htpy_formal_sol.exists_sol (𝓕₀ : R.htpy_formal_sol)
  (C : set (ℝ × E)) (hC : is_closed C) (K : set E) (hK : is_compact K)
  (h_hol : ∀ᶠ (p : ℝ × E) near C, (𝓕₀ p.1).is_holonomic_at p.2) :
  ∃ f : ℝ → E → F,
    (𝒞 ∞ $ uncurry f) ∧
    (∀ p ∈ C, f (p : ℝ × E).1 p.2 = (𝓕₀ p.1).f p.2) ∧
    (∀ x ∈ K, ∀ t ∈ I, (x, f t x, D (f t) x) ∈ R) :=
begin
  obtain ⟨𝓕, h₁, h₂, -, h₄⟩ :=
    𝓕₀.improve_htpy h_op h_ample zero_lt_one C hC (I ×ˢ K) (is_compact_Icc.prod hK) h_hol,
  refine ⟨λ s, (𝓕 (1, s)).f, _, _, _⟩,
  { exact 𝓕.f_diff.comp ((cont_diff_const.prod cont_diff_id).prod_map cont_diff_id) },
  { intros p hp, exact (prod.ext_iff.mp (h₂.nhds_set_forall_mem p hp 1)).1 },
  { intros x hx t ht,
    rw [show D (𝓕 (1, t)).f x = (𝓕 (1, t)).φ x, from
      h₄.nhds_set_forall_mem (t, x) (mk_mem_prod ht hx)],
    exact 𝓕.is_sol (1, t) x },
end

end parametric_h_principle


import Lib.Data.Array.Instances
import Lib.Data.Fold
import Lib.Data.Nat
import Lib.Data.Quot
import Lib.Data.Traversable
import Lib.Equiv
import Lib.Meta.Dump
import Lib.Tactic

import Advent.IO

section day7

def dist (i j : Nat) : Nat := max (i - j) (j - i)

def cost (input : Array Nat) (i : Nat) : Nat :=
input.foldl (λ acc pos => acc + dist pos i) 0

def resize (ar : Array α) (x : α) (i : Nat) : Array α :=
if i < ar.size
then ar
else ar ++ Array.mkArray (i - ar.size + 1) x

theorem lt_size_resize {ar : Array α} {x : α} {i : Nat} :
  i < (resize ar x i).size := by
simp [resize]; split
. assumption
next h =>
have h := Nat.le_of_not_gt h
simp [*]

def modifyResize
  (ar : Array α) (i : Nat) (f : α → α) (y : α) : Array α :=
let ar := resize ar y i
let i : Fin ar.size := ⟨i, lt_size_resize⟩
ar.set i (f <| ar.get i)

def posToCount (ar : Array Nat) : Array Nat :=
ar.foldl (λ ar x => modifyResize ar x Nat.succ 0) #[]

def linearCost : Fold Nat Nat :=
Prod.snd <$> Fold.mk (0, 0)
  (λ (num, costAccum) n =>
    (num + n, costAccum + num))

def quadCost : Fold Nat Nat :=
(Prod.snd ∘ Prod.snd) <$> Fold.mk (0, 0, 0)
  (λ (num, costAccum, totalCost) n =>
    let num' := num + n
    let costAccum' := costAccum + num'
    let totalCost' := totalCost + costAccum'
    (num', costAccum', totalCost'))

def costUp : Array Nat → Array Nat :=
quadCost.scanl
-- scanl (λ n (num, costAccum) => (costAccum + num, (num + n, costAccum + num))) (0, 0)

def costDown : Array Nat → Array Nat :=
quadCost.scanr
-- scanr (λ n (num, costAccum) => (costAccum + num, (num + n, costAccum + num)))
--   (0, 0)

def cost' (input : Array Nat) : Array Nat :=
  let counts := posToCount input
  let lr := costUp counts
  let rl := costDown counts
  Array.zipWith lr rl (.+.)

def minCost (input : Array Nat) : Option (Nat × Nat) :=
cost' input |>.foldlIdx
  (λ i x => λ
    | Option.some ((a, b) : Nat × Nat) =>
      if x < b then some (i, x)
      else some (a, b)
    | Option.none => some (i, x) )
  none

end day7

def Array.max (ar : Array Nat) : Nat :=
ar.foldl _root_.max 0

def Array.min (ar : Array Nat) : Nat :=
if h : ar.size > 0 then
  ar.foldl _root_.min <| ar.get ⟨_, h⟩
else
  0

namespace Day7

def examples :=
"16,1,2,0,4,2,7,1,2,14"

def parseInput (input : String) : IO (Array Nat) := do
return Array.mk <| (← input.splitOn "," |>.mapM parseNat)

def inputFileName := "Advent/Day7_input.txt"
-- #check elab
-- #check Lean.PrettyPrinter.ppTerm
-- #check Lean.MonadQuotation
-- #print Lean.Syntax
-- #print Lean.Elab.TermElabM
-- #print Lean.MacroM
-- #print Lean.MetaM
-- #check Lean.Syntax

-- #check Lean.MonadRef

def splitCount (ar : Array Nat) : Array <| Array Nat :=
Array.foldlIdx
  ar
  (λ i a acc =>
    if a = 0 then acc
    else acc.push <| Array.mkArray ar.size 0 |>.set! i a)
  #[]

def main : IO Unit := do
let pos ← parseInput <| (← IO.FS.lines inputFileName).get! 0
-- let pos ← parseInput examples
let count := posToCount pos
-- IO.println <| pos.size
IO.println <| dump! minCost pos
-- IO.println <| dump! count
-- IO.println <| dump! splitCount count
-- let count := splitCount count |>.get! 0
-- IO.println <| dump! count
-- IO.println <| posToCount pos
-- IO.println <| dump! costUp count
-- IO.println <| dump! costDown count
-- IO.println <| dump! cost' pos

-- #[1, 2, 3, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 1]

#eval main

end Day7


import hmem.split_cost

universe u
variables {α: Type u} [has_zero α] [decidable_eq α]

namespace hmem
namespace program

def divide_and_conquer_cost (p: program α) (fc: ℕ → ℕ): ℕ → ℕ
| 0 := p.max_internal_cost + fc 0
| (n+1) := p.max_internal_cost + fc (n + 1) + divide_and_conquer_cost n * p.max_recurse_count

theorem divide_and_conquer_cost_def {p: program α} {fc: ℕ → ℕ} {n: ℕ}:
  divide_and_conquer_cost p fc n = p.max_internal_cost + fc n + ite (n = 0) 0 (divide_and_conquer_cost p fc (n - 1) * p.max_recurse_count) :=
by cases n; simp [divide_and_conquer_cost]

theorem divide_and_conquer_cost_mono {p: program α} {fc: ℕ → ℕ} {m n: ℕ}:
  0 < p.max_recurse_count → m ≤ n → divide_and_conquer_cost p fc m ≤ divide_and_conquer_cost p fc n :=
begin
  intros hr hmn,
  induction n,
  { rw [nat.eq_zero_of_le_zero hmn] },
  cases eq_or_lt_of_le hmn,
  { rw [h] },
  exact trans (n_ih (nat.le_of_lt_succ h)) (le_add_left (nat.le_mul_of_pos_right hr)),
end

theorem divide_and_conquer_cost_sound {p: program α}
  (fr: memory α → ℕ → Prop) (hfr: ∀ inp n arg, p.recurse_arg inp arg → fr inp n → ∃ m < n, fr arg m)
  {fc: ℕ → ℕ} (hfc: ∀ inp n, fr inp n → p.call_cost inp (fc n)):
  ∀ inp n, fr inp n → p.has_time_cost inp (divide_and_conquer_cost p fc n) :=
begin
  intros inp n hn,
  cases hrc:p.max_recurse_count,
  { apply thunk.time_cost_of_split',
    apply split_cost_of_components,
    { rcases hfc _ _ hn with ⟨i, r, h⟩,
      exact internal_cost_bound ⟨fc n, r, h⟩ },
    { exact hfc _ _ hn },
    apply recurse_cost_mono,
    apply max_recurse_cost_zero,
    { rcases hfc _ _ hn with ⟨i, r, h⟩,
      exact ⟨_, thunk.time_cost_of_split h⟩ },
    exact hrc,
    apply nat.zero_le 0,
    cases n,
    { refl },
    { unfold divide_and_conquer_cost,
      rw [hrc, mul_zero] } },
  induction n using nat.strong_induction_on with n ih generalizing inp,
  cases n,
  { apply thunk.time_cost_of_split',
    { apply split_cost_of_components,
      { rcases hfc _ _ hn with ⟨i, r, h⟩,
        exact internal_cost_bound ⟨fc 0, r, h⟩ },
      { exact hfc _ _ hn },
      apply max_recurse_cost_zero',
      { rcases hfc _ _ hn with ⟨i, r, h⟩,
        exact ⟨_, thunk.time_cost_of_split h⟩ },
      intros arg harg,
      rcases hfr _ _ _ harg hn with ⟨m, hm, _⟩,
      exact absurd hm (nat.not_lt_zero _),
    },
    refl },
  { apply thunk.time_cost_of_split',
    { apply split_cost_of_components,
      { rcases hfc _ _ hn with ⟨i, r, h⟩,
        exact internal_cost_bound ⟨fc (n + 1), r, h⟩ },
      { exact hfc _ _ hn },
      apply max_recurse_cost,
      { rcases hfc _ _ hn with ⟨i, r, h⟩,
        exact ⟨_, thunk.time_cost_of_split h⟩ },
      intros arg hrec,
      rcases hfr _ (n + 1) arg hrec hn with ⟨n', hn', hfr⟩,
      apply time_cost_mono,
      specialize ih _ hn' _ hfr,
      apply ih,
      apply divide_and_conquer_cost_mono,
      rw [hrc],
      apply nat.zero_lt_succ,
      apply nat.le_of_lt_succ hn' },
    rw [mul_comm],
    refl },
end

end program
end hmem