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

module MetaBrush.Serialisable
  ( Serialisable(..)
  , encodeSequence, decodeSequence
  , encodeUniqueMap, decodeUniqueMap
  , encodeCurve, decodeCurve
  , encodeCurves, decodeCurves
  , encodeSpline, decodeSpline
  )
  where

-- base
import Control.Arrow
  ( (&&&) )
import Control.Monad.ST
  ( RealWorld, stToIO )
import Data.Foldable
  ( toList )
import Data.Functor
  ( (<&>) )
import Data.Functor.Contravariant
  ( contramap )
import Data.Functor.Identity
  ( Identity(..) )
import Data.STRef
  ( newSTRef )
import Data.Traversable
  ( for )
import GHC.Exts
  ( Word(W#) )

-- containers
import Data.Map.Strict
  ( Map )
import qualified Data.Map.Strict as Map
  ( elems, fromList )
import Data.Sequence
  ( Seq )
import qualified Data.Sequence as Seq
  ( fromList )

-- generic-lens
import Data.Generics.Product.Typed
  ( HasType(typed) )

-- lens
import Control.Lens
  ( view )

-- scientific
import qualified Data.Scientific as Scientific
  ( fromFloatDigits, toRealFloat )

-- text
import Data.Text
  ( Text )

-- transformers
import Control.Monad.IO.Class
  ( MonadIO(liftIO) )
import Control.Monad.Trans.Class
  ( MonadTrans(lift) )

-- waargonaut
import qualified Waargonaut.Decode as JSON
  ( Decoder )
import qualified Waargonaut.Decode as JSON.Decoder
  ( atKey, atKeyOptional, list, scientific, text )
import qualified Waargonaut.Encode as JSON
  ( Encoder )
import qualified Waargonaut.Encode as JSON.Encoder
  ( atKey', keyValueTupleFoldable, list, mapLikeObj, scientific, text, either )

-- meta-brushes
import Math.Bezier.Spline
  ( Spline(..), SplineType(..), SSplineType(..), SplineTypeI(..)
  , Curves(..), Curve(..), NextPoint(..)
  )
import Math.Bezier.Stroke
  ( CachedStroke(..) )
import Math.Linear
  ( ℝ(..), T(..)
  , Fin(..), Representable(tabulate, index)
  )
import MetaBrush.Records
import MetaBrush.Unique
  ( Unique )

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

class Serialisable a where
  encoder :: Monad f => JSON.Encoder f a
  decoder :: Monad m => JSON.Decoder m a

instance Serialisable Double where
  encoder = contramap Scientific.fromFloatDigits JSON.Encoder.scientific
  decoder = fmap      Scientific.toRealFloat     JSON.Decoder.scientific

instance Serialisable ( ℝ 2 ) where
  encoder = JSON.Encoder.mapLikeObj \ ( ℝ2 x y ) ->
      JSON.Encoder.atKey' "x" encoder x
    . JSON.Encoder.atKey' "y" encoder y
  decoder = ℝ2 <$> JSON.Decoder.atKey "x" decoder <*> JSON.Decoder.atKey "y" decoder

instance Serialisable ( T ( ℝ 2 ) ) where
  encoder = JSON.Encoder.mapLikeObj \ ( V2 x y ) ->
      JSON.Encoder.atKey' "x" encoder x
    . JSON.Encoder.atKey' "y" encoder y
  decoder = V2 <$> JSON.Decoder.atKey "x" decoder <*> JSON.Decoder.atKey "y" decoder
instance ( KnownSymbols ks, Representable Double ( ℝ ( Length ks ) ) )
      => Serialisable ( Record ks ) where
  encoder = contramap encodeFields ( JSON.Encoder.keyValueTupleFoldable ( encoder @Double ) )
    where
      encodeFields :: Record ks -> [ ( Text, Double ) ]
      encodeFields ( MkR r ) =
        zip [1..] ( knownSymbols @ks ) <&> \ ( W# i, fld ) ->
          ( fld, index r ( Fin i ) )

  decoder = fmap decodeFields $ for ( knownSymbols @ks ) \ k -> JSON.Decoder.atKey k ( decoder @Double )
    where
      decodeFields :: [ Double ] -> Record ks
      decodeFields coords = MkR $ tabulate \ ( Fin i# ) ->
        coords !! ( fromIntegral ( W# i# ) - 1 )

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

encodeSequence :: Applicative f => JSON.Encoder f a -> JSON.Encoder f ( Seq a )
encodeSequence enc = contramap toList ( JSON.Encoder.list enc )

decodeSequence :: Monad m => JSON.Decoder m a -> JSON.Decoder m ( Seq a )
decodeSequence dec = Seq.fromList <$> JSON.Decoder.list dec



encodeUniqueMap :: Applicative f => JSON.Encoder f a -> JSON.Encoder f ( Map Unique a )
encodeUniqueMap enc = contramap Map.elems ( JSON.Encoder.list enc )

decodeUniqueMap :: ( Monad m, HasType Unique a ) => JSON.Decoder m a -> JSON.Decoder m ( Map Unique a )
decodeUniqueMap dec = Map.fromList . map ( view typed &&& id ) <$> JSON.Decoder.list dec



{-
encodeMat22 :: Applicative f => JSON.Encoder' a -> JSON.Encoder f ( Mat22 a )
encodeMat22 enc = JSON.Encoder.mapLikeObj \ ( Mat22 m00 m01 m10 m11 ) ->
    JSON.Encoder.atKey' "m00" enc m00
  . JSON.Encoder.atKey' "m01" enc m01
  . JSON.Encoder.atKey' "m10" enc m10
  . JSON.Encoder.atKey' "m11" enc m11

decodeMat22 :: Monad m => JSON.Decoder m a -> JSON.Decoder m ( Mat22 a )
decodeMat22 dec =
  Mat22 <$> JSON.Decoder.atKey "m00" dec
        <*> JSON.Decoder.atKey "m01" dec
        <*> JSON.Decoder.atKey "m10" dec
        <*> JSON.Decoder.atKey "m11" dec



encodeAABB :: Applicative f => JSON.Encoder f AABB
encodeAABB = JSON.Encoder.mapLikeObj \ ( AABB { topLeft, botRight } ) ->
    JSON.Encoder.atKey' "topLeft"  enc topLeft
  . JSON.Encoder.atKey' "botRight" enc botRight
  where
    enc :: JSON.Encoder' ( Point2D Double )
    enc = encodePoint2D encodeDouble

decodeAABB :: forall m. Monad m => JSON.Decoder m AABB
decodeAABB = do
    topLeft  <- JSON.Decoder.atKey "topLeft"  dec
    botRight <- JSON.Decoder.atKey "botRight" dec
    pure ( AABB { topLeft, botRight } )
  where
    dec :: JSON.Decoder m ( Point2D Double )
    dec = decodePoint2D decodeDouble
-}


encodeCurve
  :: forall clo crvData ptData f
  .  ( SplineTypeI clo, Applicative f )
  => JSON.Encoder Identity ptData
  -> JSON.Encoder f ( Curve clo crvData ptData )
encodeCurve  encodePtData = case ssplineType @clo of
  SOpen   -> JSON.Encoder.mapLikeObj \case
    LineTo          ( NextPoint p1 ) _ ->
        JSON.Encoder.atKey' "p1" encodePtData p1
    Bezier2To    p1 ( NextPoint p2 ) _ ->
        JSON.Encoder.atKey' "p1" encodePtData p1
      . JSON.Encoder.atKey' "p2" encodePtData p2
    Bezier3To p1 p2 ( NextPoint p3 ) _ ->
        JSON.Encoder.atKey' "p1" encodePtData p1
      . JSON.Encoder.atKey' "p2" encodePtData p2
      . JSON.Encoder.atKey' "p3" encodePtData p3
  SClosed -> JSON.Encoder.mapLikeObj \case
    LineTo          BackToStart _ -> id
    Bezier2To    p1 BackToStart _ ->
        JSON.Encoder.atKey' "p1" encodePtData p1
    Bezier3To p1 p2 BackToStart _ ->
        JSON.Encoder.atKey' "p1" encodePtData p1
      . JSON.Encoder.atKey' "p2" encodePtData p2

decodeCurve
  :: forall clo ptData m
  .  ( SplineTypeI clo, MonadIO m )
  => JSON.Decoder m ptData
  -> JSON.Decoder m ( Curve clo ( CachedStroke RealWorld ) ptData )
decodeCurve decodePtData = do
  noCache <- lift . liftIO . stToIO $ CachedStroke <$> newSTRef Nothing
  case ssplineType @clo of
    SOpen   -> do
      p1    <- JSON.Decoder.atKey         "p1" decodePtData
      mb_p2 <- JSON.Decoder.atKeyOptional "p2" decodePtData
      case mb_p2 of
        Nothing ->
          pure ( LineTo ( NextPoint p1 ) noCache )
        Just p2 -> do
          mb_p3 <- JSON.Decoder.atKeyOptional "p3" decodePtData
          case mb_p3 of
            Nothing -> pure ( Bezier2To    p1 ( NextPoint p2 ) noCache )
            Just p3 -> pure ( Bezier3To p1 p2 ( NextPoint p3 ) noCache )
    SClosed -> do
      mb_p1 <- JSON.Decoder.atKeyOptional "p1" decodePtData
      case mb_p1 of
        Nothing ->
          pure ( LineTo BackToStart noCache )
        Just p1 -> do
          mb_p2 <- JSON.Decoder.atKeyOptional "p2" decodePtData
          case mb_p2 of
            Nothing -> pure ( Bezier2To    p1 BackToStart noCache )
            Just p2 -> pure ( Bezier3To p1 p2 BackToStart noCache )



encodeCurves
  :: forall clo crvData ptData f
  .  ( SplineTypeI clo, Applicative f )
  => JSON.Encoder Identity ptData
  -> JSON.Encoder f ( Curves clo crvData ptData )
encodeCurves encodePtData = case ssplineType @clo of
  SOpen   -> contramap ( openCurves ) ( encodeSequence $ encodeCurve @Open encodePtData )
  SClosed -> contramap ( \case { NoCurves -> Left (); ClosedCurves prevs lst -> Right ( prevs, lst ) } ) ( JSON.Encoder.either encodeL encodeR )
    where
      encodeL :: JSON.Encoder f ()
      encodeL = contramap ( const "NoCurves" ) JSON.Encoder.text
      encodeR :: JSON.Encoder f ( Seq ( Curve Open crvData ptData ), Curve Closed crvData ptData )
      encodeR = JSON.Encoder.mapLikeObj \ ( openCurves, closedCurve ) ->
          JSON.Encoder.atKey' "prevOpenCurves"  ( encodeSequence $ encodeCurve @Open   encodePtData ) openCurves
        . JSON.Encoder.atKey' "lastClosedCurve"                  ( encodeCurve @Closed encodePtData ) closedCurve

decodeCurves
  :: forall clo ptData m
  .  ( SplineTypeI clo, MonadIO m )
  => JSON.Decoder m ptData
  -> JSON.Decoder m ( Curves clo ( CachedStroke RealWorld ) ptData )
decodeCurves decodePtData = case ssplineType @clo of
  SOpen   -> OpenCurves <$> decodeSequence ( decodeCurve @Open decodePtData )
  SClosed -> do
    mbNoCurves <- JSON.Decoder.atKeyOptional "NoCurves" ( JSON.Decoder.text )
    case mbNoCurves of
      Just _ -> pure NoCurves
      Nothing -> do
        prevCurves <- JSON.Decoder.atKey "prevOpenCurves"  ( decodeSequence $ decodeCurve @Open   decodePtData )
        lastCurve  <- JSON.Decoder.atKey "lastClosedCurve"                  ( decodeCurve @Closed decodePtData )
        pure ( ClosedCurves prevCurves lastCurve )



encodeSpline
  :: forall clo crvData ptData f
  .  ( SplineTypeI clo, Applicative f )
  => JSON.Encoder Identity ptData
  -> JSON.Encoder f ( Spline clo crvData ptData )
encodeSpline encodePtData = JSON.Encoder.mapLikeObj \ ( Spline { splineStart, splineCurves } ) ->
    JSON.Encoder.atKey' "splineStart"  encodePtData                       splineStart
  . JSON.Encoder.atKey' "splineCurves" ( encodeCurves @clo encodePtData ) splineCurves

decodeSpline
  :: forall clo ptData m
  .  ( SplineTypeI clo, MonadIO m )
  => JSON.Decoder m ptData
  -> JSON.Decoder m ( Spline clo ( CachedStroke RealWorld ) ptData )
decodeSpline decodePtData = do
  splineStart  <- JSON.Decoder.atKey "splineStart"  decodePtData
  splineCurves <- JSON.Decoder.atKey "splineCurves" ( decodeCurves @clo decodePtData )
  pure ( Spline { splineStart, splineCurves } )