------------------------------------------------------------------------
-- Another implementation of the Fibonacci sequence using tail, along
-- with some other examples
------------------------------------------------------------------------

module ArbitraryChunks where

open import Codata.Musical.Notation
open import Codata.Musical.Stream as S using (Stream; _≈_; _∷_)
open import Data.Bool
open import Data.Nat
open import Data.Vec as V using (Vec; []; _∷_)
open import Relation.Binary.PropositionalEquality as P
  using (_≡_; refl)

------------------------------------------------------------------------
-- Stream programs

-- StreamP m n A encodes programs generating streams in chunks of size
-- (at least) m, where the first chunk has size (at least) n.

infixr  _∷_

data StreamP (m : ℕ) : ℕ → Set → Set₁ where
  [_]     : ∀ {A} (xs : ∞ (StreamP m m A)) → StreamP m  A
  _∷_     : ∀ {n A} (x : A) (xs : StreamP m n A) → StreamP m (suc n) A
  forget  : ∀ {n A} (xs : StreamP m (suc n) A) → StreamP m n A
  tail    : ∀ {n A} (xs : StreamP m (suc n) A) → StreamP m n A
  zipWith : ∀ {n A B C} (f : A → B → C)
            (xs : StreamP m n A) (ys : StreamP m n B) → StreamP m n C
  map     : ∀ {n A B} (f : A → B) (xs : StreamP m n A) → StreamP m n B

data StreamW (m : ℕ) : ℕ → Set → Set₁ where
  [_] : ∀ {A} (xs : StreamP m m A) → StreamW m  A
  _∷_ : ∀ {n A} (x : A) (xs : StreamW m n A) → StreamW m (suc n) A

forgetW : ∀ {m n A} → StreamW (suc m) (suc n) A → StreamW (suc m) n A
forgetW {n = zero}  (x ∷ [ xs ]) = [ x ∷ forget xs ]
forgetW {n = suc n} (x ∷ xs)     = x ∷ forgetW xs

tailW : ∀ {m n A} → StreamW m (suc n) A → StreamW m n A
tailW (x ∷ xs) = xs

zipWithW : ∀ {m n A B C} → (A → B → C) →
           StreamW m n A → StreamW m n B → StreamW m n C
zipWithW f [ xs ]   [ ys ]   = [ zipWith f xs ys ]
zipWithW f (x ∷ xs) (y ∷ ys) = f x y ∷ zipWithW f xs ys

mapW : ∀ {m n A B} → (A → B) → StreamW m n A → StreamW m n B
mapW f [ xs ]   = [ map f xs ]
mapW f (x ∷ xs) = f x ∷ mapW f xs

whnf : ∀ {m n A} → StreamP (suc m) n A → StreamW (suc m) n A
whnf [ xs ]            = [ ♭ xs ]
whnf (x ∷ xs)          = x ∷ whnf xs
whnf (forget xs)       = forgetW (whnf xs)
whnf (tail xs)         = tailW (whnf xs)
whnf (zipWith f xs ys) = zipWithW f (whnf xs) (whnf ys)
whnf (map f xs)        = mapW f (whnf xs)

mutual

  ⟦_⟧W : ∀ {m n A} → StreamW (suc m) (suc n) A → Stream A
  ⟦ x ∷ [ xs ]   ⟧W = x ∷ ♯ ⟦ xs ⟧P
  ⟦ x ∷ (y ∷ xs) ⟧W = x ∷ ♯ ⟦ y ∷ xs ⟧W

  ⟦_⟧P : ∀ {m n A} → StreamP (suc m) (suc n) A → Stream A
  ⟦ xs ⟧P = ⟦ whnf xs ⟧W

------------------------------------------------------------------------
-- Some examples

-- The Fibonacci sequence.

fib : StreamP   ℕ
fib =  ∷ [ ♯ ( ∷ zipWith _+_ (forget fib) (tail fib)) ]

-- Some sequence which is defined using (tail ∘ tail).

someSequence : StreamP   ℕ
someSequence =
   ∷  ∷ [ ♯ ( ∷  ∷ zipWith _+_ (tail (tail someSequence))
                                   (forget (forget someSequence))) ]

------------------------------------------------------------------------
-- The definition of fib is correct

-- ⟦_⟧ is homomorphic with respect to zipWith/S.zipWith.

mutual

  zipWithW-hom :
    ∀ {m n A B C} (_∙_ : A → B → C)
    (xs : StreamW (suc m) (suc n) A) ys →
    ⟦ zipWithW _∙_ xs ys ⟧W ≈ S.zipWith _∙_ ⟦ xs ⟧W ⟦ ys ⟧W
  zipWithW-hom _∙_ (x ∷ [ xs ]) (y ∷ [ ys ])   =
    refl ∷ ♯ zipWith-hom _∙_ xs ys
  zipWithW-hom _∙_ (x ∷ x′ ∷ xs) (y ∷ y′ ∷ ys) =
    refl ∷ ♯ zipWithW-hom _∙_ (x′ ∷ xs) (y′ ∷ ys)

  zipWith-hom : ∀ {m n A B C} (_∙_ : A → B → C)
                (xs : StreamP (suc m) (suc n) A) ys →
                ⟦ zipWith _∙_ xs ys ⟧P ≈ S.zipWith _∙_ ⟦ xs ⟧P ⟦ ys ⟧P
  zipWith-hom _∙_ xs ys = zipWithW-hom _∙_ (whnf xs) (whnf ys)

