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

import Control.Monad.State
import Data.List
import HGraph.Undirected
import HGraph.Utils
import qualified Data.Set as S
import System.Random

cycleGraph :: t t -> t -> t t
cycleGraph t t
g0 t
n = ((t, t) -> t t -> t t) -> t t -> [(t, t)] -> t t
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (t, t) -> t t -> t t
forall a. (a, a) -> t a -> t a
forall (t :: * -> *) a. Mutable t => (a, a) -> t a -> t a
addEdge ((t -> t t -> t t) -> t t -> [t] -> t t
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr t -> t t -> t t
forall a. a -> t a -> t a
forall (t :: * -> *) a. Mutable t => a -> t a -> t a
addVertex t t
g0 [t
0..t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
1]) [(t
x, (t
x t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) t -> t -> t
forall a. Integral a => a -> a -> a
`mod` t
n) | t
x <- [t
0..t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
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 a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((a, b), (a, b)) -> t (a, b) -> t (a, b)
forall a. (a, a) -> t a -> t a
forall (t :: * -> *) a. Mutable t => (a, a) -> t a -> t a
addEdge (((a, b) -> t (a, b) -> t (a, b))
-> t (a, b) -> [(a, b)] -> t (a, b)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (a, b) -> t (a, b) -> t (a, b)
forall a. a -> t a -> t a
forall (t :: * -> *) a. Mutable t => a -> t 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 t -> t -> Int -> t t
completeTree t t
g0 t
depth Int
arity = t t -> t -> t -> t t
completeTree' (t -> t t -> t t
forall a. a -> t a -> t a
forall (t :: * -> *) a. Mutable t => a -> t a -> t a
addVertex t
0 (t t -> t t
forall a. t a -> t a
forall (t :: * -> *) a. UndirectedGraph t => t a -> t a
empty t t
g0)) t
0 t
1
  where
    completeTree' :: t t -> t -> t -> t t
completeTree' t t
g t
root t
d
      | t
d t -> t -> Bool
forall a. Ord a => a -> a -> Bool
> t
depth = t t
g
      | Bool
otherwise = [t t] -> t t
forall a. HasCallStack => [a] -> a
head ([t t] -> t t) -> [t t] -> t t
forall a b. (a -> b) -> a -> b
$ Int -> [t t] -> [t t]
forall a. Int -> [a] -> [a]
drop Int
arity ([t t] -> [t t]) -> [t t] -> [t t]
forall a b. (a -> b) -> a -> b
$
                    (t t -> t t) -> t t -> [t t]
forall a. (a -> a) -> a -> [a]
iterate' (\t t
h -> let r1 :: t
r1 = t t -> t
forall b a. Integral b => t a -> b
forall (t :: * -> *) b a.
(UndirectedGraph t, Integral b) =>
t a -> b
numVertices t t
h
                                   in (t, t) -> t t -> t t
forall a. (a, a) -> t a -> t a
forall (t :: * -> *) a. Mutable t => (a, a) -> t a -> t a
addEdge (t
root,t
r1) (t t -> t -> t -> t t
completeTree' (t -> t t -> t t
forall a. a -> t a -> t a
forall (t :: * -> *) a. Mutable t => a -> t a -> t a
addVertex t
r1 t t
h) t
r1 (t
dt -> t -> t
forall a. Num a => a -> a -> a
+t
1) ) ) t t
g

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

randomTree :: t a -> a -> StateT g Identity (t a)
randomTree t a
g0 a
n = do
  t a -> Set a -> Set a -> StateT g Identity (t a)
