package com.dataiku.exports;

import com.dataiku.scoring.*;

public class ModelPredictor {

    public void predict(List<String[]> inputData){
        for(String[] input : inputData) {
            predict();
        }
    }

    public String[] predictBatch(List<String[]> inputData) {
        List<String[]> outputDataset = new ArrayList<>();
         for(String[] iRow : inputDataset) {
            String[] oRow = predict(header,iRow);
            outputDataset.add(oRow);
        }
        return outputDataset;
    }
    

    private String[] predict(String[] featureName, String[] inputRow) {
        List<String> outputRow = new ArrayList<>(Arrays.asList(inputRow));
        Observation.Builder obsBuilder = model.observationBuilder();
        Map<String,String> features = new HashMap<>();
        for(int i = 0 ; i < featureName.length;i++){
            features.put(featureName[i],inputRow[i]);
        }
        Observation obs = obsBuilder.withAll(features).build();
        if (obs.hasError()) {
            System.err.println("Can't build observation: " + obs.getErrorMessage());
            // maybe throw here
        }

        // For a classification model
        Try<ClassificationResult> prediction = model.predict(obs);
        if (prediction.isError()) {
            System.err.println("Can't make a prediction: " + prediction.getMessage());
            // maybe throw here
        } else {
            ClassificationResult result = prediction.get();
            // predictedClass is one of model.getClassLabels()
            String predictedClass = result.getPrediction();
            // probabilities has the same indices as model.getClassLabels()
            // i.e. 0 to (model.getNumClasses() - 1)
            double[] probabilities = result.getProbabilities();
            // Use result here
            System.out.println("Result of prediction is " + predictedClass + "("  + Arrays.toString(probabilities) + ")");
            outputRow.add(predictedClass);
            outputRow.add(Double.toString(probabilities[0]));
            outputRow.add(Double.toString(probabilities[1]));
        }
        return outputRow.toArray(new String[outputRow.size()]);
    }    

}
