{-|
Module      : BCG
Description : bi-coloured graphs
Copyright   : (c) Sebastian Tee, 2022
License     : MIT
Maintainer  : github.com/SebTee

Library for handling bi-coloured graphs
-}
module BCG where

import Data.Complex
import qualified Data.IntSet as IntSet
import qualified Data.Set as Set

-- * Data Types

-- | Bi-Colored Graph
data BCG
    = BCG
    [BCE]      -- ^ List of 'BCE's
    IntSet.IntSet     -- ^ Set of Vertices
    Int        -- ^ Number of different colours
    deriving Int -> BCG -> ShowS
[BCG] -> ShowS
BCG -> String
(Int -> BCG -> ShowS)
-> (BCG -> String) -> ([BCG] -> ShowS) -> Show BCG
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [BCG] -> ShowS
$cshowList :: [BCG] -> ShowS
show :: BCG -> String
$cshow :: BCG -> String
showsPrec :: Int -> BCG -> ShowS
$cshowsPrec :: Int -> BCG -> ShowS
Show

-- | Bi-Coloured Edge
data BCE = BCE { BCE -> Int
fromV :: Int            -- ^ From Vertex
               , BCE -> Int
fromC :: Int            -- ^ From Colour
               , BCE -> Int
toV   :: Int            -- ^ To Vertex
               , BCE -> Int
toC   :: Int            -- ^ To Colour
               , BCE -> Complex Double
w     :: Complex Double -- ^ Weight
} deriving Int -> BCE -> ShowS
[BCE] -> ShowS
BCE -> String
(Int -> BCE -> ShowS)
-> (BCE -> String) -> ([BCE] -> ShowS) -> Show BCE
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [BCE] -> ShowS
$cshowList :: [BCE] -> ShowS
show :: BCE -> String
$cshow :: BCE -> String
showsPrec :: Int -> BCE -> ShowS
$cshowsPrec :: Int -> BCE -> ShowS
Show

-- | Inherited Vertex Coloring
type IVC = Set.Set (Int, Int)

-- | Perfect Matching
type PM = BCG

-- * Constructors

-- | A 'BCG' with no edges or vertices and the number of different colors is 0
empty :: BCG
empty :: BCG
empty = [BCE] -> IntSet -> Int -> BCG
BCG [] IntSet
IntSet.empty Int
0

-- * Functions

-- | Return a value between 0 and 1 representing a 'BCG'\'s distance to the monochromatic
-- \[dist(G) = \begin{cases} 
--     0 & G \text{ has no perfect matchings} \\
--     \frac{\left|\displaystyle\sum_{c_i \in G,\ c_i\ \text{is monochromatic}} w(c_i)\right|^2}{d\cdot\left(\displaystyle\sum_{c_i \in G} |w(c_i)|^2\right)} & \text{Otherwise}
-- \end{cases}\]
-- Where \(c_i\) is an 'IVC' on \(G\), \(d\) is the number of different colors on \(G\) and \(w(c_i)\) is the colouring weight of \(c_i\).
dist :: BCG -> Double
dist :: BCG -> Double
dist BCG
g
    | Double
norm Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 = Double
0 -- return 0 in the case of no perfect matchings
    | Bool
otherwise = case BCG
g of
        (BCG [BCE]
_ IntSet
_ Int
d) -> (Complex Double -> Double
forall a. RealFloat a => Complex a -> a
magnitude ([Complex Double] -> Complex Double
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Complex Double] -> Complex Double)
-> [Complex Double] -> Complex Double
forall a b. (a -> b) -> a -> b
$ ([BCG] -> Complex Double) -> [[BCG]] -> [Complex Double]
forall a b. (a -> b) -> [a] -> [b]
map [BCG] -> Complex Double
coloringWeight [[BCG]]
monochromaticIvcs) Double -> Integer -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
2) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
d Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
norm)
    where
        pms :: [BCG]
pms = BCG -> [BCG]
enumeratePM BCG
g
        ivcs :: [[BCG]]
ivcs = [BCG] -> [[BCG]]
groupByIvc [BCG]
pms
        monochromaticIvcs :: [[BCG]]
monochromaticIvcs = ([BCG] -> Bool) -> [[BCG]] -> [[BCG]]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(BCG
pm:[BCG]
_) -> BCG -> Bool
isMonochromatic BCG
pm) [[BCG]]
ivcs
        norm :: Double
norm = [Double] -> Double
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Double] -> Double) -> [Double] -> Double
forall a b. (a -> b) -> a -> b
$ ([BCG] -> Double) -> [[BCG]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\[BCG]
c -> Complex Double -> Double
forall a. RealFloat a => Complex a -> a
magnitude ([BCG] -> Complex Double
coloringWeight [BCG]
c) Double -> Integer -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
2) [[BCG]]
ivcs

