#%% [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 math import ceil
from pyschedule import Scenario, solvers, plotters, alt
import seaborn as sns


sns.color_palette("colorblind")
sns.set_style('white')
# experiment duration in minutes
session_duration = 60
n_subjects = 2
n_sessions = 12
duration_per_question = .1666667 # minutes (~10sec)
questionnaires = {
  "study_level": {
    "xcit_demographics": 20,
    "aiss": 20,
    "arces": 12,
    "asrs": 18,
    "avg2019": 16,
    "bfi_2": 60,
    "bisbas": 24,
    "bis11": 12,
    "cfq": 25,
    "cfs": 12,
    "dyslexia": 15,
    "ehi_sf": 4,
    "grit12": 12,
    "i_panas_sf": 10,
    "ipaq_sf": 7,
    "maas": 15,
    "mfs": 12,
    "mw": 8,
    "mwq": 5,
    "ncs_6": 6,
    "nfc": 18,
    "psqi": 9,
    "rei": 40,
    "sci": 8,
    "sqs": 28,
    "upps_sf": 20,
    "webexec": 12,
    "who5": 5,
    "whoqol": 26,
    "xcit_poststudy_debrief": 20 # only in session 12
  },
  "session_level": {
  'nasa_tlx': 6, # after game
  "xcit_postgame_debrief": 15 #after game
  }
}

s = Scenario('Prolific500', horizon=session_duration)

sessions = s.Resources('Session', num=n_sessions)


# one-time questionnaires
pregame_tasks = []
for q in questionnaires['study_level'].keys():
  n_questions = questionnaires['study_level'].get(q, 0)
  duration = ceil(n_questions * duration_per_question)
  task = s.Task(q, length= duration, delay_cost=1)
  task += alt(sessions)
  # not really a pregame task
  # pregame_tasks.append(questionnaire)

# game
games = s.Task('Game',length=45, delay_cost=1, is_group=False)
games += sessions

# post-tasks
postgame_tasks = []
for q in questionnaires['session_level'].keys():
  n_questions = questionnaires['session_level'].get(q, 0)
  duration = ceil(n_questions * duration_per_question)
  task = s.Task(q, length= duration, delay_cost=1)
  task += sessions
  postgame_tasks.append(task)

# add constraints on precedences
#s += pregame_tasks < games
s += games < postgame_tasks

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

plotters.matplotlib.plot(s, fig_size=(50,5))

#%%
#%% [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)