module CoinductiveRecursion 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

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 

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) ↓

-- Safe computation and evaluation

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

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

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

eval-safe : {a : Set} -> (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 ↓=-return ↓=-return = ≡-refl 
↓=-unique {x = x} (↓=-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 ↓
↓=⇒↓ ↓=-return = ↓-return
↓=⇒↓ (↓=-bind {m}{f} 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)

unfold : {a : Set} -> Comp a -> Comp a
unfold (return x) = return x
unfold (m ⟩⟩= f) = m ⟩⟩= f

≡-unfold : {a : Set} {m : Comp a} -> unfold m ≡ m
≡-unfold {_}{return x} = ≡-refl
≡-unfold {_}{m ⟩⟩= f} = ≡-refl

half↓ : (n : ℕ) -> half n ↓
half↓ = strong-induction (\n -> half n ↓) 
               (\n f -> ≡-subst _↓ ≡-unfold (ind n f))
  where ind : forall n -> (forall m -> m <′ n -> half m ↓) -> unfold (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 m (suc n) = {! !}

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

k≤n⇒fkn↓=n : forall k n -> k ≤′′ n  -> f k n ↓= n
k≤n⇒fkn↓=n k n k≤n = ≡-subst (\m -> m ↓= n) (≡-unfold {_}{f k n}) (k≤n⇒⟨fkn⟩↓=n k n k≤n)

mutual

  n<k⇒⟨fkn⟩↓=k : forall k n -> n <′′ k -> unfold (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 (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 (n<k⇒fkn↓=k k (suc n) 2+n≤k) (k≤n⇒fkn↓=n k k ≤′′-refl)

  n<k⇒fkn↓=k : forall k n -> n <′′ k -> f k n ↓= k
  n<k⇒fkn↓=k k n n<k = 
      ≡-subst (\m -> m ↓= k) (≡-unfold {_}{f k n}) (n<k⇒⟨fkn⟩↓=k k n n<k)

f↓ : forall k n -> f k n ↓
f↓ k n with k ≤′′? n
... | inj₁ k≤n = ↓=⇒↓ (k≤n⇒fkn↓=n k n k≤n)
... | inj₂ n<k = ↓=⇒↓ (n<k⇒fkn↓=k k n n<k)



��n)
... | inj₂ n<k = ↓=⇒↓ (n<k⇒fkn↓=k k n n<k)