import argparse
import json

from connectome.builder import Sampler, DistributedFileStorer, PipelineBuilder, TimedModelWithPriorSample
from parallel import SGE, sge_available

parser = argparse.ArgumentParser(description='Sample networks.')
parser.add_argument('model_definitions', type=str, help='Model definition json')
parser.add_argument('nr_samples_per_model', type=int, help='Nr samples per model')
parser.add_argument('destination', type=str, help='Destination db')
parser.add_argument("--num_threads", type=int, default=1, help="Number of threads")
parser.add_argument("--priority", type=int, default=0, help="Number of threads")
args = parser.parse_args()


with open(args.model_definitions) as f:
    pars = json.load(f)
    modelbuilder = PipelineBuilder(pars)


def make_sample_job(name):
    model = modelbuilder.build_pipeline(name)
    prior = modelbuilder.build_prior(name)
    model_with_prior = TimedModelWithPriorSample(prior, model, name=name)
    return model_with_prior


map_engine = SGE(num_threads=args.num_threads, priority=args.priority, name="sample").map if sge_available() else map
sampler = Sampler(map_engine, DistributedFileStorer(args.destination, timeout=120))
for name in modelbuilder.defined_models():
    sampler.add_model(make_sample_job(name))
sampler.sample(args.nr_samples_per_model)
