package mb.nabl2.util.collections;

import java.util.stream.Collectors;

import org.metaborg.util.functions.Action2;

import io.usethesource.capsule.Map;

public abstract class MultiSet<E> {

    protected abstract Map<E, Integer> elements();

    public boolean isEmpty() {
        return elements().isEmpty();
    }

    public int size() {
        return elements().values().stream().mapToInt(i -> i).sum();
    }

    public int count(E e) {
        return elements().getOrDefault(e, 0);
    }

    public java.util.Set<E> elementSet() {
        return elements().keySet();
    }

    public void forEach(Action2<E, Integer> f) {
        elements().entrySet().forEach(e -> f.apply(e.getKey(), e.getValue()));
    }

    public static class Immutable<E> extends MultiSet<E> {

        private final Map.Immutable<E, Integer> elements;

        private Immutable(Map.Immutable<E, Integer> elements) {
            this.elements = elements;
        }

        @Override protected Map<E, Integer> elements() {
            return elements;
        }

        public Immutable<E> add(E e) {
            return add(e, 1);
        }

        public Immutable<E> add(E e, int n) {
            if(n < 0) {
                throw new IllegalArgumentException("count must be positive");
            }
            final int c = elements.getOrDefault(e, 0) + n;
            if(c > 0) {
                return new MultiSet.Immutable<>(elements.__put(e, c));
            } else {
                return new MultiSet.Immutable<>(elements.__remove(e));
            }
        }

        public Immutable<E> remove(E e) {
            return remove(e, 1);
        }

        public Immutable<E> remove(E e, int n) {
            if(n < 0) {
                throw new IllegalArgumentException("count must be positive");
            }
            final int c = Math.max(0, elements.getOrDefault(e, 0) - n);
            if(c > 0) {
                return new MultiSet.Immutable<>(elements.__put(e, c));
            } else {
                return new MultiSet.Immutable<>(elements.__remove(e));
            }

        }

        public Immutable<E> removeAll(E e) {
            return new MultiSet.Immutable<>(elements.__remove(e));
        }

        public MultiSet.Transient<E> melt() {
            return new MultiSet.Transient<>(elements.asTransient());
        }

        public static <E> MultiSet.Immutable<E> of() {
            return new MultiSet.Immutable<>(Map.Immutable.of());
        }

    }

    public static class Transient<E> extends MultiSet<E> {

        private Map.Transient<E, Integer> elements;

        private Transient(Map.Transient<E, Integer> elements) {
            this.elements = elements;
        }

        @Override protected Map<E, Integer> elements() {
            return elements;
        }

        /*
         * @return New count
         */
        public int add(E e) {
            return add(e, 1);
        }

        /*
         * @return New count
         */
        public int add(E e, int n) {
            if(n < 0) {
                throw new IllegalArgumentException("count must be positive");
            }
            final int c = elements.getOrDefault(e, 0) + n;
            if(c > 0) {
                elements.__put(e, c);
            } else {
                elements.__remove(e);
            }
            return c;
        }

        /*
         * @return New count
         */
        public int remove(E e) {
            return remove(e, 1);
        }

        /*
         * @return New count
         */
        public int remove(E e, int n) {
            if(n < 0) {
                throw new IllegalArgumentException("count must be positive");
            }
            final int c = Math.max(0, elements.getOrDefault(e, 0) - n);
            if(c > 0) {
                elements.__put(e, c);
            } else {
                elements.__remove(e);
            }
            return c;
        }

        /*
         * @return Old count
         */
        public int removeAll(E e) {
            final int c = elements.getOrDefault(e, 0);
            elements.__remove(e);
            return c;
        }

        public MultiSet.Immutable<E> freeze() {
            return new MultiSet.Immutable<>(elements.freeze());
        }

        public static <E> MultiSet.Transient<E> of() {
            return new MultiSet.Transient<>(Map.Transient.of());
        }


    }

    @Override public String toString() {
        return elements().entrySet().stream().map(e -> e.getKey() + ": " + e.getValue())
                .collect(Collectors.joining(", ", "{", "}"));
    }

}
@@ -2803,32 +2803,139 @@
    %7D%0A%0A        /*
+*%0A         * Add an element once.%0A         * %0A         * @param e%0A         *            Element to be added
 %0A         * @ret
@@ -2939,32 +2939,49 @@
 return New count
+ for the element.
 %0A         */%0A   
@@ -3050,32 +3050,220 @@
    %7D%0A%0A        /*
+*%0A         * Add an element n times.%0A         * %0A         * @param e%0A         *            Element to be added%0A         * @param n%0A         *            Additions, greater or equal to zero
 %0A         * @ret
@@ -3267,32 +3267,48 @@
 return New count
+ for the element
 %0A         */%0A   
@@ -3686,32 +3686,144 @@
    %7D%0A%0A        /*
+*%0A         * Remove an element once.%0A         * %0A         * @param e%0A         *            Element to be removed
 %0A         * @ret
@@ -3827,32 +3827,48 @@
 return New count
+ for the element
 %0A         */%0A   
@@ -3943,32 +3943,230 @@
    %7D%0A%0A        /*
+*%0A         * Remove an element up to n times.%0A         * %0A         * @param e%0A         *            Element to be removed%0A         * @param n%0A         *            Removals, greater or equal to zero
 %0A         * @ret
@@ -4174,24 +4174,40 @@
 rn New count
+ for the element
 %0A         */
@@ -4613,16 +4613,70 @@
       /*
+*%0A         * Remove an element completely.%0A         * 
 %0A       
@@ -4696,16 +4696,32 @@
 ld count
+ for the element
 %0A       

Add method documentation.