-- | Get a list of perfect matchings on a 'BCG'
-- The perfect matchings are represented by the 'BCG' data type
enumeratePM :: BCG -> [BCG]
enumeratePM :: BCG -> [BCG]
enumeratePM (BCG (BCE
e:[BCE]
es) IntSet
vs Int
c) = (if Bool
isPM then [[BCE] -> IntSet -> Int -> BCG
BCG [BCE
e] IntSet
vs Int
c] else (BCG -> BCG) -> [BCG] -> [BCG]
forall a b. (a -> b) -> [a] -> [b]
map BCG -> BCG
addMatchingBack ([BCG] -> [BCG]) -> [BCG] -> [BCG]
forall a b. (a -> b) -> a -> b
$ BCG -> [BCG]
enumeratePM BCG
newSubG ) [BCG] -> [BCG] -> [BCG]
forall a. [a] -> [a] -> [a]
++ BCG -> [BCG]
enumeratePM ([BCE] -> IntSet -> Int -> BCG
BCG [BCE]
es IntSet
vs Int
c)
    where
        isPM :: Bool
isPM = case BCE
e of BCE{fromV :: BCE -> Int
fromV = Int
from, toV :: BCE -> Int
toV = Int
to} -> [Int] -> IntSet
IntSet.fromList [Int
from, Int
to] IntSet -> IntSet -> Bool
forall a. Eq a => a -> a -> Bool
== IntSet
vs
        newSubG :: BCG
newSubG = BCG -> BCE -> BCG
removeEdgeAndIncidentVertices ([BCE] -> IntSet -> Int -> BCG
BCG (BCE
eBCE -> [BCE] -> [BCE]
forall a. a -> [a] -> [a]
:[BCE]
es) IntSet
vs Int
c) BCE
e
        addMatchingBack :: BCG -> BCG
addMatchingBack (BCG [BCE]
mes IntSet
mvs Int
_) = case BCE
e of BCE{fromV :: BCE -> Int
fromV = Int
from, toV :: BCE -> Int
toV = Int
to} -> [BCE] -> IntSet -> Int -> BCG
BCG (BCE
eBCE -> [BCE] -> [BCE]
forall a. a -> [a] -> [a]
:[BCE]
mes) (Int -> IntSet -> IntSet
IntSet.insert Int
from (IntSet -> IntSet) -> IntSet -> IntSet
forall a b. (a -> b) -> a -> b
$ Int -> IntSet -> IntSet
IntSet.insert Int
to IntSet
mvs) Int
c
enumeratePM BCG
_ = []

-- | Get a list of 'BCE's connected to a vertex
getIncidentEdges :: BCG -> Int -> [BCE]
getIncidentEdges :: BCG -> Int -> [BCE]
getIncidentEdges (BCG [BCE]
es IntSet
_ Int
_) Int
v = (BCE -> Bool) -> [BCE] -> [BCE]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> BCE -> Bool
isIncident Int
v) [BCE]
es

-- | Determines if a vertex is connected to an 'BCE'
isIncident :: Int -> BCE -> Bool
isIncident :: Int -> BCE -> Bool
isIncident Int
v BCE{fromV :: BCE -> Int
fromV = Int
from, toV :: BCE -> Int
toV = Int
to} = (Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
from) Bool -> Bool -> Bool
|| (Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
to)

-- | Removes a 'BCE' from a 'BCG' and the vertices it was incident to
removeEdgeAndIncidentVertices :: BCG -> BCE -> BCG
removeEdgeAndIncidentVertices :: BCG -> BCE -> BCG
removeEdgeAndIncidentVertices BCG
g BCE{fromV :: BCE -> Int
fromV = Int
from, toV :: BCE -> Int
toV = Int
to} = BCG -> [Int] -> BCG
removeVertices BCG
g [Int
from, Int
to]

-- | Removes a vertex and the connected 'BCE's from a 'BCG'
removeVertex :: BCG -> Int -> BCG
removeVertex :: BCG -> Int -> BCG
removeVertex (BCG [BCE]
es IntSet
vs Int
c) Int
v = [BCE] -> IntSet -> Int -> BCG
BCG ((BCE -> Bool) -> [BCE] -> [BCE]
forall a. (a -> Bool) -> [a] -> [a]
filter (\BCE
e -> Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Int
v Int -> BCE -> Bool
`isIncident` BCE
e) [BCE]
es) (Int -> IntSet -> IntSet
IntSet.delete Int
v IntSet
vs) Int
c

