/*
 * HPPC
 *
 * Copyright (C) 2010-2024 Carrot Search s.c. and contributors
 * All rights reserved.
 *
 * Refer to the full license file "LICENSE.txt":
 * https://github.com/carrotsearch/hppc/blob/master/LICENSE.txt
 */
package com.carrotsearch.hppc;

import java.util.Arrays;

/**
 * Optimal Piecewise Linear Approximation Model for <code>KType</code> keys.
 *
 * <p>
 * Learns a mapping that returns a position for a <code>KType</code> key which
 * is at most epsilon away from the correct one in a sorted list of keys. It is
 * optimal and piecewise because it learns the minimum number of
 * epsilon-approximate segments.
 *
 * <p>
 * The PLA-model consists of a sequence segments. A segment s is a triple
 * (key,slope,intercept) that indexes a range of keys through the function fs(k)
 * = k × slope + intercept, which provides an epsilon-approximation of the
 * position of the key k.
 */
public class PlaModel implements Accountable {

	/** Re-usable mutable (x,y) point. */
	private static class Point {

		static final int RAM_BYTES_ALLOCATED = RamUsageEstimator.NUM_BYTES_OBJECT_HEADER + Double.BYTES + Long.BYTES;

		double x;
		long y;

		Point set(double x, long y) {
			this.x = x;
			this.y = y;
			return this;
		}

		Point set(Point p) {
			return set(p.x, p.y);
		}
	}

	/**
	 * List of mutable {@link Point}. Re-uses allocated points instead of creating
	 * new instances.
	 */
	private static class PointList implements Accountable {

		Point[] points;
		int size;
		int numAllocated;

		PointList(int initialCapacity) {
			points = new Point[initialCapacity];
		}

		void add(Point point) {
			if (size == points.length) {
				int newSize = BoundedProportionalArraySizingStrategy.DEFAULT_INSTANCE.grow(points.length, size, 1);
				points = Arrays.copyOf(points, newSize);
			}
			if (size == numAllocated) {
				points[numAllocated++] = new Point();
			}
			points[size++].set(point);
		}

		void clear() {
			size = 0;
		}

		void clearFrom(int end) {
			size = end;
		}

		Point get(int index) {
			return points[index];
		}

		@Override
		public long ramBytesAllocated() {
			// int: size, numAllocated
			return RamUsageEstimator.NUM_BYTES_OBJECT_HEADER + 2 * Integer.BYTES
					+ RamUsageEstimator.shallowSizeOfArray(points) + (long) numAllocated * Point.RAM_BYTES_ALLOCATED;
		}

		@Override
		public long ramBytesUsed() {
			return ramBytesAllocated();
		}

		int size() {
			return size;
		}
	}

	/** Consumer notified when a new segment is built by the {@link PlaModel}. */
	public interface SegmentConsumer {

		/**
		 * Consumes a new segment. The segment is defined by the epsilon-approximation
		 * function fs(k) = k × slope + intercept.
		 *
		 * @param firstKey  The first key of the segment.
		 * @param slope     The segment slope.
		 * @param intercept The segment intercept.
		 */
		void accept(double firstKey, double slope, long intercept);
	}

	/** Re-usable mutable (dx,dy) slope. */
	private static class Slope {

		static final int RAM_BYTES_ALLOCATED = RamUsageEstimator.NUM_BYTES_OBJECT_HEADER + Double.BYTES + Long.BYTES;

		static double asDouble(Point p1, Point p2) {
			return (double) (p2.y - p1.y) / (p2.x - p1.x);
		}

		double dx;

		long dy;

		boolean isEqual(Slope s) {
			return Double.doubleToLongBits(dy * s.dx) == Double.doubleToLongBits(dx * s.dy);
		}

		boolean isGreaterThan(Slope s) {
			return dy * s.dx > dx * s.dy;
		}

		boolean isLessThan(Slope s) {
			return dy * s.dx < dx * s.dy;
		}

		Slope set(Point p1, Point p2) {
			dx = p2.x - p1.x;
			dy = p2.y - p1.y;
			return this;
		}

		void set(Slope s) {
			dx = s.dx;
			dy = s.dy;
		}
	}

