import torch
import zmq
import flatbuffers
from termcolor import colored

from . import util, state, __version__
from .distributions import Uniform, Normal, Categorical, Poisson, Bernoulli, Beta, Exponential, Gamma, LogNormal, Binomial, Weibull
from .ppx import Message as ppx_Message
from .ppx import MessageBody as ppx_MessageBody
from .ppx import Tensor as ppx_Tensor
from .ppx import Distribution as ppx_Distribution
from .ppx import Uniform as ppx_Uniform
from .ppx import Normal as ppx_Normal
from .ppx import Categorical as ppx_Categorical
from .ppx import Poisson as ppx_Poisson
from .ppx import Bernoulli as ppx_Bernoulli
from .ppx import Beta as ppx_Beta
from .ppx import Exponential as ppx_Exponential
from .ppx import Gamma as ppx_Gamma
from .ppx import LogNormal as ppx_LogNormal
from .ppx import Binomial as ppx_Binomial
from .ppx import Weibull as ppx_Weibull
from .ppx import Handshake as ppx_Handshake
from .ppx import HandshakeResult as ppx_HandshakeResult
from .ppx import Run as ppx_Run
from .ppx import RunResult as ppx_RunResult
from .ppx import Sample as ppx_Sample
from .ppx import SampleResult as ppx_SampleResult
from .ppx import Observe as ppx_Observe
from .ppx import ObserveResult as ppx_ObserveResult
from .ppx import Tag as ppx_Tag
from .ppx import TagResult as ppx_TagResult
from .ppx import Reset as ppx_Reset


class ZMQRequester():
    def __init__(self, server_address):
        self._server_address = server_address
        self._context = zmq.Context.instance()
        self._socket = self._context.socket(zmq.REQ)
        self._socket.setsockopt(zmq.LINGER, 100)
        print('ppx (Python): zmq.REQ socket connecting to server {}'.format(self._server_address))
        self._socket.connect(self._server_address)

    def __enter__(self):
        return self

    def __exit__(self, exception_type, exception_value, traceback):
        self.close()

    def __del__(self):
        self.close()

    def close(self):
        if not self._socket.closed:
            self._socket.close()
            self._context.term()
            print('ppx (Python): zmq.REQ socket disconnected from server {}'.format(self._server_address))

    def send_request(self, request):
        self._socket.send(request)

    def receive_reply(self):
        return self._socket.recv()


