package at.yawk.numaec;

import org.eclipse.collections.api.DoubleIterable;
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.DoubleDoubleToDoubleFunction;
import org.eclipse.collections.api.block.function.primitive.DoubleToDoubleFunction;
import org.eclipse.collections.api.block.predicate.primitive.DoublePredicate;
import org.eclipse.collections.api.block.predicate.primitive.DoubleDoublePredicate;
import org.eclipse.collections.api.iterator.MutableDoubleIterator;
import org.eclipse.collections.api.map.primitive.MutableDoubleDoubleMap;
import org.eclipse.collections.api.map.primitive.MutableDoubleDoubleMap;
import org.eclipse.collections.api.map.primitive.DoubleDoubleMap;

public class DoubleDoubleBTreeMap extends BaseDoubleDoubleMap implements DoubleDoubleBufferMap {
    protected final BTree bTree;
    protected int size = 0;

    DoubleDoubleBTreeMap(LargeByteBufferAllocator allocator, BTreeConfig config) {
        int leafSize = Double.BYTES + Double.BYTES;
        int branchSize = config.entryMustBeInLeaf ? Double.BYTES : leafSize;
        this.bTree = new BTree(allocator, config, branchSize, leafSize) {
            @Override
            protected void writeBranchEntry(LargeByteBuffer lbb, long address, long key, long value) {
                lbb.setDouble(address, fromKey(key));
                if (!config.entryMustBeInLeaf) {
                    lbb.setDouble(address + Double.BYTES, fromValue(value));
                }
            }

            @Override
            protected void writeLeafEntry(LargeByteBuffer lbb, long address, long key, long value) {
                lbb.setDouble(address, fromKey(key));
                lbb.setDouble(address + Double.BYTES, fromValue(value));
            }

            @Override
            protected long readBranchKey(LargeByteBuffer lbb, long address) {
                return toKey(lbb.getDouble(address));
            }

            @Override
            protected long readBranchValue(LargeByteBuffer lbb, long address) {
                if (config.entryMustBeInLeaf) {
                    throw new AssertionError();
                } else {
                    return toValue(lbb.getDouble(address + Double.BYTES));
                }
            }

            @Override
            protected long readLeafKey(LargeByteBuffer lbb, long address) {
                return toKey(lbb.getDouble(address));
            }

            @Override
            protected long readLeafValue(LargeByteBuffer lbb, long address) {
                return toValue(lbb.getDouble(address + Double.BYTES));
            }
        };
    }

    @Override
    protected MapStoreCursor iterationCursor() {
        BTree.Cursor cursor = bTree.allocateCursor();
        cursor.descendToStart();
        return cursor;
    }

    @Override
    protected MapStoreCursor keyCursor(double key) {
        BTree.Cursor cursor = bTree.allocateCursor();
        cursor.descendToKey(toKey(key));
        return cursor;
    }

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

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

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

    public static class Mutable extends DoubleDoubleBTreeMap implements MutableDoubleDoubleBufferMap {
        Mutable(LargeByteBufferAllocator allocator, BTreeConfig config) {
            super(allocator, config);
        }

        @Override
        public void put(double key, double value) {
            long k = toKey(key);
            long v = toValue(value);
            try (BTree.Cursor cursor = bTree.allocateCursor()) {
                cursor.descendToKey(k);
                if (cursor.elementFound()) {
                    cursor.setValue(v);
                } else {
                    cursor.simpleInsert(k, v);
                    cursor.balance();
                    size++;
                }
            }
        }

        @Override
        public void putAll(DoubleDoubleMap map) {
            map.forEachKeyValue(this::put);
        }

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

        @Override
        public void removeKey(double key) {
            long k = toKey(key);
            try (BTree.Cursor cursor = bTree.allocateCursor()) {
                cursor.descendToKey(k);
                if (cursor.elementFound()) {
                    cursor.simpleRemove();
                    cursor.balance();
                    size--;
                }
            }
        }

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

