ballistics_engine/

trajectory_sampling.rs

use crate::cli_api::BallisticsError;
use nalgebra::Vector3;
use std::collections::HashSet;
use std::fmt;

/// Hard ceiling for observations produced by one regular trajectory-sampling request.
///
/// This is initially aligned with [`crate::MAX_TRAJECTORY_POINTS`]. The limit is checked before
/// the sampler allocates its interpolation buffers or output vector.
pub const MAX_TRAJECTORY_SAMPLES: usize = 250_000;

/// Compute the candidate distance-grid size without performing an unchecked float-to-integer
/// conversion or addition.
///
/// The count intentionally matches the historical grid generator: positive intervals below
/// 0.1 m are clamped to 0.1 m, and the grid includes distance zero plus enough candidates to
/// reach the requested range. The historical endpoint-tolerance filter can remove the final
/// candidate; this function mirrors that decision so every grid with exactly the public limit is
/// accepted.
pub(crate) fn projected_sample_count(
    max_dist: f64,
    step_m: f64,
) -> Result<usize, BallisticsError> {
    if !max_dist.is_finite() || !step_m.is_finite() {
        return Err(BallisticsError::from(
            "trajectory sampling range and interval must be finite",
        ));
    }

    if step_m <= 0.0 || max_dist < 1e-9 {
        return Ok(0);
    }

    let step_size = step_m.max(0.1);
    let intervals = (max_dist / step_size).ceil();
    // The historical generator has `intervals + 1` candidates, and its filter can discard at
    // most the final candidate because `step_size >= 0.1`. If there are more than MAX intervals,
    // even discarding that candidate cannot bring the retained grid within the limit.
    if !intervals.is_finite() || intervals > MAX_TRAJECTORY_SAMPLES as f64 {
        return Err(BallisticsError::from(format!(
            "trajectory sample limit of {MAX_TRAJECTORY_SAMPLES} exceeded"
        )));
    }

    let intervals = intervals as usize;
    let candidate_count = intervals.checked_add(1).ok_or_else(|| {
        BallisticsError::from(format!(
            "trajectory sample limit of {MAX_TRAJECTORY_SAMPLES} exceeded"
        ))
    })?;
    let final_candidate_m = intervals as f64 * step_size;
    let retained_count = if final_candidate_m > max_dist + 0.1 {
        candidate_count - 1
    } else {
        candidate_count
    };

    if retained_count > MAX_TRAJECTORY_SAMPLES {
        Err(BallisticsError::from(format!(
            "trajectory sample limit of {MAX_TRAJECTORY_SAMPLES} exceeded"
        )))
    } else {
        Ok(retained_count)
    }
}

/// Trajectory flags for notable events
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum TrajectoryFlag {
    ZeroCrossing,
    MachTransition,
    Apex,
}

impl fmt::Display for TrajectoryFlag {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            TrajectoryFlag::ZeroCrossing => "zero_crossing",
            TrajectoryFlag::MachTransition => "mach_transition",
            TrajectoryFlag::Apex => "apex",
        })
    }
}

impl TrajectoryFlag {
    /// Return the stable wire-style name used by existing callers.
    ///
    /// Keep the inherent method alongside [`fmt::Display`] so fully qualified calls to the
    /// historical API continue to compile.
    #[allow(clippy::inherent_to_string_shadow_display)] // Preserve the established public method.
    pub fn to_string(&self) -> String {
        match self {
            TrajectoryFlag::ZeroCrossing => "zero_crossing".to_owned(),
            TrajectoryFlag::MachTransition => "mach_transition".to_owned(),
            TrajectoryFlag::Apex => "apex".to_owned(),
        }
    }
}

/// Single trajectory sample point
#[derive(Debug, Clone)]
pub struct TrajectorySample {
    pub distance_m: f64,
    pub drop_m: f64,
    pub wind_drift_m: f64,
    pub velocity_mps: f64,
    pub energy_j: f64,
    pub time_s: f64,
    pub flags: Vec<TrajectoryFlag>,
}

