/-
Copyright (c) 2024 Joseph Tooby-Smith. All rights reserved.
Released under Apache 2.0 license.
Authors: Joseph Tooby-Smith
-/
import HepLean.Meta.Informal.Post
import Mathlib.Lean.CoreM
/-!

# Extracting information about informal definitions and lemmas

To make the dot file for the dependency graph run:
- lake exe informal mkDot
- dot -Tsvg -o ./Docs/graph.svg ./Docs/InformalDot.dot
- or dot -Tpng  -Gdpi=300 -o ./Docs/graph.png ./Docs/InformalDot.dot
-/

open Lean

/- To make private definitions completely invisible, place them in a separate importable file. -/

structure Decl where
  name : Name
  module : Name
  lineNo : Nat
  docString : String

inductive InformalDeclKind
  | def | lemma

instance : ToString InformalDeclKind where
  toString
    | .def => "def"
    | .lemma => "lemma"

structure InformalDecl extends Decl where
  kind : InformalDeclKind
  deps : Array Decl

structure DepDecls where
  private mk::
  private decls : Std.HashMap Name Decl
  formalDecls : Array Decl
  /-- Pairs of informal declarations and their enclosing modules. -/
  informalModuleMap : Array (Name × Array InformalDecl)

private abbrev DeclsM := StateRefT (Std.HashMap Name Decl) CoreM

private def Decl.ofName (name : Name) (module : Option Name := none) : DeclsM Decl := do
  if let some decl := (← get).get? name then
    return decl
  let env ← getEnv
  let decl : Decl := {
    name
    module := module.getD (env.getModuleFor? name).get!
    lineNo := (← findDeclarationRanges? name).get!.range.pos.line
    docString := (← findDocString? env name).getD "No docstring."
  }
  modifyGet fun decls => (decl, decls.insert name decl)

private unsafe def InformalDecl.ofName? (name module : Name) : DeclsM (Option InformalDecl) := do
  unless name.isInternalDetail do
    if let some const := (← getEnv).find? name then
      if Informal.isInformalDef const then
        return ← ofName .def (← evalConst InformalDefinition name).deps
      else if Informal.isInformalLemma const then
        return ← ofName .lemma (← evalConst InformalLemma name).deps
  return none
where
  ofName (kind : InformalDeclKind) (deps : List Name) : DeclsM InformalDecl :=
    return {← Decl.ofName name module with kind, deps := ← deps.toArray.mapM Decl.ofName}

unsafe def DepDecls.ofRootModule (rootModule : Name) : CoreM DepDecls := do
  let (informalModuleMap, decls) ← getAllDecls.run ∅

  let mut informalNames : Std.HashSet Name := ∅
  for (_, informalDecls) in informalModuleMap do
    for {name, ..} in informalDecls do
      informalNames := informalNames.insert name

  let mut formalDecls : Array Decl := #[]
  for (name, decl) in decls do
    unless informalNames.contains name do
      formalDecls := formalDecls.push decl

  return {decls, formalDecls, informalModuleMap}
where
  getAllDecls : DeclsM (Array (Name × Array InformalDecl)) := do
    let ({imports, ..}, _) ← readModuleData (← findOLean rootModule)
    imports.filterMapM fun {module, ..} => do
      if module.getRoot == rootModule then
        let ({constNames, ..}, _) ← readModuleData (← findOLean module)
        let informalDecls ← constNames.filterMapM fun name => InformalDecl.ofName? name module
        unless informalDecls.isEmpty do
          let informalDecls := informalDecls.qsort fun a b => a.lineNo < b.lineNo
          printInformalDecls module informalDecls
          return (module, informalDecls)
      return none
  printInformalDecls (module : Name) (informalDecls : Array InformalDecl) : CoreM Unit := do
    println! module
    for {kind, name, lineNo, docString, deps, ..} in informalDecls do
      println! s!"
Informal {kind}: {name}
- {module.toRelativeFilePath}:{lineNo}
- Description: {docString}
- Dependencies:"
      for {name, module, lineNo, ..} in deps do
        println! s!"    * {name}: {module.toRelativeFilePath}:{lineNo}"

