module Graph where

import qualified Data.Map as M
import qualified Data.Set as S

data Vertex v =
    Vertex
    { vertexId :: Int
    , vertexValue :: v
    }

instance Eq (Vertex v) where
    Vertex l _ == Vertex r _ = l == r

instance Ord (Vertex v) where
    compare l r = compare (vertexId l) (vertexId r)

data Edge =
    Edge
    { edgeSrc :: Int
    , edgeDest :: Int
    }

instance Show Edge where
    show (Edge x y) = show (x, y)

instance Eq Edge where
    Edge ls ld == Edge rs rd = (ls, ld) == (rs, rd)

instance Ord Edge where
    compare (Edge ls ld) (Edge rs rd) = compare (ls, ld) (rs, rd)

data Graph v =
    Graph
    { graphVertices :: [Vertex v]
    , graphEdges :: [Edge]
    } deriving Show

instance Show v => Show (Vertex v) where
    show (Vertex _ n) = show n

testGraph :: Graph String
testGraph = Graph vs es
  where
    vs = [ Vertex 1 "a"
         , Vertex 2 "b"
         , Vertex 3 "c"
         ]

    es = [ Edge 1 2
         , Edge 1 3
         , Edge 2 3
         ]

adjacentList :: Graph v -> Graph [Int]
adjacentList (Graph _ es) = Graph (es >>= go) es
  where
    go (Edge src dest) =
        [ Vertex src [dest]
        , Vertex dest [src]
        ]

reduceByKey :: (v -> v -> v) -> Graph v -> Graph v
reduceByKey k (Graph vs es) = Graph vs' es
  where
    vs' = fmap (uncurry Vertex) $ M.assocs $ foldl go M.empty vs

    go m (Vertex vid v) =
        let _F (Just v') = Just $ k v v'
            _F _         = Just v in
        M.alter _F vid m

appending :: [a] -> [b] -> [Either a b]
appending xs vs = fmap Left xs ++ fmap Right vs

join :: Graph v -> Graph w -> Graph (v, w)
join (Graph lvs les) (Graph rvs res) = Graph vs' es'
  where
    vs' = (M.assocs $ foldl go M.empty (appending lvs rvs)) >>= reducing

    es' = S.toList $ foldl pruning S.empty (les ++ res)

    reducing (vid, Left _)    = []
    reducing (vid, Right tup) = [ Vertex vid tup ]

    pruning s e = S.insert e s

    go m (Left (Vertex vid v))  = M.insert vid (Left v) m
    go m (Right (Vertex wid w)) = M.adjust (\(Left v) -> Right (v, w)) wid m

vertex :: Int -> Vertex String
vertex i = Vertex i (show i)

testConnected :: Graph String
testConnected = Graph vs es
  where
    vs = [ vertex 1
         , vertex 2
         , vertex 3
         , vertex 4
         , vertex 5
         , vertex 6
         ]

    es = [ Edge 1 3
         , Edge 1 5
         , Edge 3 5
         , Edge 2 4
         , Edge 2 6
         , Edge 4 6
         ]

smallestNode :: Int -> Int -> M.Map Int Int -> M.Map Int Int
smallestNode n v m = M.alter go n m
  where
    go (Just o)
        | v < o     = Just v
        | otherwise = Just o
    go _
        | v < n     = Just v
        | otherwise = Just n

connectedComponents :: Graph v -> Graph Int
connectedComponents (Graph vs es) = Graph vs' es
  where
    vs' = fmap (uncurry Vertex) $ M.assocs $ foldl go M.empty es

    go m (Edge x y) =
        smallestNode x y (smallestNode y x m)