package com.tzld.piaoquan.ad.engine.commons.score.model;


import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.*;



public class GBDTModel extends Model {
    private static final Logger LOGGER = LoggerFactory.getLogger(GBDTModel.class);
    private static final float LINEAR_TRANSFORM_LOWER_BOUND = 0f;
    private static final float LINEAR_TRANSFORM_SLOPE = 1f / 1f;
    private int featureCount = 0;
    private Map<String, Integer> featureIdMap = null;
    private List<HashMap<Integer, Node>> boosterModel = null;

    private float transform(final float score) {
        return (score - LINEAR_TRANSFORM_LOWER_BOUND) * LINEAR_TRANSFORM_SLOPE;
    }

    private int getFeatureId(final String feature) {
        return featureIdMap.containsKey(feature) ? featureIdMap.get(feature) : -1;
    }

    @Override
    public int getModelSize() {
        return boosterModel.size();
    }

    @Override
    public boolean loadFromStream(InputStreamReader in) throws Exception {
        List<HashMap<Integer, Node>> model = new ArrayList<HashMap<Integer, Node>>();
        HashMap<Integer, Node> tree = null;

        featureIdMap = new HashMap<String, Integer>();
        featureCount = 0;

        BufferedReader input = new BufferedReader(in);
        String line;
        while ((line = input.readLine()) != null) {
            String[] tokens = line.trim().split(":");
            if (tokens.length == 1) {
                if (tokens[0].startsWith("booster")) {
                    if (tree != null && tree.size() > 0) {
                        model.add(tree);
                    }

                    tree = new HashMap<Integer, Node>();
                }
            } else if (tokens.length == 2) {
                Node node = new Node();

                Integer id = Integer.parseInt(tokens[0]);

                String[] items = tokens[1].split(" ");
                if (items.length == 1) {
                    node.isLeaf = true;

                    String[] dt = items[0].split("=");
                    node.value = Float.parseFloat(dt[1]);
                } else if (items.length == 2) {
                    node.isLeaf = false;
                    String[] fieldDescriptions = null;
                    node.compareType = NodeCompareType.LT;
                    if (items[0].substring(1, items[0].length() - 1).contains("<=")) {
                        fieldDescriptions = items[0].substring(1, items[0].length() - 1).split("<=");
                        node.compareType = NodeCompareType.LE;
                    } else {
                        fieldDescriptions = items[0].substring(1, items[0].length() - 1).split("<");
                    }
                    // feature to id
                    int featureId = getFeatureId(fieldDescriptions[0]);
                    if (featureId < 0) {
                        featureIdMap.put(fieldDescriptions[0], featureCount);
                        featureId = featureCount;
                        featureCount += 1;
                    }
                    node.splitFeatureId = featureId;

                    if (fieldDescriptions.length == 1) {
                        node.fieldType = NodeFieldType.BINARY;
                    } else {
                        node.fieldType = NodeFieldType.QUANTITATIVE;
                        node.splitCondition = Float.parseFloat(fieldDescriptions[1]);
                    }

                    String[] childrenDescriptions = items[1].split(",");
                    for (String childDescription : childrenDescriptions) {
                        String[] descs = childDescription.split("=");
                        if (descs[0].equals("yes")) {
                            node.yes = Integer.parseInt(descs[1]);
                        } else if (descs[0].equals("no")) {
                            node.no = Integer.parseInt(descs[1]);
                        } else if (descs[0].equals("missing")) {
                            node.missing = Integer.parseInt(descs[1]);
                        }
                    }
                }
                tree.put(id, node);
            }
        }
        if (tree != null && !tree.isEmpty()) {
            model.add(tree);
        }

        LOGGER.info("Boosted tree model load over and tree number is " + model.size());
        input.close();
        in.close();
        if (model != null && model.size() > 0) {
            boosterModel = model;
            return true;
        } else {
            return false;
        }
    }

    private float score(final Map<Integer, Node> tree, final boolean[] featureIsExists, final Float[] featureValues) {
        int id = 0;
        while (true) {
            Node currentNode = tree.get(id);
            if (currentNode.isLeaf) {
                return currentNode.value;
            }

            if (currentNode.fieldType == NodeFieldType.BINARY) {
                if (!featureIsExists[currentNode.splitFeatureId]) {
                    id = currentNode.no;
                } else {
                    id = currentNode.yes;
                }
            } else if (currentNode.fieldType == NodeFieldType.QUANTITATIVE) {
                if (!featureIsExists[currentNode.splitFeatureId]) {
                    id = currentNode.missing;
                } else {
                    Float value = featureValues[currentNode.splitFeatureId];

                    if ((currentNode.compareType == NodeCompareType.LT && value < currentNode.splitCondition)
                            || (currentNode.compareType == NodeCompareType.LE && value <= currentNode.splitCondition)) {
                        id = currentNode.yes;
                    } else {
                        id = currentNode.no;
                    }
                }
            } else {
                LOGGER.error("Reach undefined condition: {}", id);
            }
        }
    }

    public Float score(final Map<String, Double> features, Map<String, Double> featuresScore) {
        Float result = 0f;

        boolean[] featureIsExists = new boolean[featureCount];
        Float[] featureValues = new Float[featureCount];
        for (Map.Entry<String, Double> entry : features.entrySet()) {
            int featureId = getFeatureId(entry.getKey());
            if (featureId >= 0) {
                featureIsExists[featureId] = true;
                featureValues[featureId] = entry.getValue().floatValue();
            }
        }

        for (Map<Integer, Node> tree : boosterModel) {
            result += score(tree, featureIsExists, featureValues);
        }

        Float transformedResult = transform(result);

        LOGGER.debug("[calc_dwelltime]features: " + Arrays.toString(features.entrySet().toArray()) + ", prediction:" + result + " , transformed:" + transformedResult);

        return transformedResult;
    }

    private enum NodeFieldType {
        BINARY, QUANTITATIVE
    }

    private enum NodeCompareType {
        LT, LE
    }

    private class Node {

        boolean isLeaf;
        float value;
        NodeFieldType fieldType;
        NodeCompareType compareType;
        int splitFeatureId;
        Float splitCondition;
        Integer yes, no, missing;
    }
}
