metabrush/src/metabrushes/MetaBrush/Draw.hs

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

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

-- base
import Control.Monad
  ( when )
import Data.Foldable
  ( for_ )
import Data.Functor.Identity
  ( Identity(..) )
import GHC.Generics
  ( Generic )
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
  ( invalidateCache )
import Math.Linear
  ( ℝ(..) )

-- MetaBrush
import MetaBrush.Assert
  ( assert )
import MetaBrush.Brush
  ( NamedBrush(..), PointFields )
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
    , anchorPos     :: !( ℝ 2 )
    , anchorIndex   :: !PointIndex
    }
  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, anchorIndex = end1 } )
  ( DrawAnchor { anchorStroke = uniq2, anchorIndex = 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
  -> ℝ 2
  -> Document
  -> STM ( Document, DrawAnchor )
getOrCreateDrawAnchor uniqueSupply 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@( DrawAnchor { anchorStroke, anchorIndex }) ->
      let newSel = StrokePoints $ Map.singleton anchorStroke ( Set.singleton anchorIndex )
          newMeta :: DocumentMetadata
          newMeta = documentMetadata { selectedPoints = newSel }
      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
        newSpline :: StrokeSpline Open ( Record ( '[] :: [ Symbol ] ) )
        newSpline =
          Spline { splineStart  = PointData c ( MkR ℝ0 )
                 , splineCurves = OpenCurves Empty
                 }
        newStroke :: Stroke
        newStroke =
          Stroke
            { strokeSpline  = newSpline
            , strokeBrush   = Nothing :: Maybe ( NamedBrush ( '[] :: [ Symbol ] ) )
            }
        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
            , anchorPos     = c
            , anchorIndex   = FirstPoint
            }
      pure ( newDoc', anchor )
  where
    zoom = documentZoom documentMetadata

    findAnchor :: Unique -> Stroke -> StrokeMetadata -> Except DrawAnchor UpdateStroke
    findAnchor 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
                         , anchorPos     = p0
                         , anchorIndex   = FirstPoint
                         }
                -> 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
                         , anchorPos     = pn
                         , anchorIndex   = PointIndex
                                             ( curveIndex $ curveData lastCurve )
                                             PathPoint
                         }
                -> Just anchor
              _ -> Nothing

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

    updateStroke :: Unique -> Stroke -> UpdateStroke
    updateStroke strokeUnique stroke
      | strokeUnique == anchorStroke anchor
      , let
          updateSpline
            :: forall clo brushData
            .  SplineTypeI clo
            => StrokeSpline clo brushData -> StrokeSpline clo brushData
          updateSpline prevSpline
            | SOpen <- ssplineType @clo
            = if anchorIsAtEnd anchor
              then
                let
                  i0 = case splineCurves prevSpline of
                    OpenCurves ( _ :|> lastCurve ) ->
                      curveIndex ( curveData lastCurve ) + 1
                    _ -> 0
                  setBrushData :: PointData () -> PointData brushData
                  setBrushData = set ( field @"brushParams" ) ( brushParams ( splineEnd prevSpline ) )
                  setData = bimapSpline ( \ _ -> bimapCurve ( \ s -> CurveData i0 ( invalidateCache $ undefined s ) ) ( \ _ -> setBrushData ) ) setBrushData
                in prevSpline <> setData newSpline
              else
                let
                  i0 = case splineCurves prevSpline of
                    OpenCurves ( firstCurve :<| _ ) ->
                      curveIndex ( curveData firstCurve ) - 1
                    _ -> 0
                  setBrushData :: PointData () -> PointData brushData
                  setBrushData = set ( field @"brushParams" ) ( brushParams ( splineStart prevSpline ) )
                  setData = bimapSpline ( \ _ -> bimapCurve ( \ s -> CurveData i0 ( invalidateCache $ undefined s ) ) ( \ _ -> setBrushData ) ) setBrushData
                in setData ( reverseSpline newSpline ) <> prevSpline
            | otherwise
            = assert False ( "addToAnchor: trying to add to closed spline " <> show strokeUnique )
                prevSpline -- should never add to a closed spline
      = UpdateStrokeTo $ overStrokeSpline updateSpline stroke
      | otherwise
      = PreserveStroke

withAnchorBrushData
  :: forall r
  .  DrawAnchor
  -> Document
  -> ( forall pointParams ( pointFields :: [ Symbol ] ) ( brushFields :: [ Symbol ] )
     . ( pointParams ~ Record pointFields
       , PointFields pointFields
       )
     => Maybe ( NamedBrush brushFields )
     -> pointParams
     -> r
     )
  -> r
withAnchorBrushData anchor ( Document { documentMetadata = Metadata { layerMetadata }, documentContent = Content { strokeHierarchy } } ) f =
  splineAnchor $ Except.runExcept $ forStrokeHierarchy layerMetadata strokeHierarchy relevantStroke
  where

    relevantStroke :: Unique -> Stroke -> StrokeMetadata -> Except Stroke UpdateStroke
    relevantStroke strokeUnique stroke _
      | strokeUnique == anchorStroke anchor
      = Except.throwE stroke
      | otherwise
      = return PreserveStroke

    splineAnchor :: Either Stroke other -> r
    splineAnchor ( Left ( Stroke { strokeSpline = strokeSpline :: StrokeSpline clo pointData, strokeBrush } ) )
      | SOpen <- ssplineType @clo
      = if anchorIsAtEnd anchor
        then f strokeBrush ( brushParams ( splineEnd   strokeSpline ) )
        else f strokeBrush ( brushParams ( splineStart strokeSpline ) )
    splineAnchor _
      = f @_ @'[] @'[] Nothing ( MkR ℝ0 )