module Simulation.Aivika.Experiment.Types where
import Control.Monad
import Control.Monad.Trans
import Control.Exception
import Control.Concurrent.ParallelIO.Local
import Data.Maybe
import Data.Monoid
import Data.Either
import GHC.Conc (getNumCapabilities)
import Simulation.Aivika
import Simulation.Aivika.Trans.Exception
data Experiment = 
  Experiment { experimentSpecs         :: Specs,
               
               experimentTransform     :: ResultTransform,
               
               experimentLocalisation  :: ResultLocalisation,
               
               experimentRunCount      :: Int,
               
               experimentTitle         :: String,
               
               experimentDescription   :: String,
               
               experimentVerbose       :: Bool,
               
               experimentNumCapabilities :: IO Int
               
               
             }
defaultExperiment :: Experiment
defaultExperiment =
  Experiment { experimentSpecs         = Specs 0 10 0.01 RungeKutta4 SimpleGenerator,
               experimentTransform     = id,
               experimentLocalisation  = englishResultLocalisation,
               experimentRunCount      = 1,
               experimentTitle         = "Simulation Experiment",
               experimentDescription   = "",
               experimentVerbose       = True,
               experimentNumCapabilities = getNumCapabilities }
class ExperimentRendering r where
  
  data ExperimentContext r :: *
  
  type ExperimentEnvironment r :: *
  
  type ExperimentMonad r :: * -> *
  
  liftExperiment :: r -> ExperimentMonad r a -> IO a
  
  prepareExperiment :: Experiment -> r -> ExperimentMonad r (ExperimentEnvironment r)
  
  
  renderExperiment :: Experiment -> r -> [ExperimentReporter r] -> ExperimentEnvironment r -> ExperimentMonad r ()
  
  onExperimentCompleted :: Experiment -> r -> ExperimentEnvironment r -> ExperimentMonad r () 
  
  onExperimentFailed :: Exception e => Experiment -> r -> ExperimentEnvironment r -> e -> ExperimentMonad r ()
data ExperimentGenerator r = 
  ExperimentGenerator { generateReporter :: Experiment -> r -> ExperimentEnvironment r -> ExperimentMonad r (ExperimentReporter r)
                        
                      }
class ExperimentRendering r => ExperimentView v r where
  
  
  outputView :: v -> ExperimentGenerator r
data ExperimentData =
  ExperimentData { experimentResults :: Results,
                   
                   experimentPredefinedSignals :: ResultPredefinedSignals
                   
                 }
data ExperimentReporter r =
  ExperimentReporter { reporterInitialise :: ExperimentMonad r (),
                       
                       
                       reporterFinalise   :: ExperimentMonad r (),
                       
                       
                       reporterSimulate   :: ExperimentData -> Composite (),
                       
                       reporterContext    :: ExperimentContext r
                       
                     }
runExperiment :: (ExperimentRendering r,
                  Monad (ExperimentMonad r),
                  MonadIO (ExperimentMonad r),
                  MonadException (ExperimentMonad r))
                 => Experiment
                 
                 -> [ExperimentGenerator r]
                 
                 -> r
                 
                 -> Simulation Results
                 
                 -> IO (Either SomeException ())
runExperiment e generators r simulation =
  runExperimentWithExecutor sequence_ e generators r simulation
  
runExperimentParallel :: (ExperimentRendering r,
                          Monad (ExperimentMonad r),
                          MonadIO (ExperimentMonad r),
                          MonadException (ExperimentMonad r))
                         => Experiment
                         
                         -> [ExperimentGenerator r]
                         
                         -> r
                         
                         -> Simulation Results
                         
                         -> IO (Either SomeException ())
runExperimentParallel e generators r simulation =
  do x <- runExperimentWithExecutor executor e generators r simulation
     return (x >> return ())
       where executor tasks =
               do n <- experimentNumCapabilities e
                  withPool n $ \pool ->
                    parallel_ pool tasks
                        
runExperimentWithExecutor :: (ExperimentRendering r,
                              Monad (ExperimentMonad r),
                              MonadIO (ExperimentMonad r),
                              MonadException (ExperimentMonad r))
                             => ([IO ()] -> IO a)
                             
                             -> Experiment
                             
                             -> [ExperimentGenerator r]
                             
                             -> r
                             
                             -> Simulation Results
                             
                             -> IO (Either SomeException a)
runExperimentWithExecutor executor e generators r simulation =
  liftExperiment r $
  do let specs      = experimentSpecs e
         runCount   = experimentRunCount e
     env <- prepareExperiment e r
     let c1 =
           do reporters <- mapM (\x -> generateReporter x e r env)
                           generators
              forM_ reporters reporterInitialise
              let simulate :: Simulation ()
                  simulate =
                    do signals <- newResultPredefinedSignals
                       results <- simulation
                       let d = ExperimentData { experimentResults = experimentTransform e results,
                                                experimentPredefinedSignals = signals }
                       ((), fs) <- runDynamicsInStartTime $
                                   runEventWith EarlierEvents $
                                   flip runComposite mempty $
                                   forM_ reporters $ \reporter ->
                                   reporterSimulate reporter d
                       let m1 =
                             runEventInStopTime $
                             return ()
                           m2 =
                             runEventInStopTime $
                             disposeEvent fs
                           mh (SimulationAbort e') =
                             return ()
                       finallySimulation (catchSimulation m1 mh) m2
              a <- liftIO $
                executor $ runSimulations simulate specs runCount
              forM_ reporters reporterFinalise
              renderExperiment e r reporters env
              onExperimentCompleted e r env
              return (Right a)
         ch z@(SomeException e') =
           do onExperimentFailed e r env e'
              return (Left z)
     catchComp c1 ch