{-# LANGUAGE AllowAmbiguousTypes #-}

module Proarrow.Category.Instance.Free where

import Data.Kind (Constraint, Type)
import Prelude qualified as P

import Proarrow.Core (CAT, Category, CategoryOf (..), Is, Kind, Profunctor (..), Promonad (..), UN, dimapDefault, (:~>))

type MULTICAT k = [k] -> k -> Type

newtype FREE (str :: CAT [FK k]) (p :: CAT k) = F (FK k)
type instance UN F (F k) = k

type Free :: CAT (FREE str p)
data Free a b where
  Free
    :: (IsFK a, IsFK b)
    => (forall ctx. (IsFKs ctx) => MFree str p ctx a -> MFree str p ctx b)
    -> Free (F a :: FREE str p) (F b)

data FK k = EMB k
type instance UN EMB (EMB a) = a

type MFree :: CAT [FK k] -> CAT k -> MULTICAT (FK k)
data MFree str p as b where
  Id :: (IsFK a) => MFree str p '[a] a
  Emb :: p a b -> MFree str p ctx (EMB a) -> MFree str p ctx (EMB b)
  Str :: str ctx1 ctx2 -> MFree str p ctx2 a -> MFree str p ctx1 a

type IsFK :: forall {k}. FK k -> Constraint
class IsFK (a :: FK k) where
  type Lower (a :: FK k) :: k
  withLowerOb :: ((Ob (Lower a)) => r) -> r
instance (Ob a) => IsFK (EMB a) where
  type Lower (EMB a) = a
  withLowerOb :: forall r. (Ob (Lower ('EMB a)) => r) -> r
withLowerOb Ob (Lower ('EMB a)) => r
r = r
Ob (Lower ('EMB a)) => r
r

type IsFKs :: forall {k}. [FK k] -> Constraint
class IsFKs (as :: [FK k]) where
  type Flatten (as :: [FK k]) :: FK k
  withFlattenOb :: ((IsFK (Flatten as)) => r) -> r

instance (IsFK a) => IsFKs '[a] where
  type Flatten '[a] = a
  withFlattenOb :: forall r. (IsFK (Flatten '[a]) => r) -> r
withFlattenOb IsFK (Flatten '[a]) => r
r = r
IsFK (Flatten '[a]) => r
r

class Structure (str :: CAT [FK k]) where
  type HasStructure str :: Kind -> Constraint
  foldStructure
    :: forall (cat :: CAT k) (ctx1 :: [FK k]) (ctx2 :: [FK k])
     . (HasStructure str k)
     => str ctx1 ctx2
     -> cat (Lower (Flatten ctx1)) (Lower (Flatten ctx2))

data Plain ctx1 ctx2
class NoConstraint k
instance NoConstraint k
instance Structure Plain where
  type HasStructure Plain = NoConstraint
  foldStructure :: forall (cat :: CAT k) (ctx1 :: [FK k]) (ctx2 :: [FK k]).
HasStructure Plain k =>
Plain ctx1 ctx2
-> cat (Lower (Flatten ctx1)) (Lower (Flatten ctx2))
foldStructure = \case {}

fold
  :: forall {k} (str :: CAT [FK k]) (p :: CAT k) (cat :: CAT k) (a :: FK k) (b :: FK k)
   . (Category cat, Structure str, HasStructure str k)
  => (p :~> cat)
  -> Free (F a :: FREE str p) (F b)
  -> cat (Lower a) (Lower b)
fold :: forall {k} (str :: CAT [FK k]) (p :: CAT k) (cat :: CAT k)
       (a :: FK k) (b :: FK k).
(Category cat, Structure str, HasStructure str k) =>
(p :~> cat) -> Free ('F a) ('F b) -> cat (Lower a) (Lower b)
fold p :~> cat
pn (Free forall (ctx :: [FK k]).
IsFKs ctx =>
MFree str p ctx a -> MFree str p ctx b
f) = MFree str p '[a] b -> cat (Lower (Flatten '[a])) (Lower b)
forall (as :: [FK k]) (c :: FK k).
MFree str p as c -> cat (Lower (Flatten as)) (Lower c)
go (MFree str p '[a] a -> MFree str p '[a] b
forall (ctx :: [FK k]).
IsFKs ctx =>
MFree str p ctx a -> MFree str p ctx b
f MFree str p '[a] a
forall {k} (a :: FK k) (str :: CAT [FK k]) (p :: CAT k).
IsFK a =>
MFree str p '[a] a
Id)
  where
    go :: MFree str p as c -> cat (Lower (Flatten as)) (Lower c)
    go :: forall (as :: [FK k]) (c :: FK k).
MFree str p as c -> cat (Lower (Flatten as)) (Lower c)
go (Id @x) = forall {k} (a :: FK k) r. IsFK a => (Ob (Lower a) => r) -> r
forall (a :: FK k) r. IsFK a => (Ob (Lower a) => r) -> r
withLowerOb @x cat (Lower c) (Lower c)
Ob (Lower c) => cat (Lower c) (Lower c)
forall (a :: k). Ob a => cat a a
forall {k} (p :: k +-> k) (a :: k). (Promonad p, Ob a) => p a a
id
    go (Emb p a b
p MFree str p as ('EMB a)
g) = p a b -> cat a b
p :~> cat
pn p a b
p cat a b -> cat (Lower (Flatten as)) a -> cat (Lower (Flatten as)) b
forall (b :: k) (c :: k) (a :: k). cat b c -> cat a b -> cat a c
forall {k} (p :: k +-> k) (b :: k) (c :: k) (a :: k).
Promonad p =>
p b c -> p a b -> p a c
. MFree str p as ('EMB a)
-> cat (Lower (Flatten as)) (Lower ('EMB a))
forall (as :: [FK k]) (c :: FK k).
MFree str p as c -> cat (Lower (Flatten as)) (Lower c)
go MFree str p as ('EMB a)
g
    go (Str str as ctx2
s MFree str p ctx2 c
g) = MFree str p ctx2 c -> cat (Lower (Flatten ctx2)) (Lower c)
forall (as :: [FK k]) (c :: FK k).
MFree str p as c -> cat (Lower (Flatten as)) (Lower c)
go MFree str p ctx2 c
g cat (Lower (Flatten ctx2)) (Lower c)
-> cat (Lower (Flatten as)) (Lower (Flatten ctx2))
-> cat (Lower (Flatten as)) (Lower c)
forall (b :: k) (c :: k) (a :: k). cat b c -> cat a b -> cat a c
forall {k} (p :: k +-> k) (b :: k) (c :: k) (a :: k).
Promonad p =>
p b c -> p a b -> p a c
. str as ctx2 -> cat (Lower (Flatten as)) (Lower (Flatten ctx2))
forall k (str :: CAT [FK k]) (cat :: CAT k) (ctx1 :: [FK k])
       (ctx2 :: [FK k]).
(Structure str, HasStructure str k) =>
str ctx1 ctx2 -> cat (Lower (Flatten ctx1)) (Lower (Flatten ctx2))
forall (cat :: CAT k) (ctx1 :: [FK k]) (ctx2 :: [FK k]).
HasStructure str k =>
str ctx1 ctx2 -> cat (Lower (Flatten ctx1)) (Lower (Flatten ctx2))
foldStructure str as ctx2
s

instance (Structure str, Profunctor p) => CategoryOf (FREE str p) where
  type (~>) = Free
  type Ob a = (Is F a, IsFK (UN F a))

instance (Structure str, Profunctor p) => Promonad (Free :: CAT (FREE str p)) where
  id :: forall (a :: FREE str p). Ob a => Free a a
id = (forall (ctx :: [FK k]).
 IsFKs ctx =>
 MFree str p ctx (UN 'F a) -> MFree str p ctx (UN 'F a))
-> Free ('F (UN 'F a)) ('F (UN 'F a))
forall {k} (a :: FK k) (b :: FK k) (str :: CAT [FK k])
       (p :: CAT k).
(IsFK a, IsFK b) =>
(forall (ctx :: [FK k]).
 IsFKs ctx =>
 MFree str p ctx a -> MFree str p ctx b)
-> Free ('F a) ('F b)
Free MFree str p ctx (UN 'F a) -> MFree str p ctx (UN 'F a)
forall (ctx :: [FK k]).
IsFKs ctx =>
MFree str p ctx (UN 'F a) -> MFree str p ctx (UN 'F a)
forall a. Ob a => a -> a
forall {k} (p :: k +-> k) (a :: k). (Promonad p, Ob a) => p a a
id
  Free forall (ctx :: [FK k]).
IsFKs ctx =>
MFree str p ctx a -> MFree str p ctx b
f . :: forall (b :: FREE str p) (c :: FREE str p) (a :: FREE str p).
Free b c -> Free a b -> Free a c
. Free forall (ctx :: [FK k]).
IsFKs ctx =>
MFree str p ctx a -> MFree str p ctx b
g = (forall (ctx :: [FK k]).
 IsFKs ctx =>
 MFree str p ctx a -> MFree str p ctx b)
-> Free ('F a) ('F b)
forall {k} (a :: FK k) (b :: FK k) (str :: CAT [FK k])
       (p :: CAT k).
(IsFK a, IsFK b) =>
(forall (ctx :: [FK k]).
 IsFKs ctx =>
 MFree str p ctx a -> MFree str p ctx b)
-> Free ('F a) ('F b)
Free (MFree str p ctx a -> MFree str p ctx b
forall (ctx :: [FK k]).
IsFKs ctx =>
MFree str p ctx a -> MFree str p ctx b
f (MFree str p ctx a -> MFree str p ctx b)
-> (MFree str p ctx a -> MFree str p ctx a)
-> MFree str p ctx a
-> MFree str p ctx b
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (p :: k +-> k) (b :: k) (c :: k) (a :: k).
Promonad p =>
p b c -> p a b -> p a c
. MFree str p ctx a -> MFree str p ctx a
MFree str p ctx a -> MFree str p ctx b
forall (ctx :: [FK k]).
IsFKs ctx =>
MFree str p ctx a -> MFree str p ctx b
g)

instance (Structure str, Profunctor p) => Profunctor (Free :: CAT (FREE str p)) where
  dimap :: forall (c :: FREE str p) (a :: FREE str p) (b :: FREE str p)
       (d :: FREE str p).
(c ~> a) -> (b ~> d) -> Free a b -> Free c d
dimap = (c ~> a) -> (b ~> d) -> Free a b -> Free c d
Free c a -> Free b d -> Free a b -> Free c d
forall {k} (p :: k +-> k) (c :: k) (a :: k) (b :: k) (d :: k).
Promonad p =>
p c a -> p b d -> p a b -> p c d
dimapDefault
  (Ob a, Ob b) => r
r \\ :: forall (a :: FREE str p) (b :: FREE str p) r.
((Ob a, Ob b) => r) -> Free a b -> r
\\ Free{} = r
(Ob a, Ob b) => r
r

data Test a b where
  Show :: Test P.Int P.String
  Read :: Test P.String P.Int
  Succ :: Test P.Int P.Int
  Dup :: Test P.String P.String

type T ctx = MFree Plain Test ctx
show :: T ctx (EMB P.Int) -> T ctx (EMB P.String)
show :: forall (ctx :: [FK Type]). T ctx ('EMB Int) -> T ctx ('EMB String)
show = Test Int String
-> MFree Plain Test ctx ('EMB Int)
-> MFree Plain Test ctx ('EMB String)
forall {k} (p :: CAT k) (a :: k) (b :: k) (str :: CAT [FK k])
       (ctx :: [FK k]).
p a b -> MFree str p ctx ('EMB a) -> MFree str p ctx ('EMB b)
Emb Test Int String
Show

read :: T ctx (EMB P.String) -> T ctx (EMB P.Int)
read :: forall (ctx :: [FK Type]). T ctx ('EMB String) -> T ctx ('EMB Int)
read = Test String Int
-> MFree Plain Test ctx ('EMB String)
-> MFree Plain Test ctx ('EMB Int)
forall {k} (p :: CAT k) (a :: k) (b :: k) (str :: CAT [FK k])
       (ctx :: [FK k]).
p a b -> MFree str p ctx ('EMB a) -> MFree str p ctx ('EMB b)
Emb Test String Int
Read

succ :: T ctx (EMB P.Int) -> T ctx (EMB P.Int)
succ :: forall (ctx :: [FK Type]). T ctx ('EMB Int) -> T ctx ('EMB Int)
succ = Test Int Int
-> MFree Plain Test ctx ('EMB Int)
-> MFree Plain Test ctx ('EMB Int)
forall {k} (p :: CAT k) (a :: k) (b :: k) (str :: CAT [FK k])
       (ctx :: [FK k]).
p a b -> MFree str p ctx ('EMB a) -> MFree str p ctx ('EMB b)
Emb Test Int Int
Succ

dup :: T ctx (EMB P.String) -> T ctx (EMB P.String)
dup :: forall (ctx :: [FK Type]).
T ctx ('EMB String) -> T ctx ('EMB String)
dup = Test String String
-> MFree Plain Test ctx ('EMB String)
-> MFree Plain Test ctx ('EMB String)
forall {k} (p :: CAT k) (a :: k) (b :: k) (str :: CAT [FK k])
       (ctx :: [FK k]).
p a b -> MFree str p ctx ('EMB a) -> MFree str p ctx ('EMB b)
Emb Test String String
Dup

test :: T ctx (EMB P.String) -> T ctx (EMB P.String)
test :: forall (ctx :: [FK Type]).
T ctx ('EMB String) -> T ctx ('EMB String)
test T ctx ('EMB String)
x = T ctx ('EMB String) -> T ctx ('EMB String)
forall (ctx :: [FK Type]).
T ctx ('EMB String) -> T ctx ('EMB String)
dup (T ctx ('EMB Int) -> T ctx ('EMB String)
forall (ctx :: [FK Type]). T ctx ('EMB Int) -> T ctx ('EMB String)
show (T ctx ('EMB Int) -> T ctx ('EMB Int)
forall (ctx :: [FK Type]). T ctx ('EMB Int) -> T ctx ('EMB Int)
succ (T ctx ('EMB String) -> T ctx ('EMB Int)
forall (ctx :: [FK Type]). T ctx ('EMB String) -> T ctx ('EMB Int)
read T ctx ('EMB String)
x)))

testF :: Free (F (EMB P.String) :: FREE Plain Test) (F (EMB P.String))
testF :: Free ('F ('EMB String)) ('F ('EMB String))
testF = (forall (ctx :: [FK Type]).
 IsFKs ctx =>
 MFree Plain Test ctx ('EMB String)
 -> MFree Plain Test ctx ('EMB String))
-> Free ('F ('EMB String)) ('F ('EMB String))
forall {k} (a :: FK k) (b :: FK k) (str :: CAT [FK k])
       (p :: CAT k).
(IsFK a, IsFK b) =>
(forall (ctx :: [FK k]).
 IsFKs ctx =>
 MFree str p ctx a -> MFree str p ctx b)
-> Free ('F a) ('F b)
Free T ctx ('EMB String) -> T ctx ('EMB String)
forall (ctx :: [FK Type]).
IsFKs ctx =>
MFree Plain Test ctx ('EMB String)
-> MFree Plain Test ctx ('EMB String)
forall (ctx :: [FK Type]).
T ctx ('EMB String) -> T ctx ('EMB String)
test

-- >>> testFold "123"
-- "124124"
testFold :: P.String -> P.String
testFold :: String -> String
testFold = (Test :~> (->))
-> Free ('F ('EMB String)) ('F ('EMB String))
-> Lower ('EMB String)
-> Lower ('EMB String)
forall {k} (str :: CAT [FK k]) (p :: CAT k) (cat :: CAT k)
       (a :: FK k) (b :: FK k).
(Category cat, Structure str, HasStructure str k) =>
(p :~> cat) -> Free ('F a) ('F b) -> cat (Lower a) (Lower b)
fold Test a b -> a -> b
Test :~> (->)
interp Free ('F ('EMB String)) ('F ('EMB String))
testF
  where
    interp :: Test :~> (->)
    interp :: Test :~> (->)
interp Test a b
Show = a -> b
a -> String
forall a. Show a => a -> String
P.show
    interp Test a b
Read = a -> b
String -> b
forall a. Read a => String -> a
P.read
    interp Test a b
Succ = a -> a
a -> b
forall a. Enum a => a -> a
P.succ
    interp Test a b
Dup = \a
s -> a
String
s String -> String -> String
forall a. [a] -> [a] -> [a]
P.++ a
String
s