path-planning 0.1.0

/* Copyright (C) 2020 Dylan Staatz - All Rights Reserved. */

use nalgebra::constraint::{SameNumberOfRows, ShapeConstraint};
use nalgebra::geometry::UnitQuaternion;
use nalgebra::storage::Storage;
use nalgebra::{Const, Dim, RealField, SVector, Unit, Vector};
use num_traits::Float;
use rand::distributions::{uniform::SampleUniform, Distribution};
use rand::{thread_rng, SeedableRng};
use serde::{de::DeserializeOwned, Deserialize, Serialize};

use crate::error::{InvalidParamError, Result};
use crate::params::FromParams;
use crate::rng::{LinearCoordinates, RNG};
use crate::scalar::Scalar;
use crate::trajectories::{EuclideanTrajectory, FullTraj};
use crate::util::bounds::Bounds;
use crate::util::math::{atan2, unit_d_ball_vol};

use super::super::CSpace;
use super::polar::saturate_polar_zero;
use super::LeaderFollowerCSpace;

pub const D: usize = 3;
pub const N: usize = D * 2;

#[derive(Copy, Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(bound(
  serialize = "X: Serialize",
  deserialize = "X: DeserializeOwned"
))]
pub struct LeaderFollowerSphericalSpaceParams<X: Scalar> {
  pub bounds: Bounds<X, D>,
  pub seed: Option<u64>,
  pub sensor_range: (X, X),
}

/// A space with the state vector of [x1, y1, r2, theta2]
///
/// x1: x-coordinate of the leader robot
/// y1: y-coordinate of the leader robot
/// r2: radius polar coordinate offset of the follower relative to the leader
/// theta2: angle off of x-axis polar coordinate offset of the follower relative
/// to the leader
pub struct LeaderFollowerSphericalSpace<X>
where
  X: Scalar + SampleUniform,
{
  bounds: Bounds<X, D>,
  volume: X,
  rng: RNG,
  distribution: LinearCoordinates<X, N>,
  intial_sensor_range: (X, X),
  sensor_range: (X, X),
  sensor_range_cubed: (X, X),
}

impl<X> LeaderFollowerSphericalSpace<X>
where
  X: Scalar + SampleUniform,
{
  pub fn new(
    bounds: Bounds<X, D>,
    rng: RNG,
    sensor_range: (X, X),
  ) -> Result<Self> {
    debug_assert_eq!(D * 2, N);

    // Validate bounds
    if !bounds.is_valid() {
      Err(InvalidParamError {
        parameter_name: "bounds",
        parameter_value: format!("{:?}", bounds),
      })?;
    }

    // Validate sensor range
    if !(sensor_range.0 < sensor_range.1) {
      Err(InvalidParamError {
        parameter_name: "sensor_range",
        parameter_value: format!("{:?}", sensor_range),
      })?;
    }

    let sensor_range_cubed = (
      sensor_range.0 * sensor_range.0 * sensor_range.0,
      sensor_range.1 * sensor_range.1 * sensor_range.1,
    );

    let sensor_space_volume = (sensor_range_cubed.1 * unit_d_ball_vol(D))
      - (sensor_range_cubed.0 * unit_d_ball_vol(D));
    let volume = bounds.volume() * sensor_space_volume;

    let mins = SVector::<X, N>::from([
      bounds.mins[0],  // x1
      bounds.mins[1],  // y1
      bounds.mins[2],  // z1
      X::zero(),       // rho2
      X::zero(),       // lambda2
      -X::frac_pi_2(), // phi2
    ]);

    let maxs = SVector::<X, N>::from([
      bounds.maxs[0], // x1
      bounds.maxs[1], // y1
      bounds.maxs[2], // z1
      X::one(),       // rho2
      X::two_pi(),    // lambda2
      X::frac_pi_2(), // phi2
    ]);

    let distribution = LinearCoordinates::new(mins, maxs);

    Ok(Self {
      bounds,
      volume,
      rng,
      distribution,
      intial_sensor_range: sensor_range,
      sensor_range,
      sensor_range_cubed,
    })
  }

  pub fn intial_sensor_range(&self) -> (X, X) {
    self.intial_sensor_range
  }

  pub fn get_sensor_range(&self) -> (X, X) {
    self.sensor_range
  }

  pub fn set_sensor_range(&mut self, sensor_range: (X, X)) -> Option<()> {
    if sensor_range.0 < sensor_range.1 {
      // Valid
      self.sensor_range = sensor_range;
      self.sensor_range_cubed = (
        sensor_range.0 * sensor_range.0 * sensor_range.0,
        sensor_range.1 * sensor_range.1 * sensor_range.1,
      );
      Some(())
    } else {
      // Invalid value
      None
    }
  }
}

