module CoinductiveRecursionRight where

open import Data.Sum
open import Data.Product
open import Data.Nat
open import Data.Empty

open import Relation.Nullary
open import Relation.Binary.PropositionalEquality

open import DownwardLeq

-- The Computation Monad

codata Comp (a : Set) : Set where
  return : a -> Comp a
  _⟩⟩=_   : Comp a -> (a -> Comp a) -> Comp a
 
-- Termination arguments

mutual 

  {-
  data _↓=_ {a : Set} : Comp a -> a -> Set where
    ↓=-return : forall {x} -> (return x) ↓= x
    ↓=-bind : forall {m f x y} ->
        m ↓= x -> (f x) ↓= y -> (m ⟩⟩= f) ↓= y 
  -}
  
  _↓=_ : {a : Set} -> Comp a -> a -> Set
  return x ↓= y = ↓=-Return x y
  (m ⟩⟩= f) ↓= x = ↓=-Bind m f x 
   
  data ↓=-Return {a : Set} : a -> a -> Set where
    ↓=-return : forall {x} -> ↓=-Return x x
  
  data ↓=-Bind {a : Set} : Comp a -> (a -> Comp a) -> a -> Set where
    ↓=-bind : forall {m f x y} -> 
                   m ↓= x -> (f x) ↓= y -> ↓=-Bind m f y

mutual

  {-
  data _↓ {a : Set} : Comp a -> Set where
    ↓-return : forall {x} -> (return x) ↓
    ↓-bind : forall {m f} ->  
        m ↓ -> (forall x -> m ↓= x -> f x ↓) -> (m ⟩⟩= f) ↓
  -}

  _↓ : {a : Set} -> Comp a -> Set
  return x ↓ = ↓-Return x
  (m ⟩⟩= f) ↓ = ↓-Bind m f

  data ↓-Return {a : Set} : a -> Set where
     ↓-return : forall {x} -> ↓-Return x

  data ↓-Bind {a : Set} : Comp a -> (a -> Comp a) -> Set where
     ↓-bind : forall {m f} ->
            m ↓ -> (forall x -> m ↓= x -> f x ↓) -> ↓-Bind m f


codata co-⊥ : Set where

mutual

  {-
    codata _↑ {a : Set} : Comp a -> Set where
      ↑-prev : forall {m f} -> m ↑ -> m ⟩⟩= f ↑ 
      ↑-fun  : forall {m f x} -> 
          (forall x -> m ↓= x -> f x ↑) -> m ⟩⟩= f ↑
  -}

  _↑ : {a : Set} -> Comp a -> Set
  return x ↑ = co-⊥
  (m ⟩⟩= f) ↑ = ↑-Bind m f

  codata ↑-Bind {a : Set} : Comp a -> (a -> Comp a) -> Set where
     ↑-prev : forall {m f} -> m ↑ -> ↑-Bind m f
     ↑-fun  : forall {m f} ->
                 (forall x -> m ↓= x -> f x ↑) -> ↑-Bind m f 
 
-- Safe computation

{-
data Invoked-By {a : Set} : Comp a -> Comp a -> Set where
     invokes-prev : forall {m f} -> Invoked-By m (m ⟩⟩= f)
     invokes-fun  : forall {m f x} -> m ↓= x -> Invoked-By (f x) (m ⟩⟩= f)
-}

data Invoked-By-Bind {a : Set} : Comp a -> Comp a -> (a -> Comp a) -> Set where
     invokes-prev : forall {m f} -> Invoked-By-Bind m m f
     invokes-fun  : forall {m f x} -> m ↓= x -> Invoked-By-Bind (f x) m f

invoked-by : {a : Set} -> Comp a -> Comp a -> Set
invoked-by _ (return _) = ⊥
invoked-by m (n ⟩⟩= f) = Invoked-By-Bind m n f

