{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE OverloadedStrings   #-}

module MetaBrush.Asset.Brushes where

-- base
import Data.Coerce
  ( coerce )

-- containers
import qualified Data.Sequence as Seq
  ( fromList )

-- text
import Data.Text
  ( Text )

-- unordered-containers
import Data.HashMap.Strict
  ( HashMap )
import qualified Data.HashMap.Strict as HashMap
  ( fromList, lookup )

-- MetaBrush
import Math.Bezier.Spline
import Math.Linear
  ( ℝ(..), T(..) )
import Math.Linear.Dual
  ( D, type (~>)(..), Var(var), konst )
import Math.Module
  ( Module((^+^), (*^)) )
import MetaBrush.Brush
  ( Brush(..), SomeBrush(..) )
import MetaBrush.Records

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

lookupBrush :: Text -> Maybe SomeBrush
lookupBrush nm = HashMap.lookup nm brushes

-- | All brushes supported by this application.
brushes :: HashMap Text SomeBrush
brushes = HashMap.fromList
  [ ( nm, b )
  | b@( SomeBrush ( BrushData { brushName = nm } ) )
    <- [ SomeBrush circle, SomeBrush ellipse ]
  ]

-- | Root of @(Sqrt[2] (4 + 3 κ) - 16) (2 - 3 κ)^2 - 8 (1 - 3 κ) Sqrt[8 - 24 κ + 12 κ^2 + 8 κ^3 + 3 κ^4]@.
--
-- Used to approximate circles and ellipses with Bézier curves.
κ :: Double
κ = 0.5519150244935105707435627227925

type CircleBrushFields = '[ "r" ]
circle :: Brush CircleBrushFields
circle = BrushData "circle" ( WithParams deflts circleBrush )
  where
    deflts :: Record CircleBrushFields
    deflts = MkR ( ℝ1 1 )

type EllipseBrushFields = '[ "a", "b", "phi" ]
ellipse :: Brush EllipseBrushFields
ellipse = BrushData "ellipse" ( WithParams deflts ellipseBrush )
  where
    deflts :: Record EllipseBrushFields
    deflts = MkR ( ℝ3 1 1 0 )

--------------------------------------------------------------------------------
-- Differentiable brushes.

circleSpline :: Applicative ( D u ) => ( Double -> Double -> D u v ) -> D u ( Spline 'Closed () v )
circleSpline p = sequenceA $
  Spline { splineStart  = p 1 0
         , splineCurves = ClosedCurves crvs lastCrv }
  where
    crvs = Seq.fromList
      [ Bezier3To (p  1  κ) (p  κ  1) (NextPoint (p  0  1)) ()
      , Bezier3To (p -κ  1) (p -1  κ) (NextPoint (p -1  0)) ()
      , Bezier3To (p -1 -κ) (p -κ -1) (NextPoint (p  0 -1)) ()
      ]
    lastCrv =
        Bezier3To (p  κ -1) (p  1 -κ) BackToStart ()

circleBrush :: Record CircleBrushFields ~> Spline 'Closed () ( ℝ 2 )
circleBrush =
  D \ params ->
    let r :: D ( Record CircleBrushFields ) Double
        r   = runD ( var @1 ) params
        mkPt :: Double -> Double -> D ( Record CircleBrushFields ) ( ℝ 2 )
        mkPt ( kon -> x ) ( kon -> y )
          = fmap coerce
          $   ( x * r ) *^ e_x
          ^+^ ( y * r ) *^ e_y
    in circleSpline @( Record CircleBrushFields ) mkPt
  where
    e_x, e_y :: D ( Record CircleBrushFields ) ( T ( ℝ 2 ) )
    e_x = pure $ T $ ℝ2 1 0
    e_y = pure $ T $ ℝ2 0 1

    kon = konst @( Record CircleBrushFields )

ellipseBrush :: Record EllipseBrushFields ~> Spline 'Closed () ( ℝ 2 )
ellipseBrush =
  D \ params ->
    let a, b, phi :: D ( Record EllipseBrushFields ) Double
        a   = runD ( var @1 ) params
        b   = runD ( var @2 ) params
        phi = runD ( var @3 ) params
        mkPt :: Double -> Double -> D ( Record EllipseBrushFields ) ( ℝ 2 )
        mkPt ( kon -> x ) ( kon -> y )
          = fmap coerce
          $   ( x * a * cos phi - y * b * sin phi ) *^ e_x
          ^+^ ( y * b * cos phi + x * a * sin phi ) *^ e_y
    in circleSpline @( Record EllipseBrushFields ) mkPt
  where
    e_x, e_y :: D ( Record EllipseBrushFields ) ( T ( ℝ 2 ) )
    e_x = pure $ T $ ℝ2 1 0
    e_y = pure $ T $ ℝ2 0 1

    kon = konst @( Record EllipseBrushFields )

--ellipseArc :: ℝ 2 ~> ℝ 2
--ellipseArc = brushStroke ( linear myPath ) ( uncurryD $ fmap bezier3 myBrush )

--testing :: Double -> Double -> (# Double, T ( ℝ 2 ) #)
--testing :: Double -> Double -> (# Double, T (ℝ 2) #)
--testing t s = envelopeEquation ellipseArc t s