--------------------------------------------------------------------------------
{-# LANGUAGE TemplateHaskell #-}
module Patat.Transition.Matrix
    ( transition
    ) where


--------------------------------------------------------------------------------
import           Control.Monad             (forM_, guard, when)
import qualified Data.Aeson.Extended       as A
import qualified Data.Aeson.TH.Extended    as A
import           Data.Bifunctor            (first)
import qualified Data.Vector               as V
import qualified Data.Vector.Mutable       as VM
import           Patat.PrettyPrint.Matrix
import           Patat.Size                (Size (..))
import           Patat.Transition.Internal
import           System.Random.Stateful


--------------------------------------------------------------------------------
data Config = Config
    { Config -> Maybe (FlexibleNum Double)
cDuration  :: Maybe (A.FlexibleNum Double)
    , Config -> Maybe (FlexibleNum Int)
cFrameRate :: Maybe (A.FlexibleNum Int)
    }


--------------------------------------------------------------------------------
data Particle = Particle
    { Particle -> Double
pX        :: Double
    , Particle -> Double
pInitialY :: Double
    , Particle -> Double
pFinalY   :: Double
    , Particle -> Double
pSpeed    :: Double
    , Particle -> Cell
pCell     :: Cell
    }


--------------------------------------------------------------------------------
particleY :: Particle -> Double -> Double
particleY :: Particle -> Double -> Double
particleY Particle
p Double
t = Particle -> Double
pInitialY Particle
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t') Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Particle -> Double
pFinalY Particle
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t'
  where
    t' :: Double
t' = Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
1 (Particle -> Double
pSpeed Particle
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t)


--------------------------------------------------------------------------------
-- | Maximum speed of a particle, expressed as a factor of the minimum speed of
-- a particle.
particleMaxSpeed :: Double
particleMaxSpeed :: Double
particleMaxSpeed = Double
2


--------------------------------------------------------------------------------
-- | Number of ghosts a particle leaves behind.  Currently hardcoded but could
-- be moved to config.
particleGhosts :: Int
particleGhosts :: Int
particleGhosts = Int
3


--------------------------------------------------------------------------------
transition :: Config -> TransitionGen
transition :: Config -> TransitionGen
transition Config
config (Size Int
rows Int
cols) Matrix
initial Matrix
final StdGen
rgen =
    (Double -> Matrix) -> (Double, Duration) -> (Matrix, Duration)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Double -> Matrix
frame ((Double, Duration) -> (Matrix, Duration))
-> NonEmpty (Double, Duration) -> NonEmpty (Matrix, Duration)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
    Maybe Double -> Maybe Int -> NonEmpty (Double, Duration)
evenlySpacedFrames
        (FlexibleNum Double -> Double
forall a. FlexibleNum a -> a
A.unFlexibleNum (FlexibleNum Double -> Double)
-> Maybe (FlexibleNum Double) -> Maybe Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Config -> Maybe (FlexibleNum Double)
cDuration  Config
config)
        (FlexibleNum Int -> Int
forall a. FlexibleNum a -> a
A.unFlexibleNum (FlexibleNum Int -> Int) -> Maybe (FlexibleNum Int) -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Config -> Maybe (FlexibleNum Int)
cFrameRate Config
config)
  where
    speeds :: V.Vector Double
    speeds :: Vector Double
speeds = StdGen
-> (StateGenM StdGen -> State StdGen (Vector Double))
-> Vector Double
forall g a. RandomGen g => g -> (StateGenM g -> State g a) -> a
runStateGen_ StdGen
rgen ((StateGenM StdGen -> State StdGen (Vector Double))
 -> Vector Double)
-> (StateGenM StdGen -> State StdGen (Vector Double))
-> Vector Double
forall a b. (a -> b) -> a -> b
$ \StateGenM StdGen
g ->
        Int
-> StateT StdGen Identity Double -> State StdGen (Vector Double)
forall (m :: * -> *) a. Monad m => Int -> m a -> m (Vector a)
V.replicateM (Int
rows Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cols) ((Double, Double)
-> StateGenM StdGen -> StateT StdGen Identity Double
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *).
StatefulGen g m =>
(Double, Double) -> g -> m Double
uniformRM (Double
1, Double
particleMaxSpeed) StateGenM StdGen
g)

    up :: V.Vector Bool
    up :: Vector Bool
up = StdGen
-> (StateGenM StdGen -> State StdGen (Vector Bool)) -> Vector Bool
forall g a. RandomGen g => g -> (StateGenM g -> State g a) -> a
runStateGen_ StdGen
rgen ((StateGenM StdGen -> State StdGen (Vector Bool)) -> Vector Bool)
-> (StateGenM StdGen -> State StdGen (Vector Bool)) -> Vector Bool
forall a b. (a -> b) -> a -> b
$ \StateGenM StdGen
g ->
        Int -> StateT StdGen Identity Bool -> State StdGen (Vector Bool)
