// FlexMap, FlexSet: like built-in Map and Set, but with a
// canonicalization function which gives us the possibility of a
// coarser equivalence than the identity equivalence used in Map and
// Set.

// A Canonicalizer represents the equivalence you have in mind. For
//
//    c: Canonicalizer<V>
//    eqv: Equivalence<V>
//    v1: V
//    v2: V
//
// where `eqv` is the equivalence you want,
//
//     eqv(v1, v2) ⇔ c(v1) === c(v2)
//
export type Canonicalizer<V> = (v: V) => string;
export type Equivalence<V> = (v1: V, v2: V) => boolean;

export type IdentityMap<K, V> = Map<K, V>;
export type IdentitySet<V> = Set<V>;
export const IdentityMap = Map;
export const IdentitySet = Set;

export const IsMap = Symbol.for('IsMap');
export const IsSet = Symbol.for('IsSet');

declare global {
    interface Map<K, V> { [IsMap]: boolean; }
    interface MapConstructor { isMap<K, V>(x: any): x is Map<K, V>; }
    interface Set<T> { [IsSet]: boolean; }
    interface SetConstructor { isSet<T>(x: any): x is Set<T>; }
}
if (!(IsMap in Map.prototype)) {
    Object.defineProperty(Map.prototype, IsMap, { get() { return true; } });
    Map.isMap = <K,V> (x: any): x is Map<K, V> => !!x?.[IsMap];
}
if (!(IsSet in Set.prototype)) {
    Object.defineProperty(Set.prototype, IsSet, { get() { return true; } });
    Set.isSet = <T> (x: any): x is Set<T> => !!x?.[IsSet];
}

export function _iterMap<S,T>(i: Iterator<S>, f : (s: S) => T): IterableIterator<T> {
    const _f = (r: IteratorResult<S>): IteratorResult<T> => {
        if (r.done) {
            return { done: true, value: null };
        } else {
            return { done: false, value: f(r.value) };
        }
    };
    return {
        next: (v?: any): IteratorResult<T> => _f(i.next(v)),
        return: (v?: any): IteratorResult<T> => _f(i.return?.(v) ?? { done: true, value: null }),
        throw: (e?: any): IteratorResult<T> => _f(i.throw?.(e) ?? { done: true, value: null }),
        [Symbol.iterator]() { return this; },
    };
}

export class FlexMap<K, V> implements Map<K, V> {
    readonly items: Map<string, [K, V]>;
    readonly canonicalizer: Canonicalizer<K>;

    constructor(c: Canonicalizer<K>, items?: Iterable<readonly [K, V]>) {
        this.canonicalizer = c;
        this.items = (items === void 0)
            ? new Map()
            : new Map(_iterMap(items[Symbol.iterator](), ([k, v]) => [this._key(k), [k, v]]));
    }

    _key(k: K): string {
        return this.canonicalizer(k);
    }

    get(k: K, defaultValue?: V): V | undefined {
        const e = this.items.get(this._key(k));
        return (e === void 0) ? defaultValue : e[1];
    }

    getOrSet(k: K, initializer: () => V): V {
        const ks = this._key(k);
        let e = this.items.get(ks);
        if (e === void 0) {
            e = [k, initializer()];
            this.items.set(ks, e);
        }
        return e[1];
    }

    set(k: K, v: V): this {
        this.items.set(this._key(k), [k, v]);
        return this;
    }

    forEach(f: <T extends Map<K, V>> (v: V, k: K, map: T) => void, thisArg?: any) {
        this.items.forEach(([k, v]) => f.call(thisArg, v, k, this));
    }

    entries(): IterableIterator<[K, V]> {
        return this.items.values();
    }

    keys(): IterableIterator<K> {
        return _iterMap(this.items.values(), ([k, _v]) => k);
    }

    values(): IterableIterator<V> {
        return _iterMap(this.items.values(), ([_k, v]) => v);
    }

    delete(k: K): boolean {
        return this.items.delete(this._key(k));
    }

    getAndDelete(k: K, defaultValue?: V): V | undefined {
        const ks = this._key(k);
        const e = this.items.get(ks);
        if (e === void 0) return defaultValue;
        this.items.delete(ks);
        return e[1];
    }

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

