/-
Port of `Language/Graphs.agda`.

Representation note: `nodes : Vec (List BasicStmt) size` becomes
`nodes : Fin size → List BasicStmt`. With that, the `Data.Vec` lookup/append
lemma stack (`lookup-++ˡ/ʳ`, `cast-is-id`, …) lifts into mathlib's
`Fin.append` with `Fin.append_left` / `Fin.append_right`.

Correspondence:
  _↑ˡ_/_↑ʳ_ (on Fin)  ↦ Fin.castAdd / Fin.natAdd (mathlib)
  _↑ˡⁱ_/_↑ʳⁱ_         ↦ liftIdxL / liftIdxR
  _↑ˡᵉ_/_↑ʳᵉ_         ↦ liftEdgeL / liftEdgeR
  _∙_                 ↦ Graph.comp   (scoped notation ∙)
  _↦_                 ↦ Graph.link   (scoped notation ⤳)
  loop                ↦ Graph.loop
  _skipto_            ↦ Graph.skipto
  _[_]                ↦ Graph.nodes (plain application)
  singleton, wrap     ↦ Graph.singleton, Graph.wrap
  buildCfg            ↦ buildCfg
  indices             ↦ List.finRange (mathlib; `fins` from Utils.agda)
  indices-complete    ↦ List.mem_finRange
  indices-Unique      ↦ List.nodup_finRange
  predecessors        ↦ Graph.predecessors
  edge⇒predecessor    ↦ Graph.mem_predecessors_of_edge
  predecessor⇒edge    ↦ Graph.edge_of_mem_predecessors
-/
import Spa.Language.Base
import Mathlib.Data.Fin.Tuple.Basic
import Mathlib.Data.List.ProdSigma
import Mathlib.Data.List.FinRange

namespace Spa

structure Graph where
  size : ℕ
  nodes : Fin size → List BasicStmt
  edges : List (Fin size × Fin size)
  inputs : List (Fin size)
  outputs : List (Fin size)

namespace Graph

abbrev Index (g : Graph) : Type := Fin g.size

abbrev Edge (g : Graph) : Type := g.Index × g.Index

/-- Agda: `_↑ˡⁱ_`. -/
def liftIdxL {n : ℕ} (l : List (Fin n)) (m : ℕ) : List (Fin (n + m)) :=
  l.map (Fin.castAdd m)

/-- Agda: `_↑ʳⁱ_`. -/
def liftIdxR (n : ℕ) {m : ℕ} (l : List (Fin m)) : List (Fin (n + m)) :=
  l.map (Fin.natAdd n)

/-- Agda: `_↑ˡᵉ_` (with `_↑ˡ_` on pairs inlined). -/
def liftEdgeL {n : ℕ} (l : List (Fin n × Fin n)) (m : ℕ) :
    List (Fin (n + m) × Fin (n + m)) :=
  l.map (fun e => (e.1.castAdd m, e.2.castAdd m))

/-- Agda: `_↑ʳᵉ_` (with `_↑ʳ_` on pairs inlined). -/
def liftEdgeR (n : ℕ) {m : ℕ} (l : List (Fin m × Fin m)) :
    List (Fin (n + m) × Fin (n + m)) :=
  l.map (fun e => (e.1.natAdd n, e.2.natAdd n))

/-- Agda: `_∙_` — disjoint union. -/
def comp (g₁ g₂ : Graph) : Graph where
  size := g₁.size + g₂.size
  nodes := Fin.append g₁.nodes g₂.nodes
  edges := liftEdgeL g₁.edges g₂.size ++ liftEdgeR g₁.size g₂.edges
  inputs := liftIdxL g₁.inputs g₂.size ++ liftIdxR g₁.size g₂.inputs
  outputs := liftIdxL g₁.outputs g₂.size ++ liftIdxR g₁.size g₂.outputs

@[inherit_doc] scoped infixr:70 " ∙ " => Graph.comp