impl<X> LeaderFollowerCSpace<X, D, N> for LeaderFollowerSphericalSpace<X>
where
  X: Scalar + SampleUniform,
{
  // Default implementation okay
}

impl<X> CSpace<X, N> for LeaderFollowerSphericalSpace<X>
where
  X: Scalar + SampleUniform,
{
  type Traj = EuclideanTrajectory<X, N>;

  fn volume(&self) -> X {
    self.volume
  }

  fn cost<R1, R2, S1, S2>(
    &self,
    a: &Vector<X, R1, S1>,
    b: &Vector<X, R2, S2>,
  ) -> X
  where
    X: Scalar,
    R1: Dim,
    R2: Dim,
    S1: Storage<X, R1>,
    S2: Storage<X, R2>,
    ShapeConstraint: SameNumberOfRows<R1, R2>
      + SameNumberOfRows<R1, Const<N>>
      + SameNumberOfRows<R2, Const<N>>,
  {
    a.metric_distance(b)
  }

  fn trajectory<S1, S2>(
    &self,
    start: Vector<X, Const<N>, S1>,
    end: Vector<X, Const<N>, S2>,
  ) -> Option<FullTraj<X, Self::Traj, S1, S2, N>>
  where
    X: Scalar,
    S1: Storage<X, Const<N>>,
    S2: Storage<X, Const<N>>,
  {
    Some(FullTraj::new(start, end, EuclideanTrajectory::new()))
  }

  fn is_free<S>(&self, a: &Vector<X, Const<N>, S>) -> bool
  where
    S: Storage<X, Const<N>>,
  {
    let leader_abs = a.fixed_rows::<D>(0);
    let follower_abs = a.fixed_rows::<D>(D);

    if !self.bounds.within(&follower_abs) {
      return false;
    }

    // Determine distance between leader and follower
    let rho2 = leader_abs.metric_distance(&follower_abs);
    self.sensor_range.0 <= rho2 && rho2 <= self.sensor_range.1
  }

  fn saturate(&self, a: &mut SVector<X, N>, b: &SVector<X, N>, delta: X) {
    let delta = delta / (X::one() + X::one());

    // Saturate leader to be delta away
    let mut a_lead_mut = a.fixed_rows_mut::<D>(0);
    let b_lead = b.fixed_rows::<D>(0);

    let lead_scale = delta / a_lead_mut.metric_distance(&b_lead);

    a_lead_mut.set_column(0, &(&a_lead_mut - &b_lead));
    a_lead_mut.set_column(0, &(&a_lead_mut * lead_scale));
    a_lead_mut.set_column(0, &(&a_lead_mut + &b_lead));

    // Saturate follower to be delta away
    let a_lead = a.fixed_rows::<D>(0);
    let a_fol = a.fixed_rows::<D>(D);

    // Modifiy b follower to be an offset of the a_lead
    let mut b = b.clone(); // Local copy of b for temporary modifications
    let mut b_lead_mut = b.fixed_rows_mut::<D>(0);
    b_lead_mut.set_column(0, &a_lead); // A leader, B follower
    let mut b_rel = abs_to_rel(&b); // B follower relative to new A leader

    // Narrow down to spherical coordinates only
    let mut b_fol_rel_mut = b_rel.fixed_rows_mut::<D>(D).into_owned();
    let a_fol_rel = cartesian_to_spherical(&a_fol);

    // Bounds
    let mut delta = delta;
    if self.sensor_range.1 < b_fol_rel_mut[0] {
      delta -= b_fol_rel_mut[0] - self.sensor_range.1;
      b_fol_rel_mut[0] = self.sensor_range.1;
    }
    if b_fol_rel_mut[0] < self.sensor_range.0 {
      delta -= self.sensor_range.0 - b_fol_rel_mut[0];
      b_fol_rel_mut[0] = self.sensor_range.0;
    }

    if X::zero() <= delta {
      let mut star_fol_rel =
        saturate_spherical(&a_fol_rel, &b_fol_rel_mut, delta);

      log::debug!(
        "sensor_range: {:?}, star_fol_rel: {:?}",
        self.sensor_range,
        <[X; D]>::from(star_fol_rel)
      );

      // bound rho
      if self.sensor_range.1 < star_fol_rel[0] {
        star_fol_rel[0] = self.sensor_range.1;
      }
      if star_fol_rel[0] < self.sensor_range.0 {
        star_fol_rel[0] = self.sensor_range.0;
      }

      // Set a in relative spherical coordinates
      let mut a_fol_mut = a.fixed_rows_mut::<D>(D);
      a_fol_mut.set_column(0, &star_fol_rel);
    } else {
      // Set a in relative spherical coordinates
      let mut a_fol_mut = a.fixed_rows_mut::<D>(D);
      a_fol_mut.set_column(0, &b_fol_rel_mut);
    }

    // Convert relative spherical coordinates back to absolute
    *a = rel_to_abs(&a);
  }

  fn sample(&mut self) -> SVector<X, N> {
    let mut sample = self.distribution.sample(&mut self.rng);

    // Scale rho from [0,1] to the correct range
    let a_3 = self.sensor_range_cubed.0;
    let b_3 = self.sensor_range_cubed.1;
    sample[3] = <X as Float>::cbrt(sample[3] * (b_3 - a_3) + a_3);

    rel_to_abs(&sample)
  }
}

