# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # SEIRX model runs for prevention measures investigation in schools

# +
import pandas as pd
import numpy as np
from os.path import join
import os
import sys
import json

from scseirx import analysis_functions as af

currentdir = os.getcwd()
parentdir = os.path.dirname(currentdir)
sys.path.insert(0, parentdir) 
import data_creation_functions as dcf

# parallelisation functionality
from multiprocess import Pool
import psutil
from tqdm import tqdm
# -

# ## Simulation settings

# basic prevention measures in place. Additional measures will be added
# throughout the measure screening simulations
with open('params/intervention_screening_measures.json', 'r') as fp:
    measures = json.load(fp)
# simulation parameters, specifically the
# - base transmission risk (calibrated from household transmissions)
# - subclinical transmission modifier (literature value)
# - exposure duration, time until symtpoms and infection duration (lit. values)
# - age symptom discount (fit to empirical observations)
# - age transmission discount (from the calibration)
# - weights of transmission risks for intermediate & far contacts (calibration)
with open('params/intervention_screening_simulation_parameters.json', 'r') as fp:
    simulation_params = json.load(fp)
# characteristics (# classes, # students / class, # teachers) of the "average" 
# school, depending on school type. These characteristics were determined in 
# interviews with Austrian teachers and from statistics about Austrian schools 
# (year 2017/18, page 10: https://www.bmbwf.gv.at/Themen/schule/schulsystem/gd.html)
# NOTE: "students" indicates the number of students per class
with open('params/intervention_screening_school_characteristics.json', 'r') as fp:
    school_characteristics = json.load(fp)


# ## Simulation function

def run(params):
    '''
    Runs an ensemble of simulations and collects observable statistics. To be 
    run in parallel on many workers. Note: I/O paths and the number of runs per 
    ensemble hare hard coded here, because I only want to pass the parameter 
    values that are being screened in the simulation run to the function via the
    parallel processing interface.
    
    Parameters:
    -----------
    param_list : iterable
        Iterable that contains the values for the parameters test_type, 
        index_case, e_screen_range and r_screen_range that are passed to the
        simulation.
        
    Returns:
    --------
    row : dictionary
        Dictionary of the ensemble statistics of the observables.
    '''    

    # extract the simulation parameters from the parameter list
    N_runs, school_type, index_case, ttype, s_screen_interval, t_screen_interval,\
        student_mask, teacher_mask, half_classes, ventilation_mod = params
    
    try:
        os.mkdir(join(dst, school_type))
    except FileExistsError:
        pass

    # run the ensemble with the given parameter combination and school type
    row = dcf.run_ensemble(N_runs, school_type, measures,\
            simulation_params, school_characteristics, contact_network_src,\
            dst, index_case, ttype, s_screen_interval, t_screen_interval,\
            student_mask, teacher_mask, half_classes, ventilation_mod)
    
    row['school_type'] = school_type
    row['index_case'] = index_case
    row['test_type'] = ttype
    row['student_screen_interval'] = s_screen_interval
    row['teacher_screen_interval'] = t_screen_interval
    row['student_mask'] = student_mask
    row['teacher_mask'] = teacher_mask
    row['half_classes'] = half_classes
    row['ventilation_mod'] = ventilation_mod
    
    return row


# ## Screening parameters

# +
## parameter ranges 

# number of runs in the ensemble
# Note: this is set to 1 for testing purposes. To get properly converged 
# statistics, this has to be >= 500. Running such a high number of simulations
# will take a long time if run on just a single core. Therefore I strongly
# recommend to run the below code on several cores at once. It is easy to
# parallelize as each ensemble can be run on a different core and there are
# no interdependencies between ensembles. Results can be collected afterwards
# and evaulated together.
N_runs = 1
# different school types
school_types = ['primary', 'primary_dc', 'lower_secondary',
                'lower_secondary_dc', 'upper_secondary', 'secondary']
# specifies whether the index case will be introduced via an
# employee or a resident
index_cases = ['student', 'teacher']
# test technologies (and test result turnover times) used in the
# different scenarios
test_types = ['same_day_antigen']
# student and teacher streening intervals (in days)
s_screen_range = [None, 3, 7]
t_screen_range = [None, 3, 7]
# specifies whether teachers wear masks
student_masks = [True, False]
teacher_masks = [True, False]
half_classes = [True, False]
# specifies whether there is ventilation or not
transmission_risk_ventilation_modifiers = [1, 0.36]

params = [(N_runs, i, j, k, l, m, n, o, p, q)\
              for i in school_types \
              for j in index_cases \
              for k in test_types \
              for l in s_screen_range \
              for m in t_screen_range \
              for n in student_masks \
              for o in teacher_masks \
              for p in half_classes \
              for q in transmission_risk_ventilation_modifiers]

