{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-# OPTIONS_GHC -Wall #-}
-- {-# OPTIONS_GHC -Wno-unused-imports #-} -- TEMP

#include "ConCat/AbsTy.inc"
AbsTyPragmas

-- | A category of local approximations (and probably other uses)
module ConCat.Free.Affine where

import Prelude hiding (id,(.),curry,uncurry,const)
import Data.Key (Zip(..))

import ConCat.Misc ((:*))
import ConCat.Rep
import qualified ConCat.Category as C
import ConCat.AltCat
import ConCat.Free.VectorSpace hiding ((^+^))
import qualified ConCat.Free.VectorSpace as V
import ConCat.Free.LinearRow
import ConCat.AdditiveFun (Additive(..))

AbsTyImports

data Affine s a b = Affine (L s a b) b

linearA :: forall s a b. Ok2 (L s) a b => L s a b -> Affine s a b
linearA :: forall s a b. Ok2 (L s) a b => L s a b -> Affine s a b
linearA = (L s a b -> b -> Affine s a b) -> b -> L s a b -> Affine s a b
forall a b c. (a -> b -> c) -> b -> a -> c
flip L s a b -> b -> Affine s a b
forall s a b. L s a b -> b -> Affine s a b
Affine (forall s a. HasV s a => V s a s -> a
unV @s V s b s
forall (f :: * -> *) a. (Pointed f, Num a) => f a
zeroV)
{-# INLINE linearA #-}

applyA :: forall s a b. Ok2 (L s) a b => Affine s a b -> (a -> b)
applyA :: forall s a b. Ok2 (L s) a b => Affine s a b -> a -> b
applyA (Affine L s a b
p b
u) a
a = forall s a. (HasV s a, Zip (V s a), Num s) => a -> a -> a
add @s (L s a b -> a -> b
forall s a b. (Num s, Ok2 (L s) a b) => L s a b -> a -> b
lapply L s a b
p a
a) b
u
{-# INLINE applyA #-}

instance HasRep (Affine s a b) where
  type Rep (Affine s a b) = L s a b :* b
  abst :: Rep (Affine s a b) -> Affine s a b
abst (L s a b
m,b
b) = L s a b -> b -> Affine s a b
forall s a b. L s a b -> b -> Affine s a b
Affine L s a b
m b
b
  repr :: Affine s a b -> Rep (Affine s a b)
repr (Affine L s a b
m b
b) = (L s a b
m,b
b)

AbsTy(Affine s a b)

instance Ok2 (L s) a b => Additive (Affine s a b) where
  zero :: Affine s a b
zero = L s a b -> Affine s a b
forall s a b. Ok2 (L s) a b => L s a b -> Affine s a b
linearA L s a b
forall s a b.
(Num s, Zeroable (V s a), Zeroable (V s b)) =>
L s a b
zeroLM
  Affine L s a b
p b
u ^+^ :: Affine s a b -> Affine s a b -> Affine s a b
^+^ Affine L s a b
q b
v = L s a b -> b -> Affine s a b
forall s a b. L s a b -> b -> Affine s a b
Affine (L s a b
p Binop (L s a b)
forall s a b. Ok2 (L s) a b => Binop (L s a b)
`addLM` L s a b
q) (forall s a. (HasV s a, Zip (V s a), Num s) => a -> a -> a
add @s b
u b
v)
  {-# INLINE zero #-}
  {-# INLINE (^+^) #-}

instance Category (Affine s) where
  type Ok (Affine s) = Ok (L s)
  id :: forall a. Ok (Affine s) a => Affine s a a
id = L s a a -> Affine s a a
forall s a b. Ok2 (L s) a b => L s a b -> Affine s a b
linearA L s a a
forall (k :: * -> * -> *) a. (Category k, Ok k a) => k a a
id
  . :: forall b c a.
Ok3 (Affine s) a b c =>
Affine s b c -> Affine s a b -> Affine s a c
(.) = (Rep (Affine s b c) -> Rep (Affine s a b) -> Rep (Affine s a c))
-> Affine s b c -> Affine s a b -> Affine s a c
forall p q r.
(HasRep p, HasRep q, HasRep r) =>
(Rep p -> Rep q -> Rep r) -> p -> q -> r
inAbst2 ((Rep (Affine s b c) -> Rep (Affine s a b) -> Rep (Affine s a c))
 -> Affine s b c -> Affine s a b -> Affine s a c)
-> (Rep (Affine s b c) -> Rep (Affine s a b) -> Rep (Affine s a c))
-> Affine s b c
-> Affine s a b
-> Affine s a c
forall a b. (a -> b) -> a -> b
$ \ (L s b c
q,c
v) (L s a b
p,b
u) -> (L s b c
q L s b c -> L s a b -> L s a c
forall (k :: * -> * -> *) b c a.
(Category k, Ok3 k a b c) =>
k b c -> k a b -> k a c
. L s a b
p, forall s a. (HasV s a, Zip (V s a), Num s) => a -> a -> a
add @s (L s b c -> b -> c
forall s a b. (Num s, Ok2 (L s) a b) => L s a b -> a -> b
lapply L s b c
q b
u) c
v)
  {-# INLINE id #-}
  {-# INLINE (.) #-}

-- Semantic homomorphism: applyA g . applyA f == applyA (g . f)

-- applyA (Affine q v) . applyA (Affine p u) == applyA (Affine q v . Affine p u)

--   applyA (Affine q v) . applyA (Affine p u)
-- == \ a -> q (p a + u) + v
-- == \ a -> (q (p a) + q u) + v
-- == \ a -> (q . p) a + (q u + v)
-- == applyA (Affine (q . p) (q u + v))

instance MonoidalPCat (Affine s) where
  *** :: forall a b c d.
Ok4 (Affine s) a b c d =>
Affine s a c
-> Affine s b d
-> Affine s (Prod (Affine s) a b) (Prod (Affine s) c d)
(***) = (Rep (Affine s a c)
 -> Rep (Affine s b d) -> Rep (Affine s (a :* b) (c :* d)))
-> Affine s a c -> Affine s b d -> Affine s (a :* b) (c :* d)
forall p q r.
(HasRep p, HasRep q, HasRep r) =>
(Rep p -> Rep q -> Rep r) -> p -> q -> r
inAbst2 ((Rep (Affine s a c)
  -> Rep (Affine s b d) -> Rep (Affine s (a :* b) (c :* d)))
 -> Affine s a c -> Affine s b d -> Affine s (a :* b) (c :* d))
-> (Rep (Affine s a c)
    -> Rep (Affine s b d) -> Rep (Affine s (a :* b) (c :* d)))
-> Affine s a c
-> Affine s b d
-> Affine s (a :* b) (c :* d)
forall a b. (a -> b) -> a -> b
$ \ (L s a c
p,c
u) (L s b d
q,d
v) -> (L s a c
p L s a c -> L s b d -> L s (a :* b) (c :* d)
forall (k :: * -> * -> *) a b c d.
(MonoidalPCat k, Ok4 k a b c d) =>
k a c -> k b d -> k (Prod k a b) (Prod k c d)
*** L s b d
q, (c
u,d
v))
  {-# INLINE (***) #-}

--    applyA (Affine p u) *** applyA (Affine q v)
-- == \ (a,b) -> (applyA (Affine p u) *** applyA (Affine q v)) a
-- == \ (a,b) -> (applyA (Affine p u) a, applyA (Affine q v) b)
-- == \ (a,b) -> (p a + u, q b + v)
-- == \ (a,b) -> (p a,q b) + (u,v)
-- == \ (a,b) -> (p *** q) (a,b) + (u,v)
-- == applyA (Affine (p *** q) (u,v))

instance BraidedPCat (Affine s) where
  swapP :: forall a b.
Ok2 (Affine s) a b =>
Affine s (Prod (Affine s) a b) (Prod (Affine s) b a)
swapP = L s (Prod (Affine s) a b) (Prod (Affine s) b a)
-> Affine s (Prod (Affine s) a b) (Prod (Affine s) b a)
forall s a b. Ok2 (L s) a b => L s a b -> Affine s a b
linearA L s (Prod (Affine s) a b) (Prod (Affine s) b a)
forall (k :: * -> * -> *) a b.
(BraidedPCat k, Ok2 k a b) =>
k (Prod k a b) (Prod k b a)
swapP
  {-# INLINE swapP #-}

instance ProductCat (Affine s) where
  exl :: forall a b. Ok2 (Affine s) a b => Affine s (Prod (Affine s) a b) a
exl = L s (Prod (Affine s) a b) a -> Affine s (Prod (Affine s) a b) a
forall s a b. Ok2 (L s) a b => L s a b -> Affine s a b
linearA L s (Prod (Affine s) a b) a
forall (k :: * -> * -> *) a b.
(ProductCat k, Ok2 k a b) =>
k (Prod k a b) a
exl
  exr :: forall a b. Ok2 (Affine s) a b => Affine s (Prod (Affine s) a b) b
exr = L s (Prod (Affine s) a b) b -> Affine s (Prod (Affine s) a b) b
forall s a b. Ok2 (L s) a b => L s a b -> Affine s a b
linearA L s (Prod (Affine s) a b) b
forall (k :: * -> * -> *) a b.
(ProductCat k, Ok2 k a b) =>
k (Prod k a b) b
exr
  dup :: forall a. Ok (Affine s) a => Affine s a (Prod (Affine s) a a)
dup = L s a (Prod (Affine s) a a) -> Affine s a (Prod (Affine s) a a)
forall s a b. Ok2 (L s) a b => L s a b -> Affine s a b
linearA L s a (Prod (Affine s) a a)
forall (k :: * -> * -> *) a.
(ProductCat k, Ok k a) =>
k a (Prod k a a)
dup
  {-# INLINE exl #-}
  {-# INLINE exr #-}
  {-# INLINE dup #-}

-- instance Num s => UnitCat (Affine s)

{--------------------------------------------------------------------
    Move elsewhere
--------------------------------------------------------------------}

add :: forall s a. (HasV s a, Zip (V s a), Num s) => a -> a -> a
add :: forall s a. (HasV s a, Zip (V s a), Num s) => a -> a -> a
add = forall s a b c.
(HasV s a, HasV s b, HasV s c) =>
(V s a s -> V s b s -> V s c s) -> a -> b -> c
onV2 @s V s a s -> V s a s -> V s a s
forall (f :: * -> *) s. (Zip f, Num s) => f s -> f s -> f s
(V.^+^)
{-# INLINE add #-}