module KleeneAlgebras where

import Data.Array

infixl 6 <+>
infixl 7 <*>

infixl 6 |+|
infixl 7 |*|

-- ---------------------------------------------------------------

class SemiRing a where
  (<+>), (<*>) :: a -> a -> a
  zero, one :: a
  -- (a, (<+>), zero) is a commutative monoid
  -- (a, (<*>), one) is a monoid

class SemiRing a => IdempotentSemiRing a
  -- (<+>) is idempotent

class IdempotentSemiRing a => KleeneAlgebra a where
  many :: a -> a
  -- 1 <+> p <*> many p = many p
  -- 1 <+> many p <*> p = many p
  -- q <+> p <*> r <= r ==> many p <*> q <= r
  -- q <+> r <*> p <= r ==> q <*> many p <= r

-- ---------------------------------------------------------------

data MinPlus a = Finite a | Infinity deriving (Eq, Ord, Show)

instance (Num a, Ord a) => SemiRing (MinPlus a) where
  -- min
  Infinity <+> y = y
  x <+> Infinity = x
  Finite x <+> Finite y = Finite (min x y)

  -- plus
  Infinity <*> _ = Infinity
  _ <*> Infinity = Infinity
  Finite x <*> Finite y = Finite (x + y)

  zero = Infinity
  one  = Finite 0

instance (Num a, Ord a) => IdempotentSemiRing (MinPlus a)

instance (Num a, Ord a) => KleeneAlgebra (MinPlus a) where
  many _ = one

-- ---------------------------------------------------------------

newtype Mat a = Mat{ unMat :: Array (Int,Int) a } deriving (Eq, Show)

matrix :: SemiRing a => Int -> Int -> (Int -> Int -> a) -> Mat a
matrix m n f = Mat $ array b [((i,j), f i j) | (i,j) <- range b]
  where b = ((1,1),(m,n))

matFromList :: SemiRing a => Int -> Int -> [[a]] -> Mat a
matFromList m n xss = Mat $ array b [((i,j),x) | (i,xs) <- zip [1..] xss, (j,x) <- zip [1..] xs]
  where b = ((1,1),(m,n))

matToList :: Mat a -> [[a]]
matToList (Mat x) = [ [x!(i,j) | j<-[1..n]] | i<-[1..m]]
  where
    (_, (m,n)) = bounds x

matZero :: SemiRing a => Int -> Int -> Mat a
matZero m n = matrix m n (\_ _ -> zero)

matOne :: SemiRing a => Int -> Mat a
matOne n = matrix n n (\i j -> if i==j then one else zero)

(|+|) :: SemiRing a => Mat a -> Mat a -> Mat a
Mat x |+| Mat y
  | bounds x /= bounds y = error "Matrix size mismatch"
  | otherwise = matrix m n (\i j -> x!(i,j) <+> y!(i,j))
  where (_, (m,n)) = bounds x

(|*|) :: SemiRing a => Mat a -> Mat a -> Mat a
Mat x |*| Mat y
  | n /= n' = error "Matrix size mismatch"
  | otherwise = matrix m o (\i k -> sum [x!(i,j) <*> y!(j,k) | j<-[1..n]])
  where
    (_, (m,n))  = bounds x
    (_, (n',o)) = bounds y
    sum = foldl (<+>) zero

matMany :: KleeneAlgebra a => Mat a -> Mat a
matMany (Mat x)
    | m /= n = error "Matrix size mismatch"
    | n==1 = matrix 1 1 (\_ _ -> many (x!(1,1)))
    | otherwise =
        let n1 = n `div` 2
            n2 = n - n1
            a = matrix n1 n1 (\i j -> x!(i,j))
            b = matrix n1 n2 (\i j -> x!(i,j+n1))
            c = matrix n2 n1 (\i j -> x!(i+n1,j))
            d = matrix n2 n2 (\i j -> x!(i+n1,j+n1))
            many_a = matMany a
            many_d = matMany d
            a' = matMany (a |+| b |*| many_d |*| c)
            b' = a' |*| b |*| many_d
            c' = d' |*| c |*| many_a
            d' = matMany (d |+| c |*| many_a |*| b)
        in matrix n n $ \i j ->
            case (i<=n1, j<=n1) of
              (True, True)   -> unMat a' ! (i,j)
              (True, False)  -> unMat b' ! (i,j-n1)
              (False, True)  -> unMat c' ! (i-n1,j)
              (False, False) -> unMat d' ! (i-n1,j-n1)
  where
    (_, (m,n)) = bounds x

{-
instance SemiRing a => SemiRing (Mat a) where
  zero = undefined
  one = undefined
instance IdempotentSemiRing a => IdempotentSemiRing (Mat a)
instance KleeneAlgebra a => KleeneAlgebra (Mat a) where
  many = matMany
-}

-- ---------------------------------------------------------------

mat1 :: Mat (MinPlus Int)
mat1 = matFromList 4 4 xs
  where
    xs = [ [Finite 1, Finite 2, Finite 5, Finite 6]
         , [Infinity, Finite 4, Finite 2, Finite 4]
         , [Finite 1, Infinity, Infinity, Finite 1]
         , [Infinity, Infinity, Infinity, Infinity]
         ]

mat2 :: Mat (MinPlus Int)
mat2 = matMany mat1