ballistics_engine/

angle_calculations.rs

use crate::InternalBallisticInputs;
use std::f64;

// Constants for unit conversions
const FEET_TO_METERS: f64 = 0.3048;
const FPS_TO_MPS: f64 = FEET_TO_METERS;
const YARDS_TO_METERS: f64 = 0.9144;
const DEGREES_TO_RADIANS: f64 = std::f64::consts::PI / 180.0;
const RADIANS_TO_DEGREES: f64 = 180.0 / std::f64::consts::PI;

// Zero finding constants
const ZERO_FINDING_MAX_ITER: usize = 100;

/// Result of angle calculation
#[derive(Debug, Clone)]
pub struct AngleResult {
    pub angle_rad: f64,
    pub iterations_used: usize,
    pub final_error: f64,
    pub success: bool,
}

/// Brent's method for root finding - optimized implementation.
///
/// `x_tolerance` is an absolute tolerance in the same units as `a` and `b`.
/// [`AngleResult::final_error`] reports `|f(x)|` in the function's output units.
pub fn brent_root_find<F>(
    f: F,
    mut a: f64,
    mut b: f64,
    x_tolerance: f64,
    max_iterations: usize,
) -> Result<AngleResult, String>
where
    F: Fn(f64) -> f64,
{
    let mut fa = f(a);
    let mut fb = f(b);
    let mut iterations = 0;

    // Ensure the root is bracketed
    if fa * fb > 0.0 {
        return Err(format!("Root not bracketed: f({a}) = {fa}, f({b}) = {fb}"));
    }

    // Ensure |f(a)| >= |f(b)|
    if fa.abs() < fb.abs() {
        std::mem::swap(&mut a, &mut b);
        std::mem::swap(&mut fa, &mut fb);
    }

    let mut c = a;
    let mut fc = fa;
    let mut d = b - a;
    let mut e = d;

    while iterations < max_iterations {
        iterations += 1;

        if fb == 0.0 {
            return Ok(AngleResult {
                angle_rad: b,
                iterations_used: iterations,
                final_error: fb.abs(),
                success: true,
            });
        }

        if fc.abs() < fb.abs() {
            a = b;
            b = c;
            c = a;
            fa = fb;
            fb = fc;
            fc = fa;
        }

        let tolerance_scaled = 2.0 * f64::EPSILON * b.abs() + 0.5 * x_tolerance;
        let m = 0.5 * (c - b);

        if m.abs() <= tolerance_scaled {
            return Ok(AngleResult {
                angle_rad: b,
                iterations_used: iterations,
                final_error: fb.abs(),
                success: true,
            });
        }

        if e.abs() >= tolerance_scaled && fa.abs() > fb.abs() {
            // Check for safe division before interpolation
            if fc.abs() < f64::EPSILON || fa.abs() < f64::EPSILON {
                // Fallback to bisection if denominators are too small
                d = m;
                e = m;
            } else {
                let s = fb / fa;
                let mut p;
                let mut q;

                if (a - c).abs() < f64::EPSILON {
                    // Linear interpolation
                    p = 2.0 * m * s;
                    q = 1.0 - s;
                } else {
                    // Inverse quadratic interpolation
                    q = fa / fc;
                    let r = fb / fc;
                    p = s * (2.0 * m * q * (q - r) - (b - a) * (r - 1.0));
                    q = (q - 1.0) * (r - 1.0) * (s - 1.0);
                }

                if p > 0.0 {
                    q = -q;
                } else {
                    p = -p;
                }

                let s = e;
                e = d;

                // Check for safe division in the acceptance test
                if q.abs() > f64::EPSILON
                    && 2.0 * p < 3.0 * m * q - (tolerance_scaled * q).abs()
                    && p < (0.5 * s * q).abs()
                {
                    d = p / q;
                } else {
                    d = m;
                    e = d;
                }
            }
        } else {
            d = m;
            e = d;
        }

        a = b;
        fa = fb;

        if d.abs() > tolerance_scaled {
            b += d;
        } else if m > 0.0 {
            b += tolerance_scaled;
        } else {
            b -= tolerance_scaled;
        }

        fb = f(b);

        if (fc * fb) > 0.0 {
            c = a;
            fc = fa;
            e = b - a;
            d = e;
        }
    }

    Ok(AngleResult {
        angle_rad: b,
        iterations_used: iterations,
        final_error: fb.abs(),
        success: false,
    })
}

