{-# LANGUAGE Arrows #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTSyntax #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
-- {-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
-- {-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE NoMonoLocalBinds #-}

module Midair.Core (
   -- The FRP:
     SFlow
--   , SFlow'
--   , SFNodeRef
--   , IsNodeRef(..)

   , sMap
   , sFold
   , sCompose
   , sMapMaybe
   , sJoin
   , sFilter
   , sFilterWDefault

   , mkNodeRef
   , hotSwap
   , hotSwapSTM
--   , replaceSFlow
--   , replaceSFlowSTM
   , fireGraph
   , fireGraphIO
   ) where

import Control.Arrow
import qualified Control.Category as Cat
-- import Control.Concurrent.Async (race)
import Control.Concurrent.STM
import Data.Maybe
-- import Data.Proxy
-- import Data.Type.Equality
-- import Data.Unique

{-
data IsNodeRef = True_IsNodeRef | False_IsNodeRef
 deriving (Show, Eq, Read, Ord)

type SFlow = SFlow' 'False_IsNodeRef
type SFNodeRef = SFlow' 'True_IsNodeRef

data SFlow' (i :: IsNodeRef) a c where
-}

-- | Signal flow
-- 
--   Takes a value of type 'a' and returns one of type 'c'
data SFlow a c where
   -- Transformation:
   SF_Map :: (a -> c) -> SFlow a c
   -- Composition:
   SF_Compose :: forall a b c. SFlow b c -> SFlow a b -> SFlow a c
   -- Accumulation:
   SF_FoldP :: c -> (a -> c -> c) -> SFlow a c
   -- Recombination:
   SF_Join :: forall a b c. SFlow a b -> SFlow a c -> SFlow a (b, c)
   -- Filtration/selection:
   SF_Filter :: Maybe a -> (a -> Bool) -> SFlow a (Maybe a)
   -- Reference to a part of the graph:
   SF_NodeRef :: TVar (Maybe c) -> TVar (SFlow a c) -> SFlow a c

mkNodeRefSTM :: SFlow i o -> STM (SFlow i o)
mkNodeRefSTM sigNode = do
   sigNodeTVar <- newTVar sigNode
   foo <- newTVar Nothing
   return $ SF_NodeRef foo sigNodeTVar

mkNodeRef :: SFlow i o -> IO (SFlow i o)
mkNodeRef g = atomically $ mkNodeRefSTM g

-- | Given a node in the graph, and an input to that node, return the output of
--   that node and the \"new node\" with updated state
fireGraph :: TVar (SFlow a c) -> a -> STM c
fireGraph graphTVar inVal = do
   (newGraph, retVal) <- (flip fireGraph') inVal =<< readTVar graphTVar
   writeTVar graphTVar newGraph
   return retVal
  

-- The only action we do in STM here is to call fireGraph (no prime)
-- whenever we encounter a 'SF_NodeRef':
fireGraph' :: SFlow a c -> a -> STM (SFlow a c, c)
fireGraph' m@(SF_Map f) a =
   return (m, f a)
fireGraph' (SF_Compose bToCOld aToBOld) a = do
   (aToBNew, rightVal) <- fireGraph' aToBOld a
   (bToCNew, returnVal) <- fireGraph' bToCOld rightVal
   return (SF_Compose bToCNew aToBNew, returnVal)
fireGraph' (SF_FoldP lastVal f) a =
   let newVal = f a lastVal
   in return (SF_FoldP newVal f, newVal)
fireGraph' (SF_Join aToBOld aToCOld) a = do
   (aToBNew, b) <- fireGraph' aToBOld a
   (aToCNew, c) <- fireGraph' aToCOld a
   return (SF_Join aToBNew aToCNew, (b, c))
fireGraph' (SF_Filter previousVal filterF) newVal = return $
   if filterF newVal
      then (SF_Filter (Just newVal) filterF, Just newVal)
      else (SF_Filter previousVal filterF, previousVal)
fireGraph' (SF_NodeRef prevOutVar graphRef) newIn = do
   newOut <- fireGraph graphRef newIn
   writeTVar prevOutVar $ Just newOut
   return (SF_NodeRef prevOutVar graphRef, newOut)

fireGraphIO :: TVar (SFlow a c) -> a -> IO c
fireGraphIO graphTVar inVal =
   atomically $ fireGraph graphTVar inVal

sMap :: (a -> c) -> SFlow a c
sMap = SF_Map

sCompose :: SFlow b c -> SFlow a b -> SFlow a c
sCompose = SF_Compose

sFold :: c -> (a -> c -> c) -> SFlow a c
sFold = SF_FoldP

sMapMaybe :: state -> (update -> Maybe state) -> SFlow update state
sMapMaybe startVal maybeF =
   sFold startVal f
 where
   f update previousVal = case maybeF update of
       Just x -> x
       Nothing -> previousVal

sJoin :: SFlow a b -> SFlow a c -> SFlow a (b, c)
sJoin = SF_Join

sFilter :: (b -> Bool) -> SFlow b (Maybe b)
sFilter f = SF_Filter Nothing f

-- Want a shorter name:
sFilterWDefault :: (b -> Bool) -> b -> SFlow b b
sFilterWDefault f defaultVal =
   sMap (fromMaybe defaultVal) <<< sFilter f

instance Functor (SFlow a) where
   fmap f0 (SF_Map f1) = SF_Map (f0 . f1)
   fmap f0 sf@(SF_FoldP _ _) = SF_Compose (SF_Map f0) sf
   fmap f0 sf@(SF_Compose _ _) = SF_Compose (SF_Map f0) sf
   fmap f0 sf@(SF_Join _ _) = SF_Compose (SF_Map f0) sf
   fmap f0 sf@(SF_Filter _ _) = SF_Compose (SF_Map f0) sf
   fmap f0 sf@(SF_NodeRef _ _) = SF_Compose (SF_Map f0) sf

instance Applicative (SFlow a) where
   pure x = SF_Map $ \_ -> x
   (<*>) a b = SF_Compose (SF_Map (\(f, x) -> f x)) (SF_Join a b)

instance Cat.Category SFlow where
   id = SF_Map $ \x -> x
   (.) a b = SF_Compose a b

instance Arrow SFlow where
   arr = SF_Map
   first :: SFlow b c -> SFlow (b, d) (c, d)
   first bToC =
      sJoin (bToC `sCompose` arr fst) (arr snd)

{-
-- | Rationale: why do we do this instead of recomputing from the beginning?
-- 
--   1) Practicality: it would be expensive to save all input values
--   2) Meaning: if you e.g. have a video game that's depermining a player's
--      position based on their keypresses, if you were to e.g. recompute all steps
--      (let's say the keypresses move 3 pixels now instead of 2), your player
--      would end up in a totally new position
-}

hotSwapSTM :: SFlow a c -> SFlow a c -> STM ()
hotSwapSTM (SF_NodeRef oldValMaybeTVar graphTVar) newGraph = do
   oldValMaybe <- readTVar oldValMaybeTVar
   writeTVar graphTVar =<< tryInsertLastVal oldValMaybe newGraph
hotSwapSTM _ _ = return () -- temp

hotSwap :: SFlow a c -> SFlow a c -> IO ()
hotSwap a@(SF_NodeRef _ _) b = atomically $ hotSwapSTM a b
hotSwap _ _ = putStrLn "Warning! Attempting to hot-swap a non-SFNodeRef"

{-
replaceSFlowSTM :: SFlow a c -> SFlow a c -> STM ()
replaceSFlowSTM = undefined

-- | Like 'hotSwap' but don't preserve any state
replaceSFlow :: SFlow i o -> SFlow i o -> IO ()
replaceSFlow a b = atomically $ replaceSFlowSTM a b
-}

-- The 'STM' is only for SF_NodeRef:
tryInsertLastVal :: Maybe b -> SFlow a b -> STM (SFlow a b)
tryInsertLastVal Nothing = pure . id
tryInsertLastVal lv@(Just lastVal) = \case
   m@(SF_Map _) -> pure m
   SF_FoldP _oldVal f -> pure $ SF_FoldP lastVal f
   filt@(SF_Filter _ filtF) -> case lastVal of
      Just x -> if filtF x
         then pure $ SF_Filter (Just x) filtF
         else pure filt
      Nothing -> pure filt
   _c@(SF_Compose a b) ->
      SF_Compose <$> (tryInsertLastVal lv a) <*> pure b
{-
      | holdsState a -> SF_Compose (insertLastVal lv a) b
      -- May not have the same return type:
      -- | holdsState b -> SF_Compose a (insertLastVal lv b)
      | otherwise -> c
-}
   SF_Join a b -> case lastVal of
      (aVal,bVal) -> SF_Join
         <$> (tryInsertLastVal (Just aVal) a)
         <*> (tryInsertLastVal (Just bVal) b)
   nv@(SF_NodeRef lastValTVar graphTVar) -> do
       -- Maybe only do this if the 'tryInsert' inserts something:
      writeTVar lastValTVar $ Just lastVal

      oldGraph <- readTVar graphTVar
      writeTVar graphTVar =<< tryInsertLastVal lv oldGraph
      return nv

{-
holdsState :: SFlow a b -> Bool
holdsState = \case
   SF_Map _ -> False
   SF_FoldP _ _-> True
   -- SF_Filter _ filterF -> filterF a
   SF_Filter _ _ -> True
   SF_Compose a _b -> holdsState a -- || holdsState b
   SF_Join a b -> holdsState a && holdsState b
-}