import os
import loom.runner
from distributions.fileutil import tempdir
from distributions.io.stream import protobuf_stream_load, protobuf_stream_dump
from loom.schema_pb2 import Query

def parse_response(message):
    response = Query.Response()
    response.ParseFromString(message)
    return response

def batch_predict(
        config_in,
        model_in,
        groups_in,
        requests,
        debug=False,
        profile=None):
    root = os.getcwd()
    with tempdir(cleanup_on_error=(not debug)):
        requests_in = os.path.abspath('requests.pbs.gz')
        responses_out = os.path.abspath('responses.pbs.gz')
        protobuf_stream_dump(
            (q.SerializeToString() for q in requests),
            requests_in)

        os.chdir(root)
        loom.runner.query(
            config_in=config_in,
            model_in=model_in,
            groups_in=groups_in,
            requests_in=requests_in,
            responses_out=responses_out,
            debug=debug,
            profile=profile)

        return map(parse_response, protobuf_stream_load(responses_out))

class PreqlServer(object):
    def __init__(self, query_server):
        self.query_server = query_server

    def predict(self, to_predict, conditioning_row=None, count=10):
        return self.query_server.sample(
            to_predict,
            conditioning_row=None,
            count=10)

    def mutual_information(self, cols1, cols2, conditioning_row=None):
        raise NotImplementedError("TODO")

make pyflakes happy for (broken) preql server