-- forget is the identity on streams.

open import MapIterate using (_≈P_; _∷_; _≈⟨_⟩_; _∎; soundP)

mutual

  forgetW-lemma : ∀ {m n A} x (xs : StreamW (suc m) (suc n) A) →
                  ⟦ x ∷ forgetW xs ⟧W ≈P x ∷ ♯ ⟦ xs ⟧W
  forgetW-lemma x (x′ ∷ [ xs ])  = x ∷ ♯ (_ ≈⟨ forget-lemma x′ xs ⟩ x′ ∷ ♯ (_ ∎))
  forgetW-lemma x (x′ ∷ x″ ∷ xs) = x ∷ ♯ (_ ≈⟨ forgetW-lemma x′ (x″ ∷ xs) ⟩ x′ ∷ ♯ (_ ∎) )

  forget-lemma : ∀ {m n A} x (xs : StreamP (suc m) (suc n) A) →
                 ⟦ x ∷ forget xs ⟧P ≈P x ∷ ♯ ⟦ xs ⟧P
  forget-lemma x xs with whnf xs in eq
  ... | y ∷ y′ ∷ ys = x ∷ ♯ helper eq
    where
    helper : whnf xs ≡ y ∷ y′ ∷ ys →
             ⟦ y ∷ forgetW (y′ ∷ ys) ⟧W ≈P ⟦ xs ⟧P
    helper eq rewrite eq = _ ≈⟨ forgetW-lemma y (y′ ∷ ys) ⟩ y ∷ ♯ (_ ∎)
  ... | y ∷ [ ys ] = x ∷ ♯ helper eq
    where
    helper : whnf xs ≡ y ∷ [ ys ] →
             ⟦ y ∷ forget ys ⟧P ≈P ⟦ xs ⟧P
    helper eq rewrite eq = _ ≈⟨ forget-lemma y ys ⟩ y ∷ ♯ (_ ∎)

-- The stream ⟦ fib ⟧P satisfies its intended defining equation.

open import Relation.Binary
module SS {A : Set} = Setoid (S.setoid A)

fib-correct :
  ⟦ fib ⟧P ≈  ∷ ♯ ( ∷ ♯ S.zipWith _+_ ⟦ fib ⟧P (S.tail ⟦ fib ⟧P))
fib-correct =
  refl ∷ ♯ (refl ∷ ♯ SS.trans
    (zipWith-hom _+_ ( ∷ forget fib′) fib′)
    (S.zipWith-cong _+_ (SS.trans (soundP (forget-lemma  fib′))
                                  (refl ∷ ♯ SS.refl)) SS.refl))
  where fib′ =  ∷ zipWith _+_ (forget fib) (tail fib)

------------------------------------------------------------------------
-- map₂

-- A variant of S.map which processes streams in chunks of size 2.

map₂ : {A B : Set} → (A → B) → Stream A → Stream B
map₂ f (x ∷ xs) with ♭ xs
map₂ f (x ∷ xs) | y ∷ ys = f x ∷ ♯ (f y ∷ ♯ map₂ f (♭ ys))

-- This function is extensionally equal to S.map.

map≈map₂ : ∀ {A B} →
           (f : A → B) → (xs : Stream A) → S.map f xs ≈ map₂ f xs
map≈map₂ f (x ∷ xs) with ♭ xs in eq
map≈map₂ f (x ∷ xs) | y ∷ ys = refl ∷ ♯ helper eq
  where
  map-f-y∷ys = _

  helper : ∀ {xs} → xs ≡ y ∷ ys → S.map f xs ≈ map-f-y∷ys
  helper refl = refl ∷ ♯ map≈map₂ f (♭ ys)

-- However, as explained in "Beating the Productivity Checker Using
-- Embedded Languages", the two functions are not interchangeable
-- (assuming that pattern matching is evaluated strictly). Let us see
-- what happens when we use the language above. We can define the
-- stream of natural numbers using chunks of size 1 and 2.

nats : StreamP   ℕ
nats =  ∷ [ ♯ map suc nats ]

nats₂ : StreamP   ℕ
nats₂ =  ∷  ∷ [ ♯ map suc nats₂ ]

-- The first use of map corresponds to S.map, and the second to map₂.
-- If we try to use map₂ in the first definition, then we get a type
-- error. The following code is not accepted.

-- nats : StreamP 2 1 ℕ
-- nats = 0 ∷ [ ♯ map suc nats ]