module Encoded.Sum

import public Data.List
import public Data.List.Elem
import public Data.List.Quantifiers

import Control.Function.FunExt
import Encoded.Bool
import Encoded.Pair
import Encoded.Union
import Syntax.PreorderReasoning
import Term.Semantics
import Term.Syntax

%ambiguity_depth 4

-- Auxilliary Functions --------------------------------------------------------

public export
mapNonEmpty : NonEmpty xs -> NonEmpty (map f xs)
mapNonEmpty IsNonEmpty = IsNonEmpty

-- Types -----------------------------------------------------------------------

export
(+) : Ty -> Ty -> Ty
ty1 + ty2 = B * (ty1 <+> ty2)

export
Sum : (tys : List Ty) -> {auto 0 ok : NonEmpty tys} -> Ty
Sum = foldr1 (+)

-- Universal Properties --------------------------------------------------------

-- Binary

export
left : {ty1, ty2 : Ty} -> Term (ty1 ~> (ty1 + ty2)) ctx
left = Abs' (\t => App pair [<True, App inL [<t]])

export
right : {ty1, ty2 : Ty} -> Term (ty2 ~> (ty1 + ty2)) ctx
right = Abs' (\t => App pair [<False, App inR [<t]])

export
case' : {ty1, ty2, ty : Ty} -> Term ((ty1 + ty2) ~> (ty1 ~> ty) ~> (ty2 ~> ty) ~> ty) ctx
case' = Abs' (\t =>
  App if'
    [<App fst [<t]
    , Abs $ Const $ App (Var Here . prL . snd) [<shift t]
    , Const $ Abs $ App (Var Here . prR . snd) [<shift t]
    ])

-- Utilty

export
either : {ty1, ty2, ty : Ty} -> Term ((ty1 ~> ty) ~> (ty2 ~> ty) ~> (ty1 + ty2) ~> ty) ctx
either = Abs $ Abs $ Abs $
  let f = Var $ There $ There Here in
  let g = Var $ There Here in
  let x = Var Here in
  App case' [<x, f, g]

-- N-ary

export
any :
  {tys : List Ty} ->
  {ty : Ty} ->
  {auto 0 ok : NonEmpty tys} ->
  All (\ty' => Term (ty' ~> ty) ctx) tys ->
  Term (Sum tys ~> ty) ctx
any [f] = f
any (f :: fs@(_ :: _)) = App either [<f, any fs]

export
tag :
  {tys : List Ty} ->
  {ty : Ty} ->
  {auto 0 ok : NonEmpty tys} ->
  Elem ty tys ->
  Term (ty ~> Sum tys) ctx
tag {tys = [_]} Here = Id
tag {tys = _ :: _ :: _} Here = left
tag {tys = _ :: _ :: _} (There i) = right . tag i

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

-- Conversion to Idris

export
toEither :
  {ty1, ty2 : Ty} ->
  TypeOf (ty1 + ty2) ->
  Either (TypeOf ty1) (TypeOf ty2)
toEither x =
  if' (Sem.fst x)
    (Left (prL $ Sem.snd x))
    (Right (prR $ Sem.snd x))

export
fromEither :
  {ty1, ty2 : Ty} ->
  Either (TypeOf ty1) (TypeOf ty2) ->
  TypeOf (ty1 + ty2)
fromEither (Left x) = pair Sem.True (inL x)
fromEither (Right x) = pair Sem.False (inR x)

export
toFromId :
  FunExt =>
  {ty1, ty2 : Ty} ->
  (x : Either (TypeOf ty1) (TypeOf ty2)) ->
  toEither {ty1, ty2} (fromEither {ty1, ty2} x) = x
toFromId (Left x) =
  rewrite retractL {ty1 = B, ty2 = ty1 <+> ty2} Sem.True in
  cong Left $ Calc $
    |~ (Sem.prL {ty1, ty2} . Sem.prR . Sem.inR . Sem.inL) x
    ~~ (prL {ty1, ty2} . inL) x                             ...(cong prL $ retractR (inL x))
    ~~ x                                                    ...(retractL {ty1, ty2} x)
toFromId (Right x) =
  rewrite retractL {ty1 = B, ty2 = ty1 <+> ty2} Sem.False in
  cong Right $ Calc $
    |~ (Sem.prR {ty1, ty2} . Sem.prR . Sem.inR . Sem.inR) x
    ~~ (prR {ty1, ty2} . inR) x                             ...(cong prR $ retractR (inR x))
    ~~ x                                                    ...(retractR {ty1, ty2} x)

-- Redefinition in Idris

namespace Sem
  public export
  left : {ty1, ty2 : Ty} -> TypeOf ty1 -> TypeOf (ty1 + ty2)
  left x = pair Sem.True (inL x)

  public export
  right : {ty1, ty2 : Ty} -> TypeOf ty2 -> TypeOf (ty1 + ty2)
  right x = pair Sem.False (inR x)

  public export
  case' :
    {ty1, ty2 : Ty} ->
    TypeOf (ty1 + ty2) ->
    (TypeOf ty1 -> a) ->
    (TypeOf ty2 -> a) ->
    a
  case' x =
    if' (Sem.fst x)
      (\f, _ => f $ prL $ Sem.snd x)
      (\_, g => g $ prR $ Sem.snd x)

-- Homomorphism

export
leftHomo :
  FunExt =>
  {ty1, ty2 : Ty} ->
  (0 x : TypeOf ty1) ->
  toEither {ty1, ty2} (left {ty1, ty2} x) = Left x
leftHomo x = irrelevantEq $ toFromId (Left x)

export
rightHomo :
  FunExt =>
  {ty1, ty2 : Ty} ->
  (0 x : TypeOf ty2) ->
  toEither {ty1, ty2} (right {ty1, ty2} x) = Right x
rightHomo x = irrelevantEq $ toFromId (Right x)

export
caseHomo :
  FunExt =>
  {ty1, ty2 : Ty} ->
  (x : Either (TypeOf ty1) (TypeOf ty2)) ->
  (f : TypeOf ty1 -> a) ->
  (g : TypeOf ty2 -> a) ->
  case' {ty1, ty2} (fromEither {ty1, ty2} x) f g = either f g x
caseHomo (Left x) f g =
  rewrite retractL {ty1 = B, ty2 = ty1 <+> ty2} Sem.True in
  cong f $ Calc $
    |~ (Sem.prL {ty1, ty2} . Sem.prR . Sem.inR . Sem.inL) x
    ~~ (prL {ty1, ty2} . inL) x                             ...(cong prL $ retractR (inL x))
    ~~ x                                                    ...(retractL {ty1, ty2} x)
caseHomo (Right x) f g =
  rewrite retractL {ty1 = B, ty2 = ty1 <+> ty2} Sem.False in
  cong g $ Calc $
    |~ (Sem.prR {ty1, ty2} . Sem.prR . Sem.inR . Sem.inR) x
    ~~ (prR {ty1, ty2} . inR) x                             ...(cong prR $ retractR (inR x))
    ~~ x                                                    ...(retractR {ty1, ty2} x)