module FULE.Internal.Sparse
 ( Matrix
 , Pos
 , empty
 , eye
 , matrix
 , dims
 , expandTo
 , get
 , set
 , del
 , count
 , add
 , sub
 , mul
 , star
 , filter
 ) where

import Prelude hiding (filter)
import Control.DeepSeq
import Data.List hiding (filter)
import qualified Data.Map.Strict as Map


type Pos = (Int, Int)

maxp :: (a, b) -> (a, b) -> (a, b)
maxp (a
rl, b
cl) (a
rr, b
cr) = (a -> a -> a
forall a. Ord a => a -> a -> a
max a
rl a
rr, b -> b -> b
forall a. Ord a => a -> a -> a
max b
cl b
cr)


data Matrix a
  = M
    { forall a. Matrix a -> (Int, Int)
dimensionsOf :: (Int, Int)
    , forall a. Matrix a -> Map (Int, Int) a
matrixOf :: Map.Map Pos a
    }

instance (NFData a) => NFData (Matrix a) where
  rnf :: Matrix a -> ()
rnf (M (Int
r, Int
c) Map (Int, Int) a
m) = Int -> () -> ()
forall a b. a -> b -> b
seq Int
r (() -> ()) -> (() -> ()) -> () -> ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> () -> ()
forall a b. a -> b -> b
seq Int
c (() -> ()) -> (() -> ()) -> () -> ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map (Int, Int) a -> () -> ()
forall a b. NFData a => a -> b -> b
deepseq Map (Int, Int) a
m (() -> ()) -> () -> ()
forall a b. (a -> b) -> a -> b
$ ()

instance (Num a, Show a) => Show (Matrix a) where
  show :: Matrix a -> String
show (M (Int
0, Int
0) Map (Int, Int) a
_) = String
"[]"
  show m :: Matrix a
m@(M (Int
r, Int
c) Map (Int, Int) a
_) = [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [String
"[ ", String
mat, String
"\n]"]
    where
      row :: Int -> String
row Int
rx = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ (Int -> String) -> [Int] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
cx -> a -> String
forall a. Show a => a -> String
show (a -> String) -> a -> String
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> Matrix a -> a
forall a. Num a => (Int, Int) -> Matrix a -> a
get (Int
rx, Int
cx) Matrix a
m) [Int
1..Int
c]
      mat :: String
mat = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"\n; " ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ (Int -> String) -> [Int] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Int -> String
row [Int
1..Int
r]

instance Functor Matrix where
  fmap :: forall a b. (a -> b) -> Matrix a -> Matrix b
fmap a -> b
f (M (Int, Int)
d Map (Int, Int) a
m) = (Int, Int) -> Map (Int, Int) b -> Matrix b
forall a. (Int, Int) -> Map (Int, Int) a -> Matrix a
M (Int, Int)
d (Map (Int, Int) b -> Matrix b) -> Map (Int, Int) b -> Matrix b
forall a b. (a -> b) -> a -> b
$ (a -> b) -> Map (Int, Int) a -> Map (Int, Int) b
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map a -> b
f Map (Int, Int) a
m

empty :: Matrix a
empty :: forall a. Matrix a
empty = (Int, Int) -> Map (Int, Int) a -> Matrix a
forall a. (Int, Int) -> Map (Int, Int) a -> Matrix a
M (Int
0, Int
0) Map (Int, Int) a
forall k a. Map k a
Map.empty

eye :: (Num a) => Int -> Matrix a
eye :: forall a. Num a => Int -> Matrix a
eye Int
dim = (Int, Int) -> [((Int, Int), a)] -> Matrix a
forall a. (Int, Int) -> [((Int, Int), a)] -> Matrix a
matrix (Int
dim, Int
dim) [((Int
r,Int
c), a
1) | (Int
r, Int
c) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1..Int
dim] [Int
1..Int
dim]]

matrix :: (Int, Int) -> [(Pos, a)] -> Matrix a
matrix :: forall a. (Int, Int) -> [((Int, Int), a)] -> Matrix a
matrix (Int, Int)
dims [((Int, Int), a)]
entries = (Int, Int) -> Map (Int, Int) a -> Matrix a
forall a. (Int, Int) -> Map (Int, Int) a -> Matrix a
M (Int, Int)
dims ([((Int, Int), a)] -> Map (Int, Int) a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [((Int, Int), a)]
entries)

