Import('*')
import pandas as pd, numpy as np, re, os
from pathlib import Path
curdir = '/'.join(re.split('/|\\\\', str(Path.cwd()))[0:-4])

source = ['#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/get_weights.py',
          '#source/derived/bayes_bootstrap/infrm_index.csv']
target = ['#datastore/output/derived_large/bayes_bootstrap_weights/Dinerstein_Smith_2021_Weights.csv']
env.Python(target, source)

source = ['#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/M08_create_common_school_codes.do',
          '#datastore/output/derived_large/bayes_bootstrap_weights/Dinerstein_Smith_2021_Weights.csv',  # to wait on directories
          '#temp/DS2021/temp/model_data/school_grade8_codes_2002_2012_weights.dta']
target = ['#temp/DS2021_bayes/temp/model_data/grade_year_enroll_2002_2012_weights.csv',
          '#temp/DS2021_bayes/temp/model_data/universal_school_codes_2002_2012_weights.csv']
env.Stata(target, source)

source = ['#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/M04_wrapper_grade_model.do',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/M04b_create_grade_school_shares.do',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/M04a_create_grade_model_data.do',
          '#temp/DS2021/temp/model_data/school_locations_enrollments.dta',
          '#temp/DS2021/temp/model_data/census_block_students.dta',
          '#datastore/output/derived_large/bayes_bootstrap_weights/Dinerstein_Smith_2021_Weights.csv']  # to wait on directories
target = ['#temp/DS2021_bayes/temp/model_data/students_grade8_2002_2012.csv',
          '#temp/DS2021_bayes/temp/model_data/students_grade8_2002_2012.dta',
          '#temp/DS2021_bayes/temp/model_data/schoolyear_grade8_shares_2002_2012.csv',
          '#temp/DS2021_bayes/temp/model_data/schoolyear_grade8_shares_2002_2012.dta',
          '#temp/DS2021_bayes/temp/model_data/school_grade8_shares_2002_2012.csv',
          '#temp/DS2021_bayes/temp/model_data/school_grade8_shares_2002_2012.dta']
env.Stata(target, source)

source = ['#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/script_demand.m',
          '#temp/DS2021_bayes/temp/model_data/students_grade8_2002_2012.csv',
          '#temp/DS2021_bayes/temp/model_data/schoolyear_grade8_shares_2002_2012.csv',
          '#temp/DS2021_bayes/temp/model_data/school_grade8_shares_2002_2012.csv',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_obj_demand.m',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_cmapping.m',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_s_hat.m',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_s_hat_individual.m']
target = ['#temp/DS2021_bayes/temp/model_data/grade_8_2002_2012_weights.mat',
          '#temp/DS2021_bayes/temp/model_data/grade_8_2002_2012_shock_weights.csv']
env.Matlab(target, source)

source = ['#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/script_demand_se.m',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_obj_demand_gradient.m',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_gradient.m',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_s_hat_individual.m',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_cmapping.m',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_s_hat.m',
          '#temp/DS2021_bayes/temp/model_data/grade_8_2002_2012_weights.mat']
target = ['#temp/DS2021_bayes/temp/model_data/se_weights.csv',
          '#temp/DS2021_bayes/temp/model_data/varcov_demand_weights.mat']
env.Matlab(target, source)

source = ['#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/script_demand_bootstrap_nonFEs.m',
          '#temp/DS2021_bayes/temp/model_data/grade_8_2002_2012_weights.mat',
          '#temp/DS2021_bayes/temp/model_data/varcov_demand_weights.mat']
target = ['#temp/DS2021_bayes/temp/model_data/demand_param_non_fe_bootstrap_weights.mat',
          '#temp/DS2021_bayes/temp/model_data/bootstrap/demand_param_non_fe_bootstrap_grade8.csv']
env.Matlab(target, source)

source = ['#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/script_demand_bootstrap_FEs.m',
          '#temp/DS2021_bayes/temp/model_data/demand_param_non_fe_bootstrap_weights.mat',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_cmapping_bootstrap.m',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_cmapping.m',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_s_hat.m',
          '#temp/DS2021_bayes/temp/model_data/grade_8_2002_2012_weights.mat']
target = ['#temp/DS2021_bayes/temp/model_data/bootstrap_deltas_grade_8.mat',
          '#temp/DS2021_bayes/temp/model_data/bootstrap/demand_param_fe_bootstrap_grade8.csv']
