open HolKernel Parse boolLib bossLib;
open sumSyntax
open DefineUtils

(* A monad for sum types, roughly equivalent to Haskell's `Either` monad. It
   favors the left side of a sum type (although `bind_right` can be used to bind
   over the right side instead). Also contains several utility functions for sum
   types.
*)
val _ = new_theory "sumMonad"

val _ = define_monad("left", (fn a => mk_sum(a, mk_vartype "'right")), 
  ``INL``, ``λx fn. sum_CASE x fn INR``)

(*
val return_left_INL = store_thm("return_left_INL[simp]",
  ``return_left x = INL x``,
  simp[definition "return_left_def"])
*)

val bind_left_INR_simp = store_thm("bind_left_INR_simp[simp]",
  ``bind_left (INR x) f = INR x``, simp[])

val bind_left_INL = store_thm("bind_left_INL[simp]",
  ``(bind_left x f = INL y) = ∃y0. (x = INL y0) ∧ (f y0 = INL y)``,
  Cases_on `x` >> simp[])

val bind_left_INR = store_thm("bind_left_INR",
  ``(∃y. x = INR y) ⇒ ∀f. bind_left x f = x``,
  Cases_on `x` >> simp[])


val bind_right_def = Define`
  bind_right (sum : α + β) (fn : β -> α + γ) : α + γ =
    case sum of
    | INL a => INL a
    | INR b => fn b
`

val bind_right_INR = store_thm("bind_right_INR[simp]",
  ``(bind_right x f = INR y) = ∃y0. (x = INR y0) ∧ (f y0 = INR y)``,
  Cases_on `x` >> simp[bind_right_def])

val bind_right_INL = store_thm("bind_right_INL",
  ``(∃y. x = INL y) ⇒ ∀f. bind_right x f = x``,
  Cases_on `x` >> simp[bind_right_def])

val lift_left_INL = store_thm("lift_left_INL[simp]",
  ``(lift_left f x = INL y) = ∃y0. (x = INL y0) ∧ (y = f y0)``,
  Cases_on `x` >> simp[definition "lift_left_def", bind_left_INL] >> metis_tac[])

val lift_left_INR = store_thm("lift_left_INR[simp]",
  ``(lift_left f x = INR y) = (x = INR y)``,
  Cases_on `x` >> simp[definition "lift_left_def"])

val lift_left_id = store_thm("lift_left_id[simp]",
  ``lift_left I x = x``,
  Cases_on `x` >> simp[definition "lift_left_def"])

val lift_left_compose = store_thm("lift_left_compose[simp]",
  ``lift_left (f o g) x = lift_left f (lift_left g x)``,
  Cases_on `x` >> simp[definition "lift_left_def"])

val lift_right_def = Define`
  lift_right (f : α -> β) (sum : γ + α) : γ + β =
    bind_right sum (INR o f)
`

val lift_right_INR = store_thm("lift_right_INR[simp]",
  ``(lift_right f x = INR y) = ∃y0. (x = INR y0) ∧ (y = f y0)``,
  Cases_on `x` >> simp[lift_right_def, bind_right_INR] >> metis_tac[])

val lift_right_INL = store_thm("lift_right_INL[simp]",
  ``(lift_right f x = INL y) = (x = INL y)``,
  Cases_on `x` >> simp[lift_right_def, bind_right_def])

val lift_right_id = store_thm("lift_right_id[simp]",
  ``lift_right I x = x``,
  Cases_on `x` >> simp[lift_right_def, bind_right_def])

val lift_right_compose = store_thm("lift_right_compose[simp]",
  ``lift_right (f o g) x = lift_right f (lift_right g x)``,
  Cases_on `x` >> simp[lift_right_def, bind_right_def])

(* Merge left side of sum into right side *)
val merge_right_def = Define`
  merge_right (f : α -> β) (sum : α + β) : β =
    case sum of
    | INL a => f a
    | INR b => b
`

(* Merge right side of sum into left side *)
val merge_left_def = Define`
  merge_left (f : β -> α) (sum : α + β) : α =
    case sum of
    | INL a => a
    | INR b => f b
`

val left_set_def = Define`
  left_set (sum : α + β) = case sum of INL a => {a} | INR _ => ∅
`

val right_set_def = Define`
  right_set (sum : α + β) = case sum of INL _ => ∅ | INR b => {b}
`

val left_list_def = Define`
  left_list (sum : α + β) = case sum of INL a => [a] | INR _ => []
`

val right_list_def = Define`
  right_list (sum : α + β) = case sum of INL _ => [] | INR b => [b]
`

(* Monad laws for bind_left: left identity, right identity, and associativity *)
val bind_left_monad_left_id = store_thm("bind_left_monad_left_id[simp]",
  ``bind_left (INL a) f = f a``, simp[])
val bind_left_monad_right_id = store_thm("bind_left_monad_right_id[simp]",
  ``bind_left m INL = m``,
  Cases_on `m` >> simp[])
val bind_left_monad_assoc = store_thm("bind_left_monad_assoc",
  ``bind_left (bind_left m f) g = bind_left m (λx. bind_left (f x) g)``,
  Cases_on `m` >> simp[])
val _ = computeLib.add_funs[
  ]

(* Monad laws for bind_right: left identity, right identity, and associativity *)
val bind_right_monad_left_id = store_thm("bind_right_monad_left_id[simp]",
  ``bind_right (INR a) f = f a``, simp[bind_right_def])
val bind_right_monad_right_id = store_thm("bind_right_monad_right_id[simp]",
  ``bind_right m INR = m``, Cases_on `m` >> simp[bind_right_def])
val bind_right_monad_assoc = store_thm("bind_right_monad_assoc",
  ``bind_right (bind_right m f) g = bind_right m (λx. bind_right (f x) g)``,
  Cases_on `m` >> simp[bind_right_def])

(* Applicative laws for ap_left *)
val ap_left_identity = store_thm("ap_left_identity",
  ``INL I <*> v = v``,
  rw[definition "ap_left_def"] >> metis_tac[bind_left_monad_right_id])

(* TODO: homomorphism, interchange *)

val ap_left_composition = store_thm("ap_left_composition",
  ``INL $o <*> u <*> v <*> w = u <*> (v <*> w)``,
  Cases_on `u` >> Cases_on `v` >> Cases_on `w` >> rw[definition "ap_left_def"])

(* Other theorems that use the monad laws *)
val left_set_image = store_thm("left_set_image",
  ``left_set (lift_left f x) = IMAGE f (left_set x)``,
  Cases_on `x`
  >- simp[definition "lift_left_def", left_set_def]
  >- simp[left_set_def, METIS_PROVE [lift_left_INR] ``lift_left f (INR y) = INR y``])

val bind_left_unit_INL = store_thm("bind_left_unit_INL[simp]",
  ``monad_unitbind (INL a) b = b``, simp[])

val bind_left_unit_INR = store_thm("bind_left_unit_INR[simp]",
  ``monad_unitbind (INR a) b = INR a``, simp[])

val _ = export_theory()