module TreeGraph (LayoutGraph, 
                  flayoutToGraph, treeLayoutToGraph,
                  -- treeToGraph, 
                  orderedTreeToGraph, 
                  treeGraphNodesTree, graphToTreeOriginal,
                  graphToTreeStructure, 
                  flayoutToGraphRoots,
                  graphToOrderedTree, graphToOrderedTreeFrom,
                  orderedChildren, adjCompareEdge,
                  nextNodes, -- exported for testing; any other reason?
                  grTranslateNode, grTranslateSubtree, grTranslateGraph,

                  -- moved from Workspace.hs:
                  functoidToFunction, graphToExprTree,

                  -- Tree graph rendering
                  graphQuickView, 
                  graphWriteImageFile, graphRender, treeRender,
                  treeWriteImageFile, gtkShowTree
                 )

where

import IO
import Data.IORef
import Data.List (sort, sortBy)
import System.Cmd

import Data.Graph.Inductive as G

import LittleGtk
import Graphics.UI.Gtk.Gdk.EventM

import Graphics.Rendering.Cairo hiding (translate)

import Geometry
import Tree as T
import TreeLayout
import Expr
import Util
import Workspace.Functoid
import Workspace.WGraph

type LayoutGraph n e = Gr (LayoutNode n) e

flayoutToGraph :: FunctoidLayout -> WGraph
flayoutToGraph tlo = 
    case tlo of 
      FLayoutTree t -> treeLayoutToGraph t
      FLayoutForest ts _bbox -> 
          foldl grAddGraph wgraphNew (map treeLayoutToGraph ts)

treeLayoutToGraph :: TreeLayout ExprNode -> WGraph
treeLayoutToGraph = orderedTreeToGraph . fmap WSimple


-- flayoutToGraphRoots returns a list of graph nodes (Ints)
-- corresponding to the root of the tree, or the roots of
-- the trees in the forest.
-- The list is ordered in the same sense as the graph nodes.

flayoutToGraphRoots :: FunctoidLayout -> [G.Node]
flayoutToGraphRoots (FLayoutTree _t) = [1]
flayoutToGraphRoots (FLayoutForest trees _bbox) =
    let loop _ [] res = reverse res
        loop next (t:ts) res =
            loop (next + treeSize t) ts (next:res)
    in loop 1 trees []


-- {-# DEPRECATED treeToGraph "use ??? instead" #-}

-- treeToGraph :: Tree e -> Gr e ()
-- treeToGraph (T.Node root subtrees)  = 
--     -- Convert an ordered tree to a graph.
--     -- The graph edge labels from parent to child 
--     -- are integers corresponding to the order of the children.
--     -- So, to recover the ordered list of children, 
--     -- sort by the edge label.
--     let g0 = empty :: Gr e ()
--         g1 = insNode (1, root) g0

--         grow :: Gr e () -> [(Tree e, G.Node)] -> G.Node -> Gr e ()
--         -- grow graph [(tree, parent)] node
--         grow g [] _ = g
--         grow g ((t, p):tps) n =
--             -- insert t forming g', insert ts into g'
--             let edge = ((), p) -- predecessors, or both?
--                 g' = ([edge], n, (rootLabel t), []) & g -- assume pred only
--             in grow g' (tps ++ [(s, n) | s <- subForest t])
--                        (succ n)

--     in grow g1 [(s, 1) | s <- subtrees] 2

sprout :: G.Node -> Tree e -> [(G.Node, WEdge, Tree e)]
sprout parent (T.Node _ subtrees) =
    -- create triples (parent graph node, edge, subtree)
    let m = length subtrees - 1
    in [(parent, WEdge e, s) | (e, s) <- zip [0..m] subtrees]


orderedTreeToGraph :: Tree e -> Gr e WEdge
orderedTreeToGraph otree  = 
    -- Convert an ordered tree to a graph.
    -- The graph edge labels from parent to child 
    -- are integers corresponding to the order of the children.
    -- So, to recover the ordered list of children, 
    -- sort by the edge label.
    let g0 = empty :: Gr e WEdge
        g1 = insNode (1, rootLabel otree) g0

        grow :: Gr e WEdge -> [(G.Node, WEdge, Tree e)] -> G.Node -> Gr e WEdge
        -- grow graph [(parent, edgeLabel, subtree), ...] node
        grow g [] _ = g
        grow g ((p, e, t):pets) n =
            -- insert pet forming g', insert pets into g';
            -- n is the node id for the root of this subtree
            let adj = (e, p) -- to parent (priority, node)
                g' = ([adj], n, (rootLabel t), []) & g 
                n' = succ n
            in grow g' (pets ++ 
                        sprout n t
                        -- [(s, n) | s <- subForest t])
                        )
                       n'

    in grow g1 (sprout 1 otree) 2

