"""
Physical trajectory model for pickleball.
Accumulates ball positions across frames, fits physics-based curves,
smooths noise, and predicts ball path.

Key physics:
- Gravity: 9.81 m/s^2
- Pickleball drag coefficient: ~0.4 (perforated ball)
- Typical serve speed: 10-25 m/s
- Bounce coefficient of restitution: ~0.4
"""

import numpy as np
from collections import deque
from typing import List, Optional, Tuple, Dict
import time

GRAVITY = 9.81  # m/s^2
BALL_MASS = 0.026  # kg
BALL_RADIUS = 0.037  # meters
DRAG_COEFF = 0.4
AIR_DENSITY = 1.225  # kg/m^3
CROSS_SECTION = np.pi * BALL_RADIUS ** 2
BOUNCE_COR = 0.4  # coefficient of restitution


class TrajectoryPoint:
    __slots__ = ('x', 'y', 'z', 'timestamp', 'frame', 'sensor_id',
                 'confidence', 'interpolated')

    def __init__(self, x, y, z, timestamp, frame, sensor_id,
                 confidence=1.0, interpolated=False):
        self.x = x
        self.y = y
        self.z = z
        self.timestamp = timestamp
        self.frame = frame
        self.sensor_id = sensor_id
        self.confidence = confidence
        self.interpolated = interpolated

    def to_dict(self):
        return {
            'x': self.x, 'y': self.y, 'z': self.z,
            'timestamp': self.timestamp,
            'frame': self.frame,
            'sensor_id': self.sensor_id,
            'confidence': self.confidence,
            'interpolated': self.interpolated,
        }


class TrajectoryModel:
    """Accumulates ball observations and builds a physics-consistent trajectory."""

    def __init__(self, max_points=300, fps=30):
        self.points: deque = deque(maxlen=max_points)
        self.fps = fps
        self.velocity = None  # estimated (vx, vy, vz) m/s
        self._last_update = 0.0

    def add_observation(self, x: float, y: float, z: float,
                        timestamp: float, frame: int, sensor_id: int,
                        confidence: float = 1.0):
        """Add a new observed ball position."""
        pt = TrajectoryPoint(x, y, z, timestamp, frame, sensor_id, confidence)
        self.points.append(pt)
        self._update_velocity()

    def _update_velocity(self):
        """Estimate current velocity from last few points."""
        real_pts = [p for p in self.points if not p.interpolated and p.confidence > 0.3]
        if len(real_pts) < 2:
            self.velocity = None
            return

        p1, p2 = real_pts[-2], real_pts[-1]
        dt = p2.timestamp - p1.timestamp
        if dt <= 0:
            return

        self.velocity = np.array([
            (p2.x - p1.x) / dt,
            (p2.y - p1.y) / dt,
            (p2.z - p1.z) / dt,
        ])

    def predict_position(self, dt_ahead: float) -> Optional[Tuple[float, float, float]]:
        """Predict where the ball will be dt_ahead seconds from now.

        Uses simple ballistic model: constant horizontal velocity + gravity.
        """
        if not self.points or self.velocity is None:
            return None

        last = self.points[-1]
        vx, vy, vz = self.velocity

        # Simple drag approximation (reduces horizontal speed)
        speed = np.sqrt(vx**2 + vy**2 + vz**2)
        if speed > 0:
            drag_decel = (0.5 * DRAG_COEFF * AIR_DENSITY * CROSS_SECTION * speed**2) / BALL_MASS
            drag_factor = max(0, 1 - (drag_decel / speed) * dt_ahead)
        else:
            drag_factor = 1.0

        x = last.x + vx * dt_ahead * drag_factor
        y = last.y + vy * dt_ahead * drag_factor
        z = last.z + vz * dt_ahead - 0.5 * GRAVITY * dt_ahead**2

        # Bounce off ground
        if z < 0:
            z = abs(z) * BOUNCE_COR

        return (x, y, max(z, 0))

    def predict_landing(self) -> Optional[Tuple[float, float, float]]:
        """Predict where the ball will hit the ground (Z=0).

        Returns (x, y, time_to_landing) or None.
        """
        if not self.points or self.velocity is None:
            return None

        last = self.points[-1]
        vz = self.velocity[2]

        # Already on ground
        if last.z <= 0.05:
            return (last.x, last.y, 0.0)

        # Solve: z + vz*t - 0.5*g*t^2 = 0
        # 0.5*g*t^2 - vz*t - z = 0
        a = 0.5 * GRAVITY
        b = -vz
        c = -last.z

        discriminant = b**2 - 4*a*c
        if discriminant < 0:
            return None

        t = (-b + np.sqrt(discriminant)) / (2 * a)
        if t < 0:
            t = (-b - np.sqrt(discriminant)) / (2 * a)
        if t < 0:
            return None

        predicted = self.predict_position(t)
        if predicted:
            return (predicted[0], predicted[1], t)
        return None

    def get_smoothed_trajectory(self, window=5) -> List[Dict]:
        """Get trajectory with moving average smoothing on real observations."""
        pts = list(self.points)
        if len(pts) < window:
            return [p.to_dict() for p in pts]

        result = []
        for i, p in enumerate(pts):
            start = max(0, i - window // 2)
            end = min(len(pts), i + window // 2 + 1)
            neighbors = [pts[j] for j in range(start, end)
                         if not pts[j].interpolated]

            if neighbors:
                d = p.to_dict()
                d['x_smooth'] = np.mean([n.x for n in neighbors])
                d['y_smooth'] = np.mean([n.y for n in neighbors])
                d['z_smooth'] = np.mean([n.z for n in neighbors])
                result.append(d)
            else:
                result.append(p.to_dict())

        return result

    def get_recent(self, n=60) -> List[Dict]:
        """Get last N trajectory points as dicts."""
        pts = list(self.points)[-n:]
        return [p.to_dict() for p in pts]

    def get_speed(self) -> Optional[float]:
        """Current ball speed in m/s."""
        if self.velocity is None:
            return None
        return float(np.linalg.norm(self.velocity))

    def clear(self):
        self.points.clear()
        self.velocity = None
