packages feed

imp-ppl-0.1.0.1: 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.List as L
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 . L.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