from mesa.batchrunner import BatchRunner
import matplotlib.pyplot as plt
from mesa.space import MultiGrid
import numpy as np
from mesa.datacollection import DataCollector
import DTRegressionWithABM2 as MSM
import pylab as pl

fixed_params = {}
variable_params = {"N": range(100, 101, 2), "usshigh": np.arange(0.3, 1, 0.3), "M": np.arange(0, 0.8, 0.3), "C": np.arange(0.3, 1, 0.3)}        #"N": range(20,40, 2): from 20 to 39 (22,24,26,...); each one repeat iterations times

batch_run = BatchRunner(MSM.ABMModel,
#fixed_parameters=fixed_params,
variable_parameters=variable_params,
iterations=100,           #e.g.5
max_steps=26,        #e.g.100, number of ticks, We tell the BatchRunner to run 5 instantiations of the model with each number of agents, and to run each for 100 steps
model_reporters={ "retunrn_sum_USS_per_tick":MSM.retunrn_sum_USS_per_tick, "retunrn_sum_CSP_per_tick":MSM.retunrn_sum_CSP_per_tick, "retunrn_sum_shared_strategy_per_tick":MSM.retunrn_sum_shared_strategy_per_tick, "retunrn_USS_per_tick":MSM.return_USS_per_tick, "mod_shared_strategy_per_tick_model":MSM.mod_shared_strategy_per_tick_model, "return_strategy_per_tick":MSM.return_strategy_per_tick, "return_scenario":MSM.return_scenario, "return_number_actual_value_change_per_ticks_model":MSM.return_number_actual_value_change_per_ticks_model, "return_sum_NoLockdown_per_ticks_model":MSM.return_sum_NoLockdown_per_ticks_model, "return_sum_SocialDistancing_per_ticks_model":MSM.return_sum_SocialDistancing_per_ticks_model, "return_sum_SoftLockdown_per_ticks_model":MSM.return_sum_SoftLockdown_per_ticks_model, "return_sum_HardLockdown_per_ticks_model":MSM.return_sum_HardLockdown_per_ticks_model, "return_cases_per_tick":MSM.return_cases_per_tick, "return_unemployment_per_tick":MSM.return_unemployment_per_tick, "return_freedom_per_tick":MSM.return_freedom_per_tick, "return_ceses_thre_per_tick":MSM.return_ceses_thre_per_tick, "return_unemployment_thre_per_tick":MSM.return_unemployment_thre_per_tick, "return_freedom_thre_per_tick":MSM.return_freedom_thre_per_tick, "return_r_per_tick":MSM.return_r_per_tick, "return_unhappy_per_tick":MSM.return_unhappy_per_tick, "return_unhappy_cases_per_tick":MSM.return_unhappy_cases_per_tick, "return_unhappy_unemployment_per_tick":MSM.return_unhappy_unemployment_per_tick, "return_unhappy_freedom_per_tick":MSM.return_unhappy_freedom_per_tick, "return_region":MSM.return_region, "return_sum_USS_half": MSM.return_sum_USS_half, "return_sum_USS_last": MSM.return_sum_USS_last, "return_sum_CSP_last": MSM.return_sum_CSP_last, "return_sum_USS_initial": MSM.return_sum_USS_initial, "return_sum_CSP_initial": MSM.return_sum_CSP_initial, "return_list_institution": MSM.return_list_institution, "return_list_tick_emergeinstitution":MSM.return_list_tick_emergeinstitution, "sum_shared_strategy_half":MSM.sum_shared_strategy_half, "sum_shared_last_strategy": MSM.sum_shared_strategy_last}  ) #ct. Unlike the DataCollector, it won’t collect the data every step of the model, but only at the end of each run. lambda r: "Gini":r.x_old

#rule_compliance.append(1)
batch_run.run_all()
run_data = batch_run.get_model_vars_dataframe()
print(run_data.head())
run_data = batch_run.get_model_vars_dataframe()
run_data.to_csv('C:/CovidData/DataForTrainingDT/DataNew_20210222/result_Withoutzarb2run2v5_20210803.csv', header=True, float_format='%.3f')
