module Exercise_5 where
import Test.QuickCheck
import Data.List
import Data.Maybe
import qualified Data.HashMap.Lazy as HS

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

-- Faster O(n log n) nub specialised for integers.
-- Probably doesn't make a difference but still.
intNub :: [Int] -> [Int]
intNub = map head . group . sort

longestPath :: Graph -> Vertex -> Int
longestPath (vs, es) v = longestPath' v
  where succMap = HS.map intNub $ HS.fromListWith (++) [(v, [u]) | (u, v) <- es]
        succs v = HS.lookupDefault [] v succMap
        m = HS.fromList [(v, longestPath'' v) | v <- vs]
        longestPath' v = fromJust (HS.lookup v m)
        longestPath'' v = maximum (0 : map ((+1) . longestPath') (succs v))
{-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 u =
  do n <- choose (1, u)
     let v = [1..n]
     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])