/// Trajectory solution data for sampling
#[derive(Debug, Clone)]
pub struct TrajectoryData {
    pub times: Vec<f64>,
    pub positions: Vec<Vector3<f64>>,  // [x, y, z] positions
    pub velocities: Vec<Vector3<f64>>, // [vx, vy, vz] velocities
    pub transonic_distances: Vec<f64>, // Distances where mach transitions occur
}

/// Output data for trajectory sampling
#[derive(Debug, Clone)]
pub struct TrajectoryOutputs {
    pub target_distance_horiz_m: f64,
    pub target_vertical_height_m: f64,
    pub time_of_flight_s: f64,
    pub max_ord_dist_horiz_m: f64,
    /// Height of sight above bore (meters). Used for LOS calculation.
    /// For a flat shot, the LOS is horizontal at y = sight_height_m.
    pub sight_height_m: f64,
}

/// Sample trajectory at regular distance intervals with vectorized operations.
///
/// # Errors
///
/// Returns [`BallisticsError`] when the requested range or interval is non-finite, or when the
/// retained distance grid would exceed [`MAX_TRAJECTORY_SAMPLES`].
pub fn sample_trajectory(
    trajectory_data: &TrajectoryData,
    outputs: &TrajectoryOutputs,
    step_m: f64,
    mass_kg: f64,
) -> Result<Vec<TrajectorySample>, BallisticsError> {
    // Use the input target distance as the limit for sampling
    let max_dist = outputs.target_distance_horiz_m;
    let num_steps = projected_sample_count(max_dist, step_m)?;
    if num_steps == 0 {
        return Ok(Vec::new());
    }
    let step_size = step_m.max(0.1);

    // Extract trajectory arrays for vectorized operations (McCoy: X=downrange, Z=lateral)
    let downrange_vals: Vec<f64> = trajectory_data.positions.iter().map(|p| p.x).collect();
    let y_vals: Vec<f64> = trajectory_data.positions.iter().map(|p| p.y).collect();
    let lateral_vals: Vec<f64> = trajectory_data.positions.iter().map(|p| p.z).collect();

    // Calculate speed at each integration knot.
    let speeds: Vec<f64> = trajectory_data
        .velocities
        .iter()
        .map(|v| v.norm())
        .collect();

    // Generate sampling distances. `num_steps` was checked before any sampler allocation.
    let distances: Vec<f64> = (0..num_steps)
        .map(|i| i as f64 * step_size)
        .filter(|&d| d <= max_dist + 0.1) // Stop exactly at target (with tiny tolerance for rounding)
        .collect();

    // Vectorized interpolation for all trajectory data
    let mut samples = Vec::with_capacity(distances.len());

    for &distance in &distances {
        // Interpolate using X (downrange) as the independent variable
        // McCoy coordinate system: x=downrange, y=vertical, z=lateral (wind drift)
        let y_interp = interpolate(&downrange_vals, &y_vals, distance); // vertical at downrange distance
        let wind_drift = interpolate(&downrange_vals, &lateral_vals, distance); // lateral drift at downrange distance
        let velocity = interpolate(&downrange_vals, &speeds, distance); // velocity at downrange distance
        let time = interpolate(&downrange_vals, &trajectory_data.times, distance); // time at downrange distance
        let energy = 0.5 * mass_kg * velocity * velocity;

        // Calculate line-of-sight y-coordinate and drop
        // The LOS is a straight line from the SIGHT to the target
        // The sight is at y = sight_height_m above the bore (which starts at y = 0)
        // For a flat shot: LOS is horizontal at y = sight_height_m
        // For elevated/depressed shots: LOS slopes from sight_height_m to target_vertical_height_m
        //
        // Drop convention:
        // - Positive drop means bullet is below LOS (has dropped)
        // - Negative drop means bullet is above LOS (has risen)
        // Therefore: drop = LOS - actual (not actual - LOS)
        //
        // LOS interpolation: starts at sight_height_m (z=0), ends at target_vertical_height_m (z=max_dist)
        // Note: For a properly zeroed flat shot, target_vertical_height_m should equal sight_height_m
        // (bullet ends at LOS at target distance for a point-blank shot)
        let los_y = outputs.sight_height_m
            + (outputs.target_vertical_height_m - outputs.sight_height_m) * distance / max_dist;
        let drop = los_y - y_interp; // LOS - actual: positive when bullet is below LOS

        samples.push(TrajectorySample {
            distance_m: distance,
            drop_m: drop,
            wind_drift_m: wind_drift,
            velocity_mps: velocity,
            energy_j: energy,
            time_s: time,
            flags: Vec::new(), // Flags will be added later
        });
    }

    // Add flags using vectorized detection
    add_trajectory_flags(&mut samples, &trajectory_data.transonic_distances, max_dist);

    Ok(samples)
}

