import Table.API
import Table.ExampleTables
import Table.Notation

-- Table equality typeclass resolution requires a lot of instances
set_option synthInstance.maxSize 12000
set_option synthInstance.maxHeartbeats 0

-- TODO: come up with a better testing system
-- Modified version of `elabEvalUnsafe` (src/lean/lean/elab/builtincommand.lean)
syntax (name := test) "#test" term : command
@[commandElab test]
unsafe def elabTest : Lean.Elab.Command.CommandElab
| `(#test%$tk $term) => do
    let n := `_eval
    let ctx ← read
    let addAndCompile (value : Lean.Expr) : Lean.Elab.TermElabM Unit := do
      -- the type really should be `Bool` at this point (b/c of `mkDecide`)
      -- (but could enforcing that explicitly lead to less-graceful failures?)
      let type ← Lean.Meta.inferType value
      let us := Lean.collectLevelParams {} value |>.params
      -- let value ← Lean.Meta.instantiateMVars value
      let decl := Lean.Declaration.defnDecl {
        name        := n
        levelParams := us.toList
        type        := type
        value       := value
        hints       := Lean.ReducibilityHints.opaque
        safety      := Lean.DefinitionSafety.unsafe
      }
      Lean.Elab.Term.ensureNoUnassignedMVars decl
      Lean.addAndCompile decl
    let elabEvalTerm : Lean.Elab.TermElabM Lean.Expr := do
      let ebool ← Lean.Elab.Term.elabTerm (← `(Bool)) none
      let e ← Lean.Elab.Term.elabTerm term none
      Lean.Elab.Term.synthesizeSyntheticMVarsNoPostponing

      -- Need to do this here so we can ensure the type is correct and return
      -- a meaningful error message otherwise
      let (e, _) ← Lean.Elab.Term.levelMVarToParam
                    (← Lean.Meta.instantiateMVars e)
      let e_type ← Lean.Meta.inferType e
      if (← Lean.Meta.isProp e) then
        Lean.Meta.mkDecide e
      else if (← Lean.Meta.isDefEq e_type ebool) then
        return e
      else
        throwError m!"Tests must be of type Bool or Prop, but got '{e_type}'"
    let elabEval : Lean.Elab.Command.CommandElabM Unit :=
    Lean.Elab.Command.runTermElabM (some n) (λ _ => do
      let e ← elabEvalTerm
      let env ← Lean.getEnv
      let res ← try addAndCompile e; Lean.evalConst Bool n
                finally Lean.setEnv env
      if res
      then Lean.Elab.logInfoAt tk "Test passed"
      else Lean.Elab.logErrorAt tk "Test failed")
    elabEval
| _ => Lean.Elab.throwUnsupportedSyntax

-- Ways of making type class inference work where Lean struggles
def instHint (α : Type _) (x : α) (y : α) (inst : DecidableEq α) :=
  decide (x = y)

macro "inst" : tactic =>
  `(repeat (first
    | apply instDecidableEqTable (inst := _)
    | apply instDecidableEqRowConsHeaderMkType (it := _) (ic := _) (ir := _)
    | apply instDecidableEqRowNilHeader
    | apply instDecidableEqCell (inst := _)
    | infer_instance))

notation lhs "=(" tp ")" rhs => instHint tp lhs rhs inferInstance

notation lhs "=[" inst "]" rhs => instHint _ lhs rhs inst

notation lhs "=(" tp ")[" inst "]" rhs => instHint tp lhs rhs inst

-- `addRows`
#test
addRows students [/["Colton", 19, "blue"]]
=
Table.mk [
  /["Bob"  , 12, "blue" ],
  /["Alice", 17, "green"],
  /["Eve"  , 13, "red"  ],
  /["Colton", 19, "blue"]
]

#test
addRows gradebook []
=
Table.mk [
  /["Bob"  , 12, 8, 9, 77, 7, 9, 87],
  /["Alice", 17, 6, 8, 88, 8, 7, 85],
  /["Eve"  , 13, 7, 9, 84, 8, 8, 77]
]

-- `addColumn`
def hairColor := [some "brown", some "red", some "blonde"]
#test
addColumn students "hair-color" hairColor
=[by inst]
Table.mk [
  /[ "Bob"   , 12  , "blue"         , "brown"    ],
  /[ "Alice" , 17  , "green"        , "red"      ],
  /[ "Eve"   , 13  , "red"          , "blonde"   ]
]

def presentation := [some 9, some 9, some 6]
#test
addColumn gradebook "presentation" presentation
=[by inst]
Table.mk [
  /[ "Bob"  , 12, 8, 9, 77, 7, 9, 87, 9],
  /[ "Alice", 17, 6, 8, 88, 8, 7, 85, 9],
  /[ "Eve"  , 13, 7, 9, 84, 8, 8, 77, 6]
]

-- `buildColumn`
def isTeenagerBuilder := λ (r : Row $ schema students) =>
  match getValue r "age" (by header) with
  | some age => some (12 < age && age < 20)
  | _ => some false
#test
buildColumn students "is-teenager" isTeenagerBuilder
=[by inst]
Table.mk [
  /[ "Bob"   , 12  , "blue"         , false       ],
  /[ "Alice" , 17  , "green"        , true        ],
  /[ "Eve"   , 13  , "red"          , true        ]
]

def didWellOnFinal : Row (schema gradebook) → Option Bool := λ r =>
  match getValue r "final" (by header) with
  | some score => some $ 85 <= score
  | _ => some false
#test
buildColumn gradebook "did-well-on-final" didWellOnFinal
=[by inst]
Table.mk [
  /[ "Bob"  , 12, 8, 9, 77, 7, 9, 87, true ],
  /[ "Alice", 17, 6, 8, 88, 8, 7, 85, true ],
  /[ "Eve"  , 13, 7, 9, 84, 8, 8, 77, false]
]

-- `vcat`
def increaseAge := λ (r : Row $ schema students) =>
  ((
    Row.cons
      (Cell.fromOption $ (getValue r "age" (by header)).map (λ x => x + 1))
      Row.nil
  )
  : Row [("age", Nat)])

#test
vcat students (update [⟨("age", Nat), by header⟩] students increaseAge)
=
Table.mk [
  /[ "Bob"   , 12  , "blue"         ],
  /[ "Alice" , 17  , "green"        ],
  /[ "Eve"   , 13  , "red"          ],
  /[ "Bob"   , 13  , "blue"         ],
  /[ "Alice" , 18  , "green"        ],
  /[ "Eve"   , 14  , "red"          ]
]

def curveMidtermAndFinal := λ (r : Row $ schema gradebook) =>
  let curve := λ | some n => some (n + 5)
                 | _ => none
  let midterm := getValue r "midterm" (by header)
  let final := getValue r "final" (by header)
  (Row.cons (Cell.fromOption $ curve midterm) $
     Row.cons (Cell.fromOption $ curve final) Row.nil
   : Row [("midterm", Nat), ("final", Nat)])

#test
vcat gradebook (update [⟨("midterm", Nat), by header⟩,
                        ⟨("final", Nat), by header⟩] gradebook
                                                     curveMidtermAndFinal)
=
Table.mk [
  /[ "Bob"   , 12  , 8     , 9     , 77      , 7     , 9     , 87    ],
  /[ "Alice" , 17  , 6     , 8     , 88      , 8     , 7     , 85    ],
  /[ "Eve"   , 13  , 7     , 9     , 84      , 8     , 8     , 77    ],
  /[ "Bob"   , 12  , 8     , 9     , 82      , 7     , 9     , 92    ],
  /[ "Alice" , 17  , 6     , 8     , 93      , 8     , 7     , 90    ],
  /[ "Eve"   , 13  , 7     , 9     , 89      , 8     , 8     , 82    ]
]

-- `hcat`
#test
hcat students (dropColumns gradebook A[⟨"name", by name⟩, ⟨"age", by name⟩])
=
Table.mk [
  /[ "Bob"   , 12  , "blue"         , 8     , 9     , 77      , 7     , 9     , 87    ],
  /[ "Alice" , 17  , "green"        , 6     , 8     , 88      , 8     , 7     , 85    ],
  /[ "Eve"   , 13  , "red"          , 7     , 9     , 84      , 8     , 8     , 77    ]
]

#test
hcat (dropColumns students A[⟨"name", by name⟩, ⟨"age", by name⟩]) gradebook
=
Table.mk [
  /[ "blue"         , "Bob"   , 12  , 8     , 9     , 77      , 7     , 9     , 87    ],
  /[ "green"        , "Alice" , 17  , 6     , 8     , 88      , 8     , 7     , 85    ],
  /[ "red"          , "Eve"   , 13  , 7     , 9     , 84      , 8     , 8     , 77    ]
]

-- `values`
#test
values [/["Alice"], /["Bob"]]
=
(Table.mk [
  /["Alice"],
  /["Bob"]
] : Table [("name", String)])

#test
values [/["Alice", 12], /["Bob", 13]]
=
(Table.mk [
  /["Alice", 12],
  /["Bob", 13]
] : Table [("name", String), ("age", Nat)])

-- `crossJoin`
def petiteJelly :=
selectRows1 (selectColumns2 jellyAnon [⟨0, by simp⟩, ⟨1, by simp⟩, ⟨2, by simp⟩])
            [⟨0, by simp⟩, ⟨1, by simp⟩]
#test
crossJoin students petiteJelly
=[by simp [List.nth, List.nths, List.map]; inst]
Table.mk [
  /[ "Bob"   , 12  , "blue"         , true     , false , false ],
  /[ "Bob"   , 12  , "blue"         , true     , false , true  ],
  /[ "Alice" , 17  , "green"        , true     , false , false ],
  /[ "Alice" , 17  , "green"        , true     , false , true  ],
  /[ "Eve"   , 13  , "red"          , true     , false , false ],
  /[ "Eve"   , 13  , "red"          , true     , false , true  ]
]

#test
crossJoin emptyTable petiteJelly
=[by simp [List.nth, List.nths, List.map]; inst]
Table.mk []

-- `leftJoin`
-- TODO: we need the `header` (and probably `name`) tactic to be able to "see
-- through" the various Schema ActionList functions like `removeOtherDecCH`
#test
leftJoin students gradebook A[⟨("name", _), inferInstance, by header, by header⟩,
                               ⟨("age", _), inferInstance, by simp only [Schema.removeOtherDecCH]; header, by header⟩]
=(Table [("name", String), ("age", Nat), ("favorite color", String),
         ("quiz1", Nat), ("quiz2", Nat), ("midterm", Nat), ("quiz3", Nat),
         ("quiz4", Nat), ("final", Nat)])
Table.mk [
  /[ "Bob"   , 12  , "blue"         , 8     , 9     , 77      , 7     , 9     , 87    ],
  /[ "Alice" , 17  , "green"        , 6     , 8     , 88      , 8     , 7     , 85    ],
  /[ "Eve"   , 13  , "red"          , 7     , 9     , 84      , 8     , 8     , 77    ]
]

#test
leftJoin employees departments A[⟨("Department ID", _), inferInstance, by header, by header⟩]
=
Table.mk [
  /[ "Rafferty"   , 31            , "Sales"         ],
  /[ "Jones"      , 32            , EMP             ],
  /[ "Heisenberg" , 33            , "Engineering"   ],
  /[ "Robinson"   , 34            , "Clerical"      ],
  /[ "Smith"      , 34            , "Clerical"      ],
  /[ "Williams"   , EMP           , EMP             ]
]

#eval leftJoin
(Table.mk [
  /["name" := "Bob", "age" := 18]
])
(Table.mk [
  /["name" := "Bob", "location" := "USA"],
  /["name" := "Bob", "location" := "UK"]
])
A[⟨("name", _), inferInstance, by header, by header⟩]

-- `nrows`
#test nrows (@emptyTable String _) = 0

#test nrows studentsMissing = 3

-- `ncols`
#test ncols students = 3

#test ncols studentsMissing = 3

-- `header`
#test header students = ["name", "age", "favorite color"]

#test
header gradebook
=
["name", "age", "quiz1", "quiz2", "midterm", "quiz3", "quiz4", "final"]

-- `getRow`
#test
getRow students 0 (by simp)
=
/["Bob", 12, "blue"]

#test
getRow gradebook 1 (by simp)
=
/["Alice", 17, 6, 8, 88, 8, 7, 85]

-- `getValue`
#test
getValue /["name" := "Bob", "age" := 12] "name" (by header)
=
some "Bob"

#test
getValue /["name" := "Bob", "age" := 12] "age" (by header)
=
some 12

-- `getColumn1`
#test
getColumn1 students 1 (by simp)
=(List $ Option Nat)
[some 12, some 17, some 13]

#test
getColumn1 gradebook 0 (by simp)
=(List $ Option String)
[some "Bob", some "Alice", some "Eve"]

-- `getColumn2`
#test
getColumn2 students "age" (by header)
=
[some 12, some 17, some 13]

#test
getColumn2 gradebook "name" (by header)
=
[some "Bob", some "Alice", some "Eve"]

-- `selectRows1`
#test
selectRows1 students [⟨2, by simp⟩, ⟨0, by simp⟩, ⟨2, by simp⟩, ⟨1, by simp⟩]
=
Table.mk [
  /[ "Eve"   , 13  , "red"          ],
  /[ "Bob"   , 12  , "blue"         ],
  /[ "Eve"   , 13  , "red"          ],
  /[ "Alice" , 17  , "green"        ]
]

#test
selectRows1 gradebook [⟨2, by simp⟩, ⟨1, by simp⟩]
=
Table.mk [
  /[ "Eve"   , 13  , 7     , 9     , 84      , 8     , 8     , 77    ],
  /[ "Alice" , 17  , 6     , 8     , 88      , 8     , 7     , 85    ]
]

-- `selectRows2`
#test
selectRows2 students [true, false, true] (by simp)
=
Table.mk [
  /[ "Bob" , 12  , "blue"         ],
  /[ "Eve" , 13  , "red"          ]
]

#test
selectRows2 gradebook [false, false, true] (by simp)
=
Table.mk [/["Eve", 13, 7, 9, 84, 8, 8, 77]]

-- `selectColumns1`
#test
selectColumns1 students [true, true, false] (by simp)
=[by inst]
Table.mk [
  /[ "Bob"   , 12  ],
  /[ "Alice" , 17  ],
  /[ "Eve"   , 13  ]
]

#test
selectColumns1 gradebook [true, false, false, false, true, false, false, true] (by simp)
=[by inst]
Table.mk [
  /[ "Bob"   , 77      , 87    ],
  /[ "Alice" , 88      , 85    ],
  /[ "Eve"   , 84      , 77    ]
]

-- `selectColumns2`
#test
selectColumns2 students [⟨2, by simp⟩, ⟨1, by simp⟩]
=(Table [("favorite color", String), ("age", Nat)])
Table.mk [
  /[ "blue"         , 12  ],
  /[ "green"        , 17  ],
  /[ "red"          , 13  ]
]

#test
selectColumns2 gradebook [⟨7, by simp⟩, ⟨0, by simp⟩, ⟨4, by simp⟩]
=(Table [(_, Nat), (_, String), (_, Nat)])
Table.mk [
  /[ 87    , "Bob"   , 77      ],
  /[ 85    , "Alice" , 88      ],
  /[ 77    , "Eve"   , 84      ]
]

-- `selectColumns3`
#test
selectColumns3 students [⟨("favorite color", _), by header⟩, ⟨("age", _), by header⟩]
=(Table [("favorite color", String), ("age", Nat)])
Table.mk [
  /[ "blue"         , 12  ],
  /[ "green"        , 17  ],
  /[ "red"          , 13  ]
]

#test
selectColumns3 gradebook [⟨("final", _), by header⟩, ⟨("name", _), by header⟩, ⟨("midterm", _), by header⟩]
=(Table [("final", Nat), ("name", String), ("midterm", Nat)])
Table.mk [
  /[ 87    , "Bob"   , 77      ],
  /[ 85    , "Alice" , 88      ],
  /[ 77    , "Eve"   , 84      ]
]

-- `head`
#test
head students ⟨1, by simp⟩
=
Table.mk [ /["Bob", 12, "blue"]]

#test
head students ⟨-2, by simp⟩
=
Table.mk [ /["Bob", 12, "blue"]]

-- `distinct`
#test
distinct students
=
Table.mk [
  /[ "Bob"   , 12  , "blue"         ],
  /[ "Alice" , 17  , "green"        ],
  /[ "Eve"   , 13  , "red"          ]
]

#test
distinct (selectColumns3 gradebook [⟨("quiz3", _), by header⟩])
=(Table [("quiz3", Nat)])
Table.mk [ /[7], /[8] ]

-- `dropColumn`
#test
dropColumn students ⟨"age", by name⟩
=(Table [("name", String), ("favorite color", String)])
Table.mk [
  /[ "Bob"   , "blue"         ],
  /[ "Alice" , "green"        ],
  /[ "Eve"   , "red"          ]
]

#test
dropColumn gradebook ⟨"final", by name⟩
=(Table [("name", String), ("age", Nat), ("quiz1", Nat), ("quiz2", Nat), ("midterm", Nat), ("quiz3", Nat), ("quiz4", Nat)])
Table.mk [
  /[ "Bob"   , 12  , 8     , 9     , 77      , 7     , 9     ],
  /[ "Alice" , 17  , 6     , 8     , 88      , 8     , 7     ],
  /[ "Eve"   , 13  , 7     , 9     , 84      , 8     , 8     ]
]

-- `dropColumns`
#test
dropColumns students A[⟨"age", by name⟩]
=(Table [("name", String), ("favorite color", String)])
Table.mk [
  /[ "Bob"   , "blue"         ],
  /[ "Alice" , "green"        ],
  /[ "Eve"   , "red"          ]
]

#test
dropColumns gradebook A[⟨"final", by name⟩, ⟨"midterm", by name⟩]
=
Table.mk [
  /[ "Bob"   , 12  , 8     , 9     , 7     , 9     ],
  /[ "Alice" , 17  , 6     , 8     , 8     , 7     ],
  /[ "Eve"   , 13  , 7     , 9     , 8     , 8     ]
]

-- `tfilter`
def ageUnderFifteen : (Row $ schema students) → Bool := λ r =>
  match getValue r "age" (by header) with
  | some a => a < 15
  | _ => false

#test
tfilter students ageUnderFifteen
=
Table.mk [
  /[ "Bob" , 12  , "blue"         ],
  /[ "Eve" , 13  , "red"          ]
]

def nameLongerThan3Letters : (Row $ schema gradebook) → Bool := λ r =>
  match getValue r "name" (by header) with
  | some name => String.length name > 3
  | _ => false

#test
tfilter gradebook nameLongerThan3Letters
=
Table.mk [/["Alice", 17, 6, 8, 88, 8, 7, 85]]

-- `tsort`
#test
tsort students ⟨"age", by header⟩ true
=
Table.mk [
  /[ "Bob"   , 12  , "blue"         ],
  /[ "Eve"   , 13  , "red"          ],
  /[ "Alice" , 17  , "green"        ]
]

#test
tsort gradebook ⟨"final", by header⟩ false
=
Table.mk [
  /[ "Bob"   , 12  , 8     , 9     , 77      , 7     , 9     , 87    ],
  /[ "Alice" , 17  , 6     , 8     , 88      , 8     , 7     , 85    ],
  /[ "Eve"   , 13  , 7     , 9     , 84      , 8     , 8     , 77    ]
]

-- `sortByColumns`
#test
sortByColumns students [⟨("age", Nat), by header, inferInstance⟩]
=
Table.mk [
  /[ "Bob"   , 12  , "blue"         ],
  /[ "Eve"   , 13  , "red"          ],
  /[ "Alice" , 17  , "green"        ]
]

#test
sortByColumns gradebook [⟨("quiz2", Nat), by header, inferInstance⟩,
                         ⟨("quiz1", Nat), by header, inferInstance⟩]
=
Table.mk [
  /[ "Alice" , 17  , 6     , 8     , 88      , 8     , 7     , 85    ],
  /[ "Eve"   , 13  , 7     , 9     , 84      , 8     , 8     , 77    ],
  /[ "Bob"   , 12  , 8     , 9     , 77      , 7     , 9     , 87    ]
]

-- `orderBy`
def nameLengthOB (r : Row $ schema students) :=
  match getValue r "name" (by header) with
  | some s => String.length s
  | _ => 0

def leOB := (λ (a : Nat) b => decide $ a ≤ b)

#test
orderBy students [⟨_, nameLengthOB, leOB⟩]
=
Table.mk [
  /[ "Bob"   , 12  , "blue"         ],
  /[ "Eve"   , 13  , "red"          ],
  /[ "Alice" , 17  , "green"        ]
]

def geOB := (λ (a : Nat) b => decide $ a ≥ b)

def nameLengthOB' (r : Row $ schema gradebook) :=
  match getValue r "name" (by header) with
  | some s => String.length s
  | _ => 0

def averageOB (xs : List $ Option Nat) :=
  List.foldl (λ acc => λ | none => acc | some x => x + acc) 0 xs / xs.length

def midtermAndFinalOB (r : Row $ schema gradebook) : List $ Option Nat :=
  [getValue r "midterm" (by header), getValue r "final" (by header)]

def compareGradeOB (g1 : List $ Option Nat) (g2 : List $ Option Nat) :=
  leOB (averageOB g1) (averageOB g2)

#test
orderBy gradebook [⟨_, nameLengthOB', geOB⟩, ⟨_, midtermAndFinalOB, compareGradeOB⟩]
=
Table.mk [
  /[ "Alice" , 17  , 6     , 8     , 88      , 8     , 7     , 85    ],
  /[ "Eve"   , 13  , 7     , 9     , 84      , 8     , 8     , 77    ],
  /[ "Bob"   , 12  , 8     , 9     , 77      , 7     , 9     , 87    ]
]

-- `count`
#test
count students ⟨"favorite color", by header⟩
=
Table.mk [
  /[ some "blue"  , 1     ],
  /[ some "green" , 1     ],
  /[ some "red"   , 1     ]
]

#test
count gradebook ⟨"age", by header⟩
=
Table.mk [
  /[ some 12    , 1     ],
  /[ some 17    , 1     ],
  /[ some 13    , 1     ]
]

-- `bin`
#test
bin students ⟨"age", by header⟩ ⟨5, by simp⟩
=
Table.mk [
  /[ "10 <= age < 15" , 2     ],
  /[ "15 <= age < 20" , 1     ]
]

-- TODO: tell B2T2 that there's a typo in this test ("final," not "age")
#test
bin gradebook ⟨"final", by header⟩ ⟨5, by simp⟩
=
Table.mk [
  /[ "75 <= final < 80" , 1     ],
  /[ "80 <= final < 85" , 0     ],
  /[ "85 <= final < 90" , 2     ]
]

-- `pivotTable`
def oAverage (xs : List $ Option Nat) : Option Nat := some $
  List.foldl (λ acc => λ | none => acc | some x => x + acc) 0 xs / xs.length

#test
pivotTable students [⟨("favorite color", _), by header⟩] (by inst) [⟨("age-average", _), ⟨("age", _), by header⟩, oAverage⟩]
=[by inst]
Table.mk [
  /[ "blue"         , 12          ],
  /[ "green"        , 17          ],
  /[ "red"          , 13          ]
]

-- Slightly modified since we aren't using decimals
def proportion (bs : List $ Option Bool) : Option Nat := some $
  (100 * (bs.filter (λ | some true => true | _ => false)).length) / bs.length

-- TODO: does order matter?
#test
pivotTable
  jellyNamed
  [⟨("get acne", Bool), by header⟩, ⟨("brown", _), by header⟩]
  (by inst)
  [⟨("red-proportion", _), ⟨("red", _), by header⟩, proportion⟩,
   ⟨("pink-proportion", _), ⟨("pink", _), by header⟩, proportion⟩]
=[by inst]
Table.mk [
  /[ false    , false , 0              , 75              ],
  /[ false    , true  , 100            , 100             ],
  /[ true     , false , 0              , 25              ],
  /[ true     , true  , 0              , 0               ]
]

-- `groupBy`
-- TODO: handle `none` case?
def colorTemp : (Row $ schema students) → String := λ r =>
  match getValue r "favorite color" (by header) with
  | some "red" => "warm"
  | _ => "cool"

def nameLength : (Row $ schema students) → Nat := λ r =>
  match getValue r "name" (by header) with
  | some s => String.length s
  | _ => 0

def average (xs : List Nat) := List.foldl (·+·) 0 xs / xs.length

def aggregate := λ (k : String) vs =>
/["key" := k, "average" := average vs]

-- TODO: Need to double-check, but this seems to me like the correct output --
-- the first row matches "cool," so order preservation would tell us that this
-- order (and not the reversed order in the B2T2 docs) is most reasonable
-- (B2T2 TS is using unordered maps in the implementation of `groupBy`, which is
-- probably why its order is different)
#test
groupBy students colorTemp nameLength aggregate
=
Table.mk [
  /[ "cool" , 4       ],
  /[ "warm" , 3       ]
]

def abstractAge := λ (r : Row $ schema gradebook) =>
  match getValue r "age" (by header) with
  | some age =>
    match (age ≤ 12 : Bool), (age ≤ 19 : Bool) with
    | true, _ => "kid"
    | _, true => "teenager"
    | _, _ => "adult"
  | _ => ""

def finalGrade := λ (r : Row $ schema gradebook) =>
  match getValue r "final" (by header) with
  | some grade => grade
  | _ => 0

#test
groupBy gradebook abstractAge finalGrade aggregate
=
Table.mk [
  /[ "kid"      , 87      ],
  /[ "teenager" , 81      ]
]

-- `completeCases`
#test
completeCases students ⟨"age", by header⟩
=
[true, true, true]

#test
completeCases studentsMissing ⟨"age", by header⟩
=
[false, true, true]

-- `dropna`
#test
dropna studentsMissing
=
Table.mk [/["Alice", 17, "green"]]

#test
dropna gradebookMissing
=
Table.mk [/["Bob", 12, 8, 9, 77, 7, 9, 87]]

-- `fillna`
#test
fillna studentsMissing ⟨"favorite color", by header⟩ "white"
=
Table.mk [
  /[ "Bob"   , EMP, "blue"],
  /[ "Alice" , 17 , "green"],
  /[ "Eve"   , 13 , "white"]
]

#test
fillna gradebookMissing ⟨"quiz1", by header⟩ 0
=
Table.mk [
  /[ "Bob"   , 12  , 8     , 9     , 77      , 7     , 9     , 87    ],
  /[ "Alice" , 17  , 6     , 8     , 88      , EMP   , 7     , 85    ],
  /[ "Eve"   , 13  , 0     , 9     , 84      , 8     , 8     , 77    ]
]

-- `pivotLonger`
-- TODO: more typeclass issues...
#test
pivotLonger gradebook A[⟨"midterm", by header⟩, ⟨"final", by header⟩] "exam" "score"
=[by simp [Schema.removeTypedNames, Schema.removeTypedName, Schema.removeHeader, Schema.removeName]; inst]
Table.mk [
  /[ "Bob"   , 12  , 8     , 9     , 7     , 9     , "midterm" , 77    ],
  /[ "Bob"   , 12  , 8     , 9     , 7     , 9     , "final"   , 87    ],
  /[ "Alice" , 17  , 6     , 8     , 8     , 7     , "midterm" , 88    ],
  /[ "Alice" , 17  , 6     , 8     , 8     , 7     , "final"   , 85    ],
  /[ "Eve"   , 13  , 7     , 9     , 8     , 8     , "midterm" , 84    ],
  /[ "Eve"   , 13  , 7     , 9     , 8     , 8     , "final"   , 77    ]
]

#test
pivotLonger gradebook A[⟨"quiz1", by header⟩, ⟨"quiz2", by header⟩,
                        ⟨"quiz3", by header⟩, ⟨"quiz4", by header⟩,
                        ⟨"midterm", by header⟩, ⟨"final", by header⟩]
            "test" "score"
=(Table [("name", String), ("age", Nat), ("test", String), ("score", Nat)])
Table.mk [
  /[ "Bob"   , 12  , "quiz1"   , 8     ],
  /[ "Bob"   , 12  , "quiz2"   , 9     ],
  /[ "Bob"   , 12  , "quiz3"   , 7     ],
  /[ "Bob"   , 12  , "quiz4"   , 9     ],
  /[ "Bob"   , 12  , "midterm" , 77    ],
  /[ "Bob"   , 12  , "final"   , 87    ],
  /[ "Alice" , 17  , "quiz1"   , 6     ],
  /[ "Alice" , 17  , "quiz2"   , 8     ],
  /[ "Alice" , 17  , "quiz3"   , 8     ],
  /[ "Alice" , 17  , "quiz4"   , 7     ],
  /[ "Alice" , 17  , "midterm" , 88    ],
  /[ "Alice" , 17  , "final"   , 85    ],
  /[ "Eve"   , 13  , "quiz1"   , 7     ],
  /[ "Eve"   , 13  , "quiz2"   , 9     ],
  /[ "Eve"   , 13  , "quiz3"   , 8     ],
  /[ "Eve"   , 13  , "quiz4"   , 8     ],
  /[ "Eve"   , 13  , "midterm" , 84    ],
  /[ "Eve"   , 13  , "final"   , 77    ]
]

-- TODO: `pivotWider`

-- `flatten`
#test
flatten gradebookSeq A[⟨"quizzes", _, by header⟩]
=(Table [("name", String), ("age", Nat), ("quizzes", Nat), ("midterm", Nat), ("final", Nat)])
Table.mk [
  /[ "Bob"   , 12  , 8       , 77      , 87    ],
  /[ "Bob"   , 12  , 9       , 77      , 87    ],
  /[ "Bob"   , 12  , 7       , 77      , 87    ],
  /[ "Bob"   , 12  , 9       , 77      , 87    ],
  /[ "Alice" , 17  , 6       , 88      , 85    ],
  /[ "Alice" , 17  , 8       , 88      , 85    ],
  /[ "Alice" , 17  , 8       , 88      , 85    ],
  /[ "Alice" , 17  , 7       , 88      , 85    ],
  /[ "Eve"   , 13  , 7       , 84      , 77    ],
  /[ "Eve"   , 13  , 9       , 84      , 77    ],
  /[ "Eve"   , 13  , 8       , 84      , 77    ],
  /[ "Eve"   , 13  , 8       , 84      , 77    ]
]

def t := buildColumn gradebookSeq "quiz-pass?" (λ r =>
  let isPass : Nat → Bool := λ n => n >= 8
  (getValue r "quizzes" (by header)).map (List.map isPass)
)

#test
flatten t A[⟨"quiz-pass?", _, by header⟩, ⟨"quizzes", _, by header⟩]
=
Table.mk [
  /[ "Bob"   , 12  , 8       , 77      , 87    , true       ],
  /[ "Bob"   , 12  , 9       , 77      , 87    , true       ],
  /[ "Bob"   , 12  , 7       , 77      , 87    , false      ],
  /[ "Bob"   , 12  , 9       , 77      , 87    , true       ],
  /[ "Alice" , 17  , 6       , 88      , 85    , false      ],
  /[ "Alice" , 17  , 8       , 88      , 85    , true       ],
  /[ "Alice" , 17  , 8       , 88      , 85    , true       ],
  /[ "Alice" , 17  , 7       , 88      , 85    , false      ],
  /[ "Eve"   , 13  , 7       , 84      , 77    , false      ],
  /[ "Eve"   , 13  , 9       , 84      , 77    , true       ],
  /[ "Eve"   , 13  , 8       , 84      , 77    , true       ],
  /[ "Eve"   , 13  , 8       , 84      , 77    , true       ]
]

def unbalancedTable :
  Table [("id", Nat), ("seq1", List Nat), ("seq2", List String)] :=
Table.mk [
  /[0, [0, 1, 2], ["a", "b"]],
  /[1, [], ["c"]],
  /[2, [3, 4], ["d", "e", "f"]],
  /[3, [5, 6], []],
  /[4, [], []]
]

-- This behavior matches the B2T2 implementation, although I don't think the
-- behavior of leaving row 4 in the first eval is actually what we'd want.
-- (And I'd argue that leaving the last row in the second example actually
-- violates the spec.)
-- One potential workaround would be to check after each flattening to see if
-- the last row we get is equal (up to cell emptiness -- no need for DecEq)
-- to the clean template row and ditch it if so. (Even more ugly dynamicity, but
-- so it goes...)
-- TODO: notify B2T2 that their implementation crashes on the (valid) example
-- given by the row with ID 4.
#eval flatten unbalancedTable A[⟨"seq1", _, by header⟩]
#eval flatten unbalancedTable A[⟨"seq1", _, by header⟩, ⟨"seq2", _, by header⟩]

-- FIXME: more typeclass issues
-- `transformColumn`
def addLastName := Option.map (· ++ " Smith")

#test
transformColumn students ⟨"name", by header⟩ addLastName
=(Table [("name", String), ("age", Nat), ("favorite color", String)])
Table.mk [
  /[ "Bob Smith"   , 12  , "blue"         ],
  /[ "Alice Smith" , 17  , "green"        ],
  /[ "Eve Smith"   , 13  , "red"          ]
]

def quizScoreToPassFail := Option.map (λ n =>
  if n <= 6
  then "fail"
  else "pass")

#test
transformColumn gradebook ⟨"quiz1", by header⟩ quizScoreToPassFail
=
Table.mk [
  /[ "Bob"   , 12  , "pass" , 9     , 77      , 7     , 9     , 87    ],
  /[ "Alice" , 17  , "fail" , 8     , 88      , 8     , 7     , 85    ],
  /[ "Eve"   , 13  , "pass" , 9     , 84      , 8     , 8     , 77    ]
]

-- `renameColumns`
#test
renameColumns students A[⟨⟨"favorite color", by name⟩, "preferred color"⟩,
                         ⟨⟨"name", by name⟩, "first name"⟩]
=(Table [("first name", String), ("age", Nat), ("preferred color", String)])
Table.mk [
  /[ "Bob"      , 12  , "blue"          ],
  /[ "Alice"    , 17  , "green"         ],
  /[ "Eve"      , 13  , "red"           ]
]

#test
renameColumns gradebook A[⟨⟨"midterm", by name⟩, "final"⟩,
                          ⟨⟨"final", by name⟩, "midterm"⟩]
=(Table [("name", String), ("age", Nat), ("quiz1", Nat), ("quiz2", Nat),
         ("final", Nat), ("quiz3", Nat), ("quiz4", Nat), ("final", Nat)])
Table.mk [
  /[ "Bob"   , 12  , 8     , 9     , 77    , 7     , 9     , 87      ],
  /[ "Alice" , 17  , 6     , 8     , 88    , 8     , 7     , 85      ],
  /[ "Eve"   , 13  , 7     , 9     , 84    , 8     , 8     , 77      ]
]

-- `find`
#test
find [⟨("age", Nat), by header, inferInstance⟩] students /["age" := 13]
=
some ⟨2, by simp⟩

#test
find [⟨("age", _), by header, inferInstance⟩] students /["age" := 14]
=
none

-- `groupByRetentive`
-- Deal with ULift decidable equality
deriving instance DecidableEq for ULift

#test
groupByRetentive students ⟨"favorite color", by header⟩
=
Table.mk [
  /[ULift.up "blue" , Table.mk [
                        /["Bob"  , 12, "blue" ]]],
  /[ULift.up "green", Table.mk [
                        /["Alice", 17, "green"]]],
  /[ULift.up "red"  , Table.mk [
                        /["Eve"  , 13, "red"  ]]]
]

#test
groupByRetentive jellyAnon ⟨"brown", by header⟩
=[by inst]
Table.mk [
  /[ULift.up false, Table.mk [
    /[ true     , false , false , false , true  , false  , false , true   , false , false  ],
    /[ true     , false , true  , false , true  , true   , false , false  , false , false  ],
    /[ false    , false , false , false , true  , false  , false , false  , true  , false  ],
    /[ false    , false , false , false , false , true   , false , false  , false , false  ],
    /[ false    , false , false , false , false , true   , false , false  , true  , false  ],
    /[ true     , false , true  , false , false , false  , false , true   , true  , false  ],
    /[ false    , false , true  , false , false , false  , false , false  , true  , false  ],
    /[ true     , false , false , false , false , false  , false , true   , false , false  ]
  ]],
  /[ULift.up true, Table.mk [
    /[ true     , false , false , false , false , false  , true  , true   , false , false  ],
    /[ false    , true  , false , false , false , true   , true  , false  , true  , false  ]
  ]]
]

-- `groupBySubtractive`
-- TODO: why does the `header` tactic fail here?
-- Interestingly, only fails when we have the equality test -- evaluating
-- `groupBySubtractive` alone works just fine
#test
groupBySubtractive students ⟨"favorite color", Schema.HasCol.tl (Schema.HasCol.tl (Schema.HasCol.hd))⟩
=[by inst]
Table.mk [
  /[ULift.up "blue" , Table.mk [/["Bob"  , 12]]],
  /[ULift.up "green", Table.mk [/["Alice", 17]]],
  /[ULift.up "red", Table.mk [/["Eve"  , 13]]]
]

#test
groupBySubtractive jellyAnon ⟨"brown",
  Schema.HasCol.tl $ Schema.HasCol.tl $ Schema.HasCol.tl $ Schema.HasCol.tl $
    Schema.HasCol.tl $ Schema.HasCol.tl $ Schema.HasCol.hd⟩
=[by inst]
Table.mk [
  /[ULift.up false, Table.mk [
    /[ true     , false , false , false , true  , false  , true   , false , false  ],
    /[ true     , false , true  , false , true  , true   , false  , false , false  ],
    /[ false    , false , false , false , true  , false  , false  , true  , false  ],
    /[ false    , false , false , false , false , true   , false  , false , false  ],
    /[ false    , false , false , false , false , true   , false  , true  , false  ],
    /[ true     , false , true  , false , false , false  , true   , true  , false  ],
    /[ false    , false , true  , false , false , false  , false  , true  , false  ],
    /[ true     , false , false , false , false , false  , true   , false , false  ]
  ]],
  /[ULift.up true, Table.mk [
    /[ true     , false , false , false , false , false  , true   , false , false  ],
    /[ false    , true  , false , false , false , true   , false  , true  , false  ]
  ]]
]

-- `update`
def abstractAgeUpdate := λ (r : Row $ schema students) =>
  match getValue r "age" (by header) with
  | some age =>
    match (age ≤ 12 : Bool), (age ≤ 19 : Bool) with
    | true, _ => /["age" := "kid"]
    | _, true => /["age" := "teenager"]
    | _, _ => /["age" := "adult"]
  | _ => /["age" := EMP]

#test
update [⟨("age", String), _⟩] students abstractAgeUpdate
=(Table [("name", String), ("age", String), ("favorite color", String)])
Table.mk [
  /[ "Bob"   , "kid"      , "blue"         ],
  /[ "Alice" , "teenager" , "green"        ],
  /[ "Eve"   , "teenager" , "red"          ]
]

def didWellUpdate := λ (r : Row $ schema gradebook) =>
  match getValue r "midterm" (by header), getValue r "final" (by header) with
  | some (m : Nat), some (f : Nat) => /["midterm" := (85 ≤ m : Bool), "final" := (85 ≤ f : Bool)]
  | some m, none   => /["midterm" := (85 ≤ m : Bool), "final" := EMP]
  | none, some f   => /["midterm" := EMP, "final" := (85 ≤ f : Bool)]
  | none, none   => /["midterm" := EMP, "final" := EMP]

#test
update [⟨("midterm", Bool), _⟩, ⟨("final", Bool), _⟩] gradebook didWellUpdate
=
Table.mk [
  /[ "Bob"   , 12  , 8     , 9     , false   , 7     , 9     , true  ],
  /[ "Alice" , 17  , 6     , 8     , true    , 8     , 7     , true  ],
  /[ "Eve"   , 13  , 7     , 9     , false   , 8     , 8     , false ]
]

-- `select`
#test
select students (λ (r : Row $ schema students) (n : Fin (nrows students)) =>
  let colorCell : Cell "COLOR" String := Cell.fromOption $ getValue r "favorite color" (by header)
  let ageCell : Cell "AGE" Nat := Cell.fromOption $ getValue r "age" (by header)
  (Row.cons (Cell.val n.val : Cell "ID" Nat) $
  Row.cons colorCell $
  Row.cons ageCell
  Row.nil))
=
Table.mk [
  /[ 0  , "blue"  , 12  ],
  /[ 1  , "green" , 17  ],
  /[ 2  , "red"   , 13  ]
]

#test
select gradebook (λ (r : Row $ schema gradebook) (n : Fin (nrows gradebook)) =>
  let nameCell : Cell "full name" String :=
    Cell.fromOption $ (getValue r "name" (by header)).map (· ++ " Smith")
  let mf2 : Cell "(midterm + final) / 2" Nat :=
    match getValue r "midterm" (by header), getValue r "final" (by header) with
    | some m, some f => Cell.val ((m + f) / 2)
    | _, _ => Cell.emp
  Row.cons nameCell $ Row.cons mf2 Row.nil
)
=
Table.mk [
  /[ "Bob Smith"   , 82                    ],
  /[ "Alice Smith" , 86                    ],
  /[ "Eve Smith"   , 80                    ]
]

-- `selectMany`
-- TODO: type class resolution fails if we annotate `r : Row $ schema students`
#test
selectMany students
(λ r n =>
  if n.val % 2 == 0
  then Table.mk [r]
  else head (Table.mk [r]) ⟨0, by simp [Int.abs, nrows]⟩)
(λ r₁ r₂ => r₂)
=
Table.mk [
  /[ "Bob" , 12  , "blue"         ],
  /[ "Eve" , 13  , "red"          ]
]

def repeatRow {sch : @Schema String} : Row sch → Nat → Table sch
| r, 0 => Table.mk [r]
| r, n+1 => addRows (repeatRow r n) [r]

def decertify {sch : @Schema String}
              (f : Row sch → Nat → Table sch)
              (r : Row sch)
              (nhn : Fin (nrows gradebook)) :=
f r nhn.1

#test
selectMany gradebook (decertify repeatRow)
(λ r₁ r₂ =>
  Row.cons (Cell.fromOption (nm := "midterm") $
              getValue r₂ "midterm" (by header))
  Row.nil)
=
Table.mk [
  /[ 77      ],
  /[ 88      ],
  /[ 88      ],
  /[ 84      ],
  /[ 84      ],
  /[ 84      ]
]

-- `groupJoin`
def getName :=
λ {schema} (h : schema.HasCol ("name", String)) (r : Row schema) =>
  getValue r "name" h

def averageFinal := λ (r : Row $ schema students) (t : Table $ schema gradebook) =>
  r.addColumn "final"
              (some $ average $ List.filterMap id (getColumn2 t "final" (by header)))

#test
groupJoin students gradebook (getName (by header)) (getName (by header)) averageFinal
=[by inst]
Table.mk [
  /[ "Bob"   , 12  , "blue"         , 87    ],
  /[ "Alice" , 17  , "green"        , 85    ],
  /[ "Eve"   , 13  , "red"          , 77    ]
]

def nameLength' :=
λ {schema} (h : schema.HasCol ("name", String)) (r : Row schema) =>
  (getValue r "name" h).map String.length

def tableNRows := λ (r : Row $ schema students) (t : Table $ schema gradebook) =>
  Row.addColumn r "nrows" (some $ nrows t)

#test
groupJoin students gradebook (nameLength' (by header)) (nameLength' (by header)) tableNRows
=[by inst]
Table.mk [
  /[ "Bob"   , 12  , "blue"         , 2     ],
  /[ "Alice" , 17  , "green"        , 1     ],
  /[ "Eve"   , 13  , "red"          , 2     ]
]

-- `join`
def getName' :=
λ {schema} (h : schema.HasCol ("name", String)) (r : Row schema) =>
  getValue r "name" h

def addGradeColumn := λ (r₁ : Row $ schema students) (r₂ : Row $ schema gradebook) =>
  Row.addColumn r₁ "grade" (getValue r₂ "final" (by header))

#test
join students gradebook (getName' (by header)) (getName' (by header)) addGradeColumn
=[by inst]
Table.mk [
  /[ "Bob"   , 12  , "blue"         , 87    ],
  /[ "Alice" , 17  , "green"        , 85    ],
  /[ "Eve"   , 13  , "red"          , 77    ]
]

def nameLength'' :=
λ {schema} (h : schema.HasCol ("name", String)) (r : Row schema) =>
  (getValue r "name" h).map String.length

#test
join students gradebook (nameLength'' $ by header) (nameLength'' $ by header) addGradeColumn
=[by inst]
Table.mk [
  /[ "Bob"   , 12  , "blue"         , 87    ],
  /[ "Bob"   , 12  , "blue"         , 77    ],
  /[ "Alice" , 17  , "green"        , 85    ],
  /[ "Eve"   , 13  , "red"          , 87    ],
  /[ "Eve"   , 13  , "red"          , 77    ]
]


/-
Copyright (c) E.W.Ayers. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.

Author: E.W.Ayers

! This file was ported from Lean 3 source module init.meta.widget.basic
! leanprover-community/mathlib commit 93ae212fb944163e6df29a036182fdea83173e3f
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
prelude
import Leanbin.Init.Function
import Leanbin.Init.Data.Option.Basic
import Leanbin.Init.Util
import Leanbin.Init.Meta.Tactic
import Leanbin.Init.Meta.MkDecEqInstance
import Leanbin.Init.Meta.Json

/-! A component is a piece of UI which may contain internal state. Use component.mk to build new components.

## Using widgets.

To make a widget, you need to make a custom executor object and then instead of calling `save_info_thunk` you call `save_widget`.

Additionally, you will need a compatible build of the vscode extension or web app to use widgets in vscode.

## How it works:

The design is inspired by React.
If you are familiar with using React or Elm or a similar functional UI framework then that's helpful for this.
The [React article on reconciliation](https://reactjs.org/docs/reconciliation.html) might be helpful.

One can imagine making a UI for a particular object as just being a function `f : α → UI` where `UI` is some inductive datatype for buttons, textboxes, lists and so on.
The process of evaluating `f` is called __rendering__.
So for example `α` could be `tactic_state` and the function renders a goal view.

## HTML

For our purposes, `UI` is an HTML tree and is written `html α : Type`. I'm going to assume some familiarity with HTML for the purposes of this document.
An HTML tree is composed of elements and strings.
Each element has a tag such as "div", "span", "article" and so on and a set of attributes and child html.
Use the helper function `h : string → list (attr α) → list (html α) → html α` to build new pieces of `html`. So for example:

```lean
h "ul" [] [
     h "li" [] ["this is list item 1"],
     h "li" [style [("color", "blue")]] ["this is list item 2"],
     h "hr" [] [],
     h "li" [] [
          h "span" [] ["there is a button here"],
          h "button" [on_click (λ _, 3)] ["click me!"]
     ]
]
```
Has the type `html nat`.
The `nat` type is called the __action__ and whenever the user interacts with the UI, the html will emit an object of type `nat`.
So for example if the user clicks the button above, the html will 'emit' `3`.
The above example is compiled to the following piece of html:

```html
<ul>
  <li>this is list item 1</li>
  <li style="{ color: blue; }">this is list item 2</li>
  <hr/>
  <li>
     <span>There is a button here</span>
     <button onClick="[handler]">click me!</button>
  </li>
</ul>
```

## Components

In order for the UI to react to events, you need to be able to take these actions α and alter some state.
To do this we use __components__. `component` takes two type arguments: `π` and `α`. `α` is called the 'action' and `π` are the 'props'.
The props can be thought of as a kind of wrapped function domain for `component`. So given `C : component nat α`, one can turn this into html with
`html.of_component 4 C : html α`.

The base constructor for a component is `pure`:
```lean
meta def Hello : component string α := component.pure (λ s, ["hello, ", s, ", good day!"])

#html Hello "lean" -- renders "hello, lean, good day!"
```
So here a pure component is just a simple function `π → list (html α)`.
However, one can augment components with __hooks__.
The hooks available for compoenents are listed in the inductive definition for component.

Here we will just look at the `with_state` hook, which can be used to build components with inner state.

```
meta inductive my_action
| increment
| decrement
open my_action

meta def Counter : component unit α :=
component.with_state
     my_action          -- the action of the inner component
     int                -- the state
     (λ _, 0)           -- initialise the state
     (λ _ _ s, s)       -- update the state if the props change
     (λ _ s a,          -- update the state if an action was received
          match a with
          | increment := (s + 1, none) -- replace `none` with `some _` to emit an action
          | decrement := (s - 1, none)
          end
     )
$ component.pure (λ ⟨state, ⟨⟩⟩, [
     button "+" (λ _, increment),
     to_string state,
     button "-" (λ _, decrement)
  ])

#html Counter ()
```

You can add many hooks to a component.

- `filter_map_action` lets you filter or map actions that are emmitted by the component
- `map_props` lets you map the props.
- `with_should_update` will not re-render the child component if the given test returns false. This can be useful for efficiency.
- `with_state` discussed above.`
- `with_mouse` subscribes the component to the mouse state, for example whether or not the mouse is over the component. See the `tests/lean/widget/widget_mouse.lean` test for an example.

Given an active document, Lean (in server mode) maintains a set of __widgets__ for the document.
A widget is a component `c`, some `p : Props` and an internal state-manager which manages the states
of the component and subcomponents and also handles the routing of events from the UI.

## Reconciliation

If a parent component's state changes, this can cause child components to change position or to appear and dissappear.
However we want to preserve the state of these child components where we can.
The UI system will try to match up these child components through a process called __reconciliation__.

Reconciliation will make sure that the states are carried over correctly and will also not rerender subcomponents if they haven't changed their props or state.
To compute whether two components are the same, the system will perform a hash on their VM objects.
Not all VM objects can be hashed, so it's important to make sure that any items that you expect to change over the lifetime of the component are fed through the 'Props' argument.
This is why we need the props argument on `component`.
The reconciliation engine uses the `props_eq` predicate passed to the component constructor to determine whether the props have changed and hence whether the component should be re-rendered.

## Keys

If you have some list of components and the list changes according to some state, it is important to add keys to the components so
that if two components change order in the list their states are preserved.
If you don't provide keys or there are duplicate keys then you may get some strange behaviour in both the Lean widget engine and react.

It is possible to use incorrect HTML tags and attributes, there is (currently) no type checking that the result is a valid piece of HTML.
So for example, the client widget system will error if you add a `text_change_event` attribute to anything other than an element tagged with `input`.

## Styles with Tachyons

The widget system assumes that a stylesheet called 'tachyons' is present.
You can find documentation for this stylesheet at [Tachyons.io](http://tachyons.io/).
Tachyons was chosen because it is very terse and allows arbitrary styling without using inline styles and without needing to dynamically load a stylesheet.

## Further work (up for grabs!)

- Add type checking for html.
- Better error handling when the html tree is malformed.
- Better error handling when keys are malformed.
- Add a 'with_task' which lets long-running operations (eg running `simp`) not block the UI update.
- Timers, animation (ambitious).
- More event handlers
- Drag and drop support.
- The current perf bottleneck is sending the full UI across to the server for every update.
  Instead, it should be possible to send a smaller [JSON Patch](http://jsonpatch.com).
  Which is already supported by `json.hpp` and javascript ecosystem.

-/

namespace Widget

inductive MouseEventKind
  | on_click
  | on_mouse_enter
  | on_mouse_leave
#align widget.mouse_event_kind Widget.MouseEventKind

/-- An effect is some change that the widget makes outside of its own state.
Usually, giving instructions to the editor to perform some task.
- `insert_text_relative` will insert at a line relative to the position of the widget.
- `insert_text_absolute` will insert text at the precise position given.
- `reveal_position` will move the editor to view the given position.
- `highlight_position` will add a text highlight to the given position.
- `clear_highlighting` will remove all highlights created with `highlight_position`.
- `copy_text` will copy the given text to the clipboard.
- `custom` can be used to pass custom effects to the client without having to recompile Lean.
-/
unsafe inductive effect : Type
  | insert_text_absolute (file_name : Option String) (p : Pos) (text : String)
  | insert_text_relative (relative_line : Int) (text : String)
  | reveal_position (file_name : Option String) (p : Pos)
  | highlight_position (file_name : Option String) (p : Pos)
  | clear_highlighting
  | copy_text (text : String)
  | custom (key : String) (value : String)
#align widget.effect widget.effect

unsafe def effects :=
  List effect
#align widget.effects widget.effects

mutual
  unsafe inductive component : Type → Type → Type
    | pure {Props Action : Type} (view : Props → List (html Action)) : component Props Action
    |
    filter_map_action {Props InnerAction OuterAction}
      (action_map : Props → InnerAction → Option OuterAction) :
      component Props InnerAction → component Props OuterAction
    |
    map_props {Props1 Props2 Action} (map : Props2 → Props1) :
      component Props1 Action → component Props2 Action
    |
    with_should_update {Props Action : Type} (should_update : ∀ old new : Props, Bool) :
      component Props Action → component Props Action
    |
    with_state {Props Action : Type} (InnerAction State : Type) (init : Props → State)
      (props_changed : Props → Props → State → State)
      (update : Props → State → InnerAction → State × Option Action) :
      component (State × Props) InnerAction → component Props Action
    |
    with_effects {Props Action : Type} (emit : Props → Action → effects) :
      component Props Action → component Props Action
  unsafe inductive html : Type → Type
    | element {α : Type} (tag : String) (attrs : List (attr α)) (children : List (html α)) : html α
    | of_string {α : Type} : String → html α
    | of_component {α : Type} {Props : Type} : Props → component Props α → html α
  unsafe inductive attr : Type → Type
    | val {α : Type} (name : String) (value : json) : attr α
    | mouse_event {α : Type} (kind : MouseEventKind) (handler : Unit → α) : attr α
    | style {α : Type} : List (String × String) → attr α
    | tooltip {α : Type} : html α → attr α
    | text_change_event {α : Type} (handler : String → α) : attr α
end
#align widget.component widget.component
#align widget.html widget.html
#align widget.attr widget.attr

variable {α β : Type} {π : Type}

namespace Component

unsafe def map_action (f : α → β) : component π α → component π β
  | c => filter_map_action (fun p a => some <| f a) c
#align widget.component.map_action widget.component.map_action

/-- Returns a component that will never trigger an action. -/
unsafe def ignore_action : component π α → component π β
  | c => component.filter_map_action (fun p a => none) c
#align widget.component.ignore_action widget.component.ignore_action

unsafe def ignore_props : component Unit α → component π α
  | c => (with_should_update fun a b => false) <| component.map_props (fun p => ()) c
#align widget.component.ignore_props widget.component.ignore_props

unsafe instance : Coe (component π Empty) (component π α) :=
  ⟨component.filter_map_action fun p x => none⟩

unsafe instance : CoeFun (component π α) fun c => π → html α :=
  ⟨fun c p => html.of_component p c⟩

unsafe def stateful {π α : Type} (β σ : Type) (init : π → Option σ → σ)
    (update : π → σ → β → σ × Option α) (view : π → σ → List (html β)) : component π α :=
  with_state β σ (fun p => init p none) (fun _ p s => init p <| some s) update
    (component.pure fun ⟨s, p⟩ => view p s)
#align widget.component.stateful widget.component.stateful

unsafe def stateless {π α : Type} [DecidableEq π] (view : π → List (html α)) : component π α :=
  (component.with_should_update fun p1 p2 => p1 ≠ p2) <| component.pure view
#align widget.component.stateless widget.component.stateless

/--
Causes the component to only update on a props change when `test old_props new_props` yields `ff`. -/
unsafe def with_props_eq (test : π → π → Bool) : component π α → component π α
  | c => component.with_should_update (fun x y => not <| test x y) c
#align widget.component.with_props_eq widget.component.with_props_eq

end Component

mutual
  unsafe def attr.map_action (f : α → β) : attr α → attr β
    | attr.val k v => attr.val k v
    | attr.style s => attr.style s
    | attr.tooltip h => attr.tooltip <| html.map_action h
    | attr.mouse_event k a => attr.mouse_event k (f ∘ a)
    | attr.text_change_event a => attr.text_change_event (f ∘ a)
  unsafe def html.map_action (f : α → β) : html α → html β
    | html.element t a c => html.element t (List.map attr.map_action a) (List.map html.map_action c)
    | html.of_string s => html.of_string s
    | html.of_component p c => html.of_component p <| component.map_action f c
end
#align widget.attr.map_action widget.attr.map_action
#align widget.html.map_action widget.html.map_action

unsafe instance attr.is_functor : Functor attr where map := @attr.map_action
#align widget.attr.is_functor widget.attr.is_functor

unsafe instance html.is_functor : Functor html where map _ _ := html.map_action
#align widget.html.is_functor widget.html.is_functor

namespace Html

/-- See Note [use has_coe_t]. -/
unsafe instance to_string_coe [ToString β] : CoeTC β (html α) :=
  ⟨html.of_string ∘ toString⟩
#align widget.html.to_string_coe widget.html.to_string_coe

unsafe instance : EmptyCollection (html α) :=
  ⟨of_string ""⟩

unsafe instance list_coe : Coe (html α) (List (html α)) :=
  ⟨fun x => [x]⟩
#align widget.html.list_coe widget.html.list_coe

end Html

unsafe def as_element : html α → Option (String × List (attr α) × List (html α))
  | html.element t a c => some ⟨t, a, c⟩
  | _ => none
#align widget.as_element widget.as_element

unsafe def key [ToString β] : β → attr α
  | s => attr.val "key" <| toString s
#align widget.key widget.key

unsafe def className : String → attr α
  | s => attr.val "className" <| s
#align widget.className widget.className

unsafe def on_click : (Unit → α) → attr α
  | a => attr.mouse_event MouseEventKind.on_click a
#align widget.on_click widget.on_click

unsafe def on_mouse_enter : (Unit → α) → attr α
  | a => attr.mouse_event MouseEventKind.on_mouse_enter a
#align widget.on_mouse_enter widget.on_mouse_enter

unsafe def on_mouse_leave : (Unit → α) → attr α
  | a => attr.mouse_event MouseEventKind.on_mouse_leave a
#align widget.on_mouse_leave widget.on_mouse_leave

/-- Alias for `html.element`. -/
unsafe def h : String → List (attr α) → List (html α) → html α :=
  html.element
#align widget.h widget.h

/-- Alias for className. -/
unsafe def cn : String → attr α :=
  className
#align widget.cn widget.cn

unsafe def button : String → Thunk α → html α
  | s, t => h "button" [on_click t] [s]
#align widget.button widget.button

unsafe def textbox : String → (String → α) → html α
  | s, t => h "input" [attr.val "type" "text", attr.val "value" s, attr.text_change_event t] []
#align widget.textbox widget.textbox

unsafe structure select_item (α : Type) where
  result : α
  key : String
  view : List (html α)
#align widget.select_item widget.select_item

/-- Choose from a dropdown selection list. -/
unsafe def select {α} [DecidableEq α] : List (select_item α) → α → html α
  | items, value =>
    let k :=
      match List.filter (fun i => select_item.result i = value) items with
      | [] => ""
      | h :: _ => select_item.key h
    h "select"
        [attr.val "value" k, attr.val "key" k,
          attr.text_change_event fun k =>
            match items.filterₓ fun i => select_item.key i = k with
            | [] => undefined
            | h :: _ => h.result] <|
      items.map fun i => h "option" [attr.val "value" i.key] <| select_item.view i
#align widget.select widget.select

/-- If the html is not an of_element it will wrap it in a div. -/
unsafe def with_attrs : List (attr α) → html α → html α
  | a, x =>
    match as_element x with
    | some ⟨t, as, c⟩ => html.element t (a ++ as) c
    | none => html.element "div" a [x]
#align widget.with_attrs widget.with_attrs

/-- If the html is not an of_element it will wrap it in a div. -/
unsafe def with_attr : attr α → html α → html α
  | a, x => with_attrs [a] x
#align widget.with_attr widget.with_attr

unsafe def with_style : String → String → html α → html α
  | k, v, h => with_attr (attr.style [(k, v)]) h
#align widget.with_style widget.with_style

unsafe def with_cn : String → html α → html α
  | s, h => with_attr (className s) h
#align widget.with_cn widget.with_cn

unsafe def with_key {β} [ToString β] : β → html α → html α
  | s, h => with_attr (key s) h
#align widget.with_key widget.with_key

unsafe def effect.insert_text : String → effect :=
  effect.insert_text_relative 0
#align widget.effect.insert_text widget.effect.insert_text

end Widget

namespace Tactic

/--
Same as `tactic.save_info_thunk` except saves a widget to be displayed by a compatible infoviewer. -/
unsafe axiom save_widget : Pos → widget.component tactic_state Empty → tactic Unit
#align tactic.save_widget tactic.save_widget

/-- Outputs a widget trace position at the given position. -/
unsafe axiom trace_widget_at (p : Pos) (w : widget.component tactic_state Empty)
    (text := "(widget)") : tactic Unit
#align tactic.trace_widget_at tactic.trace_widget_at

/-- Outputs a widget trace position at the current default trace position. -/
unsafe def trace_widget (w : widget.component tactic_state Empty) (text := "(widget)") :
    tactic Unit := do
  let p ← get_trace_msg_pos
  trace_widget_at p w text
#align tactic.trace_widget tactic.trace_widget

end Tactic


syntax "my_trivial" : tactic -- extensible tactic

macro_rules | `(tactic| my_trivial) => `(tactic| decide)
macro_rules | `(tactic| my_trivial) => `(tactic| assumption)

def f (a : Nat) (h : a > 3) := a

example : True := by
  have : f 4 (by my_trivial) = 4 := rfl
  constructor

example : True :=
  have : f 4 (by my_trivial) = 4 := rfl
  ⟨⟩

example : 4 > 3 := by
  my_trivial

example : True :=
  have : f 4 (have : 4 > 3 := (by my_trivial); this) = 4 := rfl
  ⟨⟩


import LeanCodePrompts.Premises
open Lean Meta Elab Term Syntax PrettyPrinter Parser

namespace LeanAide.Meta

class Reprint(a : Type) where
    reprSyn : a → String

instance reprintString : Reprint String where
    reprSyn := id

instance reprintName : Reprint Name where
    reprSyn := toString

instance reprintNat : Reprint Nat where
    reprSyn := toString

instance reprintBool : Reprint Bool where
    reprSyn := toString

instance reprintArray {a : Type} [Reprint a] : Reprint (Array a) where
    reprSyn := fun xs => xs.toList.map Reprint.reprSyn |>.toString

instance reprintList {a : Type} [Reprint a] : Reprint (List a) where
    reprSyn := fun xs => xs.map Reprint.reprSyn |>.toString

instance reprintOption {a : Type} [Reprint a] : Reprint (Option a) where
    reprSyn := fun xs => xs.map Reprint.reprSyn |>.getD ""

instance reprintSyntax : Reprint Syntax where
    reprSyn := fun xs => xs.reprint.get!

def reprint {a : Type}[Reprint a] (x : a) : String := Reprint.reprSyn x

instance reprintTermData : Reprint TermData where
    reprSyn := fun x => s!"context: {Reprint.reprSyn x.context}; term: {Reprint.reprSyn x.value}"

instance reprintProofData : Reprint PropProofData where
    reprSyn := fun x => s!"context: {Reprint.reprSyn x.context}; prop: {Reprint.reprSyn x.prop}; proof : {Reprint.reprSyn x.proof}"

def viewSyntax (s: String) : MetaM <| Syntax × String := do
    let c := runParserCategory (← getEnv) `term s
    match c with
    | Except.error e => throwError e
    | Except.ok s => pure (s, s.reprint.get!)

def nameDefTypeSyntax (name: Name) : MetaM <| Syntax × Syntax := do
    let info? := ((← getEnv).find? name)
    let info := info?.get!
    let exp := info.value?.get!
    let type := info.type
    let (stx, _) ←  delabCore exp {} (delabVerbose)
    let (tstx, _) ←  delabCore type {} (delabVerbose)
    return (stx, tstx)

def nameDefSyntax (name: Name) : MetaM <| Option Syntax := do
    let exp? ← nameExpr? name
    match exp? with
    | none => pure none
    | some exp => do
        let stx ←  delab exp
        pure (some stx)

def premisesFromName (name : Name) : MetaM (List PremiseData) := do
    let (pf, prop) ← nameDefTypeSyntax name
    Lean.Syntax.premiseDataM #[] pf prop true name name

def _root_.PremiseData.view : PremiseData → MetaM String := fun data => do
    return s!"context: {reprint data.context}; name?: {data.name?}; defnName: {data.defnName}; type: {reprint data.type}; type-group: {reprint data.typeGroup}; sub-terms: {reprint data.terms}; sub-proofs : {reprint data.propProofs}  identifiers: {data.ids}"

def premisesViewFromName (name: Name) : MetaM <| List String := do
    let premises ← premisesFromName name
    premises.mapM (fun p => p.view)

def premisesJsonFromName (name: Name) : MetaM <| Json := do
    let premises ← premisesFromName name
    return toJson premises

#eval premisesViewFromName ``Nat.pred_le_pred

-- #eval premisesJsonFromName ``Nat.pred_le_pred

-- #eval premisesViewFromName ``Nat.le_of_succ_le_succ

def boundedDef (bound: Nat)(name: Name) : MetaM Bool := do
    let exp? ← nameExpr? name
    match exp? with
    | none => pure false
    | some exp => do
        pure (exp.approxDepth.toNat < bound)

def nameDefView (name: Name) : MetaM String := do
    let stx? ← nameDefSyntax name
    return (stx?.get!.reprint.get!)

def nameDefCleanView (name: Name) : MetaM String := do
    let stx? ← nameDefSyntax name
    return ((stx?.get!.purge).reprint.get!)

def nameDefSyntaxVerbose (name: Name) : MetaM <| Option Syntax := do
    let exp? ← nameExpr? name
    match exp? with
    | none => pure none
    | some exp => do
        let (stx, _) ←  delabCore exp {} (delabVerbose)
        pure (some stx)

def nameDefViewVerbose (name: Name) : MetaM String := do
    let stx? ← nameDefSyntaxVerbose name
    return (stx?.get!.reprint.get!)

-- #eval nameDefSyntax ``List.join

-- #eval nameDefSyntax ``Nat.le_of_succ_le_succ

-- #eval nameDefView ``Nat.gcd_eq_zero_iff

-- #eval nameDefCleanView ``Nat.gcd_eq_zero_iff

-- def egSplit : MetaM <| Option (Syntax × Array Syntax) := do
--     let stx? ← nameDefSyntax ``Nat.gcd_eq_zero_iff
--     lambdaStx? stx?.get!

-- #eval egSplit

-- def egSplitView : MetaM <| Option (String × Array String) := do
--     let stx? ← nameDefSyntax ``Nat.gcd_eq_zero_iff
--     let pair? ← lambdaStx? stx?.get!
--     let (stx, args) := pair?.get!
--     pure (stx.reprint.get!, args.map (fun s => s.reprint.get!))

-- #eval egSplitView

-- set_option pp.proofs false in
-- #eval nameDefView ``Nat.gcd_eq_zero_iff

-- set_option pp.proofs.withType true in
-- #eval nameDefView ``Nat.gcd_eq_zero_iff

-- #eval nameDefViewVerbose ``Nat.gcd_eq_zero_iff

-- #eval nameDefSyntaxVerbose ``Nat.gcd_eq_zero_iff

-- #eval nameDefViewVerbose ``Nat.gcd_eq_gcd_ab

-- set_option pp.proofs false in
-- #eval nameDefView ``Nat.gcd_eq_gcd_ab

-- #eval setDelabBound 200

-- #eval nameDefViewVerbose ``Nat.xgcd_aux_P

-- set_option pp.proofs false in
-- #eval nameDefView ``Nat.xgcd_aux_P

-- theorem oddExample : ∀ (n : Nat), ∃ m, m > n ∧ m % 2 = 1 := by
--   intro n -- introduce a variable n
--   use 2 * n + 1 -- use `m = 2 * n + 1`
--   apply And.intro -- apply the constructor of `∧` to split goals
--   · linarith -- solve the first goal using `linarith`
--   · simp [Nat.add_mod] -- solve the second goal using `simp` with the lemma `Nat.add_mod`

-- -- #eval premisesViewFromName ``oddExample

end LeanAide.Meta

open LeanAide.Meta

structure ConstsData where
    definitions : HashMap Name  Syntax
    theorems : HashMap Name  Syntax

def constsData : MetaM ConstsData := do
    let consts ← constantNameTypes
    let mut definitions := HashMap.empty
    let mut theorems := HashMap.empty
    for (c, type) in consts do
        let tstx ← delab type
        let tstx := tstx.raw.purge
        if ← Meta.isProp type then
            theorems := theorems.insert c tstx
        else
            definitions := definitions.insert c tstx
    return { definitions := definitions, theorems := theorems }

partial def Lean.Syntax.identsM (stx: Syntax)(context: Array Syntax)(maxDepth? : Option Nat := none) : MetaM <| List <| Name × Nat  := do
    if maxDepth? = some 0 then
        pure []
    else
    match ← namedArgument? stx with
    | some (arg, _) =>
        -- IO.println s!"Named: {arg}"
        let prev ←  identsM  arg context (maxDepth?.map (· -1))
        return prev.map (fun (s, m) => (s, m + 1))
    | none =>
    match ← proofWithProp? stx with
    | some (proof, _) =>
        -- IO.println s!"Proof: {proof}"
        let prev ←  identsM  proof context (maxDepth?.map (· -1))
        return prev.map (fun (s, m) => (s, m + 1))
    | none =>
    match ← lambdaStx? stx with
    | some (body, args) =>
        -- IO.println s!"Lambda: {args}"
        let prev ←  identsM  body (context ++ args) (maxDepth?.map (· -1))
        return prev.map (fun (s, m) => (s, m + args.size))
    | none =>
        match stx with
        | Syntax.node _ _ args =>
            let prev ← args.mapM (identsM · context (maxDepth?.map (· -1)))
            return prev.toList.join.map (fun (s, m) => (s, m + 1))
        | Syntax.ident _ _ name .. =>
            let contextVars := context.filterMap getVar
            -- IO.println s!"Context: {contextVars} from {context}"
            if  !(contextVars.contains name) &&
                !(excludePrefixes.any (fun pfx => pfx.isPrefixOf name)) && !(excludeSuffixes.any (fun pfx => pfx.isSuffixOf name)) then
                pure [(name, 0)]
            else pure []
        | _ => pure []

-- -- #eval termKindList

partial def Lean.Syntax.termsM (context : Array Syntax)(stx: Syntax)(maxDepth? : Option Nat := none) : MetaM <| List <| TermData   := do
    let tks ← termKindList
    let tks := tks.map (·.1)
    if maxDepth? = some 0 then
        pure []
    else
    match ← namedArgument? stx with
    | some (arg, _) =>
        -- IO.println s!"Named: {arg}"
        let prev ←  termsM  context arg (maxDepth?.map (· -1))
        return prev.map (fun s => s.increaseDepth 1 )
    | none =>
    match ← proofWithProp? stx with
    | some (proof, _) =>
        -- IO.println s!"Proof: {proof}"
        let prev ←  termsM context  proof (maxDepth?.map (· -1))
        return prev.map (fun s => s.increaseDepth 1)
    | none =>
    match ← lambdaStx? stx with
    | some (body, args) =>
        -- IO.println s!"Lambda: {args}"
        let prev ←  termsM (context ++ args) body (maxDepth?.map (· -1))
        return prev.map (fun s => s.increaseDepth args.size)
    | none =>
        match stx with
        | Syntax.node _ k args =>
            -- IO.println s!"Node: {k}"
            let prev ← args.mapM (termsM context · (maxDepth?.map (· -1)))
            let head : TermData := ⟨context, stx.purge, stx.purge.size, 0⟩
            if tks.contains k then
                return (head) :: prev.toList.join.map (fun s => s.increaseDepth 1)
            else
            return prev.toList.join.map (fun s => s.increaseDepth 1)
        | Syntax.ident .. =>
             pure []
        | _ => pure []

partial def Lean.Syntax.proofsM (context : Array Syntax)(stx: Syntax)(maxDepth? : Option Nat := none) : MetaM <| List <| PropProofData   := do
    if maxDepth? = some 0 then
        pure []
    else
    match ← namedArgument? stx with
    | some (arg, _) =>
        -- IO.println s!"Named: {arg}"
        let prev ←  proofsM  context arg (maxDepth?.map (· -1))
        return prev.map (fun s => s.increaseDepth 1)
    | none =>
    match ← proofWithProp? stx with
    | some (proof, prop) =>
        -- IO.println s!"Proof: {proof}"
        let prev ←  proofsM context  proof (maxDepth?.map (· -1))
        let head : PropProofData :=
            ⟨context, prop, proof, prop.size, proof.size, 0⟩
        return  (head) :: prev.map (fun s => s.increaseDepth 1)
    | none =>
    match ← lambdaStx? stx with
    | some (body, args) =>
        -- IO.println s!"Lambda: {args}"
        let prev ←  proofsM (context ++ args) body (maxDepth?.map (· -1))
        return prev.map (fun s => s.increaseDepth args.size)
    | none =>
        match stx with
        | Syntax.node _ _ args =>
            -- IO.println s!"Node: {k}"
            let prev ← args.mapM (proofsM context · (maxDepth?.map (· -1)))
            return prev.toList.join.map (fun s => s.increaseDepth 1)
        | Syntax.ident .. =>
             pure []
        | _ => pure []

def PremiseData.get(ctx : Array Syntax)(name: Name)(prop pf : Syntax) : MetaM PremiseData := do
    let subProofs: List (PropProofData) ←  pf.proofsM ctx
    let subTerms : List (TermData)  ← Syntax.termsM ctx pf
    let ids : List (Name  × Nat) ← pf.identsM ctx
    return ⟨ctx, some name, name, prop.purge, prop.purge, pf.purge, prop.purge.size, pf.purge.size, subTerms.toArray, subProofs.toArray, ids.toArray⟩

def viewData (name: Name) : MetaM <| String := do
    let (stx, tstx) ← nameDefTypeSyntax name
    -- IO.println s!"{stx.reprint.get!}"
    -- IO.println s!"{← proofWithProp? stx}"
    let data ←  PremiseData.get  #[] name tstx stx
    data.view

-- #eval viewData ``Nat.succ_le_succ

partial def Lean.Syntax.kinds (stx: Syntax)(maxDepth?: Option Nat := none) : List String :=
    if maxDepth? = some 0 then
        []
    else
    match stx with
    | Syntax.node _ k args =>
        let head? : Option String :=
            k.components.head?.map (· |>.toString)
        match head? with
        | some head => head :: (args.map (kinds · (maxDepth?.map (· -1))) |>.toList.join)
        | none => args.map (kinds · (maxDepth?.map (· -1))) |>.toList.join
    | _ => []

partial def Lean.Syntax.idents (stx: Syntax)(maxDepth? : Option Nat := none) : List <| Name × Nat  :=
    if maxDepth? = some 0 then
        []
    else
    match stx with
    | Syntax.node _ _ args =>
         args.map (idents · (maxDepth?.map (· -1)))
            |>.toList.join.map (fun (s, m) => (s, m + 1))
    | Syntax.ident _ _ name .. =>
        if !(excludePrefixes.any (fun pfx => pfx.isPrefixOf name)) && !(excludeSuffixes.any (fun pfx => pfx.isSuffixOf name)) then
            [(name, 0)]
        else []
    | _ => []

partial def Lean.Syntax.terms (stx: Syntax)(maxDepth?: Option Nat := none) :
     MetaM <|  List <| String × Nat × List Name := do
    let tks ← termKindList
    let tks := tks.map (·.1)
    if maxDepth? = some 0 then
        pure []
    else
    match stx with
    | Syntax.node _ k args =>
        match stx with
        | `(proved_prop| ($pf:term =: $_:term )) =>
           pf.raw.terms
        | _ =>
        let head? : Option String := do
            if
                k ∈ tks
            then
                ← stx.reprint
            else
                none
        let argTerms ← args.mapM (terms · (maxDepth?.map (· -1)))
        let argTerms := argTerms.toList.join
        match head? with
        | some head =>
            return (head.trim, 0, stx.idents.map (·.1)) ::
                 (argTerms.map (fun (s, m, l) => (s, m + 1, l)))
        | none =>
            return (argTerms.map (fun (s, m, l) => (s, m + 1, l)))

    | _ => pure []

structure Premises where
    type : String
    defTerms : List <| String × Nat × List Name
    defIdents : List <| Name × Nat
    typeTerms : List <| String × Nat × List Name
    typeIdents : List <| Name × Nat
deriving Repr, Inhabited

def Premises.defMainTerms (p: Premises) : List <| String × Nat × List Name  :=
    p.defTerms.filter (
            fun (s, _) => s.1.length < 20)

def Premises.typeMainTerms (p: Premises) : List <| String × Nat × List Name :=
    p.typeTerms.filter (fun (s, _) => (s.splitOn "=>").length == 1
                && (s.splitOn "↦").length == 1)

def getPremises (name: Name)(maxDepth? : Option Nat := none ) : MetaM <| Premises := do
    let termStx? ← nameDefSyntaxVerbose name
    let term ←  mkConstWithLevelParams name
    let type ← inferType term
    let typeView ← Meta.ppExpr type
    let typeStx ← delab type
    let defTerms ←  match termStx? with
        | none => pure []
        | some stx => stx.terms maxDepth?
    let defTerms := defTerms.filter (fun (s, _) => s.1.length < 10000
        && !s.contains '\n')
    let defIdents := match termStx? with
        | none => []
        | some stx => stx.idents maxDepth?
    pure {type := typeView.pretty 10000, defTerms := defTerms, defIdents := defIdents, typeTerms := ←  typeStx.raw |>.terms maxDepth?, typeIdents := typeStx.raw |>.idents maxDepth?}

-- Testing

def showTerms (s: String) : MetaM <| List <| String × Nat × List Name  := do
    let c := runParserCategory (← getEnv) `term s
    match c with
    | Except.error e => throwError e
    | Except.ok s => (s.terms)

def showIdents (s: String) : MetaM <| List <| Name × Nat := do
    let c := runParserCategory (← getEnv) `term s
    match c with
    | Except.error e => throwError e
    | Except.ok s => pure (s.idents)

def showKinds (s: String) : MetaM <| List String := do
    let c := runParserCategory (← getEnv) `term s
    match c with
    | Except.error e => throwError e
    | Except.ok s => pure (s.kinds)

def nameDefTerms (name: Name)(maxDepth? : Option Nat := none ) : MetaM <|
    List <| String × Nat × List Name  := do
    let stx? ← nameDefSyntax name
    match stx? with
    | none => pure []
    | some stx => (stx.terms maxDepth?)

def nameDefIdents (name: Name)(maxDepth? : Option Nat := none ) : MetaM <| List <| Name × Nat := do
    let stx? ← nameDefSyntax name
    match stx? with
    | none => pure []
    | some stx => pure (stx.idents maxDepth?)

#check List.join

-- #eval showTerms "fun n ↦ Nat.succ n"

-- #eval showIdents "fun n ↦ Nat.succ n"

-- #eval showTerms "fun n ↦ Nat.succ n = n + 1"

#eval viewSyntax "match n + 2 with | 0 => 0 | _ => 1"

-- #eval viewSyntax "f (n := m)"

def zeroOrOne : Nat → Nat
| 0 => 0
| _ => 1

#eval nameDefSyntax ``zeroOrOne

-- #eval nameDefSyntax ``Nat.succ_le_succ

#check Lean.Parser.Term.namedArgument
#check Lean.Parser.Term.matchAlt

-- #eval showKinds "n = n + 1"

-- def egTerms : MetaM <| List <| String × Nat × List Name := do
--     let p ←  getPremises ``oddExample (some 30)
--     return p.defMainTerms

-- #eval egTerms

-- def egIdents : MetaM <| List <| Name × Nat:= do
--     let p ←  getPremises ``oddExample (some 50)
--     return p.defIdents

-- #eval egIdents

-- def egGpIdents : MetaM NameGroups := do
--     let nd ← egIdents
--     return groupedNames nd.toArray

-- #eval egGpIdents

-- #check Linarith.lt_irrefl

-- -- #eval nameDefSyntax ``oddExample

def dataSize : MetaM Nat := do
    let names ← constantNames
    return names.size

-- #eval dataSize

def boundedDataSize (n: Nat) : MetaM Nat := do
    let names ← constantNames
    let names ← names.filterM (boundedDef n)
    return names.size

-- #eval boundedDataSize 50

def sampleExtract (n: Nat := 100) : MetaM <|
        List (Name × (List <| String × Nat × List Name) ×
        (List <| Name × Nat)) := do
    let names ← constantNames
    let names := names.toList.take n
    names.mapM (fun n => do
        let p ← nameDefTerms n
        let q ← nameDefIdents n
        pure (n, p, q)
        )

def batchPremiseExtractM (start stop: Nat) : MetaM Nat  := do
    let names ← constantNames
    let premisesFile := System.mkFilePath ["rawdata",
    s!"outer-premises.jsonl"]
    let h ← IO.FS.Handle.mk premisesFile IO.FS.Mode.append Bool.false
    let names := names.toList.drop start |>.take (stop - start)
    let mut cursor := start
    IO.println s!"start: {start}, stop: {stop}"
    for name in names do
        IO.println <| s!"starting: {cursor} {name}"
        let premises ← getPremises name (some 30)
        let p := premises.defMainTerms
        let pJson := p.map
            (fun (s, n, l) =>
                Json.mkObj [
                    ("term", s),
                    ("depth", n),
                    ("local-idents", Json.arr  (
                        l.toArray
                        |>.filter (fun name => !names.contains name)
                        |>.map (fun name : Name =>
                            name.toString)))
                ]
            )
        let q:= premises.defIdents
        let q := q.filter (fun (name, _) => names.contains name)
        let qJson := q.map
            (fun (name, n) =>
                Json.mkObj [
                    ("name", name.toString),
                    ("depth", n)
                ]
            )
        let js := Json.mkObj [
            ("name", name.toString),
            ("type", premises.type),
            ("terms", Json.arr
                pJson.toArray),
            ("idents", Json.arr
                qJson.toArray)
        ]
        let out := js.pretty 10000
        if out.length < 9000 then
            h.putStrLn <| out
        IO.println <| s!"{cursor}. {name} : {premises.type} ({p.length}, {q.length}, {out.length})"
        cursor := cursor + 1
    return cursor

def batchPremiseExtractCore (start stop: Nat) : CoreM Nat :=
    (batchPremiseExtractM start stop).run'

-- -- #eval sampleExtract

-- theorem imo_1964_q1b : ∀ (n : Nat), (2 ^ n + 1) % 7 ≠ 0
--     | 0 | 1 | 2 => by decide
--     | n + 3 => by
--       rw [pow_add, Nat.add_mod, Nat.mul_mod, show 2 ^ 3 % 7 = 1 from by rfl]
--       simp [imo_1964_q1b n]

-- set_option pp.proofs.withType true in
-- -- #eval nameDefView ``imo_1964_q1b

-- #eval nameDefView ``imo_1964_q1b

-- -- #eval nameDefViewVerbose ``imo_1964_q1b