import Spa.Language.Traces

namespace Spa

open Graph

theorem Fin.castAdd_ne_natAdd {n m : ℕ} (i : Fin n) (j : Fin m) :
    Fin.castAdd m i ≠ Fin.natAdd n j := by
  intro h
  have := congrArg Fin.val h
  simp only [Fin.coe_castAdd, Fin.coe_natAdd] at this
  omega

/-! ### Trace embeddings -/

section Embeddings

variable {g₁ g₂ : Graph} {ρ₁ ρ₂ : Env}

theorem Trace.comp_left {idx₁ idx₂ : g₁.Index}
    (tr : Trace g₁ idx₁ idx₂ ρ₁ ρ₂) :
    Trace (g₁ ∙ g₂) (idx₁.castAdd g₂.size) (idx₂.castAdd g₂.size) ρ₁ ρ₂ := by
  induction tr with
  | single hbs =>
    exact Trace.single (by rwa [show (g₁ ∙ g₂).nodes = Fin.append g₁.nodes g₂.nodes from rfl,
      Fin.append_left])
  | edge hbs he _ ih =>
    refine Trace.edge ?_ ?_ ih
    · rwa [show (g₁ ∙ g₂).nodes = Fin.append g₁.nodes g₂.nodes from rfl, Fin.append_left]
    · exact List.mem_append_left _ (List.mem_map_of_mem _ he)

theorem Trace.comp_right {idx₁ idx₂ : g₂.Index}
    (tr : Trace g₂ idx₁ idx₂ ρ₁ ρ₂) :
    Trace (g₁ ∙ g₂) (idx₁.natAdd g₁.size) (idx₂.natAdd g₁.size) ρ₁ ρ₂ := by
  induction tr with
  | single hbs =>
    exact Trace.single (by rwa [show (g₁ ∙ g₂).nodes = Fin.append g₁.nodes g₂.nodes from rfl,
      Fin.append_right])
  | edge hbs he _ ih =>
    refine Trace.edge ?_ ?_ ih
    · rwa [show (g₁ ∙ g₂).nodes = Fin.append g₁.nodes g₂.nodes from rfl, Fin.append_right]
    · exact List.mem_append_right _ (List.mem_map_of_mem _ he)

theorem Trace.link_left {idx₁ idx₂ : g₁.Index}
    (tr : Trace g₁ idx₁ idx₂ ρ₁ ρ₂) :
    Trace (g₁ ⤳ g₂) (idx₁.castAdd g₂.size) (idx₂.castAdd g₂.size) ρ₁ ρ₂ := by
  induction tr with
  | single hbs =>
    exact Trace.single (by rwa [show (g₁ ⤳ g₂).nodes = Fin.append g₁.nodes g₂.nodes from rfl,
      Fin.append_left])
  | edge hbs he _ ih =>
    refine Trace.edge ?_ ?_ ih
    · rwa [show (g₁ ⤳ g₂).nodes = Fin.append g₁.nodes g₂.nodes from rfl, Fin.append_left]
    · exact List.mem_append_left _ (List.mem_append_left _ (List.mem_map_of_mem _ he))

theorem Trace.link_right {idx₁ idx₂ : g₂.Index}
    (tr : Trace g₂ idx₁ idx₂ ρ₁ ρ₂) :
    Trace (g₁ ⤳ g₂) (idx₁.natAdd g₁.size) (idx₂.natAdd g₁.size) ρ₁ ρ₂ := by
  induction tr with
  | single hbs =>
    exact Trace.single (by rwa [show (g₁ ⤳ g₂).nodes = Fin.append g₁.nodes g₂.nodes from rfl,
      Fin.append_right])
  | edge hbs he _ ih =>
    refine Trace.edge ?_ ?_ ih
    · rwa [show (g₁ ⤳ g₂).nodes = Fin.append g₁.nodes g₂.nodes from rfl, Fin.append_right]
    · exact List.mem_append_left _
        (List.mem_append_right _ (List.mem_map_of_mem _ he))

theorem EndToEndTrace.comp_left (etr : EndToEndTrace g₁ ρ₁ ρ₂) :
    EndToEndTrace (g₁ ∙ g₂) ρ₁ ρ₂ := by
  obtain ⟨i₁, h₁, i₂, h₂, tr⟩ := etr
  exact ⟨i₁.castAdd g₂.size, List.mem_append_left _ (List.mem_map_of_mem _ h₁),
         i₂.castAdd g₂.size, List.mem_append_left _ (List.mem_map_of_mem _ h₂),
         tr.comp_left⟩

theorem EndToEndTrace.comp_right (etr : EndToEndTrace g₂ ρ₁ ρ₂) :
    EndToEndTrace (g₁ ∙ g₂) ρ₁ ρ₂ := by
  obtain ⟨i₁, h₁, i₂, h₂, tr⟩ := etr
  exact ⟨i₁.natAdd g₁.size, List.mem_append_right _ (List.mem_map_of_mem _ h₁),
         i₂.natAdd g₁.size, List.mem_append_right _ (List.mem_map_of_mem _ h₂),
         tr.comp_right⟩

