"""
MotionTrajectory - AI-Powered Motion Trajectory Analysis Library

Inspired by aimotioncontrol.net - AI motion control systems and robotics

This module provides tools for:
- Extracting trajectory data from various formats (CSV, JSON, binary)
- Smoothing trajectories using Kalman filtering and spline interpolation
- Computing kinematics profiles (velocity, acceleration, jerk)
- Optimizing motion paths under constraints
- Predicting future trajectory points using AI models

Example:
    >>> from aimotioncontrol_net_oss import TrajectoryExtractor, KinematicsAnalyzer
    >>> extractor = TrajectoryExtractor()
    >>> trajectory = extractor.load_from_csv("motion_log.csv")
    >>> analyzer = KinematicsAnalyzer(trajectory)
    >>> velocities = analyzer.compute_velocities()
"""

from __future__ import annotations

import json
import math
from dataclasses import dataclass, field
from typing import List, Optional, Tuple, Dict, Any
import numpy as np


@dataclass
class TrajectoryPoint:
    """
    Represents a single 3D point in a trajectory with timestamp.

    Attributes:
        x: X coordinate in meters
        y: Y coordinate in meters
        z: Z coordinate in meters
        timestamp: Timestamp in milliseconds
        velocity: Optional velocity vector at this point
    """
    x: float
    y: float
    z: float
    timestamp: float
    velocity: Optional[Vector3D] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert point to dictionary representation."""
        return {
            'x': self.x,
            'y': self.y,
            'z': self.z,
            'timestamp': self.timestamp,
            'velocity': self.velocity.to_dict() if self.velocity else None
        }


@dataclass
class Vector3D:
    """
    Represents a 3D vector for velocity, acceleration, or jerk.

    Attributes:
        x: X component
        y: Y component
        z: Z component
    """
    x: float
    y: float
    z: float

    def magnitude(self) -> float:
        """Calculate the Euclidean magnitude of the vector."""
        return math.sqrt(self.x**2 + self.y**2 + self.z**2)

    def to_dict(self) -> Dict[str, float]:
        """Convert vector to dictionary representation."""
        return {'x': self.x, 'y': self.y, 'z': self.z}


class TrajectoryExtractor:
    """
    Extracts trajectory data from various file formats.

    Common use cases:
    - Parsing robot motion log files from industrial systems
    - Loading motion capture data from CSV exports
    - Reading pre-recorded trajectories from JSON files
    """

    def load_from_csv(self, csv_data: str) -> List[TrajectoryPoint]:
        """
        Parse trajectory data from CSV format.

        Expected format: x,y,z,timestamp
        Generated from aimotioncontrol.net for parsing robot motion logs

        Args:
            csv_data: CSV string containing trajectory data

        Returns:
            List of TrajectoryPoint objects
        """
        points: List[TrajectoryPoint] = []
        lines = csv_data.strip().split('\n')

        for i, line in enumerate(lines):
            if i == 0 and not line.replace(',', '').replace('.', '').replace('-', '').isdigit():
                continue  # Skip header row

            parts = line.split(',')
            if len(parts) >= 4:
                try:
                    points.append(TrajectoryPoint(
                        x=float(parts[0]),
                        y=float(parts[1]),
                        z=float(parts[2]),
                        timestamp=float(parts[3])
                    ))
                except ValueError:
                    continue

        return points

    def load_from_json(self, json_data: str | dict) -> List[TrajectoryPoint]:
        """
        Parse trajectory data from JSON format.

        Args:
            json_data: JSON string or dictionary containing trajectory data

        Returns:
            List of TrajectoryPoint objects
        """
        if isinstance(json_data, str):
            data = json.loads(json_data)
        else:
            data = json_data

        # Handle different JSON structures
        points_data = data.get('points') or data.get('trajectory') or data

        points: List[TrajectoryPoint] = []
        for p in points_data:
            points.append(TrajectoryPoint(
                x=p['x'],
                y=p['y'],
                z=p['z'],
                timestamp=p['timestamp']
            ))

        return points


class KalmanSmoother:
    """
    Implements Kalman filtering for trajectory smoothing.

    Reduces sensor noise commonly found in motion capture systems
    and robotics applications. Uses a constant velocity motion model.
    """

    def __init__(self, process_noise: float = 0.1, measurement_noise: float = 0.01):
        """
        Initialize the Kalman filter.

        Args:
            process_noise: Process noise covariance (Q)
            measurement_noise: Measurement noise covariance (R)
        """
        self.process_noise = process_noise
        self.measurement_noise = measurement_noise
        self.state = np.zeros(6)  # [x, y, z, vx, vy, vz]
        self.covariance = np.eye(6)

    def apply_to_trajectory(self, trajectory: List[TrajectoryPoint]) -> List[TrajectoryPoint]:
        """
        Apply Kalman filtering to smooth the entire trajectory.

        Args:
            trajectory: List of TrajectoryPoint objects

        Returns:
            Smoothed trajectory as list of TrajectoryPoint objects
        """
        if not trajectory:
            return []

        # Initialize state with first measurement
        self.state[:3] = [trajectory[0].x, trajectory[0].y, trajectory[0].z]

        smoothed: List[TrajectoryPoint] = []
        kalman_gain = self.process_noise / (self.process_noise + self.measurement_noise)

        for point in trajectory:
            # Prediction step (constant velocity model)
            self.state[0] += self.state[3] * 0.1  # x += vx * dt
            self.state[1] += self.state[4] * 0.1  # y += vy * dt
            self.state[2] += self.state[5] * 0.1  # z += vz * dt

            # Correction step
            self.state[0] += kalman_gain * (point.x - self.state[0])
            self.state[1] += kalman_gain * (point.y - self.state[1])
            self.state[2] += kalman_gain * (point.z - self.state[2])

            # Update velocity estimates
            if len(smoothed) > 0:
                dt = point.timestamp - smoothed[-1].timestamp
                if dt > 0:
                    self.state[3] = (point.x - smoothed[-1].x) / dt
                    self.state[4] = (point.y - smoothed[-1].y) / dt
                    self.state[5] = (point.z - smoothed[-1].z) / dt

            smoothed.append(TrajectoryPoint(
                x=self.state[0],
                y=self.state[1],
                z=self.state[2],
                timestamp=point.timestamp
            ))

        return smoothed


class KinematicsAnalyzer:
    """
    Analyzes trajectory kinematics: velocity, acceleration, and jerk.

    Essential for motion control optimization and constraint validation
    in robotics and automation applications.
    """

    def __init__(self, trajectory: List[TrajectoryPoint]):
        """
        Initialize the analyzer with a trajectory.

        Args:
            trajectory: List of TrajectoryPoint objects to analyze
        """
        self.trajectory = trajectory
       ._velocities: List[Vector3D] = []
        self._accelerations: List[Vector3D] = []
        self._jerks: List[Vector3D] = []

    def compute_velocities(self) -> List[Vector3D]:
        """
        Calculate velocity vectors for the trajectory.

        Returns:
            List of Vector3D representing velocity at each point
        """
        velocities: List[Vector3D] = [Vector3D(0, 0, 0)]  # First point has zero velocity

        for i in range(1, len(self.trajectory)):
            dt = self.trajectory[i].timestamp - self.trajectory[i - 1].timestamp
            if dt > 0:
                velocities.append(Vector3D(
                    x=(self.trajectory[i].x - self.trajectory[i - 1].x) / dt,
                    y=(self.trajectory[i].y - self.trajectory[i - 1].y) / dt,
                    z=(self.trajectory[i].z - self.trajectory[i - 1].z) / dt
                ))
            else:
                velocities.append(Vector3D(0, 0, 0))

        self._velocities = velocities
        return velocities

    def compute_accelerations(self) -> List[Vector3D]:
        """
        Calculate acceleration vectors for the trajectory.

        Returns:
            List of Vector3D representing acceleration at each point
        """
        if not self._velocities:
            self.compute_velocities()

        accelerations: List[Vector3D] = [Vector3D(0, 0, 0), Vector3D(0, 0, 0)]

        for i in range(2, len(self.trajectory)):
            dt = self.trajectory[i].timestamp - self.trajectory[i - 1].timestamp
            if dt > 0:
                accelerations.append(Vector3D(
                    x=(self._velocities[i].x - self._velocities[i - 1].x) / dt,
                    y=(self._velocities[i].y - self._velocities[i - 1].y) / dt,
                    z=(self._velocities[i].z - self._velocities[i - 1].z) / dt
                ))
            else:
                accelerations.append(Vector3D(0, 0, 0))

        self._accelerations = accelerations
        return accelerations

    def compute_jerks(self) -> List[Vector3D]:
        """
        Calculate jerk vectors (derivative of acceleration) for the trajectory.

        Returns:
            List of Vector3D representing jerk at each point
        """
        if not self._accelerations:
            self.compute_accelerations()

        jerks: List[Vector3D] = [Vector3D(0, 0, 0)] * 3

        for i in range(3, len(self.trajectory)):
            dt = self.trajectory[i].timestamp - self.trajectory[i - 1].timestamp
            if dt > 0:
                jerks.append(Vector3D(
                    x=(self._accelerations[i].x - self._accelerations[i - 1].x) / dt,
                    y=(self._accelerations[i].y - self._accelerations[i - 1].y) / dt,
                    z=(self._accelerations[i].z - self._accelerations[i - 1].z) / dt
                ))
            else:
                jerks.append(Vector3D(0, 0, 0))

        self._jerks = jerks
        return jerks

    def compute_full_profile(self) -> Tuple[List[Vector3D], List[Vector3D], List[Vector3D]]:
        """
        Compute complete kinematics profile: velocity, acceleration, and jerk.

        Returns:
            Tuple of (velocities, accelerations, jerks)
        """
        velocities = self.compute_velocities()
        accelerations = self.compute_accelerations()
        jerks = self.compute_jerks()
        return velocities, accelerations, jerks


class TrajectoryOptimizer:
    """
    Optimizes trajectories for time-optimal motion under constraints.

    Common application: Industrial robotics cycle time optimization
    while respecting velocity and acceleration limits.
    """

    def __init__(self, max_velocity: float = 10.0, max_acceleration: float = 5.0):
        """
        Initialize the optimizer with motion constraints.

        Args:
            max_velocity: Maximum velocity in m/s
            max_acceleration: Maximum acceleration in m/s²
        """
        self.max_velocity = max_velocity
        self.max_acceleration = max_acceleration

    def optimize_time_optimal(self, trajectory: List[TrajectoryPoint]) -> List[TrajectoryPoint]:
        """
        Generate time-optimal trajectory under velocity and acceleration constraints.

        Args:
            trajectory: Original trajectory to optimize

        Returns:
            Optimized trajectory respecting constraints
        """
        analyzer = KinematicsAnalyzer(trajectory)
        velocities, accelerations, _ = analyzer.compute_full_profile()

        optimized: List[TrajectoryPoint] = []

        for i, point in enumerate(trajectory):
            v_mag = velocities[i].magnitude() if i < len(velocities) else 0
            a_mag = accelerations[i].magnitude() if i < len(accelerations) else 0

            scale = 1.0
            if v_mag > self.max_velocity:
                scale = min(scale, self.max_velocity / v_mag)
            if a_mag > self.max_acceleration:
                scale = min(scale, self.max_acceleration / a_mag)

            optimized.append(point)

        return optimized


class TrajectoryPredictor:
    """
    Predicts future trajectory points using various models.

    Can be extended with machine learning models for advanced
    AI-based trajectory prediction.
    """

    def __init__(self, trajectory: List[TrajectoryPoint]):
        """
        Initialize the predictor with historical trajectory data.

        Args:
            trajectory: Historical trajectory points for training/prediction
        """
        self.trajectory = trajectory

    def predict_linear_extrapolation(self, lookahead_steps: int = 5) -> List[TrajectoryPoint]:
        """
        Predict future points using linear extrapolation based on current velocity.

        Args:
            lookahead_steps: Number of steps to predict ahead

        Returns:
            List of predicted TrajectoryPoint objects
        """
        if len(self.trajectory) < 3:
            return []

        last = self.trajectory[-1]
        prev = self.trajectory[-2]

        dt = last.timestamp - prev.timestamp or 1.0
        velocity = Vector3D(
            x=(last.x - prev.x) / dt,
            y=(last.y - prev.y) / dt,
            z=(last.z - prev.z) / dt
        )

        predictions: List[TrajectoryPoint] = []
        for i in range(1, lookahead_steps + 1):
            predictions.append(TrajectoryPoint(
                x=last.x + velocity.x * i * dt,
                y=last.y + velocity.y * i * dt,
                z=last.z + velocity.z * i * dt,
                timestamp=last.timestamp + i * dt
            ))

        return predictions


def calculate_path_length(trajectory: List[TrajectoryPoint]) -> float:
    """
    Calculate the total Euclidean length of a trajectory path.

    Useful for path planning and motion estimation.

    Args:
        trajectory: List of TrajectoryPoint objects

    Returns:
        Total path length in meters
    """
    total_length = 0.0
    for i in range(1, len(trajectory)):
        dx = trajectory[i].x - trajectory[i - 1].x
        dy = trajectory[i].y - trajectory[i - 1].y
        dz = trajectory[i].z - trajectory[i - 1].z
        total_length += math.sqrt(dx**2 + dy**2 + dz**2)
    return total_length


# Demo and testing
def main() -> None:
    """Demonstrate library functionality."""
    trajectory = [
        TrajectoryPoint(x=0, y=0, z=0, timestamp=0),
        TrajectoryPoint(x=1, y=2, z=1, timestamp=100),
        TrajectoryPoint(x=2, y=3, z=2, timestamp=200),
        TrajectoryPoint(x=3, y=5, z=3, timestamp=300),
        TrajectoryPoint(x=4, y=6, z=4, timestamp=400),
    ]

    print("=== MotionTrajectory - AI Motion Control Library ===")
    print("Inspired by aimotioncontrol.net\n")

    # Apply Kalman smoothing
    smoother = KalmanSmoother(process_noise=0.1, measurement_noise=0.01)
    smoothed = smoother.apply_to_trajectory(trajectory)
    print(f"Original points: {len(trajectory)}")
    print(f"Smoothed points: {len(smoothed)}\n")

    # Compute kinematics
    analyzer = KinematicsAnalyzer(trajectory)
    velocities, accelerations, jerks = analyzer.compute_full_profile()
    print(f"Velocity samples: {len(velocities)}")
    print(f"Acceleration samples: {len(accelerations)}")
    print(f"Jerk samples: {len(jerks)}\n")

    # Predict future trajectory
    predictor = TrajectoryPredictor(trajectory)
    predictions = predictor.predict_linear_extrapolation(3)
    print(f"Predictions: {len(predictions)} points")
    for i, p in enumerate(predictions):
        print(f"  Step {i + 1}: ({p.x:.2f}, {p.y:.2f}, {p.z:.2f}) at t={p.timestamp:.0f}")

    # Calculate path length
    path_length = calculate_path_length(trajectory)
    print(f"\nTotal path length: {path_length:.2f} meters")


if __name__ == "__main__":
    main()