data Safe {a : Set} : Comp a -> Set where
  safe-intro : forall m -> (forall m' -> invoked-by m' m -> Safe m') -> Safe m 

↓-safe : {a : Set} -> (m : Comp a) -> m ↓ -> Safe m 
↓-safe (return x) ↓-return = safe-intro (return x) fn 
  where fn : forall k -> invoked-by k (return x) -> Safe k
        fn _ ()   
↓-safe (m ⟩⟩= f) (↓-bind m↓ m↓=x⇒fx↓) = safe-intro (m ⟩⟩= f) fn  
  where fn : forall k -> invoked-by k (m ⟩⟩= f) -> Safe k
        fn .m invokes-prev = ↓-safe m m↓ 
        fn ._ (invokes-fun m↓=x) = ↓-safe _ (m↓=x⇒fx↓ _ m↓=x)


-- Evaluation

eval-safe : forall {a} -> (m : Comp a) -> Safe m -> ∃ (\x -> m ↓= x)
eval-safe (return x) (safe-intro ._ safe) = (x , ↓=-return) 
eval-safe (m ⟩⟩= f) (safe-intro ._ safe) with eval-safe m (safe m invokes-prev)
... | (x , m↓=x) with eval-safe (f x) (safe (f x) (invokes-fun m↓=x))
...   | (y , fx↓=y) = (y , ↓=-bind m↓=x fx↓=y) 

eval : {a : Set} (m : Comp a) -> m ↓ -> a
eval m m↓ with eval-safe m (↓-safe m m↓) 
... | (x , m↓=x) = x 

-- Some more auxiliary lemmas.

↓=-unique : {a : Set} {m : Comp a} {x y : a} ->
              m ↓= x -> m ↓= y -> x ≡ y
↓=-unique {m = return x} ↓=-return ↓=-return = ≡-refl 
↓=-unique {m = m ⟩⟩= f} (↓=-bind m↓=x fx↓=y) (↓=-bind m↓=x' fx'↓=y') 
   with ↓=-unique m↓=x m↓=x' 
... | ≡-refl with ↓=-unique fx↓=y fx'↓=y'
...   | ≡-refl = ≡-refl 

↓=⇒↓ : {a : Set} {m : Comp a} {x : a} -> m ↓= x -> m ↓
↓=⇒↓ {m = return x} ↓=-return = ↓-return
↓=⇒↓ {m = m ⟩⟩= f} (↓=-bind m↓=x fx↓=y) = 
  ↓-bind (↓=⇒↓ m↓=x)
    (\z m↓=z -> ≡-subst (\w -> f w ↓) (↓=-unique m↓=x m↓=z) (↓=⇒↓ fx↓=y))

-- Strong induction

strong-induction : (P : ℕ -> Set) ->
       (forall n -> (forall m -> m <′ n -> P m) -> P n) ->
           (forall n -> P n)
strong-induction P Pind n = Pind n (ind n)
  where ind : forall n -> forall m -> m <′ n -> P m
        ind zero _ ()
        ind (suc n) ._ ≤′-refl = strong-induction P Pind n    
        ind (suc n) m (≤′-step m<n) = ind n m m<n 

-- Example: half

half : ℕ -> Comp ℕ
half zero          ~ return zero
half (suc zero)    ~ return zero
half (suc (suc n)) ~ half n ⟩⟩= \ m ->
                     return (suc m)

half↓ : (n : ℕ) -> half n ↓
half↓ = strong-induction (\n -> half n ↓) ind
  where ind : forall n -> (forall m -> m <′ n -> half m ↓) -> half n ↓
        ind zero f = ↓-return
        ind (suc zero) f = ↓-return
        ind (suc (suc n)) f = ↓-bind (f n (≤′-step ≤′-refl)) (\_ _ -> ↓-return)


-- Example: integral division

div : ℕ -> ℕ -> Comp ℕ
div zero n          ~ return zero
div (suc m) zero    ~ return (suc m)
div (suc m) (suc n) ~ div (suc m ∸ suc n) (suc n) ⟩⟩= \h ->
                      return (suc h)

sm-sn<sm : forall m n -> (m ∸ n) <′ suc m
sm-sn<sm m zero = ≤′-refl 
sm-sn<sm zero (suc n) = ≤′-refl 
sm-sn<sm (suc m) (suc n) = ≤′-step (sm-sn<sm m n) 

div↓ : forall m n -> div m n ↓
div↓ m n = strong-induction (\m -> div m n ↓) ind m
  where ind : forall m -> (forall k -> k <′ m -> div k n ↓) -> div m n ↓
        ind zero f = ↓-return  
        ind (suc m) f with n
        ... | zero = ↓-return 
        ... | suc n' = ↓-bind (f (m ∸ n') (sm-sn<sm m n')) 
                         (\_ _ -> ↓-return)

-- Example: McCarthian Maximum

f : ℕ -> ℕ -> Comp ℕ
f k n with k ≤′′? n
... | inj₁ k≤n ~ return n
... | inj₂ n<k ~ f k (suc n)  ⟩⟩= \x ->
                 f k x

k≤n⇒fkn↓=n : forall k n -> k ≤′′ n -> f k n ↓= n
k≤n⇒fkn↓=n k n k≤n with k ≤′′? n
... | inj₁ _ = ↓=-return 
... | inj₂ n<k = ⊥-elim (¬[m<n∧n≤m] (n<k , k≤n)) 

n<k⇒fkn↓=k : forall k n -> n <′′ k -> f k n ↓= k
n<k⇒fkn↓=k k n n<k with k ≤′′? n
n<k⇒fkn↓=k k n n<k | inj₁ k≤n = ⊥-elim (¬[m<n∧n≤m] (n<k , k≤n)) 
n<k⇒fkn↓=k ._ n _ | inj₂ ≤′′-refl = 
   ↓=-bind {_}{f (suc n) (suc n)}{\x -> f (suc n) x}{suc n} 
       (k≤n⇒fkn↓=n (suc n) (suc n) ≤′′-refl) 
       (k≤n⇒fkn↓=n (suc n) (suc n) ≤′′-refl) 
n<k⇒fkn↓=k ._ n ≤′′-refl | inj₂ (≤′′-step 2+n≤n) = ⊥-elim (¬1+n≤n 2+n≤n)  
n<k⇒fkn↓=k k n (≤′′-step 2+n≤k) | inj₂ (≤′′-step _) = 
   ↓=-bind {_}{f k (suc n)}{\x -> f k x}{k} 
       (n<k⇒fkn↓=k k (suc n) 2+n≤k) 
       (k≤n⇒fkn↓=n k k ≤′′-refl) 

≤′′-split : (P : ℕ -> ℕ -> Set) ->
    (forall k n -> k ≤′′ n -> P k n) -> (forall k n -> n <′′ k -> P k n) -> 
      forall k n -> P k n
≤′′-split P k≤n⇒P n<k⇒P k n with k ≤′′? n
... | inj₁ k≤n = k≤n⇒P k n k≤n
... | inj₂ n<k = n<k⇒P k n n<k

f↓ : forall k n -> f k n ↓
f↓ = ≤′′-split (\k n -> f k n ↓)
              (\k n k≤n -> ↓=⇒↓ (k≤n⇒fkn↓=n k n k≤n))
              (\k n n<k -> ↓=⇒↓ (n<k⇒fkn↓=k k n n<k))

-- Example: divergence

codata Stream (a : Set) : Set where
  [] : Stream a
  _∷_ : a -> Stream a -> Stream a

length : {a : Set} -> Stream a -> Comp ℕ
length [] ~ return zero
length (x ∷ xs) ~ length xs ⟩⟩= \n ->
                  return (suc n)

ones : Stream ℕ
ones ~ 1 ∷ ones

length-ones↑ : length ones ↑
length-ones↑ ~ ↑-prev {_}{length ones}{\n -> return (suc n)} length-ones↑