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

module MetaBrush.Asset.Brushes where

-- base
import GHC.Exts
  ( Proxy# )

-- 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
import Math.Linear.Dual
  ( D, type (~>)(..), Differentiable, Diffy(konst), var )
import Math.Module
  ( Module((^+^), (*^)) )
import MetaBrush.Brush
  ( Brush(..), SomeBrush(..), WithParams(..) )
import MetaBrush.Records
  ( Record(MkR) )

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

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 :: forall i u v
             .  Applicative ( D ( I i u ) )
             => ( Double -> Double -> D ( I i u ) ( I i v ) )
             -> D ( I i u ) ( Spline 'Closed () ( I i 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 :: forall i
            .  ( Differentiable i ( Record CircleBrushFields ) )
            => Proxy# i
            -> ( forall a. a -> I i a )
            -> I i ( Record CircleBrushFields ) ~> Spline 'Closed () ( I i ( ℝ 2 ) )
circleBrush _ mkI =
  D \ params ->
    let r :: D ( I i ( Record CircleBrushFields ) ) ( I i Double )
        r = runD ( var ( Fin 1## ) ) params
        mkPt :: Double -> Double -> D ( I i ( Record CircleBrushFields ) ) ( I i ( ℝ 2 ) )
        mkPt ( kon -> x ) ( kon -> y )
          =   ( x * r ) *^ e_x
          ^+^ ( y * r ) *^ e_y
    in circleSpline @i @( Record CircleBrushFields ) @( ℝ 2 ) mkPt
  where
    e_x, e_y :: D ( I i ( Record CircleBrushFields ) ) ( I i ( ℝ 2 ) )
    e_x = pure $ mkI $ ℝ2 1 0
    e_y = pure $ mkI $ ℝ2 0 1

    kon = konst @( I i Double ) @( I i ( Record CircleBrushFields ) ) . mkI

ellipseBrush :: forall i
             .  ( Differentiable i ( Record EllipseBrushFields ) )
             => Proxy# i
             -> ( forall a. a -> I i a )
             -> I i ( Record EllipseBrushFields ) ~> Spline 'Closed () ( I i ( ℝ 2 ) )
ellipseBrush _ mkI =
  D \ params ->
    let a, b, phi :: D ( I i ( Record EllipseBrushFields ) ) ( I i Double )
        a   = runD ( var ( Fin 1## ) ) params
        b   = runD ( var ( Fin 2## ) ) params
        phi = runD ( var ( Fin 3## ) ) params
        mkPt :: Double -> Double -> D ( I i ( Record EllipseBrushFields ) ) ( I i ( ℝ 2 ) )
        mkPt ( kon -> x ) ( kon -> y )
          =   ( x * a * cos phi - y * b * sin phi ) *^ e_x
          ^+^ ( y * b * cos phi + x * a * sin phi ) *^ e_y
    in circleSpline @i @( Record EllipseBrushFields ) @( ℝ 2 ) mkPt
  where
    e_x, e_y :: D ( I i ( Record EllipseBrushFields ) ) ( I i ( ℝ 2 ) )
    e_x = pure $ mkI $ ℝ2 1 0
    e_y = pure $ mkI $ ℝ2 0 1

    kon = konst @( I i Double ) @( I i ( Record EllipseBrushFields ) ) . mkI