#%% [markdown]
# Complex cognitive experiments require participants to perform multiple tasks across a space of time and then capture observations with different modalities. This notebook demonstrates some possible solutions for the problem of scheduling experimental tasks, and mainly concerns reformulating experiments as a resource allocation problem.

# 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
game_duration = 45
n_subjects = 2
n_sessions = 12
duration_per_question = .1666667 # minutes (~10sec)

prestudy_tasks = [
  "xcit_demographics"
]

postgame_tasks = [
  "nasa_tlx",
  "xcit_postgame_debrief"
]

poststudy_tasks = [
  "xcit_poststudy_debrief"
]

n_questions = {
  "xcit_demographics": 20,
  "aiss": 20,
  "arces": 12,
  "asrs": 18,
  "avg2019": 16,
  "bfi_2_par1": 30,
  "bfi_2_part2": 30,
  "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,
  "nasa_tlx": 6,
  "xcit_postgame_debrief": 15
}


# 1. define the study
s = Scenario('Prolific500', horizon=session_duration)
# 2. sessions
sessions = s.Resources('Session', num = n_sessions)


# 2. games
games = s.Tasks('Game',length=game_duration, num=n_sessions, is_group=False)
games += alt(sessions)

# questionnaires
for q in n_questions.keys():
  n = n_questions.get(q, 0)
  duration = ceil(n * duration_per_question)
  task = s.Task(q, length= duration, delay_cost=1)
  if q in poststudy_tasks:
    print(q, 1)
    task += sessions[-1]
    s += games < task
  elif q in prestudy_tasks:
    print(q, 2)
    task += sessions[0]
    s += task < games
  elif q in postgame_tasks:
    print(q, 3)
    task += sessions
    s += games < task
  else:
    print(q, 4)
    task += alt(sessions)

print(s.tasks)

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)