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

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

cycleGraph :: t a -> a -> t a
cycleGraph t a
g0 a
n = ((a, a) -> t a -> t a) -> t a -> [(a, a)] -> t a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((t a -> (a, a) -> t a) -> (a, a) -> t a -> t a
forall a b c. (a -> b -> c) -> b -> a -> c
flip t a -> (a, a) -> t a
forall (t :: * -> *) a. Mutable t => t a -> (a, a) -> t a
addEdge) ((a -> t a -> t a) -> t a -> [a] -> t a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((t a -> a -> t a) -> a -> t a -> t a
forall a b c. (a -> b -> c) -> b -> a -> c
flip t a -> a -> t a
forall (t :: * -> *) a. Mutable t => t a -> a -> t a
addVertex) t a
g0 [a
0..a
na -> a -> a
forall a. Num a => a -> a -> a
-a
1]) [(a
x, (a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
n) | a
x <- [a
0..a
na -> a -> a
forall a. Num a => a -> a -> a
-a
1]]

grid :: t (a, b) -> a -> b -> t (a, b)
grid t (a, b)
g0 a
w b
h = (((a, b), (a, b)) -> t (a, b) -> t (a, b))
-> t (a, b) -> [((a, b), (a, b))] -> t (a, b)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((t (a, b) -> ((a, b), (a, b)) -> t (a, b))
-> ((a, b), (a, b)) -> t (a, b) -> t (a, b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip t (a, b) -> ((a, b), (a, b)) -> t (a, b)
forall (t :: * -> *) a. Mutable t => t a -> (a, a) -> t a
addEdge) (((a, b) -> t (a, b) -> t (a, b))
-> t (a, b) -> [(a, b)] -> t (a, b)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((t (a, b) -> (a, b) -> t (a, b)) -> (a, b) -> t (a, b) -> t (a, b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip t (a, b) -> (a, b) -> t (a, b)
forall (t :: * -> *) a. Mutable t => t a -> a -> t a
addVertex) t (a, b)
g0 [(a, b)]
vs) [((a, b), (a, b))]
es
  where
    vs :: [(a, b)]
vs = [(a
x,b
y) | a
x <- [a
1..a
w], b
y <- [b
1..b
h]]
    es :: [((a, b), (a, b))]
es = [[((a, b), (a, b))]] -> [((a, b), (a, b))]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[((a, b)
v,(a
xa -> a -> a
forall a. Num a => a -> a -> a
+a
1,b
y)), ((a, b)
v, (a
x,b
yb -> b -> b
forall a. Num a => a -> a -> a
+b
1))] | a
x <- [a
1..a
wa -> a -> a
forall a. Num a => a -> a -> a
-a
1], b
y <- [b
1..b
hb -> b -> b
forall a. Num a => a -> a -> a
-b
1], let v :: (a, b)
v = (a
x,b
y)]
      [((a, b), (a, b))] -> [((a, b), (a, b))] -> [((a, b), (a, b))]
forall a. [a] -> [a] -> [a]
++ [((a
x, b
h), (a
xa -> a -> a
forall a. Num a => a -> a -> a
+a
1, b
h)) | a
x <- [a
1..a
wa -> a -> a
forall a. Num a => a -> a -> a
-a
1]]
      [((a, b), (a, b))] -> [((a, b), (a, b))] -> [((a, b), (a, b))]
forall a. [a] -> [a] -> [a]
++ [((a
w, b
y), (a
w, b
yb -> b -> b
forall a. Num a => a -> a -> a
+b
1)) | b
y <- [b
1..b
hb -> b -> b
forall a. Num a => a -> a -> a
-b
1]]

completeTree :: t a -> t -> Int -> t a
completeTree t a
g0 t
depth Int
arity = t a -> a -> t -> t a
completeTree' (t a -> a -> t a
forall (t :: * -> *) a. Mutable t => t a -> a -> t a
addVertex (t a -> t a
forall (t :: * -> *) a. UndirectedGraph t => t a -> t a
empty t a
g0) a
0) a
0 t
1
  where
    completeTree' :: t a -> a -> t -> t a
completeTree' t a
g a
root t
d
      | t
d t -> t -> Bool
forall a. Ord a => a -> a -> Bool
> t
depth = t a
g
      | Bool
otherwise = [t a] -> t a
forall a. [a] -> a
head ([t a] -> t a) -> [t a] -> t a
forall a b. (a -> b) -> a -> b
$ Int -> [t a] -> [t a]
forall a. Int -> [a] -> [a]
drop Int
arity ([t a] -> [t a]) -> [t a] -> [t a]
forall a b. (a -> b) -> a -> b
$
                    (t a -> t a) -> t a -> [t a]
forall a. (a -> a) -> a -> [a]
iterate' (\t a
h -> let r1 :: a
r1 = t a -> a
forall (t :: * -> *) b a.
(UndirectedGraph t, Integral b) =>
t a -> b
numVertices t a
h
                                   in t a -> (a, a) -> t a
forall (t :: * -> *) a. Mutable t => t a -> (a, a) -> t a
addEdge (t a -> a -> t -> t a
completeTree' (t a -> a -> t a
forall (t :: * -> *) a. Mutable t => t a -> a -> t a
addVertex t a
h a
r1) a
r1 (t
dt -> t -> t
forall a. Num a => a -> a -> a
+t
1) ) (a
root,a
r1)) t a
g

