"""
Trajectory Predictor — Polynomial Curve Fitting for Path Prediction
====================================================================
Fits a polynomial curve to historical positions and extrapolates
the trajectory 1-3 seconds into the future.

Features:
    - 2nd-degree polynomial (quadratic) fitting by default
    - R² goodness-of-fit confidence scoring
    - Configurable prediction horizon and minimum data requirements
    - Pure NumPy — no external dependencies

Author: ByIbos
License: MIT
"""

import numpy as np


class TrajectoryPredictor:
    """
    Predicts future positions by fitting a polynomial curve to
    historical (x, y) position data.

    Uses numpy.polyfit to fit separate polynomials for x(t) and y(t),
    then extrapolates forward. The confidence score is based on the
    R² (coefficient of determination) of the fit.

    Args:
        min_points: Minimum number of historical points required
            before a prediction is generated (default: 10).
        predict_frames: How many frames into the future to predict
            (default: 45, ~1.5 seconds at 30 FPS).
        poly_degree: Degree of the polynomial fit. 2 = quadratic
            (captures acceleration), 1 = linear (constant velocity).

    Example:
        >>> tp = TrajectoryPredictor(min_points=10, predict_frames=45)
        >>> history = [(100, 200), (102, 198), (105, 195), ...]  # 10+ points
        >>> future_path = tp.predict(history)
        >>> print(future_path)  # [(148, 155), (150, 152), ...]
        >>> print(tp.confidence)  # 0.92
    """

    def __init__(self, min_points=10, predict_frames=45, poly_degree=2):
        self.min_points = min_points
        self.predict_frames = predict_frames
        self.poly_degree = poly_degree

        self.predicted_path = []   # Predicted path [(x, y), ...]
        self.confidence = 0.0      # Prediction confidence [0.0 - 1.0]

    def predict(self, history):
        """
        Predict the future trajectory from historical position data.

        Fits a polynomial curve to the most recent positions (up to 40)
        and extrapolates forward by `predict_frames` steps.

        Args:
            history: An iterable of (x, y) tuples, e.g. a deque or list.
                Must contain at least `min_points` entries.

        Returns:
            list: A list of (x, y) integer tuples representing the
            predicted future path. Empty list if insufficient data.
        """
        pts = list(history)

        if len(pts) < self.min_points:
            self.predicted_path = []
            self.confidence = 0.0
            return []

        # Use only the most recent points (distant past can be misleading)
        pts = pts[-40:]

        t = np.arange(len(pts))
        xs = np.array([p[0] for p in pts])
        ys = np.array([p[1] for p in pts])

        try:
            # Polynomial fitting for x(t) and y(t)
            coeff_x = np.polyfit(t, xs, self.poly_degree)
            coeff_y = np.polyfit(t, ys, self.poly_degree)

            poly_x = np.poly1d(coeff_x)
            poly_y = np.poly1d(coeff_y)

            # Goodness of fit (R² score)
            xs_fit = poly_x(t)
            ys_fit = poly_y(t)
            ss_res_x = np.sum((xs - xs_fit) ** 2)
            ss_tot_x = np.sum((xs - np.mean(xs)) ** 2)
            ss_res_y = np.sum((ys - ys_fit) ** 2)
            ss_tot_y = np.sum((ys - np.mean(ys)) ** 2)

            r2_x = 1 - (ss_res_x / max(ss_tot_x, 1e-6))
            r2_y = 1 - (ss_res_y / max(ss_tot_y, 1e-6))
            self.confidence = max(0, min(1, (r2_x + r2_y) / 2))

            # Extrapolate into the future
            future_t = np.arange(len(pts), len(pts) + self.predict_frames)
            future_x = poly_x(future_t)
            future_y = poly_y(future_t)

            self.predicted_path = [(int(x), int(y)) for x, y in zip(future_x, future_y)]

        except (np.RankWarning, np.linalg.LinAlgError):
            self.predicted_path = []
            self.confidence = 0.0

        return self.predicted_path

    def reset(self):
        """Reset the predictor, clearing all stored predictions."""
        self.predicted_path = []
        self.confidence = 0.0
