--
-- Copyright (c) 2009-2010, ERICSSON AB All rights reserved.
-- 
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions are met:
-- 
--     * Redistributions of source code must retain the above copyright notice,
--       this list of conditions and the following disclaimer.
--     * Redistributions in binary form must reproduce the above copyright
--       notice, this list of conditions and the following disclaimer in the
--       documentation and/or other materials provided with the distribution.
--     * Neither the name of the ERICSSON AB nor the names of its contributors
--       may be used to endorse or promote products derived from this software
--       without specific prior written permission.
-- 
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-- AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-- IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-- ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS
-- BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
-- OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-- SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-- INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-- CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-- ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
-- THE POSSIBILITY OF SUCH DAMAGE.
--

-- |Fs2dot is to help us create a visualisation of an algorithm written in
-- Feldspar by converting its graph into dot format -- which can be further
-- processed by the Graphviz suite.
module Feldspar.Fs2dot
  ( fs2dot
  , writeDot
  , DOTSource
  )
    where

import Feldspar.Core.Types
import Feldspar.Core.Graph
import Feldspar.Core.Reify (reify, Program)
import Prelude hiding (id)

{- frontend -}

-- |'fs2dot' takes a Feldspar function as its argument and produces DOT language
-- source.
fs2dot :: (Program prg)
       => prg       -- ^Feldspar function
       -> DOTSource -- ^DOT language source
fs2dot = toDot . fromGraph . makeHierarchical . reify

-- |'writeDot' creates a DOT language format source file. Expected arguments
-- are the desired filename and the Feldspar function to be output in DOT
-- language.
writeDot :: (Program prg)
         => FilePath  -- ^output filename
         -> prg       -- ^Feldspar function
         -> IO ()
writeDot filename prg = writeFile filename $ fs2dot prg

-- |This is for clarity.
type DOTSource = String

{- data types -}

data DGraph =
  DGraph
  { inputs  :: [NodeId]
  , outputs :: [NodeId]
  , nodes   :: [DNode]
  , edges   :: [DEdge]
  }
    deriving (Eq, Show)

data DNode =
  DNode
  { id        :: Int
  , role      :: Function
  , subgraphs :: [DGraph]
  , label     :: String
  }
    deriving (Eq, Show)

data DEdge =
  DEdge
  { start :: DConnector
  , end   :: DConnector
  }
    deriving (Eq, Show)

data DConnector =
    DNodeConn (NodeId, Int)
  | DConstConn PrimitiveData
    deriving (Eq, Show)

{- core -}

fromGraph :: HierarchicalGraph
          -> DGraph
fromGraph graph =
    DGraph
    { inputs = enumerateInputs graph
    , outputs = enumerateOutputs graph
    , nodes = (\(Hierarchy h) -> enumerateNodes h) $ graphHierarchy graph
    , edges = (\(Hierarchy h) -> enumerateEdges h) $ graphHierarchy graph
    }

  where
    enumerateInputs graph = [interfaceInput $ hierGraphInterface graph]
    enumerateOutputs graph = graph
      |> tuple2list . interfaceOutput . hierGraphInterface
      |> map (\(Variable (n, _)) -> n) . filter isVariable

    enumerateNodes = map
      (\(node, hiers) ->
        DNode
        { id = nodeId node
        , role = function node
        , subgraphs = hiers |> map
          (\hier -> DGraph
            { inputs = []
            , outputs = []
            , nodes = (\(Hierarchy h) -> enumerateNodes h) hier
            , edges = []
            }
          )
        , label = (fun2label (function node)
            ++ " (" ++ show (nodeId node) ++ ")") |> subst '"' '\''
        }
      )

    enumerateEdges :: [(Node, [Hierarchy])] -> [DEdge]
    enumerateEdges = concatMap
      (\(node, hiers) ->
        [ DEdge
          { start = DNodeConn (inputnode, 0)
          , end = DNodeConn (nodeId node, 0)
          }
        | inputnode <-
          (tuple2list $ input node)
            |> filter isVariable |> map (\(Variable (n, _)) -> n)
        ] ++
        [ DEdge
          { start = DConstConn (constval)
          , end = DNodeConn (nodeId node, 0)
          }
        | constval <-
          (tuple2list $ input node)
            |> filter (not.isVariable) |> map (\(Constant val) -> val)
        ] ++
        concatMap (\(Hierarchy h) -> enumerateEdges h) hiers
      )

    isVariable src = case src of
      Variable _ -> True
      _          -> False

