{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE LambdaCase #-}

import Diagrams.Prelude
import Diagrams.Backend.Cairo.CmdLine

main = mainWith $ pad 1.05 . centerXY $ getWordDiagV

getWordDiagV =
    [ vsep 1
        [ anArrow Ne "accum' Nothing mplus" # named "accum"
        , aFork Ne "|||" # named "merge"
        ]
    , vsep 1
        [ aFunction Le "Just" # named "just"
        , anArrow Le "pickWord rng" # named "pickWord"
        , aFunction Le "const ()" # named "unit"
        ] # centerY
        |||
    strutX 1
        |||
    aFunction Ri "const Nothing" # named "nothing" # centerY
    , vsep 1
        [ aFork Ne "if" # named "if"
        , anArrow Ne "oneShot" # named "oneShot"
        ]
    ] # fmap centerX # vsep 0.5

    # connectNodes Ne "oneShot" "if" (1/4) (3/4)
    # connectNodes Le "if" "unit" (3/8) (7/8)
    # connectNodes Le "unit" "pickWord" (1/4) (3/4)
    # connectNodes Le "pickWord" "just" (1/4) (3/4)
    # connectNodes Le "just" "merge" (1/8) (5/8)
    # connectNodes Ri "if" "nothing" (1/8) (5/8)
    # connectNodes Ri "nothing" "merge" (3/8) (7/8)
    # connectNodes Le "merge" "accum" (1/4) (3/4)

data ForkFlag = Le | Ne | Ri

forkColour :: ForkFlag -> Colour Double
forkColour = \case
    Le -> blue
    Ne -> mediumorchid
    Ri -> red

connectNodes :: ForkFlag -> String -> String
                         -> Double -> Double -> Diagram B -> Diagram B
connectNodes ffl nam1 nam2 tur1 tur2 =
    connectPerim' (baseConnOpts ffl)
        nam1 nam2 (tur1 @@ turn) (tur2 @@ turn)

baseConnOpts ffl = with & lengths .~ large
    & shaftStyle %~ lc colo . lw veryThick
    & headStyle %~ fc (darken 0.2 colo)
    where
    colo = forkColour ffl

nodeSkeleton :: Diagram B -> ForkFlag -> String -> Diagram B
nodeSkeleton d ffl s = d # scaleX (max 1 $ fromIntegral (length s) / 7)
        # applyStyle initStyle # lc colo
    <> nudgedText s
        # fc colo
    where
    colo = forkColour ffl

aFunction :: ForkFlag -> String -> Diagram B
aFunction = nodeSkeleton $ roundedRect 1.6 1.2 0.4

anArrow :: ForkFlag -> String -> Diagram B
anArrow = nodeSkeleton $ rect 1.6 1.2

aFork :: ForkFlag -> String -> Diagram B
aFork = nodeSkeleton $ circle 0.5

initStyle :: Style V2 Double
initStyle = mempty # lwL 0.1

nudgedText :: String -> Diagram B
nudgedText s = alignedText 0.5 0.5 s
    # font "monospace" # fontSize (local 0.25) # bold