/-- Making the Markdown file for dependency graph. -/
def mkMarkdown (depDecls : DepDecls) : IO Unit := do
  println! "
# Informal definitions and lemmas

See [informal definitions and lemmas as a dependency graph](https://heplean.github.io/HepLean/graph.svg).

This file is automatically generated using `informal_lemma` and `informal_definition` commands
appearing in the raw Lean code of HepLean.

There is an implicit invitation to the reader to contribute to the formalization of
 informal definitions and lemmas. You may also wish to contribute to HepLean by including
 informal definitions and lemmas in the raw Lean code, especially if you do not have a
 background in Lean.

"
  for (module, informalDecls) in depDecls.informalModuleMap do
    println! s!"## {module}"
    for {kind, name, lineNo, docString, deps, ..} in informalDecls do
      println! s!"
**Informal {kind}**: [{name}]({module.toGitHubLink lineNo}) :=
  *{docString}*
- Dependencies:"
      for {name, module, lineNo, ..} in deps do
        println! s!"    * {name}: {module.toRelativeFilePath}:{lineNo}"

/-- Making the DOT file for dependency graph. -/
def mkDOT (depDecls : DepDecls) : IO Unit := do
  IO.println "strict digraph G {
    graph [
    pack=true;
    packmode=\"array1\";
    ];
    tooltip = \"Informal HepLean graph\";
    node [margin=\"0.2,0.05\";  fontsize=10;  fontname=\"Georgia\", height=0.1];
    bgcolor=\"white\";
    labelloc=\"t\";
    labeljust=\"l\";
    edge [arrowhead=vee];"

  let mut clusters : Std.HashSet Name := ∅
  for (_, informalDecls) in depDecls.informalModuleMap do
    for {name, ..} in informalDecls do
      let cluster := name.getPrefix
      unless cluster.isAnonymous do
        clusters := clusters.insert cluster
        println! s!"subgraph cluster_{cluster.toString.replace "." "_"}" ++ "
      {
          label=\"" ++ cluster.toString ++ "\";
          color=steelblue;
              }"

  for decl in depDecls.formalDecls do
    println! toNode clusters decl "box" "steelblue"

  for (_, informalDecls) in depDecls.informalModuleMap do
    for {toDecl := decl, kind, ..} in informalDecls do
      match kind with
      | .def => println! toNode clusters decl "box" "lightgray"
      | .lemma => println! toNode clusters decl "ellipse" "lightgray"

  for (_, informalDecls) in depDecls.informalModuleMap do
    for informalDecl in informalDecls do
      for dep in informalDecl.deps do
        println! s!"\"{dep.name}\" -> \"{informalDecl.name}\""

  println! "}"
where
  toNode (clusters : Std.HashSet Name) (decl : Decl) (shape color : String) : String :=
    let {name, docString, ..} := decl
    let prefixName := if clusters.contains name then name else name.getPrefix
    let nodeStr := s!"\"{name}\"[label=\"{name}\", shape={shape}, style=filled, fillcolor={color},
    tooltip=\"{docString}\"]"
    if prefixName.isAnonymous then
      nodeStr
    else
      s!"subgraph cluster_{prefixName.toString.replace "." "_"} \{ {nodeStr}; }"

/-- Making the HTML file for dependency graph. -/
def mkHTML (depDecls : DepDecls) : IO Unit := do
  IO.println "---
layout: default
---
<!DOCTYPE html>
<html>
<head>
    <title>Informal dependency graph for HepLean</title>
    <style>
        /* Optional styling for the message display */
        #message {
            margin-top: 20px;
            font-size: 1.2em;
            color: #333;
        }

        /* Optional: Hide the text content divs */
        .node-text {
            display: none;
            font-family: 'Times New Roman', Times, serif;
        }
    </style>
</head>
<body>
    <!-- Include D3.js -->
    <script src=\"https://d3js.org/d3.v5.min.js\"></script>
    <!-- Include hpcc-js/wasm for Graphviz layout -->
    <script src=\"https://unpkg.com/@hpcc-js/wasm@0.3.11/dist/index.min.js\"></script>
    <!-- Include d3-graphviz for rendering Graphviz graphs -->
    <script src=\"https://unpkg.com/d3-graphviz@3.0.5/build/d3-graphviz.js\"></script>

    <!-- Add a title to the page -->
    <h1 style=\"text-align: center;\">Informal dependency graph for HepLean</h1>
    <p style=\"text-align: center;\">Click on a node to display the text associated with it.</p>
    <p style=\"text-align: center;\">This graph only shows informal results (gray) and their direct formal dependencies (blue).
    It does not show all results formalised into HepLean.</p>
    <!-- Div to display the graph -->
    <div id=\"graph\" style=\"width: 75vw; height: 75vh; border: 1px solid black\"></div>

    <!-- Hidden divs containing the text to display when nodes are clicked -->

    <!-- Add more divs for other nodes as needed, each with id equal to the node name -->
    <div id=\"includedContent\"></div>
    <!-- Div to display the message when a node is clicked -->
    <div id=\"message\"></div>"

  for decl in depDecls.formalDecls do
    println! toHTML decl

  for (_, informalDecls) in depDecls.informalModuleMap do
    for {toDecl := decl, ..} in informalDecls do
      println! toHTML decl

  IO.println "
  <script type=\"text/javascript\">
        // Load the DOT file and render the graph
        d3.text(\"InformalDot.dot\").then(function(dotSrc) {
            var width = 0.75 * window.innerWidth;
            var height = 0.75 * window.innerHeight;

            // Initialize the graphviz renderer
            var graphviz = d3.select(\"#graph\")
                .graphviz()
                .width(width)
                .height(height)
                .fit(true);

            // Render the graph and attach event listeners
            graphviz.renderDot(dotSrc).on(\"end\", function() {
                // Attach click event to nodes
                d3.selectAll('.node').on('click', function() {
                    // Get the node's name from the title element
                    var nodeName = d3.select(this).select('title').text();

                    // Get the text associated with the nodeName
                    var nodeTextElement = document.getElementById(nodeName);

                    if (nodeTextElement) {
                        var nodeText = nodeTextElement.innerHTML;
                        // Display the text in the #message div
                        document.getElementById('message').innerHTML = nodeText;
                    } else {
                        // If no element with the nodeName ID exists
                        document.getElementById('message').innerHTML = 'No information available for node: ' + nodeName;
                    }
                });
            });
        });
    </script>
