/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.prestosql.geospatial;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.base.Predicate;
import com.google.common.collect.ComparisonChain;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;

import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.OptionalInt;

import static com.google.common.base.Preconditions.checkArgument;
import static io.prestosql.geospatial.KdbTree.Node.newInternal;
import static io.prestosql.geospatial.KdbTree.Node.newLeaf;
import static java.util.Objects.requireNonNull;

/**
 * 2-dimensional K-D-B Tree
 * see https://en.wikipedia.org/wiki/K-D-B-tree
 */
public class KdbTree
{
    private static final int MAX_LEVELS = 10_000;

    private final Node root;

    public static final class Node
    {
        private final Rectangle extent;
        private final OptionalInt leafId;
        private final Optional<Node> left;
        private final Optional<Node> right;

        public static Node newLeaf(Rectangle extent, int leafId)
        {
            return new Node(extent, OptionalInt.of(leafId), Optional.empty(), Optional.empty());
        }

        public static Node newInternal(Rectangle extent, Node left, Node right)
        {
            return new Node(extent, OptionalInt.empty(), Optional.of(left), Optional.of(right));
        }

        @JsonCreator
        public Node(
                @JsonProperty("extent") Rectangle extent,
                @JsonProperty("leafId") OptionalInt leafId,
                @JsonProperty("left") Optional<Node> left,
                @JsonProperty("right") Optional<Node> right)
        {
            this.extent = requireNonNull(extent, "extent is null");
            this.leafId = requireNonNull(leafId, "leafId is null");
            this.left = requireNonNull(left, "left is null");
            this.right = requireNonNull(right, "right is null");
            if (leafId.isPresent()) {
                checkArgument(leafId.getAsInt() >= 0, "leafId must be >= 0");
                checkArgument(!left.isPresent(), "Leaf node cannot have left child");
                checkArgument(!right.isPresent(), "Leaf node cannot have right child");
            }
            else {
                checkArgument(left.isPresent(), "Intermediate node must have left child");
                checkArgument(right.isPresent(), "Intermediate node must have right child");
            }
        }

        @JsonProperty
        public Rectangle getExtent()
        {
            return extent;
        }

        @JsonProperty
        public OptionalInt getLeafId()
        {
            return leafId;
        }

        @JsonProperty
        public Optional<Node> getLeft()
        {
            return left;
        }

        @JsonProperty
        public Optional<Node> getRight()
        {
            return right;
        }

        @Override
        public boolean equals(Object obj)
        {
            if (obj == null) {
                return false;
            }

            if (!(obj instanceof Node)) {
                return false;
            }

            Node other = (Node) obj;
            return this.extent.equals(other.extent)
                    && Objects.equals(this.leafId, other.leafId)
                    && Objects.equals(this.left, other.left)
                    && Objects.equals(this.right, other.right);
        }

        @Override
        public int hashCode()
        {
            return Objects.hash(extent, leafId, left, right);
        }
    }

    @JsonCreator
    public KdbTree(@JsonProperty("root") Node root)
    {
        this.root = requireNonNull(root, "root is null");
    }

    @JsonProperty
    public Node getRoot()
    {
        return root;
    }

    @Override
    public boolean equals(Object obj)
    {
        if (obj == null) {
            return false;
        }

        if (!(obj instanceof KdbTree)) {
            return false;
        }

        KdbTree other = (KdbTree) obj;
        return this.root.equals(other.root);
    }

    @Override
    public int hashCode()
    {
        return Objects.hash(root);
    }

    public Map<Integer, Rectangle> getLeaves()
    {
        ImmutableMap.Builder<Integer, Rectangle> leaves = ImmutableMap.builder();
        addLeaves(root, leaves, node -> true);
        return leaves.build();
    }

    public Map<Integer, Rectangle> findIntersectingLeaves(Rectangle envelope)
    {
        ImmutableMap.Builder<Integer, Rectangle> leaves = ImmutableMap.builder();
        addLeaves(root, leaves, node -> node.extent.intersects(envelope));
        return leaves.build();
    }

    private static void addLeaves(Node node, ImmutableMap.Builder<Integer, Rectangle> leaves, Predicate<Node> predicate)
    {
        if (!predicate.apply(node)) {
            return;
        }

        if (node.leafId.isPresent()) {
            leaves.put(node.leafId.getAsInt(), node.extent);
        }
        else {
            addLeaves(node.left.get(), leaves, predicate);
            addLeaves(node.right.get(), leaves, predicate);
        }
    }

    private interface SplitDimension
    {
        Comparator<Rectangle> getComparator();

        double getValue(Rectangle rectangle);

        SplitResult<Rectangle> split(Rectangle rectangle, double value);
    }

    private static final SplitDimension BY_X = new SplitDimension()
    {
        private final Comparator<Rectangle> comparator = (first, second) -> ComparisonChain.start()
                .compare(first.getXMin(), second.getXMin())
                .compare(first.getYMin(), second.getYMin())
                .result();

        @Override
        public Comparator<Rectangle> getComparator()
        {
            return comparator;
        }

        @Override
        public double getValue(Rectangle rectangle)
        {
            return rectangle.getXMin();
        }

        @Override
        public SplitResult<Rectangle> split(Rectangle rectangle, double x)
        {
            checkArgument(rectangle.getXMin() < x && x < rectangle.getXMax());
            return new SplitResult<>(
                    new Rectangle(rectangle.getXMin(), rectangle.getYMin(), x, rectangle.getYMax()),
                    new Rectangle(x, rectangle.getYMin(), rectangle.getXMax(), rectangle.getYMax()));
        }
    };