completeGraph :: t a -> a -> t a
completeGraph t a
g0 a
k = ((a, a) -> t a -> t a) -> t a -> [(a, a)] -> t a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((t a -> (a, a) -> t a) -> (a, a) -> t a -> t a
forall a b c. (a -> b -> c) -> b -> a -> c
flip t a -> (a, a) -> t a
forall (t :: * -> *) a. Mutable t => t a -> (a, a) -> t a
addEdge) ((a -> t a -> t a) -> t a -> [a] -> t a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((t a -> a -> t a) -> a -> t a -> t a
forall a b c. (a -> b -> c) -> b -> a -> c
flip t a -> a -> t a
forall (t :: * -> *) a. Mutable t => t a -> a -> t a
addVertex) (t a -> t a
forall (t :: * -> *) a. UndirectedGraph t => t a -> t a
empty t a
g0) [a
0..a
ka -> a -> a
forall a. Num a => a -> a -> a
-a
1]) [(a
u,a
v) | a
u <- [a
0..a
ka -> a -> a
forall a. Num a => a -> a -> a
-a
1], a
v <- [a
ua -> a -> a
forall a. Num a => a -> a -> a
+a
1..a
ka -> a -> a
forall a. Num a => a -> a -> a
-a
1] ]

randomGraph :: t t -> t -> t -> StateT g Identity (t t)
randomGraph t t
g0 t
n t
m
  | t
m t -> t -> Bool
forall a. Ord a => a -> a -> Bool
> (t
n t -> t -> t
forall a. Num a => a -> a -> a
* (t
n t -> t -> t
forall a. Num a => a -> a -> a
- t
1)) t -> t -> t
forall a. Integral a => a -> a -> a
`div` t
4 = do -- dense graph
    let g1 :: t t
g1 = ((t, t) -> t t -> t t) -> t t -> [(t, t)] -> t t
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((t t -> (t, t) -> t t) -> (t, t) -> t t -> t t
forall a b c. (a -> b -> c) -> b -> a -> c
flip t t -> (t, t) -> t t
forall (t :: * -> *) a. Mutable t => t a -> (a, a) -> t a
addEdge) ((t -> t t -> t t) -> t t -> [t] -> t t
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((t t -> t -> t t) -> t -> t t -> t t
forall a b c. (a -> b -> c) -> b -> a -> c
flip t t -> t -> t t
forall (t :: * -> *) a. Mutable t => t a -> a -> t a
addVertex) (t t -> t t
forall (t :: * -> *) a. UndirectedGraph t => t a -> t a
empty t t
g0) [t
1..t
n]) [(t
v,t
u) | t
v <- [t
1..t
n], t
u <- [t
vt -> t -> t
forall a. Num a => a -> a -> a
+t
1..t
n]]
    t t -> t -> StateT g Identity (t t)
forall (t :: * -> *) a g t.
(Random a, RandomGen g, Integral t, Integral a, Adjacency t,
 Mutable t) =>
t a -> t -> StateT g Identity (t a)
removeRandomEdges t t
g1 t
m
  | Bool
otherwise = do -- spare graph
    let g1 :: t t
