{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE FlexibleInstances,MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DeriveGeneric #-}
module Algebra.Module.Free
  (
    T(Free)
  , first
  , second
    
  , Map
  , Domain
  , Codomain
  , apply
  
  , matrix
  , vector
  , fromVector
  , linear
  , filter
  , filterZeroes
  , toList_
  
  , Enumerable
  , enumerate
  , coefficient
  , basisVector
  , fromList
  , Basis(Basis)
  ) where


import           Algebra.Linear (Matrix,Vector)

import qualified Algebra.Additive
import qualified Algebra.Module
import qualified Algebra.ModuleBasis
import qualified Algebra.Ring
import           NumericPrelude hiding (filter)

-- import           Data.Bifunctor (Bifunctor,first,second)
import           Data.Binary  (Binary)
import           Data.List    (intercalate,foldl')
import qualified Data.Map.Strict as Map
import           Data.Proxy   (Proxy(Proxy))
import           GHC.Generics (Generic)


newtype T k t
  = Free (Map.Map t k)
  deriving (Generic)

type Free k t = T k t

-- The natural Bifunctor Free instance is not possible, because of the constraints on 'second'.
first :: (k₁ -> k₂) -> Free k₁ t -> Free k₂ t
first f (Free m) = Free (fmap f m)
second :: (Algebra.Additive.C k,Ord t₂) => (t₁ -> t₂) -> Free k t₁ -> Free k t₂
second g (Free m) = Free (Map.mapKeysWith (+) g m)

instance (Algebra.Additive.C k,Ord t) => Algebra.Additive.C (Free k t) where
  zero = Free Map.empty
  negate (Free m) = Free (fmap negate m)
  Free m₁ + Free m₂ = Free $ Map.unionWith (+) m₁ m₂

instance (Algebra.Ring.C k,Ord t) => Algebra.Module.C k (Free k t) where
  c *> (Free m) = Free (fmap (c *) m)

instance (Show t,Ord t,Show k) => Show (Free k t) where
  show (Free m) = intercalate " + " . map (\ (x,c) -> show c ++ " · " ++ show x) $ Map.toList m

instance (Binary t,Binary k) => Binary (Free k t)

-- 'b' is a basis for the module 'm'
class ModuleBasis b m where
  type Scalar b m
  type BasisElement b m
  coefficient :: b -> BasisElement b m -> m -> Scalar b m
  basisVector :: b -> BasisElement b m -> m

fromList :: (ModuleBasis b m,Algebra.Module.C (Scalar b m) m) => b -> [(BasisElement b m,Scalar b m)] -> m
fromList b = foldl' (+) zero . map (\ (x,c) -> c *> basisVector b x)

toList_ :: Free k t -> [(t,k)]
toList_ (Free m) = Map.toList m

data Basis
  = Basis

instance (Algebra.Ring.C k,Ord t) => ModuleBasis Basis (Free k t) where
  type Scalar Basis (Free k t) = k
  type BasisElement Basis (Free k t) = t
  coefficient Basis x (Free m) = Map.findWithDefault zero x m
  basisVector Basis x = Free (Map.singleton x one)

class (Ord t) => Enumerable p t where
  enumerate :: p -> [t]

-- instance forall k t. (Enumerable t,Algebra.Ring.C k) => Algebra.ModuleBasis.C k (Free k t) where
--   dimension _ _ = length $ enumerate (Proxy :: Proxy t)
--   flatten (Free m) = map (flip coefficient m) $ enumerate (Proxy :: Proxy t)
--   basis _ = map basisVector $ enumerate (Proxy :: Proxy t)


class Map k m where
  type Domain m
  type Codomain m
  apply :: m -> Domain m -> Codomain m

instance (Algebra.Ring.C k,Ord t₁,Ord t₂) => Map k (Free k (t₁,t₂)) where
  type Domain (Free k (t₁,t₂)) = Free k t₁
  type Codomain (Free k (t₁,t₂)) = Free k t₂
  apply (Free m) (Free v₁) = Free $ Map.foldrWithKey
    (\ (x₁,x₂) c -> case Map.lookup x₁ v₁ of
      Nothing -> id
      Just d  -> Map.insertWith (+) x₂ (c * d)
    )
    Map.empty
    m

matrix :: forall k t₁ t₂ p₁ p₂. (Enumerable p₁ t₁,Enumerable p₂ t₂,Algebra.Ring.C k) => p₁ -> p₂ -> Free k (t₁,t₂) -> Matrix k
matrix p₁ p₂ m = [[coefficient Basis (x₁,x₂) m | x₁ <- enumerate p₁] | x₂ <- enumerate p₂]

vector :: forall k t p. (Enumerable p t,Algebra.Ring.C k) => p -> Free k t -> Vector k
vector p v = [coefficient Basis x v | x <- enumerate p]

fromVector :: forall k t p. (Enumerable p t,Algebra.Ring.C k) => p -> Vector k -> Free k t
fromVector p = Free . Map.fromList . zip (enumerate p)

linear :: (Algebra.Module.C k a) => (t -> a) -> Free k t -> a
linear f (Free m) = sum . map (\ (t,c) -> c *> f t) . Map.toList $ m

filterZeroes :: (Algebra.Additive.C k,Eq k,Ord t) => Free k t -> Free k t
filterZeroes = filter $ const $ (/=) zero

filter :: (Algebra.Additive.C k,Ord t) => (t -> k -> Bool) -> Free k t -> Free k t
filter f (Free m) = Free $ Map.filterWithKey f m

-- data X2 = S1 | S2 deriving (Eq,Ord)
-- instance Enumerable X2 where
--   enumerate _ = [S1,S2]
-- data X3 = T1 | T2 | T3 deriving (Eq,Ord)
-- instance Enumerable X3 where
--   enumerate _ = [T1,T2,T3]
-- 
-- m :: Free Integer (X2,X3)
-- m = (2 :: Integer) *> basisVector FreeBasis (S2,T2) - (3 :: Integer) *> basisVector FreeBasis (S1,T3) :: Free Integer (X2,X3)