dims :: Matrix a -> (Int, Int)
dims :: forall a. Matrix a -> (Int, Int)
dims = Matrix a -> (Int, Int)
forall a. Matrix a -> (Int, Int)
dimensionsOf

expandTo :: (Int, Int) -> Matrix a -> Matrix a
expandTo :: forall a. (Int, Int) -> Matrix a -> Matrix a
expandTo (Int, Int)
d (M (Int, Int)
_ Map (Int, Int) a
m) = (Int, Int) -> Map (Int, Int) a -> Matrix a
forall a. (Int, Int) -> Map (Int, Int) a -> Matrix a
M (Int, Int)
d Map (Int, Int) a
m

get :: Num a => Pos -> Matrix a -> a
get :: forall a. Num a => (Int, Int) -> Matrix a -> a
get (Int, Int)
pos Matrix a
m = a -> (Int, Int) -> Map (Int, Int) a -> a
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault a
0 (Int, Int)
pos (Map (Int, Int) a -> a) -> Map (Int, Int) a -> a
forall a b. (a -> b) -> a -> b
$ Matrix a -> Map (Int, Int) a
forall a. Matrix a -> Map (Int, Int) a
matrixOf Matrix a
m

set :: (Eq a, Num a) => Pos -> a -> Matrix a -> Matrix a
set :: forall a. (Eq a, Num a) => (Int, Int) -> a -> Matrix a -> Matrix a
set (Int, Int)
pos a
v (M (Int, Int)
d Map (Int, Int) a
m) = (Int, Int) -> Map (Int, Int) a -> Matrix a
forall a. (Int, Int) -> Map (Int, Int) a -> Matrix a
M ((Int, Int) -> (Int, Int) -> (Int, Int)
forall {a} {b}. (Ord a, Ord b) => (a, b) -> (a, b) -> (a, b)
maxp (Int, Int)
pos (Int, Int)
d) Map (Int, Int) a
m'
  where m' :: Map (Int, Int) a
m' = if a
v a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then (Int, Int) -> Map (Int, Int) a -> Map (Int, Int) a
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete (Int, Int)
pos Map (Int, Int) a
m else (Int, Int) -> a -> Map (Int, Int) a -> Map (Int, Int) a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (Int, Int)
pos a
v Map (Int, Int) a
m

del :: Pos -> Matrix a -> Matrix a
del :: forall a. (Int, Int) -> Matrix a -> Matrix a
del (Int, Int)
pos Matrix a
m = Matrix a
m { matrixOf = Map.delete pos $ matrixOf m }

-- remove row/col (both?)

count :: Matrix a -> Int
count :: forall a. Matrix a -> Int
count = Map (Int, Int) a -> Int
forall k a. Map k a -> Int
Map.size (Map (Int, Int) a -> Int)
-> (Matrix a -> Map (Int, Int) a) -> Matrix a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Matrix a -> Map (Int, Int) a
forall a. Matrix a -> Map (Int, Int) a
matrixOf

add :: (Num a) => Matrix a -> Matrix a -> Matrix a
add :: forall a. Num a => Matrix a -> Matrix a -> Matrix a
add (M (Int, Int)
dl Map (Int, Int) a
ml) (M (Int, Int)
dr Map (Int, Int) a
mr) = (Int, Int) -> Map (Int, Int) a -> Matrix a
forall a. (Int, Int) -> Map (Int, Int) a -> Matrix a
M ((Int, Int) -> (Int, Int) -> (Int, Int)
forall {a} {b}. (Ord a, Ord b) => (a, b) -> (a, b) -> (a, b)
maxp (Int, Int)
dl (Int, Int)
dr) (Map (Int, Int) a -> Matrix a) -> Map (Int, Int) a -> Matrix a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a)
-> Map (Int, Int) a -> Map (Int, Int) a -> Map (Int, Int) a
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith a -> a -> a
forall a. Num a => a -> a -> a
(+) Map (Int, Int) a
ml Map (Int, Int) a
mr

sub :: (Num a) => Matrix a -> Matrix a -> Matrix a
sub :: forall a. Num a => Matrix a -> Matrix a -> Matrix a
sub (M (Int, Int)
dl Map (Int, Int) a
ml) (M (Int, Int)
dr Map (Int, Int) a
mr) = (Int, Int) -> Map (Int, Int) a -> Matrix a
forall a. (Int, Int) -> Map (Int, Int) a -> Matrix a
M ((Int, Int) -> (Int, Int) -> (Int, Int)
forall {a} {b}. (Ord a, Ord b) => (a, b) -> (a, b) -> (a, b)
maxp (Int, Int)
dl (Int, Int)
dr) (Map (Int, Int) a -> Matrix a) -> Map (Int, Int) a -> Matrix a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a)
-> Map (Int, Int) a -> Map (Int, Int) a -> Map (Int, Int) a
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith (-) Map (Int, Int) a
ml Map (Int, Int) a
mr

