module Coalgebras where

open import Data.Nat.Base
open import Data.Nat.DivMod
open import Data.Fin using (Fin; zero; suc)

-- Coalgebras in Agda

-- Given a functor F, an F-coalgebra is an object S with a morphism
-- t : S → F S.
-- A typical interpretation would be S is the set of states and
-- t is the transition function which takes us from a state s : S
-- to (t s), which can, depending on F, do some input and/or output
-- and take us (non-deterministically) to a next state or terminate.

-- Examples for coalgebras:
--
--   F S = Bool × (A → S)  automaton accepting A-words  (True = accepting state)
--   F S = A × S            output automaton, generating a stream of outputs of type A
--   F S = 1 + A × S        output automaton, generating a possibly infinite sequence
--
--   Pr S = probability distribution over S
--   F S = A → Pr S

-- The functor generating lists.

data ListF A S : Set where
  nil  : ListF A S
  cons : A → S → ListF A S

mapF : ∀{A S T} (f : S → T) → ListF A S → ListF A T
mapF f nil        = nil
mapF f (cons a s) = cons a (f s)

-- Example: Automaton that outputs Collatz sequence starting with n
--   n --> n/2  if n is even
--   n --> 3n+1 if n is odd

collatzStep : ℕ → ListF ℕ ℕ
collatzStep  = nil
collatzStep n with n divMod 
collatzStep n | result q zero _ = cons n q
collatzStep n | result _ _    _ = cons n ( +  * n)

-- Final/terminal F-coalgebra
--
-- The full computation tree of an automaton.
--
-- In Agda, we have only weakly final coalgebras.
-- Obtained as greates fixed-point of F

record Colist A : Set where
  coinductive
  field force : ListF A (Colist A)

open Colist using (force)

-- Coiteration is defined via copattern matching.

collatz : ℕ → Colist ℕ
force (collatz n) with collatzStep n
... | nil       = nil
... | cons m n' = cons m (collatz n')

-- abstracting out mapF upsets the termination checker.

{-# TERMINATING #-}
collatz' : ℕ → Colist ℕ
force (collatz' n) = mapF collatz' (collatzStep n)

-- General coiteration

{-# TERMINATING #-}
unfold : ∀{A S} (t : S → ListF A S) → S → Colist A
force (unfold t s) = mapF (unfold t) (t s)