/-- Agda: `_↦_` — sequencing: all outputs of `g₁` feed all inputs of `g₂`. -/
def link (g₁ g₂ : Graph) : Graph where
  size := g₁.size + g₂.size
  nodes := Fin.append g₁.nodes g₂.nodes
  edges := liftEdgeL g₁.edges g₂.size ++ liftEdgeR g₁.size g₂.edges ++
    (liftIdxL g₁.outputs g₂.size).product (liftIdxR g₁.size g₂.inputs)
  inputs := liftIdxL g₁.inputs g₂.size
  outputs := liftIdxR g₁.size g₂.outputs

@[inherit_doc] scoped infixr:70 " ⤳ " => Graph.link

/-- The entry node of a `loop` graph. -/
def loopIn (g : Graph) : Fin (2 + g.size) := (0 : Fin 2).castAdd g.size

/-- The exit node of a `loop` graph. -/
def loopOut (g : Graph) : Fin (2 + g.size) := (1 : Fin 2).castAdd g.size

/-- Agda: `loop`. -/
def loop (g : Graph) : Graph where
  size := 2 + g.size
  nodes := Fin.append (fun _ : Fin 2 => []) g.nodes
  edges := liftEdgeR 2 g.edges ++
    (liftIdxR 2 g.inputs).map (g.loopIn, ·) ++
    (liftIdxR 2 g.outputs).map (·, g.loopOut) ++
    [(g.loopOut, g.loopIn), (g.loopIn, g.loopOut)]
  inputs := [g.loopIn]
  outputs := [g.loopOut]

@[simp] theorem loop_inputs (g : Graph) : (loop g).inputs = [g.loopIn] := rfl

@[simp] theorem loop_outputs (g : Graph) : (loop g).outputs = [g.loopOut] := rfl

/-- Agda: `_skipto_` (unused by `buildCfg`, ported for parity). -/
def skipto (g₁ g₂ : Graph) : Graph where
  size := g₁.size + g₂.size
  nodes := Fin.append g₁.nodes g₂.nodes
  edges := liftEdgeL g₁.edges g₂.size ++ liftEdgeR g₁.size g₂.edges ++
    (liftIdxL g₁.inputs g₂.size).product (liftIdxR g₁.size g₂.inputs)
  inputs := liftIdxL g₁.inputs g₂.size
  outputs := liftIdxR g₁.size g₂.inputs

/-- Agda: `singleton`. -/
def singleton (bss : List BasicStmt) : Graph where
  size := 1
  nodes := fun _ => bss
  edges := []
  inputs := [0]
  outputs := [0]

/-- Agda: `wrap`. -/
def wrap (g : Graph) : Graph :=
  singleton [] ⤳ g ⤳ singleton []

end Graph

open Graph in
/-- Agda: `buildCfg`. -/
def buildCfg : Stmt → Graph
  | .basic bs => Graph.singleton [bs]
  | .andThen s₁ s₂ => buildCfg s₁ ⤳ buildCfg s₂
  | .ifElse _ s₁ s₂ => buildCfg s₁ ∙ buildCfg s₂
  | .whileLoop _ s => Graph.loop (buildCfg s)

namespace Graph

variable (g : Graph)

/-- Agda: `indices` (`fins` is mathlib's `List.finRange`). -/
def indices : List g.Index := List.finRange g.size

/-- Agda: `indices-complete`. -/
theorem mem_indices (idx : g.Index) : idx ∈ g.indices :=
  List.mem_finRange idx

/-- Agda: `indices-Unique`. -/
theorem nodup_indices : g.indices.Nodup :=
  List.nodup_finRange g.size

/-- Agda: `predecessors`. -/
def predecessors (idx : g.Index) : List g.Index :=
  g.indices.filter (fun idx' => (idx', idx) ∈ g.edges)

/-- Agda: `edge⇒predecessor`. -/
theorem mem_predecessors_of_edge {idx₁ idx₂ : g.Index}
    (h : (idx₁, idx₂) ∈ g.edges) : idx₁ ∈ g.predecessors idx₂ :=
  List.mem_filter.mpr ⟨g.mem_indices idx₁, by simpa using h⟩

/-- Agda: `predecessor⇒edge`. -/
theorem edge_of_mem_predecessors {idx₁ idx₂ : g.Index}
    (h : idx₁ ∈ g.predecessors idx₂) : (idx₁, idx₂) ∈ g.edges := by
  simpa using (List.mem_filter.mp h).2

end Graph

end Spa