theorem EndToEndTrace.concat {ρ₃ : Env} (etr₁ : EndToEndTrace g₁ ρ₁ ρ₂)
    (etr₂ : EndToEndTrace g₂ ρ₂ ρ₃) : EndToEndTrace (g₁ ⤳ g₂) ρ₁ ρ₃ := by
  obtain ⟨i₁, h₁, i₂, h₂, tr₁⟩ := etr₁
  obtain ⟨j₁, k₁, j₂, k₂, tr₂⟩ := etr₂
  refine ⟨i₁.castAdd g₂.size, List.mem_map_of_mem _ h₁,
          j₂.natAdd g₁.size, List.mem_map_of_mem _ k₂,
          Trace.concat tr₁.link_left ?_ tr₂.link_right⟩
  exact List.mem_append_right _
    (List.mem_product.mpr ⟨List.mem_map_of_mem _ h₂, List.mem_map_of_mem _ k₁⟩)

end Embeddings

/-! ### Loops -/

section Loop

variable {g : Graph} {ρ₁ ρ₂ ρ₃ : Env}

theorem Trace.loop {idx₁ idx₂ : g.Index} (tr : Trace g idx₁ idx₂ ρ₁ ρ₂) :
    Trace (Graph.loop g) (idx₁.natAdd 2) (idx₂.natAdd 2) ρ₁ ρ₂ := by
  induction tr with
  | single hbs =>
    exact Trace.single (by
      rwa [show (Graph.loop g).nodes = Fin.append (fun _ : Fin 2 => []) g.nodes from rfl,
        Fin.append_right])
  | edge hbs he _ ih =>
    refine Trace.edge ?_ ?_ ih
    · rwa [show (Graph.loop g).nodes = Fin.append (fun _ : Fin 2 => []) g.nodes from rfl,
        Fin.append_right]
    · exact List.mem_append_left _ (List.mem_append_left _
        (List.mem_append_left _ (List.mem_map_of_mem _ he)))

private theorem loop_nodes_at_in :
    (Graph.loop g).nodes g.loopIn = [] :=
  Fin.append_left (fun _ : Fin 2 => []) g.nodes 0

private theorem loop_nodes_at_out :
    (Graph.loop g).nodes g.loopOut = [] :=
  Fin.append_left (fun _ : Fin 2 => []) g.nodes 1

theorem EndToEndTrace.loop (etr : EndToEndTrace g ρ₁ ρ₂) :
    EndToEndTrace (Graph.loop g) ρ₁ ρ₂ := by
  obtain ⟨i₁, h₁, i₂, h₂, tr⟩ := etr
  -- the edge in → (2 ↑ʳ i₁), reached through the second edge group
  have hin : (g.loopIn, i₁.natAdd 2) ∈ (Graph.loop g).edges := by
    refine List.mem_append_left _ (List.mem_append_left _ (List.mem_append_right _ ?_))
    exact List.mem_map_of_mem _ (List.mem_map_of_mem _ h₁)
  -- the edge (2 ↑ʳ i₂) → out, reached through the third edge group
  have hout : (i₂.natAdd 2, g.loopOut) ∈ (Graph.loop g).edges := by
    refine List.mem_append_left _ (List.mem_append_right _ ?_)
    exact List.mem_map_of_mem _ (List.mem_map_of_mem _ h₂)
  refine ⟨g.loopIn, List.mem_singleton_self _, g.loopOut, List.mem_singleton_self _, ?_⟩
  exact Trace.concat (Trace.single (loop_nodes_at_in ▸ EvalBasicStmts.nil)) hin
    (Trace.concat tr.loop hout (Trace.single (loop_nodes_at_out ▸ EvalBasicStmts.nil)))

private theorem loop_edge_out_in :
    ((g.loopOut, g.loopIn) : (Graph.loop g).Edge) ∈ (Graph.loop g).edges := by
  refine List.mem_append_right _ ?_
  exact List.mem_cons_self _ _

theorem EndToEndTrace.loop_concat (etr₁ : EndToEndTrace (Graph.loop g) ρ₁ ρ₂)
    (etr₂ : EndToEndTrace (Graph.loop g) ρ₂ ρ₃) :
    EndToEndTrace (Graph.loop g) ρ₁ ρ₃ := by
  obtain ⟨i₁, h₁, i₂, h₂, tr₁⟩ := etr₁
  obtain ⟨j₁, k₁, j₂, k₂, tr₂⟩ := etr₂
  simp only [Graph.loop_inputs, Graph.loop_outputs, List.mem_singleton] at h₁ h₂ k₁ k₂
  subst h₁; subst h₂; subst k₁; subst k₂
  exact ⟨g.loopIn, List.mem_singleton_self _, g.loopOut, List.mem_singleton_self _,
    Trace.concat tr₁ loop_edge_out_in tr₂⟩

