Break.hs

module KB.Text.Layout.Break
  ( -- * Breaking prepared text into lines
    layoutGreedy
  , layoutOptimal
  , LineRange (..)
  , Cursor (..)
  , LineEnd (..)

    -- * Breaking one line at a time
  , layoutNextLineRange

    -- * Extracting line contents
  , materializeSlices
  , LineSlice (..)
  , materializeLineRange

    -- * Inspecting layouts
  , layoutStats
  , LayoutStats (..)
  , lineStretch
  , LineStretch (..)
  , shrinkwrap
  ) where

import Data.Foldable (foldl')
import Data.Maybe (mapMaybe)
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Vector (Vector)
import Data.Vector qualified as Vector

import KB.Text.Layout.Analysis (BreakKind (..))
import KB.Text.Layout.Measure (MeasuredSegment (..), PreparedText (..))

layoutGreedy :: PreparedText -> Float -> [LineRange]
layoutGreedy prepared maxWidth = go (Cursor 0 0)
  where
    go cursor =
      case layoutNextLineRange prepared maxWidth cursor of
        Nothing -> []
        Just (line, Nothing) -> [line]
        Just (line, Just next) -> line : go next

layoutOptimal :: PreparedText -> Float -> [LineRange]
layoutOptimal prepared target =
  case chosen of
    Just lns -> lns
    Nothing -> layoutGreedy prepared target
  where
    segs = prepared.segments
    total = Vector.length segs

    flowWidth seg = case seg.kind of
      SoftHyphen -> 0
      _ -> seg.width
    spaceWidth seg = case seg.kind of
      Space -> seg.width
      _ -> 0
    prefixFlow = Vector.scanl' (\acc seg -> acc + flowWidth seg) 0 segs
    prefixSpace = Vector.scanl' (\acc seg -> acc + spaceWidth seg) 0 segs

    candidates :: Vector Breakpoint
    candidates =
      Vector.fromList $
        concatMap breakable [0 .. total - 1] <> [Breakpoint{to = total, next = total, hyphen = False, forced = True}]
      where
        breakable i = case (segs Vector.! i).kind of
          Space -> [plain i]
          Tab -> [plain i]
          PreservedSpace -> [plain i]
          ZeroWidthBreak -> [plain i]
          SoftHyphen -> [Breakpoint{to = i, next = i + 1, hyphen = True, forced = False}]
          HardBreak -> [Breakpoint{to = i, next = i + 1, hyphen = False, forced = True}]
          Atomic -> [Breakpoint{to = i, next = i, hyphen = False, forced = False}, Breakpoint{to = i + 1, next = i + 1, hyphen = False, forced = False}]
          Word -> []
          Glue -> []
        plain i = Breakpoint{to = i, next = i + 1, hyphen = False, forced = False}

    skipSpaces i
      | i < total, (segs Vector.! i).kind == Space = skipSpaces (i + 1)
      | otherwise = i

    lineStart i = skipSpaces i

    lineOf from cand =
      let
        a = lineStart from
        hyphenWidth = if cand.hyphen then (segs Vector.! cand.to).width else 0
        natural = prefixFlow Vector.! cand.to - prefixFlow Vector.! a + hyphenWidth
        stretchable = prefixSpace Vector.! cand.to - prefixSpace Vector.! a
      in
        (a, natural, stretchable)

    badness natural stretchable free
      | slack >= 0 =
          if free then
            Just 0
          else
            if stretchable <= 0 then
              Just $! if slack < 1e-4 then 0 else looseNoGlue + slack * 1000
            else
              Just $! 100 * ratio * ratio * ratio
      | negate slack > shrink + 1e-4 = Nothing
      | otherwise = Just $! 100 * tightness * tightness * tightness
      where
        looseNoGlue :: Float
        looseNoGlue = 1e7

        slack = target - natural
        ratio = slack / (0.5 * stretchable)
        shrink = stretchable / 3
        tightness = negate slack / max shrink 1e-6

    demerits b hyphen = (10 + b) * (10 + b) + (if hyphen then 2500 else 0)

    lastForcedBefore :: Vector Int
    lastForcedBefore = Vector.fromList $ scanl step (-1) [0 .. Vector.length candidates - 2]
      where
        step acc j = if (candidates Vector.! j).forced then j else acc

    best :: Vector (Maybe (Float, Int))
    best = Vector.generate (Vector.length candidates) bestAt
      where
        bestAt j =
          let
            cand = candidates Vector.! j
            floor_ = lastForcedBefore Vector.! j
            fromStart
              | floor_ == -1 = attempt (-1) 0
              | otherwise = Nothing
            fromPrev i
              | i < max floor_ 0 = []
              | otherwise = case best Vector.! i of
                  Nothing -> rest
                  Just (cost, _) -> maybe rest (: rest) do
                    d <- lineCost (candidates Vector.! i).next
                    pure (cost + d, i)
              where
                rest = fromPrev (i - 1)
            attempt i from = (\d -> (d, i)) <$> lineCost from
            lineCost from =
              let (_a, natural, stretchable) = lineOf from cand
              in demerits <$> badness natural stretchable cand.forced <*> pure cand.hyphen
            options = maybe [] (: []) fromStart <> fromPrev (j - 1)
          in
            case options of
              [] -> Nothing
              _ -> Just $! minimum options

    chosen = do
      let final = Vector.length candidates - 1
      _ <- best Vector.! final
      let walk j acc
            | j == -1 = acc
            | otherwise = case best Vector.! j of
                Nothing -> acc
                Just (_, i) ->
                  let
                    cand = candidates Vector.! j
                    from = if i == -1 then 0 else (candidates Vector.! i).next
                    (a, natural, _) = lineOf from cand
                    ended
                      | cand.hyphen = Hyphenated
                      | cand.to >= total = Finished
                      | cand.forced = HardBroken
                      | otherwise = Wrapped
                    line =
                      LineRange
                        { from = Cursor a 0
                        , to = Cursor cand.to 0
                        , width = natural
                        , ended
                        }
                  in
                    walk i $ if a >= total then acc else (line : acc)
      pure $! walk final []

data LineRange = LineRange
  { from :: Cursor
  , to :: Cursor
  , width :: Float
  , ended :: LineEnd
  }
  deriving stock (Eq, Show)

data Cursor = Cursor
  { segment :: Int
  , grapheme :: Int
  }
  deriving stock (Eq, Ord, Show)

data LineEnd
  = Wrapped
  | Hyphenated
  | Overflowed
  | HardBroken
  | Finished
  deriving stock (Eq, Ord, Show)

data Breakpoint = Breakpoint
  { to :: Int
  , next :: Int
  , hyphen :: Bool
  , forced :: Bool
  }

layoutNextLineRange :: PreparedText -> Float -> Cursor -> Maybe (LineRange, Maybe Cursor)
layoutNextLineRange prepared maxWidth cursor
  | start.segment >= total = Nothing
  | otherwise = Just (scanSeg start.segment start.grapheme 0 0 Nothing)
  where
    segs = prepared.segments
    total = Vector.length segs
    start = skipSpaces cursor

    skipSpaces c
      | c.segment < total
      , c.grapheme == 0
      , (segs Vector.! c.segment).kind == Space =
          skipSpaces Cursor{segment = c.segment + 1, grapheme = 0}
      | otherwise = c

    emit to width ended = LineRange{from = start, to, width, ended}

    scanSeg i g acc visible cand
      | i >= total = (emit (Cursor total 0) visible Finished, Nothing)
      | otherwise =
          let
            seg = segs Vector.! i
            segWidth = if g == 0 then seg.width else sum (map snd (dropChars g seg.graphemeWidths))
            after = Cursor (i + 1) 0
            spaceLike =
              scanSeg
                (i + 1)
                0
                (acc + segWidth)
                visible
                (Just Candidate{to = Cursor i 0, width = visible, end = Wrapped, next = after})
            fitOrBreak
              | acc + segWidth <= maxWidth + epsilon =
                  scanSeg (i + 1) 0 (acc + segWidth) (acc + segWidth) cand
              | otherwise =
                  case cand of
                    Just c -> (emit c.to c.width c.end, Just c.next)
                    Nothing -> emergency seg i g acc
          in
            case seg.kind of
              HardBreak -> (emit (Cursor i 0) visible HardBroken, Just after)
              Space -> spaceLike
              PreservedSpace -> spaceLike
              Tab -> spaceLike
              ZeroWidthBreak -> spaceLike
              SoftHyphen
                | visible + segWidth <= maxWidth + epsilon ->
                    scanSeg (i + 1) 0 acc visible $
                      Just
                        Candidate
                          { to = Cursor i 0
                          , width = visible + segWidth
                          , end = Hyphenated
                          , next = after
                          }
                | otherwise -> scanSeg (i + 1) 0 acc visible cand
              Atomic ->
                let candBefore
                      | acc > 0
                      , maybe True (\c -> c.width < visible) cand =
                          Just Candidate{to = Cursor i 0, width = visible, end = Wrapped, next = Cursor i 0}
                      | otherwise = cand
                in if acc + segWidth <= maxWidth + epsilon then
                     scanSeg (i + 1) 0 (acc + segWidth) (acc + segWidth) $
                       Just Candidate{to = after, width = acc + segWidth, end = Wrapped, next = after}
                   else case candBefore of
                     Just c -> (emit c.to c.width c.end, Just c.next)
                     Nothing -> (emit after (acc + segWidth) Overflowed, Just after)
              Word -> fitOrBreak
              Glue -> fitOrBreak

    emergency seg i g acc =
      let
        ws = dropChars g seg.graphemeWidths
        counted = countFits acc ws
        count
          | counted == 0 && acc <= epsilon = 1
          | otherwise = counted
        taken = take count ws
        g' = g + sum (map fst taken)
        takenWidth = sum (map snd taken)
      in
        if count == 0 then
          (emit (Cursor i g) acc Overflowed, Just (Cursor i g))
        else
          (emit (Cursor i g') (acc + takenWidth) Overflowed, Just (advanceCursor seg i g'))

    countFits = fits 0
      where
        fits n w = \case
          [] -> n
          (_, cw) : rest
            | w + cw <= maxWidth + epsilon -> fits (n + 1) (w + cw) rest
            | otherwise -> n

    advanceCursor seg i g'
      | g' >= Text.length seg.text = Cursor (i + 1) 0
      | otherwise = Cursor i g'

data Candidate = Candidate
  { to :: Cursor
  , width :: Float
  , end :: LineEnd
  , next :: Cursor
  }

epsilon :: Float
epsilon = 1e-4

materializeSlices :: PreparedText -> LineRange -> [LineSlice]
materializeSlices prepared line = mergeSlices $ mapMaybe slice [line.from.segment .. line.to.segment] <> hyphen
  where
    segs = prepared.segments
    hyphen =
      [ LineSlice
          { text = "-"
          , style = seg.style
          , width = seg.width
          , atom = False
          }
      | line.ended == Hyphenated
      , let seg = segs Vector.! line.to.segment
      ]
    slice j
      | j >= Vector.length segs = Nothing
      | otherwise =
          let
            seg = segs Vector.! j
            startG = if j == line.from.segment then line.from.grapheme else 0
            endG = if j == line.to.segment then line.to.grapheme else Text.length seg.text
            piece = Text.take (endG - startG) (Text.drop startG seg.text)
            pieceWidth
              | startG == 0 && endG == Text.length seg.text =
                  seg.width
              | otherwise =
                  sum . map snd $
                    takeChars (endG - startG) $
                      dropChars startG seg.graphemeWidths
          in
            case seg.kind of
              SoftHyphen -> Nothing
              HardBreak -> Nothing
              ZeroWidthBreak -> Nothing
              _
                | Text.null piece -> Nothing
                | otherwise ->
                    Just
                      LineSlice
                        { text = piece
                        , style = seg.style
                        , width = pieceWidth
                        , atom = seg.kind == Atomic
                        }

data LineSlice = LineSlice
  { text :: Text
  , style :: Int
  , width :: Float
  , atom :: Bool
  }
  deriving stock (Eq, Show)

mergeSlices :: [LineSlice] -> [LineSlice]
mergeSlices = \case
  [] -> []
  a : b : rest
    | not a.atom
    , not b.atom
    , a.style == b.style ->
        mergeSlices $
          LineSlice
            { text = a.text <> b.text
            , style = a.style
            , width = a.width + b.width
            , atom = False
            }
            : rest
  a : rest ->
    a : mergeSlices rest

materializeLineRange :: PreparedText -> LineRange -> Text
materializeLineRange prepared line = visible <> suffix
  where
    segs = prepared.segments
    suffix = if line.ended == Hyphenated then "-" else ""
    visible = Text.concat (mapMaybe piece [line.from.segment .. line.to.segment])
    piece j
      | j >= Vector.length segs = Nothing
      | otherwise =
          let
            seg = segs Vector.! j
            startG = if j == line.from.segment then line.from.grapheme else 0
            endG = if j == line.to.segment then line.to.grapheme else Text.length seg.text
          in
            case seg.kind of
              SoftHyphen -> Nothing
              HardBreak -> Nothing
              ZeroWidthBreak -> Nothing
              _ -> Just (Text.take (endG - startG) (Text.drop startG seg.text))

layoutStats :: [LineRange] -> LayoutStats
layoutStats = foldl' step LayoutStats{lineCount = 0, maxLineWidth = 0}
  where
    step stats line =
      LayoutStats
        { lineCount = stats.lineCount + 1
        , maxLineWidth = max stats.maxLineWidth line.width
        }

data LayoutStats = LayoutStats
  { lineCount :: Int
  , maxLineWidth :: Float
  }
  deriving stock (Eq, Show)

lineStretch :: PreparedText -> LineRange -> LineStretch
lineStretch prepared line = LineStretch{spaces = length inside, width = sum inside}
  where
    segs = prepared.segments
    upper
      | line.to.grapheme == 0 = line.to.segment - 1
      | otherwise = line.to.segment
    inside =
      [ seg.width
      | j <- [line.from.segment .. min upper (Vector.length segs - 1)]
      , let seg = segs Vector.! j
      , seg.kind == Space
      ]

data LineStretch = LineStretch
  { spaces :: Int
  , width :: Float
  }
  deriving stock (Eq, Show)

shrinkwrap :: PreparedText -> Float -> Float
shrinkwrap prepared maxWidth = go (40 :: Int) 0 maxWidth
  where
    countAt w = (layoutStats (layoutGreedy prepared w)).lineCount
    target = countAt maxWidth
    go k lo hi
      | k <= 0 = (layoutStats (layoutGreedy prepared hi)).maxLineWidth
      | countAt mid <= target = go (k - 1) lo mid
      | otherwise = go (k - 1) mid hi
      where
        mid = (lo + hi) / 2

dropChars :: Int -> [(Int, Float)] -> [(Int, Float)]
dropChars n xs
  | n <= 0 = xs
  | otherwise = case xs of
      [] -> []
      (len, _) : rest -> dropChars (n - len) rest

takeChars :: Int -> [(Int, Float)] -> [(Int, Float)]
takeChars n xs
  | n <= 0 = []
  | otherwise = case xs of
      [] -> []
      x@(len, _) : rest -> x : takeChars (n - len) rest