/-
Copyright (c) 2018 Mario Carneiro. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Mario Carneiro

An experimental variant on the `ring` tactic that uses computational
reflection instead of proof generation. Useful for kernel benchmarking.
-/
import tactic.ring data.num.lemmas
import tactic.converter.interactive

namespace tactic.ring2

@[derive has_reflect]
inductive {u} tree (α : Type u) : Type u
| nil {} : tree
| node : α → tree → tree → tree

def tree.get {α} [has_zero α] : pos_num → tree α → α
| _                tree.nil            := 0
| pos_num.one      (tree.node a t₁ t₂) := a
| (pos_num.bit0 n) (tree.node a t₁ t₂) := t₁.get n
| (pos_num.bit1 n) (tree.node a t₁ t₂) := t₂.get n

def tree.of_rbnode {α} : rbnode α → tree α
| rbnode.leaf               := tree.nil
| (rbnode.red_node l a r)   :=
  tree.node a (tree.of_rbnode l) (tree.of_rbnode r)
| (rbnode.black_node l a r) :=
  tree.node a (tree.of_rbnode l) (tree.of_rbnode r)

def tree.index_of {α} (lt : α → α → Prop) [decidable_rel lt]
  (x : α) : tree α → option pos_num
| tree.nil := none
| (tree.node a t₁ t₂) :=
  match cmp_using lt x a with
  | ordering.lt := pos_num.bit0 <$> tree.index_of t₁
  | ordering.eq := some pos_num.one
  | ordering.gt := pos_num.bit1 <$> tree.index_of t₂
  end

meta def tree.reflect' (u : level) (α : expr) : tree expr → expr
| tree.nil := (expr.const ``tree.nil [u] : expr) α
| (tree.node a t₁ t₂) :=
  (expr.const ``tree.node [u] : expr) α a t₁.reflect' t₂.reflect'

@[derive has_reflect]
inductive csring_expr
| atom : pos_num → csring_expr
| const : num → csring_expr
| add : csring_expr → csring_expr → csring_expr
| mul : csring_expr → csring_expr → csring_expr
| pow : csring_expr → num → csring_expr

namespace csring_expr

def eval {α} [comm_semiring α] (t : tree α) : csring_expr → α
| (atom n)  := @tree.get _ ⟨0⟩ n t
| (const n) := n
| (add x y) := eval x + eval y
| (mul x y) := eval x * eval y
| (pow x n) := eval x ^ (n : ℕ)

end csring_expr

@[derive decidable_eq]
inductive horner_expr
| const : znum → horner_expr
| horner : horner_expr → pos_num → num → horner_expr → horner_expr

namespace horner_expr

def is_cs : horner_expr → Prop
| (const n) := ∃ m:num, n = m.to_znum
| (horner a x n b) := is_cs a ∧ is_cs b

instance : has_zero horner_expr := ⟨const 0⟩
instance : has_one horner_expr := ⟨const 1⟩

def atom (n : pos_num) : horner_expr := horner 1 n 1 0

def repr : horner_expr → string
| (const n) := _root_.repr n
| (horner a x n b) :=
    "(" ++ repr a ++ ") * x" ++ _root_.repr x ++ "^"
        ++ _root_.repr n ++ " + " ++ repr b
instance : has_repr horner_expr := ⟨repr⟩

def horner' (a : horner_expr)
  (x : pos_num) (n : num) (b : horner_expr) : horner_expr :=
match a with
| const q := if q = 0 then b else horner a x n b
| horner a₁ x₁ n₁ b₁ :=
  if x₁ = x ∧ b₁ = 0 then horner a₁ x (n₁ + n) b
  else horner a x n b
end

def add_const (k : znum) (e : horner_expr) : horner_expr :=
if k = 0 then e else begin
  induction e with n a x n b A B,
  { exact const (k + n) },
  { exact horner a x n B }
end

def add_aux (a₁ : horner_expr) (A₁ : horner_expr → horner_expr) (x₁ : pos_num) :
  horner_expr → num → horner_expr → (horner_expr → horner_expr) → horner_expr
| (const n₂)           n₁ b₁ B₁ := add_const n₂ (horner a₁ x₁ n₁ b₁)
| (horner a₂ x₂ n₂ b₂) n₁ b₁ B₁ :=
  let e₂ := horner a₂ x₂ n₂ b₂ in
  match pos_num.cmp x₁ x₂ with
  | ordering.lt := horner a₁ x₁ n₁ (B₁ e₂)
  | ordering.gt := horner a₂ x₂ n₂ (add_aux b₂ n₁ b₁ B₁)
  | ordering.eq :=
    match num.sub' n₁ n₂ with
    | znum.zero := horner' (A₁ a₂) x₁ n₁ (B₁ b₂)
    | (znum.pos k) := horner (add_aux a₂ k 0 id) x₁ n₂ (B₁ b₂)
    | (znum.neg k) := horner (A₁ (horner a₂ x₁ k 0)) x₁ n₁ (B₁ b₂)
    end
  end

def add : horner_expr → horner_expr → horner_expr
| (const n₁)           e₂ := add_const n₁ e₂
| (horner a₁ x₁ n₁ b₁) e₂ := add_aux a₁ (add a₁) x₁ e₂ n₁ b₁ (add b₁)
/-begin
  induction e₁ with n₁ a₁ x₁ n₁ b₁ A₁ B₁ generalizing e₂,
  { exact add_const n₁ e₂ },
  exact match e₂ with e₂ := begin
    induction e₂ with n₂ a₂ x₂ n₂ b₂ A₂ B₂ generalizing n₁ b₁;
    let e₁ := horner a₁ x₁ n₁ b₁,
    { exact add_const n₂ e₁ },
    let e₂ := horner a₂ x₂ n₂ b₂,
    exact match pos_num.cmp x₁ x₂ with
    | ordering.lt := horner a₁ x₁ n₁ (B₁ e₂)
    | ordering.gt := horner a₂ x₂ n₂ (B₂ n₁ b₁)
    | ordering.eq :=
      match num.sub' n₁ n₂ with
      | znum.zero := horner' (A₁ a₂) x₁ n₁ (B₁ b₂)
      | (znum.pos k) := horner (A₂ k 0) x₁ n₂ (B₁ b₂)
      | (znum.neg k) := horner (A₁ (horner a₂ x₁ k 0)) x₁ n₁ (B₁ b₂)
      end
    end
  end end
end-/

def neg (e : horner_expr) : horner_expr :=
begin
  induction e with n a x n b A B,
  { exact const (-n) },
  { exact horner A x n B }
end

def mul_const (k : znum) (e : horner_expr) : horner_expr :=
if k = 0 then 0 else if k = 1 then e else begin
  induction e with n a x n b A B,
  { exact const (n * k) },
  { exact horner A x n B }
end

def mul_aux (a₁ x₁ n₁ b₁) (A₁ B₁ : horner_expr → horner_expr) :
  horner_expr → horner_expr
| (const n₂) := mul_const n₂ (horner a₁ x₁ n₁ b₁)
| e₂@(horner a₂ x₂ n₂ b₂) :=
  match pos_num.cmp x₁ x₂ with
  | ordering.lt := horner (A₁ e₂) x₁ n₁ (B₁ e₂)
  | ordering.gt := horner (mul_aux a₂) x₂ n₂ (mul_aux b₂)
  | ordering.eq := let haa := horner' (mul_aux a₂) x₁ n₂ 0 in
    if b₂ = 0 then haa else haa.add (horner (A₁ b₂) x₁ n₁ (B₁ b₂))
  end

def mul : horner_expr → horner_expr → horner_expr
| (const n₁)              := mul_const n₁
| e₁@(horner a₁ x₁ n₁ b₁) := mul_aux a₁ x₁ n₁ b₁ (mul a₁) (mul b₁).
/-begin
  induction e₁ with n₁ a₁ x₁ n₁ b₁ A₁ B₁ generalizing e₂,
  { exact mul_const n₁ e₂ },
  induction e₂ with n₂ a₂ x₂ n₂ b₂ A₂ B₂;
  let e₁ := horner a₁ x₁ n₁ b₁,
  { exact mul_const n₂ e₁ },
  let e₂ := horner a₂ x₂ n₂ b₂,
  cases pos_num.cmp x₁ x₂,
  { exact horner (A₁ e₂) x₁ n₁ (B₁ e₂) },
  { let haa := horner' A₂ x₁ n₂ 0,
    exact if b₂ = 0 then haa else
      haa.add (horner (A₁ b₂) x₁ n₁ (B₁ b₂)) },
  { exact horner A₂ x₂ n₂ B₂ }
end-/

instance : has_add horner_expr := ⟨add⟩
instance : has_neg horner_expr := ⟨neg⟩
instance : has_mul horner_expr := ⟨mul⟩

def pow (e : horner_expr) : num → horner_expr
| 0 := 1
| (num.pos p) := begin
  induction p with p ep p ep,
  { exact e },
  { exact (ep.mul ep).mul e },
  { exact ep.mul ep }
end

def inv (e : horner_expr) : horner_expr := 0

def of_csexpr : csring_expr → horner_expr
| (csring_expr.atom n)  := atom n
| (csring_expr.const n) := const n.to_znum
| (csring_expr.add x y) := (of_csexpr x).add (of_csexpr y)
| (csring_expr.mul x y) := (of_csexpr x).mul (of_csexpr y)
| (csring_expr.pow x n) := (of_csexpr x).pow n