N_configs = len(params)
print('there are {} different parameter combinations'.format(N_configs))
# -

# ## Run simulations

# +
# %%time
# paths for data I/O
contact_network_src = '../../../data/contact_networks/representative_schools'
dst = '../../../data/intervention_measures_wild_type/simulation_results'

number_of_cores = psutil.cpu_count(logical=True) - 2
pool = Pool(number_of_cores)

rows = []
for row in tqdm(pool.imap_unordered(func=run, iterable=params[0:10]), total=len(params)):
        rows.append(row)

# turn off your parallel workers 
pool.close()

# format and dump the results to disk
results = pd.DataFrame()
for row in rows:
    results = results.append(row, ignore_index=True)
    
results = results.reset_index(drop=True)
index_cols = ['school_type', 'index_case', 'test_type',
              'student_screen_interval', 'teacher_screen_interval',
              'student_mask', 'teacher_mask', 'half_classes',
              'ventilation_mod']
other_cols = [c for c in results.columns if c not in index_cols]
results = results[index_cols + other_cols]

results.to_csv(join(dst,'intervention_measures_{}.csv'\
                   .format(N_runs)), index=False)
results.head(3)
# -

# ## Aggregate results

# + tags=[]
src = '../../../data/intervention_measures_wild_type/simulation_results'
dst = src
for stype in school_types:
    print(stype)
    data = dcf.get_data(stype, src)
    data['measure'] = np.nan
    dcf.set_individual_measures(data)
    dcf.set_measure_packages(data)
    data.to_csv(join(dst, '{}_combined_ensembles.csv'.format(stype)), index=False)
# -

# ## Extract observables

# +
dst = '../../../data/intervention_measures_wild_type/simulation_results'

turnovers = {'same':0, 'one':1, 'two':2, 'three':3}
bmap = {True:'T', False:'F'}

for stype in school_types:
    print(stype)
    spath_ensmbl = join(dst,'{}'.format(stype))
    
    observables = pd.DataFrame()
    for N_runs, pstype, index_case, ttype, s_screen_interval, t_screen_interval,\
        student_mask, teacher_mask, half_classes, ventilation_mod in params:
        
        if pstype != stype:
            continue
        
        turnover, _, test = ttype.split('_')
        turnover = turnovers[turnover]
        
        measure_string = '{}_test-{}_turnover-{}_index-{}_tf-{}_sf-{}_tmask-{}'\
        .format(stype, test, turnover, index_case[0], t_screen_interval,
                s_screen_interval, bmap[teacher_mask]) +\
                '_smask-{}_half-{}_vent-{}'\
        .format(bmap[student_mask], bmap[half_classes], ventilation_mod)
        
        ensmbl = pd.read_csv(join(spath_ensmbl, measure_string + '.csv'))
        ensmbl = ensmbl.drop(columns=['Unnamed: 0'])
        
        row = {'test_type':test,
               'turnover':turnover,
               'index_case':index_case,
               'student_screen_interval':s_screen_interval,
               'teacher_screen_interval':t_screen_interval,
               'student_mask':student_mask,
               'teacher_mask':teacher_mask,
               'half_classes':half_classes,
               'ventilation_modification':ventilation_mod}
        
        N = len(ensmbl)
        ensmbl = ensmbl[ensmbl['infected_agents'] > 0]
        N_red = len(ensmbl)
        row.update({'no_outbreak':(N - N_red) / N})
        
        for col in ensmbl.columns:
            row.update(af.get_statistics(ensmbl, col))
        observables = observables.append(row, ignore_index=True)

    # calculate the number of tests per day and agent in the school
    observables['N_tests_per_day_per_agent'] = \
    (observables['N_diagnostic_tests_median'] + observables['N_preventive_tests_median']) /\
    observables['duration_median'] / observables['N_school_agents_median']
    
    screen_cols = ['test_type', 'turnover', 'index_case', 'student_screen_interval',
            'teacher_screen_interval', 'student_mask', 'teacher_mask',
            'half_classes', 'ventilation_modification']

    other_cols = [c for c in observables.columns if c not in screen_cols]
    observables = observables[screen_cols + other_cols]
    
    for col in ['infected_teachers_median', 'infected_students_median', 
                        'infected_family_members_median',
                        'infected_teachers_0.90', 'infected_students_0.90', 
                        'infected_family_members_0.90']:
        observables[col] = observables[col].round(0).astype(int)
    
    observables.to_csv(join(dst, '{}_observables'.format(stype) + '.csv'))

StarcoderdataJupyterScriptsDedupFiltered