impl<X> FromParams for LeaderFollowerSphericalSpace<X>
where
  X: Scalar + SampleUniform,
{
  type Params = LeaderFollowerSphericalSpaceParams<X>;
  fn from_params(params: Self::Params) -> Result<Self> {
    let rng = match params.seed {
      Some(seed) => RNG::seed_from_u64(seed),
      None => RNG::from_rng(thread_rng())?,
    };

    LeaderFollowerSphericalSpace::new(params.bounds, rng, params.sensor_range)
  }
}

/// Converts [x1, y1, z1, rho2, lambda2, phi2] to [x1, y1, z1, x2, y2, z2]
///
/// Inverse of [`abs_to_rel`]
pub fn rel_to_abs<X, R, S>(v: &Vector<X, R, S>) -> SVector<X, N>
where
  X: RealField + Copy,
  R: Dim,
  S: Storage<X, R>,
  ShapeConstraint: SameNumberOfRows<R, Const<N>>,
{
  let x1 = v[0];
  let y1 = v[1];
  let z1 = v[2];
  let x2 = x1 + (v[3] * v[5].cos() * v[4].cos());
  let y2 = y1 + (v[3] * v[5].cos() * v[4].sin());
  let z2 = z1 + (v[3] * v[5].sin());

  [x1, y1, z1, x2, y2, z2].into()
}

/// Converts [rho, lambda, phi] to [x, y, z]
pub fn spherical_to_cartesian<X, R, S>(v: &Vector<X, R, S>) -> SVector<X, D>
where
  X: RealField + Copy,
  R: Dim,
  S: Storage<X, R>,
  ShapeConstraint: SameNumberOfRows<R, Const<D>>,
{
  [
    v[0] * v[2].cos() * v[1].cos(),
    v[0] * v[2].cos() * v[1].sin(),
    v[0] * v[2].sin(),
  ]
  .into()
}

/// Converts [x1, y1, z1, x2, y2, z2] to [x1, y1, z1, rho2, lambda2, phi2]
///
/// Inverse of [`rel_to_abs`]
pub fn abs_to_rel<X, R, S>(v: &Vector<X, R, S>) -> SVector<X, N>
where
  X: RealField + Copy,
  R: Dim,
  S: Storage<X, R>,
  ShapeConstraint: SameNumberOfRows<R, Const<N>>,
{
  let x1 = v[0];
  let y1 = v[1];
  let z1 = v[2];
  let x2 = v[3] - x1; // offset cartesian
  let y2 = v[4] - y1; // offset cartesian
  let z2 = v[5] - z1; // offset cartesian

  let rho2 = (x2 * x2 + y2 * y2 + z2 * z2).sqrt();
  let lambda2 = atan2(y2, x2).unwrap_or(X::zero());
  let phi2 = atan2(z2, (x2 * x2 + y2 * y2).sqrt()).unwrap_or(X::zero());

  [x1, y1, z1, rho2, bound_lambda(lambda2), bound_phi(phi2)].into()
}

