
from .utils import DeferredActionCollection


class BatchRunner(object):
    '''
    Run multiple simulations and post-process summary results
    '''

    def __init__(self, endpoints):
        self.parameters = DeferredActionCollection('$', 'v.')
        self.parameters.model.defer()
        self.endpoints = endpoints

        self._retrieval = DeferredActionCollection('$', 'v.')
#    self._retrieval.instructions.append('result={}')

    def _run_single(self, parameter_set, endpoint, **kwargs):
        self.parameters.eval_script(parameter_set, {'v': endpoint})
        endpoint.drop_all_runs()
        ticket = endpoint.run_model(run_async=True, **kwargs)
        return (ticket, endpoint)

    def retrieve(self, var_name='y'):
        self._retrieval.instruction_prefix = "results['%s'] = v." % var_name
        return self._retrieval

    def _retrieve(self, ticket, endpoint):
        results = {}
        self._retrieval.eval_script({}, {'v': endpoint, 'results': results})
        return results

    def run(self, parameter_sets, **kwargs):
        if not len(self.endpoints):
            raise Exception('No model runners available')

        self.parameters.model.flush()

        available_endpoints = self.endpoints[:]
        jobs = []
        results = [None] * len(parameter_sets)

        for ix, row in parameter_sets.iterrows():
            print(ix, row)
            if len(available_endpoints) == 0:
                # Wait for something to be finished
                for i, (jx, row, ticket, endpoint) in enumerate(jobs):
                    if not ticket:
                        continue
                    results[i] = self._retrieve(ticket, endpoint)
                    # print(ix,row,results[i])
                    # print('Recycling %d',endpoint.port)
                    available_endpoints.append(endpoint)
                    jobs[i] = (jx, row, None, endpoint)
                    break

            endpoint_to_use = available_endpoints.pop()

            ticket, endpoint = self._run_single(row, endpoint_to_use, **kwargs)
            jobs.append((ix, row, ticket, endpoint))

        for i, (jx, row, ticket, endpoint) in enumerate(jobs):
            if not ticket:
                continue
            results[i] = self._retrieve(ticket, endpoint)
        return jobs, results