/// Linear interpolation function optimized for trajectory data
fn interpolate(x_vals: &[f64], y_vals: &[f64], x: f64) -> f64 {
    if x_vals.is_empty() || y_vals.is_empty() {
        return 0.0;
    }

    if x_vals.len() != y_vals.len() {
        return 0.0;
    }

    if x <= x_vals[0] {
        return y_vals[0];
    }

    if x >= x_vals[x_vals.len() - 1] {
        return y_vals[y_vals.len() - 1];
    }

    // Binary search for the correct interval
    let mut left = 0;
    let mut right = x_vals.len() - 1;

    while right - left > 1 {
        let mid = (left + right) / 2;
        if x_vals[mid] <= x {
            left = mid;
        } else {
            right = mid;
        }
    }

    // Linear interpolation
    let x1 = x_vals[left];
    let x2 = x_vals[right];
    let y1 = y_vals[left];
    let y2 = y_vals[right];

    if (x2 - x1).abs() < f64::EPSILON {
        return y1;
    }

    y1 + (y2 - y1) * (x - x1) / (x2 - x1)
}

/// Add trajectory flags using vectorized detection algorithms
fn add_trajectory_flags(
    samples: &mut [TrajectorySample],
    transonic_distances: &[f64],
    target_distance_input_m: f64,
) {
    let tolerance = 1e-6;

    // 1. Zero crossings - vectorized detection
    detect_zero_crossings(samples, tolerance);

    // 2. Mach transitions
    for &transonic_dist in transonic_distances {
        if let Some(idx) = find_closest_sample_index(samples, transonic_dist) {
            samples[idx].flags.push(TrajectoryFlag::MachTransition);
        }
    }

    // 3. Apex - find the point with maximum height between muzzle and target
    // Since drop is positive when bullet is below LOS and negative when above,
    // the apex is where drop is minimum (most negative)
    if samples.len() > 2 {
        // Use the target distance passed as parameter
        let target_distance_m = target_distance_input_m;

        // Find the index of maximum height (minimum drop, most negative) within target distance.
        // Only mark an interior apex if it is actually above the muzzle/first sample.
        let first_drop = samples[0].drop_m;
        let mut min_drop = first_drop;
        let mut apex_idx: Option<usize> = None;

        // Search from index 1, but stop at target distance
        for (i, sample) in samples.iter().enumerate().skip(1) {
            // Only consider points up to target distance
            if sample.distance_m > target_distance_m {
                break;
            }

            if sample.drop_m < min_drop {
                min_drop = sample.drop_m;
                apex_idx = Some(i);
            }
        }

        if let Some(idx) = apex_idx {
            samples[idx].flags.push(TrajectoryFlag::Apex);
        }
    }
}

/// Detect zero crossings in trajectory drop values using vectorized operations
fn detect_zero_crossings(samples: &mut [TrajectorySample], tolerance: f64) {
    if samples.len() < 2 {
        return;
    }

    let drops: Vec<f64> = samples.iter().map(|s| s.drop_m).collect();

    // Find crossing indices where drop changes sign
    for i in 0..(drops.len() - 1) {
        let current = drops[i];
        let next = drops[i + 1];

        // Check for sign change crossings
        let crosses_zero = (current < -tolerance && next >= -tolerance)
            || (current > tolerance && next <= tolerance);

        if crosses_zero {
            samples[i + 1].flags.push(TrajectoryFlag::ZeroCrossing);
        }
    }

    // Find points very close to zero
    for (i, &drop) in drops.iter().enumerate() {
        if drop.abs() <= tolerance {
            samples[i].flags.push(TrajectoryFlag::ZeroCrossing);
        }
    }

    // Remove duplicate zero crossing flags
    for sample in samples.iter_mut() {
        let mut unique_flags = Vec::new();
        let mut seen = HashSet::new();

        for flag in &sample.flags {
            if seen.insert(flag.clone()) {
                unique_flags.push(flag.clone());
            }
        }
        sample.flags = unique_flags;
    }
}