forall {g} {t :: * -> *} {a}.
(RandomGen g, Mutable t, Ord a) =>
t a -> Set a -> Set a -> StateT g Identity (t a)
randomTree' ((a -> t a -> t a) -> t a -> [a] -> t a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> t a -> t a
forall a. a -> t a -> t a
forall (t :: * -> *) a. Mutable t => a -> t a -> t a
addVertex (t a -> t a
forall a. t a -> t a
forall (t :: * -> *) a. UndirectedGraph t => t a -> t a
empty t a
g0) [a
0..a
na -> a -> a
forall a. Num a => a -> a -> a
-a
1]) (a -> Set a
forall a. a -> Set a
S.singleton a
0) ([a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList [a
1..a
na -> a -> a
forall a. Num a => a -> a -> a
-a
1])
  where
    randomTree' :: t a -> Set a -> Set a -> StateT g Identity (t a)
randomTree' t a
g Set a
added Set a
missing
      | Set a -> Bool
forall a. Set a -> Bool
S.null Set a
missing = t a -> StateT g Identity (t a)
forall a. a -> StateT g Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return t a
g
      | Bool
otherwise = do
        a
v <- (Int -> a) -> StateT g Identity Int -> StateT g Identity a
forall a b. (a -> b) -> StateT g Identity a -> StateT g Identity b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
i -> Int -> Set a -> a
forall a. Int -> Set a -> a
S.elemAt Int
i Set a
missing) (StateT g Identity Int -> StateT g Identity a)
-> StateT g Identity Int -> StateT g Identity a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> StateT g Identity Int
forall a g. (Random a, RandomGen g) => a -> a -> State g a
randomN Int
0 (Set a -> Int
forall a. Set a -> Int
S.size Set a
missing Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        a
u <- (Int -> a) -> StateT g Identity Int -> StateT g Identity a
forall a b. (a -> b) -> StateT g Identity a -> StateT g Identity b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
i -> Int -> Set a -> a
forall a. Int -> Set a -> a
S.elemAt Int
i Set a
added) (StateT g Identity Int -> StateT g Identity a)
-> StateT g Identity Int -> StateT g Identity a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> StateT g Identity Int
forall a g. (Random a, RandomGen g) => a -> a -> State g a
randomN Int
0 (Set a -> Int
forall a. Set a -> Int
S.size Set a
added Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        t a -> Set a -> Set a -> StateT g Identity (t a)
randomTree' ((a, a) -> t a -> t a
forall a. (a, a) -> t a -> t a
forall (t :: * -> *) a. Mutable t => (a, a) -> t a -> t a
addEdge (a
v,a
u) t a
g) (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
S.insert a
v Set a
added) (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
S.delete a
v Set a
missing)

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 a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (t, t) -> t t -> t t
forall a. (a, a) -> t a -> t a
forall (t :: * -> *) a. Mutable t => (a, a) -> t a -> t a
addEdge ((t -> t t -> t t) -> t t -> [t] -> t t
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr t -> t t -> t t
forall a. a -> t a -> t a
forall (t :: * -> *) a. Mutable t => a -> t a -> t a
addVertex (t t -> t t
forall a. t a -> t a
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 -- sparse graph
    let g1 :: t t
g1 = (t -> t t -> t t) -> t t -> [t] -> t t
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr t -> t t -> t t
forall a. a -> t a -> t a
forall (t :: * -> *) a. Mutable t => a -> t a -> t a
addVertex (t t -> t t
forall a. t a -> t a
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 b a. Integral b => t a -> b
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 a. a -> StateT g Identity 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 b a. Integral b => t a -> b
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 b a. Integral b => t a -> b
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 ((a, a) -> t a -> t a
forall a. (a, a) -> t a -> t a
forall (t :: * -> *) a. Mutable t => (a, a) -> t a -> t a
addEdge (a
v,a
u) t a
g) 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 b a. Integral b => t a -> b
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 a. a -> StateT g Identity 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 b a. Integral b => t a -> b
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 b a. Integral b => t a -> b
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 a. 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 ((a, a) -> t a -> t a
forall a. (a, a) -> t a -> t a
forall (t :: * -> *) a. Mutable t => (a, a) -> t a -> t a
removeEdge (a
v,a
u) t a
g) t
m
    else
      t a -> t -> StateT g Identity (t a)
removeRandomEdges t a
g t
m