-- | Defines printer types and functions for generating traceability graphs (as .dot files).
module Language.Drasil.DOT.Print where

import Data.List (intercalate)
import System.IO (Handle, IOMode(WriteMode), openFile, hPutStrLn, hClose)
import System.Directory (setCurrentDirectory)

import Drasil.Database (UID)
import Utils.Drasil (createDirIfMissing)

-- * Types

-- | Type synonym for clarity.
type Colour = String
-- | Type synonym for clarity.
type Label = String

-- | A node family contains a list of 'UID's, their display labels, general subgraph label, and colour.
data NodeFamily = NF {
    -- | Node 'UID's.
    NodeFamily -> [UID]
nodeUIDs :: [UID]
    -- | Display labels for nodes. We use the reference addresses from the 'UID's.
    , NodeFamily -> [Label]
nodeLabels :: [Label]
    -- | Individual subgraph labels. These labels do not show on the
    -- final generated pdf or png files.
    , NodeFamily -> Label
nfLabel :: Label
    -- | Gives the ability to change colours of bubbles within the graph.
    , NodeFamily -> Label
nfColour :: Colour
}

-- | Holds all important and relevant information for generating a traceability graph.
-- Includes nodes, graph edges, and node family information.
data GraphInfo = GI {
    --------------- graph node families -------------------------------
    -- | Assumptions.
    GraphInfo -> NodeFamily
assumpNF :: NodeFamily
    -- | Data definitions.
    , GraphInfo -> NodeFamily
ddNF :: NodeFamily
    -- | General definitions.
    , GraphInfo -> NodeFamily
gdNF :: NodeFamily
    -- | Theory models.
    , GraphInfo -> NodeFamily
tmNF :: NodeFamily
    -- | Instance models.
    , GraphInfo -> NodeFamily
imNF :: NodeFamily
    -- | Requirements (both functional and non-functional).
    , GraphInfo -> NodeFamily
reqNF :: NodeFamily
    -- | Goal statement.
    , GraphInfo -> NodeFamily
gsNF :: NodeFamily
    -- | Changes (both likely and unlikely).
    , GraphInfo -> NodeFamily
chgNF :: NodeFamily

    -------------- graph edges  ---------------------------
    -- | Assumptions dependent on assumptions.
    , GraphInfo -> [(UID, [UID])]
edgesAvsA     :: [(UID, [UID])]
    -- | Definitions, models, requirements, and changes dependent on assumptions.
    , GraphInfo -> [(UID, [UID])]
edgesAvsAll   :: [(UID, [UID])]
    -- | Definitions and models that are dependent on other definitions and models.
    , GraphInfo -> [(UID, [UID])]
edgesRefvsRef :: [(UID, [UID])]
    -- | Goals and requirements dependent on definitions, models, and other requirements.
    , GraphInfo -> [(UID, [UID])]
edgesAllvsR   :: [(UID, [UID])]
    -- | Definitions, models, requirements, goals, and changes that are dependent on one another.
    , GraphInfo -> [(UID, [UID])]
edgesAllvsAll :: [(UID, [UID])]

    -- may need more information regarding ranking & ordering, but for now I'm just keeping it simple
}

-- * Functions
-- ** Main Outputs

-- | Creates the directory for output, gathers all individual graph output functions and calls them.
outputDot :: FilePath -> GraphInfo -> IO ()
outputDot :: Label -> GraphInfo -> IO ()
outputDot Label
outputFilePath GraphInfo
gi = do
    Bool -> Label -> IO ()
createDirIfMissing Bool
False Label
outputFilePath
    Label -> IO ()
setCurrentDirectory Label
outputFilePath
    GraphInfo -> IO ()
mkOutputAvsA GraphInfo
gi
    GraphInfo -> IO ()
mkOutputAvsAll GraphInfo
gi
    GraphInfo -> IO ()
mkOutputRefvsRef GraphInfo
gi
    GraphInfo -> IO ()
mkOutputAllvsR GraphInfo
gi
    GraphInfo -> IO ()
mkOutputAllvsAll GraphInfo
gi

-- | Output function for assumptions dependent on assumptions.
mkOutputAvsA :: GraphInfo -> IO ()
mkOutputAvsA :: GraphInfo -> IO ()
mkOutputAvsA GraphInfo
gi = do
    let labels :: [GraphInfo -> NodeFamily]
labels = [GraphInfo -> NodeFamily
assumpNF]
    GraphInfo
