Monoid2.lean

import Mathlib.Util.AtomM
import Mathlib.Data.Nat.Basic

open Lean Elab Meta Tactic Mathlib.Tactic
open Lean.Expr

namespace Monoid
structure Context where
  α : Expr -- The type of the expression
  univ : Level -- The universe level of α
  inst : Expr -- The Monoid instance
  α1 : Expr -- the identity of the monoid

def mkContext (e : Expr) : MetaM Context := do
  let α ← inferType e
  let c ← synthInstance (← mkAppM ``Monoid #[α])
  let u ← mkFreshLevelMVar
  let α1 ← Expr.ofNat α 1
  _ ← isDefEq (.sort (.succ u)) (← inferType α)
  return ⟨ α, u, c, α1 ⟩

-- applies functions of form ∀ {u} {a : Type u}, _
def Context.app (c : Context) (n : Name) : Array Expr → Expr :=
  mkAppN ((@Expr.const n [c.univ]).app c.α)

-- applies functions of form ∀ {u} {a : Type u} [Monoid α], _
def Context.appInst (c : Context) (n : Name) : Array Expr → Expr :=
  mkAppN (((@Expr.const n [c.univ]).app c.α).app c.inst)

abbrev M := ReaderT Context AtomM

inductive NF : Type u where
  | mult : Expr → Expr → NF → NF
  | mempty : Expr → NF
  deriving Inhabited
open NF

def NF.e : NF → Expr
  | mult e .. => e
  | mempty e .. => e

def e1 : M NF := return NF.mempty (← read).α1

def mkSingle (e : Expr) : M (NF × Expr) := do
  let exp ← mkAppM ``HMul.hMul #[e, (← read).α1]
  let nf := NF.mult exp e (← e1)
  let pf := (← read).appInst ``Monoid.mul_one #[e]
  return (nf, ← mkEqSymm pf)

lemma left_mult [Monoid α] (x y z w : α) :
  y * z = w →
  (x * y) * z = x * w := by
  intros H
  rw [mul_assoc, H]

partial def evalMult : NF → NF → M (NF × Expr)
  | mempty _, e₂ => do
    let p := (← read).appInst ``Monoid.one_mul #[e₂.e]
    return (e₂, p)
  | e₁, mempty _ => do
    let p := (← read).appInst ``Monoid.mul_one #[e₁.e]
    return (e₁, p)
  | mult _ x₁ y₁, e₂ => do
    let (nf, p) ← evalMult y₁ e₂
    let aaa ← mkAppM ``HMul.hMul #[x₁, nf.e]
    let nf' := mult aaa x₁ nf
    -- p is a proof that y₁.e * e₂.e = nf.e
    -- we want a proof that (x₁ * y₁.e) * e₂.e = x₁ * nf.e
    let p' := (← read).appInst ``left_mult #[x₁, y₁.e, e₂.e, nf.e, p]
    return (nf', p')

lemma subst_into_mult {α} [Monoid α] (l r tl tr t)
  (prl : (l : α) = tl)
  (prr : r = tr)
  (prt : tl * tr = t) :
  l * r = t := by
  rw [prl, prr, prt]

partial def eval (e : Expr) : M (NF × Expr) := do
  match getAppFnArgs e with
  | (``HMul.hMul, #[_, _, _, _, e₁, e₂]) => do
    let (e₁', p₁) ← eval e₁
    let (e₂', p₂) ← eval e₂
    let (e', p') ← evalMult e₁' e₂'
    return (e', (← read).appInst ``subst_into_mult #[e₁, e₂, e₁'.e, e₂'.e, e'.e, p₁, p₂, p'])
  | _ =>
    if ← isDefEq e (← read).α1 then
      return (NF.mempty e, ← mkEqRefl (← read).α1)
    else
      mkSingle e

syntax (name := monoid) "monoid" : tactic
elab_rules : tactic | `(tactic| monoid) => withMainContext do
    let some (_, e₁, e₂) := (← whnfR <| ← getMainTarget).eq?
      | throwError "monoid: requires an equality goal"
    let c ← mkContext e₁
    closeMainGoal <| ← AtomM.run .default <| ReaderT.run (r := c) do
      let (e₁', p₁) ← eval e₁
      -- dbg_trace f!"found `{p₁}`, a proof that `{e₁} = {e₁'.e}`"
      let (e₂', p₂) ← eval e₂
      -- dbg_trace f!"found `{p₂}`, a proof that `{e₂} = {e₂'.e}`"
      unless ← isDefEq e₁'.e e₂'.e do
        throwError "monoid: could not find that the two sides are equal"
      mkEqTrans p₁ (← mkEqSymm p₂)


lemma ex1' [Monoid α] (a b : α) : a * b = 1 * a * b := by
  monoid

--
-- -- Examples
instance : Monoid (List α) where
  mul := List.append
  one := []
  mul_assoc := List.append_assoc
  one_mul := List.nil_append
  mul_one := List.append_nil

def f := 3
lemma ex1 : 1 * 2 * 3 = 1 * (2 * f) * 1 * 1 := by
  monoid

lemma ex2 : ∀ (a b c : ℕ), a * 1 * b * (a + c) = a * (b * (a + c)):= by
  intros
  monoid

lemma ex3 : ∀ (a b c d : List Bool), [] * a * (b * c) * d = a * (b * c * d) := by
  intros
  monoid

lemma ex4 : ∀ (a b c d : ℕ), a * b * (b * d * c * a * 1 * d) = 1 * 1 * a * (1 * b) * (b * d) * c * (a * d) * 1 := by
  intros
  monoid

lemma ex5 : ∀ (a b c d : List α), a * b * c * d * [] = a * (b * c) * (1 * 1 * d) := by
  intros
  monoid

lemma ex7 [Monoid α]: ∀ (a b c d : α), a * b * c * d * 1 = a * (b * c) * (1 * 1 * d) := by
  intros
  monoid