blueprintgraph.lean

import ProofWidgets.Demos.Graph.ExprGraph
import ProofWidgets.Component.GraphDisplay
import ProofWidgets.Component.HtmlDisplay
import Lean.Data.Json
import Mathlib.Tactic.Ring
import Architect

open Lean ProofWidgets Jsx

--------------------------------------------------------------------------------
-- 1. JSON parsing structures for LeanArchitect output
--------------------------------------------------------------------------------
structure StatementData where
  uses : Array String := #[]
  text : String := ""
  latexEnv : String := "theorem" 
  deriving FromJson, Inhabited

structure ProofData where
  uses : Array String := #[]
  text : String := ""
  deriving FromJson, Inhabited

structure NodeData where
  name : String
  title : Option String := none
  uses : Option (Array String) := none
  statement : StatementData
  proof : Option ProofData := none
  notReady : Bool := false       
  hasLean : Bool := true         
  deriving FromJson, Inhabited

structure NodeInfo where
  type : String
  data : NodeData
  deriving FromJson, Inhabited

--------------------------------------------------------------------------------
-- 2. ProofWidgets Graph Display Builder
--------------------------------------------------------------------------------
def buildArchitectGraph (nodes : Array NodeInfo) : Html := Id.run do
  let mut vertices : Array GraphDisplay.Vertex := #[]
  let mut edges : Array GraphDisplay.Edge := #[]
  let mut maxRadius := 10

  -- Safelist of valid IDs to prevent D3 from crashing on external links
  let validIds := nodes.map (fun n => n.data.name)

  for node in nodes do
    let d := node.data
    let id := d.name
    
    let labelText := d.title.getD d.name
    let rx := labelText.length * 4
    maxRadius := Nat.max maxRadius rx

    let detailsHtml :=
      <div>
        <b>{.text d.name}</b>
        <hr/>
        <MarkdownDisplay contents={d.statement.text} />
        {match d.proof with
         | some p => <div><b>{.text "Proof:"}</b><MarkdownDisplay contents={p.text} /></div>
         | none => <span/>}
      </div>

    let isDef := d.statement.latexEnv == "definition"
    let hasProof := d.proof.isSome

    let (fillColor, strokeColor) :=
      if d.notReady then
        ("#bbdefb", "#0d47a1")
      else if isDef || hasProof then
        ("#a8e6cf", "#1b5e20")
      else
        ("#ffffff", "#1b5e20")

    let shapeHtml := if isDef then
        <rect
          fill={fillColor}
          stroke={strokeColor}
          strokeWidth="2"
          x={s!"-{rx*2}"}
          y="-15"
          width={s!"{rx*4}"}
          height="30"
          rx="3"
        />
      else
        <ellipse
          fill={fillColor}
          stroke={strokeColor}
          strokeWidth="2"
          rx={(rx*2 : Nat)}
          ry="15"
        />

    let v : GraphDisplay.Vertex := {
      id := id
      label :=
        <g>
          {shapeHtml}
          <text x="0" y="5" textAnchor="middle" className="font-code">{.text labelText}</text>
        </g>
      boundingShape := .rect (rx*4).toFloat 30
      details? := some detailsHtml
    }
    vertices := vertices.push v

    -- FIX: Safely extract top-level uses using `getD` to default to an empty array if null
    let mut allUses := d.uses.getD #[] ++ d.statement.uses
    if let some p := d.proof then
      allUses := allUses ++ p.uses

    -- Only draw the edge if the target node exists in this graph
    for u in allUses do
      if validIds.contains u then
        edges := edges.push { source := u, target := id }

  return <GraphDisplay
    vertices={vertices}
    edges={edges}
    forces={#[
      .link { distance? := Float.ofNat (maxRadius * 3) },
      .collide { radius? := Float.ofNat (maxRadius + 10) },
      .x { strength? := some 0.05 },
      .y { strength? := some 0.05 }
    ]}
    showDetails={true}
  />

--------------------------------------------------------------------------------
-- 3. Intercept Command to Automate Rendering
--------------------------------------------------------------------------------
syntax (name := architectGraphCmd) "#architect_graph" : command