def cseval {α} [comm_semiring α] (t : tree α) : horner_expr → α
| (const n)        := n.abs
| (horner a x n b) := tactic.ring.horner (cseval a) (t.get x) n (cseval b)

theorem cseval_atom {α} [comm_semiring α] (t : tree α)
  (n : pos_num) : (atom n).is_cs ∧ cseval t (atom n) = t.get n :=
⟨⟨⟨1, rfl⟩, ⟨0, rfl⟩⟩, (tactic.ring.horner_atom _).symm⟩

theorem cseval_add_const {α} [comm_semiring α] (t : tree α)
  (k : num) {e : horner_expr} (cs : e.is_cs) :
  (add_const k.to_znum e).is_cs ∧
    cseval t (add_const k.to_znum e) = k + cseval t e :=
begin
  simp [add_const],
  cases k; simp! *,
  simp [show znum.pos k ≠ 0, from dec_trivial],
  induction e with n a x n b A B; simp *,
  { rcases cs with ⟨n, rfl⟩,
    refine ⟨⟨n + num.pos k, by simp; refl⟩, _⟩,
    cases n; simp! },
  { rcases B cs.2 with ⟨csb, h⟩, simp! [*, cs.1],
    rw [← tactic.ring.horner_add_const, add_comm], refl }
end

theorem cseval_horner' {α} [comm_semiring α] (t : tree α)
  (a x n b) (h₁ : is_cs a) (h₂ : is_cs b) :
  (horner' a x n b).is_cs ∧ cseval t (horner' a x n b) =
    tactic.ring.horner (cseval t a) (t.get x) n (cseval t b) :=
begin
  cases a with n₁ a₁ x₁ n₁ b₁; simp [horner']; split_ifs,
  { simp! [*, tactic.ring.horner] },
  { exact ⟨⟨h₁, h₂⟩, rfl⟩ },
  { refine ⟨⟨h₁.1, h₂⟩, eq.symm _⟩, simp! *,
    apply tactic.ring.horner_horner, simp },
  { exact ⟨⟨h₁, h₂⟩, rfl⟩ }
end

theorem cseval_add {α} [comm_semiring α] (t : tree α)
  {e₁ e₂ : horner_expr} (cs₁ : e₁.is_cs) (cs₂ : e₂.is_cs) :
  (add e₁ e₂).is_cs ∧
  cseval t (add e₁ e₂) = cseval t e₁ + cseval t e₂ :=
begin
  induction e₁ with n₁ a₁ x₁ n₁ b₁ A₁ B₁ generalizing e₂; simp!,
  { rcases cs₁ with ⟨n₁, rfl⟩,
    simpa using cseval_add_const t n₁ cs₂ },
  induction e₂ with n₂ a₂ x₂ n₂ b₂ A₂ B₂ generalizing n₁ b₁,
  { rcases cs₂ with ⟨n₂, rfl⟩,
    simp! [cseval_add_const t n₂ cs₁] },
  cases cs₁ with csa₁ csb₁, cases id cs₂ with csa₂ csb₂,
  simp!, have C := pos_num.cmp_to_nat x₁ x₂,
  cases pos_num.cmp x₁ x₂; simp!,
  { rcases B₁ csb₁ cs₂ with ⟨csh, h⟩,
    refine ⟨⟨csa₁, csh⟩, eq.symm _⟩,
    apply tactic.ring.horner_add_const,
    exact h.symm },
  { cases C,
    have B0 : is_cs 0 → ∀ {e₂ : horner_expr}, is_cs e₂ →
      is_cs (add 0 e₂) ∧ cseval t (add 0 e₂) = cseval t 0 + cseval t e₂ :=
      λ _ e₂ c, ⟨c, (zero_add _).symm⟩,
     cases e : num.sub' n₁ n₂ with k k; simp!,
    { have : n₁ = n₂,
      { have := congr_arg (coe : znum → ℤ) e,
        simp at this,
        have := sub_eq_zero.1 this,
        rw [← num.to_nat_to_int, ← num.to_nat_to_int] at this,
        exact num.to_nat_inj.1 (int.coe_nat_inj this) },
      subst n₂,
      rcases cseval_horner' _ _ _ _ _ _ _ with ⟨csh, h⟩,
      { refine ⟨csh, h.trans (eq.symm _)⟩,
        simp *,
        apply tactic.ring.horner_add_horner_eq; try {refl},
        simp },
      all_goals {simp! *} },
    { simp [B₁ csb₁ csb₂],
      rcases A₂ csa₂ _ _ B0 ⟨csa₁, 0, rfl⟩ with ⟨csh, h⟩,
      refine ⟨csh, eq.symm _⟩,
      rw [show id = add 0, from rfl, h],
      apply tactic.ring.horner_add_horner_gt,
      { change (_ + k : ℕ) = _,
        rw [← int.coe_nat_inj', int.coe_nat_add,
          eq_comm, ← sub_eq_iff_eq_add'],
        simpa using congr_arg (coe : znum → ℤ) e },
      all_goals { apply add_comm } },
    { have : (horner a₂ x₁ (num.pos k) 0).is_cs := ⟨csa₂, 0, rfl⟩,
      simp [B₁ csb₁ csb₂, A₁ csa₁ this],
      symmetry, apply tactic.ring.horner_add_horner_lt,
      { change (_ + k : ℕ) = _,
          rw [← int.coe_nat_inj', int.coe_nat_add,
            eq_comm, ← sub_eq_iff_eq_add', ← neg_inj', neg_sub],
        simpa using congr_arg (coe : znum → ℤ) e },
      { refl }, { apply add_comm } } },
  { rcases B₂ csb₂ _ _ B₁ ⟨csa₁, csb₁⟩ with ⟨csh, h⟩,
    refine ⟨⟨csa₂, csh⟩, eq.symm _⟩,
    apply tactic.ring.const_add_horner,
    simp [h] }
end

theorem cseval_mul {α} [comm_semiring α] (t : tree α)
  {e₁ e₂ : horner_expr} (cs₁ : e₁.is_cs) (cs₂ : e₂.is_cs) :
  (mul e₁ e₂).is_cs ∧
  cseval t (mul e₁ e₂) = cseval t e₁ * cseval t e₂ :=
begin
  induction e₁ with n₁ a₁ x₁ n₁ b₁ A₁ B₁ generalizing e₂; simp!,
  { rcases cs₁ with ⟨n₁, rfl⟩,
    simpa [mul_comm] using cseval_mul_const t n₁ cs₂ },
  induction e₂ with n₂ a₂ x₂ n₂ b₂ A₂ B₂,
  { rcases cs₂ with ⟨n₂, rfl⟩,
    simpa! using cseval_mul_const t n₂ cs₁ },
  cases cs₁ with csa₁ csb₁, cases id cs₂ with csa₂ csb₂,
  simp!, have C := pos_num.cmp_to_nat x₁ x₂,
  cases A₂ csa₂ with csA₂ hA₂,
  cases pos_num.cmp x₁ x₂; simp!,
  { simp [A₁ csa₁ cs₂, B₁ csb₁ cs₂],
    symmetry, apply tactic.ring.horner_mul_const; refl },
  { cases cseval_horner' t _ x₁ n₂ 0 csA₂ ⟨0, rfl⟩ with csh₁ h₁,
    cases C, split_ifs,
    { subst b₂,
      refine ⟨csh₁, h₁.trans (eq.symm _)⟩,
      apply tactic.ring.horner_mul_horner_zero; try {refl},
      simp! [hA₂] },
    { cases A₁ csa₁ csb₂ with csA₁ hA₁,
      cases cseval_add t csh₁ _ with csh₂ h₂,
      { refine ⟨csh₂, h₂.trans (eq.symm _)⟩,
        apply tactic.ring.horner_mul_horner; try {refl},
        simp! * },
      exact ⟨csA₁, (B₁ csb₁ csb₂).1⟩ } },
  { simp [A₂ csa₂, B₂ csb₂], rw [mul_comm, eq_comm],
    apply tactic.ring.horner_const_mul,
    {apply mul_comm}, {refl} },
end

theorem cseval_pow {α} [comm_semiring α] (t : tree α)
  {x : horner_expr} (cs : x.is_cs) :
  ∀ (n : num), (pow x n).is_cs ∧
    cseval t (pow x n) = cseval t x ^ (n : ℕ)
| 0 := ⟨⟨1, rfl⟩, (pow_zero _).symm⟩
| (num.pos p) := begin
  simp [pow], induction p with p ep p ep,
  { simp * },
  { simp [pow_bit1],
    cases cseval_mul t ep.1 ep.1 with cs₀ h₀,
    cases cseval_mul t cs₀ cs with cs₁ h₁,
    simp * },
  { simp [pow_bit0],
    cases cseval_mul t ep.1 ep.1 with cs₀ h₀,
    simp * }
end

theorem cseval_of_csexpr {α} [comm_semiring α] (t : tree α) :
  ∀ (r : csring_expr), (of_csexpr r).is_cs ∧ cseval t (of_csexpr r) = r.eval t
| (csring_expr.atom n)  := cseval_atom _ _
| (csring_expr.const n) := ⟨⟨n, rfl⟩, by cases n; refl⟩
| (csring_expr.add x y) :=
  let ⟨cs₁, h₁⟩ := cseval_of_csexpr x,
      ⟨cs₂, h₂⟩ := cseval_of_csexpr y,
      ⟨cs, h⟩ := cseval_add t cs₁ cs₂ in ⟨cs, by simp! [h, *]⟩
| (csring_expr.mul x y) :=
  let ⟨cs₁, h₁⟩ := cseval_of_csexpr x,
      ⟨cs₂, h₂⟩ := cseval_of_csexpr y,
      ⟨cs, h⟩ := cseval_mul t cs₁ cs₂ in ⟨cs, by simp! [h, *]⟩
| (csring_expr.pow x n) :=
  let ⟨cs, h⟩ := cseval_of_csexpr x,
      ⟨cs, h⟩ := cseval_pow t cs n in ⟨cs, by simp! [h, *]⟩

end horner_expr

theorem correctness {α} [comm_semiring α] (t : tree α) (r₁ r₂ : csring_expr)
  (H : horner_expr.of_csexpr r₁ = horner_expr.of_csexpr r₂) :
  r₁.eval t = r₂.eval t :=
by repeat {rw ← (horner_expr.cseval_of_csexpr t _).2}; rw H

meta def reflect_expr : expr → csring_expr × dlist expr
| `(%%e₁ + %%e₂) :=
  let (r₁, l₁) := reflect_expr e₁, (r₂, l₂) := reflect_expr e₂ in
  (r₁.add r₂, l₁ ++ l₂)
/-| `(%%e₁ - %%e₂) :=
  let (r₁, l₁) := reflect_expr e₁, (r₂, l₂) := reflect_expr e₂ in
  (r₁.add r₂.neg, l₁ ++ l₂)
| `(- %%e) := let (r, l) := reflect_expr e in (r.neg, l)-/
| `(%%e₁ * %%e₂) :=
  let (r₁, l₁) := reflect_expr e₁, (r₂, l₂) := reflect_expr e₂ in
  (r₁.mul r₂, l₁ ++ l₂)
/-| `(has_inv.inv %%e) := let (r, l) := reflect_expr e in (r.neg, l)
| `(%%e₁ / %%e₂) :=
  let (r₁, l₁) := reflect_expr e₁, (r₂, l₂) := reflect_expr e₂ in
  (r₁.mul r₂.inv, l₁ ++ l₂)-/
| e@`(%%e₁ ^ %%e₂) :=
  match reflect_expr e₁, expr.to_nat e₂ with
  | (r₁, l₁), some n₂ := (r₁.pow (num.of_nat' n₂), l₁)
  | (r₁, l₁), none := (csring_expr.atom 1, dlist.singleton e)
  end
| e := match expr.to_nat e with
  | some n := (csring_expr.const (num.of_nat' n), dlist.empty)
  | none := (csring_expr.atom 1, dlist.singleton e)
  end

meta def csring_expr.replace (t : tree expr) : csring_expr → state_t (list expr) option csring_expr
| (csring_expr.atom _)  := do e ← get,
  p ← monad_lift (t.index_of (<) e.head),
  put e.tail, pure (csring_expr.atom p)
| (csring_expr.const n) := pure (csring_expr.const n)
| (csring_expr.add x y) := csring_expr.add <$> x.replace <*> y.replace
| (csring_expr.mul x y) := csring_expr.mul <$> x.replace <*> y.replace
| (csring_expr.pow x n) := (λ x, csring_expr.pow x n) <$> x.replace
--| (csring_expr.neg x)   := csring_expr.neg <$> x.replace
--| (csring_expr.inv x)   := csring_expr.inv <$> x.replace

end tactic.ring2

namespace tactic
namespace interactive
open interactive interactive.types lean.parser
open tactic.ring2

local postfix `?`:9001 := optional

/-- Tactic for solving equations in the language of rings.
  This variant on the `ring` tactic uses kernel computation instead
  of proof generation. -/
meta def ring2 : tactic unit :=
do `[repeat {rw ← nat.pow_eq_pow}],
  `(%%e₁ = %%e₂) ← target,
  α ← infer_type e₁,
  expr.sort (level.succ u) ← infer_type α,
  let (r₁, l₁) := reflect_expr e₁,
  let (r₂, l₂) := reflect_expr e₂,
  let L := (l₁ ++ l₂).to_list,
  let s := tree.of_rbnode (rbtree_of L).1,
  (r₁, L) ← (state_t.run (r₁.replace s) L : option _),
  (r₂, _) ← (state_t.run (r₂.replace s) L : option _),
  let se : expr := s.reflect' u α,
  let er₁ : expr := reflect r₁,
  let er₂ : expr := reflect r₂,
  cs ← mk_app ``comm_semiring [α] >>= mk_instance,
  e ← to_expr ```(correctness %%se %%er₁ %%er₂ rfl)
    <|> fail ("ring2 tactic failed, using abstract equality:\n"
      ++ repr (horner_expr.of_csexpr r₁) ++
      "\n  =?=\n" ++ repr (horner_expr.of_csexpr r₂)),
  tactic.exact e

end interactive
end tactic

namespace conv.interactive
open conv

meta def ring2 : conv unit := discharge_eq_lhs tactic.interactive.ring2

end conv.interactive


def casesTFOn {motive : Prop → Sort _} (P) [inst : Decidable P] : (T : motive True) → (F : motive False) → motive P :=
  λ ht hf => match inst with
  | isTrue H => eqTrue H ▸ ht
  | isFalse H => eqFalse H ▸ hf

example (P) [Decidable P] : ¬¬P → P := by
  induction P using casesTFOn
  admit
  admit


/-
Copyright (c) 2017 Daniel Selsam. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author: Daniel Selsam

Certified graph transformation that integrates out a specific KL divergence term.
-/
import .util .tensor .compute_grad .graph .tactics .predicates .lemmas .env .expected_value .ops

namespace certigrad
open list

def integrate_mvn_kl (eloss : ID) : list node → list node
| [] := []

| (⟨(z, .(shape)), [(μ, .(shape)), (σ, .(shape))], operator.rand (rand.op.mvn shape)⟩
 ::⟨(el, []), [(μ', .(shape')), (σ', .(shape')), (z', .(shape'))], operator.det (det.op.mvn_empirical_kl shape')⟩
 ::nodes) :=
 (⟨(el, []), [⟨μ, shape⟩, ⟨σ, shape⟩], operator.det (ops.mvn_kl shape)⟩
::⟨(z, shape), [⟨μ, shape⟩, ⟨σ, shape⟩], operator.rand (rand.op.mvn shape)⟩
::nodes)

