module Encoded.Union

import Control.Function.FunExt
import Syntax.PreorderReasoning
import Term.Semantics
import Term.Syntax

-- Type ------------------------------------------------------------------------

export
(<+>) : Ty -> Ty -> Ty
N <+> N = N
N <+> (ty2 ~> ty2') = ty2 ~> (N <+> ty2')
(ty1 ~> ty1') <+> N = ty1 ~> (ty1' <+> N)
(ty1 ~> ty1') <+> (ty2 ~> ty2') = (ty1 <+> ty2) ~> (ty1' <+> ty2')

-- Universal Property ----------------------------------------------------------

export
swap : {ty1, ty2 : Ty} -> Term ((ty1 <+> ty2) ~> (ty2 <+> ty1)) ctx
swap {ty1 = N, ty2 = N} = Id
swap {ty1 = N, ty2 = ty2 ~> ty2'} = Abs' (\f => swap . f)
swap {ty1 = ty1 ~> ty1', ty2 = N} = Abs' (\f => swap . f)
swap {ty1 = ty1 ~> ty1', ty2 = ty2 ~> ty2'} = Abs' (\f => swap . f . swap)

export
inL : {ty1, ty2 : Ty} -> Term (ty1 ~> (ty1 <+> ty2)) ctx
export
prL : {ty1, ty2 : Ty} -> Term ((ty1 <+> ty2) ~> ty1) ctx

inL {ty1 = N, ty2 = N} = Id
inL {ty1 = N, ty2 = ty2 ~> ty2'} = Abs' (\n => Const (App inL [<n]))
inL {ty1 = ty1 ~> ty1', ty2 = N} = Abs' (\f => inL . f)
inL {ty1 = ty1 ~> ty1', ty2 = ty2 ~> ty2'} = Abs' (\f => inL . f . prL)

prL {ty1 = N, ty2 = N} = Id
prL {ty1 = N, ty2 = ty2 ~> ty2'} = Abs' (\t => App prL [<App t [<Arb]])
prL {ty1 = ty1 ~> ty1', ty2 = N} = Abs' (\t => prL . t)
prL {ty1 = ty1 ~> ty1', ty2 = ty2 ~> ty2'} = Abs' (\t => prL . t . inL)

export
inR : {ty1, ty2 : Ty} -> Term (ty2 ~> (ty1 <+> ty2)) ctx
inR = swap . inL

export
prR : {ty1, ty2 : Ty} -> Term ((ty1 <+> ty2) ~> ty2) ctx
prR = prL . swap

-- Semantics -------------------------------------------------------------------

-- Redefinition in Idris

namespace Sem
  public export
  swap : {ty1, ty2: Ty} -> TypeOf (ty1 <+> ty2) -> TypeOf (ty2 <+> ty1)
  swap {ty1 = N, ty2 = N} = id
  swap {ty1 = N, ty2 = ty2 ~> ty2'} = \f => Sem.swap . f
  swap {ty1 = ty1 ~> ty1', ty2 = N} = \f => Sem.swap . f
  swap {ty1 = ty1 ~> ty1', ty2 = ty2 ~> ty2'} = \f => Sem.swap . f . Sem.swap

  public export
  inL : {ty1, ty2 : Ty} -> TypeOf ty1 -> TypeOf (ty1 <+> ty2)

  public export
  prL : {ty1, ty2 : Ty} -> TypeOf (ty1 <+> ty2) -> TypeOf ty1

  inL {ty1 = N, ty2 = N} = id
  inL {ty1 = N, ty2 = ty2 ~> ty2'} = \n => const (Sem.inL n)
  inL {ty1 = ty1 ~> ty1', ty2 = N} = \f => Sem.inL . f
  inL {ty1 = ty1 ~> ty1', ty2 = ty2 ~> ty2'} = \f => Sem.inL . f . Sem.prL

  prL {ty1 = N, ty2 = N} = id
  prL {ty1 = N, ty2 = ty2 ~> ty2'} = \f => Sem.prL (f $ arb ty2)
  prL {ty1 = ty1 ~> ty1', ty2 = N} = \f => Sem.prL . f
  prL {ty1 = ty1 ~> ty1', ty2 = ty2 ~> ty2'} = \f => Sem.prL . f . Sem.inL

  public export
  inR : {ty1, ty2 : Ty} -> TypeOf ty2 -> TypeOf (ty1 <+> ty2)
  inR = Sem.swap . inL

  public export
  prR : {ty1, ty2 : Ty} -> TypeOf (ty1 <+> ty2) -> TypeOf ty2
  prR = prL . Sem.swap

-- Proofs

export
swapInvolutive :
  FunExt =>
  {ty1, ty2 : Ty} ->
  (0 x : TypeOf (ty1 <+> ty2)) ->
  Sem.swap {ty1 = ty2, ty2 = ty1} (Sem.swap {ty1, ty2} x) = x
swapInvolutive {ty1 = N, ty2 = N} x = Refl
swapInvolutive {ty1 = N, ty2 = ty2 ~> ty2'} x = funExt (\y => swapInvolutive (x y))
swapInvolutive {ty1 = ty1 ~> ty1', ty2 = N} x = funExt (\y => swapInvolutive (x y))
swapInvolutive {ty1 = ty1 ~> ty1', ty2 = ty2 ~> ty2'} x = funExt (\y =>
  Calc $
    |~ (Sem.swap . Sem.swap . x . Sem.swap . Sem.swap) y
    ~~ (x . Sem.swap . Sem.swap) y                       ...(swapInvolutive (x _))
    ~~ x y                                               ...(cong x $ swapInvolutive y))

export
retractL :
  FunExt =>
  {ty1, ty2 : Ty} ->
  (0 x : TypeOf ty1) ->
  prL {ty1, ty2} (inL {ty1, ty2} x) = x
retractL {ty1 = N, ty2 = N} x = Refl
retractL {ty1 = N, ty2 = ty2 ~> ty2'} x = retractL {ty1 = N, ty2 = ty2'} x
retractL {ty1 = ty1 ~> ty1', ty2 = N} x = funExt (\y => retractL (x y))
retractL {ty1 = ty1 ~> ty1', ty2 = ty2 ~> ty2'} x = funExt (\y =>
  Calc $
    |~ (Sem.prL . Sem.inL . x . Sem.prL {ty2} . Sem.inL) y
    ~~ (x . Sem.prL {ty2} . Sem.inL) y                     ...(retractL (x _))
    ~~ x y                                                 ...(cong x $ retractL {ty2} y))

export
retractR :
  FunExt =>
  {ty1, ty2 : Ty} ->
  (0 x : TypeOf ty2) ->
  prR {ty1, ty2} (inR {ty1, ty2} x) = x
retractR x = Calc $
  |~ (Sem.prL . Sem.swap . Sem.swap . Sem.inL) x
  ~~ (Sem.prL . Sem.inL) x                       ...(cong prL $ swapInvolutive (inL x))
  ~~ x                                           ...(retractL x)