-> Label
-> (GraphInfo -> [(UID, [UID])])
-> [GraphInfo -> NodeFamily]
-> IO ()
mkOutput GraphInfo
gi Label
"avsa" GraphInfo -> [(UID, [UID])]
edgesAvsA [GraphInfo -> NodeFamily]
labels

-- | Output function for definitions, models, requirements, and changes dependent on assumptions.
mkOutputAvsAll :: GraphInfo -> IO ()
mkOutputAvsAll :: GraphInfo -> IO ()
mkOutputAvsAll GraphInfo
gi = do
    let labels :: [GraphInfo -> NodeFamily]
labels = [GraphInfo -> NodeFamily
assumpNF, GraphInfo -> NodeFamily
ddNF, GraphInfo -> NodeFamily
tmNF, GraphInfo -> NodeFamily
gdNF, GraphInfo -> NodeFamily
imNF, GraphInfo -> NodeFamily
reqNF, GraphInfo -> NodeFamily
chgNF]
    GraphInfo
-> Label
-> (GraphInfo -> [(UID, [UID])])
-> [GraphInfo -> NodeFamily]
-> IO ()
mkOutput GraphInfo
gi Label
"avsall" GraphInfo -> [(UID, [UID])]
edgesAvsAll [GraphInfo -> NodeFamily]
labels

-- | Output function for definitions and models that are dependent on other definitions and models.
mkOutputRefvsRef :: GraphInfo -> IO ()
mkOutputRefvsRef :: GraphInfo -> IO ()
mkOutputRefvsRef GraphInfo
gi = do
    let labels :: [GraphInfo -> NodeFamily]
labels = [GraphInfo -> NodeFamily
ddNF, GraphInfo -> NodeFamily
tmNF, GraphInfo -> NodeFamily
gdNF, GraphInfo -> NodeFamily
imNF]
    GraphInfo
-> Label
-> (GraphInfo -> [(UID, [UID])])
-> [GraphInfo -> NodeFamily]
-> IO ()
mkOutput GraphInfo
gi Label
"refvsref" GraphInfo -> [(UID, [UID])]
edgesRefvsRef [GraphInfo -> NodeFamily]
labels

-- | Output function for goals and requirements dependent on definitions, models, and other requirements.
mkOutputAllvsR :: GraphInfo -> IO ()
mkOutputAllvsR :: GraphInfo -> IO ()
mkOutputAllvsR GraphInfo
gi = do
    let labels :: [GraphInfo -> NodeFamily]
labels = [GraphInfo -> NodeFamily
assumpNF, GraphInfo -> NodeFamily
ddNF, GraphInfo -> NodeFamily
tmNF, GraphInfo -> NodeFamily
gdNF, GraphInfo -> NodeFamily
imNF, GraphInfo -> NodeFamily
reqNF, GraphInfo -> NodeFamily
gsNF]
    GraphInfo
-> Label
-> (GraphInfo -> [(UID, [UID])])
-> [GraphInfo -> NodeFamily]
-> IO ()
mkOutput GraphInfo
gi Label
"allvsr" GraphInfo -> [(UID, [UID])]
edgesAllvsR [GraphInfo -> NodeFamily]
labels

-- | Output function for definitions, models, requirements, goals, and changes that are dependent on one another.
mkOutputAllvsAll :: GraphInfo -> IO ()
mkOutputAllvsAll :: GraphInfo -> IO ()
mkOutputAllvsAll GraphInfo
gi = do
    let labels :: [GraphInfo -> NodeFamily]
labels = [GraphInfo -> NodeFamily
assumpNF, GraphInfo -> NodeFamily
ddNF, GraphInfo -> NodeFamily
tmNF, GraphInfo -> NodeFamily
gdNF, GraphInfo -> NodeFamily
imNF, GraphInfo -> NodeFamily
reqNF, GraphInfo -> NodeFamily
gsNF, GraphInfo -> NodeFamily
chgNF]
    GraphInfo
-> Label
-> (GraphInfo -> [(UID, [UID])])
-> [GraphInfo -> NodeFamily]
-> IO ()
mkOutput GraphInfo
gi Label
"allvsall" GraphInfo -> [(UID, [UID])]
edgesAllvsAll [GraphInfo -> NodeFamily]
labels

-- ** Helpers

