{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}

module Day
  ( -- Day convolutions
    Day(..)
  ) where

import Data.Kind (Type)
import Control.Comonad (Comonad(..), ComonadApply(..))
import Control.Comonad.Trans.Class (ComonadTrans(..))
import Data.Distributive (Distributive(..))
import Data.Functor.Rep (Representable(..))

data Day f g a = forall x y. Day (x -> y -> a) (f x) (g y)

instance Functor (Day f g) where
  fmap :: forall a b. (a -> b) -> Day f g a -> Day f g b
fmap a -> b
g (Day x -> y -> a
f f x
x g y
y) = forall (f :: * -> *) (g :: * -> *) a x y.
(x -> y -> a) -> f x -> g y -> Day f g a
Day (\x
a y
b -> let fab :: a
fab = x -> y -> a
f x
a y
b in a -> b
g a
fab) f x
x g y
y
  {-# INLINE fmap #-}

instance (Comonad f, Comonad g) => Comonad (Day f g) where
  extract :: forall a. Day f g a -> a
extract (Day x -> y -> a
f f x
x g y
y) = x -> y -> a
f (forall (w :: * -> *) a. Comonad w => w a -> a
extract f x
x) (forall (w :: * -> *) a. Comonad w => w a -> a
extract g y
y)
  {-# INLINE extract #-}
  duplicate :: forall a. Day f g a -> Day f g (Day f g a)
duplicate (Day x -> y -> a
f f x
x g y
y) =
    let xx :: f (f x)
xx = forall (w :: * -> *) a. Comonad w => w a -> w (w a)
duplicate f x
x
        yy :: g (g y)
yy = forall (w :: * -> *) a. Comonad w => w a -> w (w a)
duplicate g y
y
    in forall (f :: * -> *) (g :: * -> *) a x y.
(x -> y -> a) -> f x -> g y -> Day f g a
Day (forall (f :: * -> *) (g :: * -> *) a x y.
(x -> y -> a) -> f x -> g y -> Day f g a
Day x -> y -> a
f) f (f x)
xx g (g y)
yy
  {-# INLINE duplicate #-}

instance (ComonadApply f, ComonadApply g) => ComonadApply (Day f g) where
  Day x -> y -> a -> b
u f x
fa g y
fb <@> :: forall a b. Day f g (a -> b) -> Day f g a -> Day f g b
<@> Day x -> y -> a
v f x
gc g y
gd =
    forall (f :: * -> *) (g :: * -> *) a x y.
(x -> y -> a) -> f x -> g y -> Day f g a
Day
      (\(x
a,x
c) (y
b, y
d) -> x -> y -> a -> b
u x
a y
b (x -> y -> a
v x
c y
d))
      ((,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f x
fa forall (w :: * -> *) a b.
ComonadApply w =>
w (a -> b) -> w a -> w b
<@> f x
gc)
      ((,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> g y
fb forall (w :: * -> *) a b.
ComonadApply w =>
w (a -> b) -> w a -> w b
<@> g y
gd)

instance (Representable f, Representable g) => Distributive (Day f g) where
  distribute :: forall (f :: * -> *) a. Functor f => f (Day f g a) -> Day f g (f a)
distribute f (Day f g a)
f = forall (f :: * -> *) (g :: * -> *) a x y.
(x -> y -> a) -> f x -> g y -> Day f g a
Day Rep f -> Rep g -> f a
fn (forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate forall a. a -> a
id) (forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate forall a. a -> a
id) where
    fn :: Rep f -> Rep g -> f a
fn Rep f
x Rep g
y = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Day x -> y -> a
o f x
m g y
n) -> x -> y -> a
o (forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index f x
m Rep f
x) (forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index g y
n Rep g
y)) f (Day f g a)
f
  {-# INLINE distribute #-}

  collect :: forall (f :: * -> *) a b.
Functor f =>
(a -> Day f g b) -> f a -> Day f g (f b)
collect a -> Day f g b
g f a
f = forall (f :: * -> *) (g :: * -> *) a x y.
(x -> y -> a) -> f x -> g y -> Day f g a
Day Rep f -> Rep g -> f b
fn (forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate forall a. a -> a
id) (forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate forall a. a -> a
id) where
    fn :: Rep f -> Rep g -> f b
fn Rep f
x Rep g
y = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
q -> case a -> Day f g b
g a
q of Day x -> y -> b
o f x
m g y
n  -> x -> y -> b
o (forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index f x
m Rep f
x) (forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index g y
n Rep g
y)) f a
f
  {-# INLINE collect #-}

instance (Representable f, Representable g) => Representable (Day f g) where
  type Rep (Day f g) = (Rep f, Rep g)
  tabulate :: forall a. (Rep (Day f g) -> a) -> Day f g a
tabulate Rep (Day f g) -> a
f = forall (f :: * -> *) (g :: * -> *) a x y.
(x -> y -> a) -> f x -> g y -> Day f g a
Day (forall a b c. ((a, b) -> c) -> a -> b -> c
curry Rep (Day f g) -> a
f) (forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate forall a. a -> a
id) (forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate forall a. a -> a
id)
  {-# INLINE tabulate #-}
  index :: forall a. Day f g a -> Rep (Day f g) -> a
index (Day x -> y -> a
o f x
m g y
n ) (Rep f
x,Rep g
y) = x -> y -> a
o (forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index f x
m Rep f
x) (forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index g y
n Rep g
y)
  {-# INLINE index #-}