/// Converts [x, y, z] to [rho, lambda, phi]
pub fn cartesian_to_spherical<X, R, S>(v: &Vector<X, R, S>) -> SVector<X, D>
where
  X: RealField + Copy,
  R: Dim,
  S: Storage<X, R>,
  ShapeConstraint: SameNumberOfRows<R, Const<D>>,
{
  let rho2 = (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt();
  let lambda2 = atan2(v[1], v[0]).unwrap_or(X::zero());
  let phi2 =
    atan2(v[2], (v[0] * v[0] + v[1] * v[1]).sqrt()).unwrap_or(X::zero());
  [rho2, bound_lambda(lambda2), bound_phi(phi2)].into()
}

fn saturate_spherical<X, R1, S1, R2, S2>(
  a: &Vector<X, R1, S1>,
  b: &Vector<X, R2, S2>,
  delta: X,
) -> SVector<X, D>
where
  X: RealField + Copy,
  R1: Dim,
  R2: Dim,
  S1: Storage<X, R1>,
  S2: Storage<X, R2>,
  ShapeConstraint: SameNumberOfRows<R1, R2>
    + SameNumberOfRows<R1, Const<D>>
    + SameNumberOfRows<R2, Const<D>>,
{
  let rho_a = a[0];
  let lambda_a = a[1];
  let phi_a = a[2];
  let rho_b = b[0];
  let lambda_b = b[1];
  let phi_b = b[2];

  // Normal Vectors: https://en.wikipedia.org/wiki/N-vector
  let n_a = SVector::<X, D>::from([
    phi_a.cos() * lambda_a.cos(),
    phi_a.cos() * lambda_a.sin(),
    phi_a.sin(),
  ]);

  let n_b = SVector::<X, D>::from([
    phi_b.cos() * lambda_b.cos(),
    phi_b.cos() * lambda_b.sin(),
    phi_b.sin(),
  ]);

  // Slicing planes normal (b -> a)
  let n = n_b.cross(&n_a);
  let (n, magnitude) = Unit::new_and_get(n);

  // The angle formed by the great circle path: https://en.wikipedia.org/wiki/Great-circle_distance#Vector_version
  let omega = atan2(magnitude, n_b.dot(&n_a)).unwrap_or(X::zero());

  // Setup 2D problem in polar coordinates on slicing plane
  let r_a = rho_a;
  let theta_a = omega;
  let r_b = rho_b;

  log::debug!("a: [{:?}, {:?}]", r_a, theta_a);
  log::debug!("b: [{:?}, {:?}]", r_b, 0.0);

  let (r_star, theta_star) = saturate_polar_zero(r_a, theta_a, r_b, delta);

  log::debug!("*: [{:?}, {:?}]", r_star, theta_star);

  // Convert 2D solution back to 3D
  let rotation = UnitQuaternion::from_axis_angle(&n, theta_star);
  let v = rotation * n_b;

  // Return bounded
  let rho_star = r_star;
  let lambda_star = atan2(v.y, v.x).unwrap_or(X::zero());
  let phi_star =
    atan2(v.z, (v.x.powi(2) + v.y.powi(2)).sqrt()).unwrap_or(X::zero());

  [rho_star, bound_lambda(lambda_star), bound_phi(phi_star)].into()
}

/// Longitude
fn bound_lambda<X>(mut lambda: X) -> X
where
  X: RealField + Copy,
{
  // Lower bound (Inclusive)
  while lambda < X::zero() {
    lambda += X::two_pi()
  }
  // Upper bound (Exclusive)
  while X::two_pi() <= lambda {
    lambda -= X::two_pi()
  }
  lambda
}

/// Latitude
fn bound_phi<X>(mut phi: X) -> X
where
  X: RealField + Copy,
{
  // Lower bound (Inclusive)
  while phi < -X::frac_pi_2() {
    phi += X::pi()
  }
  // Upper bound (Inclusive)
  while X::frac_pi_2() < phi {
    phi -= X::pi()
  }
  phi
}