package at.yawk.numaec;

import org.eclipse.collections.api.ShortIterable;
import org.eclipse.collections.api.bag.MutableBag;
import org.eclipse.collections.api.bag.primitive.MutableDoubleBag;
import org.eclipse.collections.api.block.function.primitive.DoubleFunction;
import org.eclipse.collections.api.block.function.primitive.DoubleFunction0;
import org.eclipse.collections.api.block.function.primitive.DoubleToDoubleFunction;
import org.eclipse.collections.api.block.function.primitive.DoubleToObjectFunction;
import org.eclipse.collections.api.block.function.primitive.ShortDoubleToDoubleFunction;
import org.eclipse.collections.api.block.function.primitive.ShortToDoubleFunction;
import org.eclipse.collections.api.block.predicate.primitive.DoublePredicate;
import org.eclipse.collections.api.block.predicate.primitive.ShortDoublePredicate;
import org.eclipse.collections.api.iterator.MutableDoubleIterator;
import org.eclipse.collections.api.map.primitive.MutableDoubleShortMap;
import org.eclipse.collections.api.map.primitive.MutableShortDoubleMap;
import org.eclipse.collections.api.map.primitive.ShortDoubleMap;

class ShortDoubleLinearHashMap extends BaseShortDoubleMap implements ShortDoubleBufferMap {
    private final float loadFactor;
    private final long sipHashK0, sipHashK1;
    private final long hashMask;

    protected final LinearHashTable table;
    protected int size;

    ShortDoubleLinearHashMap(LargeByteBufferAllocator allocator, LinearHashMapConfig config) {
        this.sipHashK0 = config.sipHashK0.getAsLong();
        this.sipHashK1 = config.sipHashK1.getAsLong();
        this.loadFactor = config.loadFactor;
        int hashLength = config.hashLength;
        this.hashMask = hashLength == 0 ? -1L : ~(-1L >>> hashLength);
        this.table = new LinearHashTable(allocator, config, hashLength + Short.BYTES + Double.BYTES) {
            @Override
            protected void write(LargeByteBuffer lbb, long address, long hash, long key, long value) {
                if (hashLength != 0) {
                    if ((hash & ~hashMask) != 0) {
                        throw new AssertionError();
                    }
                    BTree.uset(lbb, address, hashLength, Long.reverse(hash));
                }
                lbb.setShort(address + hashLength, fromKey(key));
                lbb.setDouble(address + hashLength + Short.BYTES, fromValue(value));
            }

            @Override
            protected long readHash(LargeByteBuffer lbb, long address) {
                if (hashLength == 0) {
                    return hash(fromKey(readKey(lbb, address)));
                } else {
                    return Long.reverse(BTree.uget(lbb, address, hashLength));
                }
            }

            @Override
            protected long readKey(LargeByteBuffer lbb, long address) {
                return toKey(lbb.getShort(address + hashLength));
            }

            @Override
            protected long readValue(LargeByteBuffer lbb, long address) {
                return toValue(lbb.getDouble(address + hashLength + Short.BYTES));
            }
        };
    }

    protected void ensureCapacity(int capacity) {
        table.expandToFullLoadCapacity((long) (capacity / loadFactor));
    }

    @Override
    protected MapStoreCursor iterationCursor() {
        return table.allocateCursor();
    }

    @Override
    protected MapStoreCursor keyCursor(short key) {
        LinearHashTable.Cursor cursor = table.allocateCursor();
        cursor.seek(hash(key), toKey(key));
        return cursor;
    }

    @DoNotMutate
    @Override
    void checkInvariants() {
        super.checkInvariants();
        table.checkInvariants();
    }

    protected long hash(short key) {
        return SipHash.sipHash2_4_8_to_8(sipHashK0, sipHashK1, toKey(key)) & hashMask;
    }

    @Override
    public void close() {
        table.close();
    }

    @Override
    public int size() {
        return size;
    }