        @Override
        public double removeKeyIfAbsent(double key, double value) {
            long k = toKey(key);
            try (BTree.Cursor cursor = bTree.allocateCursor()) {
                cursor.descendToKey(k);
                if (cursor.elementFound()) {
                    double v = fromValue(cursor.getValue());
                    cursor.simpleRemove();
                    cursor.balance();
                    size--;
                    return v;
                } else {
                    return value;
                }
            }
        }

        @Override
        public double getIfAbsentPut(double key, double value) {
            long k = toKey(key);
            try (BTree.Cursor cursor = bTree.allocateCursor()) {
                cursor.descendToKey(k);
                if (cursor.elementFound()) {
                    return fromValue(cursor.getValue());
                } else {
                    cursor.simpleInsert(k, toValue(value));
                    cursor.balance();
                    size++;
                    return value;
                }
            }
        }

        @Override
        public double getIfAbsentPut(double key, DoubleFunction0 function) {
            long k = toKey(key);
            try (BTree.Cursor cursor = bTree.allocateCursor()) {
                cursor.descendToKey(k);
                if (cursor.elementFound()) {
                    return fromValue(cursor.getValue());
                } else {
                    double v = function.value();
                    cursor.simpleInsert(k, toValue(v));
                    cursor.balance();
                    size++;
                    return v;
                }
            }
        }

        @Override
        public double getIfAbsentPutWithKey(double key, DoubleToDoubleFunction function) {
            long k = toKey(key);
            try (BTree.Cursor cursor = bTree.allocateCursor()) {
                cursor.descendToKey(k);
                if (cursor.elementFound()) {
                    return fromValue(cursor.getValue());
                } else {
                    double v = function.valueOf(key);
                    cursor.simpleInsert(k, toValue(v));
                    cursor.balance();
                    size++;
                    return v;
                }
            }
        }

        @Override
        public <P> double getIfAbsentPutWith(double key, DoubleFunction<? super P> function, P parameter) {
            long k = toKey(key);
            try (BTree.Cursor cursor = bTree.allocateCursor()) {
                cursor.descendToKey(k);
                if (cursor.elementFound()) {
                    return fromValue(cursor.getValue());
                } else {
                    double v = function.doubleValueOf(parameter);
                    cursor.simpleInsert(k, toValue(v));
                    cursor.balance();
                    size++;
                    return v;
                }
            }
        }

        @Override
        public double updateValue(double key, double initialValueIfAbsent, DoubleToDoubleFunction function) {
            long k = toKey(key);
            try (BTree.Cursor cursor = bTree.allocateCursor()) {
                cursor.descendToKey(k);
                if (cursor.elementFound()) {
                    double updated = function.valueOf(fromValue(cursor.getValue()));
                    cursor.setValue(toValue(updated));
                    return updated;
                } else {
                    double updated = function.valueOf(initialValueIfAbsent);
                    cursor.simpleInsert(k, toValue(updated));
                    cursor.balance();
                    size++;
                    return updated;
                }
            }
        }

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

        @Override
        public MutableDoubleDoubleMap withoutKey(double key) {
            removeKey(key);
            return this;
        }

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

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

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

        @Override
        public double addToValue(double key, double toBeAdded) {
            long k = toKey(key);
            try (BTree.Cursor cursor = bTree.allocateCursor()) {
                cursor.descendToKey(k);
                if (cursor.elementFound()) {
                    double updated = (double) (fromValue(cursor.getValue()) + toBeAdded);
                    cursor.setValue(toValue(updated));
                    return updated;
                } else {
                    cursor.simpleInsert(k, toValue(toBeAdded));
                    cursor.balance();
                    size++;
                    return toBeAdded;
                }
            }
        }

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

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

        @Override
        public MutableDoubleDoubleMap select(DoubleDoublePredicate predicate) {
            throw new UnsupportedOperationException("DoubleDoubleBufferMap.Mutable.select not implemented yet");
        }

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

        @Override
        public MutableDoubleDoubleMap reject(DoubleDoublePredicate predicate) {
            throw new UnsupportedOperationException("DoubleDoubleBufferMap.Mutable.reject not implemented yet");
        }

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

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

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