forall (m :: * -> *) a. Monad m => Int -> m a -> m (Vector a)
V.replicateM (Int
rows Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cols) (StateGenM StdGen -> StateT StdGen Identity Bool
forall a g (m :: * -> *). (Uniform a, StatefulGen g m) => g -> m a
forall g (m :: * -> *). StatefulGen g m => g -> m Bool
uniformM StateGenM StdGen
g)

    ghosts :: Double -> [Double]
    ghosts :: Double -> [Double]
ghosts Double
baseSpeed =
        [ Double
baseSpeed Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
particleGhosts)
        | Int
i <- [Int
0 .. Int
particleGhosts]
        ]

    initialParticles :: [Particle]
    initialParticles :: [Particle]
initialParticles = do
        (Int
x, Int
y, Cell
cell) <- Matrix -> [(Int, Int, Cell)]
posCells Matrix
initial
        let idx :: Int
idx = Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cols Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x
        Double
speed <- Double -> [Double]
ghosts (Double -> [Double]) -> Double -> [Double]
forall a b. (a -> b) -> a -> b
$ Vector Double
speeds Vector Double -> Int -> Double
forall a. Vector a -> Int -> a
V.! Int
idx
        Particle -> [Particle]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure Particle
            { pX :: Double
pX        = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x
            , pInitialY :: Double
pInitialY = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
y
            , pFinalY :: Double
pFinalY   = if Vector Bool
up Vector Bool -> Int -> Bool
forall a. Vector a -> Int -> a
V.! Int
idx then Double
0 else Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
rows
            , pSpeed :: Double
pSpeed    = Double
speed
            , pCell :: Cell
pCell     = Cell
cell
            }

    finalParticles :: [Particle]
    finalParticles :: [Particle]
finalParticles = do
        (Int
x, Int
y, Cell
cell) <- Matrix -> [(Int, Int, Cell)]
posCells Matrix
final
        let idx :: Int
idx = Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cols Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x
        Double
speed <- Double -> [Double]
ghosts (Double -> [Double]) -> Double -> [Double]
forall a b. (a -> b) -> a -> b
$ Vector Double
speeds Vector Double -> Int -> Double
forall a. Vector a -> Int -> a
V.! Int
idx
        Particle -> [Particle]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure Particle
            { pX :: Double
pX        = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x
            , pInitialY :: Double
pInitialY = if Vector Bool
up Vector Bool -> Int -> Bool
forall a. Vector a -> Int -> a
V.! Int
idx then -Double
1 else Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
rows
            , pFinalY :: Double
pFinalY   = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
y
            , pSpeed :: Double
pSpeed    = Double
speed
            , pCell :: Cell
pCell     = Cell
cell
            }

    posCells :: Matrix -> [(Int, Int, Cell)]
    posCells :: Matrix -> [(Int, Int, Cell)]
posCells Matrix
mat = do
        Int
y <- [Int
0 .. Int
rows Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        Int
x <- [Int
0 .. Int
cols Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        let cell :: Cell
cell = Matrix
mat Matrix -> Int -> Cell
forall a. Vector a -> Int -> a
V.! (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cols Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x)
        Bool -> [()]
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> [()]) -> (Bool -> Bool) -> Bool -> [()]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Bool
not (Bool -> [()]) -> Bool -> [()]
forall a b. (a -> b) -> a -> b
$ Cell
cell Cell -> Cell -> Bool
forall a. Eq a => a -> a -> Bool
== Cell
emptyCell
        (Int, Int, Cell) -> [(Int, Int, Cell)]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
x, Int
y, Cell
cell)

    frame :: Double -> Matrix
    frame :: Double -> Matrix
frame Double
t = (forall s. ST s (MVector s Cell)) -> Matrix
forall a. (forall s. ST s (MVector s a)) -> Vector a
V.create ((forall s. ST s (MVector s Cell)) -> Matrix)
-> (forall s. ST s (MVector s Cell)) -> Matrix
forall a b. (a -> b) -> a -> b
$ do
        MVector s Cell
mat <- Int -> Cell -> ST s (MVector (PrimState (ST s)) Cell)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MVector (PrimState m) a)
VM.replicate (Int
rows Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cols) Cell
emptyCell
        [Particle] -> (Particle -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Particle]
initialParticles [Particle] -> [Particle] -> [Particle]
forall a. [a] -> [a] -> [a]
++ [Particle]
finalParticles) ((Particle -> ST s ()) -> ST s ())
-> (Particle -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Particle
particle ->
            let y :: Int
y = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Particle -> Double -> Double
particleY Particle
particle Double
t
                x :: Int
x = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Particle -> Double
pX Particle
particle
                idx :: Int
idx = Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cols Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x in
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
cols Bool -> Bool -> Bool
&& Int
y Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
y Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
rows) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
                MVector (PrimState (ST s)) Cell -> Int -> Cell -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
VM.write MVector s Cell
MVector (PrimState (ST s)) Cell
mat Int
idx (Cell -> ST s ()) -> Cell -> ST s ()
forall a b. (a -> b) -> a -> b
$ Particle -> Cell
pCell Particle
particle
        MVector s Cell -> ST s (MVector s Cell)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Cell
mat


--------------------------------------------------------------------------------
$(A.deriveFromJSON A.dropPrefixOptions ''Config)