-- | Removes a list of vertices and their connected 'BCE's from a 'BCG'
removeVertices :: BCG -> [Int] -> BCG
removeVertices :: BCG -> [Int] -> BCG
removeVertices = (BCG -> Int -> BCG) -> BCG -> [Int] -> BCG
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl BCG -> Int -> BCG
removeVertex

-- | Checks if all edges on a 'PM' are monochromatic and have the same color
isMonochromatic :: PM -> Bool
isMonochromatic :: BCG -> Bool
isMonochromatic (BCG [BCE]
es IntSet
_ Int
_) = (Int -> Bool -> Bool) -> Bool -> [Int] -> Bool
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Bool -> Bool -> Bool
(&&) (Bool -> Bool -> Bool) -> (Int -> Bool) -> Int -> Bool -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (\Int
c -> Int
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> a
head [Int]
colors)) Bool
True ([Int] -> Bool) -> [Int] -> Bool
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int]
forall a. [a] -> [a]
tail [Int]
colors
    where
        getColors :: BCE -> [Int]
getColors BCE{fromC :: BCE -> Int
fromC = Int
from, toC :: BCE -> Int
toC = Int
to} = [Int
from, Int
to]
        colors :: [Int]
colors = (BCE -> [Int]) -> [BCE] -> [Int]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap BCE -> [Int]
getColors [BCE]
es

-- | Takes a list of 'PM's with the same 'IVC' and returns the 'IVC'\'s weight
coloringWeight :: [PM] -> Complex Double
coloringWeight :: [BCG] -> Complex Double
coloringWeight ((BCG [BCE]
es IntSet
_ Int
_):[BCG]
gs) = [BCE] -> Complex Double
edgeWeightProduct [BCE]
es Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
+ [BCG] -> Complex Double
coloringWeight [BCG]
gs
    where
        edgeWeightProduct :: [BCE] -> Complex Double
edgeWeightProduct (BCE{w :: BCE -> Complex Double
w = Complex Double
w}:[BCE]
es') = Complex Double
w Complex Double -> Complex Double -> Complex Double
forall a. Num a => a -> a -> a
* [BCE] -> Complex Double
edgeWeightProduct [BCE]
es'
        edgeWeightProduct [BCE]
_ = Complex Double
1
coloringWeight [BCG]
_ = Complex Double
0

-- | Get the 'IVC' of a 'PM'
pmsIvc :: PM -> IVC
pmsIvc :: BCG -> IVC
pmsIvc (BCG [BCE]
es IntSet
_ Int
_) = [(Int, Int)] -> IVC
forall a. Ord a => [a] -> Set a
Set.fromList ([(Int, Int)] -> IVC) -> [(Int, Int)] -> IVC
forall a b. (a -> b) -> a -> b
$ (BCE -> [(Int, Int)]) -> [BCE] -> [(Int, Int)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\BCE{fromV :: BCE -> Int
fromV = Int
fv, fromC :: BCE -> Int
fromC = Int
fc, toV :: BCE -> Int
toV = Int
tv, toC :: BCE -> Int
toC = Int
tc} -> [(Int
fv, Int
fc), (Int
tv, Int
tc)]) [BCE]
es

-- | Groups a list of 'PM'\'s by 'IVC'
groupByIvc :: [PM] -> [[PM]]
groupByIvc :: [BCG] -> [[BCG]]
groupByIvc (BCG
p:[BCG]
ps) = (BCG
p BCG -> [BCG] -> [BCG]
forall a. a -> [a] -> [a]
: [BCG]
sameIvc) [BCG] -> [[BCG]] -> [[BCG]]
forall a. a -> [a] -> [a]
: [BCG] -> [[BCG]]
groupByIvc [BCG]
diffIvc
    where
        ivc :: IVC
ivc = BCG -> IVC
pmsIvc BCG
p
        hasSameIvc :: BCG -> Bool
hasSameIvc BCG
p' = BCG -> IVC
pmsIvc BCG
p' IVC -> IVC -> Bool
forall a. Eq a => a -> a -> Bool
== IVC
ivc
        sameIvc :: [BCG]
sameIvc = (BCG -> Bool) -> [BCG] -> [BCG]
forall a. (a -> Bool) -> [a] -> [a]
filter BCG -> Bool
hasSameIvc [BCG]
ps
        diffIvc :: [BCG]
diffIvc = (BCG -> Bool) -> [BCG] -> [BCG]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (BCG -> Bool) -> BCG -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BCG -> Bool
hasSameIvc) [BCG]
ps
groupByIvc [BCG]
_ = []