    has(k: K): boolean {
        return this.items.has(this._key(k));
    }

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

    [Symbol.iterator](): IterableIterator<[K, V]> {
        return this.items.values();
    }

    get [Symbol.toStringTag]() { return 'FlexMap'; }

    equals(other: any, eqv: Equivalence<V> = (v1, v2) => v1 === v2): boolean {
        if (!('size' in other && 'has' in other && 'get' in other)) return false;
        if (this.size !== other.size) return false;
        for (let [k, v] of this.items.values()) {
            if (!other.has(k)) return false;
            if (!eqv(v, other.get(k))) return false;
        }
        return true;
    }

    update(key: K,
           f: (oldValue?: V) => V | undefined,
           defaultValue?: V,
           eqv: Equivalence<V> = (v1, v2) => v1 === v2): number
    {
        const ks = this._key(key);
        if (this.items.has(ks)) {
            const oldValue = this.items.get(ks)![1];
            const newValue = f(oldValue);
            if (newValue === void 0) {
                this.items.delete(ks);
                return -1;
            } else {
                if (!eqv(newValue, oldValue)) this.items.set(ks, [key, newValue]);
                return 0;
            }
        } else {
            const newValue = f(defaultValue);
            if (newValue === void 0) {
                return 0;
            } else {
                this.items.set(ks, [key, newValue]);
                return 1;
            }
        }
    }

    canonicalKeys(): IterableIterator<string> {
        return this.items.keys();
    }

    get [IsMap](): boolean {
        return true;
    }
}

export class FlexSet<V> implements Set<V> {
    readonly items: Map<string, V>;
    readonly canonicalizer: Canonicalizer<V>;

    constructor(c: Canonicalizer<V>, items?: Iterable<V>) {
        this.canonicalizer = c;
        this.items = (items === void 0)
            ? new Map()
            : new Map(_iterMap(items[Symbol.iterator](), (v) => [this._key(v), v]));
    }

    _key(v: V): string {
        return this.canonicalizer(v);
    }

    has(v: V): boolean {
        return this.items.has(this._key(v));
    }

    get(v: V): {item: V} | null {
        const vs = this._key(v);
        if (this.items.has(vs)) {
            return { item: this.items.get(vs)! };
        } else {
            return null;
        }
    }

    add(v: V): this {
        this.items.set(this._key(v), v);
        return this;
    }

    forEach(f: <T extends Set<V>>(v: V, v2: V, set: T) => void, thisArg?: any) {
        this.items.forEach((v) => f.call(thisArg, v, v, this));
    }

    entries(): IterableIterator<[V, V]> {
        return _iterMap(this.items.values(), (v) => [v, v]);
    }

    keys(): IterableIterator<V> {
        return this.items.values();
    }

    values(): IterableIterator<V> {
        return this.items.values();
    }

    delete(v: V): boolean {
        return this.items.delete(this._key(v));
    }

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

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

    [Symbol.iterator](): IterableIterator<V> {
        return this.items.values();
    }

    get [Symbol.toStringTag]() { return 'FlexSet'; }

    equals(other: any): boolean {
        if (!('size' in other && 'has' in other)) return false;
        if (this.size !== other.size) return false;
        for (let v of this.items.values()) {
            if (!other.has(v)) return false;
        }
        return true;
    }

    canonicalValues(): IterableIterator<string> {
        return this.items.keys();
    }

    union(other: Set<V>): FlexSet<V> {
        const result = new FlexSet(this.canonicalizer, this);
        for (let k of other) result.add(k);
        return result;
    }

    intersect(other: Set<V>): FlexSet<V> {
        const result = new FlexSet(this.canonicalizer);
        for (let k of this) if (other.has(k)) result.add(k);
        return result;
    }

    subtract(other: Set<V>): FlexSet<V> {
        const result = new FlexSet(this.canonicalizer);
        for (let k of this) if (!other.has(k)) result.add(k);
        return result;
    }

    get [IsSet](): boolean {
        return true;
    }
}