package com.tzld.piaoquan.recommend.server.service.score.model;


import com.tzld.piaoquan.recommend.server.util.CompressUtil;
import com.tzld.piaoquan.recommend.server.util.PropertiesUtil;
import ml.dmlc.xgboost4j.scala.DMatrix;
import ml.dmlc.xgboost4j.scala.spark.XGBoostClassificationModel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.List;
import java.util.Map;
import java.util.UUID;


public class XGBoostModel extends Model {
    private static final Logger LOGGER = LoggerFactory.getLogger(XGBoostModel.class);
    private XGBoostClassificationModel model;

    private String[] features;

    public void setFeatures(String[] features) {
        this.features = features;
    }

    @Override
    public int getModelSize() {
        if (this.model == null)
            return 0;
        return 1;
    }

    @Override
    public boolean loadFromStream(InputStreamReader in) throws Exception {
        return false;
    }

    public void cleanModel() {
        this.model = null;
    }

    public Float score(Map<String, Float> featureMap) {
        try {
            float[] values = new float[features.length];
            for (int i = 0; i < features.length; i++) {
                float v = featureMap.getOrDefault(features[i], 0.0f);
                values[i] = v;
            }
            DMatrix dm = new DMatrix(values, 1, features.length, 0.0f);
            float[][] result = model._booster().predict(dm, false, 0);
            return result[0][0];
        } catch (Exception e) {
            return 0f;
        }
    }

    @Override
    public boolean loadFromStream(InputStream in) throws Exception {
        Object localDir = params.getOrDefault("localDir",
                PropertiesUtil.getString("model.xgboost.path") + "/" + UUID.randomUUID());
        String modelDir = String.valueOf(localDir);
        CompressUtil.decompressGzFile(in, modelDir);
        String absolutePath = new File(modelDir).getAbsolutePath();
        XGBoostClassificationModel model2 = XGBoostClassificationModel.load("file://" + absolutePath);
        model2.setMissing(0.0f);
        if (params != null) {
            Object value = params.get("features");
            List<String> features = (List<String>) value;
            setFeatures(features.toArray(new String[features.size()]));
        }
        this.model = model2;
        return true;
    }

}