/// Calculate adjusted muzzle velocity for powder temperature sensitivity.
///
/// `powder_temp_sensitivity` is an additive velocity slope in m/s per degree Celsius; both
/// temperature fields are Celsius.
pub fn adjusted_muzzle_velocity(inputs: &InternalBallisticInputs) -> f64 {
    let mut mv = inputs.muzzle_velocity;

    if inputs.use_powder_sensitivity {
        mv += inputs.powder_temp_sensitivity * (inputs.temperature - inputs.powder_temp);
    }

    mv
}

/// Calculate zero angle using Brent's method and Rust trajectory integration.
///
/// `inputs.target_distance` is meters and `inputs.shooting_angle` is radians.
/// `trajectory_func` must return the projectile's fixed-frame vertical height
/// at the horizontal downrange distance `inputs.target_distance`.
pub fn zero_angle(
    inputs: &InternalBallisticInputs,
    trajectory_func: impl Fn(&InternalBallisticInputs, f64) -> Result<f64, String> + Copy,
) -> Result<AngleResult, String> {
    // Set up the target vertical position based on shooting angle
    let vert = if inputs.shooting_angle.abs() > 1e-6 {
        inputs.target_distance * inputs.shooting_angle.tan()
    } else {
        0.0
    };

    // Define the height difference function
    // MBA-192: Use NaN instead of -999.0 on trajectory failure to prevent false roots
    let height_diff = |look_angle_rad: f64| -> f64 {
        // Calculate bullet height at target distance minus target height
        match trajectory_func(inputs, look_angle_rad) {
            Ok(bullet_height) => bullet_height - vert,
            Err(_) => f64::NAN, // NaN causes Brent's method to fail gracefully
        }
    };

    // Reasonable bounds for the zero angle in radians
    // Most rifle zeroing will be within +/- 10 degrees
    let lower_bound = -10.0 * DEGREES_TO_RADIANS;
    let upper_bound = 10.0 * DEGREES_TO_RADIANS;

    // Try primary bounds first
    match brent_root_find(height_diff, lower_bound, upper_bound, 1e-6, 100) {
        Ok(result) if result.success => Ok(result),
        _ => {
            // Fallback to wider search range
            let wider_lower = -45.0 * DEGREES_TO_RADIANS;
            let wider_upper = 45.0 * DEGREES_TO_RADIANS;

            match brent_root_find(height_diff, wider_lower, wider_upper, 1e-5, 150) {
                Ok(result) if result.success => Ok(result),
                Ok(result) => {
                    // Return best attempt even if not fully successful
                    Ok(AngleResult {
                        angle_rad: result.angle_rad,
                        iterations_used: result.iterations_used,
                        final_error: result.final_error,
                        success: false,
                    })
                }
                Err(_) => {
                    // If all else fails, return 0 as a safe default
                    Ok(AngleResult {
                        angle_rad: 0.0,
                        iterations_used: 0,
                        final_error: f64::INFINITY,
                        success: false,
                    })
                }
            }
        }
    }
}

/// Solve muzzle angle using Brent's method optimization
pub fn solve_muzzle_angle(
    inputs: &InternalBallisticInputs,
    zero_distance_los_m: f64,
    trajectory_func: impl Fn(&InternalBallisticInputs) -> Result<f64, String> + Copy, // Returns drop_m
    angle_lower_deg: f64,
    angle_upper_deg: f64,
    rtol: f64,
) -> Result<AngleResult, String> {
    if angle_lower_deg >= angle_upper_deg {
        return Err("angle_lower_deg must be less than angle_upper_deg".to_string());
    }

    let lower = angle_lower_deg * DEGREES_TO_RADIANS;
    let mut upper = angle_upper_deg * DEGREES_TO_RADIANS;

    // Define the vertical error function
    let vertical_error = |angle_rad: f64| -> f64 {
        // Create modified inputs with new angle and target distance
        let mut candidate = inputs.clone();
        candidate.muzzle_angle = angle_rad * RADIANS_TO_DEGREES;
        candidate.target_distance = zero_distance_los_m / YARDS_TO_METERS; // Convert back to yards

        // NaN (not a finite sentinel) on trajectory failure: a large finite value like 1e6
        // can fake a sign change and make the bracket/Brent solver lock onto the
        // discontinuity and report a spurious root as success. NaN makes those comparisons
        // false, so the solver fails gracefully (mirrors the zero_angle MBA-192 fix).
        trajectory_func(&candidate).unwrap_or(f64::NAN)
    };

    // Check bounds
    let f_lower = vertical_error(lower);
    if f_lower.abs() < 1e-9 {
        return Ok(AngleResult {
            angle_rad: lower,
            iterations_used: 1,
            final_error: f_lower.abs(),
            success: true,
        });
    }

    let f_upper = vertical_error(upper);
    if f_upper.abs() < 1e-9 {
        return Ok(AngleResult {
            angle_rad: upper,
            iterations_used: 1,
            final_error: f_upper.abs(),
            success: true,
        });
    }

    // Expand upper bound if needed to get a sign change
    if f_lower * f_upper > 0.0 {
        let step = 5.0 * DEGREES_TO_RADIANS;
        let max_angle = 45.0 * DEGREES_TO_RADIANS;
        let mut current = upper;
        let mut f_current = f_upper;

        while current < max_angle && f_lower * f_current > 0.0 {
            current += step;
            f_current = vertical_error(current);
        }

        if f_lower * f_current > 0.0 {
            return Err("Unable to bracket zero; widen angle bounds or check inputs".to_string());
        }

        upper = current;
    }

    // Use Brent's method to find the root with safe tolerance calculation
    let range = (upper - lower).abs();
    let tolerance = if range > f64::EPSILON {
        rtol * range
    } else {
        rtol * 1e-12 // Minimum tolerance for very small ranges
    };
    brent_root_find(
        vertical_error,
        lower,
        upper,
        tolerance,
        ZERO_FINDING_MAX_ITER,
    )
}

