{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE DeriveFunctor #-}
module Graph (
      Graph(..)
    , Vertex(..)
    , editVertex
    , follow
    , rewind
    , singleton
    , stitch
    , weave
  ) where

import Data.Map ((!), Map)
import qualified Data.Map as Map (adjust, empty, insert, lookup, singleton, size, toList)
import Data.Set (Set)
import qualified Data.Set as Set (insert, member, singleton)
import Tree (Tree(..), Structure(..))

type VertexID = Int

data Vertex edge label = Vertex {
      label :: label
    , edges :: Map edge VertexID
  } deriving (Functor, Show)

data Graph edge label = Graph {
      vertices :: Map VertexID (Vertex edge label)
    , focus :: VertexID
    , root :: VertexID
  } deriving (Functor, Show)

getStructureWithoutLoop :: Show edge => Set VertexID -> Graph edge String -> Structure
getStructureWithoutLoop visitedIDs graph@(Graph {focus, vertices}) = Node [(
      "(" ++ show focus ++ ":" ++ label ++ ")"
    , Node $ recurseOnEdge <$> Map.toList edges
  )]
  where
    Vertex label edges = vertices ! focus
    recurseOnEdge (edge, vertexID) = (
          show edge
        , if Set.member vertexID visitedIDs
          then Node [("(:" ++ show vertexID ++ ")", Node [])]
          else getStructureWithoutLoop (Set.insert vertexID visitedIDs) $ graph {focus = vertexID}
      )

instance Show edge => Tree (Graph edge) where
  getStructure graph = 
    getStructureWithoutLoop (Set.singleton $ root graph) (rewind graph)

vertex :: label -> Vertex edge label
vertex label = Vertex {label, edges = Map.empty}

singleton :: label -> Graph edge label
singleton label = Graph {
      vertices = Map.singleton 0 $ vertex label
    , focus = 0
    , root = 0
  }

rewind :: Graph edge label -> Graph edge label
rewind graph = graph {focus = root graph}

follow :: Ord edge => Graph edge label -> edge -> Maybe (Graph edge label)
follow graph@(Graph {vertices, focus}) edge =
  setFocus <$> Map.lookup edge (edges $ vertices ! focus)
  where
    setFocus vertexID = graph {focus = vertexID}

stitch :: (Monoid label, Ord edge) => Graph edge label -> edge -> Graph edge label
stitch graph edge =
  case graph `follow` edge of
    Nothing ->
      let newVertexID = Map.size $ vertices graph in
      let link aVertex = aVertex {edges = Map.insert edge newVertexID $ edges aVertex} in
      graph {
          vertices =
            Map.adjust link (focus graph) . Map.insert newVertexID (vertex mempty) $ vertices graph
        , focus = newVertexID
      }
    Just newGraph -> newGraph

weave :: (Monoid label, Ord edge) => Graph edge label -> [edge] -> Graph edge label
weave = foldl stitch

editVertex :: Graph edge label -> (Vertex edge label -> Vertex edge label) -> Graph edge label
editVertex graph@(Graph {vertices, focus}) vertexEditor =
  graph {vertices = Map.adjust vertexEditor focus vertices}