-- And what about this (e, Node) type?  Is it some sort of monad?

treeGraphNodesTree :: Tree e -> Tree Node
treeGraphNodesTree atree =
    -- returns a tree of Nodes of the Graph (treeToGraph atree)
    let gnTree :: Tree e -> Node -> Node -> (Tree Node, Node)
        gnTree (T.Node _root subtrees) rootNode next =
            -- rootNode = Node (number) for the root,
            -- next = next unused Node
            let (nNodes, next') = nextNodes subtrees next
                (subtrees', next'') = gnSubtrees subtrees nNodes next'
            in (T.Node rootNode subtrees', next'')

        gnSubtrees :: [Tree e] -> [Node] -> Node -> ([Tree Node], Node)
        gnSubtrees [] [] next = ([], next)
        gnSubtrees (t:ts) (n:ns) next = 
            let (t', next') = gnTree t n next
                (ts', next'') = gnSubtrees ts ns next'
            in ((t' : ts'), next'')
        gnSubtrees _ _ _ = error "gnSubtrees: list lengths do not match"

    in fst (gnTree atree 1 2)

nextNodes :: [e] -> Node -> ([Node], Node)
nextNodes items next = 
  -- next is the next unused Node (number).
  -- Returns list of new nodes, and a new "next" node
  -- E.g., nextNodes [a, b] 3 = ([3, 4], 5)
  --       nextNodes [c, d, e] 5 = ([5, 6, 7], 8)
  --       nextNodes [] 8 = ([], 8)
  let n = length items
      next' = next + n
  in ([next .. (next' - 1)], next')

-- When a tree is converted to a graph,
-- each tree node's ordered children get graph node numbers
-- in ascending order.  Therefore, when reconstructing the tree,
-- sorting the graph nodes restores the order of the children
-- as in the tree.

graphToOrderedTree :: Gr e WEdge -> Tree e
-- inverse of orderedTreeToGraph
graphToOrderedTree g = graphToOrderedTreeFrom g 1

graphToOrderedTreeFrom :: Gr e WEdge -> G.Node -> Tree e
graphToOrderedTreeFrom g n =
    case lab g n of
      Just label -> 
          T.Node label (map (graphToOrderedTreeFrom g) (orderedChildren g n))
      Nothing -> 
          errcats ["missing label for node", show n]

-- | List of the nodes children, ordered by edge number
orderedChildren :: Gr e WEdge -> G.Node -> [G.Node]
orderedChildren g = map fst . sortBy adjCompareEdge . lsuc g

adjCompareEdge :: (Node, WEdge) -> (Node, WEdge) -> Ordering
adjCompareEdge (_n1, e1) (_n2, e2) = compare e1 e2

{-# DEPRECATED graphToTreeOriginal "use ??? instead" #-}

graphToTreeOriginal :: Gr e () -> G.Node -> Tree e
-- (\g -> graphToTreeOriginal g 1) is the inverse of treeToGraph
graphToTreeOriginal g n =
    case lab g n of
      Just label -> T.Node label (map (graphToTreeOriginal g) 
                                      (sort (suc g n)))
      _ -> errcats ["missing label for node", show n]

{-# DEPRECATED graphToTreeStructure "use ??? instead" #-}

graphToTreeStructure :: Gr n e -> G.Node -> Tree G.Node
-- This is *not* an inverse of treeToGraph.
-- Rather, graphToTreeStructure (treeToGraph t 1) is a tree t' of Nodes
-- (i.e., integer identifiers of nodes in the graph)
-- which parallels the structure of t.

graphToTreeStructure g n = T.Node n (map (graphToTreeStructure g)
                                         (sort (suc g n)))

grTranslateNode :: 
  Node -> Double -> Double -> LayoutGraph n e -> LayoutGraph n e
grTranslateNode node dx dy graph = 
    grUpdateNodeLabel graph node (translate dx dy)

grTranslateSubtree :: 
  Node -> Double -> Double -> LayoutGraph n e -> LayoutGraph n e
grTranslateSubtree root dx dy graph = 
    let trSubtrees :: [Node] -> LayoutGraph n e -> LayoutGraph n e
        trSubtrees [] g = g
        trSubtrees (r:rs) g = trSubtrees (rs ++ suc g r)
                              (grTranslateNode r dx dy g)
    in trSubtrees [root] graph

grTranslateGraph :: Double -> Double -> LayoutGraph n e -> LayoutGraph n e
grTranslateGraph dx dy graph = nmap (translate dx dy) graph

grUpdateNodeLabel :: (DynGraph g) => g a b -> Node -> (a -> a) -> g a b
grUpdateNodeLabel graph node updater =
  case match node graph of
    (Nothing, _) -> error "no such node"
    (Just (preds, jnode, label, succs), graph') ->
        -- jnode == node
        (preds, jnode, updater label, succs) & graph'

functoidToFunction :: 
    Functoid -> WGraph -> G.Node -> Env -> SuccFail Function
functoidToFunction functoid graph frameNode env =
    case functoid of
      FunctoidFunc f -> Succ f
      FunctoidParts {fpName = name, fpArgs = args} ->
          let roots = suc graph frameNode
              expr = treeToExpr $ graphToExprTree graph (head roots)
              impl = Compound args expr
          in 
            -- Check whether the frame contains a single tree
            if length roots /= 1
            then Fail "The graph structure is not a tree!"
            else case decideTypes expr args env of
                   Left errmsg -> Fail errmsg
                   Right (atypes, rtype) -> 
                       Succ (Function (Just name) atypes rtype impl)


graphToExprTree :: WGraph -> G.Node -> Tree ExprNode
graphToExprTree g root =
    let extractExprNode wnode =
            case wnode of
              WSimple layoutNode -> gnodeValue (nodeGNode layoutNode)
              WFrame _ -> error "graphToExprTreeFrom: unexpected WFrame node"
    in fmap extractExprNode (graphToOrderedTreeFrom g root)



-- ============================================================
-- GRAPH VIEWING AND RENDERING
-- ============================================================

-- Quick view of a graph using GraphViz

graphQuickView :: (Graph g, Show a, Show b) => g a b -> IO ()
graphQuickView g = 
    let dot_src = graphviz g "graphQuickView" (6, 4) (1, 1) Portrait
        dot_file = "tmp.dot"
        png_file = "tmp.png"
    in do
      h <- openFile dot_file WriteMode
      hPutStr h dot_src
      hClose h
      _ <- system ("dot -Tpng -o" ++ png_file ++ " " ++ dot_file)
      _ <- system ("feh " ++ png_file)
      return ()


graphWriteImageFile :: (Repr n) => 
                       Style -> Maybe Node -> Maybe Node ->
                       Double -> Double -> 
                       LayoutGraph n e -> String -> 
                       IO String 
graphWriteImageFile style mactive mselected dwidth dheight graph file = do
    withImageSurface FormatARGB32 (round dwidth) (round dheight) $ \surf -> do
      renderWith surf $ 
        graphRender style mactive mselected graph
      surfaceWriteToPNG surf file
    return file

graphRender :: (Repr n) => 
               Style -> Maybe Node -> Maybe Node -> LayoutGraph n e -> 
               Render ()
graphRender style mactive mselected graph = do
  let renderNode :: Node -> Render ()
      renderNode node = do
        let Just layoutNode = lab graph node 
            -- ^^ is this safe? it can't be Nothing?
            nodeBB = gnodeNodeBB (nodeGNode layoutNode)
            active = (mactive == Just node)
            selected = case mselected of
                         Nothing -> False
                         Just sel -> sel == node
            mode = if active then DrawActive
                   else if selected then DrawSelectedNode -- !!!
                        else DrawNormal
            xcenter = bbXCenter nodeBB
        draw style mode layoutNode
        connectParent node xcenter (bbTop nodeBB)
        return ()

      connectParent node x y = 
        let parents = pre graph node in
            case parents of
              [] -> return ()
              [parent] -> do
                let Just playoutNode = lab graph parent
                    parentBB = gnodeNodeBB (nodeGNode playoutNode)
                    px = bbXCenter parentBB
                    py = bbBottom parentBB
                setColor (styleNormalEdgeColor style)
                moveTo px (py + snd (vtinypad style)) -- bottom of parent
                lineTo x (y - fst (vtinypad style))   -- top of node
                stroke

              _ -> error "Too many parents"
        
  setAntialias AntialiasDefault

  -- canvas background
  setColor (styleNormalFillColor style)
  let Just layoutNode = lab graph 1 -- root node, represents whole tree
      BBox x y bwidth bheight = nodeTreeBB layoutNode
  rectangle x y bwidth bheight
  fill

  -- draw the graph/tree
  setLineWidth (lineWidth style)
  mapM_ renderNode (nodes graph)

treeRender :: (Repr e) => Style -> TreeLayout e -> Render ()
treeRender style = graphRender style Nothing Nothing . orderedTreeToGraph

treeWriteImageFile :: (Repr e) => 
  Style -> IoletCounter e -> Tree e -> String -> IO String

treeWriteImageFile style counter atree filename = do
  let tlo = treeLayout style counter atree
      Size surfWidth surfHeight = treeLayoutPaddedSize style tlo
  withImageSurface FormatARGB32 (round surfWidth) (round surfHeight) $ 
    \ surf -> do
      renderWith surf $ treeRender style tlo
      surfaceWriteToPNG surf filename
  return filename

-- ============================================================
-- Simple tree viewing.
-- gtkShowTree displays a single tree very simply
-- Works for any kind of (Repr e, Show e) => Tree e.

gtkShowTree :: (Repr e, Show e) => 
  Style -> IoletCounter e -> Tree e -> IO ()
gtkShowTree style counter atree = do
  let tlo = treeLayout style counter atree
      Size dwidth dheight = treeLayoutPaddedSize style tlo

  tloRef <- newIORef tlo

  _ <- initGUI

  -- init window
  window <- windowNew
  set window [windowTitle := "Test Cairo Tree"]
  _ <- onDestroy window mainQuit

  -- init vbox
  vbox <- vBoxNew False 5 -- width not homogeneous; spacing
  set window [containerChild := vbox]

  -- init canvas
  canvas <- layoutNew Nothing Nothing
  _ <- onSizeRequest canvas 
       (return (Requisition (round dwidth) (round dheight)))
  widgetSetCanFocus canvas True -- to receive key events

  -- event handlers
  _ <- on canvas exposeEvent (updateCanvas style canvas tloRef)
  _ <- on canvas keyPressEvent (keyPress window)

  -- pack, show, and run
  boxPackStartDefaults vbox canvas
  widgetShowAll window
  mainGUI

updateCanvas :: (Repr e) => Style -> Layout -> IORef (TreeLayout e) 
             -> EventM EExpose Bool
updateCanvas style canvas tloRef = 
    tryEvent $ liftIO $ do
      {       
        tlo <- readIORef tloRef
      ; win <- layoutGetDrawWindow canvas
      ; renderWithDrawable win (treeRender style tlo)
      }

keyPress :: Window -> EventM EKey Bool
keyPress window =
    tryEvent $ do
      {
        kname <- eventKeyName
      ; case kname of
          "q" -> liftIO $ widgetDestroy window  -- implies mainQuit
          _ -> stopEvent
      }