import Spa.Language.Base
import Mathlib.Data.Fin.Tuple.Basic
import Mathlib.Data.List.ProdSigma
import Mathlib.Data.List.FinRange

def List.finCastAdd {n : ℕ} (l : List (Fin n)) (m : ℕ) : List (Fin (n + m)) :=
  l.map (Fin.castAdd m)

def List.finNatAdd {m : ℕ} (l : List (Fin m)) (n : ℕ) : List (Fin (n + m)) :=
  l.map (Fin.natAdd n)

def List.finCastAddProd {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))

def List.finNatAddProd {m : ℕ} (l : List (Fin m × Fin m)) (n : ℕ) :
    List (Fin (n + m) × Fin (n + m)) :=
  l.map (fun e => (e.1.natAdd n, e.2.natAdd n))

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

def comp (g₁ g₂ : Graph) : Graph where
  size := g₁.size + g₂.size
  nodes := Fin.append g₁.nodes g₂.nodes
  edges := g₁.edges.finCastAddProd g₂.size ++ g₂.edges.finNatAddProd g₁.size
  inputs := g₁.inputs.finCastAdd g₂.size ++ g₂.inputs.finNatAdd g₁.size
  outputs := g₁.outputs.finCastAdd g₂.size ++ g₂.outputs.finNatAdd g₁.size

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

def link (g₁ g₂ : Graph) : Graph where
  size := g₁.size + g₂.size
  nodes := Fin.append g₁.nodes g₂.nodes
  edges := g₁.edges.finCastAddProd g₂.size ++ g₂.edges.finNatAddProd g₁.size ++
    (g₁.outputs.finCastAdd g₂.size).product (g₂.inputs.finNatAdd g₁.size)
  inputs := g₁.inputs.finCastAdd g₂.size
  outputs := g₂.outputs.finNatAdd g₁.size

@[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

def loop (g : Graph) : Graph where
  size := 2 + g.size
  nodes := Fin.append (fun _ : Fin 2 => []) g.nodes
  edges := g.edges.finNatAddProd 2 ++
    (g.inputs.finNatAdd 2).map (g.loopIn, ·) ++
    (g.outputs.finNatAdd 2).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

def skipto (g₁ g₂ : Graph) : Graph where
  size := g₁.size + g₂.size
  nodes := Fin.append g₁.nodes g₂.nodes
  edges := g₁.edges.finCastAddProd g₂.size ++ g₂.edges.finNatAddProd g₁.size ++
    (g₁.inputs.finCastAdd g₂.size).product (g₂.inputs.finNatAdd g₁.size)
  inputs := g₁.inputs.finCastAdd g₂.size
  outputs := g₂.inputs.finNatAdd g₁.size

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

def wrap (g : Graph) : Graph :=
  singleton [] ⤳ g ⤳ singleton []

end Graph

open Graph in
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)

def indices : List g.Index := List.finRange g.size

theorem mem_indices (idx : g.Index) : idx ∈ g.indices :=
  List.mem_finRange idx

theorem nodup_indices : g.indices.Nodup :=
  List.nodup_finRange g.size

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

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⟩

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