-- | General output function for making a traceability graph. Takes in the graph information, title, edge generator functions, and node family functions.
mkOutput :: GraphInfo -> String -> (GraphInfo -> [(UID, [UID])]) -> [GraphInfo -> NodeFamily] -> IO ()
mkOutput :: GraphInfo
-> Label
-> (GraphInfo -> [(UID, [UID])])
-> [GraphInfo -> NodeFamily]
-> IO ()
mkOutput GraphInfo
gi Label
ttl GraphInfo -> [(UID, [UID])]
getDirections [GraphInfo -> NodeFamily]
getLabels = do
    Handle
handle <- Label -> IOMode -> IO Handle
openFile (Label
ttl Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
".dot") IOMode
WriteMode
    Handle -> Label -> IO ()
hPutStrLn Handle
handle (Label -> IO ()) -> Label -> IO ()
forall a b. (a -> b) -> a -> b
$ Label
"digraph " Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
ttl Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
" {"
    let labels :: [NodeFamily]
labels = GraphInfo -> [GraphInfo -> NodeFamily] -> [NodeFamily]
filterAndGI GraphInfo
gi [GraphInfo -> NodeFamily]
getLabels
    Handle -> [(UID, [UID])] -> [NodeFamily] -> IO ()
outputSub Handle
handle (GraphInfo -> [(UID, [UID])]
getDirections GraphInfo
gi) [NodeFamily]
labels
    Handle -> Label -> IO ()
hPutStrLn Handle
handle Label
"}"
    Handle -> IO ()
hClose Handle
handle

-- | Graph output helper. Takes in the file handle, edges, and node families.
outputSub :: Handle -> [(UID, [UID])] -> [NodeFamily] -> IO ()
outputSub :: Handle -> [(UID, [UID])] -> [NodeFamily] -> IO ()
outputSub Handle
handle [(UID, [UID])]
edges [NodeFamily]
nodes = do
    ((UID, [UID]) -> IO ()) -> [(UID, [UID])] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Handle -> (UID, [UID]) -> IO ()
mkDirections Handle
handle) [(UID, [UID])]
edges
    Handle -> Label -> IO ()
hPutStrLn Handle
handle Label
"\n"
    (NodeFamily -> IO ()) -> [NodeFamily] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Handle -> NodeFamily -> IO ()
mkNodes Handle
handle) [NodeFamily]
nodes

-- | Prints graph edges (directions) onto a given file handle.
mkDirections :: Handle -> (UID, [UID]) -> IO ()
mkDirections :: Handle -> (UID, [UID]) -> IO ()
mkDirections Handle
handle (UID, [UID])
ls = do
    (Label -> IO ()) -> [Label] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Handle -> Label -> IO ()
hPutStrLn Handle
handle) ([Label] -> IO ()) -> [Label] -> IO ()
forall a b. (a -> b) -> a -> b
$ Label -> [Label] -> [Label]
makeEdgesSub (UID -> Label
forall a. Show a => a -> Label
show (UID -> Label) -> UID -> Label
forall a b. (a -> b) -> a -> b
$ (UID, [UID]) -> UID
forall a b. (a, b) -> a
fst (UID, [UID])
ls) ((Label -> Bool) -> [Label] -> [Label]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Label -> Bool) -> Label -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Label -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) ([Label] -> [Label]) -> [Label] -> [Label]
forall a b. (a -> b) -> a -> b
$ (UID -> Label) -> [UID] -> [Label]
forall a b. (a -> b) -> [a] -> [b]
map UID -> Label
forall a. Show a => a -> Label
show ([UID] -> [Label]) -> [UID] -> [Label]
forall a b. (a -> b) -> a -> b
$ (UID, [UID]) -> [UID]
forall a b. (a, b) -> b
snd (UID, [UID])
ls)
    where
       -- Creates an edge between a type and its dependency (indented for subgraphs)
        makeEdgesSub :: String -> [String] -> [String]
        makeEdgesSub :: Label -> [Label] -> [Label]
makeEdgesSub Label
_ [] = []
        makeEdgesSub Label
nm (Label
c:[Label]
cs) = (Label
"\t" Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label -> Label
filterInvalidChars Label
nm Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
" -> " Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label -> Label
filterInvalidChars Label
c Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
";")Label -> [Label] -> [Label]
forall a. a -> [a] -> [a]
: Label -> [Label] -> [Label]
makeEdgesSub Label
nm [Label]
cs