/// Find the closest sample index to a given distance
fn find_closest_sample_index(samples: &[TrajectorySample], target_distance: f64) -> Option<usize> {
    if samples.is_empty() {
        return None;
    }

    // Binary search for the closest distance
    let distances: Vec<f64> = samples.iter().map(|s| s.distance_m).collect();

    let mut left = 0;
    let mut right = distances.len();

    while left < right {
        let mid = (left + right) / 2;
        if distances[mid] < target_distance {
            left = mid + 1;
        } else {
            right = mid;
        }
    }

    // Find the closest point (could be left-1 or left)
    let mut best_idx = left.min(distances.len() - 1);

    if left > 0 {
        let left_dist = (distances[left - 1] - target_distance).abs();
        let right_dist = (distances[best_idx] - target_distance).abs();

        // Prefer earlier index in case of tie
        if left_dist <= right_dist {
            best_idx = left - 1;
        }
    }

    Some(best_idx)
}

/// Convert trajectory samples to Python-compatible format
pub fn trajectory_samples_to_dicts(samples: &[TrajectorySample]) -> Vec<TrajectoryDict> {
    samples
        .iter()
        .map(|sample| TrajectoryDict {
            distance_m: sample.distance_m,
            drop_m: sample.drop_m,
            wind_drift_m: sample.wind_drift_m,
            velocity_mps: sample.velocity_mps,
            energy_j: sample.energy_j,
            time_s: sample.time_s,
            flags: sample.flags.iter().map(|f| f.to_string()).collect(),
        })
        .collect()
}

/// Python-compatible trajectory sample structure
#[derive(Debug, Clone)]
pub struct TrajectoryDict {
    pub distance_m: f64,
    pub drop_m: f64,
    pub wind_drift_m: f64,
    pub velocity_mps: f64,
    pub energy_j: f64,
    pub time_s: f64,
    pub flags: Vec<String>,
}

#[cfg(test)]
mod tests {
    use super::*;

    fn linear_fixture(max_dist: f64) -> (TrajectoryData, TrajectoryOutputs) {
        (
            TrajectoryData {
                times: vec![0.0, 1.0],
                positions: vec![
                    Vector3::new(0.0, -1.0, 0.0),
                    Vector3::new(max_dist, -1.0, 0.0),
                ],
                velocities: vec![
                    Vector3::new(800.0, 0.0, 0.0),
                    Vector3::new(700.0, 0.0, 0.0),
                ],
                transonic_distances: vec![],
            },
            TrajectoryOutputs {
                target_distance_horiz_m: max_dist,
                target_vertical_height_m: 0.0,
                time_of_flight_s: 1.0,
                max_ord_dist_horiz_m: 0.0,
                sight_height_m: 0.0,
            },
        )
    }

    #[test]
    fn mba1299_projected_sample_count_checks_exact_limit_and_overflow() {
        assert_eq!(MAX_TRAJECTORY_SAMPLES, crate::MAX_TRAJECTORY_POINTS);
        assert_eq!(
            projected_sample_count((MAX_TRAJECTORY_SAMPLES - 1) as f64, 1.0)
                .expect("the exact sample cap should be accepted"),
            MAX_TRAJECTORY_SAMPLES
        );
        assert_eq!(
            projected_sample_count(MAX_TRAJECTORY_SAMPLES as f64 - 0.5, 1.0)
                .expect("a filtered final candidate must not reject an exact-cap grid"),
            MAX_TRAJECTORY_SAMPLES
        );
        assert_eq!(
            projected_sample_count(0.2, 0.01)
                .expect("the historical 0.1 meter interval floor should remain valid"),
            3
        );

        for (range, interval) in [
            (MAX_TRAJECTORY_SAMPLES as f64, 1.0),
            (f64::MAX, 0.1),
        ] {
            let error = projected_sample_count(range, interval)
                .expect_err("a grid above the sample cap must fail");
            assert!(
                error
                    .to_string()
                    .contains("trajectory sample limit of 250000 exceeded"),
                "unexpected sampling limit error: {error}"
            );
        }
    }

