package net.coderodde.stat;

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

public class StandardDeviationNumberSet {

    private double sum;
    private double squareSum;
    private int numberOfElements;
    private final Map<Double, Integer> countMap = new HashMap<>();

    public void add(Number number) {
        double numberValue = number.doubleValue();
        Integer count = countMap.get(numberValue);

        if (count != null) {
            countMap.put(numberValue, count + 1);
        } else {
            countMap.put(numberValue, 1);
        }

        sum += numberValue;
        squareSum += numberValue * numberValue;
        numberOfElements++;
    }

    public double remove(Number number) {
        double numberValue = number.doubleValue();
        Integer count = countMap.get(numberValue);

        if (count != null) {
            if (count > 1) {
                countMap.put(numberValue, count - 1);
            } else {
                countMap.remove(numberValue);
            }

            sum -= numberValue;
            squareSum -= numberValue * numberValue;
            numberOfElements--;
            return numberValue;
        } else {
            return Double.NaN;
        }
    }

    public double getStandardDeviation() {
        checkSetHasAtleastTwoElements();
        double step1 = squareSum - sum * sum / numberOfElements;
        double step2 = step1 / (numberOfElements - 1);
        return Math.sqrt(step2);
    }

    private void checkSetHasAtleastTwoElements() {
        if (numberOfElements < 2) {
            throw new IllegalStateException(
                    "Computing the standard deviation requires at least two " +
                    "elements.");
        }
    }

    public static void main(String[] args) {
        StandardDeviationNumberSet set = new StandardDeviationNumberSet();

        set.add(1);
        set.add(1);
        set.add(3);

        System.out.println(set.getStandardDeviation());

        set.remove(1);

        System.out.println(set.getStandardDeviation());
    }
}