glance/app/Rendering.hs

{-# LANGUAGE NoMonomorphismRestriction, FlexibleContexts, TypeFamilies, PartialTypeSignatures, ScopedTypeVariables #-}

module Rendering (
  renderDrawing,
  customLayoutParams,
  renderIngSyntaxGraph
) where

import qualified Diagrams.Prelude as DIA
import Diagrams.Prelude(toName, shaftStyle, global, arrowShaft, noTail
                       , arrowTail, arrowHead, scale, r2, bezier3
                       , fromSegments, Angle, P2, V2, Point, Name, ArrowOpts, N
                       , TrailLike, V, height, width, (*^), reflectX, rotate
                       , centerXY, place
                       , roundedRect, dashingG, lwG, lightgreen, lc, centerPoint
                       , moveTo, turn, (@@), unitX, signedAngleBetween, (.-.)
                       , applyAll, names, angleV, rad, (^.), angleBetween, (.>)
                       , connectOutside', connect', with, (%~), lengths, (^+^)
                       , (.~))
import Diagrams.TwoD.GraphViz(mkGraph, getGraph, layoutGraph')

import qualified Data.GraphViz as GV
import qualified Data.GraphViz.Attributes.Complete as GVA
import qualified Data.IntMap as IM
import qualified Data.Map as Map
import qualified Data.Set as Set

import Control.Arrow(first)
import Data.Function(on)
import qualified Data.Graph.Inductive as ING
import Data.Graph.Inductive.PatriciaTree (Gr)
import Data.List(find, minimumBy)
import Data.Maybe(isNothing, fromMaybe, mapMaybe)
import Data.Typeable(Typeable)
import GHC.Stack(HasCallStack)

--import qualified Data.GraphViz.Types
--import Data.GraphViz.Commands

import Icons(colorScheme, iconToDiagram, defaultLineWidth, ColorStyle(..)
            , getPortAngles, TransformParams(..), circleRadius
            , findIconFromName)
import TranslateCore(nodeToIcon)
import Types(EmbedInfo(..), AnnotatedGraph, Edge(..)
            , Drawing(..), NameAndPort(..)
            , SpecialQDiagram, SpecialBackend, SpecialNum, NodeName(..)
            , Port(..), NamedIcon, Icon(..), NodeInfo(..), IconInfo
            , Named(..))

import Util(nodeNameToInt, fromMaybeError, namedToTuple)

-- If the inferred types for these functions becomes unweildy,
-- try using PartialTypeSignitures.

-- CONSTANT
graphvizScaleFactor :: (Fractional a) => a

-- For Neato
graphvizScaleFactor = 0.12

-- For Fdp
--scaleFactor = 0.09

--scaleFactor = 0.04

drawingToGraphvizScaleFactor :: Fractional a => a
-- For Neato, ScaleOverlaps
--drawingToGraphvizScaleFactor = 0.08

-- For Neato, PrismOverlap
drawingToGraphvizScaleFactor = 0.15

-- TODO Refactor with syntaxGraphToFglGraph in TranslateCore
-- TODO Make this work with nested icons now that names are not qualified.
drawingToIconGraph :: Drawing -> Gr NamedIcon (EmbedInfo Edge)
drawingToIconGraph (Drawing nodes edges) =
  mkGraph nodes labeledEdges where
    labeledEdges = fmap makeLabeledEdge edges

    makeLabeledEdge :: Edge -> (NamedIcon, NamedIcon, EmbedInfo Edge)
    makeLabeledEdge e@(Edge _ (NameAndPort n1 _, NameAndPort n2 _))
      = (Named n1 (lookupInNodes n1)
        , Named n2 (lookupInNodes n2)
        , EmbedInfo Nothing e)
      where
        lookupInNodes name = fromMaybeError
                             errorString
                             (lookup name (fmap namedToTuple nodes))
          where
            errorString =
              "syntaxGraphToFglGraph edge connects to non-existent node. Node NodeName ="
              ++ show name ++ " Edge=" ++ show e


bezierShaft :: (V t ~ V2, TrailLike t) =>
  Angle (N t) -> Angle (N t) -> t
bezierShaft angle1 angle2 = fromSegments [bezier3 c1 c2 x] where
  scaleFactor = 0.5
  x = r2 (1,0)
  c1 = rotate angle1 (scale scaleFactor unitX)
  c2 = rotate angle2 (scale scaleFactor unitX) ^+^ x

getArrowOpts :: (RealFloat n, Typeable n) =>
  (Angle n, Angle n)
  -> NameAndPort
  -> (ArrowOpts n, DIA.Colour Double)
getArrowOpts
  (fromAngle, toAngle)
  (NameAndPort (NodeName nodeNum) mPort)
  = (arrowOptions, shaftColor)
  where
    edgeColors = edgeListC colorScheme
    Port portNum = fromMaybe (Port 0) mPort
    namePortHash = mod (portNum + (503 * nodeNum)) (length edgeColors)
    shaftColor = edgeColors !! namePortHash
    arrowOptions =
      arrowHead .~ DIA.noHead
      -- arrowHead .~ DIA.tri
      $ DIA.headStyle %~ DIA.fc shaftColor
      $ arrowTail .~ noTail
      $ arrowShaft .~ bezierShaft fromAngle toAngle
      -- TODO Don't use a magic number for lengths (headLength and tailLength)
      $ lengths .~ global 0.5
      $ with

-- | Given an Edge, return a transformation on Diagrams that will draw a line.
connectMaybePorts :: SpecialBackend b n =>
  (Angle n, Angle n)
  -> EmbedInfo Edge
  -> SpecialQDiagram b n
  -> SpecialQDiagram b n
connectMaybePorts portAngles
  (EmbedInfo embedDir
    (Edge
      _
      (fromNamePort@(NameAndPort name0 mPort1)
      , NameAndPort name1 mPort2)))
  origDia
  -- In order to give arrows a "shadow" effect, draw a thicker semi-transparent
  -- line shaft the same color as the background underneath the normal line
  -- shaft.
  = -- if DIA.location (DIA.lookupName qPort0 origDia) == DIA.location (DIA.lookupName qPort1 origDia
  -- if nameToPoint qPort0 == nameToPoint qPort1
  --   then error "connectMaybePorts: fromNamePort equals toNamePort!"
  case pointsTheSame of
    Nothing -> origDia
    Just True -> origDia
    _ ->
      (connectFunc normalOpts qPort0 qPort1
       . connectFunc arrOptsShadow qPort0 qPort1) origDia
  where
    pointsTheSame = do
      p0 <- nameToPoint qPort0
      p1 <- nameToPoint qPort1
      return $ p0 == p1
    nameToPoint n = case DIA.lookupName n origDia of
      --Nothing -> DIA.r2 (0, 0)--error "Name does not exist!"
      Nothing -> Nothing-- error $ "Name does not exist! name=" <> show n <> "\neInfo=" <> show eInfo
      Just subDia -> Just $ DIA.location subDia


    lineWidth = 2 * defaultLineWidth
    (baseArrOpts, shaftCol) = getArrowOpts portAngles fromNamePort
    -- TODO Use a color from the color scheme for un-embedded shafts.
    shaftCol' = if isNothing embedDir then shaftCol else DIA.lime
    normalOpts = (shaftStyle %~ (lwG lineWidth . lc shaftCol'))
                 baseArrOpts
    arrOptsShadow = (shaftStyle
                     %~ (lwG (1.9 * lineWidth)
                         . DIA.lcA
                          $ DIA.withOpacity (backgroundC colorScheme) 0.5))
                    baseArrOpts
    (connectFunc, qPort0, qPort1) = case (mPort1, mPort2) of
      (Just port0, Just port1) -> (connect', name0 .> port0, name1 .> port1)
      (Nothing, Just port1) -> (connectOutside', toName name0, name1 .> port1)
      (Just port0, Nothing) -> (connectOutside', name0 .> port0, toName name1)
      (_, _) -> (connectOutside', toName name0, toName name1)


-- START addEdges --
nameAndPortToName :: NameAndPort -> Name
nameAndPortToName (NameAndPort name mPort) = case mPort of
  Nothing -> toName name
  Just port -> name .> port

findPortAngles :: SpecialNum n
  => IconInfo -> NamedIcon -> NameAndPort -> [Angle n]
findPortAngles iconInfo (Named nodeName nodeIcon) (NameAndPort diaName mPort)
  = case mPort of
      Nothing -> []
      Just port -> foundAngles where
        mName = if nodeName == diaName then Nothing else Just diaName
        foundAngles = getPortAngles iconInfo nodeIcon port mName

-- TODO Clean up the Angle arithmatic
pickClosestAngle :: SpecialNum n =>
  (Bool, Angle n)
  -> Angle n
  -> Angle n
  -> Angle n
  -> [Angle n]
  -> Angle n
pickClosestAngle (nodeFlip, nodeAngle) emptyCase target shaftAngle angles
  = case angles of
      [] -> emptyCase
      _ -> (-) <$>
        fst (minimumBy (compare `on` snd) $ fmap angleDiff adjustedAngles)
        <*>
        shaftAngle
        where
          adjustedAngles = fmap adjustAngle angles
          angleDiff angle = (angle, angleBetween (angleV target) (angleV angle))

          adjustAngle angle = if nodeFlip
                              then signedAngleBetween
                                   (rotate nodeAngle $ reflectX (angleV angle))
                                   unitX
                              else (+) <$> angle <*> nodeAngle


-- TODO Refactor with pickClosestAngle
smallestAngleDiff :: SpecialNum n =>
  (Bool, Angle n) -> Angle n -> [Angle n] -> n
smallestAngleDiff (nodeFlip, nodeAngle) target angles = case angles of
  [] -> 0
  _ -> minimum $ fmap angleDiff adjustedAngles
    where
      adjustedAngles = fmap adjustAngle angles
      angleDiff angle = angleBetween (angleV target) (angleV angle) ^. rad

      adjustAngle angle = if nodeFlip then
        signedAngleBetween (rotate nodeAngle $ reflectX (angleV angle)) unitX
        else
        (+) <$> angle <*> nodeAngle


lookupNodeAngle ::  Show n =>
  [(NamedIcon, (Bool, Angle n))] -> NamedIcon -> (Bool, Angle n)
lookupNodeAngle rotationMap key
  = fromMaybeError
  ("nodeVector: key not in rotaionMap. key = " ++ show key
    ++ "\n\n rotationMap = " ++ show rotationMap)
  $ lookup key rotationMap

makeEdge :: (HasCallStack, SpecialBackend b n, ING.Graph gr) =>
  String  -- ^ Debugging information
  -> IconInfo
  -> gr NamedIcon (EmbedInfo Edge)
  -> SpecialQDiagram b n
  -> [(NamedIcon, (Bool, Angle n))]
  -> ING.LEdge (EmbedInfo Edge)
  -> SpecialQDiagram b n
  -> SpecialQDiagram b n
makeEdge debugInfo iconInfo graph dia rotationMap
  (node0, node1, edge@(EmbedInfo _ (Edge _ (namePort0, namePort1))))
  = connectMaybePorts portAngles edge
  where
    node0label = fromMaybeError
                 ("makeEdge: node0 is not in graph. node0: " ++ show node0)
                 $ ING.lab graph node0
    node1label = fromMaybeError
                 ("makeEdge: node1 is not in graph. node1: " ++ show node1)
                 $ ING.lab graph node1

    node0Angle = lookupNodeAngle rotationMap node0label
    node1Angle = lookupNodeAngle rotationMap node1label

    diaNodeNamePointMap = names dia
    port0Point = getPortPoint $ nameAndPortToName namePort0
    port1Point = getPortPoint $ nameAndPortToName namePort1
    shaftVector = if port0Point == port1Point
                  then error "makeEdge: points are equal!"
                  else port1Point .-. port0Point
    shaftAngle = signedAngleBetween shaftVector unitX

    icon0PortAngle = pickClosestAngle node0Angle mempty shaftAngle shaftAngle
                     $ findPortAngles iconInfo node0label namePort0

    shaftAnglePlusOneHalf = (+) <$> shaftAngle <*> (1/2 @@ turn)
    icon1PortAngle = pickClosestAngle
                     node1Angle
                     (1/2 @@ turn)
                     shaftAnglePlusOneHalf
                     shaftAngle
                     (findPortAngles iconInfo node1label namePort1)

    getPortPoint n = case foundPoints of
      [point] -> point
      -- (p1:_) -> p1
      _ -> error $ "Multiple points. Debug info: " <> debugInfo
                     <> "\nn: " <> show n
      where
        foundPoints = fromMaybeError
          ( "makeEdge: port not found. Debug info: " <> debugInfo
            <> "\nPort: " ++ show n ++ "\nValid ports: "
            ++ show diaNodeNamePointMap)
          (lookup n diaNodeNamePointMap)

    portAngles = (icon0PortAngle, icon1PortAngle)

-- | addEdges draws the edges underneath the nodes.
addEdges :: (HasCallStack, SpecialBackend b n, ING.Graph gr) =>
  String  -- ^ Debugging information
  -> IconInfo
  -> gr NamedIcon (EmbedInfo Edge)
  -> SpecialQDiagram b n
  -> [(NamedIcon, (Bool, Angle n))]
  -> SpecialQDiagram b n
addEdges debugInfo iconInfo graph dia rotationMap = applyAll connections dia
  where
    connections
      = makeEdge debugInfo iconInfo graph dia rotationMap <$> ING.labEdges graph

-- BEGIN rotateNodes --

-- TODO May want to use a power other than 2 for the edgeAngleDiffs
scoreAngle :: SpecialNum n =>
  Point V2 n
  -> [(Point V2 n, [Angle n])]
  -> Bool
  -> Angle n
  -> n
scoreAngle iconPosition edges reflected angle
  = sum $ (^(2 :: Int)) <$> fmap edgeAngleDiff edges
  where
    edgeAngleDiff (otherNodePosition, portAngles) = angleDiff
      where
        shaftVector = otherNodePosition .-. iconPosition
        shaftAngle = signedAngleBetween shaftVector unitX
        angleDiff = smallestAngleDiff (reflected, angle) shaftAngle portAngles

bestAngleForIcon :: (HasCallStack, SpecialNum n, ING.Graph gr) =>
  IconInfo
  -> Map.Map NamedIcon (Point V2 n)
  -> gr NamedIcon (EmbedInfo Edge)
  -> NamedIcon
  -> Bool
  -> (Angle n, n)
bestAngleForIcon iconInfo positionMap graph key@(Named (NodeName nodeId) _) reflected
  = minimumBy (compare `on` snd)
    ( (\angle -> (angle
                 , scoreAngle iconPosition edges reflected angle))
      <$> fmap (@@ turn) possibleAngles)
  where
    possibleAngles = [0,(1/24)..1]
    -- possibleAngles = [0, 1/2] -- (uncomment this line and comment out the line above to disable rotation)
    iconPosition = positionMap Map.! key
    edges = getPositionAndAngles
            <$> ( fmap getSucEdge (ING.lsuc graph nodeId)
                  <> fmap getPreEdge (ING.lpre graph nodeId))

    getPositionAndAngles (node, nameAndPort)
      = (positionMap Map.! nodeLabel, portAngles)
      where
        nodeLabel = fromMaybeError
                    "getPositionAndAngles: node not found"
                    (ING.lab graph node)
        portAngles = findPortAngles iconInfo key nameAndPort

  -- Edge points from id to otherNode
    getSucEdge (otherNode, EmbedInfo _ edge) = (otherNode, nameAndPort) where
      (nameAndPort, _) = edgeConnection edge

  -- Edge points from otherNode to id
    getPreEdge (otherNode, EmbedInfo _ edge) = (otherNode, nameAndPort) where
      (_, nameAndPort) = edgeConnection edge

findIconRotation :: (HasCallStack, SpecialNum n, ING.Graph gr) =>
  IconInfo
  -> Map.Map NamedIcon (Point V2 n)
  -> gr NamedIcon (EmbedInfo Edge)
  -> NamedIcon
  -> (NamedIcon, (Bool, Angle n))
findIconRotation iconInfo positionMap graph key = (key, (reflected, angle)) where
  -- Smaller scores are better
  (reflectedAngle, reflectedScore) = bestAngleForIcon iconInfo positionMap graph key True
  (nonReflectedAngle, nonReflectedScore)
    = bestAngleForIcon iconInfo positionMap graph key False
  reflected = reflectedScore < nonReflectedScore
  angle = if reflected then reflectedAngle else nonReflectedAngle

rotateNodes :: (HasCallStack, SpecialNum n, ING.Graph gr) =>
  IconInfo
  -> Map.Map NamedIcon (Point V2 n)
  -> gr NamedIcon (EmbedInfo Edge)
  -> [(NamedIcon, (Bool, Angle n))]
rotateNodes iconInfo positionMap graph
  = findIconRotation iconInfo positionMap graph <$> Map.keys positionMap

-- END rotateNodes --

drawLambdaRegions :: forall b . SpecialBackend b Double =>
  IconInfo
  -> [(NamedIcon, SpecialQDiagram b Double)]
  -> SpecialQDiagram b Double
drawLambdaRegions iconInfo placedNodes
  = mconcat $ fmap (drawRegion Set.empty . fst) placedNodes
  where
    findDia :: NodeName -> SpecialQDiagram b Double
    findDia n1
      = maybe mempty snd
        (find (\(Named n2 _, _) -> n1 == n2) placedNodes)

    -- Also draw the region around the icon the lambda is in.
    drawRegion :: Set.Set NodeName -> NamedIcon -> SpecialQDiagram b Double
    drawRegion parentNames icon = case icon of
      Named _ (LambdaIcon _ _ enclosedNames)
        -> regionRect $ findDia <$> Set.toList (parentNames <> enclosedNames)
      Named parentName (NestedApply _ headIcon icons)
        -> mconcat
           $ drawRegion (Set.insert parentName parentNames)
           <$> mapMaybe
           (fmap (findIconFromName iconInfo))
           (headIcon:icons)
      _ -> mempty

    -- TODO Use something better than a rectangle
    regionRect dias
      = moveTo (centerPoint combinedDia)
        $ lc lightgreen (lwG defaultLineWidth contentsRect)
      where
        combinedDia = mconcat dias
        rectPadding = 3 * circleRadius
        contentsRect = dashingG [0.4 * circleRadius, 0.8 * circleRadius] 0
                       $ roundedRect
                       (rectPadding + width combinedDia)
                       (rectPadding + height combinedDia)
                       (3 * circleRadius)

placeNodes :: SpecialBackend b Double =>
  IconInfo
  -> Map.Map NamedIcon (P2 Double)
  -> [(NamedIcon, (Bool, Angle Double))]
  -> [(NamedIcon, SpecialQDiagram b Double)]
placeNodes namedIcons positionMap = fmap placeNode
  where
    placeNode (key@(Named name icon), (reflected, angle))
      = (key, place transformedDia diaPosition)
      where
        origDia = centerXY
                  $ iconToDiagram
                  namedIcons
                  icon
                  (TransformParams name 0 reflected angle)
        transformedDia = centerXY $ rotate angle
                         $ (if reflected then reflectX else id) origDia
        diaPosition = graphvizScaleFactor *^ (positionMap Map.! key)


customLayoutParams :: GV.GraphvizParams n v e () v
customLayoutParams = GV.defaultParams{
  GV.globalAttributes = [
    GV.NodeAttrs [GVA.Shape GVA.BoxShape]
    --GV.NodeAttrs [GVA.Shape GVA.Circle]
    , GV.GraphAttrs
      [
      --GVA.Overlap GVA.KeepOverlaps,
      --GVA.Overlap GVA.ScaleOverlaps,
      GVA.Overlap $ GVA.PrismOverlap (Just 5000),
      GVA.Splines GVA.LineEdges,
      GVA.OverlapScaling 8,
      --GVA.OverlapScaling 4,
      GVA.OverlapShrink True
      ]
    ],
  GV.fmtEdge = const [GV.arrowTo GV.noArrow]
  }

renderIconGraph :: forall b.
  SpecialBackend b Double =>
  String  -- ^ Debugging information
  -> Gr (NodeInfo NamedIcon) (EmbedInfo Edge)
  -> IO (SpecialQDiagram b Double)
renderIconGraph debugInfo fullGraphWithInfo = do
    -- graph = ING.nmap niVal fullGraphWithInfo
  layoutResult <- layoutGraph' layoutParams GVA.Neato parentGraph
  --  layoutResult <- layoutGraph' layoutParams GVA.Fdp graph
  let
    positionMap = fst $ getGraph layoutResult
    rotationMap = rotateNodes iconInfo positionMap parentGraph
    placedNodeList = placeNodes iconInfo positionMap rotationMap
    placedNodes = mconcat $ fmap snd placedNodeList
    edges = addEdges debugInfo iconInfo parentGraph placedNodes rotationMap
    placedRegions = drawLambdaRegions iconInfo placedNodeList
  pure (placedNodes <> edges <> placedRegions)
  where
    parentGraph
      = ING.nmap niVal $ ING.labfilter (isNothing . niParent) fullGraphWithInfo
    fullGraph = ING.nmap niVal fullGraphWithInfo
    iconInfo = IM.fromList
                 $ first nodeNameToInt . namedToTuple . snd
                 <$> ING.labNodes fullGraph

    layoutParams :: GV.GraphvizParams Int NamedIcon e () NamedIcon
    --layoutParams :: GV.GraphvizParams Int l el Int l
    layoutParams = customLayoutParams{
      GV.fmtNode = nodeAttribute
      }
    nodeAttribute :: (Int, NamedIcon) -> [GV.Attribute]
    nodeAttribute (_, Named _ nodeIcon) =
      -- GVA.Width and GVA.Height have a minimum of 0.01
      --[GVA.Width diaWidth, GVA.Height diaHeight]
      [GVA.Width circleDiameter, GVA.Height circleDiameter]
      where
        -- This type annotation (:: SpecialQDiagram b n) requires Scoped Typed
        -- Variables, which only works if the function's
        -- type signiture has "forall b e."
        dia :: SpecialQDiagram b Double
        dia = iconToDiagram
              iconInfo
              nodeIcon
              (TransformParams (NodeName (-1)) 0 False mempty)

        diaWidth = drawingToGraphvizScaleFactor * width dia
        diaHeight = drawingToGraphvizScaleFactor * height dia
        circleDiameter' = max diaWidth diaHeight
        circleDiameter
          = if circleDiameter' <= 0.01
            then error ("circleDiameter too small: " ++ show circleDiameter')
            else circleDiameter'

-- | Given a Drawing, produce a Diagram complete with rotated/flipped icons and
-- lines connecting ports and icons. IO is needed for the GraphViz layout.
renderDrawing :: SpecialBackend b Double
  => String  -- ^ Debugging information
  -> Drawing
  -> IO (SpecialQDiagram b Double)
renderDrawing debugInfo drawing
  = renderIconGraph debugInfo graph
  where
    graph = ING.nmap (NodeInfo Nothing) . drawingToIconGraph $ drawing

renderIngSyntaxGraph :: (HasCallStack, SpecialBackend b Double)
  => String -> AnnotatedGraph Gr -> IO (SpecialQDiagram b Double)
renderIngSyntaxGraph debugInfo gr
  = renderIconGraph debugInfo
    $ ING.nmap (fmap (fmap nodeToIcon)) gr