    #[test]
    fn mba1299_public_sampler_accepts_the_exact_cap() {
        for max_dist in [
            (MAX_TRAJECTORY_SAMPLES - 1) as f64,
            MAX_TRAJECTORY_SAMPLES as f64 - 0.5,
        ] {
            let (trajectory_data, outputs) = linear_fixture(max_dist);
            let samples = sample_trajectory(&trajectory_data, &outputs, 1.0, 0.01)
                .expect("an exact-cap sample grid should succeed");

            assert_eq!(samples.len(), MAX_TRAJECTORY_SAMPLES);
            assert_eq!(samples.first().expect("muzzle sample").distance_m, 0.0);
            assert_eq!(
                samples.last().expect("terminal sample").distance_m,
                (MAX_TRAJECTORY_SAMPLES - 1) as f64
            );
        }
    }

    #[test]
    fn mba1299_public_sampler_rejects_oversized_grids_before_allocation() {
        for max_dist in [MAX_TRAJECTORY_SAMPLES as f64, f64::MAX] {
            let (trajectory_data, outputs) = linear_fixture(max_dist);
            let error = sample_trajectory(&trajectory_data, &outputs, 1.0, 0.01)
                .expect_err("an oversized public sampling request must fail");
            assert!(
                error
                    .to_string()
                    .contains("trajectory sample limit of 250000 exceeded"),
                "unexpected sampling limit error: {error}"
            );
        }
    }

    #[test]
    fn test_interpolate() {
        let x_vals = vec![0.0, 1.0, 2.0, 3.0];
        let y_vals = vec![0.0, 10.0, 20.0, 30.0];

        assert_eq!(interpolate(&x_vals, &y_vals, 0.5), 5.0);
        assert_eq!(interpolate(&x_vals, &y_vals, 1.5), 15.0);
        assert_eq!(interpolate(&x_vals, &y_vals, 2.5), 25.0);

        // Test boundary conditions
        assert_eq!(interpolate(&x_vals, &y_vals, -1.0), 0.0); // Below range
        assert_eq!(interpolate(&x_vals, &y_vals, 4.0), 30.0); // Above range
    }

    #[test]
    fn test_find_closest_sample_index() {
        let samples = vec![
            TrajectorySample {
                distance_m: 0.0,
                drop_m: 0.0,
                wind_drift_m: 0.0,
                velocity_mps: 100.0,
                energy_j: 1000.0,
                time_s: 0.0,
                flags: Vec::new(),
            },
            TrajectorySample {
                distance_m: 10.0,
                drop_m: -1.0,
                wind_drift_m: 0.1,
                velocity_mps: 95.0,
                energy_j: 950.0,
                time_s: 0.1,
                flags: Vec::new(),
            },
            TrajectorySample {
                distance_m: 20.0,
                drop_m: -4.0,
                wind_drift_m: 0.2,
                velocity_mps: 90.0,
                energy_j: 900.0,
                time_s: 0.2,
                flags: Vec::new(),
            },
        ];

        assert_eq!(find_closest_sample_index(&samples, 5.0), Some(0));
        assert_eq!(find_closest_sample_index(&samples, 12.0), Some(1));
        assert_eq!(find_closest_sample_index(&samples, 18.0), Some(2));
    }

    #[test]
    fn test_detect_zero_crossings() {
        let mut samples = vec![
            TrajectorySample {
                distance_m: 0.0,
                drop_m: 1.0, // Positive
                wind_drift_m: 0.0,
                velocity_mps: 100.0,
                energy_j: 1000.0,
                time_s: 0.0,
                flags: Vec::new(),
            },
            TrajectorySample {
                distance_m: 10.0,
                drop_m: -0.5, // Negative - crossing here
                wind_drift_m: 0.1,
                velocity_mps: 95.0,
                energy_j: 950.0,
                time_s: 0.1,
                flags: Vec::new(),
            },
            TrajectorySample {
                distance_m: 20.0,
                drop_m: -2.0, // Still negative
                wind_drift_m: 0.2,
                velocity_mps: 90.0,
                energy_j: 900.0,
                time_s: 0.2,
                flags: Vec::new(),
            },
        ];

        detect_zero_crossings(&mut samples, 1e-6);

        // Should have a zero crossing flag at index 1
        assert!(!samples[0].flags.contains(&TrajectoryFlag::ZeroCrossing));
        assert!(samples[1].flags.contains(&TrajectoryFlag::ZeroCrossing));
        assert!(!samples[2].flags.contains(&TrajectoryFlag::ZeroCrossing));
    }

