{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Exercise_5 where

import Control.Monad
import Control.Monad.ST.Strict
import Data.Array.ST.Safe
import qualified Data.Array as A
import qualified Data.IntMap.Strict as IM
import Data.List
import Data.Maybe
import Data.Ord
import Data.STRef
import Test.QuickCheck

{- H5.3 -}
{-WETT-}
type Vertex = Int
type Edge = (Vertex, Vertex)
type Graph = ([Vertex], [Edge])

type CVertex = Int
-- Since we work on contiguous indices, a map will eventually consume at least as much RAM as arrays, so we opt for the latter
type CAdjMatrix s = STArray s CVertex [CVertex]
type CIncMatrix s = STArray s CVertex [CVertex]

-- DFS-based topological sort
topoSort :: CVertex -> Int -> CAdjMatrix s -> ST s [CVertex]
topoSort !start !size adj = do
    used <- newArray (0, size - 1) False
    res <- newSTRef []
    aux start used adj res
    readSTRef res
  where
    aux :: CVertex -> STUArray s Int Bool -> CAdjMatrix s -> STRef s [CVertex] -> ST s ()
    aux !s used adj res = do
        writeArray used s True
        adjs <- readArray adj s
        forM_ adjs $ \ !n -> do
            !visited <- readArray used n
            unless visited $ aux n used adj res
        modifySTRef res (s:)

-- The used graph representation has the following two properties over what is supplied to longestPath:
-- - Node indices are contiguous starting from 0
-- - Adjacency is saved in a map of lists instead of the list of tuples, allowing faster lookup
-- Building this representation might be quite some overhead for small graphs, but hopefully gets amortised when graphs get sufficiently big.
toCGraph :: Graph -> ST s (IM.IntMap CVertex, Int, CVertex, CAdjMatrix s, CIncMatrix s)
toCGraph (v, e) = do
    hasIncoming <- newArray (0, size - 1) False
    adj <- newArray (0, size - 1) []
    inc <- newArray (0, size - 1) []
    forM_ e $ \(!fV, !tV) -> do
        let (!fC, !tC) = (toC IM.! fV, toC IM.! tV)
        writeArray adj fC =<< (tC:) <$> readArray adj fC
        writeArray inc tC =<< (fC:) <$> readArray inc tC
        writeArray hasIncoming tC True
    !s <- smallestUnused 0 hasIncoming
    pure (toC, size, s, adj, inc)
  where
    !size = length v
    toC = IM.fromList $ zip v [0..]
    smallestUnused :: Int -> STUArray s Int Bool -> ST s Int
    smallestUnused !i arr = do
        !isUsed <- readArray arr i
        if not isUsed then pure i else smallestUnused (i+1) arr

longestPath :: Graph -> Vertex -> Int
longestPath g targetV = runST stBlock
  where
    -- God, this produced quite some type errors until I finally got it under control.
    -- Guess I no longer have to fear on students eventually picking up ST-based imperative code, as the type errors are enough to scare any sane person away.
    stBlock :: ST s Int
    stBlock = do
        (toC, !size, !s, adj, inc) <- toCGraph g
        let !targetC = toC IM.! targetV
        soFar <- newArray_ (0, size - 1) :: ST s (STUArray s Int Int)
        writeArray soFar s 0
        topoSorted <- topoSort s size adj
        aux (tail topoSorted) inc targetC soFar
    aux :: [CVertex] -> CIncMatrix s -> CVertex -> STUArray s CVertex Int -> ST s Int
    aux [] _ _ _ = pure 0 -- single-node graph
    aux (!t:ts) inc !targetC soFar = do
        incs <- readArray inc t
        !pathLength <- succ . maximum <$> mapM (readArray soFar) incs
        if t == targetC
            then pure pathLength
            else do
                writeArray soFar t pathLength
                aux ts inc targetC soFar
{-TTEW-}

-- generates a DAG with u vertices and only one node without incoming edges
-- you can use this function to test your implementation using QuickCheck
genDag :: Int -> Gen Graph
genDag n = let v = [1..n] in
  do b <- mapM (\i -> choose (1,n-i)) [1..n-1]
     t <- mapM (\(c,i) -> vectorOf c (choose (i+1, n))) (zip b [1..n])
     let e = nub $ ([(1, i) | i<-[2..n]] ++ edges t 1 [])
     return $ (v,e)
  where
    edges [] _ acc = acc
    edges (ts:xs) i acc = edges xs (i+1) (acc ++ [(i,t) | t<-ts])

-- Shamelessly stolen from Piazza
dotGraph :: Graph -> String
dotGraph (v, e) = "digraph{" ++ ((listToString . map (toString)) v) ++ ";" ++ ((listToString . map (\(a, b) -> (toString a) ++ "->" ++ (toString b))) e) ++ "}"
  where
      toString s = show s :: String
      listToString = intercalate ";"

-- vim: set expandtab shiftwidth=4 :