#%% [markdown]
# Complex cognitive experiments require participants to perform multiple tasks across a span of limited lab time or online crowdsourcing time, hence demand flexible scheduling of tasks. 
# This notebook demonstrates a possible solution for the problem of scheduling experimental tasks, and mainly concerns reformulating experiments as a resource allocation problem.
#
# Typical experiments are comprised of several session, in each session participants view a set of questionnaires, perform a test, and then respond to another set of questionnaires. This structure is repeated for each srssion throughout the experiment.

# To orgnaize questionniare and tasks in a limited time, we first model a study as an ordered list of sessions (here, for example, 12 session). Sessions consist of several questionnaires before and after a set of experimental tasks. We then define constraints on sueverys and tasks (e.g., pre-task, post-task, end of study, and start of study surveys), and solve for a schedule for the given scenario. Note that his example scenario only plans the experiment for a single subject.
#
# Note: Use the following command to install dependencies: `pip install pyschedule`

# %%
from math import ceil

from pyschedule import Scenario, plotters, alt
from pyschedule.solvers.mip import solve
from pyschedule.plotters.matplotlib import plot

import seaborn as sns

# 1. set theme and style
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 = .16 # minutes (~10sec)

prestudy_tasks = [
  "xcit_demographics"
]

pregame_tasks = [

]

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,
  #"mmi_part1": 66,
  "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)


def _duration(questions):
  t = ceil(questions * duration_per_question)
  return t


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

# 3. questionnaires
for q in n_questions.keys():
  n = n_questions.get(q, 0)
  task = s.Task(q, length=_duration(n), delay_cost=1)

  if q in prestudy_tasks:
    task += sessions[0] # first session
    task.delay_cost = 2
    s += task < games
  elif q in poststudy_tasks:
    task += sessions[-1] # last session
    task.delay_cost = -2
    s += games < task
  elif q in postgame_tasks:
    task += sessions
    s += games < task
  elif q in pregame_tasks:
    task += sessions
    s += task < games
  else:
    task += alt(sessions)

print(s.tasks)

# 4. solve resource allocation
solve(s)

# 5. print last computed solution
print(s.solution())

# 6. pot gantt
plot(s, fig_size=(50,5))

#%%
from pyschedule import Scenario, plotters, alt
from pyschedule.solvers.mip import solve
from pyschedule.plotters.matplotlib import plot

import pandas as pd
from math import ceil


# load the study design data
activities = pd.read_csv('p500_activities.csv')
activities['activity'] = activities.apply(
    lambda r:
    r['activity'] + '_g' + str(r['group']) if r['type'] == 'assessment' else r['activity'], axis=1)


# 1. define the study
n_sessions = 10
session_duration = 60

# OR infer session duration from data
# session_duration = ceil(activities['duration'].sum() / n_sessions)

scenario = Scenario('p500_activities', horizon=session_duration)

# 2. define the sessions
sessions = scenario.Resources('Bucket', num=n_sessions)

# 3. define the tasks
for _, _, activity, activity_type, duration in activities.itertuples():
  # fixes for the pyschedule resolver to work
  duration = ceil(duration)
  activity = activity.replace('-', '_')

  # create the task
  task = scenario.Task(activity, length=duration, delay_cost=1)

  # naive constraint (all buckets)
  task += alt(sessions)

# 4. solve resource allocation
solve(scenario)

# 5. print last computed solution
print(scenario.solution())

# 6. pot gantt
plot(scenario, fig_size=(50, 10))

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