-- | You should not need to import this module unless you're adding support
-- for a new model of Arduino, or an Arduino library.

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Copilot.Arduino.Internals where

import Language.Copilot
import Control.Monad.Writer
import Control.Monad.State.Strict
import Data.Functor.Identity
import qualified Data.Map as M
import qualified Data.Set as S
import Data.Type.Bool
import Data.Proxy
import GHC.TypeLits

-- | A value that changes over time.
--
-- This is implemented as a `Stream` in the Copilot DSL.
-- Copilot provides many operations on streams, for example
-- `Language.Copilot.&&` to combine two streams of Bools.
-- 
-- For documentation on using the Copilot DSL, see
-- <https://copilot-language.github.io/>
type Behavior t = Stream t

-- | A Behavior with an additional phantom type `p`.
--
-- The Compilot DSL only lets a Stream contain basic C types,
-- a limitation that `Behavior` also has. When more type safely
-- is needed, this can be used.
data TypedBehavior p t = TypedBehavior (Behavior t)

-- | A discrete event, that occurs at particular points in time.
data Event p v = Event v (Stream Bool)

-- | An Arduino sketch, implemented using Copilot.
--
-- It's best to think of the `Sketch` as a description of the state of the
-- Arduino at any point in time.
--
-- Under the hood, the `Sketch` is run in a loop. On each iteration, it first
-- reads inputs and then updates outputs as needed.
--
-- While it is a monad, a Sketch's outputs are not updated in any
-- particular order, because Copilot does not guarantee any order.
newtype Sketch t = Sketch (WriterT [(TriggerLimit -> Spec, TriggerLimit -> Framework)] (State UniqueIds) t)
	deriving 
		( Monad
		, Applicative
		, Functor
		, MonadWriter [(TriggerLimit -> Spec, TriggerLimit -> Framework)]
		, MonadState UniqueIds
		)

instance Monoid (Sketch ()) where
	mempty = Sketch (return ())

instance Semigroup (Sketch t) where
	(Sketch a) <> (Sketch b) = Sketch (a >> b)

newtype UniqueIds = UniqueIds (M.Map String Integer)

newtype UniqueId = UniqueId Integer

data TriggerLimit
	= TriggerLimit (Behavior Bool)
	| NoTriggerLimit

getTriggerLimit :: TriggerLimit -> Behavior Bool
getTriggerLimit (TriggerLimit b) = b
getTriggerLimit NoTriggerLimit = true

addTriggerLimit :: TriggerLimit -> Behavior Bool -> Behavior Bool
addTriggerLimit tl c = getTriggerLimit (tl <> TriggerLimit c)

instance Monoid TriggerLimit where
	mempty = NoTriggerLimit

instance Semigroup TriggerLimit where
	TriggerLimit a <> TriggerLimit b =
		TriggerLimit (a Language.Copilot.&& b)
	a <> NoTriggerLimit = a
	NoTriggerLimit <> b = b

evalSketch :: Sketch a -> (Maybe Spec, Framework)
evalSketch (Sketch s) = (spec, f)
  where
	(is, fs) = unzip $ 
		runIdentity $ evalStateT (execWriterT s) (UniqueIds mempty)
	f = mconcat (map (\f' -> f' NoTriggerLimit) fs)
	-- Copilot will throw an ugly error if given a spec that does
	-- nothing at all, so return Nothing to avoid that.
	spec :: Maybe Spec
	spec = if null is
		then Nothing
		else Just $ sequence_ $ map (\i -> i NoTriggerLimit) is

-- | Limit the effects of a `Sketch` to times when a `Behavior` `Bool` is True.
--
-- When applied to `=:`, this does the same thing as `@:` but without
-- the FRP style conversion the input `Behavior` into an `Event`. So `@:`
-- is generally better to use than this.
--
-- But, this can also be applied to `input`, to limit how often input
-- gets read. Useful to avoid performing slow input operations on every
-- iteration of a Sketch.
--
-- > v <- whenB (frequency 10) $ input pin12
--
-- (It's best to think of the value returned by that as an Event,
-- but it's currently represented as a Behavior, since the Copilot DSL
-- cannot operate on Events.)
whenB :: Behavior Bool -> Sketch t -> Sketch t
whenB c (Sketch s) = do
	ids <- get
	let ((r, w), ids') = runIdentity $ runStateT (runWriterT s) ids
	put ids'
	let (is, fs) = unzip w
	let spec = combinetl $ \c' -> sequence_ (map (\i -> i c') is)
	tell [(spec, mempty)]
	forM_ fs $ \f -> tell [(const (return ()), combinetl f)]
	return r
  where
	combinetl :: (TriggerLimit -> a) -> TriggerLimit -> a
	combinetl g tl = g (TriggerLimit c <> tl)

