mass-driver/crates/mass-driver-core/src/transfer_matrix.rs

use crate::station::{Station, station_position};
use orbital_mechanics::bodies;
use orbital_mechanics::constants::*;
use orbital_mechanics::lambert;
use nalgebra::Vector3;
use serde::{Deserialize, Serialize};

/// Precomputed transfer costs between station pairs over time.
///
/// Stored as a sparse list of feasible transfers to save memory.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransferMatrix {
    pub station_count: usize,
    pub week_count: usize,
    /// Start Julian Date
    pub start_jd: f64,
    /// Sparse entries: (from, to, departure_week, travel_weeks)
    pub entries: Vec<TransferEntry>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransferEntry {
    pub from: u16,
    pub to: u16,
    pub departure_week: u16,
    pub travel_weeks: u16,
}

/// Compute the transfer matrix for a set of stations.
///
/// For each station pair and each weekly departure time, solve Lambert's problem
/// to find the minimum transfer time achievable at the given launch velocity.
///
/// # Arguments
/// * `stations` - The station network
/// * `max_launch_v_kms` - Maximum launch velocity in km/s
/// * `start_jd` - Start of the time window (Julian Date)
/// * `week_count` - Number of weeks to compute
/// * `progress_callback` - Called with (completed, total) for progress reporting
pub fn compute_transfer_matrix(
    stations: &[Station],
    max_launch_v_kms: f64,
    start_jd: f64,
    week_count: usize,
    mut progress_callback: impl FnMut(usize, usize),
) -> TransferMatrix {
    let all_bodies = bodies::all_bodies();
    let n = stations.len();
    let total_work = n * (n - 1) * week_count;
    let mut completed = 0;

    let mut entries = Vec::new();

    // Candidate TOFs to try (in days) - logarithmic spacing from weeks to years
    let candidate_tofs_days: Vec<f64> = {
        let mut tofs = Vec::new();
        // From 1 week to 5 years, ~20 samples
        let min_days: f64 = 7.0;
        let max_days: f64 = 5.0 * 365.25;
        for i in 0..20 {
            let t: f64 = i as f64 / 19.0;
            let days = min_days * (max_days / min_days).powf(t);
            tofs.push(days);
        }
        tofs
    };

    for week in 0..week_count {
        let jd = start_jd + (week as f64) * DAYS_PER_WEEK;

        // Precompute all station positions at this epoch
        let station_positions: Vec<Vector3<f64>> = stations
            .iter()
            .map(|s| station_position(s, &all_bodies, jd))
            .collect();

        for from in 0..n {
            for to in 0..n {
                if from == to {
                    continue;
                }

                let r1 = station_positions[from] * AU_KM; // Convert AU to km
                let r2 = station_positions[to] * AU_KM;

                // Try different TOFs and find the minimum that works
                let mut best_travel_weeks: Option<u16> = None;

                for &tof_days in &candidate_tofs_days {
                    let tof_seconds = tof_days * SECONDS_PER_DAY;

                    if let Some(sol) = lambert::solve_lambert(r1, r2, tof_seconds, MU_SUN, true) {
                        // Check if departure delta-v is within launch capability
                        // For simplicity, we check if the departure velocity magnitude
                        // is achievable (station is co-moving with its parent body)
                        let v1_mag = sol.v1.norm();

                        // Station orbital velocity (approximate)
                        let r1_au = station_positions[from].norm();
                        let v_station = if r1_au > 0.01 {
                            (MU_SUN / (r1_au * AU_KM)).sqrt() // km/s
                        } else {
                            0.0
                        };

                        // The departure delta-v is roughly v1 - v_station
                        // This is simplified — in reality we'd need the vector difference
                        let dv_approx = (v1_mag - v_station).abs();

                        if dv_approx <= max_launch_v_kms {
                            let travel_weeks = (tof_days / 7.0).ceil() as u16;
                            if travel_weeks > 0 && travel_weeks < 5200 {
                                match best_travel_weeks {
                                    None => best_travel_weeks = Some(travel_weeks),
                                    Some(prev) if travel_weeks < prev => {
                                        best_travel_weeks = Some(travel_weeks);
                                    }
                                    _ => {}
                                }
                            }
                        }
                    }
                }

                if let Some(travel_weeks) = best_travel_weeks {
                    entries.push(TransferEntry {
                        from: from as u16,
                        to: to as u16,
                        departure_week: week as u16,
                        travel_weeks,
                    });
                }

                completed += 1;
                if completed % 1000 == 0 {
                    progress_callback(completed, total_work);
                }
            }
        }
    }

    progress_callback(total_work, total_work);

    TransferMatrix {
        station_count: n,
        week_count,
        start_jd,
        entries,
    }
}