/*
 * Copyright (c)  2016, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
 *
 * WSO2 Inc. licenses this file to you under the Apache License,
 * Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.wso2.extension.siddhi.execution.ml;

import org.apache.commons.io.IOUtils;
import org.wso2.carbon.ml.commons.constants.MLConstants;
import org.wso2.carbon.ml.commons.domain.Feature;
import org.wso2.carbon.ml.commons.domain.MLModel;
import org.wso2.carbon.ml.core.exceptions.MLInputAdapterException;
import org.wso2.carbon.ml.core.exceptions.MLModelHandlerException;
import org.wso2.carbon.ml.core.factories.DatasetType;
import org.wso2.carbon.ml.core.h2o.POJOPredictor;
import org.wso2.carbon.ml.core.impl.MLIOFactory;
import org.wso2.carbon.ml.core.impl.Predictor;
import org.wso2.carbon.ml.core.interfaces.MLInputAdapter;
import org.wso2.carbon.ml.core.utils.MLCoreServiceValueHolder;

import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;

/**
 * Model Handler
 */
public class ModelHandler {

    public static final String FILE_STORAGE_PREFIX = "file";
    public static final String REGISTRY_STORAGE_PREFIX = "registry";
    public static final String PATH_TO_GOVERNANCE_REGISTRY = "/_system/governance";

    private MLModel mlModel;
    private long modelId;

    /**
     *
     * @param modelStorageLocation MLModel storage location
     * @throws ClassNotFoundException class not found exception
     * @throws java.net.URISyntaxException URI syntax exception
     * @throws MLInputAdapterException ML input adapter exception
     * @throws java.io.IOException java IO exception
     */
    public ModelHandler(String modelStorageLocation)
            throws ClassNotFoundException, URISyntaxException, MLInputAdapterException,
            IOException {
        mlModel = retrieveModel(modelStorageLocation);
        Random random = new Random();
        modelId = random.nextLong();
    }

    /**
     * Retrieve the MLModel from the storage location.
     * 
     * @param modelStorageLocation model storage location (file path or registry path)
     * @return the deserialized MLModel object
     * @throws URISyntaxException
     * @throws MLInputAdapterException
     * @throws IOException
     * @throws ClassNotFoundException
     */
    private static MLModel retrieveModel(String modelStorageLocation)
            throws URISyntaxException, MLInputAdapterException, IOException,
            ClassNotFoundException {

        String[] modelStorage = modelStorageLocation.trim().split(":");
        String storageType = modelStorage[0];
        if (storageType.equals(REGISTRY_STORAGE_PREFIX)) {
            if (modelStorage[1].startsWith(PATH_TO_GOVERNANCE_REGISTRY)) {
                modelStorageLocation = modelStorage[1].substring(
                        PATH_TO_GOVERNANCE_REGISTRY.length());
            } else {
                modelStorageLocation = modelStorage[1];
            }
        } else if (storageType.equals(FILE_STORAGE_PREFIX)) {
            modelStorageLocation = modelStorage[1];
            storageType = DatasetType.FILE.getValue();
        } else {
            storageType = DatasetType.FILE.getValue();
        }

        MLIOFactory ioFactory = new MLIOFactory(
                MLCoreServiceValueHolder.getInstance().getMlProperties());
        MLInputAdapter inputAdapter = ioFactory.getInputAdapter(
                storageType + MLConstants.IN_SUFFIX);
        InputStream in = inputAdapter.read(modelStorageLocation);
        ObjectInputStream ois = new ObjectInputStream(in);
        MLModel mlModel = null;
        try {
            mlModel = (MLModel) ois.readObject();
        } finally {
            IOUtils.closeQuietly(in);
            IOUtils.closeQuietly(ois);
        }

        in.close();
        ois.close();
        return mlModel;

    }

    /**
     * Predict the value using the feature values.
     * 
     * @param data feature values array
     * @param outputType data type of the output
     * @return predicted value
     * @throws MLModelHandlerException throws ML model handler exception
     */
    public Object predict(String[] data, String outputType) throws MLModelHandlerException {
        ArrayList<String[]> list = new ArrayList<String[]>();
        list.add(data);
        Predictor predictor = new Predictor(modelId, mlModel, list);
        List<?> predictions = predictor.predict();
        String predictionStr = predictions.get(0).toString();
        Object prediction = castValue(outputType, predictionStr);
        return prediction;
    }

    /**
     * Predict the value using the feature values.
     * 
     * @param data feature values array
     * @param outputType output type
     * @param percentile percentile value for predictions
     * @return predicted value
     * @throws MLModelHandlerException throws ML model handler exception
     */
    public Object predict(String[] data, String outputType, double percentile)
            throws MLModelHandlerException {
        ArrayList<String[]> list = new ArrayList<String[]>();
        list.add(data);
        Predictor predictor = new Predictor(modelId, mlModel, list, percentile, false);
        List<?> predictions = predictor.predict();
        String predictionStr = predictions.get(0).toString();
        Object prediction = castValue(outputType, predictionStr);
        return prediction;
    }

    /**
     * Predict the value using the feature values with POJO predictor.
     * 
     * @param data feature values array
     * @param outputType data type of the output
     * @param pojoPredictor POJO predictor
     * @return predicted value
     * @throws MLModelHandlerException ML model handler exception
     */
    public Object predict(String[] data, String outputType, POJOPredictor pojoPredictor)
            throws MLModelHandlerException {
        String predictionStr = pojoPredictor.predict(data).toString();
        Object prediction = castValue(outputType, predictionStr);
        return prediction;
    }

    /**
     * Cast the given value to the given output type.
     * 
     * @param outputType Output data type
     * @param value value to be casted in String
     * @return Value casted to output type object
     */
    private Object castValue(String outputType, String value) {
        if (outputType.equalsIgnoreCase("double")) {
            return Double.parseDouble(value);
        } else if (outputType.equalsIgnoreCase("float")) {
            return Float.parseFloat(value);
        } else if (outputType.equalsIgnoreCase("integer") || outputType.equalsIgnoreCase("int")) {
            return Integer.parseInt(value);
        } else if (outputType.equalsIgnoreCase("long")) {
            return Long.parseLong(value);
        } else if (outputType.equalsIgnoreCase("boolean") || outputType.equalsIgnoreCase("bool")) {
            return Boolean.parseBoolean(value);
        } else {
            return value;
        }
    }

    /**
     * Return the map containing &lt;feature-name, index&gt; pairs
     *
     * @return the &lt;feature-name, feature-index&gt; map of the MLModel
     */
    public Map<String, Integer> getFeatures() {
        List<Feature> features = mlModel.getFeatures();
        Map<String, Integer> featureIndexMap = new HashMap<String, Integer>();
        for (Feature feature : features) {
            featureIndexMap.put(feature.getName(), feature.getIndex());
        }
        return featureIndexMap;
    }

    /**
     * Get new to old indices list of this model.
     * 
     * @return the new to old indices list of the MLModel
     */
    public List<Integer> getNewToOldIndicesList() {
        return mlModel.getNewToOldIndicesList();
    }

    /**
     * Return the response variable of the model
     * 
     * @return the response variable of the MLModel
     */
    public String getResponseVariable() {
        return mlModel.getResponseVariable();
    }

    /**
     * Returns the algorithm class - classification, numerical prediction or clustering
     * 
     * @return the algorithm class
     */
    public String getAlgorithmClass() {
        return mlModel.getAlgorithmClass();
    }

    /**
     * @return the model
     */
    public MLModel getMlModel() {
        return mlModel;
    }
}
