import argparse
import json
from connectome.abctask import ABCConnectomeTask, execute_task_list, GroundTruthData
from parallel import SGE, sge_available
from connectome.model import EXP
import scipy as sp

argument_parser = argparse.ArgumentParser(description="SMC ABC")
argument_parser.add_argument("db", type=str)
argument_parser.add_argument("--nr_parallel_jobs", type=int, default=7, help="Default: 7")
argument_parser.add_argument("--models", type=str, default="model_definitions.json", help="Default: model_definitions.json")
args = argument_parser.parse_args()

mapper = SGE(memory='3G', name="abc").map if sge_available() else map


def model_definitions():
    with open(args.models) as f:
        definitions = json.load(f)
    return definitions


exp = EXP(nr_neurons=2000, inh_ratio=.1, p_exc=.2, p_inh=.6)

ground_truth_samples = [GroundTruthData(exp(decay=decay), -1, {"decay": float(decay)})
                        for decay in sp.arange(.1, 1, .1)]


print("Start ABC Tasks", flush=True)
abc_tasks = [ABCConnectomeTask(model_definitions(), gt_sample, args.db, mapper) for gt_sample in ground_truth_samples]
execute_task_list(abc_tasks, args.nr_parallel_jobs)
