packages feed

imp-ppl-0.1.0.0: viz/Viz.hs

module Viz
  ( -- * Simplex visualization
    CredalLayer(..)
  , simplexSVG
    -- * BDD visualization
  , bddSVG
    -- * HTML page
  , VizEntry(..)
  , htmlPage
  , srcBlock
  ) where

import Data.List (sort)
import qualified Data.IntMap.Strict as IntMap
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

import Imp.BDD (BDD(..), BDDNode(..), VarLabel(..), NodeId(..))
import Imp.BDD.Builder (BDDManager, nodeTable)

import Viz.ConvexHull

lookupNodeTable :: NodeId -> BDDManager -> Maybe BDDNode
lookupNodeTable (NodeId n) mgr = IntMap.lookup n (nodeTable mgr)

-- | A layer on the simplex: legend name, colour, and barycentric points.
data CredalLayer = CredalLayer !String !String ![(Double, Double, Double)]

-- | Generate SVG of credal sets on the probability 2-simplex.
--   The three labels name the vertices (bottom-left, bottom-right, top).
simplexSVG :: (String, String, String) -> [CredalLayer] -> String
simplexSVG (label1, label2, label3) layers = unlines $
  [ "<svg xmlns=\"http://www.w3.org/2000/svg\" viewBox=\"0 0 450 420\">"
  , "  <defs><style>"
  , "    .vtx { font: bold 16px sans-serif; }"
  , "    .leg { font: 13px sans-serif; }"
  , "  </style></defs>"
  , "  <rect width=\"100%\" height=\"100%\" fill=\"white\"/>"
  -- simplex outline
  , "  <polygon points=\"" ++ triPts ++ "\""
  , "    fill=\"#f5f5f5\" stroke=\"#bbb\" stroke-width=\"1\"/>"
  ] ++
  -- grid lines at 0.25, 0.5, 0.75
  gridLines ++
  -- credal-set layers (first in list drawn first = behind)
  concatMap drawLayer layers ++
  -- vertex labels
  [ "  <text x=\"" ++ showCoord (fst v1 - 5) ++ "\" y=\"" ++ showCoord (snd v1 + 22)
    ++ "\" text-anchor=\"middle\" class=\"vtx\">" ++ label1 ++ "</text>"
  , "  <text x=\"" ++ showCoord (fst v2 + 5) ++ "\" y=\"" ++ showCoord (snd v2 + 22)
    ++ "\" text-anchor=\"middle\" class=\"vtx\">" ++ label2 ++ "</text>"
  , "  <text x=\"" ++ showCoord (fst v3) ++ "\" y=\"" ++ showCoord (snd v3 - 10)
    ++ "\" text-anchor=\"middle\" class=\"vtx\">" ++ label3 ++ "</text>"
  ] ++
  -- legend
  [ "  <g transform=\"translate(15,15)\">" ] ++
  [ "    <rect x=\"0\" y=\"" ++ show (i*24) ++ "\" width=\"16\" height=\"16\""
    ++ " fill=\"" ++ color ++ "\" fill-opacity=\"0.4\""
    ++ " stroke=\"" ++ color ++ "\" stroke-width=\"1.5\"/>"
    ++ "<text x=\"22\" y=\"" ++ show (i*24+13) ++ "\" class=\"leg\">"
    ++ name ++ "</text>"
  | (i, CredalLayer name color _) <- zip [0::Int ..] layers ] ++
  [ "  </g>"
  , "</svg>"
  ]
  where
    side = 320.0
    cx   = 225.0
    base = 380.0
    top' = base - side * sqrt 3 / 2

    v1 = (cx - side / 2, base)   -- bottom-left
    v2 = (cx + side / 2, base)   -- bottom-right
    v3 = (cx,            top')   -- top

    triPts = showCoord (fst v1) ++ "," ++ showCoord (snd v1) ++ " "
          ++ showCoord (fst v2) ++ "," ++ showCoord (snd v2) ++ " "
          ++ showCoord (fst v3) ++ "," ++ showCoord (snd v3)

    bary (p1, p2, p3) =
      ( p1 * fst v1 + p2 * fst v2 + p3 * fst v3
      , p1 * snd v1 + p2 * snd v2 + p3 * snd v3 )

    gridLines = concatMap (\t ->
        [ gridLine t v1 v2 v3
        , gridLine t v2 v3 v1
        , gridLine t v3 v1 v2
        ]) [0.25, 0.5, 0.75]

    gridLine t (ax,ay) (bx,by) (ex,ey) =
      let x1 = t*ax + (1-t)*bx; y1 = t*ay + (1-t)*by
          x2 = t*ax + (1-t)*ex; y2 = t*ay + (1-t)*ey
      in "  <line x1=\"" ++ showCoord x1 ++ "\" y1=\"" ++ showCoord y1
         ++ "\" x2=\"" ++ showCoord x2 ++ "\" y2=\"" ++ showCoord y2
         ++ "\" stroke=\"#ddd\" stroke-width=\"0.5\"/>"

    drawLayer (CredalLayer _ color pts) =
      let xys     = map bary pts
          outline = convexHull2D xys
          polyPts = unwords [showCoord x ++ "," ++ showCoord y | (x,y) <- outline]
      in [ "  <polygon points=\"" ++ polyPts ++ "\""
           ++ " fill=\"" ++ color ++ "\" fill-opacity=\"0.25\""
           ++ " stroke=\"" ++ color ++ "\" stroke-width=\"2.5\"/>" ]
         ++ [ "  <circle cx=\"" ++ showCoord x ++ "\" cy=\"" ++ showCoord y
              ++ "\" r=\"4\" fill=\"" ++ color ++ "\"/>"
            | (x,y) <- xys ]

-- | Generate an inline SVG of a BDD graph, laid out top-to-bottom.
bddSVG :: BDDManager
       -> Map.Map VarLabel String   -- ^ human-readable variable names
       -> [(String, BDD)]           -- ^ named root BDDs
       -> String
bddSVG mgr varNames roots =
  let -- Collect reachable nodes
      allNids  = collectNids mgr (map snd roots)
      nodeList = [ (nid, node) | nid <- allNids
                                , Just node <- [lookupNodeTable nid mgr] ]

      -- Group by variable, sorted ascending
      byVar = Map.fromListWith (flip (++))
                [(bddVar node, [nid]) | (nid, node) <- nodeList]
      layers = Map.toAscList byVar             -- [(VarLabel, [NodeId])]
      nLayers = length layers

      -- Layout constants
      nodeR   = 18.0
      layerH  = 80.0
      colW    = 90.0
      topPad  = 45.0
      rootH   = 30.0

      -- Determine SVG width from widest row
      maxCols   = maximum $ [length nids | (_, nids) <- layers]
                          ++ [length roots, 2]
      svgW      = max 260 (fromIntegral maxCols * colW + 60)

      -- Center a row of items horizontally, return list of (item, x)
      centerRow :: [a] -> [(a, Double)]
      centerRow []    = []
      centerRow [a]   = [(a, svgW / 2)]
      centerRow items =
        let n  = length items
            tw = fromIntegral (n - 1) * colW
            x0 = (svgW - tw) / 2
        in [(item, x0 + fromIntegral i * colW) | (i, item) <- zip [0 :: Int ..] items]

      -- Root label positions
      rootRow = centerRow roots

      -- Node positions (by layer)
      nodePos :: Map.Map NodeId (Double, Double)
      nodePos = Map.fromList $ concat
        [ [(nid, (x, y))
          | (nid, x) <- centerRow (sort nids)]
        | (li, (_, nids)) <- zip [0::Int ..] layers
        , let y = topPad + rootH + fromIntegral li * layerH
        ]

      -- Terminal row
      termY  = topPad + rootH + fromIntegral nLayers * layerH
      trueX  = svgW / 2 - colW / 2
      falseX = svgW / 2 + colW / 2
      svgH   = termY + 55

      -- Resolve a BDD to its (x,y) position
      posOf :: BDD -> (Double, Double)
      posOf BDDTrue      = (trueX,  termY)
      posOf BDDFalse     = (falseX, termY)
      posOf (BDDRef nid)  = Map.findWithDefault (svgW/2, 0) nid nodePos
      posOf (BDDComp nid) = Map.findWithDefault (svgW/2, 0) nid nodePos

      isComp (BDDComp _) = True
      isComp _           = False

      mkLine x1 y1 x2 y2 dashed comp =
        "  <line x1=\"" ++ showCoord x1 ++ "\" y1=\"" ++ showCoord y1
        ++ "\" x2=\"" ++ showCoord x2 ++ "\" y2=\"" ++ showCoord y2
        ++ "\" stroke=\"" ++ (if comp then "#d32f2f" else if dashed then "#999" else "#444")
        ++ "\" stroke-width=\"" ++ (if dashed then "1" else "1.5")
        ++ "\"" ++ (if dashed then " stroke-dasharray=\"5,3\"" else "") ++ "/>"

      mkCompDot x y =
        "<circle cx=\"" ++ showCoord x ++ "\" cy=\"" ++ showCoord y
        ++ "\" r=\"4\" fill=\"white\" stroke=\"#d32f2f\" stroke-width=\"1.5\"/>"

      mkCircle x y label =
        "  <circle cx=\"" ++ showCoord x ++ "\" cy=\"" ++ showCoord y ++ "\" r=\"" ++ showCoord nodeR
        ++ "\" fill=\"white\" stroke=\"#333\" stroke-width=\"1.5\"/>"
        ++ "<text x=\"" ++ showCoord x ++ "\" y=\"" ++ showCoord (y + 5)
        ++ "\" text-anchor=\"middle\" font-family=\"sans-serif\" font-size=\"12\">"
        ++ label ++ "</text>"

      mkSquare x y label color =
        "  <rect x=\"" ++ showCoord (x - 14) ++ "\" y=\"" ++ showCoord (y - 14)
        ++ "\" width=\"28\" height=\"28\" rx=\"3\" fill=\""
        ++ color ++ "\" stroke=\"#333\" stroke-width=\"1\"/>"
        ++ "<text x=\"" ++ showCoord x ++ "\" y=\"" ++ showCoord (y + 5)
        ++ "\" text-anchor=\"middle\" font-family=\"sans-serif\""
        ++ " font-size=\"13\" font-weight=\"bold\">"
        ++ label ++ "</text>"

      -- Draw an edge from a source point to a child BDD: the connecting
      -- line plus, when the child is a complemented reference, the little
      -- red complemented-edge dot at the child end.
      edgeFrom x1 y1 child dashed =
        let (x2, y2) = posOf child
            r2 = if child == BDDTrue || child == BDDFalse then 14 else nodeR
        in mkLine x1 y1 x2 (y2 - r2) dashed (isComp child)
           ++ if isComp child then mkCompDot x2 (y2 - r2) else ""

      header = "<svg xmlns=\"http://www.w3.org/2000/svg\""
               ++ " viewBox=\"0 0 " ++ showCoord svgW ++ " " ++ showCoord svgH ++ "\""
               ++ " width=\"" ++ showCoord (min svgW 400) ++ "\">"
               ++ "<rect width=\"100%\" height=\"100%\" fill=\"white\"/>"
      edgeElems = concat
        [ case lookupNodeTable nid mgr of
            Just (BDDNode _ lo hi) ->
              let (x, y) = Map.findWithDefault (0,0) nid nodePos
              in [ edgeFrom x (y + nodeR) lo True     -- lo = dashed (0-branch)
                 , edgeFrom x (y + nodeR) hi False    -- hi = solid  (1-branch)
                 ]
            Nothing -> []
        | nid <- allNids
        ]
      rootEdges =
        [ edgeFrom rx (rootH + 8) bdd False
        | ((_, bdd), (_, rx)) <- zip roots rootRow
        ]
      nodeElems =
        [ mkCircle x y (vName varNames (bddVar node))
        | (nid, node) <- nodeList
        , let (x, y) = Map.findWithDefault (0,0) nid nodePos
        ]
      termElems =
        [ mkSquare trueX  termY "T" "#c8e6c9"
        , mkSquare falseX termY "F" "#ffcdd2"
        ]
      rootLabels =
        [ "  <text x=\"" ++ showCoord rx ++ "\" y=\"" ++ showCoord rootH
          ++ "\" text-anchor=\"middle\" font-family=\"sans-serif\""
          ++ " font-size=\"14\" font-weight=\"bold\">" ++ name ++ "</text>"
        | ((name, _), (_, rx)) <- zip roots rootRow
        ]
      edgeAnnotations = concat
        [ case lookupNodeTable nid mgr of
            Just (BDDNode _ lo hi) ->
              let (x, y)   = Map.findWithDefault (0,0) nid nodePos
                  (lx, _)  = posOf lo
                  (hx, _)  = posOf hi
                  labelY   = y + nodeR + 12
                  loLabelX = x + (lx - x) * 0.3 - 8
                  hiLabelX = x + (hx - x) * 0.3 + 8
                  mkLbl lx' ly lbl clr =
                    "<text x=\"" ++ showCoord lx' ++ "\" y=\"" ++ showCoord ly
                    ++ "\" font-family=\"sans-serif\" font-size=\"10\""
                    ++ " fill=\"" ++ clr ++ "\">" ++ lbl ++ "</text>"
              in [ mkLbl loLabelX labelY "0" "#999"
                 , mkLbl hiLabelX labelY "1" "#444"
                 ]
            Nothing -> []
        | nid <- allNids
        ]

  in unlines $
       [header]
       ++ edgeElems ++ rootEdges
       ++ nodeElems ++ termElems
       ++ rootLabels ++ edgeAnnotations
       ++ ["</svg>"]

-- | An entry for the HTML page: title, source code, and BDD SVG.
data VizEntry = VizEntry
  { veTitle  :: !String
  , veSource :: !String   -- ^ Haskell source snippet
  , veBddSvg :: !String   -- ^ inline SVG
  }

-- | Generate a self-contained HTML page with embedded SVG visualizations.
htmlPage :: String       -- ^ simplex SVG (inline)
         -> [VizEntry]   -- ^ example entries
         -> [String]     -- ^ extra HTML sections appended after examples
         -> String
htmlPage simplex entries extras = unlines $
  [ "<!DOCTYPE html>"
  , "<html lang=\"en\"><head><meta charset=\"utf-8\">"
  , "<title>imp — credal set visualization</title>"
  , "<style>"
  , "  body { font-family: system-ui, sans-serif; max-width: 960px;"
  , "         margin: 0 auto; padding: 24px; background: #fafafa; color: #222; }"
  , "  h1 { font-size: 1.6rem; }"
  , "  h2 { font-size: 1.2rem; margin-top: 2rem; color: #555; }"
  , "  .row { display: flex; gap: 24px; flex-wrap: wrap; justify-content: center;"
  , "         align-items: flex-start; }"
  , "  .card { background: white; border: 1px solid #e0e0e0; border-radius: 8px;"
  , "          padding: 16px; text-align: center; }"
  , "  .card h3 { margin: 0 0 8px; font-size: 1rem; color: #333; }"
  , "  svg { max-width: 100%; height: auto; }"
  , "  .legend { font-size: 0.85rem; color: #666; margin-top: 12px; }"
  , "  pre.src { background: #1e1e2e; color: #cdd6f4; padding: 12px 16px;"
  , "            border-radius: 6px; text-align: left; font-size: 0.82rem;"
  , "            line-height: 1.45; overflow-x: auto; margin: 10px 0 0; }"
  , "  table.api { border-collapse: collapse; width: 100%; margin: 12px 0; font-size: 0.9rem; }"
  , "  table.api th, table.api td { border: 1px solid #ddd; padding: 6px 10px; text-align: left; }"
  , "  table.api th { background: #f0f0f0; }"
  , "  table.api code { background: #eee; padding: 1px 4px; border-radius: 3px;"
  , "                    font-size: 0.85em; }"
  , "</style>"
  , "</head><body>"
  , "<h1>imp &mdash; imprecise probabilistic programming</h1>"
  , "<p>Credal sets and BDD compilation for examples from"
  , "  Liell-Cock &amp; Staton (POPL 2025), Figure 1.</p>"
  , ""
  , "<h2>Credal sets on the probability simplex</h2>"
  , "<div class=\"row\"><div class=\"card\">"
  , simplex
  , "<div class=\"legend\">Dashed grid at 0.25 / 0.5 / 0.75 probability levels.</div>"
  , "</div></div>"
  , ""
  , "<h2>Source &amp; BDD compilation</h2>"
  , "<p class=\"legend\">Solid lines = hi (1) branch &nbsp; Dashed = lo (0) branch"
  , "  &nbsp; <span style=\"color:#d32f2f\">Red &#x25cb;</span> = complemented edge</p>"
  , "<div class=\"row\">"
  ] ++
  concatMap (\(VizEntry title src svg) ->
    [ "<div class=\"card\">"
    , "  <h3>" ++ title ++ "</h3>"
    , "  " ++ srcBlock src
    , svg
    , "</div>"
    ]) entries ++
  [ "</div>" ] ++
  extras ++
  [ "</body></html>" ]

-- | A source snippet as a styled, HTML-escaped @<pre>@ block.
srcBlock :: String -> String
srcBlock src = "<pre class=\"src\">" ++ escapeHtml src ++ "</pre>"

escapeHtml :: String -> String
escapeHtml = concatMap esc
  where
    esc '&' = "&amp;"
    esc '<' = "&lt;"
    esc '>' = "&gt;"
    esc c   = [c]

-- | Render an SVG coordinate: rounded to two decimals.
showCoord :: Double -> String
showCoord x = show (fromIntegral (round (x * 100) :: Int) / 100 :: Double)

vName :: Map.Map VarLabel String -> VarLabel -> String
vName varNames vl = case Map.lookup vl varNames of
  Just n  -> n
  Nothing -> "v" ++ show (unVarLabel vl)

-- | All internal node ids reachable from the given roots, each visited once.
collectNids :: BDDManager -> [BDD] -> [NodeId]
collectNids mgr = Set.toList . foldl' go Set.empty
  where
    go seen BDDTrue       = seen
    go seen BDDFalse      = seen
    go seen (BDDRef nid)  = visit seen nid
    go seen (BDDComp nid) = visit seen nid

    visit seen nid
      | nid `Set.member` seen = seen
      | otherwise = case lookupNodeTable nid mgr of
          Nothing                -> Set.insert nid seen
          Just (BDDNode _ lo hi) -> go (go (Set.insert nid seen) lo) hi