g1 = (t -> t t -> t t) -> t t -> [t] -> t t
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((t t -> t -> t t) -> t -> t t -> t t
forall a b c. (a -> b -> c) -> b -> a -> c
flip t t -> t -> t t
forall (t :: * -> *) a. Mutable t => t a -> a -> t a
addVertex) (t t -> t t
forall (t :: * -> *) a. UndirectedGraph t => t a -> t a
empty t t
g0) [t
0..t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
1]
    t t -> t -> StateT g Identity (t t)
forall (t :: * -> *) a g t.
(UndirectedGraph t, Random a, RandomGen g, Integral t, Integral a,
 Mutable t) =>
t a -> t -> StateT g Identity (t a)
addRandomEdges t t
g1 t
m

addRandomEdges :: t a -> t -> StateT g Identity (t a)
addRandomEdges t a
g t
m
  | t a -> t
forall (t :: * -> *) b a.
(UndirectedGraph t, Integral b) =>
t a -> b
numEdges t a
g t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
m = t a -> StateT g Identity (t a)
forall (m :: * -> *) a. Monad m => a -> m a
return t a
g
  | Bool
otherwise = do
    a
v <- a -> a -> State g a
forall a g. (Random a, RandomGen g) => a -> a -> State g a
randomN a
0 (t a -> a
forall (t :: * -> *) b a.
(UndirectedGraph t, Integral b) =>
t a -> b
numVertices t a
g a -> a -> a
forall a. Num a => a -> a -> a
- a
1)
    a
u <- a -> a -> State g a
forall a g. (Random a, RandomGen g) => a -> a -> State g a
randomN a
0 (t a -> a
forall (t :: * -> *) b a.
(UndirectedGraph t, Integral b) =>
t a -> b
numVertices t a
g a -> a -> a
forall a. Num a => a -> a -> a
- a
1)
    if a
u a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
v then
      t a -> t -> StateT g Identity (t a)
addRandomEdges (t a -> (a, a) -> t a
forall (t :: * -> *) a. Mutable t => t a -> (a, a) -> t a
addEdge t a
g (a
v,a
u)) t
m
    else
      t a -> t -> StateT g Identity (t a)
addRandomEdges t a
g t
m

removeRandomEdges :: t a -> t -> StateT g Identity (t a)
removeRandomEdges t a
g t
m
  | t a -> t
forall (t :: * -> *) b a.
(UndirectedGraph t, Integral b) =>
t a -> b
numEdges t a
g t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
m = t a -> StateT g Identity (t a)
forall (m :: * -> *) a. Monad m => a -> m a
return t a
g
  | Bool
otherwise = do
    a
v <- a -> a -> State g a
forall a g. (Random a, RandomGen g) => a -> a -> State g a
randomN a
0 (t a -> a
forall (t :: * -> *) b a.
(UndirectedGraph t, Integral b) =>
t a -> b
numVertices t a
g a -> a -> a
forall a. Num a => a -> a -> a
- a
1)
    a
u <- a -> a -> State g a
forall a g. (Random a, RandomGen g) => a -> a -> State g a
randomN a
0 (t a -> a
forall (t :: * -> *) b a.
(UndirectedGraph t, Integral b) =>
t a -> b
numVertices t a
g a -> a -> a
forall a. Num a => a -> a -> a
- a
1)
    if a
u a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
v Bool -> Bool -> Bool
&& t a -> (a, a) -> Bool
forall (t :: * -> *) a. Adjacency t => t a -> (a, a) -> Bool
edgeExists t a
g (a
v,a
u) then
      t a -> t -> StateT g Identity (t a)
removeRandomEdges (t a -> (a, a) -> t a
forall (t :: * -> *) a. Mutable t => t a -> (a, a) -> t a
removeEdge t a
g (a
v,a
u)) t
m
    else
      t a -> t -> StateT g Identity (t a)
removeRandomEdges t a
g t
m

randomN :: (Random a, RandomGen g) => a -> a -> State g a
randomN :: a -> a -> State g a
randomN a
n0 a
n1 = do
  g
gen <- StateT g Identity g
forall s (m :: * -> *). MonadState s m => m s
get
  let (a
r,g
gen') = (a, a) -> g -> (a, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (a
n0,a
n1) g
gen
  g -> StateT g Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put g
gen'
  a -> State g a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r