-- | Gets a unique id.
getUniqueId :: String -> Sketch UniqueId
getUniqueId s = do
	UniqueIds m <- get
	let u = maybe 1 succ (M.lookup s m)
	put $ UniqueIds $ M.insert s u m
	return (UniqueId u)

-- | Generates a unique name.
uniqueName :: String -> UniqueId -> String
uniqueName s (UniqueId i)
	| i Prelude.== 1 = s
	| otherwise = s <> "_" <>  show i

uniqueName' :: String -> UniqueId -> String
uniqueName' s (UniqueId i) = s <> "_" <>  show i

-- | The framework of an Arduino sketch.
data Framework = Framework
	{ defines :: [CChunk]
	-- ^ Things that come before the C code generated by Copilot.
	, setups :: [CChunk]
	-- ^ Things to do at setup, not including configuring pins.
	, earlySetups :: [CChunk]
	-- ^ Things to do at setup, before the setups.
	, pinmodes :: M.Map PinId (S.Set PinMode)
	-- ^ How pins are used.
	, loops :: [CChunk]
	-- ^ Things to run in `loop`.
	}

instance Semigroup Framework where
	a <> b = Framework
		{ defines = defines a <> defines b
		, setups = setups a <> setups b
		, earlySetups = earlySetups a <> earlySetups b
		, pinmodes = M.unionWith S.union (pinmodes a) (pinmodes b)
		, loops = loops a  <> loops b
		}

instance Monoid Framework where
	mempty = Framework mempty mempty mempty mempty mempty

-- | A line of C code.
newtype CLine = CLine { fromCLine :: String }
	deriving (Eq, Show, Ord)

-- | A chunk of C code. Identical chunks get deduplicated.
newtype CChunk = CChunk [CLine]
	deriving (Eq, Show, Ord, Semigroup, Monoid)

mkCChunk :: [CLine] -> [CChunk]
mkCChunk l = [CChunk l]

-- | Copilot only supports calling a trigger with a given name once
-- per Spec; the generated C code will fail to build if the same name is
-- used in two triggers. This generates a unique alias that can be 
-- used in a trigger.
defineTriggerAlias :: String -> Framework -> Sketch (Framework, String)
defineTriggerAlias = defineTriggerAlias' ""

defineTriggerAlias' :: String -> String -> Framework -> Sketch (Framework, String)
defineTriggerAlias' suffix cfuncname f = do
	let basetname = if null suffix 
		then cfuncname
		else cfuncname <> "_" <> suffix
	u <- getUniqueId basetname
	let triggername = uniqueName basetname u
	let define = if cfuncname Prelude./= triggername
		then mkCChunk [ CLine $ "#define " <> triggername <> " " <> cfuncname	]
		else mempty
	return (f { defines = define <> defines f }, triggername)

data InputSource t = InputSource
	{ defineVar :: [CChunk]
	-- ^ Added to the `Framework`'s `defines`, this typically
	-- defines a C variable.
	, setupInput :: [CChunk]
	-- ^ How to set up the input, not including pin mode.
	, inputPinmode :: M.Map PinId PinMode
	-- ^ How pins are used by the input.
	, readInput :: [CChunk]
	-- ^ How to read a value from the input, this typically
	-- reads a value into a C variable.
	, inputStream :: Stream t
	-- ^ How to use Copilot's extern to access the input values.
	}

