import ultra.Iris.UFL_Types;
import ultra.ws.Iris.PRF_WebService.Xtraining;

consume {
    debug(input);
    if (instanceOf(input, WSCycle_predict)) {
        WSCycle_predict cycle = (WSCycle_predict) input;
        PredictObservation obs = udrCreate(PredictObservation);
        obs.observation = listCreate(double, cycle.param.sepal_length, cycle.param.sepal_width, 
            cycle.param.petal_length, cycle.param.petal_width);
        obs.context = cycle;
        udrRoute(obs, "observation");
    }
    if (instanceOf(input, PredictObservation)) {
        PredictObservation obs = (PredictObservation) input;
        WSCycle_predict cycle = (WSCycle_predict) obs.context;
        Response response = udrCreate(Response);
        response.iris = obs.prediction;
        cycle.response = response;
        udrRoute(cycle, "response");
    }
}

import pickle

model = None
targets = None

def consume(input):
    if isinstance(input, InstallModel):
        global model, targets
        model = pickle.loads(input.model)
        targets = input.targets
    elif isinstance(input, PredictObservation):
        if model:
            idx = model.predict([input.observation])[0]
            input.prediction = targets[idx]
        else:
            input.prediction = 'please install model'
        if input.testing:
            debug(input)
            udrRoute(input, 'test_prediction')
        else:
            udrRoute(input, 'prediction')

import pickle
from .UFL_Types import InstallModel
from .UFL_Types import PredictObservation

def install_model(model, targets):
    udr = InstallModel(
        model=pickle.dumps(model),
        targets=targets)
    udrRoute(udr)

def predict(observation):
    udr = PredictObservation(
        observation=observation,
        testing=True)
    udrRoute(udr)
    return udrConsume().prediction

__all__ = ['install_model', 'predict']