module Iteratee where

open import Function
open import Data.Nat
open import Data.Fin hiding (_+_)
open import Data.Vec hiding ([_])
open import Data.Vec.Properties
open import Data.Product renaming (map to mapΣ)
open import Relation.Binary.PropositionalEquality

data Input (I : Set) : Set where
  eoi   : Input I
  chunk : {n : _}(xs : Vec I n) → Input I

mutual
  -- Todo: encode chunk-invariance into the consumer? Tricky but it'd be super sweet
  data Consumer (I O : Set) : Set where
    satisfied : (x : O) → Consumer I O
    hungry    : (f : (xs : Input I) → Output O xs) → Consumer I O

  Output : ∀ {I} O → Input I → Set
  Output O eoi = O
  Output {I} O (chunk {n} xs) = Fin (1 + n) × Consumer I O

mutual
  data _≈_ {I O : Set} : (C D : Consumer I O) → Set where
    satisfied : ∀ x → satisfied x ≈ satisfied x
    hungry    : ∀ f g → (eq : ∀ xs → Equal xs (f xs) (g xs)) → hungry f ≈ hungry g

  data Equal {I O : Set} : (i : Input I) (X Y : Output O i) → Set where
    satisfied : ∀ x → Equal eoi x x
    hungry    : ∀ {l}{xs : Vec I l} (n : Fin (1 + l)) {C D} (eq : C ≈ D) → Equal (chunk xs) (n , C) (n , D)


  ≈-refl : ∀ {I O} (C : Consumer I O) → C ≈ C
  ≈-refl (satisfied x) = satisfied x
  ≈-refl (hungry f) = hungry f f (λ xs → Equal-refl (f xs))

  ≈-sym : ∀ {I O} {C D : Consumer I O} → C ≈ D → D ≈ C
  ≈-sym (satisfied x) = satisfied x
  ≈-sym (hungry f g eq) = hungry g f (λ xs → Equal-sym (eq xs))

  ≈-trans : ∀ {I O} {C D E : Consumer I O} → C ≈ D → D ≈ E → C ≈ E
  ≈-trans (satisfied x) (satisfied .x) = satisfied x
  ≈-trans (hungry f g eq) (hungry .g h eq′) = hungry f h (λ xs → Equal-trans (eq xs) (eq′ xs))

  Equal-refl : ∀ {I O} {i : Input I} (X : Output O i) → Equal i X X
  Equal-refl {i = eoi} x = satisfied x
  Equal-refl {i = chunk xs} (n , C) = hungry n (≈-refl C)

  Equal-sym : ∀ {I O} {i : Input I} {X Y : Output O i} → Equal i X Y → Equal i Y X
  Equal-sym (satisfied x) = satisfied x
  Equal-sym (hungry n eq) = hungry n (≈-sym eq)

  Equal-trans : ∀ {I O} {i : Input I} {X Y Z : Output O i} → Equal i X Y → Equal i Y Z → Equal i X Z
  Equal-trans (satisfied x) (satisfied .x) = satisfied x
  Equal-trans (hungry n eq) (hungry .n eq′) = hungry n (≈-trans eq eq′)


terminate : ∀ {I O} → Consumer I O → O
terminate (satisfied x) = x
terminate (hungry f) = f eoi

mapC : ∀ {I A B} → (A → B) → Consumer I A → Consumer I B
mapC f (satisfied x) = satisfied (f x)
mapC f (hungry g) = hungry helper
  module mapC where
  helper : (xs : Input _) → Output _ xs
  helper eoi = f (g eoi)
  helper (chunk xs) with g (chunk xs)
  ...               | n , C = n , mapC f C

mapC-id : ∀ {I A} (C : Consumer I A) → mapC id C ≈ C
mapC-id (satisfied x) = satisfied x
mapC-id (hungry f) = hungry (mapC.helper id f) f helper
 module mapC-id where
 helper : (xs : Input _) → Equal xs (mapC.helper id f xs) (f xs)
 helper eoi = Equal-refl (f eoi)
 helper (chunk xs) = hungry (proj₁ (f (chunk xs))) (mapC-id (proj₂ (f (chunk xs))))

mapC-comp : ∀ {I A B C} (f : B → C) (g : A → B) (C : Consumer I A) → mapC (f ∘ g) C ≈ mapC f (mapC g C)
mapC-comp f g (satisfied x) = satisfied (f (g x))
mapC-comp f g (hungry h) = hungry (mapC.helper (f ∘ g) h) (mapC.helper f (mapC.helper g h)) helper
  module mapC-comp where
  helper : (xs : Input _) → Equal xs (mapC.helper (f ∘ g) h xs) (mapC.helper f (mapC.helper g h) xs)
  helper eoi = satisfied (f (g (h eoi)))
  helper (chunk xs) = hungry (proj₁ (h (chunk xs))) (mapC-comp f g (proj₂ (h (chunk xs))))



cmapC : ∀ {A B O} → (A → B) → Consumer B O → Consumer A O
cmapC f (satisfied x) = satisfied x
cmapC f (hungry g) = hungry helper
  module cmapC where
  helper : (xs : Input _) → Output _ xs
  helper eoi = g eoi
  helper (chunk xs) with (g (chunk (map f xs)))
  ...               | n , C = n , cmapC f C

  helper-chunk₂ : ∀ {l} (xs : Vec _ l) → proj₂ (helper (chunk xs)) ≡ cmapC f (proj₂ (g (chunk (map f xs))))
  helper-chunk₂ xs = refl

cmapC-id : ∀ {I A} (C : Consumer I A) → cmapC id C ≈ C
cmapC-id (satisfied x) = satisfied x
cmapC-id (hungry f) = hungry (cmapC.helper id f) f helper
  module cmapC-id where
  helper : (xs : Input _) → Equal xs (cmapC.helper id f xs) (f xs)
  helper eoi = satisfied (f eoi)
  helper (chunk xs) rewrite map-id xs with f (chunk xs) | cmapC-id (proj₂ (f (chunk xs)))
  ... | x , y | eq = hungry x eq