astrodynamics-gnss 0.14.0

//! Regular-grid TEC ionosphere delay variant.

use std::f64::consts::PI;

pub const EARTH_RADIUS_M: f64 = 6_371_000.0;
pub const IONOSPHERE_HEIGHT_M: f64 = 450_000.0;
pub const IONOSPHERE_RADIUS_M: f64 = EARTH_RADIUS_M + IONOSPHERE_HEIGHT_M;
pub const IONOSPHERE_CONSTANT: f64 = 40.308193 * 1e16;
pub const L1_FREQUENCY_HZ: f64 = 1_575_420_000.0;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TecGridEpoch {
    pub unix_nanos: i64,
    pub day_of_year: u16,
}

impl TecGridEpoch {
    pub fn new(unix_nanos: i64, day_of_year: u16) -> Self {
        Self {
            unix_nanos,
            day_of_year,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct TecGridEvalOptions {
    pub epoch: TecGridEpoch,
    pub min_elevation_rad: f64,
    pub nan_pierce_point_height_m: f64,
    pub frequency_hz: f64,
}

impl TecGridEvalOptions {
    pub fn l1(epoch: TecGridEpoch) -> Self {
        Self {
            epoch,
            min_elevation_rad: 5.0 * PI / 180.0,
            nan_pierce_point_height_m: IONOSPHERE_HEIGHT_M,
            frequency_hz: L1_FREQUENCY_HZ,
        }
    }
}

#[derive(Clone, Debug)]
pub struct TecGrid {
    epochs_ns: Vec<f64>,
    latitudes_deg: Vec<f64>,
    longitudes_deg: Vec<f64>,
    values: Vec<f64>,
}

impl TecGrid {
    pub fn new(
        epochs_ns: Vec<f64>,
        latitudes_deg: Vec<f64>,
        longitudes_deg: Vec<f64>,
        values: Vec<f64>,
    ) -> Result<Self, String> {
        if epochs_ns.len() < 2 || latitudes_deg.len() < 2 || longitudes_deg.len() < 2 {
            return Err("TEC grid axes must each contain at least two entries".to_string());
        }
        if !strictly_increasing(&epochs_ns)
            || !strictly_increasing(&latitudes_deg)
            || !strictly_increasing(&longitudes_deg)
        {
            return Err("TEC grid axes must be strictly increasing".to_string());
        }
        let expected = epochs_ns
            .len()
            .checked_mul(latitudes_deg.len())
            .and_then(|v| v.checked_mul(longitudes_deg.len()))
            .ok_or_else(|| "TEC grid dimensions overflow".to_string())?;
        if values.len() != expected {
            return Err(format!(
                "TEC grid has {} values but expected {}",
                values.len(),
                expected
            ));
        }
        Ok(Self {
            epochs_ns,
            latitudes_deg,
            longitudes_deg,
            values,
        })
    }

    pub fn from_json_str(input: &str) -> Result<Self, String> {
        let epochs_ns = parse_json_number_array(input, "epochs")?;
        let latitudes_deg = parse_json_number_array(input, "lats")?;
        let longitudes_deg = parse_json_number_array(input, "lons")?;
        let values = parse_json_number_array(input, "values")?;
        Self::new(epochs_ns, latitudes_deg, longitudes_deg, values)
    }

    pub fn vtec_at_pierce_point(
        &self,
        epoch: TecGridEpoch,
        longitude_deg: f64,
        latitude_deg: f64,
    ) -> Result<f64, String> {
        let latitude_deg = if latitude_deg.abs() > 87.5 {
            clamp(latitude_deg, -87.5, 87.5)
        } else {
            latitude_deg
        };
        self.interpolate_vtec(epoch.unix_nanos as f64, latitude_deg, longitude_deg)
    }

    pub(crate) fn interpolate_vtec(
        &self,
        epoch_ns: f64,
        latitude_deg: f64,
        longitude_deg: f64,
    ) -> Result<f64, String> {
        let (epoch_i, epoch_y) = interval(&self.epochs_ns, epoch_ns, "timestamp")?;
        let (lat_i, lat_y) = interval(&self.latitudes_deg, latitude_deg, "latitude")?;
        let (lon_i, lon_y) = interval(&self.longitudes_deg, longitude_deg, "longitude")?;

        let indices = [epoch_i, lat_i, lon_i];
        let norm_distances = [epoch_y, lat_y, lon_y];
        let shift_norm_distances = [
            1.0 - norm_distances[0],
            1.0 - norm_distances[1],
            1.0 - norm_distances[2],
        ];
        let shift_indices = [indices[0] + 1, indices[1] + 1, indices[2] + 1];

        let mut value = 0.0;
        for a in 0..2 {
            for b in 0..2 {
                for c in 0..2 {
                    let i0 = if a == 0 { indices[0] } else { shift_indices[0] };
                    let i1 = if b == 0 { indices[1] } else { shift_indices[1] };
                    let i2 = if c == 0 { indices[2] } else { shift_indices[2] };
                    let w0 = if a == 0 {
                        shift_norm_distances[0]
                    } else {
                        norm_distances[0]
                    };
                    let w1 = if b == 0 {
                        shift_norm_distances[1]
                    } else {
                        norm_distances[1]
                    };
                    let w2 = if c == 0 {
                        shift_norm_distances[2]
                    } else {
                        norm_distances[2]
                    };

                    let mut weight = 1.0;
                    weight *= w0;
                    weight *= w1;
                    weight *= w2;
                    let term = self.value_at(i0, i1, i2) * weight;
                    value += term;
                }
            }
        }
        Ok(value)
    }

    fn value_at(&self, epoch_i: usize, lat_i: usize, lon_i: usize) -> f64 {
        let n_lat = self.latitudes_deg.len();
        let n_lon = self.longitudes_deg.len();
        self.values[(epoch_i * n_lat + lat_i) * n_lon + lon_i]
    }
}

pub fn iono_delay_xyz<F>(
    grid: &TecGrid,
    options: TecGridEvalOptions,
    sat_xyz: &[f64; 3],
    receiver_xyz: &[f64; 3],
    ecef_to_lla: F,
) -> Result<f64, String>
where
    F: Fn(&[f64; 3]) -> [f64; 3],
{
    let (_pp_xyz, pp_lonlatalt, mut elevation_rad) =
        pierce_point(sat_xyz, receiver_xyz, &ecef_to_lla);
    if elevation_rad < options.min_elevation_rad {
        elevation_rad = options.min_elevation_rad;
    }

    let pp_lonlatalt = if pp_lonlatalt.iter().any(|v| v.is_nan()) {
        let receiver_lonlatalt = ecef_to_lla(receiver_xyz);
        [
            receiver_lonlatalt[0],
            receiver_lonlatalt[1],
            options.nan_pierce_point_height_m,
        ]
    } else {
        pp_lonlatalt
    };

    let vtec = grid.vtec_at_pierce_point(options.epoch, pp_lonlatalt[0], pp_lonlatalt[1])?;
    let obliquity_arg = EARTH_RADIUS_M * elevation_rad.cos() / IONOSPHERE_RADIUS_M;
    let stec = vtec / (1.0 - obliquity_arg * obliquity_arg).sqrt();
    Ok(IONOSPHERE_CONSTANT * stec / (options.frequency_hz * options.frequency_hz))
}

pub fn tec_xyz<F>(
    grid: &TecGrid,
    options: TecGridEvalOptions,
    sat_xyz: &[f64; 3],
    receiver_xyz: &[f64; 3],
    ecef_to_lla: F,
) -> Result<(f64, f64), String>
where
    F: Fn(&[f64; 3]) -> [f64; 3],
{
    let (_pp_xyz, pp_lonlatalt, mut elevation_rad) =
        pierce_point(sat_xyz, receiver_xyz, &ecef_to_lla);
    if elevation_rad < options.min_elevation_rad {
        elevation_rad = options.min_elevation_rad;
    }

    let pp_lonlatalt = if pp_lonlatalt.iter().any(|v| v.is_nan()) {
        let receiver_lonlatalt = ecef_to_lla(receiver_xyz);
        [
            receiver_lonlatalt[0],
            receiver_lonlatalt[1],
            options.nan_pierce_point_height_m,
        ]
    } else {
        pp_lonlatalt
    };

    let vtec = grid.vtec_at_pierce_point(options.epoch, pp_lonlatalt[0], pp_lonlatalt[1])?;
    let obliquity_arg = EARTH_RADIUS_M * elevation_rad.cos() / IONOSPHERE_RADIUS_M;
    let stec = vtec / (1.0 - obliquity_arg * obliquity_arg).sqrt();
    Ok((vtec, stec))
}

pub fn pierce_point<F>(
    sat_xyz: &[f64; 3],
    receiver_xyz: &[f64; 3],
    ecef_to_lla: F,
) -> ([f64; 3], [f64; 3], f64)
where
    F: Fn(&[f64; 3]) -> [f64; 3],
{
    let receiver_sat_vector = [
        sat_xyz[0] - receiver_xyz[0],
        sat_xyz[1] - receiver_xyz[1],
        sat_xyz[2] - receiver_xyz[2],
    ];

    let receiver_up = unit_vector(receiver_xyz);
    let sat_unit = unit_vector(&receiver_sat_vector);
    let elevation_rad = dot_three_fused(&sat_unit, &receiver_up).asin();

    let a = 1.0;
    let b = 2.0 * dot_three_fused(receiver_xyz, &sat_unit);
    let c = dot_three_fused(receiver_xyz, receiver_xyz) - IONOSPHERE_RADIUS_M * IONOSPHERE_RADIUS_M;
    let t = (-b + (b * b - 4.0 * a * c).sqrt()) / (2.0 * a);

    let pp_xyz = [
        receiver_xyz[0] + t * sat_unit[0],
        receiver_xyz[1] + t * sat_unit[1],
        receiver_xyz[2] + t * sat_unit[2],
    ];
    let pp_lonlatalt = ecef_to_lla(&pp_xyz);
    (pp_xyz, pp_lonlatalt, elevation_rad)
}

fn unit_vector(vector: &[f64; 3]) -> [f64; 3] {
    let norm = dot_three(vector, vector).sqrt();
    [vector[0] / norm, vector[1] / norm, vector[2] / norm]
}

fn dot_three(a: &[f64; 3], b: &[f64; 3]) -> f64 {
    a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
}

fn dot_three_fused(a: &[f64; 3], b: &[f64; 3]) -> f64 {
    a[2].mul_add(b[2], a[1].mul_add(b[1], a[0] * b[0]))
}

fn clamp(v: f64, lo: f64, hi: f64) -> f64 {
    if v < lo {
        lo
    } else if v > hi {
        hi
    } else {
        v
    }
}

fn strictly_increasing(values: &[f64]) -> bool {
    values.windows(2).all(|w| w[1] > w[0])
}

fn interval(axis: &[f64], x: f64, name: &str) -> Result<(usize, f64), String> {
    if x < axis[0] || x > axis[axis.len() - 1] {
        return Err(format!("{name} {x} is out of TEC grid bounds"));
    }
    let upper = axis.partition_point(|v| *v <= x);
    let mut lower = upper.saturating_sub(1);
    if lower >= axis.len() - 1 {
        lower = axis.len() - 2;
    }
    let y = (x - axis[lower]) / (axis[lower + 1] - axis[lower]);
    Ok((lower, y))
}

fn parse_json_number_array(input: &str, key: &str) -> Result<Vec<f64>, String> {
    let quoted_key = format!("\"{key}\"");
    let key_pos = input
        .find(&quoted_key)
        .ok_or_else(|| format!("missing JSON key {quoted_key}"))?;
    let open_rel = input[key_pos + quoted_key.len()..]
        .find('[')
        .ok_or_else(|| format!("missing array for JSON key {quoted_key}"))?;
    let open = key_pos + quoted_key.len() + open_rel;
    let close = matching_bracket(input, open)?;
    parse_json_numbers(&input[open + 1..close])
}

fn matching_bracket(input: &str, open: usize) -> Result<usize, String> {
    let mut depth = 0usize;
    for (idx, ch) in input[open..].char_indices() {
        match ch {
            '[' => depth += 1,
            ']' => {
                depth -= 1;
                if depth == 0 {
                    return Ok(open + idx);
                }
            }
            _ => {}
        }
    }
    Err("unterminated JSON array".to_string())
}

fn parse_json_numbers(input: &str) -> Result<Vec<f64>, String> {
    let bytes = input.as_bytes();
    let mut out = Vec::new();
    let mut i = 0usize;
    while i < bytes.len() {
        if bytes[i] == b'-' || bytes[i] == b'+' || bytes[i] == b'.' || bytes[i].is_ascii_digit() {
            let start = i;
            i += 1;
            while i < bytes.len()
                && (bytes[i].is_ascii_digit()
                    || bytes[i] == b'.'
                    || bytes[i] == b'e'
                    || bytes[i] == b'E'
                    || bytes[i] == b'-'
                    || bytes[i] == b'+')
            {
                i += 1;
            }
            let s = std::str::from_utf8(&bytes[start..i]).map_err(|e| e.to_string())?;
            out.push(s.parse::<f64>().map_err(|e| format!("{e}: {s}"))?);
        } else {
            i += 1;
        }
    }
    Ok(out)
}