| (⟨(z, .(shape)), [(μ, .(shape)), (σ, .(shape))], operator.rand (rand.op.mvn shape)⟩ :: nodes) :=
⟨(z, shape), [(μ, shape), (σ, shape)], operator.rand (rand.op.mvn shape)⟩ :: integrate_mvn_kl nodes

| (n::nodes) := n :: integrate_mvn_kl nodes

def integrate_mvn_kl_pre (eloss : ID) : list node → env → Prop
-- EQ1
| [] _ := true
-- EQ2 (a)
| (⟨(rname, rshape), [], operator.det op⟩::nodes) inputs :=
integrate_mvn_kl_pre nodes (env.insert (rname, rshape) (op^.f (env.get_ks [] inputs)) inputs)
-- EQ2 (b)
| (⟨(rname, rshape), [], operator.rand op⟩::nodes) inputs := false
-- EQ3 (a)
| (⟨(rname, rshape), [(pname, pshape)], operator.det op⟩::nodes) inputs :=
integrate_mvn_kl_pre nodes (env.insert (rname, rshape) (op^.f (env.get_ks [(pname, pshape)] inputs)) inputs)
-- EQ3 (b)
| (⟨(rname, rshape), [(pname, pshape)], operator.rand op⟩::nodes) inputs := false
-- EQ4
| (⟨(rname, rshape), [(pname₁, pshape₁), (pname₂, pshape₂)], operator.det op⟩::nodes) inputs :=
integrate_mvn_kl_pre nodes (env.insert (rname, rshape) (op^.f (env.get_ks [(pname₁, pshape₁), (pname₂, pshape₂)] inputs)) inputs)
-- EQ5
| [⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩] inputs := false
-- EQ6
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, []), [], op⟩::nodes) inputs := false
-- EQ7
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, []), [(pname₃, shape₃)], op⟩::nodes) inputs := false
-- EQ8
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, []), [(pname₃, shape₃), (pname₄, shape₄)], op⟩::nodes) inputs := false
-- EQ9
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, []), [(pname₃, shape₃), (pname₄, shape₄), (pname₅, shape₅)], operator.det (det.op.mk _ _ _ _ _ _ _)⟩::nodes) inputs := false
-- EQ10
| (⟨(z, .(shape)), [(μ, .(shape)), (σ, .(shape))], operator.rand (rand.op.mvn shape)⟩
 ::⟨(el, []), [(μ', .(shape')), (σ', .(shape')), (z', .(shape'))], operator.det (det.op.mvn_empirical_kl shape')⟩
 ::nodes) inputs :=
