module HGraph.Undirected.Generator
       ( grid
       , cycleGraph
       , completeTree
       , completeGraph
       , randomGraph
       )
where

import HGraph.Undirected
import Data.List
import Control.Monad.State
import System.Random

cycleGraph g0 n = foldr (flip addEdge) (foldr (flip addVertex) g0 [0..n-1]) [(x, (x + 1) `mod` n) | x <- [0..n-1]]

grid g0 w h = foldr (flip addEdge) (foldr (flip addVertex) g0 vs) es
  where
    vs = [(x,y) | x <- [1..w], y <- [1..h]]
    es = concat [[(v,(x+1,y)), (v, (x,y+1))] | x <- [1..w-1], y <- [1..h-1], let v = (x,y)]
      ++ [((x, h), (x+1, h)) | x <- [1..w-1]]
      ++ [((w, y), (w, y+1)) | y <- [1..h-1]]

completeTree g0 depth arity = completeTree' (addVertex (empty g0) 0) 0 1
  where
    completeTree' g root d
      | d > depth = g
      | otherwise = head $ drop arity $
                    iterate' (\h -> let r1 = numVertices h
                                   in addEdge (completeTree' (addVertex h r1) r1 (d+1) ) (root,r1)) g

completeGraph g0 k = foldr (flip addEdge) (foldr (flip addVertex) (empty g0) [0..k-1]) [(u,v) | u <- [0..k-1], v <- [u+1..k-1] ]

randomGraph g0 n m
  | m > (n * (n - 1)) `div` 4 = do -- dense graph
    let g1 = foldr (flip addEdge) (foldr (flip addVertex) (empty g0) [1..n]) [(v,u) | v <- [1..n], u <- [v+1..n]]
    removeRandomEdges g1 m
  | otherwise = do -- spare graph
    let g1 = foldr (flip addVertex) (empty g0) [0..n-1]
    addRandomEdges g1 m

addRandomEdges g m
  | numEdges g == m = return g
  | otherwise = do
    v <- randomN 0 (numVertices g - 1)
    u <- randomN 0 (numVertices g - 1)
    if u /= v then
      addRandomEdges (addEdge g (v,u)) m
    else
      addRandomEdges g m

removeRandomEdges g m
  | numEdges g == m = return g
  | otherwise = do
    v <- randomN 0 (numVertices g - 1)
    u <- randomN 0 (numVertices g - 1)
    if u /= v && edgeExists g (v,u) then
      removeRandomEdges (removeEdge g (v,u)) m
    else
      removeRandomEdges g m

randomN :: (Random a, RandomGen g) => a -> a -> State g a
randomN n0 n1 = do
  gen <- get
  let (r,gen') = randomR (n0,n1) gen
  put gen'
  return r