{-# LANGUAGE FlexibleContexts,
             RankNTypes,
             ScopedTypeVariables,
             NoMonomorphismRestriction,
             RelaxedPolyRec,
             NoMonoLocalBinds #-}

module Fregel where

import Data.Maybe

type Vid = Int
data Vertex a b = V Vid a [Edge a b] [Edge a b] (Graph a b)
type Edge a b = (b, Vertex a b)
type Graph a b = [Vertex a b]
data Termination a = Fix | Iter Int | Until (a -> Bool)

{- Vertex equivalence: its values and neighbors' ids only.-}
instance (Eq a) => Eq (Vertex a b) where
  (==) (V i1 a1 is1 rs1 g1) (V i2 a2 is2 rs2 g2) =
     i1 == i2 && a1 == a2 &&
     (map (vid . snd) is1) == (map (vid . snd) is2) && (map (vid . snd) rs1) == (map (vid . snd) rs2)

instance (Show a, Show b) => Show (Vertex a b) where
  show (V vid a is rs g) =
    "V " ++ show vid ++ " " ++ show a ++
    " [" ++ showEdges is ++ "] [" ++ showEdges rs ++ "]"
    where showEdges [] = ""
          showEdges [e] = showE e
          showEdges (e:es) = showE e ++ ", " ++ showEdges es
          showE (b, V k _ _ _ _) = "(" ++ show b ++ ",v" ++ show k ++ ")"

getVertexId :: Vertex a b -> Vid
getVertexId (V vid _ _ _ _) = vid

getVertexValue :: Vertex a b -> a
getVertexValue (V _ a _ _ _) = a

getGraph :: Vertex a b -> Graph a b
getGraph (V _ _ _ _ g) = g

makeGraph :: (a -> r -> c) -> Graph a b -> [r] -> Graph c b
makeGraph vf vs rs = newvs
  where newvs = zipWith cf vs rs
        vps = zip (map getVertexId vs) newvs
        convE (b,v) = (b, fromJust (lookup (getVertexId v) vps))
        cf (V vid a is rs g) r = V vid (vf a r) (map convE is) (map convE rs) newvs

graphy :: Graph a b -> [r] -> Graph r b
graphy g = makeGraph (\a r -> r) g

-- short-hands
val = getVertexValue
vid = getVertexId
is (V _ _ es _ _) = es
rs (V _ _ _ es _) = es
gof = getGraph

-- field access notation
(.^) :: forall a c . a -> (a->c) -> c
(.^) a f = f a
infixl .^

(!=) :: forall a . Eq a => a -> a -> Bool
(!=) a b = not (a == b)
infixl !=

termination :: Eq a => Termination a -> [a] -> a
termination Fix xs =
  fst . head . dropWhile (\(a,b) -> (a /= b)) $ zip xs (tail xs)
termination (Iter n) xs = head (drop n xs)
termination (Until p) xs = head $ dropWhile (not.p) xs

-- fregel computation with two tables of the previous values and current values
fregel :: (Eq r) =>
          (Vertex a b -> r) ->
          (Vertex a b -> (Vertex a' b -> r) -> (Vertex a' b -> r) -> r) ->
          Termination (Graph r b) ->
          Graph a b ->
          Graph r b
fregel h f t g = 
  let rs0 = map h g
      step rs = let rs' = map (\v -> f v prev curr) g
                    prev u = rs !! ((getVertexId u)-1)
                    curr u = rs' !! ((getVertexId u)-1)
                in rs'
      rss = iterate step rs0
  in termination t (map (graphy g) rss)

-- iterative fregel!
giter :: (Eq r) =>
         (Vertex a b -> r) -> (Graph r b -> Graph r b) ->
         Termination (Graph r b) -> Graph a b -> Graph r b
giter h f t g = 
  let g0 = makeGraph (\a r -> r) g (map h g)
      gss = iterate f g0
  in termination t gss

-- gmap is a special case of fregel; only initialization (one superstep?)
gmap :: (Eq r) => (Vertex a b -> r) -> Graph a b -> Graph r b
gmap f g = fregel f ft Fix g
           where ft v prev curr = prev v

-- gzip; two graphs have to be of the same shape
data Pair a b = Pair {_fst :: a, _snd :: b} deriving (Show, Eq)
gzip :: Graph a1 b -> Graph a2 b -> Graph (Pair a1 a2) b
gzip g1 g2 = makeGraph (\a r -> Pair a r) g1 (map getVertexValue g2)