(μ = μ' ∧ σ = σ' ∧ z = z' ∧ shape = shape' ∧ eloss = el ∧ σ ≠ μ)
∧ ((¬ env.has_key (eloss, []) inputs) ∧ (eloss, []) ∉ (z, shape) :: map node.ref nodes ∧ 0 < env.get (σ, shape) inputs ∧ ∀ (y : T shape), all_parents_in_env (env.insert (z, shape) y inputs) nodes)
-- EQ11
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, []), [(pname₃, shape₃), (pname₄, shape₄), (pname₅, shape₅)], operator.rand op⟩::nodes) inputs := false
-- EQ12
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, []), ((pname₃, shape₃)::(pname₄, shape₄)::(pname₅, shape₅)::parent::parents), op⟩::nodes) inputs := false
-- EQ13
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, (shape₂ :: shapes₂)), parents, op⟩::nodes) inputs := false
-- EQ14 (a)
| (⟨(rname₂, shape₂), ((pname₃, shape₃)::(pname₄, shape₄)::parent::parents), operator.det op⟩::nodes) inputs :=
integrate_mvn_kl_pre nodes (env.insert (rname₂, shape₂) (op^.f $ env.get_ks ((pname₃, shape₃)::(pname₄, shape₄)::parent::parents) inputs) inputs)
-- EQ14 (b)
| (⟨(rname₂, shape₂), ((pname₃, shape₃)::(pname₄, shape₄)::parent::parents), operator.rand op⟩::nodes) inputs := false

theorem integrate_mvn_kl_correct (eloss : ID) (costs : list ID) :
∀ (nodes : list node) (inputs : env),
  eloss ∉ costs →
  integrate_mvn_kl_pre eloss nodes inputs →
  uniq_ids nodes inputs →
  all_parents_in_env inputs nodes →
  pdfs_exist_at nodes inputs →
  is_gintegrable (λ m, ⟦sum_costs m (eloss::costs)⟧) inputs (integrate_mvn_kl eloss nodes) dvec.head →
  is_gintegrable (λ m, ⟦sum_costs m (eloss::costs)⟧) inputs nodes dvec.head →
E (graph.to_dist (λ env₀, ⟦sum_costs env₀ (eloss::costs)⟧) inputs (integrate_mvn_kl eloss nodes)) dvec.head
=
E (graph.to_dist (λ env₀, ⟦sum_costs env₀ (eloss::costs)⟧) inputs nodes) dvec.head

-- EQ1
| [] _ _ _ _ _ _ _ _ := rfl

-- EQ2 (a)
| (⟨(rname, rshape), [], operator.det op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint :=
begin
dunfold graph.to_dist operator.to_dist integrate_mvn_kl,
dunfold integrate_mvn_kl_pre at H_pre,
simp [E.E_bind, E.E_ret],
assertv H_pre_next : integrate_mvn_kl_pre eloss nodes (env.insert (rname, rshape) (op^.f (env.get_ks nil inputs)) inputs) := H_pre,
assertv H_ps_in_env_next : all_parents_in_env (env.insert (rname, rshape) (op^.f (env.get_ks nil inputs)) inputs) nodes := H_ps_in_env^.right _,
exact (integrate_mvn_kl_correct _ _ H_eloss_not_cost H_pre_next (H_uids^.right _) H_ps_in_env_next H_pdfs_exist_at H_kl_gint H_gint)
end

-- EQ2 (b)
| (⟨(rname, rshape), [], operator.rand op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint := false.rec _ H_pre

-- EQ3 (a)
| (⟨(rname, rshape), [(pname, pshape)], operator.det op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint :=
begin
dunfold graph.to_dist operator.to_dist integrate_mvn_kl,
dunfold integrate_mvn_kl_pre at H_pre,
simp [E.E_bind, E.E_ret],
assertv H_pre_next : integrate_mvn_kl_pre eloss nodes (env.insert (rname, rshape) (op^.f (env.get_ks [(pname, pshape)] inputs)) inputs) := H_pre,
assertv H_ps_in_env_next : all_parents_in_env (env.insert (rname, rshape) (op^.f (env.get_ks [(pname, pshape)] inputs)) inputs) nodes := H_ps_in_env^.right _,
exact (integrate_mvn_kl_correct _ _ H_eloss_not_cost H_pre_next (H_uids^.right _) H_ps_in_env_next H_pdfs_exist_at H_kl_gint H_gint)
end

-- EQ3 (b)
| (⟨(rname, rshape), [(pname, pshape)], operator.rand op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint := false.rec _ H_pre

-- EQ4
| (⟨(rname, rshape), [(pname₁, pshape₁), (pname₂, pshape₂)], operator.det op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint :=
begin
dunfold graph.to_dist operator.to_dist integrate_mvn_kl,
dunfold integrate_mvn_kl_pre at H_pre,
simp [E.E_bind, E.E_ret],
assertv H_pre_next : integrate_mvn_kl_pre eloss nodes (env.insert (rname, rshape) (op^.f (env.get_ks [(pname₁, pshape₁), (pname₂, pshape₂)] inputs)) inputs) := H_pre,
assertv H_ps_in_env_next : all_parents_in_env (env.insert (rname, rshape) (op^.f (env.get_ks [(pname₁, pshape₁), (pname₂, pshape₂)] inputs)) inputs) nodes := H_ps_in_env^.right _,
exact (integrate_mvn_kl_correct _ _ H_eloss_not_cost H_pre_next (H_uids^.right _) H_ps_in_env_next H_pdfs_exist_at H_kl_gint H_gint)
end

-- EQ5
| [⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩] inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint := false.rec _ H_pre

-- EQ6
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, []), [], op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint := false.rec _ H_pre

-- EQ7
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, []), [(pname₃, shape₃)], op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint := false.rec _ H_pre

-- EQ8
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, []), [(pname₃, shape₃), (pname₄, shape₄)], op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint := false.rec _ H_pre

-- EQ9
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, []), [(pname₃, shape₃), (pname₄, shape₄), (pname₅, shape₅)], operator.det (det.op.mk _ _ _ _ _ _ _)⟩::nodes) inputs
H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint := false.rec _ H_pre

-- EQ10
| (⟨(z, .(shape)), [(μ, .(shape)), (σ, .(shape))], operator.rand (rand.op.mvn shape)⟩
 ::⟨(el, []), [(μ', .(shape')), (σ', .(shape')), (z', .(shape'))], operator.det (det.op.mvn_empirical_kl shape')⟩
 ::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint :=
begin
assertv H_ok : μ = μ' ∧ σ = σ' ∧ z = z' ∧ shape = shape' ∧ eloss = el ∧ σ ≠ μ := H_pre^.left,

cases H_ok with H_μ H_ok, subst H_μ,
cases H_ok with H_σ H_ok, subst H_σ,
cases H_ok with H_z H_ok, subst H_z,
cases H_ok with H_shape H_ok, subst H_shape,
cases H_ok with H_eloss_eq_el H_σ_neq_μ, subst H_eloss_eq_el,

dunfold graph.to_dist operator.to_dist dvec.head integrate_mvn_kl,
simp [E.E_bind, E.E_ret],
dunfold dvec.head, dsimp,

assertv H_μ_in : env.has_key (μ, shape) inputs := H_ps_in_env^.left (μ, shape) (mem_cons_self _ _),
assertv H_σ_in : env.has_key (σ, shape) inputs := H_ps_in_env^.left (σ, shape) (mem_cons_of_mem _ (mem_cons_self _ _)),
assertv H_z_nin : ¬ env.has_key (z, shape) inputs := H_uids^.left,
assertv H_eloss_nin : ¬ env.has_key (eloss, []) inputs := H_pre^.right^.left,

assertv H_μ_neq_z : (μ, shape) ≠ (z, shape) := env_in_nin_ne H_μ_in H_z_nin,
assertv H_σ_neq_z : (σ, shape) ≠ (z, shape) := env_in_nin_ne H_σ_in H_z_nin,
assertv H_μ_neq_eloss : (μ, shape) ≠ (eloss, []) := env_in_nin_ne H_μ_in H_eloss_nin,
assertv H_σ_neq_eloss : (σ, shape) ≠ (eloss, []) := env_in_nin_ne H_σ_in H_eloss_nin,

assertv H_eloss_neq_z : (eloss, []) ≠ (z, shape) := ne_of_not_mem_cons H_pre^.right^.right^.left,
assertv H_eloss_nin_nodes : (eloss, []) ∉ map node.ref nodes := not_mem_of_not_mem_cons H_pre^.right^.right^.left,

dunfold env.get_ks,
tactic.dget_dinsert,
dunfold det.op.f ops.mvn_kl ops.mvn_kl.f dvec.head dvec.head2 dvec.head3,
-- TODO(dhs): unbelievable that this won't unfold det.op.f
-- I tried simp as well
-- I copy-pasted to a new file (to avoid the 45 second delay) and all variants worked
dsimp [det.op.f, ops.mvn_kl.f, dvec.head, dvec.head2, dvec.head3],

-- TODO(dhs): I think this is a bug in LEAN. (z, shape) and (eloss, []) are in the type of the last argument.
simp only [λ (x : dvec T [shape]),
               @env.insert_insert_flip (z, shape) (eloss, []) x^.head
                                       (T.mvn_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape))
                                       inputs (ne.symm H_eloss_neq_z)],

-- Two steps:
-- 1. Split out eloss, and prove that it equals the lookup val
-- 2. For RHS, use eloss not appearing to remove them both
dunfold det.op.f ops.mvn_kl ops.mvn_kl.f dvec.head dvec.head2 dvec.head3,
dunfold sum_costs map sumr,

-- Step 1

definev k₁ : env → ℝ := λ (m : env), env.get (eloss, @nil ℕ) m,
definev k₂ : env → ℝ := λ (m : env), sumr (map (λ (cost : ID), env.get (cost, @nil ℕ) m) costs),

