package com.thebrokenrail.energonrelics.util;

import java.util.HashMap;
import java.util.Map;
import java.util.Random;

/**
 * Weighted List
 * @param <T> Entry Type
 */
public class WeightedList<T> {
    private final Map<T, Integer> map = new HashMap<>();

    /**
     * Add Entry
     * @param weight Weight
     * @param obj Entry
     */
    public void add(int weight, T obj) {
        map.put(obj, weight);
    }

    /**
     * Pick Entry
     * @param random Random
     * @return Entry
     */
    public T pick(Random random) {
        int totalWeight = 0;
        for (Map.Entry<T, Integer> entry : map.entrySet()) {
            totalWeight += entry.getValue();
        }
        int i = random.nextInt(totalWeight + 1);
        int countWeight = 0;
        for (Map.Entry<T, Integer> entry : map.entrySet()) {
            countWeight += entry.getValue();
            if (countWeight >= i) {
                return entry.getKey();
            }
        }
        throw new UnsupportedOperationException();
    }
}