	/** Initial capacity of the lower and upper point lists. */
	private static final int INITIAL_CAPACITY = 1 << 8;

	private static double cross(Point o, Point a, Point b) {
		return (a.x - o.x) * (b.y - o.y) - (a.y - o.y) * (b.x - o.x);
	}

	/** Epsilon precision of the PLA-model. */
	private int epsilon;

	/** First key of the current segment. */
	private double firstKey;

	/**
	 * Previous key used to check that keys are added in strictly increasing
	 * sequence.
	 */
	private double previousKey;

	/** Number of points in the convex hull for the current segment. */
	private int numPointsInHull;

	/** Enclosing rectangle for the current segment. */
	private final Point[] rect = new Point[4];
	/**
	 * Ordered list of lower points for the current segment. Inside the list,
	 * allocated points are re-used.
	 */
	private final PointList lower = new PointList(INITIAL_CAPACITY);
	/**
	 * Ordered list of upper points for the current segment. Inside the list,
	 * allocated points are re-used.
	 */
	private final PointList upper = new PointList(INITIAL_CAPACITY);
	/** Index of the first lower point to compare to. */
	private int lowerStart;
	/** Index of the first upper point to compare to. */
	private int upperStart;
	// Re-used mutable points and slopes.
	private final Point point1 = new Point();
	private final Point point2 = new Point();

	private final Slope slope1 = new Slope();

	private final Slope slope2 = new Slope();

	private final Slope slopeTmp = new Slope();

	private final Slope slopeMin = new Slope();

	private final Slope slopeMax = new Slope();

	/**
	 * Creates an optimal PLA-model with the provided epsilon precision.
	 *
	 * @param epsilon must be greater than or equal to 0.
	 */
	public PlaModel(int epsilon) {
		setEpsilon(epsilon);
		for (int i = 0; i < rect.length; i++) {
			rect[i] = new Point();
		}
		reset();
	}

	private int addEpsilon(int index) {
		try {
			return Math.addExact(index, epsilon);
		} catch (ArithmeticException e) {
			return Integer.MAX_VALUE;
		}
	}

	/**
	 * Adds a key to this PLA-model. The keys must be provided in a strictly
	 * increasing sequence. That is, the key must be greater than the previous key.
	 *
	 * @param index           The index of the key in the sorted key list.
	 * @param segmentConsumer The consumer to call when a new segment is built in
	 *                        the PLA-model.
	 */
	public void addKey(double key, int index, SegmentConsumer segmentConsumer) {
		if (key <= previousKey) {
			throw new IllegalArgumentException("Keys must be increasing");
		}
		previousKey = key;
		point1.set(key, addEpsilon(index));
		point2.set(key, subtractEpsilon(index));

		if (numPointsInHull > 1) {
			slope1.set(rect[0], rect[2]);
			slope2.set(rect[1], rect[3]);
			boolean outside_line1 = slopeTmp.set(rect[2], point1).isLessThan(slope1);
			boolean outside_line2 = slopeTmp.set(rect[3], point2).isGreaterThan(slope2);
			if (outside_line1 || outside_line2) {
				produceSegment(segmentConsumer);
				numPointsInHull = 0;
			}
		}
		if (numPointsInHull == 0) {
			firstKey = key;
			rect[0].set(point1);
			rect[1].set(point2);
			upper.clear();
			lower.clear();
			upper.add(point1);
			lower.add(point2);
			upperStart = lowerStart = 0;
			numPointsInHull++;
			return;
		}
		if (numPointsInHull == 1) {
			rect[2].set(point2);
			rect[3].set(point1);
			upper.add(point1);
			lower.add(point2);
			numPointsInHull++;
			return;
		}

		if (slopeTmp.set(rect[1], point1).isLessThan(slope2)) {
			// Find extreme slope.
			slopeMin.set(point1, lower.get(lowerStart));
			int min_i = lowerStart;
			for (int i = lowerStart + 1; i < lower.size(); i++) {
				slopeTmp.set(point1, lower.get(i));
				if (slopeTmp.isGreaterThan(slopeMin)) {
					break;
				}
				slopeMin.set(slopeTmp);
				min_i = i;
			}
			rect[1].set(lower.get(min_i));
			rect[3].set(point1);
			lowerStart = min_i;

			// Hull update.
			int end = upper.size();
			while (end >= upperStart + 2 && cross(upper.get(end - 2), upper.get(end - 1), point1) <= 0) {
				end--;
			}
			upper.clearFrom(end);
			upper.add(point1);
		}

		if (slopeTmp.set(rect[0], point2).isGreaterThan(slope1)) {
			// Find extreme slope.
			slopeMax.set(point2, upper.get(upperStart));
			int max_i = upperStart;
			for (int i = upperStart + 1; i < upper.size(); i++) {
				slopeTmp.set(point2, upper.get(i));
				if (slopeTmp.isLessThan(slopeMax)) {
					break;
				}
				slopeMax.set(slopeTmp);
				max_i = i;
			}
			rect[0].set(upper.get(max_i));
			rect[2].set(point2);
			upperStart = max_i;

			// Hull update.
			int end = lower.size();
			while (end >= lowerStart + 2 && cross(lower.get(end - 2), lower.get(end - 1), point2) >= 0) {
				end--;
			}
			lower.clearFrom(end);
			lower.add(point2);
		}

		numPointsInHull++;
	}