    private static final SplitDimension BY_Y = new SplitDimension()
    {
        private final Comparator<Rectangle> comparator = (first, second) -> ComparisonChain.start()
                .compare(first.getYMin(), second.getYMin())
                .compare(first.getXMin(), second.getXMin())
                .result();

        @Override
        public Comparator<Rectangle> getComparator()
        {
            return comparator;
        }

        @Override
        public double getValue(Rectangle rectangle)
        {
            return rectangle.getYMin();
        }

        @Override
        public SplitResult<Rectangle> split(Rectangle rectangle, double y)
        {
            checkArgument(rectangle.getYMin() < y && y < rectangle.getYMax());
            return new SplitResult<>(
                    new Rectangle(rectangle.getXMin(), rectangle.getYMin(), rectangle.getXMax(), y),
                    new Rectangle(rectangle.getXMin(), y, rectangle.getXMax(), rectangle.getYMax()));
        }
    };

    private static final class LeafIdAllocator
    {
        private int nextId;

        public int next()
        {
            return nextId++;
        }
    }

    public static KdbTree buildKdbTree(int maxItemsPerNode, Rectangle extent, List<Rectangle> items)
    {
        checkArgument(maxItemsPerNode > 0, "maxItemsPerNode must be > 0");
        requireNonNull(extent, "extent is null");
        requireNonNull(items, "items is null");
        return new KdbTree(buildKdbTreeNode(maxItemsPerNode, 0, extent, items, new LeafIdAllocator()));
    }

    private static Node buildKdbTreeNode(int maxItemsPerNode, int level, Rectangle extent, List<Rectangle> items, LeafIdAllocator leafIdAllocator)
    {
        checkArgument(maxItemsPerNode > 0, "maxItemsPerNode must be > 0");
        checkArgument(level >= 0, "level must be >= 0");
        checkArgument(level <= MAX_LEVELS, "level must be <= 10,000");
        requireNonNull(extent, "extent is null");
        requireNonNull(items, "items is null");

        if (items.size() <= maxItemsPerNode || level == MAX_LEVELS) {
            return newLeaf(extent, leafIdAllocator.next());
        }

        // Split over longer side
        boolean splitVertically = extent.getWidth() >= extent.getHeight();
        Optional<SplitResult<Node>> splitResult = trySplit(splitVertically ? BY_X : BY_Y, maxItemsPerNode, level, extent, items, leafIdAllocator);
        if (!splitResult.isPresent()) {
            // Try spitting by the other side
            splitResult = trySplit(splitVertically ? BY_Y : BY_X, maxItemsPerNode, level, extent, items, leafIdAllocator);
        }

        if (!splitResult.isPresent()) {
            return newLeaf(extent, leafIdAllocator.next());
        }

        return newInternal(extent, splitResult.get().getLeft(), splitResult.get().getRight());
    }

    private static final class SplitResult<T>
    {
        private final T left;
        private final T right;

        private SplitResult(T left, T right)
        {
            this.left = requireNonNull(left, "left is null");
            this.right = requireNonNull(right, "right is null");
        }

        public T getLeft()
        {
            return left;
        }

        public T getRight()
        {
            return right;
        }
    }

    private static Optional<SplitResult<Node>> trySplit(SplitDimension splitDimension, int maxItemsPerNode, int level, Rectangle extent, List<Rectangle> items, LeafIdAllocator leafIdAllocator)
    {
        checkArgument(items.size() > 1, "Number of items to split must be > 1");

        // Sort envelopes by xMin or yMin
        List<Rectangle> sortedItems = ImmutableList.sortedCopyOf(splitDimension.getComparator(), items);

        // Find a mid-point
        int middleIndex = (sortedItems.size() - 1) / 2;
        Rectangle middleEnvelope = sortedItems.get(middleIndex);
        double splitValue = splitDimension.getValue(middleEnvelope);
        int splitIndex = middleIndex;

        // skip over duplicate values
        while (splitIndex < sortedItems.size() && splitDimension.getValue(sortedItems.get(splitIndex)) == splitValue) {
            splitIndex++;
        }

        // all values between left-of-middle and the end are the same, so can't split
        if (splitIndex == sortedItems.size()) {
            return Optional.empty();
        }

        // about half of the objects are <= splitValue, the rest are >= next value
        // assuming the input set of objects is a sample from a much larger set,
        // let's split in the middle; this way objects from the larger set with values
        // between splitValue and next value will get split somewhat evenly into left
        // and right partitions
        splitValue = (splitValue + splitDimension.getValue(sortedItems.get(splitIndex))) / 2;

        SplitResult<Rectangle> childExtents = splitDimension.split(extent, splitValue);

        return Optional.of(new SplitResult<>(
                buildKdbTreeNode(maxItemsPerNode, level + 1, childExtents.getLeft(), sortedItems.subList(0, splitIndex), leafIdAllocator),
                buildKdbTreeNode(maxItemsPerNode, level + 1, childExtents.getRight(), sortedItems.subList(splitIndex, sortedItems.size()), leafIdAllocator)));
    }
}
