package com.usoftchina.bi.core.utils;

import com.google.common.collect.HashBasedTable;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.google.common.collect.Table;

import java.util.*;
import java.util.function.Function;

/**
 * @author ChenWei
 * @date 2019/04/22
 */
public abstract class CollectionUtils extends org.springframework.util.CollectionUtils {
    /**
     * 按指定方法，将list分组返回map
     *
     * @param sources
     * @param keyGetter
     * @param <K>
     * @param <V>
     * @return
     */
    public static <K, V> Map<K, List<V>> groupBy(List<V> sources, Function<V, K> keyGetter) {
        Map<K, List<V>> map = new HashMap<>(1);
        if (!isEmpty(sources)) {
            sources.forEach(source -> {
                K key = keyGetter.apply(source);
                if (map.containsKey(key)) {
                    map.get(key).add(source);
                } else {
                    List<V> childList = new ArrayList<>();
                    childList.add(source);
                    map.put(key, childList);
                }
            });
        }
        return map;
    }

    /**
     * 按两个指定方法，将list分组返回复合键map
     *
     * @param sources
     * @param rKeyGetter
     * @param cKeyGetter
     * @param <R>
     * @param <C>
     * @param <V>
     * @return
     */
    public static <R, C, V> Table<R, C, List<V>> groupBy(List<V> sources, Function<V, R> rKeyGetter, Function<V, C> cKeyGetter) {
        Table<R, C, List<V>> table = HashBasedTable.create();
        if (!isEmpty(sources)) {
            sources.forEach(source -> {
                R r = rKeyGetter.apply(source);
                C c = cKeyGetter.apply(source);
                if (table.contains(r, c)) {
                    table.get(r, c).add(source);
                } else {
                    List<V> childList = new ArrayList<>();
                    childList.add(source);
                    table.put(r, c, childList);
                }
            });
        }
        return table;
    }

    /**
     * 按三个指定方法，将list分组返回复合键map
     *
     * @param sources
     * @param rKeyGetter
     * @param cKeyGetter
     * @param valueGetter
     * @param <R>
     * @param <C>
     * @param <V>
     * @param <S>
     * @return
     */
    public static <R, C, V, S> Map<R, Map<C, List<V>>> groupBy(List<S> sources, Function<S, R> rKeyGetter, Function<S, C> cKeyGetter, Function<S, V> valueGetter) {
        Map<R, Map<C, List<V>>> table = new HashMap<>(1);
        if (!isEmpty(sources)) {
            sources.forEach(source -> {
                R r = rKeyGetter.apply(source);
                C c = cKeyGetter.apply(source);
                V v = valueGetter.apply(source);

                if (table.containsKey(r)) {
                    Map<C, List<V>> rows = table.get(r);
                    if (rows.containsKey(c)) {
                        rows.get(c).add(v);
                    } else {
                        rows.put(c, Lists.newArrayList(v));
                    }
                } else {
                    Map<C, List<V>> rows = new HashMap<>(1);
                    rows.put(c, Lists.newArrayList(v));
                    table.put(r, rows);
                }
            });
        }
        return table;
    }

    /**
     * 按三个指定方法，将list分组返回复合键map
     *
     * @param sources
     * @param rKeyGetter
     * @param cKeyGetter
     * @param valueGetter
     * @param <R>
     * @param <C>
     * @param <V>
     * @param <S>
     * @return
     */
    public static <R, C, V, S> Map<R, Map<C, Set<V>>> distinctBy(List<S> sources, Function<S, R> rKeyGetter, Function<S, C> cKeyGetter, Function<S, V> valueGetter) {
        Map<R, Map<C, Set<V>>> table = new HashMap<>(1);
        if (!isEmpty(sources)) {
            sources.forEach(source -> {
                R r = rKeyGetter.apply(source);
                C c = cKeyGetter.apply(source);
                V v = valueGetter.apply(source);

                if (table.containsKey(r)) {
                    Map<C, Set<V>> rows = table.get(r);
                    if (rows.containsKey(c)) {
                        rows.get(c).add(v);
                    } else {
                        rows.put(c, Sets.newHashSet(v));
                    }
                } else {
                    Map<C, Set<V>> rows = new HashMap<>(1);
                    rows.put(c, Sets.newHashSet(v));
                    table.put(r, rows);
                }
            });
        }
        return table;
    }
}
