metabrush/src/metabrushes/MetaBrush/Draw.hs

{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}

module MetaBrush.Draw
  ( DrawAnchor(..), anchorsAreComplementary
  , getOrCreateDrawAnchor, addToAnchor
  , anchorPos
  )
  where

-- base
import Control.Monad
  ( when )
import Control.Monad.ST
  ( RealWorld, ST, runST )
import Data.Foldable
  ( for_ )
import GHC.Generics
  ( Generic )
import GHC.Stack
  ( HasCallStack )
import GHC.TypeLits
  ( Symbol )

-- containers
import Data.Sequence
  ( Seq(..) )
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

-- deepseq
import Control.DeepSeq
  ( NFData )

-- generic-lens
import Data.Generics.Product.Fields
  ( field, field' )

-- lens
import Control.Lens
  ( set, over )

-- stm
import Control.Concurrent.STM
  ( STM )

-- transformers
import Control.Monad.Trans.Except
  ( Except )
import qualified Control.Monad.Trans.Except as Except
import Control.Monad.Trans.Reader
  ( runReaderT )

-- brush-strokes
import Math.Bezier.Spline
import Math.Bezier.Stroke
  ( newCache )
import Math.Linear
  ( ℝ(..) )

-- MetaBrush
import MetaBrush.Brush
  ( SomeBrush(..), NamedBrush (..), WithParams (defaultParams) )
import MetaBrush.Document
import MetaBrush.Hover
  ( inPointClickRange )
import MetaBrush.Layer
import MetaBrush.Records
import MetaBrush.Stroke
import MetaBrush.Unique
  ( Unique, UniqueSupply, freshUnique )

--------------------------------------------------------------------------------

-- | A draw anchor, to continue drawing from one end of an existing stroke.
data DrawAnchor
  = DrawAnchor
    { anchorIsNew   :: !Bool
    , anchorStroke  :: !Unique
    , anchorIsAtEnd :: !Bool
    }
  deriving stock ( Show, Eq, Generic )
  deriving anyclass NFData

-- | Computes whether two anchors are the two ends of the same stroke.
anchorsAreComplementary :: DrawAnchor -> DrawAnchor -> Bool
anchorsAreComplementary
  ( DrawAnchor { anchorStroke = uniq1, anchorIsAtEnd = end1 } )
  ( DrawAnchor { anchorStroke = uniq2, anchorIsAtEnd = end2 } )
  = uniq1 == uniq2 && end1 /= end2

-- | Compute a draw anchor at the given position, e.g. to continue
-- drawing a stroke or to start a new one.
getOrCreateDrawAnchor
  :: UniqueSupply
  -> Maybe SomeBrush
  -> ℝ 2
  -> Document
  -> STM ( Document, DrawAnchor )
getOrCreateDrawAnchor uniqueSupply mbBrush c doc@( Document { documentContent = oldContent, documentMetadata } ) =

  -- Deselect all points, and try to find a valid anchor for drawing
  -- (a path start/end point at mouse click point).

  case
    Except.runExcept $
      forStrokeHierarchy
        ( layerMetadata documentMetadata )
        ( strokeHierarchy oldContent )
        findAnchor
  of
    -- Anchor found: use it.
    Left anchor ->
      let newMeta :: DocumentMetadata
          newMeta = documentMetadata { selectedPoints = mempty }
      in pure ( doc { documentMetadata = newMeta }
              , anchor )
    -- No anchor found: start a new stroke (on a new stroke layer).
    Right {} -> do
      newStrokeUnique <- runReaderT freshUnique uniqueSupply
      let
        newStroke :: Stroke
        newStroke =
          case mbBrush of
            Nothing ->
              let
                newSpline :: StrokeSpline Open ( Record @Symbol '[] )
                newSpline =
                  Spline
                    { splineStart  = PointData c ( MkR ℝ0 )
                    , splineCurves = OpenCurves Empty
                    }
              in
                Stroke
                  { strokeSpline  = newSpline
                  , strokeBrush   = Nothing
                  }
            Just ( SomeBrush ( b@( NamedBrush { brushFunction = fn } ) :: NamedBrush brushFields ) ) ->
              let
                newSpline :: StrokeSpline Open ( Record brushFields )
                newSpline =
                  Spline
                    { splineStart  = PointData c ( MkR $ defaultParams fn )
                    , splineCurves = OpenCurves Empty
                    }
              in
                Stroke
                  { strokeSpline  = newSpline
                  , strokeBrush   = Just b
                  }
        newSel = StrokePoints $ Map.singleton newStrokeUnique ( Set.singleton FirstPoint )
        newMeta :: DocumentMetadata
        newMeta =
            set ( field' @"selectedPoints" ) newSel
          . over ( field' @"layerMetadata" . field' @"layerNames" ) ( Map.insert newStrokeUnique "Stroke" )
          $ documentMetadata
        newHierarchy :: StrokeHierarchy
        newHierarchy =
          ( strokeHierarchy oldContent )
            { topLevel = newStrokeUnique : topLevel ( strokeHierarchy oldContent )
            , content = Map.insert newStrokeUnique newStroke ( content ( strokeHierarchy oldContent ) )
            }
        newContent :: DocumentContent
        newContent = oldContent { strokeHierarchy = newHierarchy }
        newDoc' :: Document
        newDoc' =
          doc { documentMetadata = newMeta
              , documentContent  = newContent
              }
        anchor =
          DrawAnchor
            { anchorIsNew   = True
            , anchorStroke  = newStrokeUnique
            , anchorIsAtEnd = True
            }
      pure ( newDoc', anchor )
  where
    zoom = documentZoom documentMetadata

    findAnchor :: WithinParent Unique -> Stroke -> StrokeMetadata -> Except DrawAnchor UpdateStroke
    findAnchor ( WithinParent _ strokeUnique ) ( Stroke { strokeSpline }) ( StrokeMetadata { strokeVisible, strokeLocked } ) = do
      when ( strokeVisible && not strokeLocked ) $
        findAnchorSpline strokeSpline
      return PreserveStroke

      where
        findAnchorSpline
          :: forall clo brushParams
          .  SplineTypeI clo
          => StrokeSpline clo brushParams
          -> Except DrawAnchor ()
        findAnchorSpline spline =
          for_ ( endpointAnchor strokeUnique spline ) Except.throwE
          where
            -- See if we can anchor a drawing operation on a given (visible) stroke.
            endpointAnchor :: Unique -> StrokeSpline clo brushParams -> Maybe DrawAnchor
            endpointAnchor uniq ( Spline { splineStart, splineCurves } ) = case ssplineType @clo of
              SOpen
                | let
                    p0 = coords splineStart
                , inPointClickRange zoom c p0
                , let anchor =
                       DrawAnchor
                         { anchorIsNew   = False
                         , anchorStroke  = uniq
                         , anchorIsAtEnd = False
                         }
                -> Just anchor
                | OpenCurves ( _ :|> lastCurve ) <- splineCurves
                , let
                    pn :: ℝ 2
                    pn = coords ( openCurveEnd lastCurve )
                , inPointClickRange zoom c pn
                , let anchor =
                       DrawAnchor
                         { anchorIsNew   = False
                         , anchorStroke  = uniq
                         , anchorIsAtEnd = True
                         }
                -> Just anchor
              _ -> Nothing

anchorPos :: HasCallStack => DrawAnchor -> Document -> ℝ 2
anchorPos ( DrawAnchor { anchorStroke, anchorIsAtEnd } )
  ( Document { documentContent = Content { strokeHierarchy = Hierarchy { content = strokes } } } ) =
    case Map.lookup anchorStroke strokes of
      Nothing -> error "draw anchor: stroke not present in hierarchy"
      Just ( Stroke { strokeSpline } ) ->
        coords $
          if anchorIsAtEnd
          then splineEnd   ( adjustSplineType @Open strokeSpline )
          else splineStart ( adjustSplineType @Open strokeSpline )

addToAnchor :: DrawAnchor -> Spline Open () ( PointData () ) -> Document -> Document
addToAnchor anchor newSpline doc@( Document { documentMetadata, documentContent = oldContent } ) =
  let
    strokes' =
      runST $
        forStrokeHierarchy
          ( layerMetadata documentMetadata )
          ( strokeHierarchy oldContent )
          ( \ u s _ -> updateStroke u s )
  in doc { documentContent = oldContent { strokeHierarchy = strokes' } }
  where

    updateStroke :: forall s. WithinParent Unique -> Stroke -> ST s UpdateStroke
    updateStroke ( WithinParent _ strokeUnique ) stroke@( Stroke { strokeSpline = ( oldSpline :: StrokeSpline clo brushParams ) })
      | strokeUnique == anchorStroke anchor
      , SOpen <- ssplineType @clo
      , let prevSpline0 = coCache @RealWorld @s oldSpline
      = do
        finalSpline <-
          if anchorIsAtEnd anchor
          then do
            ( i0, prevSpline ) <-
              case splineCurves prevSpline0 of
                OpenCurves ( prevCurves :|> lastCurve ) -> do
                  cache <- newCache
                  -- Invalidate the point we are connecting to (stroke endpoint).
                  let lastCurve' = set ( field' @"curveData" . field' @"cachedStroke" ) cache lastCurve
                  return
                    ( curveIndex ( curveData lastCurve ) + 1

                    , prevSpline0 { splineCurves =
                        OpenCurves ( prevCurves :|> lastCurve' ) }
                    )
                OpenCurves _ ->
                  return ( 0, prevSpline0 )
            newSpline' <- newCaches ( \ i -> i0 + fromIntegral i ) ( brushParams ( splineEnd prevSpline ) ) newSpline
            return $ prevSpline <> newSpline'
          else do
            ( i0, prevSpline ) <-
              case splineCurves prevSpline0 of
                OpenCurves ( firstCurve :<| nextCurves ) -> do
                  cache <- newCache
                  -- Invalidate the point we are connecting to (stroke endpoint).
                  let firstCurve' = set ( field' @"curveData" . field' @"cachedStroke" ) cache firstCurve
                  return
                    ( curveIndex ( curveData firstCurve ) - 1
                    -- Invalidate the point we are connecting to (stroke endpoint).
                    , prevSpline0 { splineCurves = OpenCurves ( firstCurve' :<| nextCurves ) }
                    )
                OpenCurves _ ->
                  return ( 0, prevSpline0 )
            newSpline' <- newCaches ( \ i -> i0 - fromIntegral i - 1 ) ( brushParams ( splineStart prevSpline ) ) ( reverseSpline newSpline )
            return $ newSpline' <> prevSpline
        return $ UpdateStrokeTo ( stroke { strokeSpline = coCache @s @RealWorld finalSpline } )
      | otherwise
      = return PreserveStroke

newCaches :: ( Int -> Rational )
          -> brushParams
          -> Spline Open () ( PointData () )
          -> ST s ( Spline Open ( CurveData s ) ( PointData brushParams ) )
newCaches mkCrvIx brushParams =
  ibitraverseSpline
    ( \ i _ -> bitraverseCurve
                  ( \ _ -> CurveData ( mkCrvIx i ) <$> newCache )
                  ( \ _ -> return . setBrushData )
    )
    ( return . setBrushData )
  where
    setBrushData = set ( field @"brushParams" ) brushParams