    public static class Mutable extends ShortDoubleLinearHashMap implements MutableShortDoubleBufferMap {

        Mutable(LargeByteBufferAllocator allocator, LinearHashMapConfig config) {
            super(allocator, config);
        }

        @Override
        public void put(short key, double value) {
            ensureCapacity(1);
            long h = hash(key);
            long k = toKey(key);
            long v = toValue(value);
            try (LinearHashTable.Cursor cursor = table.allocateCursor()) {
                cursor.seek(h, k);
                if (cursor.elementFound()) {
                    cursor.setValue(v);
                } else {
                    cursor.insert(h, k, v);
                    size++;
                    ensureCapacity(size);
                }
            }
        }

        @Override
        public void putAll(ShortDoubleMap map) {
            // this is too pessimistic when the given map's keys overlap with outs but probably covers the main use
            // cases just fine
            ensureCapacity(size + map.size());
            map.forEachKeyValue(this::put);
        }

        @Override
        public void updateValues(ShortDoubleToDoubleFunction function) {
            try (LinearHashTable.Cursor cursor = table.allocateCursor()) {
                while (cursor.next()) {
                    double updated = function.valueOf(fromKey(cursor.getKey()), fromValue(cursor.getValue()));
                    cursor.setValue(toValue(updated));
                }
            }
        }

        @Override
        public void removeKey(short key) {
            try (LinearHashTable.Cursor cursor = table.allocateCursor()) {
                cursor.seek(hash(key), toKey(key));
                if (cursor.elementFound()) {
                    cursor.remove();
                    size--;
                }
            }
        }

        @Override
        public void remove(short key) {
            removeKey(key);
        }

        @Override
        public double removeKeyIfAbsent(short key, double value) {
            try (LinearHashTable.Cursor cursor = table.allocateCursor()) {
                cursor.seek(hash(key), toKey(key));
                if (cursor.elementFound()) {
                    double v = fromValue(cursor.getValue());
                    cursor.remove();
                    size--;
                    return v;
                } else {
                    return value;
                }
            }
        }

        @Override
        public double getIfAbsentPut(short key, double value) {
            // only resize map if we really need to down below
            ensureCapacity(1);
            try (LinearHashTable.Cursor cursor = table.allocateCursor()) {
                long hash = hash(key);
                cursor.seek(hash, toKey(key));
                if (cursor.elementFound()) {
                    return fromValue(cursor.getValue());
                } else {
                    cursor.insert(hash, toKey(key), toValue(value));
                    size++;
                    ensureCapacity(size);
                    return value;
                }
            }
        }

        @Override
        public double getIfAbsentPut(short key, DoubleFunction0 function) {
            // only resize map if we really need to down below
            ensureCapacity(1);
            try (LinearHashTable.Cursor cursor = table.allocateCursor()) {
                long hash = hash(key);
                cursor.seek(hash, toKey(key));
                if (cursor.elementFound()) {
                    return fromValue(cursor.getValue());
                } else {
                    double v = function.value();
                    cursor.insert(hash, toKey(key), toValue(v));
                    size++;
                    ensureCapacity(size);
                    return v;
                }
            }
        }

        @Override
        public double getIfAbsentPutWithKey(short key, ShortToDoubleFunction function) {
            // only resize map if we really need to down below
            ensureCapacity(1);
            try (LinearHashTable.Cursor cursor = table.allocateCursor()) {
                long hash = hash(key);
                cursor.seek(hash, toKey(key));
                if (cursor.elementFound()) {
                    return fromValue(cursor.getValue());
                } else {
                    double v = function.valueOf(key);
                    cursor.insert(hash, toKey(key), toValue(v));
                    size++;
                    ensureCapacity(size);
                    return v;
                }
            }
        }

