{-# LANGUAGE GADTs               #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- |
-- Copyright : (c) NASA, 2024-2025
-- License   : BSD-style (see the LICENSE file in the distribution)
--
-- Allow loading Copilot specifications from strings and modifying them live.
--
-- An example of a spec that can be passed as argument to the live simulator
-- follows:
--
-- @
--   spec :: String
--   spec = unlines
--     [ "let temperature :: Stream Word8"
--     , "    temperature = extern \"temperature\" (Just [0, 15, 20, 25, 30])"
--     , ""
--     , "    ctemp :: Stream Float"
--     , "    ctemp = (unsafeCast temperature) * (150.0 / 255.0) - 50.0"
--     , ""
--     , "    trueFalse :: Stream Bool"
--     , "    trueFalse = [True] ++ not trueFalse"
--     , ""
--     , "in do trigger \"heaton\"  (temperature < 18) [arg ctemp, arg (constI16 1), arg trueFalse]"
--     , "      trigger \"heatoff\" (temperature > 21) [arg (constI16 1), arg ctemp]"
--     , "      observer \"temperature\" temperature"
--     , "      observer \"temperature2\" (temperature + 1)"
--     ]
-- @
--
-- The imports are predefined.
module Copilot.Visualize.Dynamic
    ( simInit
    , simStep
    , SimulationSettings(..)
    , mkDefaultSimulationSettings
    , SimData(..)
    )
  where

-- External imports
import qualified Language.Haskell.Interpreter as HI
import           Prelude                      hiding (div, not, (++), (<), (>))
import qualified Prelude

-- External imports: Copilot
import qualified Copilot.Core     as Core
import           Copilot.Language hiding (interpret, typeOf)
import           Language.Copilot hiding (interpret, typeOf)

-- Internal imports
import Copilot.Visualize.TypedTrace

-- | Simulation data.
data SimData = SimData
  { simSteps  :: Int        -- ^ Number of steps to simulate.
  , simSpec   :: Core.Spec  -- ^ Copilot spec to simulate.
  , simString :: String     -- ^ Textual representation of Copilot spec.
  }

-- | Initialize a simulation.
simInit :: SimulationSettings -> String -> IO SimData
simInit simulationSettings spec = do
  -- Load the spec, and save it so that it doesn't need to be reloaded if it
  -- doesn't change.
  let numSteps = simulationSettingsInitialSteps simulationSettings
  spec' <- readSpec simulationSettings spec
  let simData = SimData numSteps spec' spec
  return simData

-- | Update the spec based on the input command received.
simStep :: SimulationSettings
        -> SimData
        -> (Command, String)
        -> IO SimData
simStep _ simData (StepUp, _) =
  pure simData { simSteps = simSteps simData + 1 }

simStep _ simData (StepDown, _) =
  pure simData { simSteps = simSteps simData + 1 }

simStep simulationSettings simData (AddStream name expr, _) = do
  let specN = concat [ simString simData, "\n", "      ", completeExpr]

      completeExpr = concat [ "observer ", show name, " (", expr, ")" ]

  let trace = extractTrace $ simSpec simData
  spec2 <- readSpec simulationSettings specN
  let spec3 = updateWithTrace trace spec2

  pure simData { simSpec = spec3, simString = specN }

simStep simulationSettings simData (command, name) = do
  spec' <- apply simulationSettings (simSpec simData) name command
  pure simData { simSpec = spec' }

-- | Settings used for a simulation.
data SimulationSettings = SimulationSettings
  { simulationSettingsInitialSteps :: Int
  , simulationSettingsImports      :: [(String, Maybe String)]
  }

-- | Default settings that simulates 3 steps and provides default imports of
-- the main Copilot modules.
mkDefaultSimulationSettings :: SimulationSettings
mkDefaultSimulationSettings =
  SimulationSettings
    3
    [ ("Control.Monad.Writer",  Nothing)
    , ("Copilot.Language",      Nothing)
    , ("Copilot.Language.Spec", Nothing)
    , ("Data.Functor.Identity", Nothing)
    , ("Language.Copilot",      Nothing)
    , ("Prelude",               Just "P")
    ]

-- * Commands

-- * Load specs from strings

-- | Read a specification from a string and reify it.
readSpec :: SimulationSettings -> String -> IO Core.Spec
readSpec settings spec = do
  r <- HI.runInterpreter (loadSpec settings spec)
  case r of
    Left err -> do putStrLn $ "Error: " Prelude.++ show err
                   error $ show err
    Right s  -> return s

-- | Load a spec from a string, adding standard imports.
loadSpec :: SimulationSettings -> String -> HI.Interpreter Core.Spec
loadSpec settings spec = do
  HI.setImportsQ (simulationSettingsImports settings)
  spec' <- HI.interpret spec (HI.as :: Spec)
  HI.liftIO $ reify spec'

-- | Commands to apply to a running simulation.
data Command = StepUp
             | StepDown
             | Up Int
             | Down Int
             | AddStream String String
             | Noop
  deriving (Eq, Read, Show)

-- | Apply a command to a simulation.
apply :: SimulationSettings -> Core.Spec -> String -> Command -> IO Core.Spec
apply settings spec _name (AddStream sName sExpr) = do
  spec' <- addStream settings sName sExpr
  -- TODO: I need to bring the streams from the other spec too, otherwise the
  -- streams to include may refer to streams by ID that are in a different
  -- scope.
  let observers' = Core.specObservers spec'
      observers  = Core.specObservers spec
  return $ spec { Core.specObservers = observers Prelude.++ observers' }

apply _settings spec name command = pure $ spec
  { Core.specStreams =
      map (updateStream name command) (Core.specStreams spec)
  , Core.specObservers =
      map (updateObserver name command) (Core.specObservers spec)
  , Core.specTriggers =
      map (updateTrigger name command) (Core.specTriggers spec)
  }

-- | Apply a command to an observer.
updateObserver :: String -> Command -> Core.Observer -> Core.Observer
updateObserver name command (Core.Observer i e ty) =
  Core.Observer i (updateExpr name command e) ty

-- | Apply a command to a trigger.
updateTrigger :: String -> Command -> Core.Trigger -> Core.Trigger
updateTrigger name command (Core.Trigger i e es) = Core.Trigger i e' es'
  where
    e'  = updateExpr name command e
    es' = map (updateUExpr name command) es

-- | Apply a command to a core expression.
updateExpr :: String -> Command -> Core.Expr a -> Core.Expr a
updateExpr name command expr = case expr of
  Core.ExternVar ty nameE vs
    | nameE Prelude.== name
    -> Core.ExternVar ty nameE (updateValues vs ty command)

    | otherwise
    -> expr

  Core.Op1 op expr1 ->
    Core.Op1 op (updateExpr name command expr1)

  Core.Op2 op expr1 expr2 ->
    Core.Op2 op (updateExpr name command expr1) (updateExpr name command expr2)

  Core.Op3 op expr1 expr2 expr3 ->
    Core.Op3
      op
      (updateExpr name command expr1)
      (updateExpr name command expr2)
      (updateExpr name command expr3)

  _ -> expr

-- | Apply a command to an untyped expression that carries the information
-- about the type of the expression as a value (existential).
updateUExpr :: String -> Command -> Core.UExpr -> Core.UExpr
updateUExpr name cmd (Core.UExpr ty expr) =
  Core.UExpr ty (updateExpr name cmd expr)

-- | Apply a command to a stream.
updateStream :: String -> Command -> Core.Stream -> Core.Stream
updateStream name command (Core.Stream i buf expr ty) =
  Core.Stream i buf (updateExpr name command expr) ty

-- | Apply a command to a series of typed values, if any.
updateValues :: Maybe [a] -> Type a -> Command -> Maybe [a]
updateValues vsM ty command =
  fmap (fmap (updateValue command ty) . zip [0..]) vsM

-- | Apply a command to a typed value.
updateValue :: Command -> Type a -> (Int, a) -> a
updateValue (Up n)   Core.Bool  (ix, a) =
  if n Prelude.== ix then Prelude.not a else a
updateValue (Down n) Core.Bool  (ix, a) =
  if n Prelude.== ix then Prelude.not a else a
updateValue (Up n)   Core.Word8 (ix, a) =
  if n Prelude.== ix then a + 1 else a
updateValue (Down n) Core.Word8 (ix, a) =
  if n Prelude.== ix then a - 1 else a
updateValue _        _          (_,  a) = a

-- * Sample spec

-- | Produce a specification from the expression for a stream.
addStream :: SimulationSettings -> String -> String -> IO Core.Spec
addStream settings name expr = do
  r <- HI.runInterpreter (addStream' settings name expr)
  case r of
    Left err   -> do putStrLn $ "Error: " Prelude.++ show err
                     error $ show err
    Right spec -> return spec

-- | Produce a specification from the expression for a stream, in an
-- interpreter context.
--
-- Observe that Interpreter () is an alias for InterpreterT IO ()
addStream' :: SimulationSettings -> String -> String -> HI.Interpreter Core.Spec
addStream' settings name expr = do
  HI.setImportsQ (simulationSettingsImports settings)
  let completeExpr = concat [ "observer ", show name, " (", expr, ")" ]

  spec <- HI.interpret completeExpr (HI.as :: Spec)
  HI.liftIO $ reify spec