env.Matlab(target, source)

next_source = []
base_source = ['#temp/DS2021_bayes/temp/model_data/grade_year_enroll_2002_2012_weights.csv',
          '#temp/DS2021/temp/model_data/supply2002_annualexit_2002_2012_weights.csv',
          '#temp/DS2021_bayes/temp/model_data/grade_8_2002_2012_weights.mat',
          '#temp/DS2021_bayes/temp/model_data/bootstrap_deltas_grade_8.mat',
          '#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/func_supply_bootstrap.m',
          '#datastore/output/derived_large/bayes_bootstrap_weights/Dinerstein_Smith_2021_Weights.csv']

dir = f'{curdir}/temp/parallel'
if not os.path.exists(dir):
    os.makedirs(dir)

inpath, rep = 'source/derived/bayes_bootstrap', 100
outpath, paper = 'temp/parallel', 'Dinerstein_Smith_2021'
df_index = pd.read_csv(f'{curdir}/{inpath}/infrm_index.csv')
template = open(f'{curdir}/{inpath}/{paper}/wrapper_template.m', 'r')
lines = template.readlines()
template = template.close()
if not os.path.exists('temp/DS2021_bayes/temp/model_data/bootstrap'):
    os.makedirs('temp/DS2021_bayes/temp/model_data/bootstrap')

def get_rem_failures(wgt_dists, curdir):
    bootdir = f'{curdir}/temp/DS2021_bayes/temp/model_data/bootstrap'
    files = os.listdir(bootdir)
    repfile = np.array(['supply_full_annual_boot' in file for file in files])
    mats = np.array(['.mat' in file for file in files])
    using = np.empty((0, 1))
    for i in range(mats.shape[0]):
        if repfile[i] == True and mats[i] == True:
            using = np.append(using, files[i])
    rawnames = np.array([file.rstrip('.mat') for file in using])
    infrms = np.array([file.split('boot')[-1] for file in rawnames])
    wgtnum = np.array([int(index[0]) for index in infrms])
    reps = np.array([int(index.split('_')[-1]) for index in infrms])
    targets = np.array(range(1, 101))
    output = np.empty((0, 2))

    for curnum in wgt_dists:
        cor_reps = np.delete(reps, wgtnum != curnum)
        still_need = np.array([target not in cor_reps for target in targets])
        recomp_i = targets[still_need]
        to_comp = np.stack(([curnum for i in range(recomp_i.shape[0])], recomp_i), axis = 1)
        output = np.concatenate((output, to_comp), axis = 0).astype(int)
    return output

wgtarray = get_rem_failures([1, 2], curdir)#df_index['index'].unique())

for i, s in wgtarray:
    filename = f'{curdir}/{outpath}/wrapper_{i}_{s}.m'
    copy = open(filename, 'w')
    for line in lines:
        if '<infrm>' in line:
            line = line.replace('<infrm>', str(i))
        if '<index>' in line:
            line = line.replace('<index>', str(s))
        copy.write(line)
    copy.close()

    source = [filename] + base_source
    target = [f'#temp/DS2021_bayes/temp/model_data/bootstrap/supply_full_annual_boot{i}_{s}.mat',
            f'#temp/DS2021_bayes/temp/model_data/bootstrap/supply_full_annual_boot{i}_{s}.csv']
    next_source += target
    env.Matlab(target, source)

source = ['#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/script_export_supply.m',
          '#temp/DS2021/temp/model_data/supply_full_annual.mat',
          '#source/derived/bayes_bootstrap/infrm_index.csv']
source += next_source
target = [f'#temp/DS2021_bayes/model_data/supply_estimates_{num}.csv' for num in [1, 2]]# df_index['index'].unique()]
env.Matlab(target, source)

source = ['#source/derived/bayes_bootstrap/Dinerstein_Smith_2021/clean_weights.py',
          '#source/derived/bayes_bootstrap/infrm_index.csv']
source += [f'temp/DS2021_bayes/model_data/supply_estimates_{num}.csv' for num in [1, 2]]
target = ['#output/derived/bayes_bootstrap/Dinerstein_Smith_2021_Marginal.csv']
for i in df_index.loc[:, 'index']:
    source += [f'#temp/DS2021_bayes/temp/model_data/supply_estimates_{i}.csv']
env.Python(target, source)
