export class MultiMap<K, V> {
    private _map = new Map<K, Set<V>>()

    add(key: K, value: V) {
        if (!this._map.has(key)) {
            this._map.set(key, new Set());
        }
        this._map.get(key).add(value);
    }

    add_all(key: K, values: Iterable<V>) {
        for (let v of values) {
            this.add(key, v);
        }
    }

    delete(key: K, value: V): V | undefined {
        const set = this._map.get(key);
        if (!set) return;
        set.delete(value);
        if (set.size == 0) this._map.delete(key);
    }

    delete_all(key: K): Set<V> | undefined {
        const set = this._map.get(key);
        this._map.delete(key);
        return set;
    }

    clear() {
        this._map.clear();
    }

    group(key: K): Set<V> {
        return this._map.get(key) || new Set();
    }

    has(key: K, value: V): boolean {
        return this._map.get(key)?.has(value);
    }

    hasKey(key: K): boolean {
        return this._map.has(key);
    }

    get keys() {
        return this._map.keys();
    }

    get keyCount(): number {
        return this._map.size;
    }

    get size(): number {
        let total_count = 0;
        for (let set of this._map.values()) {
            total_count += set.size;
        }
        return total_count;
    }

    *grouped(): IterableIterator<[K, Set<V>]> {
        for (let [key, set] of this._map) {
            yield [key, set];
        }
    }

    *items(): IterableIterator<[K, V]> {
        for (let [key, set] of this._map) {
            for (let v of set) {
                yield [key, v];
            }
        }
    }
}