    #[test]
    fn test_sample_trajectory_basic() {
        // Create simple test trajectory data
        // McCoy coordinate system: x=downrange, y=vertical, z=lateral (wind drift)
        let trajectory_data = TrajectoryData {
            times: vec![0.0, 1.0, 2.0],
            positions: vec![
                Vector3::new(0.0, 0.0, 0.0), // x=0 (start), y=0 (vertical), z=0 (no drift)
                Vector3::new(100.0, 10.0, 1.0), // x=100 (mid downrange), y=10 (apex height), z=1 (drift)
                Vector3::new(200.0, 5.0, 2.0), // x=200 (end downrange), y=5 (below apex), z=2 (drift)
            ],
            velocities: vec![
                Vector3::new(1.0, 10.0, 100.0),
                Vector3::new(1.0, 5.0, 95.0),
                Vector3::new(1.0, 0.0, 90.0),
            ],
            transonic_distances: vec![150.0],
        };

        let outputs = TrajectoryOutputs {
            target_distance_horiz_m: 200.0,
            target_vertical_height_m: 0.0,
            time_of_flight_s: 2.0,
            max_ord_dist_horiz_m: 100.0,
            sight_height_m: 0.0, // For test: assume bore-referenced coordinates
        };

        let samples = sample_trajectory(&trajectory_data, &outputs, 50.0, 0.1)
            .expect("normal sampling should succeed");

        // Should have samples at 0, 50, 100, 150, 200 meters
        assert_eq!(samples.len(), 5);
        assert_eq!(samples[0].distance_m, 0.0);
        assert_eq!(samples[1].distance_m, 50.0);
        assert_eq!(samples[2].distance_m, 100.0);
        assert_eq!(samples[3].distance_m, 150.0);
        assert_eq!(samples[4].distance_m, 200.0);

        // Check that interpolation is working
        assert!(samples[1].velocity_mps > 90.0 && samples[1].velocity_mps < 100.0);

        // Check flags
        assert!(samples[2].flags.contains(&TrajectoryFlag::Apex)); // At apex distance
        assert!(samples[3].flags.contains(&TrajectoryFlag::MachTransition)); // At transonic distance
    }

    #[test]
    fn sampled_energy_is_derived_from_interpolated_speed() {
        let mass_kg = 0.01;
        let trajectory_data = TrajectoryData {
            times: vec![0.0, 1.0],
            positions: vec![Vector3::zeros(), Vector3::new(100.0, 0.0, 0.0)],
            velocities: vec![Vector3::new(800.0, 0.0, 0.0), Vector3::new(700.0, 0.0, 0.0)],
            transonic_distances: vec![],
        };
        let outputs = TrajectoryOutputs {
            target_distance_horiz_m: 100.0,
            target_vertical_height_m: 0.0,
            time_of_flight_s: 1.0,
            max_ord_dist_horiz_m: 0.0,
            sight_height_m: 0.0,
        };

        let samples = sample_trajectory(&trajectory_data, &outputs, 50.0, mass_kg)
            .expect("normal sampling should succeed");
        assert_eq!(samples.len(), 3);
        assert_eq!(samples[1].velocity_mps.to_bits(), 750.0_f64.to_bits());
        assert_eq!(samples[1].energy_j.to_bits(), 2812.5_f64.to_bits());
        for sample in samples {
            let expected_energy = 0.5 * mass_kg * sample.velocity_mps * sample.velocity_mps;
            assert_eq!(sample.energy_j.to_bits(), expected_energy.to_bits());
        }
    }
}