-- | Prints graph nodes (labels) onto a given file handle.
mkNodes :: Handle -> NodeFamily -> IO ()
mkNodes :: Handle -> NodeFamily -> IO ()
mkNodes Handle
handle NF{nodeUIDs :: NodeFamily -> [UID]
nodeUIDs = [UID]
u, nodeLabels :: NodeFamily -> [Label]
nodeLabels = [Label]
ls, nfLabel :: NodeFamily -> Label
nfLabel = Label
lbl, nfColour :: NodeFamily -> Label
nfColour = Label
col} = do
    ((Label, Label) -> IO ()) -> [(Label, Label)] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Handle -> Label -> IO ()
hPutStrLn Handle
handle (Label -> IO ())
-> ((Label, Label) -> Label) -> (Label, Label) -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Label -> Label -> Label) -> (Label, Label) -> Label
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (Label -> Label -> Label -> Label
makeNodesSub Label
col)) ([(Label, Label)] -> IO ()) -> [(Label, Label)] -> IO ()
forall a b. (a -> b) -> a -> b
$ [Label] -> [Label] -> [(Label, Label)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Label]
ls ([Label] -> [(Label, Label)]) -> [Label] -> [(Label, Label)]
forall a b. (a -> b) -> a -> b
$ (UID -> Label) -> [UID] -> [Label]
forall a b. (a -> b) -> [a] -> [b]
map UID -> Label
forall a. Show a => a -> Label
show [UID]
u
    Handle -> Label -> [UID] -> IO ()
mkSubgraph Handle
handle Label
lbl [UID]
u
    where
        -- Creates a node based on the kind of datatype (indented for subgraphs)
        makeNodesSub :: Colour -> String -> String -> String
        makeNodesSub :: Label -> Label -> Label -> Label
makeNodesSub Label
c Label
l Label
nm  = Label
"\t" Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label -> Label
filterInvalidChars Label
nm Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
"\t[shape=box, color=black, style=filled, fillcolor=" Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
c Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
", label=\"" Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
l Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
"\"];"

-- | Helper that only makes a subgraph if there are elements in the subgraph. Otherwise, it returns nothing.
mkSubgraph :: Handle -> Label -> [UID] -> IO ()
mkSubgraph :: Handle -> Label -> [UID] -> IO ()
mkSubgraph Handle
handle Label
l [UID]
u
    | Label -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Label
l = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
forall a. Monoid a => a
mempty
    | Bool
otherwise = do
             Handle -> Label -> IO ()
hPutStrLn Handle
handle (Label -> IO ()) -> Label -> IO ()
forall a b. (a -> b) -> a -> b
$ Label
"\n\tsubgraph " Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
l Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
" {"
             Handle -> Label -> IO ()
hPutStrLn Handle
handle Label
"\trank=\"same\""
             Handle -> Label -> IO ()
hPutStrLn Handle
handle (Label -> IO ()) -> Label -> IO ()
forall a b. (a -> b) -> a -> b
$ Label
"\t{" Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label -> [Label] -> Label
forall a. [a] -> [[a]] -> [a]
intercalate Label
", " ((UID -> Label) -> [UID] -> [Label]
forall a b. (a -> b) -> [a] -> [b]
map (Label -> Label
filterInvalidChars (Label -> Label) -> (UID -> Label) -> UID -> Label
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UID -> Label
forall a. Show a => a -> Label
show) [UID]
u) Label -> Label -> Label
forall a. [a] -> [a] -> [a]
++ Label
"}"
             Handle -> Label -> IO ()
hPutStrLn Handle
handle Label
"\t}\n"

-- | Gets graph labels.
filterAndGI :: GraphInfo -> [GraphInfo -> NodeFamily] -> [NodeFamily]
filterAndGI :: GraphInfo -> [GraphInfo -> NodeFamily] -> [NodeFamily]
filterAndGI GraphInfo
gi [GraphInfo -> NodeFamily]
toNodes = [NodeFamily]
labels
    where
        labels :: [NodeFamily]
labels = ((GraphInfo -> NodeFamily) -> NodeFamily)
-> [GraphInfo -> NodeFamily] -> [NodeFamily]
forall a b. (a -> b) -> [a] -> [b]
map ((GraphInfo -> NodeFamily) -> GraphInfo -> NodeFamily
forall a b. (a -> b) -> a -> b
$ GraphInfo
gi) [GraphInfo -> NodeFamily]
toNodes

-- | Helper to remove invalid characters.
filterInvalidChars :: String -> String
filterInvalidChars :: Label -> Label
filterInvalidChars = (Char -> Bool) -> Label -> Label
forall a. (a -> Bool) -> [a] -> [a]
filter (Char -> Label -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` Label
invalidChars)
  where
    invalidChars :: Label
invalidChars = Label
"^[]!} (){->,$"