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


import com.tzld.piaoquan.recommend.feature.model.sample.BaseFeature;
import com.tzld.piaoquan.recommend.feature.model.sample.GroupedFeature;
import com.tzld.piaoquan.recommend.feature.model.sample.LRSamples;
import it.unimi.dsi.fastutil.longs.Long2FloatMap;
import it.unimi.dsi.fastutil.longs.Long2FloatOpenHashMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

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


public class LRModel extends Model {
    protected static final int MODEL_FIRST_LOAD_COUNT = 1 << 25; // 32M
    private static final Logger LOGGER = LoggerFactory.getLogger(LRModel.class);
    private final int bucketBits = 10;  // power(2, 10) => 1024 个槽位
    private List<Long2FloatMap> lrModel;

    public LRModel() {
        //配置不同环境的hdfs conf
        this.lrModel = constructModel();
    }

    public List<Long2FloatMap> getLrModel() {
        return lrModel;
    }

    public List<Long2FloatMap> constructModel() {
        List<Long2FloatMap> initModel = new ArrayList<Long2FloatMap>();
        int buckets = (int) Math.pow(2, bucketBits);
        for (int i = 0; i < buckets; i++) {
            Long2FloatMap internalModel = new Long2FloatOpenHashMap();
            internalModel.defaultReturnValue(0.0f);
            initModel.add(internalModel);
        }
        return initModel;
    }

    public int getBucket(long featureHash) {
        return (int) (((featureHash >> bucketBits) << bucketBits) ^ featureHash);
    }

    public void putFeature(List<Long2FloatMap> model, long featureHash, float weight) {
        model.get(getBucket(featureHash)).put(featureHash, weight);
    }

    public float getWeight(List<Long2FloatMap> model, long featureHash) {
        return model.get(getBucket(featureHash)).get(featureHash);
    }

    @Override
    public int getModelSize() {
        if (this.lrModel == null)
            return 0;
        int sum = 0;
        for (Map<Long, Float> model : this.lrModel) {
            sum += model.size();
        }
        return sum;
    }

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

    public Float score(LRSamples lrSamples) {
        float sum = 0.0f;
        for (int i = 0; i < lrSamples.getFeaturesCount(); i++) {
            GroupedFeature groupedFeature = lrSamples.getFeatures(i);
            if (groupedFeature != null && groupedFeature.getFeaturesCount() != 0) {
                for (int j = 0; j < groupedFeature.getFeaturesCount(); j++) {
                    BaseFeature baseFeature = groupedFeature.getFeatures(j);
                    if (baseFeature != null) {
                        float weight = getWeight(this.lrModel, baseFeature.getIdentifier());
                        baseFeature.toBuilder().setWeight(weight);
                        sum += weight;
                    }
                }
            }
        }

        float pro = (float) (1.0f / (1 + Math.exp(-sum)));
        lrSamples.toBuilder().setPredictCtr(pro);
        return pro;
    }

    public Float getWeights(LRSamples lrSamples) {
        float sum = 0.0f;

        for (int i = 0; i < lrSamples.getFeaturesCount(); i++) {
            GroupedFeature gf = lrSamples.getFeatures(i);
            if (gf != null && gf.getFeatures(i) != null) {
                for (int j = 0; j < gf.getFeaturesCount(); j++) {
                    BaseFeature fea = gf.getFeatures(j);
                    if (fea != null) {
                        float tmp = getWeight(this.lrModel, fea.getIdentifier());
                        fea.toBuilder().setWeight(tmp);
                        sum += tmp;
                    }
                }
            }
        }
        lrSamples.toBuilder().setWeight(sum);
        return sum;
    }

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

        List<Long2FloatMap> model = constructModel();
        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: {}", lrModel.size(), curTime);
        //first stage
        while ((line = input.readLine()) != null) {
            String[] items = line.split("\t");
            if (items.length < 2) {
                continue;
            }

            putFeature(model, new BigInteger(items[0].trim()).longValue(), Float.valueOf(items[1].trim()).floatValue());
            if (cnt++ < 10) {
                LOGGER.debug("fea: " + items[0] + ", weight: " + items[1]);
            }
            if (cnt > MODEL_FIRST_LOAD_COUNT) {
                break;
            }
        }
        //model update
        this.lrModel = model;
        LOGGER.info("[MODELLOAD] after first stage model load, key size: {}, current time: {}", lrModel.size(), curTime);
        //final stage
        while ((line = input.readLine()) != null) {
            String[] items = line.split("\t");
            if (items.length < 2) {
                continue;
            }
            putFeature(model, new BigInteger(items[0]).longValue(), Float.valueOf(items[1]).floatValue());
        }
        LOGGER.info("[MODELLOAD] after model load, key size: {}, current time: {}", lrModel.size(), curTime);

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

}