-- | Check if 2 vertices are adjacent
isAdjacent :: BCG -> Int -> Int -> Bool
isAdjacent :: BCG -> Int -> Int -> Bool
isAdjacent BCG
g Int
from Int
to = Int -> IntSet -> Bool
IntSet.member Int
to (IntSet -> Bool) -> IntSet -> Bool
forall a b. (a -> b) -> a -> b
$ BCG -> Int -> IntSet
getAdjacent BCG
g Int
from

-- | Get a set of adjacent nodes
getAdjacent :: BCG -> Int -> IntSet.IntSet
getAdjacent :: BCG -> Int -> IntSet
getAdjacent (BCG [BCE]
es IntSet
_ Int
_) Int
v = Int -> IntSet -> IntSet
IntSet.delete Int
v (IntSet -> IntSet) -> IntSet -> IntSet
forall a b. (a -> b) -> a -> b
$ [Int] -> IntSet
IntSet.fromList ([Int] -> IntSet) -> [Int] -> IntSet
forall a b. (a -> b) -> a -> b
$ [[Int]] -> [Int]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Int]] -> [Int]) -> [[Int]] -> [Int]
forall a b. (a -> b) -> a -> b
$ ([Int] -> Bool) -> [[Int]] -> [[Int]]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem Int
v) ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ (BCE -> [Int]) -> [BCE] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map (\BCE{fromV :: BCE -> Int
fromV = Int
fv, toV :: BCE -> Int
toV = Int
tv} -> [Int
fv, Int
tv]) [BCE]
es

-- | get list of vertices connected to a vertex on a 'BCG'
getConnected :: BCG -> Int -> IntSet.IntSet
getConnected :: BCG -> Int -> IntSet
getConnected BCG
g Int
v = Int -> IntSet -> IntSet
IntSet.delete Int
v (IntSet -> IntSet) -> IntSet -> IntSet
forall a b. (a -> b) -> a -> b
$ IntSet -> BCG -> Int -> IntSet
gc (Int -> IntSet
IntSet.singleton Int
v) BCG
g Int
v
    where
        gc :: IntSet.IntSet -> BCG -> Int -> IntSet.IntSet
        gc :: IntSet -> BCG -> Int -> IntSet
gc IntSet
r BCG
g Int
v = (Int -> IntSet -> IntSet) -> IntSet -> [Int] -> IntSet
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (IntSet -> IntSet -> IntSet
IntSet.union (IntSet -> IntSet -> IntSet)
-> (Int -> IntSet) -> Int -> IntSet -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSet -> BCG -> Int -> IntSet
gc IntSet
nr BCG
g) IntSet
IntSet.empty (IntSet -> [Int]
IntSet.toList IntSet
yetToVisit) IntSet -> IntSet -> IntSet
`IntSet.union` IntSet
nr
            where
                adj :: IntSet
adj = BCG -> Int -> IntSet
getAdjacent BCG
g Int
v
                yetToVisit :: IntSet
yetToVisit = (Int -> Bool) -> IntSet -> IntSet
IntSet.filter (\Int
e -> Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Int -> IntSet -> Bool
IntSet.member Int
e IntSet
r) IntSet
adj
                nr :: IntSet
nr = IntSet
adj IntSet -> IntSet -> IntSet
`IntSet.union` IntSet
r

-- | Determine if a given 'BCG' is a fully connected graph
isConnectedGraph :: BCG -> Bool
isConnectedGraph :: BCG -> Bool
isConnectedGraph BCG
g = case BCG
g of
    BCG [BCE]
_ IntSet
vs Int
_ -> Int -> IntSet -> IntSet
IntSet.delete Int
start IntSet
vs IntSet -> IntSet -> Bool
forall a. Eq a => a -> a -> Bool
== BCG -> Int -> IntSet
getConnected BCG
g Int
start
        where start :: Int
start = [Int] -> Int
forall a. [a] -> a
head ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ IntSet -> [Int]
IntSet.toList IntSet
vs

-- | Counts the number of different colors on a 'BCG'
numberDifferentColors :: BCG -> Int
numberDifferentColors :: BCG -> Int
numberDifferentColors (BCG [BCE]
es IntSet
_ Int
_) = IntSet -> Int
IntSet.size (IntSet -> Int) -> IntSet -> Int
forall a b. (a -> b) -> a -> b
$ [Int] -> IntSet
IntSet.fromList ([Int] -> IntSet) -> [Int] -> IntSet
forall a b. (a -> b) -> a -> b
$ (BCE -> [Int]) -> [BCE] -> [Int]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\BCE{fromC :: BCE -> Int
fromC = Int
f, toC :: BCE -> Int
toC = Int
t} -> [Int
f, Int
t]) [BCE]
es