definev m_lhs_k_add : dvec T [shape] → env := λ (x : dvec T [shape]), env.insert (eloss, @nil ℕ) (T.mvn_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape))
                                                                             (env.insert (z, shape) (dvec.head x) inputs),

assert H_lhs_kint₁ : ∀ (x : dvec T [shape]), is_gintegrable (λ m, ⟦k₁ m⟧) (m_lhs_k_add x) nodes dvec.head,
 { dsimp [is_gintegrable, integrate_mvn_kl, dvec.head] at H_kl_gint, dsimp, intro x,
   cases x with xx x xxx xnil, cases xnil,
   simp only [λ a1 a2 a3, @env.insert_insert_flip (eloss, []) (z, shape) a1 a2 a3 H_eloss_neq_z],
   simp only [det.op.f, ops.mvn_kl, ops.mvn_kl.f, dvec.head, env.get_ks, sum_costs] at H_kl_gint,
   dunfold sumr map at H_kl_gint,
   exact (iff.mpr (is_gintegrable_k_add _ _ _ _) (H_kl_gint^.right x))^.left
  },

assert H_lhs_kint₂ : ∀ (x : dvec T [shape]), is_gintegrable (λ m, ⟦k₂ m⟧) (m_lhs_k_add x) nodes dvec.head,
 { dsimp [is_gintegrable, integrate_mvn_kl, dvec.head] at H_kl_gint, dsimp, intro x,
   cases x with xx x xxx xnil, cases xnil,
   simp only [λ a1 a2 a3, @env.insert_insert_flip (eloss, []) (z, shape) a1 a2 a3 H_eloss_neq_z],
   simp only [det.op.f, ops.mvn_kl, ops.mvn_kl.f, dvec.head, env.get_ks, sum_costs] at H_kl_gint,
   dunfold sumr map at H_kl_gint,
   exact (iff.mpr (is_gintegrable_k_add _ _ _ _) (H_kl_gint^.right x))^.right },

simp only [(λ (x : dvec T [shape]), E.E_k_add k₁ k₂ (m_lhs_k_add x) nodes (H_lhs_kint₁ x) (H_lhs_kint₂ x))],

clear H_lhs_kint₁ H_lhs_kint₂,

definev m_rhs_k_add : dvec T [shape] → env := λ x, env.insert (eloss, @nil ℕ)
                                                              (T.mvn_empirical_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape) (dvec.head x))
                                                              (env.insert (z, shape) (dvec.head x) inputs),

assert H_rhs_kint₁ : ∀ (x : dvec T [shape]), is_gintegrable (λ m, ⟦k₁ m⟧) (m_rhs_k_add x) nodes dvec.head,
 { dsimp [is_gintegrable, integrate_mvn_kl, dvec.head] at H_gint, dsimp, intro x,
   cases x with xx x xxx xnil, cases xnil,
   simp only [det.op.f, dvec.head, env.get_ks, sum_costs] at H_gint,
   tactic.dget_dinsert_at `H_gint,
   dunfold sumr map dvec.head dvec.head2 dvec.head3 at H_gint,
   exact (iff.mpr (is_gintegrable_k_add _ _ _ _) (H_gint^.right x))^.left },

assert H_rhs_kint₂ : ∀ (x : dvec T [shape]), is_gintegrable (λ m, ⟦k₂ m⟧) (m_rhs_k_add x) nodes dvec.head,
 { dsimp [is_gintegrable, integrate_mvn_kl, dvec.head] at H_gint, dsimp, intro x,
   cases x with xx x xxx xnil, cases xnil,
   simp only [det.op.f, dvec.head, env.get_ks, sum_costs] at H_gint,
   tactic.dget_dinsert_at `H_gint,
   dunfold sumr map dvec.head dvec.head2 dvec.head3 at H_gint,
   exact (iff.mpr (is_gintegrable_k_add _ _ _ _) (H_gint^.right x))^.right },

simp only [(λ (x : dvec T [shape]), E.E_k_add k₁ k₂ (m_rhs_k_add x) nodes (H_rhs_kint₁ x) (H_rhs_kint₂ x))],

clear H_rhs_kint₁ H_rhs_kint₂,

pose d_base := sprog.prim (rand.op.mvn shape) ⟦env.get (μ, shape) inputs, env.get (σ, shape) inputs⟧,
pose lhs_f₁ := λ x, E (graph.to_dist (λ (m : env), ⟦k₁ m⟧) (m_lhs_k_add x) nodes) dvec.head,
pose lhs_f₂ := λ x, E (graph.to_dist (λ (m : env), ⟦k₂ m⟧) (m_lhs_k_add x) nodes) dvec.head,

assert H_E_kl_add :
∀ x, E (graph.to_dist (λ (m : env),
                         ⟦env.get (eloss, @nil ℕ) m + sumr (map (λ (cost : ID), (env.get (cost, @nil ℕ) m : ℝ)) costs)⟧)
            (env.insert (z, shape) x
                        (env.insert (eloss, @nil ℕ) (T.mvn_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape)) inputs))
            nodes)
       dvec.head
=
E (graph.to_dist (λ (m : env), ⟦k₁ m⟧)
            (env.insert (z, shape) x
                        (env.insert (eloss, @nil ℕ) (T.mvn_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape)) inputs))
            nodes)
       dvec.head
+
E (graph.to_dist (λ (m : env), ⟦k₂ m⟧)
            (env.insert (z, shape) x
                        (env.insert (eloss, @nil ℕ) (T.mvn_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape)) inputs))
            nodes)
       dvec.head,
{ intro x,
  exact E.E_k_add k₁ k₂ _ _ (iff.mpr (is_gintegrable_k_add _ _ _ _) (H_kl_gint^.right x))^.left (iff.mpr (is_gintegrable_k_add _ _ _ _) (H_kl_gint^.right x))^.right },

assert H_lhs_eint₁ : E.is_eintegrable d_base lhs_f₁,
{
dsimp [E.is_eintegrable, dvec.head],
dsimp [integrate_mvn_kl, is_gintegrable] at H_kl_gint,
simp only [λ a1 a2 a3, @env.insert_insert_flip (eloss, []) (z, shape) a1 a2 a3 H_eloss_neq_z],
simp only [det.op.f, ops.mvn_kl, p1, p2, ops.mvn_kl.f, dvec.head, env.get_ks, sum_costs] at H_kl_gint,
dunfold det.op.f ops.mvn_kl p1 p2 ops.mvn_kl.f dvec.head env.get_ks sum_costs at H_kl_gint,
tactic.dget_dinsert_at `H_kl_gint,
dunfold det.op.f ops.mvn_kl ops.mvn_kl.f sumr map dvec.head dvec.head2 dvec.head3 at H_kl_gint,
simp only [H_E_kl_add] at H_kl_gint,
exact (iff.mpr (T.is_integrable_add_middle _ _ _) H_kl_gint^.left)^.left
},

assert H_lhs_eint₂ : E.is_eintegrable d_base lhs_f₂,
{
dsimp [E.is_eintegrable, dvec.head],
dsimp [integrate_mvn_kl, is_gintegrable] at H_kl_gint,
simp only [λ a1 a2 a3, @env.insert_insert_flip (eloss, []) (z, shape) a1 a2 a3 H_eloss_neq_z],
simp only [det.op.f, ops.mvn_kl, p1, p2, ops.mvn_kl.f, dvec.head, env.get_ks, sum_costs] at H_kl_gint,
dunfold det.op.f ops.mvn_kl p1 p2 ops.mvn_kl.f dvec.head env.get_ks sum_costs at H_kl_gint,
tactic.dget_dinsert_at `H_kl_gint,
dunfold sumr map dvec.head dvec.head2 dvec.head3 at H_kl_gint,
simp only [H_E_kl_add] at H_kl_gint,
exact (iff.mpr (T.is_integrable_add_middle _ _ _) H_kl_gint^.left)^.right
},

erw E.E_add d_base lhs_f₁ lhs_f₂ H_lhs_eint₁ H_lhs_eint₂,

pose rhs_f₁ := λ x, E (graph.to_dist (λ (m : env), ⟦k₁ m⟧) (m_rhs_k_add x) nodes) dvec.head,
pose rhs_f₂ := λ x, E (graph.to_dist (λ (m : env), ⟦k₂ m⟧) (m_rhs_k_add x) nodes) dvec.head,