open Lean Elab Command in
@[command_elab architectGraphCmd]
def elabArchitectGraphCmd : CommandElab := fun
  | stx@`(#architect_graph) => do
    let env ← getEnv
    
    let showCmd ← match Parser.runParserCategory env `command "#show_blueprint_json" with
      | .ok s => pure s
      | .error err => throwError "Failed to parse command: {err}"

    let initMsgLog := (← get).messages
    modify fun s => { s with messages := {} }

    try elabCommand showCmd
    catch e => 
      modify fun s => { s with messages := initMsgLog }
      throw e

    let msgs := (← get).messages
    modify fun s => { s with messages := initMsgLog }

    let msgArr := msgs.msgs.toArray
    if msgArr.isEmpty then
      throwError "No output captured from #show_blueprint_json"
    else
      let mut jsonStr := ""
      for msg in msgArr do
        let fmt ← msg.data.format
        let s := fmt.pretty 999999 
        if s.trim.startsWith "[" then
          jsonStr := s
          break
      
      if jsonStr == "" then
        throwError "Could not find a valid JSON array in the output."

      match Json.parse jsonStr with
      | .error err => throwError "Failed to parse JSON: {err}"
      | .ok j =>
        match fromJson? (α := Array NodeInfo) j with
        | .error err => throwError "Failed to decode JSON: {err}"
        | .ok nodes =>
          let html := buildArchitectGraph nodes
          runTermElabM fun _ => do
            Widget.savePanelWidgetInfo
              (hash HtmlDisplayPanel.javascript)
              (return json% { html : $(← Server.rpcEncode html) })
              stx
  | _ => throwUnsupportedSyntax

--------------------------------------------------------------------------------
-- 4. Your Sample Code
--------------------------------------------------------------------------------

def a2 : Nat := 0
def b2 : Nat := 1
def foobar (c : Nat) : Nat × Int := (a2 + b2 * c, a2 / b2)

#expr_graph foobar 

@[blueprint]
inductive MyNat : Type where
  | zero : MyNat
  | succ : MyNat → MyNat

namespace MyNat

@[blueprint
  -- You may manually specify a \label
  "def:nat-add"
  (statement := /-- Natural number addition. -/)]
def add (a b : MyNat) : MyNat :=
  match b with
  | zero => a
  | succ b => succ (add a b)

@[simp, blueprint
  (statement := /-- For any natural number $a$, $0 + a = a$, where $+$ is \cref{def:nat-add}. -/)]
theorem zero_add (a : MyNat) : add zero a = a := by
  /-- The proof follows by induction. -/
  induction a <;> simp [*, add]

@[blueprint
  (statement := /-- For any natural numbers $a, b$, $(a + 1) + b = (a + b) + 1$. -/)]
theorem succ_add (a b : MyNat) : add (succ a) b = succ (add a b) := by
  /-- Proof by induction on $b$. -/
  -- If the proof contains sorry, the `\leanok` command will not be added
  sorry

@[blueprint
  (statement := /-- For any natural numbers $a, b$, $a + b = b + a$. -/)]
theorem add_comm (a b : MyNat) : add a b = add b a := by
  induction b with
  | zero =>
    have := trivial
    /-- The base case follows from \cref{MyNat.zero_add}. -/
    simp [add]
  | succ b ih =>
    /-- The inductive case follows from \cref{MyNat.succ_add}. -/
    sorry_using [succ_add]  -- the `sorry_using` tactic declares that the proof uses succ_add

/-! ## Multiplication -/

@[blueprint
  (uses := [add]) -- Manually added dependency
  (statement := /-- Natural number multiplication. -/)]
def mul (a b : MyNat) : MyNat := sorry

@[blueprint
  (statement := /-- For any natural numbers $a, b$, $a * b = b * a$. -/)]
theorem mul_comm (a b : MyNat) : mul a b = mul b a := by sorry

/-! ## Fermat's Last Theorem -/

@[blueprint "thm:flt"
  (statement := /-- Fermat's last theorem. -/)
  (title := "Taylor-Wiles")
  -- You may override the inferred statement dependencies by `uses`.
  (uses := [mul])
  -- Alternatively to docstring tactics and `using` tactics, proof metadata can be specified
  -- by `proof` and `proofUses`.
  (proof := /-- See \cite{Wiles1995, Taylor-Wiles1995}. -/) (proofUses := [mul_comm])
  (notReady := true) (discussion := 1)]
theorem flt : (sorry : Prop) := sorry

end MyNat

--------------------------------------------------------------------------------
-- 5. View the Dependency Graph!
--------------------------------------------------------------------------------
#show_blueprint_json
#architect_graph