#%% [markdown]
# Complex experiments in cognitive science require participants to perform multiple tasks and capture observations with different modalities. This notebook demonstrates some possible solutions for the problem of scheduling experimental tasks.

# First, we model a study with 12 sessions. Sessions consist of several surveys before and after a set of experimental tasks. We then define constraints and solve for a schedule for the given scenario. This scenario only plans the experiment for a single subject.

# Use the following command to install dependencies: `pip install pyschedule`

#%%

from pyschedule import Scenario, solvers, plotters, alt

# experiment duration in minutes
duration = 60
questionnaires = ["xcit-demographics", "aiss", "arces", "asrs", "avg2019", "bfi-2", "bisbas", "bis11", "cfq", "cfs", "dyslexia", "ehi-sf", "grit12", "i-panas-sf", "ipaq-sf", "maas", "mfs", "mw", "mwq", "ncs-6", "nfc", "psqi", "rei", "sci", "sqs", "upps-sf", "webexec", "who5", "whoqol","nasa-tlx", "xcit-postgame-debrief"]


# the planning horizon has 
s = Scenario('Prolific500', horizon=duration)

session = s.Resource('session', num=12)

mmi = s.Task('mmi', length=3, delay_cost=1)
mmi += alt(session)

games = s.Task('games',length=45, delay_cost=1, is_group=True)
games += alt(session)

nasa_tlx = s.Task('nasa_tlx', length=1, delay_cost=1)
nasa_tlx += alt(session)

# add constraints on precedences
s += mmi < games 
s += games < nasa_tlx


# compute and print session schedules
solvers.mip.solve(s)
print(s.solution())

plotters.matplotlib.plot(s)

#%%
#%% [markdown]
# Complex experiments in cognitive science require participants to perform multiple tasks and capture observations with different modalities. This notebook demonstrates some possible solutions for the problem of scheduling experimental tasks.

# First, we model a study with 12 sessions. Sessions consist of several surveys before and after a set of experimental tasks. We then define constraints and solve for a schedule for the given scenario. This scenario only plans the experiment for a single subject.

#%%
# MPILX Scheduling
from pyschedule import Scenario, solvers, plotters, alt

# experiment duration for each subject in minutes
study_duration = 12 * 24 * 60 # days * hours * minutes
n_subjects = 2
task_durations = {
  'pretest_presurvey': 30,
  'pretest_DWI': 15,
  'pretest_rsfMRI': 15,
  'pretest_anatomical': 15,
  'pretest_taskfMRI': 45,
  'pretest_postsurvey': 30,
  'pretest_presurvey': 30,
  'training_presurvey': 30,
  'training_': 10 * 24 * 60, #TODO expand to daily
  'training_postsurvey': 30,
  'posttest_DWI': 15,
  'posttest_rsfMRI': 15,
  'posttest_anatomical': 15,
  'posttest_taskfMRI': 45,
  'posttest_postsurvey': 30
}

# the planning horizon has 
s = Scenario('MPILX50', horizon=study_duration)

subject = s.Resource('subject', num=n_subjects)

tasks = list()

for t in task_durations.keys():
  duration = task_durations[t]
  task = s.Task(t, length=duration, delay_cost=1)
  task += alt(subject)
  tasks.append(task)

for i in range(len(tasks)-1):
  s += tasks[i] < tasks[i+1]

# compute and print session schedules
solvers.mip.solve(s)
print(s.solution())

plotters.matplotlib.plot(s)