mkInput :: InputSource t -> Sketch (Behavior t)
mkInput i = do
	u <- getUniqueId "input"
	tell [(mkspec u, f u)]
	return (inputStream i)
  where
	f u ratelimited = Framework
		{ defines = defineVar i <> mkdefine u ratelimited
		, setups = setupInput i
		, earlySetups = mempty
		, pinmodes = M.map S.singleton (inputPinmode i)
		, loops = mkloops u ratelimited (readInput i)
		}

	varname = uniqueName "update_input"
	triggername = uniqueName "input"
	
	mkdefine _ NoTriggerLimit = []
	mkdefine u (TriggerLimit _) = mkCChunk $ map CLine
		[ "bool " <> varname u <> " = true;"
		, "void " <> triggername u <> " (bool v) {"
		, "  " <> varname u <> " = v;"
		, "}"
		]
	
	mkloops _ NoTriggerLimit reader = reader
	mkloops u (TriggerLimit _) reader = mkCChunk $ concat
		[ [ CLine $ "if (" <> varname u <> ") {" ]
		, map (\(CLine l) -> CLine $ "  " <> l ) readerlines
		, [ CLine "}" ]
		]
	  where
		readerlines = concatMap (\(CChunk l) -> l) reader

	mkspec _ NoTriggerLimit = return ()
	mkspec u (TriggerLimit c) = trigger (triggername u) true [arg c]

-- | A pin on the Arduino board.
--
-- For definitions of pins like `Copilot.Arduino.Uno.pin12`, 
-- load a module such as Copilot.Arduino.Uno, which provides the pins of a
-- particular board.
--
-- A type-level list indicates how a Pin can be used, so the haskell
-- compiler will detect impossible uses of pins.
newtype Pin t = Pin PinId
	deriving (Show, Eq, Ord)

newtype PinId = PinId Int16
	deriving (Show, Eq, Ord)

data PinCapabilities
	= DigitalIO
	| AnalogInput
	| PWM
	deriving (Show, Eq, Ord)