dsimp [graph.to_dist, operator.to_dist, is_gintegrable, integrate_mvn_kl, dvec.head] at H_gint,
simp only [E.E_bind, E.E_ret, det.op.f, dvec.head, env.get_ks, sum_costs] at H_gint,
tactic.dget_dinsert_at `H_gint,

assert H_E_add : ∀ x,
E
         (graph.to_dist
            (λ (m : env), ⟦(env.get (eloss, @nil ℕ) m : ℝ) + sumr (map (λ (cost : ID), (env.get (cost, @nil ℕ) m : ℝ)) costs)⟧)
            (env.insert (eloss, @nil ℕ)
               (T.mvn_empirical_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape) x)
               (env.insert (z, shape) x inputs))
            nodes)
         dvec.head
=
E
         (graph.to_dist
            (λ (m : env), ⟦k₁ m⟧)
            (env.insert (eloss, @nil ℕ)
               (T.mvn_empirical_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape) x)
               (env.insert (z, shape) x inputs))
            nodes)
         dvec.head
+
E
         (graph.to_dist
            (λ (m : env), ⟦k₂ m⟧)
            (env.insert (eloss, @nil ℕ)
               (T.mvn_empirical_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape) x)
               (env.insert (z, shape) x inputs))
            nodes)
         dvec.head,
{
intro x,
exact E.E_k_add k₁ k₂ _ _ (iff.mpr (is_gintegrable_k_add _ _ _ _) (H_gint^.right x))^.left (iff.mpr (is_gintegrable_k_add _ _ _ _) (H_gint^.right x))^.right
},

assert H_rhs_eint₁ : E.is_eintegrable d_base rhs_f₁,
{
dsimp [E.is_eintegrable, dvec.head],
dunfold sumr map dvec.head dvec.head2 dvec.head3 at H_gint,
simp only [H_E_add] at H_gint,
exact (iff.mpr (T.is_integrable_add_middle _ _ _) H_gint^.left)^.left
},

assert H_rhs_eint₂ : E.is_eintegrable d_base rhs_f₂,
{
dsimp [E.is_eintegrable, dvec.head],
dunfold sumr map dvec.head dvec.head2 dvec.head3 at H_gint,
simp only [H_E_add] at H_gint,
exact (iff.mpr (T.is_integrable_add_middle _ _ _) H_gint^.left)^.right
},

erw E.E_add d_base rhs_f₁ rhs_f₂ H_rhs_eint₁ H_rhs_eint₂,

clear integrate_mvn_kl_correct,

assert H_term₁_lhs :
∀ (x : dvec T [shape]),
E (graph.to_dist (λ (m : env), ⟦(λ (m : env), env.get (eloss, []) m) m⟧)
                 (env.insert (eloss, []) (T.mvn_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape))
                              (env.insert (z, shape) (dvec.head x) inputs))
                 nodes)
   dvec.head
=
T.mvn_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape),
{ intro x, apply (E.E_of_lookup H_eloss_nin_nodes),
dsimp [pdfs_exist_at] at H_pdfs_exist_at,
dsimp [all_parents_in_env] at H_ps_in_env,
exact (pdfs_exist_at_ignore (H_pre^.right^.right^.right^.right _) (env_not_has_key_insert H_eloss_neq_z H_eloss_nin) H_eloss_nin_nodes (H_pdfs_exist_at^.right _))
},

assert H_term₁_rhs :
∀ (x : dvec T [shape]),
E (graph.to_dist (λ (m : env), ⟦(λ (m : env), env.get (eloss, @nil ℕ) m) m⟧)
                 (env.insert (eloss, @nil ℕ)
                   (T.mvn_empirical_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape) (dvec.head x))
               (env.insert (z, shape) (dvec.head x) inputs))
            nodes)
         dvec.head
=
T.mvn_empirical_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape) (dvec.head x),
{ intro x, apply (E.E_of_lookup H_eloss_nin_nodes),
dsimp [pdfs_exist_at] at H_pdfs_exist_at,
dsimp [all_parents_in_env] at H_ps_in_env,
exact (pdfs_exist_at_ignore (H_pre^.right^.right^.right^.right _) (env_not_has_key_insert H_eloss_neq_z H_eloss_nin) H_eloss_nin_nodes (H_pdfs_exist_at^.right _))
 },

assert H_term₁ :
E (sprog.prim (rand.op.mvn shape) ⟦env.get (μ, shape) inputs, env.get (σ, shape) inputs⟧)
    (λ (x : dvec T [shape]),
       E
         (graph.to_dist (λ (m : env), ⟦(λ (m : env), env.get (eloss, @nil ℕ) m) m⟧)
            (env.insert (eloss, @nil ℕ) (T.mvn_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape))
               (env.insert (z, shape) (dvec.head x) inputs))
            nodes)
         dvec.head)
=
E (sprog.prim (rand.op.mvn shape) ⟦env.get (μ, shape) inputs, env.get (σ, shape) inputs⟧)
    (λ (x : dvec T [shape]),
         E (graph.to_dist (λ (m : env), ⟦(λ (m : env), env.get (eloss, @nil ℕ) m) m⟧)
                          (env.insert (eloss, @nil ℕ)
                                       (T.mvn_empirical_kl (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape) (dvec.head x))
                                       (env.insert (z, shape) (dvec.head x) inputs))
            nodes)
         dvec.head),
{
simp [H_term₁_lhs, H_term₁_rhs],
dunfold E T.dintegral dvec.head rand.op.pdf dvec.head2 dvec.head3,
dsimp,
dunfold dvec.head,
erw T.integral_fscale,
erw (@T.mvn_kl_identity shape (env.get (μ, shape) inputs) (env.get (σ, shape) inputs) H_pre^.right^.right^.right^.left),
assertv H_pdf_1 : ∫ (λ (x : T shape), T.mvn_pdf (env.get (μ, shape) inputs : T shape) (env.get (σ, shape) inputs : T shape) x) = 1 :=
  T.mvn_pdf_int1 _ _ H_pre^.right^.right^.right^.left,
delta rand.pdf.mvn,
dsimp,
rw H_pdf_1,
rw T.one_smul
},

erw H_term₁, clear H_term₁, apply congr_arg,
apply congr_arg, apply funext, intro x,
assertv H_ps_in_env : all_parents_in_env (env.insert (z, shape) x^.head inputs) nodes := by apply H_pre^.right^.right^.right^.right,
dsimp,
erw (to_dist_congr_insert H_ps_in_env (env_not_has_key_insert H_eloss_neq_z H_eloss_nin) H_eloss_nin_nodes H_eloss_not_cost),
erw (to_dist_congr_insert H_ps_in_env (env_not_has_key_insert H_eloss_neq_z H_eloss_nin) H_eloss_nin_nodes H_eloss_not_cost)
end

-- EQ11
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, []), [(pname₃, shape₃), (pname₄, shape₄), (pname₅, shape₅)], operator.rand op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint := false.rec _ H_pre

-- EQ12
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, []), ((pname₃, shape₃)::(pname₄, shape₄)::(pname₅, shape₅)::parent::parents), op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint := false.rec _ H_pre

-- EQ13
| (⟨(rname, .(shape)), [(pname₁, .(shape)), (pname₂, .(shape))], operator.rand (rand.op.mvn shape)⟩
  ::⟨(rname₂, (shape₂ :: shapes₂)), parents, op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint := false.rec _ H_pre

-- EQ14 (a)
| (⟨(rname₂, shape₂), ((pname₃, shape₃)::(pname₄, shape₄)::parent::parents), operator.det op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint :=
begin
dunfold graph.to_dist operator.to_dist integrate_mvn_kl,
dunfold integrate_mvn_kl_pre at H_pre,
rw [E.E_bind, E.E_bind],
simp [E.E_ret],

definev next_inputs : env := env.insert (rname₂, shape₂) (op^.f (env.get_ks ((pname₃, shape₃)::(pname₄, shape₄)::parent::parents) inputs)) inputs,
definev next_nodes : list node := nodes,

assertv H_pre_next : integrate_mvn_kl_pre eloss next_nodes next_inputs := H_pre,
assertv H_ps_in_env_next : all_parents_in_env next_inputs next_nodes := H_ps_in_env^.right _,
exact (integrate_mvn_kl_correct _ _ H_eloss_not_cost H_pre_next (H_uids^.right _) H_ps_in_env_next H_pdfs_exist_at H_kl_gint H_gint)
end

-- EQ14 (b)
| (⟨(rname₂, shape₂), ((pname₃, shape₃)::(pname₄, shape₄)::parent::parents), operator.rand op⟩::nodes) inputs H_eloss_not_cost H_pre H_uids H_ps_in_env H_pdfs_exist_at H_kl_gint H_gint := false.rec _ H_pre

-- More useful API

def integrate_kl_pre : graph → env → Prop
| g m := integrate_mvn_kl_pre g^.costs^.head g^.nodes m

def integrate_kl : graph → graph
| g := ⟨integrate_mvn_kl g^.costs^.head g^.nodes, g^.costs, g^.targets, g^.inputs⟩

end certigrad


/-
Copyright (c) 2020 Scott Morrison. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Scott Morrison
-/
import category_theory.shift
import category_theory.concrete_category
import category_theory.pi.basic
import algebra.group.basic

/-!
# The category of graded objects

For any type `β`, a `β`-graded object over some category `C` is just
a function `β → C` into the objects of `C`.
We put the "pointwise" category structure on these, as the non-dependent specialization of
`category_theory.pi`.

We describe the `comap` functors obtained by precomposing with functions `β → γ`.

As a consequence a fixed element (e.g. `1`) in an additive group `β` provides a shift
functor on `β`-graded objects

When `C` has coproducts we construct the `total` functor `graded_object β C ⥤ C`,
show that it is faithful, and deduce that when `C` is concrete so is `graded_object β C`.
-/

open category_theory.pi
open category_theory.limits

namespace category_theory

universes w v u

/-- A type synonym for `β → C`, used for `β`-graded objects in a category `C`. -/
def graded_object (β : Type w) (C : Type u) : Type (max w u) := β → C

-- Satisfying the inhabited linter...
instance inhabited_graded_object (β : Type w) (C : Type u) [inhabited C] :
  inhabited (graded_object β C) :=
⟨λ b, inhabited.default C⟩

/--
A type synonym for `β → C`, used for `β`-graded objects in a category `C`
with a shift functor given by translation by `s`.
-/
@[nolint unused_arguments] -- `s` is here to distinguish type synonyms asking for different shifts
abbreviation graded_object_with_shift {β : Type w} [add_comm_group β] (s : β) (C : Type u) :
  Type (max w u) := graded_object β C

namespace graded_object

variables {C : Type u} [category.{v} C]

instance category_of_graded_objects (β : Type w) : category.{(max w v)} (graded_object β C) :=
category_theory.pi (λ _, C)

section
variable (C)

/--
The natural isomorphism comparing between
pulling back along two propositionally equal functions.
-/
@[simps]
def comap_eq {β γ : Type w} {f g : β → γ} (h : f = g) : comap (λ _, C) f ≅ comap (λ _, C) g :=
{ hom := { app := λ X b, eq_to_hom begin dsimp [comap], subst h, end },
  inv := { app := λ X b, eq_to_hom begin dsimp [comap], subst h, end }, }

lemma comap_eq_symm {β γ : Type w} {f g : β → γ} (h : f = g) :
  comap_eq C h.symm = (comap_eq C h).symm :=
by tidy

lemma comap_eq_trans {β γ : Type w} {f g h : β → γ} (k : f = g) (l : g = h) :
  comap_eq C (k.trans l) = comap_eq C k ≪≫ comap_eq C l :=
begin
  ext X b,
  simp,
end

/--
The equivalence between β-graded objects and γ-graded objects,
given an equivalence between β and γ.
-/
@[simps]
def comap_equiv {β γ : Type w} (e : β ≃ γ) :
  (graded_object β C) ≌ (graded_object γ C) :=
{ functor := comap (λ _, C) (e.symm : γ → β),
  inverse := comap (λ _, C) (e : β → γ),
  counit_iso := (comap_comp (λ _, C) _ _).trans (comap_eq C (by { ext, simp } )),
  unit_iso := (comap_eq C (by { ext, simp } )).trans (comap_comp _ _ _).symm,
  functor_unit_iso_comp' := λ X, by { ext b, dsimp, simp, }, }  -- See note [dsimp, simp].

end

instance has_shift {β : Type*} [add_comm_group β] (s : β) :
  has_shift (graded_object_with_shift s C) :=
{ shift := comap_equiv C
  { to_fun := λ b, b-s,
    inv_fun := λ b, b+s,
    left_inv := λ x, (by simp),
    right_inv := λ x, (by simp), } }

@[simp] lemma shift_functor_obj_apply {β : Type*} [add_comm_group β] (s : β) (X : β → C) (t : β) :
  (shift (graded_object_with_shift s C)).functor.obj X t = X (t + s) :=
rfl

@[simp] lemma shift_functor_map_apply {β : Type*} [add_comm_group β] (s : β)
  {X Y : graded_object_with_shift s C} (f : X ⟶ Y) (t : β) :
  (shift (graded_object_with_shift s C)).functor.map f t = f (t + s) :=
rfl

instance has_zero_morphisms [has_zero_morphisms C] (β : Type w) :
  has_zero_morphisms.{(max w v)} (graded_object β C) :=
{ has_zero := λ X Y,
  { zero := λ b, 0 } }

@[simp]
lemma zero_apply [has_zero_morphisms C] (β : Type w) (X Y : graded_object β C) (b : β) :
  (0 : X ⟶ Y) b = 0 := rfl

section
local attribute [instance] has_zero_object.has_zero

instance has_zero_object [has_zero_object C] [has_zero_morphisms C] (β : Type w) :
  has_zero_object.{(max w v)} (graded_object β C) :=
{ zero := λ b, (0 : C),
  unique_to := λ X, ⟨⟨λ b, 0⟩, λ f, (by ext)⟩,
  unique_from := λ X, ⟨⟨λ b, 0⟩, λ f, (by ext)⟩, }
end

end graded_object

namespace graded_object
-- The universes get a little hairy here, so we restrict the universe level for the grading to 0.
-- Since we're typically interested in grading by ℤ or a finite group, this should be okay.
-- If you're grading by things in higher universes, have fun!
variables (β : Type)
variables (C : Type u) [category.{v} C]
variables [has_coproducts C]

/--
The total object of a graded object is the coproduct of the graded components.
-/
noncomputable def total : graded_object β C ⥤ C :=
{ obj := λ X, ∐ (λ i : ulift.{v} β, X i.down),
  map := λ X Y f, limits.sigma.map (λ i, f i.down) }.

variables [has_zero_morphisms C]

/--
The `total` functor taking a graded object to the coproduct of its graded components is faithful.
To prove this, we need to know that the coprojections into the coproduct are monomorphisms,
which follows from the fact we have zero morphisms and decidable equality for the grading.
-/
instance : faithful (total β C) :=
{ map_injective' := λ X Y f g w,
  begin
    classical,
    ext i,
    replace w := sigma.ι (λ i : ulift β, X i.down) ⟨i⟩ ≫= w,
    erw [colimit.ι_map, colimit.ι_map] at w,
    exact mono.right_cancellation _ _ w,
  end }

end graded_object

namespace graded_object

noncomputable theory

variables (β : Type)
variables (C : Type (u+1)) [large_category C] [concrete_category C]
  [has_coproducts C] [has_zero_morphisms C]

instance : concrete_category (graded_object β C) :=
{ forget := total β C ⋙ forget C }

instance : has_forget₂ (graded_object β C) C :=
{ forget₂ := total β C }

end graded_object

end category_theory


/-
Copyright (c) 2017 Johannes Hölzl. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Johannes Hölzl, Mario Carneiro

The module `lc α β` of linear combinations over `β` (`α` is the scalar ring)
-/
import linear_algebra.basic
noncomputable theory

open classical function lattice
local attribute [instance] prop_decidable

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

/-- The type of linear coefficients, which are simply the finitely supported functions
from the module `β` to the scalar ring `α`. -/
@[reducible] def lc (α β) [ring α] [add_comm_group β] [module α β] : Type* := β →₀ α

namespace lc
variables [ring α] [add_comm_group β] [add_comm_group γ] [add_comm_group δ]
variables [module α β] [module α γ] [module α δ]
open submodule linear_map

instance : add_comm_group (lc α β) := finsupp.add_comm_group

instance : has_scalar α (lc α β) := finsupp.to_has_scalar

instance : module α (lc α β) := finsupp.to_module β α

variables (α)
def supported (s : set β) : submodule α (lc α β) :=
{ carrier := {l | ↑l.support ⊆ s},
  zero := by simp,
  add := λ l₁ l₂ h₁ h₂, set.subset.trans (finset.coe_subset.2 finsupp.support_add)
    (by simpa using set.union_subset h₁ h₂),
  smul := λ a l h, set.subset.trans (finset.coe_subset.2 finsupp.support_smul)
    (by simpa using h) }
variables {α}

theorem mem_supported {s : set β} {l : lc α β} :
  l ∈ supported α s ↔ ↑l.support ⊆ s := iff.rfl

theorem mem_supported' {s : set β} {l : lc α β} :
  l ∈ supported α s ↔ ∀ x ∉ s, l x = 0 :=
by simp [mem_supported, set.subset_def, not_imp_comm]

theorem single_mem_supported {s : set β} (a : α) {b : β} (h : b ∈ s) :
  finsupp.single b a ∈ supported α s :=
set.subset.trans finsupp.support_single_subset (set.singleton_subset_iff.2 h)

theorem supported_eq_span_single (s : set β) :
  lc.supported α s = span α ((λ x, finsupp.single x (1:α)) '' s) :=
begin
  refine (span_eq_of_le _ _ (le_def'.2 $ λ l hl, _)).symm,
  { rintro _ ⟨l, hl, rfl⟩, exact single_mem_supported _ hl },
  { rw ← l.sum_single,
    refine sum_mem _ (λ x xl, _),
    rw (by simp : finsupp.single x (l x) = l x • finsupp.single x 1),
    exact smul_mem _ _ (subset_span ⟨_, hl xl, rfl⟩) }
end

variables (α)
def restrict_dom (s : set β) : lc α β →ₗ supported α s :=
linear_map.cod_restrict _
  { to_fun := finsupp.filter (∈ s),
    add := λ l₁ l₂, finsupp.filter_add,
    smul := λ a l, finsupp.filter_smul }
  (λ l, mem_supported'.2 $ λ x, finsupp.filter_apply_neg (∈ s) l)
variables {α}

@[simp] theorem restrict_dom_apply (s : set β) (l : lc α β) :
  ↑((restrict_dom α s : lc α β →ₗ supported α s) l) = finsupp.filter (∈ s) l := rfl

theorem restrict_dom_comp_subtype (s : set β) :
  (restrict_dom α s).comp (submodule.subtype _) = linear_map.id :=
by ext l; apply subtype.coe_ext.2; simp; ext a;
   by_cases a ∈ s; simp [h]; exact (mem_supported'.1 l.2 _ h).symm

theorem range_restrict_dom (s : set β) : (restrict_dom α s).range = ⊤ :=
begin
  have := linear_map.range_comp (submodule.subtype _) (restrict_dom α s),
  rw [restrict_dom_comp_subtype, linear_map.range_id] at this,
  exact eq_top_mono (submodule.map_mono le_top) this.symm
end

theorem supported_mono {s t : set β} (st : s ⊆ t) :
  supported α s ≤ supported α t :=
λ l h, set.subset.trans h st

@[simp] theorem supported_empty : supported α (∅ : set β) = ⊥ :=
eq_bot_iff.2 $ λ l h, (submodule.mem_bot α).2 $
by ext; simp [*, mem_supported'] at *

@[simp] theorem supported_univ : supported α (set.univ : set β) = ⊤ :=
eq_top_iff.2 $ λ l _, set.subset_univ _

theorem supported_Union {ι : Type*} (s : ι → set β) :
  supported α (⋃ i, s i) = ⨆ i, supported α (s i) :=
begin
  refine le_antisymm _ (supr_le $ λ i, supported_mono $ set.subset_Union _ _),
  suffices : ((submodule.subtype _).comp (restrict_dom α (⋃ i, s i))).range ≤ ⨆ i, supported α (s i),
  { rwa [range_comp, range_restrict_dom, map_top, range_subtype] at this },
  rw [range_le_iff_comap, eq_top_iff],
  rintro l ⟨⟩, rw mem_coe,
  apply finsupp.induction l, {exact zero_mem _},
  refine λ x a l hl a0, add_mem _ _,
  by_cases (∃ i, x ∈ s i); simp [h],
  cases h with i hi,
  exact le_supr (λ i, supported α (s i)) i (single_mem_supported _ hi)
end

theorem supported_union (s t : set β) :
  supported α (s ∪ t) = supported α s ⊔ supported α t :=
by erw [set.union_eq_Union, supported_Union, supr_bool_eq]; refl

theorem supported_Inter {ι : Type*} (s : ι → set β) :
  supported α (⋂ i, s i) = ⨅ i, supported α (s i) :=
begin
  refine le_antisymm (le_infi $ λ i, supported_mono $ set.Inter_subset _ _) _,
  simp [le_def, infi_coe, set.subset_def],
  exact λ l, set.subset_Inter
end

def apply (x : β) : lc α β →ₗ α :=
⟨λ l, l x, λ _ _, finsupp.add_apply, λ _ _, finsupp.smul_apply⟩

@[simp] theorem apply_apply (x : β) (l : lc α β) :
  (lc.apply x : lc α β →ₗ α) l = l x := rfl

protected def lsum (f : β → α →ₗ[α] γ) : lc α β →ₗ[α] γ :=
⟨λ d, d.sum (λ b, f b),
  assume d₁ d₂, by simp [finsupp.sum_add_index],
  assume a d, by simp [finsupp.sum_smul_index, finsupp.smul_sum,
    -smul_eq_mul, smul_eq_mul.symm]⟩

@[simp] theorem lsum_apply (f : β → α →ₗ γ) (l : lc α β) :
  (lc.lsum f : lc α β →ₗ γ) l = l.sum (λ b, f b) := rfl

section
variables (α β)
protected def total : lc α β →ₗ β := lc.lsum linear_map.id.smul_right
end

theorem total_apply (l : lc α β) :
  lc.total α β l = l.sum (λ b a, a • b) := rfl

@[simp] theorem total_single (a : α) (x : β) :
  lc.total α β (finsupp.single x a) = a • x :=
by simp [total_apply, finsupp.sum_single_index]

@[simp] theorem total_range : (lc.total α β).range = ⊤ :=
range_eq_top.2 $ λ x, ⟨finsupp.single x 1, by simp⟩

variables (α)
protected def map (f : β → γ) : lc α β →ₗ[α] lc α γ :=
{ to_fun := finsupp.map_domain f,
  add := λ l₁ l₂, finsupp.map_domain_add,
  smul := λ a l, finsupp.map_domain_smul _ _ }
variables {α}

@[simp] theorem map_apply (f : β → γ) (l : lc α β) :
  (lc.map α f : _ →ₗ _) l = finsupp.map_domain f l := rfl

@[simp] theorem map_id : (lc.map α id : lc α β →ₗ[α] lc α β) = linear_map.id :=
linear_map.ext $ λ l, finsupp.map_domain_id

theorem map_comp (f : β → γ) (g : γ → δ) :
  lc.map α (g ∘ f) = (lc.map α g).comp (lc.map α f) :=
linear_map.ext $ λ l, finsupp.map_domain_comp

theorem supported_comap_map (f : β → γ) (s : set γ) :
  supported α (f ⁻¹' s) ≤ (supported α s).comap (lc.map α f) :=
λ l (hl : ↑l.support ⊆ f ⁻¹' s),
show ↑(finsupp.map_domain f l).support ⊆ s, begin
  rw [← set.image_subset_iff, ← finset.coe_image] at hl,
  exact set.subset.trans finsupp.map_domain_support hl
end

theorem map_supported (f : β → γ) (s : set β) :
  (supported α s).map (lc.map α f) = supported α (f '' s) :=
begin
  refine le_antisymm (map_le_iff_le_comap.2 $
    le_trans (supported_mono $ set.subset_preimage_image _ _)
       (supported_comap_map _ _)) _,
  intros l hl, haveI : inhabited β := ⟨0⟩,
  refine ⟨(lc.map α (inv_fun_on f s) : lc α γ →ₗ lc α β) l, λ x hx, _, _⟩,
  { rcases finset.mem_image.1 (finsupp.map_domain_support hx) with ⟨c, hc, rfl⟩,
    exact inv_fun_on_mem (by simpa using hl hc) },
  { rw [← linear_map.comp_apply, ← map_comp],
    refine (finsupp.map_domain_congr $ λ c hc, _).trans finsupp.map_domain_id,
    exact inv_fun_on_eq (by simpa using hl hc) }
end

theorem map_disjoint_ker (f : β → γ) {s : set β}
  (H : ∀ a b ∈ s, f a = f b → a = b) :
  disjoint (supported α s) (lc.map α f).ker :=
begin
  rintro l ⟨h₁, h₂⟩,
  rw [mem_coe, mem_ker, map_apply, finsupp.map_domain] at h₂,
  simp, ext x,
  by_cases xs : x ∈ s,
  { have : finsupp.sum l (λ a, finsupp.single (f a)) (f x) = 0, {rw h₂, refl},
    rw [finsupp.sum_apply, finsupp.sum, finset.sum_eq_single x] at this,
    { simpa [finsupp.single_apply] },
    { intros y hy xy, simp [mt (H _ _ (h₁ hy) xs) xy] },
    { simp {contextual := tt} } },
  { by_contra h, exact xs (h₁ $ finsupp.mem_support_iff.2 h) }
end

theorem map_total (f : β →ₗ[α] γ) :
  (lc.total α γ).comp (lc.map α f) = f.comp (lc.total α β) :=
by ext; simp [total_apply, finsupp.sum_map_domain_index, add_smul]

end lc

namespace lc
variables [discrete_field α] [add_comm_group β] [vector_space α β]

instance : vector_space α (lc α β) := { .. lc.module }

end lc

section module
variables [ring α] [add_comm_group β] [add_comm_group γ] [add_comm_group δ]
variables [module α β] [module α γ] [module α δ]
variables {a b : α} {s t : set β} {x y : β}
include α
open submodule

theorem span_eq_map_lc : span α s = (lc.supported α s).map (lc.total α β) :=
begin
  apply span_eq_of_le,
  { exact λ x hx, ⟨_, lc.single_mem_supported 1 hx, by simp⟩ },
  { refine map_le_iff_le_comap.2 (λ v hv, _),
    have : ∀c, v c • c ∈ span α s,
    { intro c, by_cases c ∈ s,
      { exact smul_mem _ _ (subset_span h) },
      { simp [lc.mem_supported'.1 hv _ h] } },
    refine sum_mem _ _, simp [this] }
end

theorem mem_span_iff_lc : x ∈ span α s ↔ ∃ l ∈ lc.supported α s, lc.total α β l = x :=
by rw span_eq_map_lc; simp

variables (α)
def lc.total_on (s : set β) : lc.supported α s →ₗ span α s :=
linear_map.cod_restrict _ ((lc.total α _).comp (submodule.subtype _)) $
  λ ⟨l, hl⟩, mem_span_iff_lc.2 ⟨l, hl, rfl⟩
variables {α}

theorem lc.total_on_range (s : set β) : (lc.total_on α s).range = ⊤ :=
by rw [lc.total_on, linear_map.range, linear_map.map_cod_restrict, ← linear_map.range_le_iff_comap,
  range_subtype, map_top, linear_map.range_comp, range_subtype]; exact le_of_eq span_eq_map_lc

lemma linear_eq_on {f g : β →ₗ[α] γ} (H : ∀x∈s, f x = g x) {x} (h : x ∈ span α s) : f x = g x :=
by apply span_induction h H; simp {contextual := tt}

end module


/-
Copyright (c) 2021 Jannis Limperg. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Jannis Limperg
-/

import Aesop

set_option aesop.check.all true

example : True := by
  aesop

example : Unit := by
  aesop

example : PUnit.{u} := by
  aesop

example (h : False) : α := by
  aesop

example (h : Empty) : α := by
  aesop

example (h : PEmpty.{u}) : α := by
  aesop


import linear_algebra.matrix
import data.matrix.basis
import data.polynomial

def curry (a b y: Type*) (f: a × b -> y) : a -> b -> y :=
  begin
    intros aa bb,
    exact f (aa, bb),
  end
#print curry

def curry2 (a b y: Type*) (f: a × b -> y) : a -> b -> y :=
  λ (x1: a) (x2: b), f (x1, x2)

constants (x: nat)
variables (y: nat)
section testSection
  #print x
  variable (p: Prop)
  theorem test1: p \/ ¬p := begin
    exact classical.em p
  end
end testSection

namespace testSection2
  variable (p: Prop)
  theorem test1: p \/ ¬p := begin
    exact classical.em p
  end
end testSection2