	/**
	 * Finishes the PLA-model construction. Declares that no additional keys will be
	 * added. Builds the last segment and calls the provided
	 * {@link SegmentConsumer}.
	 */
	public void finish(SegmentConsumer segmentConsumer) {
		produceSegment(segmentConsumer);
		reset();
	}

	private void produceSegment(SegmentConsumer segmentConsumer) {
		double slope;
		long intercept;

		if (numPointsInHull == 1) {
			slope = 0d;
			intercept = ((long) rect[0].y + rect[1].y) >>> 1;

		} else {
			Point p0 = rect[0];
			Point p1 = rect[1];
			Point p2 = rect[2];
			Point p3 = rect[3];

			// Compute the slope intersection point.
			double intersectX;
			double intersectY;
			slope1.set(p0, p2);
			slope2.set(p1, p3);
			if (slope1.isEqual(slope2)) {
				intersectX = p0.x;
				intersectY = p0.y;
			} else {
				slopeTmp.set(p0, p1);
				double a = slope1.dx * slope2.dy - slope1.dy * slope2.dx;
				double b = (slopeTmp.dx * slope2.dy - slopeTmp.dy * slope2.dx) / a;
				intersectX = p0.x + b * slope1.dx;
				intersectY = p0.y + b * slope1.dy;
			}

			// Compute the slope range.
			double minSlope = Slope.asDouble(p0, p2);
			double maxSlope = Slope.asDouble(p1, p3);

			// Compute the segment slope and intercept.
			slope = (minSlope + maxSlope) / 2d;
			intercept = (long) (intersectY - (intersectX - firstKey) * slope);
		}

		segmentConsumer.accept(firstKey, slope, intercept);
	}

	@Override
	public long ramBytesAllocated() {
		// int: epsilon, numPointsInHull, lowerStart, upperStart
		// double: firstKey, previousKey
		// Point: rect[4], point1, point2
		// Slope: slope1, slope2, slopeTmp, slopeMin, slopeMax
		return RamUsageEstimator.NUM_BYTES_OBJECT_HEADER + 4 * Integer.BYTES + 2 * Double.BYTES
				+ 6L * Point.RAM_BYTES_ALLOCATED + lower.ramBytesAllocated() + upper.ramBytesAllocated()
				+ 5L * Slope.RAM_BYTES_ALLOCATED;
	}

	@Override
	public long ramBytesUsed() {
		return ramBytesAllocated();
	}

	private void reset() {
		previousKey = Double.NEGATIVE_INFINITY;
		numPointsInHull = 0;
		lower.clear();
		upper.clear();
	}

	/** Sets epsilon precision which must be greater than or equal to 0. */
	public void setEpsilon(int epsilon) {
		if (epsilon < 0) {
			throw new IllegalArgumentException("epsilon must be >= 0");
		}
		this.epsilon = epsilon;
	}

	private int subtractEpsilon(int index) {
		try {
			return Math.subtractExact(index, epsilon);
		} catch (ArithmeticException e) {
			return Integer.MIN_VALUE;
		}
	}
}