type family IsDigitalIOPin t where
	IsDigitalIOPin t = 
		'True ~ If (HasPinCapability 'DigitalIO t)
			('True)
			(TypeError ('Text "This Pin does not support digital IO"))

type family IsAnalogInputPin t where
	IsAnalogInputPin t = 
		'True ~ If (HasPinCapability 'AnalogInput t)
			('True)
			(TypeError ('Text "This Pin does not support analog input"))

type family IsPWMPin t where
	IsPWMPin t = 
		'True ~ If (HasPinCapability 'PWM t)
			('True)
			(TypeError ('Text "This Pin does not support PWM"))

type family HasPinCapability (c :: t) (list :: [t]) :: Bool where
	HasPinCapability c '[] = 'False
	HasPinCapability c (x ': xs) = SameCapability c x || HasPinCapability c xs

type family SameCapability a b :: Bool where
	SameCapability 'DigitalIO 'DigitalIO = 'True
	SameCapability 'AnalogInput 'AnalogInput = 'True
	SameCapability 'PWM 'PWM = 'True
	SameCapability _ _ = 'False

data PinMode = InputMode | InputPullupMode | OutputMode
	deriving (Show, Eq, Ord)

-- | Things that can have a `Behavior` or `Event` output to them.
class Output o t where
	(=:) :: o -> t -> Sketch ()
	-- ^ Connect a `Behavior` or `Event` to an `Output`
	-- 
	-- > led =: blinking
	-- 
	-- When a `Behavior` is used, its current value is written on each
	-- iteration of the `Sketch`. 
	-- 
	-- For example, this constantly turns on the LED, even though it will
	-- already be on after the first iteration, because `true`
	-- is a `Behavior` (that is always True).
	-- 
	-- > led =: true
	-- 
	-- To avoid unncessary work being done, you can use an `Event`
	-- instead. Then the write only happens at the points in time
	-- when the `Event` occurs. To turn a `Behavior` into an `Event`,
	-- use `@:`
	-- 
	-- So to make the LED only be turned on in the first iteration,
	-- and allow it to remain on thereafter without doing extra work:
	-- 
	-- > led =: true @: firstIteration

-- Same fixity as =<<
infixr 1 =:

instance Output o (Event () (Stream v)) => Output o (Behavior v) where
	(=:) o b = o =: te
	  where
	  	te :: Event () (Stream v)
		te = Event b true

instance Output o (Event p (Stream v)) => Output o (TypedBehavior p v) where
	(=:) o (TypedBehavior b) = o =: te
	  where
		te :: Event p (Stream v)
		te = Event b true

-- | This type family is open, so it can be extended when adding other data
-- types to the IsBehavior class.
type family BehaviorToEvent a
type instance BehaviorToEvent (Behavior v) = Event () (Stream v)
type instance BehaviorToEvent (TypedBehavior p v) = Event p (Stream v)

class IsBehavior behavior where
	-- | Generate an Event, from some type of behavior,
	-- that only occurs when the `Behavior` Bool is True.
	(@:) :: behavior -> Behavior Bool -> BehaviorToEvent behavior

instance IsBehavior (Behavior v) where
	b @: c = Event b c

instance IsBehavior (TypedBehavior p v) where
	(@:) (TypedBehavior b) c = Event b c

instance IsDigitalIOPin t => Output (Pin t) (Event () (Stream Bool)) where
	(Pin p@(PinId n)) =: (Event b c) = do
		(f, triggername) <- defineTriggerAlias' ("pin_" <> show n) "digitalWrite" $
			mempty { pinmodes = M.singleton p (S.singleton OutputMode) }
		tell [(go triggername, const f)]
	  where
		go triggername tl = 
			let c' = addTriggerLimit tl c
			in trigger triggername c' [arg (constant n), arg b]

instance IsPWMPin t => Output (Pin t) (Event 'PWM (Stream Word8)) where
	(Pin (PinId n)) =: (Event v c) = do
		(f, triggername) <- defineTriggerAlias' ("pin_" <> show n) "analogWrite" mempty
		tell [(go triggername, const f)]
	  where
		go triggername tl = 
			let c' = addTriggerLimit tl c
			in trigger triggername c' [arg (constant n), arg v]
		-- analogWrite does not need any pinmodes set up

class Input o t where
	-- | The list is input to use when simulating the Sketch.
	input' :: o -> [t] -> Sketch (Behavior t)

-- | Use this to read a value from a component of the Arduino.
--
-- For example, to read a digital value from pin12 and turn on the 
-- led when the pin is high:
--
-- > buttonpressed <- input pin12
-- > led =: buttonpressed
--
-- Some pins support multiple types of reads, for example pin a0
-- supports a digital read (`Bool`), and an analog to digital converter
-- read (`ADC`). In such cases you may need to specify the type of
-- data to read:
--
-- > v <- input a0 :: Sketch (Behavior ADC)
input :: Input o t => o -> Sketch (Behavior t)
input o = input' o []

instance IsDigitalIOPin t => Input (Pin t) Bool where
	input' (Pin p@(PinId n)) interpretvalues = mkInput $ InputSource
		{ defineVar = mkCChunk [CLine $ "bool " <> varname <> ";"]
		, setupInput = mempty
		, inputPinmode = M.singleton p InputMode
		, readInput = mkCChunk
			[CLine $ varname <> " = digitalRead(" <> show n <> ");"]
		, inputStream = extern varname interpretvalues'
		}
	  where
		varname = "arduino_digital_pin_input" <> show n
		interpretvalues'
			| null interpretvalues = Nothing
			| otherwise = Just interpretvalues

-- | Value read from an Arduino's ADC. Ranges from 0-1023.
type ADC = Int16

instance IsAnalogInputPin t => Input (Pin t) ADC where
	input' (Pin (PinId n)) interpretvalues = mkInput $ InputSource
		{ defineVar = mkCChunk [CLine $ "int " <> varname <> ";"]
		, setupInput = mempty
		, inputPinmode = mempty
		, readInput = mkCChunk
			[CLine $ varname <> " = analogRead(" <> show n <> ");"]
		, inputStream = extern varname interpretvalues'
		}
	  where
		varname = "arduino_analog_pin_input" <> show n
		interpretvalues'
			| null interpretvalues = Nothing
			| otherwise = Just interpretvalues

class ShowCType t where
	showCType :: Proxy t -> String

instance ShowCType Bool where showCType _ = "bool"
instance ShowCType Int8 where showCType _ = "int8_t"
instance ShowCType Int16 where showCType _ = "int16_t"
instance ShowCType Int32 where showCType _ = "int32_t"
instance ShowCType Int64 where showCType _ = "int64_t"
instance ShowCType Word8 where showCType _ = "uint8_t"
instance ShowCType Word16 where showCType _ = "uint16_t"
instance ShowCType Word32 where showCType _ = "uint32_t"
instance ShowCType Word64 where showCType _ = "uint64_t"
instance ShowCType Float where showCType _ = "float"
instance ShowCType Double where showCType _ = "double"