package test.random;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Random;
import java.util.TreeMap;
import java.util.function.Function;
import java.util.function.Supplier;

/**
 * Implements a random weighted chooser for the given result class. The weights are doubles,
 * but the class can be used supplying purely integer weights. The class supports the Java 8 functional
 * interfaces Supplier<E> and Function<Random, E> in addition to its own newChoice() methods.
 * <p>
 * The class supports giving <i>null</i> a weight. It's thread-safe.
 * <p>
 * The range of usable weights depends on the quality of the random number generator used.
 * Given a sufficiently well-coded Random subclass, the full precision range of the <i>double</i>
 * class (53 bits, or almost 16 orders of magnitude) can be utilized.
 * <p>
 * Negative weights are possible, and result in the specific value never being chosen.
 * The one exception to it is <i>null</i> if there is no other valid choice.
 * <p>
 * Nothing in the class prevents from using <i>infinity</i> or <i>NaN</i> values as weights, even though
 * it makes no sense to do so. Go right ahead.
 */
public class WeightedChoice<E> implements Supplier<E>, Function<Random, E> {
    private final transient Object[] LOCK = new Object[0];

    private Map<E, Double> weights;
    private transient TreeMap<Double, E> choices;
    private boolean valid = false;
    private double nullWeight = 0.0;
    private double maxWeight = 0.0;
    private Random rnd;

    public WeightedChoice(Random random) {
        weights = new HashMap<E, Double>();
        rnd = random;
    }

    public WeightedChoice(long seed) {
        this(new Random(seed));
    }

    public WeightedChoice() {
        this(new Random());
    }

    /**
     * Clears all current weights.
     */
    public void clear() {
        synchronized(LOCK) {
            weights = new HashMap<E, Double>();
            nullWeight = 0.0;
            valid = false;
        }
    }

    /**
     * Add a (potentially negative) weight to the given choice.
     * <p>
     * If the choice doesn't exist yet, it's created with the supplied weight.
     * If it exists, the weight is modified by the value.
     * A resulting negative weight means that the choice is never used.
     * <p>
     * Adding a weight to <i>null</i> is a valid operation - it means that there is a
     * non-zero chance for the randomizer to not return any value.
     */
    public void add(E choice, double weight) {
        if( weight == 0.0 ) {
            // No need to do anything
            return;
        }
        synchronized(LOCK) {
            internalAdd(choice, weight);
        }
    }

    /**
     * Set a (potentially negative) weight to the given choice.
     * Any existing weight gets overwritten.
     * <p>
     * A resulting negative weight means that the choice is never used.
     * <p>
     * Setting a weight for <i>null</i> is a valid operation - it means that there is a
     * non-zero chance for the randomizer to not return any value.
     */
    public void set(E choice, double weight) {
        synchronized(LOCK) {
            internalSet(choice, weight);
        }
    }

    /**
     * Remove a choice. Same as setting its weight to 0, but also cleans up the entry table.
     * <p>
     * Calling this method with <i>null</i> removes the choice to obtain a null result.
     */
    public void remove(E choice) {
        synchronized(LOCK) {
            internalRemove(choice);
        }
    }

    /**
     * Add the given entry at the given choice. This is a convenience method, meant for quickly
     * adding parts of a map.
     */
    public void add(Entry<E, Double> entry) {
        synchronized(LOCK) {
            internalAdd(entry.getKey(), entry.getValue());
        }
    }

    /**
     * Add the contents of the map to the current weight map.
     */
    public void add(Map<E, Double> map) {
        synchronized(LOCK) {
            for(Entry<E, Double> entry : map.entrySet()) {
                internalAdd(entry.getKey(), entry.getValue());
            }
        }
    }

    /**
     * Sets the given entry at the given choice. This is a convenience method, meant for quickly
     * setting parts of a map.
     */
    public void set(Entry<E, Double> entry) {
        synchronized(LOCK) {
            internalSet(entry.getKey(), entry.getValue());
        }
    }

    /**
     * Set the contents of the map as values in the current weight map.
     */
    public void set(Map<E, Double> map) {
        synchronized(LOCK) {
            for(Entry<E, Double> entry : map.entrySet()) {
                internalSet(entry.getKey(), entry.getValue());
            }
        }
    }

    /**
     * Remove all weights in the given collection.
     */
    public void remove(Collection<E> collection) {
        synchronized(LOCK) {
            for(E choice : collection) {
                internalRemove(choice);
            }
        }
    }

    /**
     * Remove all weights returned by the given iterable class.
     */
    public void remove(Iterable<E> iterable) {
        synchronized(LOCK) {
            for(E choice : iterable) {
                internalRemove(choice);
            }
        }
    }

    private void internalAdd(E choice, double weight) {
        if( choice == null ) {
            nullWeight += weight;
        }
        else {
            double currentWeight = ( weights.containsKey(choice) ? weights.get(choice) : 0.0 );
            weights.put(choice, currentWeight + weight);
        }
        valid = false;
    }

    private void internalRemove(E choice) {
        if( choice == null ) {
            nullWeight = 0.0;
        }
        else {
            if( !weights.containsKey(choice) ) {
                // NO-OP
                return;
            }
            weights.remove(choice);
        }
        valid = false;
    }

    private void internalSet(E choice, double weight) {
        if( choice == null ) {
            nullWeight = weight;
        }
        else if( weights.containsKey(choice) && weights.get(choice) == weight ) {
            // No need to do anything
            return;
        }
        else {
            weights.put(choice, weight);
        }
        valid = false;
    }

    private void validate() {
        if( !valid ) {
            // Generate the distribution TreeMap
            choices = new TreeMap<Double, E>();
            double count = 0.0;
            for(Entry<E, Double> entry : weights.entrySet()) {
                if( entry.getValue() > 0.0 ) {
                    choices.put(count, entry.getKey());
                    count += entry.getValue();
                }
            }
            maxWeight = count;
            valid = true;
        }
    }

    /**
     * Return a random weighted choice, using the built-in or supplied during creation randomizer.
     */
    @Override public E get() {
        return apply(rnd);
    }

    /**
     * Return a random weighted choice, using the built-in or supplied during creation randomizer.
     */
    public E newChoice() {
        return apply(rnd);
    }

    /**
     * Return a random weighted choice, using the supplied randomizer.
     */
    public E newChoice(Random random) {
        return apply(random);
    }

    /**
     * Return a random weighted choice, using the supplied randomizer.
     */
    @Override public E apply(Random random) {
        synchronized(LOCK) {
            validate();

            double randomVal = (maxWeight + (nullWeight > 0.0 ? nullWeight : 0.0)) * random.nextDouble();
            if( nullWeight > 0.0 ) {
                if( randomVal <= nullWeight ) {
                    // We have chosen to not choose.
                    return null;
                }
                randomVal -= nullWeight;
            }
            return choices.floorEntry(randomVal).getValue();
        }
    }

    /**
     * Returns a read-only copy of the current weights.
     */
    public Map<E, Double> weights() {
        synchronized(LOCK) {
            return Collections.<E, Double>unmodifiableMap(weights);
        }
    }

    public void setSeed(long seed) {
        rnd.setSeed(seed);
    }
}