class ModelServer(object):
    def __init__(self, server_address):
        self._requester = ZMQRequester(server_address)
        self.system_name, self.model_name = self._handshake()
        print('ppx (Python): This system        : {}'.format(colored('pyprob {}'.format(__version__), 'green')))
        print('ppx (Python): Connected to system: {}'.format(colored(self.system_name, 'green')))
        print('ppx (Python): Model name         : {}'.format(colored(self.model_name, 'green', attrs=['bold'])))

    def __enter__(self):
        return self

    def __exit__(self, exception_type, exception_value, traceback):
        self.close()

    def __del__(self):
        self.close()

    def close(self):
        self._requester.close()

    def _protocol_tensor_to_variable(self, protocol_tensor):
        if protocol_tensor is None:
            return None
        data = protocol_tensor.DataAsNumpy()
        shape = protocol_tensor.ShapeAsNumpy()
        if len(data) == 0:
            return None
        else:
            t = torch.from_numpy(data)
        if len(shape) != 0:
            t = t.view(shape.tolist())
        return util.to_tensor(t)

    def _variable_to_protocol_tensor(self, builder, variable):
        if variable is None:
            variable = util.to_tensor(torch.zeros(0))
        variable_numpy = util.to_numpy(variable)
        data = variable_numpy.flatten().tolist()
        shape = list(variable_numpy.shape)

        # pack data
        ppx_Tensor.TensorStartDataVector(builder, len(data))
        for d in reversed(data):
            builder.PrependFloat64(d)
        data = builder.EndVector(len(data))

        # pack shape
        ppx_Tensor.TensorStartShapeVector(builder, len(shape))
        for s in reversed(shape):
            builder.PrependInt32(s)
        shape = builder.EndVector(len(shape))

        ppx_Tensor.TensorStart(builder)
        ppx_Tensor.TensorAddData(builder, data)
        ppx_Tensor.TensorAddShape(builder, shape)
        return ppx_Tensor.TensorEnd(builder)

    def _get_message_body(self, message_buffer):
        message = ppx_Message.Message.GetRootAsMessage(message_buffer, 0)
        body_type = message.BodyType()
        if body_type == ppx_MessageBody.MessageBody().HandshakeResult:
            message_body = ppx_HandshakeResult.HandshakeResult()
        elif body_type == ppx_MessageBody.MessageBody().RunResult:
            message_body = ppx_RunResult.RunResult()
        elif body_type == ppx_MessageBody.MessageBody().Sample:
            message_body = ppx_Sample.Sample()
        elif body_type == ppx_MessageBody.MessageBody().Observe:
            message_body = ppx_Observe.Observe()
        elif body_type == ppx_MessageBody.MessageBody().Tag:
            message_body = ppx_Tag.Tag()
        elif body_type == ppx_MessageBody.MessageBody().Reset:
            message_body = ppx_Reset.Reset()
        else:
            raise RuntimeError('ppx (Python): Received unexpected message body type: {}'.format(body_type))
        message_body.Init(message.Body().Bytes, message.Body().Pos)
        return message_body

    def _handshake(self):
        builder = flatbuffers.Builder(64)
        # consturct MessageBody
        system_name = builder.CreateString('pyprob {}'.format(__version__))
        ppx_Handshake.HandshakeStart(builder)
        ppx_Handshake.HandshakeAddSystemName(builder, system_name)
        message_body = ppx_Handshake.HandshakeEnd(builder)

        # construct Message
        ppx_Message.MessageStart(builder)
        ppx_Message.MessageAddBodyType(builder, ppx_MessageBody.MessageBody().Handshake)
        ppx_Message.MessageAddBody(builder, message_body)
        message = ppx_Message.MessageEnd(builder)
        builder.Finish(message)

        message = builder.Output()
        self._requester.send_request(message)

        reply = self._requester.receive_reply()
        message_body = self._get_message_body(reply)
        if isinstance(message_body, ppx_HandshakeResult.HandshakeResult):
            system_name = message_body.SystemName().decode('utf-8')
            model_name = message_body.ModelName().decode('utf-8')
            return system_name, model_name
        else:
            raise RuntimeError('ppx (Python): Unexpected reply to handshake.')

    def forward(self):
        builder = flatbuffers.Builder(64)

        # construct MessageBody
        ppx_Run.RunStart(builder)
        message_body = ppx_Run.RunEnd(builder)

        # construct Message
        ppx_Message.MessageStart(builder)
        ppx_Message.MessageAddBodyType(builder, ppx_MessageBody.MessageBody().Run)
        ppx_Message.MessageAddBody(builder, message_body)
        message = ppx_Message.MessageEnd(builder)
        builder.Finish(message)

        message = builder.Output()
        self._requester.send_request(message)

        while True:
            reply = self._requester.receive_reply()
            message_body = self._get_message_body(reply)

            if isinstance(message_body, ppx_RunResult.RunResult):
                result = self._protocol_tensor_to_variable(message_body.Result())
                return result
            elif isinstance(message_body, ppx_Sample.Sample):
                address = message_body.Address().decode('utf-8')
                name = message_body.Name().decode('utf-8')
                if name == '':
                    name = None
                control = bool(message_body.Control())
                replace = bool(message_body.Replace())
                distribution_type = message_body.DistributionType()
                if distribution_type == ppx_Distribution.Distribution().Uniform:
                    uniform = ppx_Uniform.Uniform()
                    uniform.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    low = self._protocol_tensor_to_variable(uniform.Low())
                    high = self._protocol_tensor_to_variable(uniform.High())
                    dist = Uniform(low, high)
                elif distribution_type == ppx_Distribution.Distribution().Normal:
                    normal = ppx_Normal.Normal()
                    normal.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    mean = self._protocol_tensor_to_variable(normal.Mean())
                    stddev = self._protocol_tensor_to_variable(normal.Stddev())
                    dist = Normal(mean, stddev)
                elif distribution_type == ppx_Distribution.Distribution().Categorical:
                    categorical = ppx_Categorical.Categorical()
                    categorical.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    probs = self._protocol_tensor_to_variable(categorical.Probs())
                    dist = Categorical(probs)
                elif distribution_type == ppx_Distribution.Distribution().Poisson:
                    poisson = ppx_Poisson.Poisson()
                    poisson.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    rate = self._protocol_tensor_to_variable(poisson.Rate())
                    dist = Poisson(rate)
                elif distribution_type == ppx_Distribution.Distribution().Bernoulli:
                    bernoulli = ppx_Bernoulli.Bernoulli()
                    bernoulli.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    probs = self._protocol_tensor_to_variable(bernoulli.Probs())
                    dist = Bernoulli(probs)
                elif distribution_type == ppx_Distribution.Distribution().Beta:
                    beta = ppx_Beta.Beta()
                    beta.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    concentration1 = self._protocol_tensor_to_variable(beta.Concentration1())
                    concentration0 = self._protocol_tensor_to_variable(beta.Concentration0())
                    dist = Beta(concentration1, concentration0)
                elif distribution_type == ppx_Distribution.Distribution().Exponential:
                    exponential = ppx_Exponential.Exponential()
                    exponential.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    rate = self._protocol_tensor_to_variable(exponential.Rate())
                    dist = Exponential(rate)
                elif distribution_type == ppx_Distribution.Distribution().Gamma:
                    gamma = ppx_Gamma.Gamma()
                    gamma.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    concentration = self._protocol_tensor_to_variable(gamma.Concentration())
                    rate = self._protocol_tensor_to_variable(gamma.Rate())
                    dist = Gamma(concentration, rate)
                elif distribution_type == ppx_Distribution.Distribution().LogNormal:
                    log_normal = ppx_LogNormal.LogNormal()
                    log_normal.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    loc = self._protocol_tensor_to_variable(log_normal.Loc())
                    scale = self._protocol_tensor_to_variable(log_normal.Scale())
                    dist = LogNormal(loc, scale)
                elif distribution_type == ppx_Distribution.Distribution().Binomial:
                    binomial = ppx_Binomial.Binomial()
                    binomial.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    total_count = self._protocol_tensor_to_variable(binomial.TotalCount())
                    probs = self._protocol_tensor_to_variable(binomial.Probs())
                    dist = Binomial(total_count, probs)
                elif distribution_type == ppx_Distribution.Distribution().Weibull:
                    weibull = ppx_Weibull.Weibull()
                    weibull.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    scale = self._protocol_tensor_to_variable(weibull.Scale())
                    concentration = self._protocol_tensor_to_variable(weibull.Concentration())
                    dist = Weibull(scale, concentration)
                else:
                    raise RuntimeError('ppx (Python): Sample from an unexpected distribution requested.')
                result = state.sample(distribution=dist, control=control, replace=replace, name=name, address=address)
                builder = flatbuffers.Builder(64)
                result = self._variable_to_protocol_tensor(builder, result)
                ppx_SampleResult.SampleResultStart(builder)
                ppx_SampleResult.SampleResultAddResult(builder, result)
                message_body = ppx_SampleResult.SampleResultEnd(builder)

                # construct Message
                ppx_Message.MessageStart(builder)
                ppx_Message.MessageAddBodyType(builder, ppx_MessageBody.MessageBody().SampleResult)
                ppx_Message.MessageAddBody(builder, message_body)
                message = ppx_Message.MessageEnd(builder)
                builder.Finish(message)

                message = builder.Output()
                self._requester.send_request(message)
            elif isinstance(message_body, ppx_Observe.Observe):
                address = message_body.Address().decode('utf-8')
                name = message_body.Name().decode('utf-8')
                if name == '':
                    name = None
                value = self._protocol_tensor_to_variable(message_body.Value())
                distribution_type = message_body.DistributionType()
                if distribution_type == ppx_Distribution.Distribution().NONE:
                    dist = None
                elif distribution_type == ppx_Distribution.Distribution().Uniform:
                    uniform = ppx_Uniform.Uniform()
                    uniform.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    low = self._protocol_tensor_to_variable(uniform.Low())
                    high = self._protocol_tensor_to_variable(uniform.High())
                    dist = Uniform(low, high)
                elif distribution_type == ppx_Distribution.Distribution().Normal:
                    normal = ppx_Normal.Normal()
                    normal.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    mean = self._protocol_tensor_to_variable(normal.Mean())
                    stddev = self._protocol_tensor_to_variable(normal.Stddev())
                    dist = Normal(mean, stddev)
                elif distribution_type == ppx_Distribution.Distribution().Categorical:
                    categorical = ppx_Categorical.Categorical()
                    categorical.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    probs = self._protocol_tensor_to_variable(categorical.Probs())
                    dist = Categorical(probs)
                elif distribution_type == ppx_Distribution.Distribution().Poisson:
                    poisson = ppx_Poisson.Poisson()
                    poisson.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    rate = self._protocol_tensor_to_variable(poisson.Rate())
                    dist = Poisson(rate)
                elif distribution_type == ppx_Distribution.Distribution().Bernoulli:
                    bernoulli = ppx_Bernoulli.Bernoulli()
                    bernoulli.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    probs = self._protocol_tensor_to_variable(bernoulli.Probs())
                    dist = Bernoulli(probs)
                elif distribution_type == ppx_Distribution.Distribution().Beta:
                    beta = ppx_Beta.Beta()
                    beta.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    concentration1 = self._protocol_tensor_to_variable(beta.Concentration1())
                    concentration0 = self._protocol_tensor_to_variable(beta.Concentration0())
                    dist = Beta(concentration1, concentration0)
                elif distribution_type == ppx_Distribution.Distribution().Exponential:
                    exponential = ppx_Exponential.Exponential()
                    exponential.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    rate = self._protocol_tensor_to_variable(exponential.Rate())
                    dist = Exponential(rate)
                elif distribution_type == ppx_Distribution.Distribution().Gamma:
                    gamma = ppx_Gamma.Gamma()
                    gamma.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    concentration = self._protocol_tensor_to_variable(gamma.Concentration())
                    rate = self._protocol_tensor_to_variable(gamma.Rate())
                    dist = Gamma(concentration, rate)
                elif distribution_type == ppx_Distribution.Distribution().LogNormal:
                    log_normal = ppx_LogNormal.LogNormal()
                    log_normal.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    loc = self._protocol_tensor_to_variable(log_normal.Loc())
                    scale = self._protocol_tensor_to_variable(log_normal.Scale())
                    dist = LogNormal(loc, scale)
                elif distribution_type == ppx_Distribution.Distribution().Binomial:
                    binomial = ppx_Binomial.Binomial()
                    binomial.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    total_count = self._protocol_tensor_to_variable(binomial.TotalCount())
                    probs = self._protocol_tensor_to_variable(binomial.Probs())
                    dist = Binomial(total_count, probs)
                elif distribution_type == ppx_Distribution.Distribution().Weibull:
                    weibull = ppx_Weibull.Weibull()
                    weibull.Init(message_body.Distribution().Bytes, message_body.Distribution().Pos)
                    scale = self._protocol_tensor_to_variable(weibull.Scale())
                    concentration = self._protocol_tensor_to_variable(weibull.Concentration())
                    dist = Weibull(scale, concentration)
                else:
                    raise RuntimeError('ppx (Python): Sample from an unexpected distribution requested: {}'.format(distribution_type))

                state.observe(distribution=dist, value=value, name=name, address=address)
                builder = flatbuffers.Builder(64)
                ppx_ObserveResult.ObserveResultStart(builder)
                message_body = ppx_ObserveResult.ObserveResultEnd(builder)

                # construct Message
                ppx_Message.MessageStart(builder)
                ppx_Message.MessageAddBodyType(builder, ppx_MessageBody.MessageBody().ObserveResult)
                ppx_Message.MessageAddBody(builder, message_body)
                message = ppx_Message.MessageEnd(builder)
                builder.Finish(message)

                message = builder.Output()
                self._requester.send_request(message)
            elif isinstance(message_body, ppx_Tag.Tag):
                address = message_body.Address().decode('utf-8')
                name = message_body.Name().decode('utf-8')
                if name == '':
                    name = None
                value = self._protocol_tensor_to_variable(message_body.Value())
                state.tag(value=value, name=name, address=address)
                builder = flatbuffers.Builder(64)
                ppx_TagResult.TagResultStart(builder)
                message_body = ppx_TagResult.TagResultEnd(builder)

                # construct Message
                ppx_Message.MessageStart(builder)
                ppx_Message.MessageAddBodyType(builder, ppx_MessageBody.MessageBody().TagResult)
                ppx_Message.MessageAddBody(builder, message_body)
                message = ppx_Message.MessageEnd(builder)
                builder.Finish(message)

                message = builder.Output()
                self._requester.send_request(message)
            elif isinstance(message_body, ppx_Reset.Reset):
                raise RuntimeError('ppx (Python): Received a reset request. Protocol out of sync.')
            else:
                raise RuntimeError('ppx (Python): Received unexpected message.')
