import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.List (transpose)
import Control.Monad (zipWithM_)

newtype Vector a = Vector { unVector :: IntMap a }

vector :: Num a => IntMap a -> Vector a
vector = Vector

newtype Vector1 a = Vector1 { unVector1 :: Int }

vector1 :: Num a => Int -> Vector1 a
vector1 = Vector1

class VectorLike v where
    toVector :: Num a => v a -> Vector a

instance VectorLike Vector where
    toVector = id

instance VectorLike Vector1 where
    toVector (Vector1 p) = Vector (IntMap.singleton p 1)

instance Functor Vector where
    fmap f = Vector . IntMap.map f . unVector

unVector' :: (Num a, VectorLike v) => v a -> IntMap a
unVector' = unVector . toVector

(<+>) :: (Num a, VectorLike v1, VectorLike v2) => v1 a -> v2 a -> Vector a
a <+> b = Vector $ IntMap.unionWith (+) (unVector' a) (unVector' b)

(<->) :: (Num a, VectorLike v1, VectorLike v2) => v1 a -> v2 a -> Vector a
a <-> b = a <+> fmap negate (toVector b)

(*>) :: (Num a, VectorLike v) => v a -> a -> Vector a
a *> b = fmap (*b) (toVector a)

(<*) :: (Num a, VectorLike v) => a -> v a -> Vector a
a <* b = fmap (a*) (toVector b)

infixl 6 <+>,<->
infixl 7 *>
infixr 7 <*

emptyVector :: Num a => Vector a
emptyVector = Vector IntMap.empty





data Matrix a = Matrix { unMatrix :: [[a]] }
              | Diagonal { unDiagonal :: [a] }
    deriving (Show, Eq)

toMatrix :: Num a => Matrix a -> Matrix a
toMatrix (Matrix a) = Matrix a
toMatrix (Diagonal a) = Matrix [replicate i 0 ++ [aii] ++ repeat 0 | (i, aii) <- zip [0..] a]

unMatrix' :: Num a => Matrix a -> [[a]]
unMatrix' = unMatrix . toMatrix

matrix :: [[a]] -> Matrix a
matrix = Matrix

diagonal :: [a] -> Matrix a
diagonal = Diagonal

instance Num a => Num (Matrix a) where
    Diagonal a + Diagonal b = diagonal (zipWith (+) a b)
    a + b = matrix (zipWith (zipWith (+)) (unMatrix' a) (unMatrix' b))

    negate (Matrix a) = matrix (map (map negate) a)
    negate (Diagonal a) = diagonal (map negate a)

    fromInteger = diagonal . repeat . fromInteger

    Matrix a * Matrix b = let tb = transpose b
                              c = [[sum (zipWith (*) ra cb) | cb <- tb] | ra <- a]
                          in
                              matrix c
    Diagonal a * Diagonal b = diagonal (zipWith (*) a b)
    Diagonal a * Matrix b = matrix (zipWith (\v row -> map (v*) row) a b)
    Matrix a * Diagonal b = matrix (map (\row -> zipWith (*) row b) a)

    abs = error "Matrix: abs undefined"
    signum = error "Matrix: abs undefined"





data LRVariable a = LRV { initialValue :: a, dependency :: Vector a }

dmap :: Num a => (Vector a -> Vector a) -> LRVariable a -> LRVariable a
dmap f (LRV val dep) = LRV val (f dep)

type LRVariables a = IntMap (LRVariable a)

data LinearRecursive a b = LR { unLR :: Int -> (b, Int, LRVariables a -> LRVariables a) }

instance Num a => Monad (LinearRecursive a) where
    return a = LR (const (a, 0, id))
    a >>= b = LR $ \v -> let (ra, nva, ma) = unLR a v
                             (rb, nvb, mb) = unLR (b ra) (v + nva)
                         in
                             (rb, nva + nvb, mb . ma)

newVariable :: Num a => a -> LinearRecursive a (Vector1 a)
newVariable val0 = LR $ \v -> (vector1 v, 1, IntMap.insert v variable)
  where
    variable = LRV { initialValue = val0, dependency = emptyVector }

newVariables :: Num a => [a] -> LinearRecursive a [Vector1 a]
newVariables vals = do
    ret <- mapM newVariable vals
    zipWithM_ (<:-) (tail ret) ret
    return ret

newConstant :: Num a => a -> LinearRecursive a (Vector a)
newConstant val = do
    ret <- newVariable val
    ret <:- ret
    return (toVector ret)

(<+-) :: (Num a, VectorLike v) => Vector1 a -> v a -> LinearRecursive a ()
(<+-) var dep = LR (const ((), 0, IntMap.adjust (dmap (<+>toVector dep)) (unVector1 var)))

(<:-) :: (Num a, VectorLike v) => Vector1 a -> v a -> LinearRecursive a ()
(<:-) var dep = LR (const ((), 0, IntMap.adjust (dmap (const (toVector dep))) (unVector1 var)))

infix 1 <:-,<+-

buildMatrix :: Num a => LRVariables a -> (Matrix a, Matrix a)
buildMatrix mapping = (Matrix trans, Matrix $ map (\x -> [x]) initValues)
  where
    initValues = map initialValue (IntMap.elems mapping)
    rawDep = map (unVector'.dependency) (IntMap.elems mapping)
    varCount = length initValues
    trans = map (\m -> [IntMap.findWithDefault 0 i m | i <- [0..varCount-1]]) rawDep

runLinearRecursive :: (Num a, Integral b, VectorLike v) => LinearRecursive a (v a) -> b -> a
runLinearRecursive monad steps = sum [head (res !! i) * ai | (i, ai) <- IntMap.assocs (unVector' target)]
  where
    (target, nv, g) = unLR monad 0
    dep = g IntMap.empty
    (trans, init) = buildMatrix dep

    Matrix res = trans^steps * init

fib = do
    a <- newVariable 1
    b <- newVariable 1
    a <:- b
    b <:- a <+> b
    return a

fib2 = do
    [f0,f1] <- newVariables [1, 1]
    f0 <:- f0 <+> f1
    return f1

a004146 = do
    two <- newConstant 2
    [a0, a1] <- newVariables [1, 0]
    a0 <:- a0 *> 3 <-> a1 <+> two
    return a1

main = print $ map (runLinearRecursive fib2) [0..10]