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


import org.apache.commons.collections4.MapUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


public class FMModel extends Model {
    protected static final int MODEL_FIRST_LOAD_COUNT = 1 << 25; // 32M
    private static final Logger LOGGER = LoggerFactory.getLogger(FMModel.class);
    private Map<String, List<Float>> model;

    public void putFeature(Map<String, List<Float>> model, String[] items) {

        String featureKey = items[0];
        List<Float> weights = new ArrayList<>();
        for (int i = 1; i < items.length; i++) {
            weights.add(Float.valueOf(items[i]));
        }

        model.put(featureKey, weights);
    }

    public float getWeight(Map<String, List<Float>> model, String featureKey, int index) {
        if (!model.containsKey(featureKey)) {
            return 0.0f;
        }

        return model.get(featureKey).get(index);

    }

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

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

    public Float score(Map<String, Float> featureMap) {
        float sum = 0.0f;

        if (MapUtils.isNotEmpty(featureMap)) {
            // 计算 sum w*x
            float sum0 = 0.0f;
            Map<String, Float> validFeatureMap = new HashMap<>(featureMap.size());
            for (Map.Entry<String, Float> e : featureMap.entrySet()) {
                if (!this.model.containsKey(e.getKey())) {
                    continue;
                }
                float x = e.getValue();
                float w = getWeight(this.model, e.getKey(), 0);
                sum0 += w * x;
                validFeatureMap.put(e.getKey(), x);
            }
            sum += sum0;

            // 计算 sum v*v*x*X
            float sum1 = 0.0f;
            for (int i = 1; i < 9; i++) {
                float sum10 = 0.0f;
                float sum11 = 0.0f;
                for (Map.Entry<String, Float> e : validFeatureMap.entrySet()) {
                    float x = e.getValue();
                    float v = getWeight(this.model, e.getKey(), i);
                    float d = v * x;
                    sum10 += d;
                    sum11 += d * d;
                }
                sum1 += sum10 * sum10 - sum11;
            }
            sum1 = 0.5f * sum1;
            float biasW = model.get("bias").get(0);
            sum = biasW + sum0 + sum1;
        }

        return (float) (1.0f / (1 + Math.exp(-sum)));
    }

    /**
     * 目前模型比较大,分两个阶段load模型
     * (1). load 8M 模型, 并更新;
     * (2). load 剩余的模型
     * 中间提供一段时间有损的打分服务
     *
     * @param in
     * @return
     * @throws IOException
     */
    @Override
    public boolean loadFromStream(InputStreamReader in) throws IOException {

        Map<String, List<Float>> model = new HashMap<>();
        BufferedReader input = new BufferedReader(in);
        String line = null;
        int cnt = 0;

        long curTime = System.currentTimeMillis() / 1000;
        LOGGER.info("[MODELLOAD] before model load, key size: {}, current time: {}", model.size(), curTime);
        //first stage
        while ((line = input.readLine()) != null) {
            String[] items = line.split("\t");
            if (items.length < 9) {
                if (items[0].equals("bias")) {
                    putFeature(model, items);
                }
                continue;
            }

            putFeature(model, items);
            if (cnt > MODEL_FIRST_LOAD_COUNT) {
                break;
            }
        }
        //model update
        this.model = model;

        LOGGER.info("[MODELLOAD] after first stage model load, key size: {}, current time: {}", model.size(), curTime);
        //final stage
        while ((line = input.readLine()) != null) {
            String[] items = line.split("\t");
            if (items.length < 9) {
                continue;
            }
            putFeature(model, items);
        }
        LOGGER.info("[MODELLOAD] after model load, key size: {}, current time: {}", model.size(), curTime);

        LOGGER.info("[MODELLOAD] model load over and size " + cnt);
        input.close();
        in.close();
        return true;
    }

}
