rolling/src/Main.hs

module Main where

import qualified Control.Monad.IO.Class as Monad
import qualified Control.Monad as Monad
import GHC.Exts (groupWith)
import qualified Graphics.Rendering.Cairo as Cairo
import qualified Graphics.Rendering.Pango as Pango

data P = P Int Int deriving (Eq, Ord)

instance Show P where
  show (P n t) = show n ++ " in " ++ show t

times :: P -> P -> P
P a at `times` P b bt = P (a * b) (at * bt)

combine :: [P] -> P
combine ps =
  let tots = [ t | P _ t <- ps ]
      newTot = foldr lcm 1 tots
      newNums = [ if t == newTot then n else error "???" | P n t <- ps ]
  in P (sum newNums) newTot

newtype Prob a = Prob { unprob :: [(a, P)] } deriving (Eq, Show)

instance Functor Prob where
  fmap f (Prob xs) = Prob [(f x, p) | (x, p) <- xs]

instance Applicative Prob where
  pure x = Prob [(x, P 1 1)]
  Prob fs <*> Prob xs = Prob [(f x, pf `times` px) | (f, pf) <- fs, (x, px) <- xs]

instance Monad Prob where
  Prob xs >>= f =
    Prob [ (y, p `times` p')
         | (x, p) <- xs
         , (y, p') <- unprob (f x)
         ]

die :: Int -> Prob Int
die n = Prob [(d, P 1 n) | d <- [1..n]]

data Result = Success | Partial | Failure deriving (Eq, Show, Ord)

judge :: Int -> Result
judge n
  | n <= 6 = Failure
  | n <= 9 = Partial
  | otherwise = Success

normalize :: (Ord a) => Prob a -> Prob a
normalize (Prob xs) =
  Prob [(fst (head rs), combine (map snd rs)) | rs <- groupWith fst xs ]

roll :: Int -> Prob Result
roll m = normalize $ do
  a <- die 6
  b <- die 6
  return (judge (a + b + m))

modRoll :: Int -> Prob Result
modRoll m = normalize $ do
  a <- die 20
  let r = a + m
  return (if r < 10 then Failure else if r < 16 then Partial else Success)

find :: Result -> Prob Result -> Double
find r (Prob rs) = case [ fromIntegral n / fromIntegral t | (r', P n t) <- rs, r == r' ] of
  [x] -> x
  _   -> 0.0

renderRow :: Prob Result -> Int -> Cairo.Render ()
renderRow results offset = do
  let f = find Failure results
      p = find Partial results
      s = find Success results
      o = 50.0 * fromIntegral offset + 5

  Cairo.rectangle 5 o (f * 290) 40
  Cairo.setSourceRGB 1.0 0.0 0.0
  Cairo.fill

  Cairo.rectangle (5 + f * 290) o (p * 290) 40
  Cairo.setSourceRGB 1.0 1.0 0.0
  Cairo.fill

  Cairo.rectangle (5 + (f + p) * 290) o (s * 290) 40
  Cairo.setSourceRGB 0.0 0.0 1.0
  Cairo.fill

  Cairo.setSourceRGB 0.0 0.0 0.0
  layout <- text ("+" ++ show offset)
  Cairo.moveTo 310.0 (50.0 * fromIntegral offset)
  Pango.showLayout layout

text :: String -> Cairo.Render Pango.PangoLayout
text t = do
  layout <- Pango.createLayout t
  Monad.liftIO $ do
    font <- Pango.fontDescriptionNew
    Pango.fontDescriptionSetFamily font "Fira Sans"
    Pango.fontDescriptionSetSize font 30
    Pango.layoutSetFontDescription layout (Just font)
  return layout


main :: IO ()
main = do
  let results = [ (n, roll n) | n <- [0..5] ]
      width = 400
      height = length results * 50 + 50
  surface <- Cairo.createImageSurface Cairo.FormatRGB24 width height
  Cairo.renderWith surface $ do
    Cairo.setSourceRGB 1.0 1.0 1.0
    Cairo.rectangle 0.0 0.0 (fromIntegral width) (fromIntegral height)
    Cairo.fill

    Monad.forM_ results $ \(n, r) ->
      renderRow r n

    Cairo.setSourceRGB 0.0 0.0 0.0
    layout <- text "PbtA rolls (2d6)"
    Cairo.moveTo 20 (fromIntegral height - 50)
    Pango.showLayout layout
  Cairo.surfaceWriteToPNG surface "output.png"