mul :: (Num a) => Matrix a -> Matrix a -> Matrix a
mul :: forall a. Num a => Matrix a -> Matrix a -> Matrix a
mul (M (Int
r, Int
_) Map (Int, Int) a
ml) (M (Int
_, Int
c) Map (Int, Int) a
mr) =
  -- not the most efficient algorithm but very concise
  (Int, Int) -> Map (Int, Int) a -> Matrix a
forall a. (Int, Int) -> Map (Int, Int) a -> Matrix a
M (Int
r, Int
c) (Map (Int, Int) a -> Matrix a) -> Map (Int, Int) a -> Matrix a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> [((Int, Int), a)] -> Map (Int, Int) a
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith a -> a -> a
forall a. Num a => a -> a -> a
(+) ([((Int, Int), a)] -> Map (Int, Int) a)
-> [((Int, Int), a)] -> Map (Int, Int) a
forall a b. (a -> b) -> a -> b
$
  [((Int
rl,Int
cr), a
vla -> a -> a
forall a. Num a => a -> a -> a
*a
vr) | ((Int
rl,Int
cl), a
vl) <- Map (Int, Int) a -> [((Int, Int), a)]
forall {k} {a}. Map k a -> [(k, a)]
ps Map (Int, Int) a
ml, ((Int
rr,Int
cr), a
vr) <- Map (Int, Int) a -> [((Int, Int), a)]
forall {k} {a}. Map k a -> [(k, a)]
ps Map (Int, Int) a
mr, Int
clInt -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
rr]
  where ps :: Map k a -> [(k, a)]
ps = Map k a -> [(k, a)]
forall {k} {a}. Map k a -> [(k, a)]
Map.toList

star :: (Num a) => Matrix a -> Matrix a -> Matrix a
star :: forall a. Num a => Matrix a -> Matrix a -> Matrix a
star (M (Int
r, Int
_) Map (Int, Int) a
ml) (M (Int
_, Int
c) Map (Int, Int) a
mr) =
  -- not the most efficient algorithm but very concise
  (Int, Int) -> Map (Int, Int) a -> Matrix a
forall a. (Int, Int) -> Map (Int, Int) a -> Matrix a
M (Int
r, Int
c) (Map (Int, Int) a -> Matrix a) -> Map (Int, Int) a -> Matrix a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> [((Int, Int), a)] -> Map (Int, Int) a
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith a -> a -> a
forall a. Num a => a -> a -> a
(*) ([((Int, Int), a)] -> Map (Int, Int) a)
-> [((Int, Int), a)] -> Map (Int, Int) a
forall a b. (a -> b) -> a -> b
$
  [((Int
rl,Int
cr), a
vla -> a -> a
forall a. Num a => a -> a -> a
*a
vr) | ((Int
rl,Int
cl), a
vl) <- Map (Int, Int) a -> [((Int, Int), a)]
forall {k} {a}. Map k a -> [(k, a)]
ps Map (Int, Int) a
ml, ((Int
rr,Int
cr), a
vr) <- Map (Int, Int) a -> [((Int, Int), a)]
forall {k} {a}. Map k a -> [(k, a)]
ps Map (Int, Int) a
mr, Int
clInt -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
rr]
  where ps :: Map k a -> [(k, a)]
ps = Map k a -> [(k, a)]
forall {k} {a}. Map k a -> [(k, a)]
Map.toList

filter :: (a -> Bool) -> Matrix a -> Matrix a
filter :: forall a. (a -> Bool) -> Matrix a -> Matrix a
filter a -> Bool
f (M (Int, Int)
d Map (Int, Int) a
m) = (Int, Int) -> Map (Int, Int) a -> Matrix a
forall a. (Int, Int) -> Map (Int, Int) a -> Matrix a
M (Int, Int)
d (Map (Int, Int) a -> Matrix a) -> Map (Int, Int) a -> Matrix a
forall a b. (a -> b) -> a -> b
$ (a -> Bool) -> Map (Int, Int) a -> Map (Int, Int) a
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter a -> Bool
f Map (Int, Int) a
m