theorem EndToEndTrace.loop_empty {ρ : Env} : EndToEndTrace (Graph.loop g) ρ ρ := by
  have hedge : ((g.loopIn, g.loopOut) : (Graph.loop g).Edge) ∈ (Graph.loop g).edges :=
    List.mem_append_right _ (List.mem_cons_of_mem _ (List.mem_cons_self _ _))
  exact ⟨g.loopIn, List.mem_singleton_self _, g.loopOut, List.mem_singleton_self _,
    Trace.concat (Trace.single (loop_nodes_at_in ▸ EvalBasicStmts.nil)) hedge
      (Trace.single (loop_nodes_at_out ▸ EvalBasicStmts.nil))⟩

end Loop

/-! ### Singletons, wrap, and the main result -/

theorem EndToEndTrace.singleton {bss : List BasicStmt} {ρ₁ ρ₂ : Env}
    (h : EvalBasicStmts ρ₁ bss ρ₂) : EndToEndTrace (Graph.singleton bss) ρ₁ ρ₂ :=
  ⟨(0 : Fin 1), List.mem_singleton_self _, (0 : Fin 1), List.mem_singleton_self _,
   Trace.single h⟩

theorem EndToEndTrace.singleton_nil (ρ : Env) :
    EndToEndTrace (Graph.singleton []) ρ ρ :=
  EndToEndTrace.singleton EvalBasicStmts.nil

theorem EndToEndTrace.wrap {g : Graph} {ρ₁ ρ₂ : Env}
    (etr : EndToEndTrace g ρ₁ ρ₂) : EndToEndTrace (Graph.wrap g) ρ₁ ρ₂ :=
  (EndToEndTrace.singleton_nil ρ₁).concat (etr.concat (EndToEndTrace.singleton_nil ρ₂))

theorem buildCfg_sufficient {s : Stmt} {ρ₁ ρ₂ : Env}
    (h : EvalStmt ρ₁ s ρ₂) : EndToEndTrace (buildCfg s) ρ₁ ρ₂ := by
  induction h with
  | basic ρ₁ ρ₂ bs hbs =>
    exact EndToEndTrace.singleton (EvalBasicStmts.cons hbs EvalBasicStmts.nil)
  | andThen ρ₁ ρ₂ ρ₃ s₁ s₂ _ _ ih₁ ih₂ =>
    exact ih₁.concat ih₂
  | ifTrue ρ₁ ρ₂ e z s₁ s₂ _ _ _ ih =>
    exact ih.comp_left
  | ifFalse ρ₁ ρ₂ e s₁ s₂ _ _ ih =>
    exact ih.comp_right
  | whileTrue ρ₁ ρ₂ ρ₃ e z s _ _ _ _ ih₁ ih₂ =>
    exact (ih₁.loop).loop_concat ih₂
  | whileFalse ρ e s _ =>
    exact EndToEndTrace.loop_empty

/-! ### The wrapped graph's entry has no predecessors (Agda's "ugly" block) -/

def Graph.wrapInput (g : Graph) : (Graph.wrap g).Index :=
  (0 : Fin 1).castAdd ((g ⤳ Graph.singleton []).size)

def Graph.wrapOutput (g : Graph) : (Graph.wrap g).Index :=
  Fin.natAdd 1 ((Fin.natAdd g.size (0 : Fin 1)))

theorem Graph.wrap_inputs (g : Graph) :
    (Graph.wrap g).inputs = [g.wrapInput] := rfl

theorem Graph.wrap_outputs (g : Graph) :
    (Graph.wrap g).outputs = [g.wrapOutput] := rfl

private theorem not_mem_edges_castAdd_link {g₂ : Graph} (i : Fin 1)
    (idx : (Graph.singleton [] ⤳ g₂).Index) :
    ((idx, i.castAdd g₂.size) : (Graph.singleton [] ⤳ g₂).Edge)
      ∉ (Graph.singleton [] ⤳ g₂).edges := by
  intro h
  rcases List.mem_append.mp h with h' | h'
  · rcases List.mem_append.mp h' with h'' | h''
    · -- lifted edges of `singleton []`: there are none
      simp [Graph.singleton, Graph.liftEdgeL] at h''
    · -- lifted edges of g₂: targets are natAdd
      obtain ⟨e, _, heq⟩ := List.mem_map.mp h''
      exact Fin.castAdd_ne_natAdd i e.2 (congrArg Prod.snd heq).symm
  · -- product edges: targets are natAdd'd inputs of g₂
    obtain ⟨-, hb⟩ := List.mem_product.mp h'
    obtain ⟨j, -, heq⟩ := List.mem_map.mp hb
    exact Fin.castAdd_ne_natAdd i j heq.symm

theorem Graph.wrap_predecessors_eq_nil (g : Graph) (idx : (Graph.wrap g).Index)
    (h : idx ∈ (Graph.wrap g).inputs) :
    (Graph.wrap g).predecessors idx = [] := by
  rw [Graph.wrap_inputs, List.mem_singleton] at h
  subst h
  rw [Graph.predecessors, List.filter_eq_nil_iff]
  intro idx' _
  simpa using not_mem_edges_castAdd_link (g₂ := g ⤳ Graph.singleton []) 0 idx'

end Spa
