{- SSGIP Examples, Stephanie Weirich, March 24-26, 2010 -}
{-# OPTIONS --type-in-type --no-termination-check #-}
module ssgip.nmap where
open import Data.Unit
open import Data.Sum hiding (map)
open import Data.Nat 
open import Data.Product hiding (map)
open import Relation.Binary.PropositionalEquality
open import Data.Vec hiding (_⊛_)
open import Data.Bool

open import ssgip.aritygen

error : ∀ {A} → A
error = error

eq-nat : ℕ → ℕ → Bool
eq-nat zero     zero     = true
eq-nat (suc n)  (suc m)  = eq-nat n m 
eq-nat _        _        = false

map-sum : ∀ {A₁ A₂ B₁ B₂} → (A₁ → B₁) 
           → (A₂ → B₂) → (A₁ ⊎ A₂) → B₁ ⊎ B₂
map-sum ra rb (inj₁ x) = inj₁ (ra x)
map-sum ra rb (inj₂ x) = inj₂ (rb x)

map-prod : ∀ {A₁ A₂ B₁ B₂} → (A₁ → B₁) 
           → (A₂ → B₂) → (A₁ × A₂) → B₁ × B₂
map-prod ra rb (x , y) = (ra x , rb y)

zip-sum : ∀ {A₁ A₂ A₃ B₁ B₂ B₃} → (A₁ → A₂ → A₃) 
           → (B₁ → B₂ → B₃) → (A₁ ⊎ B₁) → A₂ ⊎ B₂ → A₃ ⊎ B₃
zip-sum ra rb (inj₁ x)(inj₁ y) = inj₁ (ra x y)
zip-sum ra rb (inj₂ x)(inj₂ y) = inj₂ (rb x y) 
zip-sum ra rb _ _ = error

zip-prod : ∀ {A₁ A₂ A₃ B₁ B₂ B₃} → (A₁ → A₂ → A₃) 
           → (B₁ → B₂ → B₃) → (A₁ × B₁) → (A₂ × B₂) → A₃ × B₃
zip-prod ra rb (x , y) ( w , z ) = (ra x w , rb y z)


Repeat : Vec Set  → Set
Repeat ( A ∷ [] ) = A 

grepeat : {k : Kind} → (t : Ty k) → Repeat ⟨ k ⟩ (ι ⌊ t ⌋)
grepeat t = ngen t re (\ {As} → rb {As}) where
  re : ConstEnv Repeat
  re Unit = tt
  re Sum  = g where
    g : Repeat ⟨ ⋆ ⇒ ⋆ ⇒ ⋆ ⟩ (ι _⊎_ )
    g {A ∷ []} ra {B ∷ []} rb = inj₁ ra
  re Prod = g where 
    g : Repeat ⟨ ⋆ ⇒ ⋆ ⇒ ⋆ ⟩ (ι _×_ )
    g {A ∷ []} ra {B ∷ []} rb = (ra , rb)
  
  rb : MuGen Repeat
  rb {A ∷ []} = roll

Map : Vec Set  → Set 
Map ( A ∷ B ∷ [] ) = A → B
gmap : ∀ {k : Kind} → (t : Ty k) → Map ⟨ k ⟩ (ι ⌊ t ⌋)
gmap t = ngen t re rb where
  re : ConstEnv Map
  re Unit = \ x → x
  re Sum  = g where
     g : Map ⟨ ⋆ ⇒ ⋆ ⇒ ⋆ ⟩ (ι _⊎_)
     g {_ ∷ _ ∷ []} ra {_ ∷ _ ∷ []} rb = map-sum ra rb
  re Prod = g where
     g : Map ⟨ ⋆ ⇒ ⋆ ⇒ ⋆ ⟩ (ι _×_)
     g {_ ∷ _ ∷ []} ra {_ ∷ _ ∷ []} rb = map-prod ra rb 
  
  rb : ∀ {As} → Map (As ⊛ (ι μ ⊛ As)) → Map (ι μ ⊛ As)
  rb {_ ∷ _ ∷ []} = \ x y → roll (x (unroll y))
ZW : Vec Set  → Set 
ZW ( A ∷ B ∷ C ∷ [] ) = A → B → C
gzipWith : ∀ {k} → (t : Ty k) → ZW ⟨ k ⟩ (ι ⌊ t ⌋ )
gzipWith t = ngen t re rb where
  re : ConstEnv ZW
  re Unit = \ x y → x
  re Sum  = g where
     g : ZW ⟨ ⋆ ⇒ ⋆ ⇒ ⋆ ⟩ (ι _⊎_)
     g {_ ∷ _ ∷ _ ∷ []} ra {_ ∷ _ ∷ _ ∷ []} rb = zip-sum ra rb 
  re Prod = g where
     g : ZW ⟨ ⋆ ⇒ ⋆ ⇒ ⋆ ⟩ (ι _×_)
     g {_ ∷ _ ∷ _ ∷ []} ra {_ ∷ _ ∷ _ ∷ []} rb = zip-prod ra rb   
  rb : ∀ {As} → ZW (As ⊛ (ι μ ⊛ As)) → ZW (ι μ ⊛ As)
  rb {_ ∷ _ ∷ _ ∷ []} = \ x y z → roll (x (unroll y) (unroll z))
NGmap : {n : ℕ} → Vec Set (suc n) → Set
NGmap (A ∷ []) = A
NGmap (A ∷ B ∷ As ) = A → NGmap (B ∷ As)
defUnit : (n : ℕ) → NGmap {n} ⟨ ⋆ ⟩ (ι ⊤)
-- (n : ℕ) → ⊤ → ⊤ → ... → ⊤
defUnit zero     = tt
defUnit (suc n)  = \ x → (defUnit n)

defPair  : (n : ℕ) 
         → {As : Vec Set (suc n)} → NGmap As
         → {Bs : Vec Set (suc n)} → NGmap Bs 
         → NGmap (ι _×_ ⊛ As ⊛ Bs)
-- (n : ℕ) → (A₁ → A₂ → .... An)
--  → (B₁ → B₂ → .... Bn) 
--  → (A₁ × B₁ → A₂ × B₂ → .... An × Bn)
defPair zero     {A ∷ []}        a  {B ∷ []}        b  = (a , b)
defPair (suc n)  {A₁ ∷ A₂ ∷ As}  a  {B₁ ∷ B₂ ∷ Bs}  b  =
 \ p →
  defPair n  {A₂ ∷ As} (a (proj₁ p)) 
             {B₂ ∷ Bs} (b (proj₂ p))
defSum : (n : ℕ) 
         → {As : Vec Set (suc n)} → NGmap As
         → {Bs : Vec Set (suc n)} → NGmap Bs
         → NGmap (ι _⊎_ ⊛ As ⊛ Bs)
defSum zero     {(A ∷ [])}      a {B ∷ []}        b  = 
     (inj₂ b) 
defSum (suc n)  {A₁ ∷ A₂ ∷ As}  a {B₁ ∷ B₂ ∷ Bs}  b  = f
  where 
    f : A₁ ⊎ B₁ → NGmap (ι _⊎_ ⊛ (A₂ ∷ As) ⊛ (B₂ ∷ Bs))
    f (inj₁ a₁)  = defSum n {A₂ ∷ As} (a a₁) {B₂ ∷ Bs} (b error)
    f (inj₂ b₁)  = defSum n {A₂ ∷ As} (a error) {B₂ ∷ Bs} (b b₁)

ngmap-const : {n : ℕ} → ConstEnv {n} NGmap
ngmap-const {n} Unit  = defUnit n
ngmap-const {n} Prod  = defPair n
ngmap-const {n} Sum   = defSum n

ngmap-mu : ∀ {n} → MuGen {n} NGmap
ngmap-mu {zero}   {A ∷ []}        = roll
ngmap-mu {suc n}  {A₁ ∷ A₂ ∷ As}  = \ f x → 
  ngmap-mu {n}{A₂ ∷ As} (f (unroll x))
ngmap  : (n : ℕ) → {k : Kind} → (e : Ty k) 
       → NGmap {n} ⟨ k ⟩ (ι ⌊ e ⌋)
ngmap n e = ngen e  ngmap-const
                    (\ {As} → ngmap-mu {n} {As})
 

repeat-ml    : ∀ {B} → B → List B
repeat-ml    = ngmap  list {_ ∷ []} 

map-ml       : ∀ {A₁ B} → (A₁ → B) → List A₁ → List B 
map-ml       = ngmap  list {_ ∷ _ ∷ []} 

zipWith-ml   : ∀ {A₁ A₂ B} → (A₁ → A₂ → B)
             → List A₁ → List A₂ → List B
zipWith-ml   = ngmap  list {_ ∷ _ ∷ _ ∷ []} 

zipWith3-ml  : ∀ {A₁ A₂ A₃ B} → (A₁ → A₂ → A₃ → B)
             → List A₁ → List A₂ → List A₃ → List B 
zipWith3-ml  = ngmap  list {_ ∷ _ ∷ _ ∷ _ ∷ []}