{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}

module MetaBrush.Stroke where

-- base
import Control.Arrow
  ( (***) )
import Control.Monad.ST
  ( ST, RealWorld, runST )
import Data.Coerce
  ( coerce )
import Data.Foldable
  ( foldr' )
import Data.Functor.Identity
  ( Identity(..) )
import GHC.Generics
  ( Generic, Generic1 )
import GHC.Stack
  ( HasCallStack )
import GHC.TypeLits
  ( Symbol )
import Unsafe.Coerce
  ( unsafeCoerce )

-- acts
import Data.Act
  ( Act(..), Torsor(..) )

-- containers
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

-- deepseq
import Control.DeepSeq
  ( NFData(..), NFData1 )

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

-- groups
import Data.Group
  ( Group(..) )

-- lens
import Control.Lens
  ( Lens'
  , view, over
  )

-- text
import Data.Text
  ( Text )

-- transformers
import Control.Monad.IO.Class
  ( MonadIO )
import Control.Monad.Trans.Reader
  ( ReaderT )
import Control.Monad.Trans.State.Strict
  ( StateT )
import qualified Control.Monad.Trans.State.Strict as State

-- brush-strokes
import Math.Bezier.Spline
  ( Spline(..), KnownSplineType
  , PointType(..), bitraverseSpline, bitraverseCurve
  )
import Math.Bezier.Stroke
  ( CachedStroke, newCache )
import Math.Module
  ( Module
    ( origin, (^+^), (^-^), (*^) )
  )
import Math.Linear
  ( ℝ(..), T(..) )

-- MetaBrush
import MetaBrush.Brush
  ( NamedBrush, PointFields )
import MetaBrush.Layer
    ( Hierarchy(..), LayerMetadata(..), emptyHierarchy )
import MetaBrush.Records
import MetaBrush.Unique
  ( Unique, UniqueSupply, freshUnique )
import MetaBrush.Util
  ( (!) )

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

-- | Data attached to each point on a spline.
data PointData params
  = PointData
  { pointCoords :: !( ℝ 2 )
  , brushParams :: !params
  }
  deriving stock ( Show, Generic, Functor, Foldable, Traversable )
  deriving anyclass NFData

-- | Data attached to each curve in a spline.
data CurveData s =
  CurveData
    { curveIndex   :: !Rational
    , cachedStroke :: !( CachedStroke s )
    }
  deriving stock Generic
  deriving anyclass NFData

instance Show ( CurveData s ) where
  show ( CurveData { curveIndex } ) = show curveIndex
instance Eq  ( CurveData s ) where
  ( CurveData { curveIndex = i1 } ) == ( CurveData { curveIndex = i2 } )
    = i1 == i2
instance Ord  ( CurveData s ) where
  compare ( CurveData { curveIndex = i1 } ) ( CurveData { curveIndex = i2 } )
    = compare i1 i2

-- | An index for a point on a spline.
data PointIndex
  = FirstPoint
  | PointIndex
    -- | Which curve the point belongs to.
    { pointCurve :: !Rational
    -- | Index within a curve.
    , pointType  :: !PointType
    }
  deriving stock ( Show, Eq, Ord, Generic )
  deriving anyclass NFData

_coords :: Lens' ( PointData brushParams ) ( ℝ 2 )
_coords = field' @"pointCoords"

coords :: PointData brushParams -> ℝ 2
coords = view _coords

type StrokeSpline clo brushParams =
  Spline clo ( CurveData RealWorld ) ( PointData brushParams )

data Stroke where
  Stroke
    :: forall clo ( pointFields :: [ Symbol ] ) ( brushFields :: [ Symbol ] )
    .  ( KnownSplineType clo , PointFields pointFields )
    =>
      { strokeBrush  :: !( Maybe ( NamedBrush brushFields ) )
      , strokeSpline :: !( StrokeSpline clo ( Record pointFields ) )
      }
    -> Stroke
deriving stock instance Show Stroke
instance NFData Stroke where
  rnf ( Stroke { strokeBrush, strokeSpline } )
    = rnf strokeBrush `seq` rnf strokeSpline

_strokeSpline
  :: forall f
  .  Functor f
  => (  forall clo pointParams ( pointFields :: [ Symbol ] )
     .  ( KnownSplineType clo
        , pointParams ~ Record pointFields
        , PointFields pointFields
        )
     => StrokeSpline clo pointParams
     -> f ( StrokeSpline clo pointParams )
     )
  -> Stroke -> f Stroke
_strokeSpline f ( Stroke { strokeSpline = oldStrokeSpline, .. } )
  = ( \ newSpline -> Stroke { strokeSpline = newSpline, .. } ) <$> f oldStrokeSpline

overStrokeSpline
  :: (  forall clo pointParams ( pointFields :: [ Symbol ] )
     .  ( KnownSplineType clo
        , pointParams ~ Record pointFields
        , PointFields pointFields
        )
     => StrokeSpline clo pointParams
     -> StrokeSpline clo pointParams
     )
  -> Stroke -> Stroke
overStrokeSpline f = coerce ( _strokeSpline @Identity ( coerce . f ) )

setStrokeBrush
  :: ( Maybe ( NamedBrush brushFields ) )
  -> Stroke -> Stroke
setStrokeBrush brush
  ( Stroke { strokeSpline = ( oldStrokeSpline :: StrokeSpline clo pointParams ) } ) =
    -- Invalidate all of the cached brush strokes.
  let spline' :: ST s ( Spline clo ( CurveData s ) ( PointData pointParams ) )
      spline' = bitraverseSpline
                  ( \ _ -> bitraverseCurve invalidateCurve ( const return ) )
                  return
                oldStrokeSpline
  in
    Stroke
      { strokeSpline = runST $ coCache <$> spline'
      , strokeBrush = brush }
  where
    invalidateCurve :: CurveData RealWorld -> ST s ( CurveData s )
    invalidateCurve crv = do
      noCache <- newCache
      return $ crv { cachedStroke = noCache }

{-# NOINLINE coCache #-}
coCache :: forall s t clo ptData. Spline clo ( CurveData s ) ptData -> Spline clo ( CurveData t ) ptData
coCache = unsafeCoerce

instance Act ( T ( ℝ 2 ) ) ( PointData params ) where
  v • ( dat@( PointData { pointCoords = p } ) ) =
    dat { pointCoords = v • p }

data DiffPointData diffBrushParams
  = DiffPointData
  { diffVector :: !( T ( ℝ 2 ) )
  , diffParams :: !diffBrushParams
  }
  deriving stock    ( Show, Generic, Generic1, Functor, Foldable, Traversable )
  deriving anyclass ( NFData, NFData1 )

instance Module Double diffBrushParams => Semigroup ( DiffPointData diffBrushParams ) where
  DiffPointData v1 p1 <> DiffPointData v2 p2 =
    DiffPointData ( v1 <> v2 ) ( p1 ^+^ p2 )
instance Module Double diffBrushParams => Monoid ( DiffPointData diffBrushParams ) where
  mempty = DiffPointData mempty origin
instance Module Double diffBrushParams => Group ( DiffPointData diffBrushParams ) where
  invert ( DiffPointData v1 p1 ) =
    DiffPointData ( invert v1 ) ( -1 *^ p1 )

instance ( Module Double diffBrushParams, Act diffBrushParams brushParams )
      => Act    ( DiffPointData diffBrushParams ) ( PointData brushParams ) where
  (•) ( DiffPointData { diffVector = dp, diffParams = db } )
    = over _coords ( dp • )
    . over ( field' @"brushParams" ) ( db • )
instance ( Module Double diffBrushParams, Torsor diffBrushParams brushParams )
      => Torsor ( DiffPointData diffBrushParams ) ( PointData brushParams ) where
  ( PointData
    { pointCoords = p1
    , brushParams = b1
    } ) <--
    ( PointData
      { pointCoords = p2
      , brushParams = b2 } ) =
    DiffPointData
      { diffVector = p1 <-- p2
      , diffParams = b1 <-- b2
      }

instance Module Double brushParams => Module Double ( DiffPointData brushParams ) where
  origin = mempty
  (^+^) = (<>)
  x ^-^ y = x <> invert y
  d *^ DiffPointData v1 p1 =
    DiffPointData ( d *^ v1 ) ( d *^ p1 )

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

-- | Metadata about a stroke, such as its name or its visibility.
data StrokeMetadata
  = StrokeMetadata
  { strokeName    :: !Text
  , strokeVisible :: !Bool
  , strokeLocked  :: !Bool
  }

type StrokeHierarchy = Hierarchy Stroke

data UpdateStroke
  = PreserveStroke
  | DeleteStroke
  | UpdateStrokeTo !Stroke
  deriving stock Show

-- | Traverse through a stroke hierarchy.
forStrokeHierarchy
  :: forall f
  .  ( HasCallStack, Applicative f )
  => LayerMetadata
  -> StrokeHierarchy
  -> ( Unique -> Stroke -> StrokeMetadata -> f UpdateStroke )
  -> f StrokeHierarchy
forStrokeHierarchy
  ( LayerMetadata { layerNames, invisibleLayers, lockedLayers } ) hierarchy0 f =
    foldr' ( g Nothing ( True, False ) ) ( pure hierarchy0 ) ( topLevel hierarchy0 )
  where

    insertMaybe :: Maybe Unique -> Unique -> StrokeHierarchy -> UpdateStroke -> StrokeHierarchy
    insertMaybe mbPar u old@( Hierarchy oldTl oldGps oldStrokes ) = \case
      PreserveStroke   -> old
      UpdateStrokeTo s -> Hierarchy oldTl oldGps ( Map.insert u s oldStrokes )
      DeleteStroke     ->
        let newStrokes = Map.delete u oldStrokes
         in case mbPar of
              Nothing ->
                Hierarchy ( filter ( /= u ) oldTl ) oldGps newStrokes
              Just par ->
                Hierarchy oldTl ( Map.adjust ( filter ( /= u ) ) par oldGps ) newStrokes


    g :: Maybe Unique -> ( Bool, Bool ) -> Unique -> f StrokeHierarchy -> f StrokeHierarchy
    g par ( vis, lock ) u acc =
      let vis'  = vis && not ( u `Set.member` invisibleLayers )
          lock' = lock || u `Set.member` lockedLayers
      in
        case Map.lookup u ( groups hierarchy0 ) of
          Nothing ->
            let
              meta =
                StrokeMetadata
                  { strokeName    = layerNames ! u
                  , strokeVisible = vis'
                  , strokeLocked  = lock'
                  }
            in
              insertMaybe par u <$> acc <*> f u ( content hierarchy0 ! u ) meta
          Just ds ->
            foldr' ( g ( Just u ) ( vis', lock' ) ) acc ds

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

-- | Recursive representation of a stroke hierarchy.
--
-- Used for serialisation/deserialisation only.
type Layers = [ Layer ]

-- | Layer in a recursive representation of a stroke hierarchy.
--
-- Used for serialisation/deserialisation only.
data Layer
  = StrokeLayer
      { layerName    :: !Text
      , layerVisible :: !Bool
      , layerLocked  :: !Bool
      , layerStroke  :: !Stroke
      }
  | GroupLayer
      { layerName     :: !Text
      , layerVisible  :: !Bool
      , layerLocked   :: !Bool
      , groupChildren :: !Layers
      }
  deriving stock Show

strokeHierarchyLayers :: LayerMetadata -> StrokeHierarchy -> Layers
strokeHierarchyLayers
  ( LayerMetadata { layerNames, invisibleLayers, lockedLayers } )
  ( Hierarchy topLevel hierarchy content ) = map go topLevel
    where
      go :: Unique -> Layer
      go layerUnique =
        let
          layerName    = layerNames ! layerUnique
          layerVisible = not $ layerUnique `Set.member` invisibleLayers
          layerLocked  = layerUnique `Set.member` lockedLayers
        in
          case Map.lookup layerUnique hierarchy of
            Nothing ->
              StrokeLayer
                { layerName, layerVisible, layerLocked
                , layerStroke = content ! layerUnique
                }
            Just cs ->
              GroupLayer
                { layerName, layerVisible, layerLocked
                , groupChildren = map go cs
                }

{-# INLINEABLE layersStrokeHierarchy #-}
layersStrokeHierarchy :: forall m. MonadIO m => Layers -> ReaderT UniqueSupply m ( LayerMetadata, StrokeHierarchy )
layersStrokeHierarchy lays = ( `State.execStateT` ( mempty, emptyHierarchy ) ) $ do
  us <- traverse go lays
  State.modify' ( \ ( meta, hierarchy ) -> ( meta, hierarchy { topLevel = us } ) )
  where
    go :: Layer -> StateT ( LayerMetadata, StrokeHierarchy ) ( ReaderT UniqueSupply m ) Unique
    go l = do
      u <- freshUnique
      let updMeta ( LayerMetadata nms invis locked ) =
            LayerMetadata
              { layerNames      = Map.insert u ( layerName l ) nms
              , invisibleLayers = if layerVisible l then invis else Set.insert u invis
              , lockedLayers    = if layerLocked l then Set.insert u locked else locked
              }
      updHierarchy <- case l of
        StrokeLayer { layerStroke } ->
          return $ \ h -> h { content = Map.insert u layerStroke ( content h ) }
        GroupLayer { groupChildren } -> do
          us <- traverse go groupChildren
          return $ \ h -> h { groups = Map.insert u us ( groups h ) }
      State.modify' ( updMeta *** updHierarchy )
      return u

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