import Batteries.Data.Rat.Basic
import Batteries.Control.Lawful.MonadLift

inductive Ast
| lit (n : Rat)
| add (l r : Ast)
| sub (l r : Ast)
| mul (l r : Ast)
| div (l r : Ast)

instance : Coe Rat Ast where coe := .lit
instance : OfScientific Ast where
  ofScientific a b c := .lit (.ofScientific a b c)
instance : OfNat Ast n where ofNat := .lit n

instance : Add Ast where add := .add
instance : Sub Ast where sub := .sub
instance : Mul Ast where mul := .mul
instance : Div Ast where div := .div

def Ast.interpret : Ast → Rat
| lit n => n
| add l r => l.interpret + r.interpret
| sub l r => l.interpret - r.interpret
| mul l r => l.interpret * r.interpret
| div l r => l.interpret / r.interpret

#eval Ast.interpret <| 3/2 + 0.5

inductive Op
| push (n : Rat)
| add
| sub
| mul
| div
  deriving Repr

-- TODO: prove that the stack manip is all good throughout
def Op.pop_push : Op → Nat × Nat
| push _ => (0, 1)
| _ => (2, 1)

abbrev Ops := Array Op

def pushM (x : α) : StateM (Array α) Unit :=
  modify <| (·.push x)

def popM? : StateM (Array α) (Option α) := do
  let top ← Array.back? <$> get
  modify Array.pop
  return top

abbrev M := OptionT (StateM (Array Rat))
instance : LawfulMonad M := sorry

def Op.action (op : Op) : M Unit := do
  match op with
  | .push n => pushM n
  | .add => pushM <| (← popM?) + (← popM?)
  | .sub => pushM <| (← popM?) - (← popM?)
  | .mul => pushM <| (← popM?) * (← popM?)
  | .div => pushM <| (← popM?) / (← popM?)

def Ops.action (ops : Ops) : M Unit := ops.forM Op.action

def Ops.run (ops : Ops) : Option Rat :=
  (do
    ops.action
    popM?
  ).run #[] |>.fst

def Ast.compile : Ast → Ops
| .lit n => #[.push n]
| .add l r => l.compile ++ r.compile ++ #[.add]
| .sub l r => l.compile ++ r.compile ++ #[.sub]
| .mul l r => l.compile ++ r.compile ++ #[.mul]
-- Note that l and r are flipped here. compilation!
| .div l r => r.compile ++ l.compile ++ #[.div]

#eval (3/2 + 0.5 : Ast) |>.compile.run

@[simp]
theorem Ops.append_action (a b : Ops)
    : (a ++ b).action = (do a.action; b.action) := by
  rw [action]
  exact Array.forM_append

@[simp]
theorem Ops.push_action (ops : Ops) (op : Op)
    : action (ops.push op) = (do ops.action; op.action) := by
  rw [Array.push_eq_append, append_action]
  have : op.action = action #[op] := by simp [action]
  rw [this]

theorem Ast.compile_sound (ast : Ast)
    : ast.compile.run = some ast.interpret := by
  induction ast
  · rfl
  simp [compile]