{-# LANGUAGE CPP                        #-}
{-# LANGUAGE Safe                       #-}
{-# LANGUAGE PolyKinds                  #-}
{-# LANGUAGE ConstraintKinds            #-}
{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE NoImplicitPrelude          #-}
{-# LANGUAGE RebindableSyntax           #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE RankNTypes               #-}

module Data.Algebra (
  -- * Algebras 
    type FreeAlgebra
  , Algebra(..)
  , (.*.)
  , type FreeUnital
  , Unital(..)
  , unital
  , unit
  -- * Coalgebras 
  , type FreeCoalgebra
  , Coalgebra(..)
  , type FreeCounital
  , Counital(..)
  , counital
  -- * Bialgebras 
  , type FreeBialgebra
  , Bialgebra
) where

import safe Data.Bool
import safe Data.Functor.Rep
import safe Data.Semimodule
import safe Data.Semiring
import safe Prelude (Ord, reverse)
import safe qualified Data.IntSet as IntSet
import safe qualified Data.Set as Set
import safe qualified Data.Sequence as Seq
import safe Data.Sequence hiding (reverse,index)
import safe Prelude hiding (Num(..), Fractional(..), negate, sum, product)

-------------------------------------------------------------------------------
-- Algebras
-------------------------------------------------------------------------------

-- | An algebra over a free module /f/.
--
-- Note that this is distinct from a < https://en.wikipedia.org/wiki/Free_algebra free algebra >.
--
type FreeAlgebra a f = (FreeSemimodule a f, Algebra a (Rep f))

-- | An algebra < https://en.wikipedia.org/wiki/Algebra_over_a_field#Generalization:_algebra_over_a_ring algebra > over a semiring.
--
-- Note that the algebra < https://en.wikipedia.org/wiki/Non-associative_algebra needn't be associative >.
--
class Semiring a => Algebra a b where
  append :: (b -> b -> a) -> b -> a

infixl 7 .*.

-- | Multiplication operator on an algebra over a free semimodule.
--
-- /Caution/ in general (.*.) needn't be commutative, nor associative.
--
(.*.) :: FreeAlgebra a f => f a -> f a -> f a
(.*.) x y = tabulate $ append (\i j -> index x i * index y j)

-- | A unital algebra over a free semimodule /f/.
--
type FreeUnital a f = (FreeAlgebra a f, Unital a (Rep f))

-- | A < https://en.wikipedia.org/wiki/Algebra_over_a_field#Unital_algebra unital algebra > over a semiring.
--
class Algebra a b => Unital a b where
  aempty :: a -> b -> a

-- | Insert an element into an algebra.
--
-- >>> V4 1 2 3 4 .*. unital two :: V4 Int
-- V4 2 4 6 8
unital :: FreeUnital a f => a -> f a
unital = tabulate . aempty

-- | Unital element of a unital algebra over a free semimodule.
--
-- >>> unit :: Complex Int
-- 1 :+ 0
-- >>> unit :: QuatD
-- Quaternion 1.0 (V3 0.0 0.0 0.0)
--
unit :: FreeUnital a f => f a
unit = unital one

-------------------------------------------------------------------------------
-- Coalgebras
-------------------------------------------------------------------------------

-- | A coalgebra over a free semimodule /f/.
--
type FreeCoalgebra a f = (FreeSemimodule a f, Coalgebra a (Rep f))

-- | A coalgebra over a semiring.
--
-- ( id *** coempty ) . coappend = id = ( coempty *** id ) . coappend
class Semiring a => Coalgebra a c where
  coappend :: (c -> a) -> c -> c -> a

-- | A counital coalgebra over a free semimodule /f/.
--
type FreeCounital a f = (FreeCoalgebra a f, Counital a (Rep f))

-- | A counital coalgebra over a semiring.
--
class Coalgebra a c => Counital a c where
  coempty :: (c -> a) -> a

-- | Obtain an element from a coalgebra over a free semimodule.
--
counital :: FreeCounital a f => f a -> a
counital = coempty . index

-------------------------------------------------------------------------------
-- Bialgebras
-------------------------------------------------------------------------------

-- | A bialgebra over a free semimodule /f/.
--
type FreeBialgebra a f = (FreeAlgebra a f, FreeCoalgebra a f, Bialgebra a (Rep f))

-- | A < https://en.wikipedia.org/wiki/Bialgebra bialgebra > over a semiring.
--
class (Unital a b, Counital a b) => Bialgebra a b

-------------------------------------------------------------------------------
-- Instances
-------------------------------------------------------------------------------


--instance (Semiring a, Algebra a b) => Algebra a (a -> r) where
--  aempty = aempty one

--instance (Semiring a, Division a b) => Division r (a -> r) where
--  reciprocalWith = reciprocalWith

-- incoherent
-- instance Algebra () a where aempty _ _ = ()
-- instance (Algebra a b, Algebra a c) => Algebra (a -> r) b where aempty f b a = aempty (f a) b
--instance (Algebra r a, Algebra r b) => Algebra (a -> r) b where aempty f b a = aempty (f a) b

--instance (Algebra r b, Algebra r a) => Algebra (b -> r) a where append f a b = append (\a1 a2 -> f a1 a2 b) a


instance Semiring a => Algebra a () where
  append f = f ()

instance Semiring a => Unital a () where
  aempty r () = r

instance (Algebra a b, Algebra a c) => Algebra a (b, c) where
  append f (a,b) = append (\a1 a2 -> append (\b1 b2 -> f (a1,b1) (a2,b2)) b) a

instance (Unital a b, Unital a c) => Unital a (b, c) where
  aempty r (a,b) = aempty r a * aempty r b

instance (Algebra a b, Algebra a c, Algebra a d) => Algebra a (b, c, d) where
  append f (a,b,c) = append (\a1 a2 -> append (\b1 b2 -> append (\c1 c2 -> f (a1,b1,c1) (a2,b2,c2)) c) b) a

instance (Unital a b, Unital a c, Unital a d) => Unital a (b, c, d) where
  aempty r (a,b,c) = aempty r a * aempty r b * aempty r c

-- | Tensor algebra
--
-- >>> append (<>) [1..3 :: Int]
-- [1,2,3,1,2,3,1,2,3,1,2,3]
--
-- >>> append (\f g -> fold (f ++ g)) [1..3] :: Int
-- 24
--
instance Semiring a => Algebra a [a] where
  append f = go [] where
    go ls rrs@(r:rs) = f (reverse ls) rrs + go (r:ls) rs
    go ls [] = f (reverse ls) []

instance Semiring a => Unital a [a] where
  aempty a [] = a
  aempty _ _ = zero


-- | The tensor algebra
instance Semiring r => Algebra r (Seq a) where
  append f = go Seq.empty where
    go ls s = case viewl s of
       EmptyL -> f ls s
       r :< rs -> f ls s + go (ls |> r) rs

instance Semiring r => Unital r (Seq a) where
  aempty r a | Seq.null a = r
             | otherwise = zero

instance (Semiring r, Ord a) => Algebra r (Set.Set a) where
  append f = go Set.empty where
    go ls s = case Set.minView s of
       Nothing -> f ls s
       Just (r, rs) -> f ls s + go (Set.insert r ls) rs

instance (Semiring r, Ord a) => Unital r (Set.Set a) where
  aempty r a | Set.null a = r
           | otherwise = zero

instance Semiring r => Algebra r IntSet.IntSet where
  append f = go IntSet.empty where
    go ls s = case IntSet.minView s of
       Nothing -> f ls s
       Just (r, rs) -> f ls s + go (IntSet.insert r ls) rs

instance Semiring r => Unital r IntSet.IntSet where
  aempty r a | IntSet.null a = r
             | otherwise = zero

---------------------------------------------------------------------
-- Coalgebra instances
---------------------------------------------------------------------


instance Semiring r => Coalgebra r () where
  coappend = const

instance Semiring r => Counital r () where
  coempty f = f ()

instance (Coalgebra r a, Coalgebra r b) => Coalgebra r (a, b) where
  coappend f (a1,b1) (a2,b2) = coappend (\a -> coappend (\b -> f (a,b)) b1 b2) a1 a2

instance (Counital r a, Counital r b) => Counital r (a, b) where
  coempty k = coempty $ \a -> coempty $ \b -> k (a,b)

instance (Coalgebra r a, Coalgebra r b, Coalgebra r c) => Coalgebra r (a, b, c) where
  coappend f (a1,b1,c1) (a2,b2,c2) = coappend (\a -> coappend (\b -> coappend (\c -> f (a,b,c)) c1 c2) b1 b2) a1 a2

instance (Counital r a, Counital r b, Counital r c) => Counital r (a, b, c) where
  coempty k = coempty $ \a -> coempty $ \b -> coempty $ \c -> k (a,b,c)

instance (Algebra r a) => Coalgebra r (a -> r) where
  coappend k f g = k (f * g)

instance (Algebra r a) => Counital r (a -> r) where
  coempty k = k one

{-
instance (Semiring r, FreeAlgebra r f) => Coalgebra r (f r) where
  coappend k f g = k (f .*. g)

instance (Semiring r, FreeUnital r f) => Counital r (f r) where
  coempty k = k unit
-}


-- incoherent
-- instance (UnitalAlgebra r a, Coalgebra r c) => Coalgebra (a -> r) c where coempty k a = coempty (`k` a)
-- instance Coalgebra () a where coempty _ = ()


-- | The tensor Hopf algebra
-- Δ(x) = x ⊗ 1 + 1 ⊗ x, x in V, Δ(1) = 1 ⊗ 1
instance Semiring r => Coalgebra r [a] where
  coappend f as bs = f (mappend as bs)

instance Semiring r => Counital r [a] where
  coempty k = k []

-- | The tensor Hopf algebra
instance Semiring r => Coalgebra r (Seq a) where
  coappend f as bs = f (mappend as bs)

instance Semiring r => Counital r (Seq a) where
  coempty k = k (Seq.empty)

-- | the free commutative band coalgebra
instance (Semiring r, Ord a) => Coalgebra r (Set.Set a) where
  coappend f as bs = f (Set.union as bs)

instance (Semiring r, Ord a) => Counital r (Set.Set a) where
  coempty k = k (Set.empty)

-- | the free commutative band coalgebra over Int
instance Semiring r => Coalgebra r IntSet.IntSet where
  coappend f as bs = f (IntSet.union as bs)

instance Semiring r => Counital r IntSet.IntSet where
  coempty k = k (IntSet.empty)

{-
-- | the free commutative coalgebra over a set and a given semigroup
instance (Semiring r, Ord a, Additive b) => Coalgebra r (Map a b) where
  coappend f as bs = f (Map.unionWith (+) as bs)
  coempty k = k (Map.empty)

-- | the free commutative coalgebra over a set and Int
instance (Semiring r, Additive b) => Coalgebra r (IntMap b) where
  coappend f as bs = f (IntMap.unionWith (+) as bs)
  coempty k = k (IntMap.empty)
-}