/// Estimate bore-line drop for a horizontal shot in ICAO sea-level conditions.
///
/// Inputs are muzzle velocity in feet per second, distance in yards, bullet mass in grains, and
/// a G1 ballistic coefficient; the returned drop is meters. Bullet mass is retained for API
/// compatibility but does not enter the retardation calculation because G1 BC already includes
/// the projectile's sectional density and form factor.
pub fn quick_drop_estimate(
    muzzle_velocity_fps: f64,
    distance_yards: f64,
    _bullet_mass_grains: f64,
    bc: f64,
) -> f64 {
    if muzzle_velocity_fps <= 0.0 || distance_yards <= 0.0 {
        return 0.0; // No drop if no velocity or distance
    }

    let bc_safe = bc.max(0.1);
    let distance_ft = distance_yards * 3.0;
    let step_count = ((distance_yards / 5.0).ceil() as usize).clamp(32, 4096);
    let step_ft = distance_ft / step_count as f64;
    let gravity_ft_s2 = crate::constants::G_ACCEL_MPS2 / FEET_TO_METERS;
    let speed_of_sound_fps = crate::constants::SPEED_OF_SOUND_MPS / FPS_TO_MPS;

    // State is vertical position, downrange velocity, and vertical velocity, all in imperial
    // units. Integrating over downrange distance keeps the work bounded and still captures the
    // Mach-dependent G1 retardation and vertical drag that control flight time and drop.
    let derivatives = |state: &[f64; 3]| -> Option<[f64; 3]> {
        let [_, vx, vy] = *state;
        if !(vx.is_finite() && vy.is_finite() && vx > 1e-9) {
            return None;
        }
        let speed = vx.hypot(vy);
        let mach = speed / speed_of_sound_fps;
        let cd = crate::drag::get_drag_coefficient(mach, &crate::DragModel::G1);
        let drag_accel = speed.powi(2) * cd * crate::constants::CD_TO_RETARD / bc_safe;

        Some([
            vy / vx,
            -drag_accel / speed,
            (-gravity_ft_s2 - drag_accel * vy / speed) / vx,
        ])
    };

    let mut state = [0.0, muzzle_velocity_fps, 0.0];
    for _ in 0..step_count {
        let Some(k1) = derivatives(&state) else {
            return f64::NAN;
        };
        let midpoint: [f64; 3] = std::array::from_fn(|i| state[i] + 0.5 * step_ft * k1[i]);
        let Some(k2) = derivatives(&midpoint) else {
            return f64::NAN;
        };
        state = std::array::from_fn(|i| state[i] + step_ft * k2[i]);
    }

    -state[0] * FEET_TO_METERS
}

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

    fn create_test_inputs() -> InternalBallisticInputs {
        InternalBallisticInputs {
            muzzle_velocity: 823.0, // 2700 fps in m/s
            bc_value: 0.5,
            bullet_mass: 0.0109,      // 168 grains in kg
            bullet_diameter: 0.00782, // 0.308 inches in meters
            target_distance: 500.0,
            temperature: 21.1, // 70°F in Celsius
            powder_temp: 21.1, // Match temperature for default case (70°F)
            ..Default::default()
        }
    }

    #[test]
    fn test_brent_root_find_quadratic() {
        // Test with simple quadratic: x^2 - 4 = 0, root at x = 2
        let f = |x: f64| x * x - 4.0;
        let result = brent_root_find(f, 1.0, 3.0, 1e-6, 100).unwrap();

        assert!(result.success);
        assert!((result.angle_rad - 2.0).abs() < 1e-6);
        assert!(result.iterations_used > 0);
        assert!(result.final_error < 1e-6);
    }

    #[test]
    fn brent_uses_inverse_quadratic_interpolation() {
        let f = |x: f64| x * x - 2.0;
        let result = brent_root_find(f, 1.0, 2.0, 1e-12, 100).unwrap();

        assert!(result.success);
        assert!((result.angle_rad - 2.0_f64.sqrt()).abs() < 1e-12);
        assert!(
            result.iterations_used <= 10,
            "smooth quadratic should converge superlinearly, took {} iterations",
            result.iterations_used
        );
    }

    #[test]
    fn test_brent_root_find_linear() {
        // Test with linear function: 2x - 6 = 0, root at x = 3
        let f = |x: f64| 2.0 * x - 6.0;
        let result = brent_root_find(f, 0.0, 5.0, 1e-6, 100).unwrap();

        assert!(result.success);
        assert!((result.angle_rad - 3.0).abs() < 1e-6);
    }

    #[test]
    fn brent_angle_tolerance_is_invariant_to_residual_units() {
        let expected_root = 2.0_f64.sqrt();
        let x_tolerance = 1e-8;

        for scale in [1.0, 1e-9] {
            let result = brent_root_find(
                |angle_rad| scale * (angle_rad * angle_rad - 2.0),
                1.0,
                2.0,
                x_tolerance,
                100,
            )
            .unwrap();

            assert!(result.success);
            assert!(
                (result.angle_rad - expected_root).abs() <= x_tolerance,
                "residual scale {scale} loosened the angular tolerance: {result:?}"
            );
            assert!(result.iterations_used > 1);
        }

        let exhausted = brent_root_find(|x| 1e-9 * (x * x - 2.0), 1.0, 2.0, 1e-6, 0).unwrap();
        assert!(!exhausted.success);
    }

    #[test]
    fn test_brent_root_find_no_bracket() {
        // Test with function that doesn't change sign in the interval
        let f = |x: f64| x * x + 1.0; // Always positive
        let result = brent_root_find(f, 1.0, 3.0, 1e-6, 100);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Root not bracketed"));
    }

    #[test]
    fn test_adjusted_muzzle_velocity_no_sensitivity() {
        let inputs = create_test_inputs();

        let result = adjusted_muzzle_velocity(&inputs);
        assert_eq!(result, 823.0); // muzzle_velocity in m/s
    }

    #[test]
    fn test_adjusted_muzzle_velocity_with_sensitivity() {
        let mut inputs = create_test_inputs();
        inputs.use_powder_sensitivity = true;
        inputs.powder_temp_sensitivity = 0.6; // m/s per degree Celsius
        inputs.temperature = 31.1;
        inputs.powder_temp = 21.1;

        let result = adjusted_muzzle_velocity(&inputs);
        assert!((result - 829.0).abs() < 1e-12);

        inputs.temperature = 11.1;
        let colder_result = adjusted_muzzle_velocity(&inputs);
        assert!((colder_result - 817.0).abs() < 1e-12);
    }

    #[test]
    fn test_quick_drop_estimate() {
        let drop = quick_drop_estimate(2700.0, 500.0, 168.0, 0.5);

        // Should be a reasonable drop value (a few meters for 500 yards)
        assert!(drop > 0.0);
        assert!(drop < 50.0); // Sanity check - shouldn't be more than 50m drop

        // Test that higher BC gives less drop
        let drop_high_bc = quick_drop_estimate(2700.0, 500.0, 168.0, 0.8);
        assert!(drop_high_bc < drop);
    }

    #[test]
    fn quick_drop_tracks_g1_point_mass_reference() {
        let distance_yards = 500.0;
        let distance_m = distance_yards * YARDS_TO_METERS;
        let inputs = InternalBallisticInputs {
            muzzle_velocity: 1200.0 * FPS_TO_MPS,
            bc_value: 0.5,
            bc_type: crate::DragModel::G1,
            bullet_mass: 168.0 * crate::constants::GRAINS_TO_KG,
            bullet_diameter: 0.308 * 0.0254,
            muzzle_height: 0.0,
            ground_threshold: f64::NEG_INFINITY,
            ..Default::default()
        };

        let mut solver =
            crate::TrajectorySolver::new(inputs, Default::default(), Default::default());
        solver.set_max_range(distance_m);
        let reference = solver.solve().unwrap();
        let position = reference.position_at_range(distance_m).unwrap();
        let reference_drop_m = -position.y;
        assert!(
            (9.2..=9.6).contains(&reference_drop_m),
            "unexpected canonical G1 fixture: {reference_drop_m} m"
        );

        let estimated_drop_m = quick_drop_estimate(1200.0, distance_yards, 168.0, 0.5);
        let relative_error = (estimated_drop_m - reference_drop_m).abs() / reference_drop_m;
        assert!(
            relative_error < 0.1,
            "quick G1 drop {estimated_drop_m} m differs from reference {reference_drop_m} m by {:.1}%",
            relative_error * 100.0
        );
    }

    #[test]
    fn test_zero_angle_uses_si_distance_and_radians() {
        let mut inputs = create_test_inputs();
        inputs.target_distance = 800.0;
        inputs.shooting_angle = 1.0_f64.to_radians();

        let result = zero_angle(&inputs, |trajectory_inputs, look_angle_rad| {
            Ok(trajectory_inputs.target_distance * look_angle_rad)
        })
        .unwrap();

        assert!(result.success);
        assert!(
            (result.angle_rad - inputs.shooting_angle).abs() < 1e-4,
            "SI zero target should solve near {} rad, got {}",
            inputs.shooting_angle,
            result.angle_rad
        );
    }

    #[test]
    fn zero_angle_uses_horizontal_range_for_incline_geometry() {
        let mut inputs = create_test_inputs();
        inputs.target_distance = 800.0;
        inputs.shooting_angle = 30.0_f64.to_radians();

        let result = zero_angle(&inputs, |trajectory_inputs, look_angle_rad| {
            Ok(trajectory_inputs.target_distance * look_angle_rad.tan())
        })
        .unwrap();

        assert!(result.success);
        assert!(
            (result.angle_rad - inputs.shooting_angle).abs() < 1e-4,
            "horizontal-range zero should solve at {} rad, got {}",
            inputs.shooting_angle,
            result.angle_rad
        );
    }

    #[test]
    fn test_zero_angle_bounds() {
        // Test that angle bounds are reasonable
        let lower = -10.0 * DEGREES_TO_RADIANS;
        let upper = 10.0 * DEGREES_TO_RADIANS;

        assert!(lower < 0.0);
        assert!(upper > 0.0);
        assert!((upper - lower).abs() > 0.1); // Reasonable search range
    }

    #[test]
    fn test_brent_root_find_near_zero_function_values() {
        // Test with function that has very small values near the root
        // This exercises the EPSILON guards against division by zero
        // The algorithm should not panic and should converge to some result
        let f = |x: f64| (x - 1.0) * 1e-10; // Small function values
        let result = brent_root_find(f, 0.0, 2.0, 1e-12, 100).unwrap();

        // Main goal: no panic from division by zero
        // With very small function values, the algorithm should still work
        assert!(result.success || result.iterations_used > 0);
        // The result should be reasonably close to the root
        assert!((result.angle_rad - 1.0).abs() < 0.1);
    }

    #[test]
    fn test_brent_root_find_steep_function() {
        // Test with steep function that could cause large intermediate values
        let f = |x: f64| (x - 0.5).powi(3) * 1e6;
        let result = brent_root_find(f, 0.0, 1.0, 1e-9, 100).unwrap();

        assert!(result.success);
        assert!((result.angle_rad - 0.5).abs() < 1e-6);
    }

    #[test]
    fn test_brent_root_find_oscillating_convergence() {
        // Test function that might cause the interpolation to struggle
        // forcing fallback to bisection (tests the safety guards)
        let f = |x: f64| x.sin() - 0.5;
        let result = brent_root_find(f, 0.0, 1.0, 1e-10, 100).unwrap();

        assert!(result.success);
        // Root is at arcsin(0.5) ≈ 0.5236
        assert!((result.angle_rad - std::f64::consts::FRAC_PI_6).abs() < 1e-6);
    }

    #[test]
    fn test_brent_root_find_flat_region() {
        // Test with function that has flat regions (derivative near zero)
        // which could cause issues in interpolation
        let f = |x: f64| (x - 2.0).powi(5);
        let result = brent_root_find(f, 1.0, 3.0, 1e-8, 100).unwrap();

        assert!(result.success);
        assert!((result.angle_rad - 2.0).abs() < 1e-4);
    }
}