        @Override
        public <P> double getIfAbsentPutWith(short key, DoubleFunction<? super P> function, P parameter) {
            // only resize map if we really need to down below
            ensureCapacity(1);
            try (LinearHashTable.Cursor cursor = table.allocateCursor()) {
                long hash = hash(key);
                cursor.seek(hash, toKey(key));
                if (cursor.elementFound()) {
                    return fromValue(cursor.getValue());
                } else {
                    double v = function.doubleValueOf(parameter);
                    cursor.insert(hash, toKey(key), toValue(v));
                    size++;
                    ensureCapacity(size);
                    return v;
                }
            }
        }

        @Override
        public double updateValue(short key, double initialValueIfAbsent, DoubleToDoubleFunction function) {
            // only resize map if we really need to down below
            ensureCapacity(1);
            try (LinearHashTable.Cursor cursor = table.allocateCursor()) {
                long hash = hash(key);
                cursor.seek(hash, toKey(key));
                if (cursor.elementFound()) {
                    double updated = function.valueOf(fromValue(cursor.getValue()));
                    cursor.setValue(toValue(updated));
                    return updated;
                } else {
                    double updated = function.valueOf(initialValueIfAbsent);
                    cursor.insert(hash, toKey(key), toValue(updated));
                    size++;
                    ensureCapacity(size);
                    return updated;
                }
            }
        }

        @Override
        public MutableShortDoubleMap withKeyValue(short key, double value) {
            put(key, value);
            return this;
        }

        @Override
        public MutableShortDoubleMap withoutKey(short key) {
            removeKey(key);
            return this;
        }

        @Override
        public MutableShortDoubleMap withoutAllKeys(ShortIterable keys) {
            keys.forEach(this::removeKey);
            return this;
        }

        @Override
        public MutableShortDoubleMap asUnmodifiable() {
            throw new UnsupportedOperationException("Mutable.asUnmodifiable not implemented yet");
        }

        @Override
        public MutableShortDoubleMap asSynchronized() {
            throw new UnsupportedOperationException("Mutable.asSynchronized not implemented yet");
        }

        @Override
        public double addToValue(short key, double toBeAdded) {
            // only resize map if we really need to down below
            ensureCapacity(1);
            try (LinearHashTable.Cursor cursor = table.allocateCursor()) {
                long hash = hash(key);
                cursor.seek(hash, toKey(key));
                if (cursor.elementFound()) {
                    double updated = (double) (fromValue(cursor.getValue()) + toBeAdded);
                    cursor.setValue(toValue(updated));
                    return updated;
                } else {
                    cursor.insert(hash, toKey(key), toValue(toBeAdded));
                    size++;
                    ensureCapacity(size);
                    return toBeAdded;
                }
            }
        }

        @Override
        public void clear() {
            table.clear();
            size = 0;
        }

        @Override
        public MutableDoubleShortMap flipUniqueValues() {
            throw new UnsupportedOperationException("ShortDoubleBufferMap.Mutable.flipUniqueValues not implemented yet");
        }

        @Override
        public MutableShortDoubleMap select(ShortDoublePredicate predicate) {
            throw new UnsupportedOperationException("ShortDoubleBufferMap.Mutable.select not implemented yet");
        }

        @Override
        public MutableDoubleBag select(DoublePredicate predicate) {
            throw new UnsupportedOperationException("ShortDoubleBufferMap.Mutable.select not implemented yet");
        }

        @Override
        public MutableShortDoubleMap reject(ShortDoublePredicate predicate) {
            throw new UnsupportedOperationException("ShortDoubleBufferMap.Mutable.reject not implemented yet");
        }

        @Override
        public MutableDoubleBag reject(DoublePredicate predicate) {
            throw new UnsupportedOperationException("ShortDoubleBufferMap.Mutable.reject not implemented yet");
        }

        @Override
        public <V> MutableBag<V> collect(DoubleToObjectFunction<? extends V> function) {
            throw new UnsupportedOperationException("ShortDoubleBufferMap.Mutable.collect not implemented yet");
        }

        @Override
        public MutableDoubleIterator doubleIterator() {
            return super.doubleIterator();
        }
    }
}