toDot :: DGraph
      -> DOTSource
toDot graph =
  [ dGraphHead
  , dGraphOptions
  , dGraphNodes graph
  , dGraphEdges graph
  , dGraphOutputs graph
  , dGraphTail
  ] |> unlines
    |> unlines . filter (not.null) . lines

  where
    dGraphHead = "digraph G {"
    dGraphOptions =
      [ "node [shape=box]"
      , "compound=true bgcolor=\"lightgray\""
      , "node [style=filled color=\"black\" fillcolor=\"steelblue\"]"
      , "edge []"
      ] |> unlines

    dGraphNodes graph =
      nodes graph
        |> map
          (\node -> 
            if compound node
            then
              [ "subgraph cluster" ++ show (id node) ++ " {"
              , "label =\"" ++ label node ++ "\""
              , subgraphs node |> map
                  (\subgraph ->
                    [ dGraphNodes subgraph
                    , dGraphEdges subgraph
                    ] |> unlines
                  ) |> unlines
              , "}"
              ] |> unlines
            else
              [ "node" ++ show (id node)
              , "[label=\"" ++ label node ++ "\""
              , "href=\"#node" ++ show (id node) ++ "\"]"
              ] |> unwords
          )
        |> unlines

    dGraphEdges graph =
      zip [1..] (edges graph)
        |> map
          (\(n, edge) ->
            if constEdge edge
            then   "const" ++ show ((\(DNodeConn (i, _)) -> i) $ end edge)
                  ++ "_" ++ show n
                ++ " [label=\""
                  ++ show ((\(DEdge (DConstConn val) _) -> val) edge)
                ++ "\"]\n"
                ++ "const" ++ show ((\(DNodeConn (i, _)) -> i) $ end edge)
                  ++ "_" ++ show n
                ++ " -> "
                ++ "node" ++ show ((\(DNodeConn (i, _)) -> i) $ end edge)
            else   "node" ++ show ((\(DNodeConn (i, _)) -> i) $ start edge)
                ++ " -> "
                ++ "node" ++ show ((\(DNodeConn (i, _)) -> i) $ end edge)
          )
        |> unlines

      where
        label edge = ""
        constEdge edge = case edge of
          DEdge (DConstConn _) _ -> True
          _                      -> False

    dGraphOutputs graph = zip [0 ..] (outputs graph) |> map
      (\(n, opid) ->
        [ "node" ++ show opid ++ " -> output" ++ show n
        , "output" ++ show n ++ " [label=\"Output " ++ show n ++ "\"]"
        ] |> unlines
      ) |> unlines
    dGraphTail = "}"
    compound = \n -> (not.null) $ subgraphs n

fun2label :: Function
          -> String
fun2label (Input)                = "Input"
fun2label (Array sd)             = "Array " ++ (show sd)
fun2label (Function str)         = "Function " ++ (show str)
fun2label (NoInline str ifc)     = "NoInLine " ++ (show str)
fun2label (IfThenElse ifc1 ifc2) = "IfThenElse"
fun2label (While ifc1 ifc2)      = "While"
fun2label (Parallel ifc)       = "Parallel" -- ++ (show i)

{- utility functions -}

tupleCount :: Tuple a -> Int
tupleCount (One a) = 1
tupleCount (Tup as) = sum $ map tupleCount as

tuple2list :: Tuple a -> [a]
tuple2list (One a) = [a]
tuple2list (Tup as) = concatMap tuple2list as

subst :: (Eq a) => a -> a -> [a] -> [a]
subst _ _ [] = []
subst a b (x:xs) = (if a == x then b else x) : subst a b xs

infixl 1 |>
(|>) :: a -> (a -> b) -> b
(|>) x f = f x