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

import com.aliyun.openservices.eas.predict.http.HttpConfig;
import com.aliyun.openservices.eas.predict.http.PredictClient;
import com.aliyun.openservices.eas.predict.proto.PredictProtos;
import com.aliyun.openservices.eas.predict.request.TFDataType;
import com.aliyun.openservices.eas.predict.request.TFRequest;
import com.aliyun.openservices.eas.predict.response.TFResponse;
import com.tzld.piaoquan.recommend.server.common.base.RankItem;
import com.tzld.piaoquan.recommend.server.model.Tuple4;
import com.tzld.piaoquan.recommend.server.service.rank.tansform.FGEncoder;
import com.tzld.piaoquan.recommend.server.service.rank.tansform.FGEncoderManager;
import com.tzld.piaoquan.recommend.server.service.score.DnnConfigParam;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.math.NumberUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.Assert;

import java.util.*;

public class PAIModel {

    private static final Logger LOGGER = LoggerFactory.getLogger(PAIModel.class);

    private final String fgConfigFilePath;
    private final String easToken;
    private final String easEndpoint;
    private final String easDirectEndpoint;
    private final String modelName;
    private final List<String> fetchs;

    private PredictClient predictClient;
    private final List<Tuple4<String, String, String, TFDataType>> featureKeys = new ArrayList<>();

    public PAIModel(DnnConfigParam param) {
        this(param.getFgConfigFilePath(), param.getEasToken(), param.getEasEndpoint(), param.getEasDirectEndpoint(), param.getModelName(), param.getFetchs());
    }

    public PAIModel(String fgConfigFilePath, String easToken, String easEndpoint, String easDirectEndpoint, String modelName, List<String> fetchs) {
        Assert.hasText(fgConfigFilePath, "fgConfigFilePath must not be empty");
        Assert.hasText(easToken, "easToken must not be empty");
        Assert.hasText(modelName, "modelName must not be empty");
        Assert.notEmpty(fetchs, "fetch must not be empty");
        Assert.noNullElements(fetchs, "fetch all item must not be empty");

        Assert.isTrue(StringUtils.isNotBlank(easEndpoint) || StringUtils.isNotBlank(easDirectEndpoint), "at least one of easEndpoint or easDirectEndpoint must not be empty");

        this.fgConfigFilePath = fgConfigFilePath;
        this.easToken = easToken;
        this.easEndpoint = easEndpoint;
        this.easDirectEndpoint = easDirectEndpoint;
        this.modelName = modelName;
        this.fetchs = fetchs;

        this.parseFeatureKey();
        this.initPredictClient();
    }

    private void parseFeatureKey() {
        // 初始化fg key与特征生产时的key的映射
        FGEncoder fgEncoder = FGEncoderManager.getInstance().getFGEncoder(fgConfigFilePath);
        for (FGEncoder.FeatureDef featureDef : fgEncoder.getFeatureDefs()) {
            Tuple4<String, String, String, TFDataType> featureKeyTuple = new Tuple4<>();

            featureKeyTuple.setV1(featureDef.feature_key);
            featureKeyTuple.setV2(featureDef.feature_name);
            featureKeyTuple.setV3(featureDef.default_value);
            featureKeyTuple.setV4(featureDef.typeCoverToTFDataType());
            featureKeys.add(featureKeyTuple);
        }

    }

    private void initPredictClient() {
        predictClient = new PredictClient(new HttpConfig());
        predictClient.setModelName(this.modelName);
        predictClient.setToken(this.easToken);
        if (StringUtils.isNotBlank(this.easDirectEndpoint)) {
            predictClient.setDirectEndpoint(this.easDirectEndpoint);
        } else {
            predictClient.setEndpoint(this.easEndpoint);
        }
    }

    public Map<String, List<Float>> score(List<RankItem> rankItems) {

        int size = rankItems.size();
        String currentName = "";
        // 生成特征名与特征值数组的映射
        Map<String, double[]> doubleFeatureArr = new HashMap<>();
        Map<String, String[]> stringFeatureArr = new HashMap<>();

        try {

            List<Map<String, Float>> featureMapList = new ArrayList<>(rankItems.size());
            List<Map<String, String>> featureMapStringList = new ArrayList<>(rankItems.size());

            for (RankItem rankItem : rankItems) {
                featureMapList.add(Optional.ofNullable(rankItem.featureMap).orElse(new HashMap<>()));
                featureMapStringList.add(Optional.ofNullable(rankItem.featureMapString).orElse(new HashMap<>()));
            }

            for (Tuple4<String, String, String, TFDataType> featureKeyTuple : featureKeys) {
                TFDataType dataType = featureKeyTuple.getV4();
                String featureName = featureKeyTuple.getV2();
                String featureKey = featureKeyTuple.getV1();
                String defaultValue = featureKeyTuple.getV3();
                currentName = String.format("1@%s@%s", featureKey, featureName);
                if (TFDataType.DT_STRING.equals(dataType)) {
                    for (int i = 0; i < size; i++) {
                        Map<String, String> featureMapString = featureMapStringList.get(i);
                        String[] strArr = stringFeatureArr.computeIfAbsent(featureName, s -> new String[size]);
                        String value = featureMapString.get(featureKey);
                        strArr[i] = Objects.nonNull(value) ? value : defaultValue;
                    }
                } else {
                    for (int i = 0; i < size; i++) {
                        Map<String, Float> featureMap = featureMapList.get(i);
                        double[] doubleArr = doubleFeatureArr.computeIfAbsent(featureName, s -> new double[size]);
                        double value = featureMap.getOrDefault(featureKey, NumberUtils.toFloat(defaultValue, 0f));
                        doubleArr[i] = value;
                    }
                }
            }

            TFRequest request = new TFRequest();
            long[] shape = new long[]{size};
            for (Tuple4<String, String, String, TFDataType> featureKey : featureKeys) {
                String featureName = featureKey.getV2();
                TFDataType dataType = featureKey.getV4();
                currentName = String.format("2@%s@%s", featureKey.getV1(), featureName);
                if (TFDataType.DT_STRING.equals(dataType)) {
                    request.addFeed(featureName, TFDataType.DT_STRING, shape, stringFeatureArr.get(featureName));
                } else {
                    request.addFeed(featureName, TFDataType.DT_DOUBLE, shape, doubleFeatureArr.get(featureName));
                }
            }

            currentName = "";

            for (String fetch : this.fetchs) {
                request.addFetch(fetch);
            }

            TFResponse response = predictClient.predict(request);
            Map<String, PredictProtos.ArrayProto> outputsMap = response.getOutputsMap();
            // LOGGER.info("response keys: {}", outputsMap.keySet());
            Map<String, List<Float>> resultMap = new HashMap<>();
            for (String scoreKey : this.fetchs) {
                resultMap.put(scoreKey, response.getFloatVals(scoreKey));
            }
            return resultMap;
        } catch (Exception e) {
            LOGGER.error("pai eas error currentName: {}", currentName, e);
            Map<String, List<Float>> resultMap = new HashMap<>();
            for (String scoreKey : this.fetchs) {
                resultMap.put(scoreKey, Collections.nCopies(size, 0.0f));
            }
            return resultMap;
        }
    }

}
