package com.uas.search.grouping;

import com.uas.search.util.StringUtils;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.SimpleCollector;

import java.io.IOException;
import java.util.*;

/**
 * 批次统计
 *
 * @author sunyj
 * @since 2017年7月12日 下午2:12:56
 */
public class GoodsGroupCollector extends SimpleCollector {

    private Set<String> fieldsToLoad = new HashSet<>();

    private List<String> uniqueValues = new ArrayList<>();

    private String uniqueField;

    private LeafReader reader;

    private List<Map<String, Object>> values = new ArrayList<>();

    public GoodsGroupCollector(String uniqueField, Set<String> fieldsToLoad) {
        this.uniqueField = uniqueField;
        this.fieldsToLoad.addAll(fieldsToLoad);
    }

    @Override
    public boolean needsScores() {
        return false;
    }

    @Override
    protected void doSetNextReader(LeafReaderContext context) throws IOException {
        super.doSetNextReader(context);
        reader = context.reader();
    }

    @Override
    public void collect(int doc) throws IOException {
        if (null != reader) {
            Document document = reader.document(doc);
            String uniqueValue = document.get(uniqueField);
            // 排除重复值
            if (!StringUtils.isEmpty(uniqueValue)) {
                if (!uniqueValues.contains(uniqueValue)) {
                    Map<String, Object> map = new HashMap<>();
                    uniqueValues.add(uniqueValue);
                    for (String field : fieldsToLoad) {
                        String value = document.get(field);
                        map.put(field, value);
                    }
                    map.put("count", 1L);
                    values.add(map);
                } else {
                    // 统计数目
                    int index = uniqueValues.indexOf(uniqueValue);
                    Map<String, Object> map = values.get(index);
                    Long count = (Long) map.get("count");
                    map.put("count", count + 1);
                }
            }
        }
    }

    public List<Map<String, Object>> getValues() {
        return values;
    }

}