</body>
</html>
"
where
  toHTML (decl : Decl) : String :=
    let {name, module, lineNo, docString} := decl
    s!"
<div id=\"{name}\" class=\"node-text\">
<b><a href=\"{module.toGitHubLink lineNo}\">{name}</a></b><br>
<b>Description: </b>{docString}

</div>"

/-!

## Main function

-/

def IO.withStdoutRedirectedTo {α} (filePath : System.FilePath) (action : IO α) : IO α :=
  FS.withFile filePath .write fun handle => withStdout (.ofHandle handle) action

unsafe def main (args : List String) : IO Unit := do
  initSearchPath (← findSysroot)
  let rootModule := `HepLean
  CoreM.withImportModules #[rootModule] do
    /- Build informal dependencies with all information necessary. -/
    let depDecls ← DepDecls.ofRootModule rootModule
    /- Stay inside `CoreM.withImportModules` so that references to objects in `env` are not yet
    destroyed. -/
    if "mkFile" ∈ args then
      IO.withStdoutRedirectedTo "./docs/Informal.md" do mkMarkdown depDecls
      println! "Markdown file made."
    if "mkDot" ∈ args then
      IO.withStdoutRedirectedTo "./docs/InformalDot.dot" do mkDOT depDecls
      println! "DOT file made."
    if "mkHTML" ∈ args then
      IO.withStdoutRedirectedTo "./docs/InformalGraph.html" do mkHTML depDecls
      println! "HTML file made."