// The rapidly-exploring random tree and its refinements. Each one grows the same
// tree and differs in how it picks the parent and when it stops.

///|
/// A rapidly-exploring random tree.
///
/// `step` is the length of the edge a new sample is allowed to reach, and
/// `goal_bias` the chance that a sample is the goal itself rather than a random
/// state. The generator is seeded from the planner, so a run replays exactly.
pub(all) struct Rrt {
  space : &@contract.Space
  validity : &@contract.Validity
  step : Double
  goal_bias : Double
  rng : &@rand_contract.Rng
}

///|
/// An RRT that extends by `step` and steers to the goal once every
/// `1 / goal_bias` samples.
pub fn Rrt::new(
  space~ : &@contract.Space,
  validity~ : &@contract.Validity,
  step~ : Double = 0.1,
  goal_bias~ : Double = 0.05,
  seed~ : Int32 = 0,
) -> Rrt raise @linear.Refused {
  if step <= 0.0 {
    raise @linear.Refused::Gain(got=step)
  }
  if goal_bias < 0.0 || goal_bias > 1.0 {
    raise @linear.Refused::Range(got=goal_bias, low=0.0, high=1.0)
  }
  { space, validity, step, goal_bias, rng: @rand.seeded(seed_bytes(seed)), }
}

///|
pub impl @contract.Planner for Rrt with fn plan(self, start, goal, budget) {
  let tree = Tree::new(start)
  for _ in 0.. RrtStar raise @linear.Refused {
  if step <= 0.0 {
    raise @linear.Refused::Gain(got=step)
  }
  if neighbourhood <= 0.0 {
    raise @linear.Refused::Gain(got=neighbourhood)
  }
  {
    space,
    validity,
    step,
    goal_bias,
    neighbourhood,
    rng: @rand.seeded(seed_bytes(seed)),
  }
}

///|
pub impl @contract.Planner for RrtStar with fn plan(self, start, goal, budget) {
  let tree = Tree::new(start)
  let mut best : Array[@linear.Vec]? = None
  let mut best_cost = 1.0e30
  for _ in 0.. self.neighbourhood {
        continue
      }
      let through = tree.costs[i] +
        self.space.distance(tree.nodes[i], candidate)
      if through < cost &&
        self.validity.motion(self.space, tree.nodes[i], candidate) {
        parent = i
        cost = through
      }
    }
    let index = tree.push(candidate, parent, cost)
    rewire(
      self.space,
      self.validity,
      tree,
      candidate,
      index,
      parent,
      cost,
      self.neighbourhood,
    )
    let to_goal = self.space.distance(candidate, goal)
    if cost + to_goal < best_cost &&
      self.validity.motion(self.space, candidate, goal) {
      let finish = tree.push(goal, index, cost + to_goal)
      best = Some(tree.path(finish))
      best_cost = cost + to_goal
    }
  }
  match best {
    Some(path) => path
    None =>
      raise @linear.Refused::Exceeded(
        limit=budget.iterations,
        got=budget.iterations,
      )
  }
}

///|
pub extend RrtStar with @contract.Planner::{plan}

///|
/// An RRT grown from both ends at once: the two trees take turns reaching for
/// each other, which closes a narrow passage far faster than one tree alone.
pub(all) struct RrtConnect {
  space : &@contract.Space
  validity : &@contract.Validity
  step : Double
  rng : &@rand_contract.Rng
}

///|
/// A bidirectional RRT.
pub fn RrtConnect::new(
  space~ : &@contract.Space,
  validity~ : &@contract.Validity,
  step~ : Double = 0.1,
  seed~ : Int32 = 0,
) -> RrtConnect raise @linear.Refused {
  if step <= 0.0 {
    raise @linear.Refused::Gain(got=step)
  }
  { space, validity, step, rng: @rand.seeded(seed_bytes(seed)), }
}

///|
/// Extend `tree` towards `target` one step; `Some` with the new index when the
/// step was allowed.
fn extend_once(
  tree : Tree,
  space : &@contract.Space,
  validity : &@contract.Validity,
  target : @linear.Vec,
  step : Double,
) -> Int32? {
  let nearest = tree.nearest(space, target)
  let candidate = steer(space, tree.nodes[nearest], target, step)
  if !validity.motion(space, tree.nodes[nearest], candidate) {
    return None
  }
  let cost = tree.costs[nearest] +
    space.distance(tree.nodes[nearest], candidate)
  Some(tree.push(candidate, nearest, cost))
}

///|
/// Extend `tree` towards `target` as far as it can go, one step at a time.
fn connect_towards(
  tree : Tree,
  space : &@contract.Space,
  validity : &@contract.Validity,
  target : @linear.Vec,
  step : Double,
) -> Int32? {
  let mut last : Int32? = None
  let mut reached = false
  while !reached {
    let nearest = tree.nearest(space, target)
    let distance = space.distance(tree.nodes[nearest], target)
    let candidate = if distance <= step {
      target.clone()
    } else {
      space.interpolate(tree.nodes[nearest], target, step / distance)
    }
    if !validity.motion(space, tree.nodes[nearest], candidate) {
      break
    }
    let cost = tree.costs[nearest] +
      space.distance(tree.nodes[nearest], candidate)
    last = Some(tree.push(candidate, nearest, cost))
    reached = distance <= step
  }
  last
}

///|
pub impl @contract.Planner for RrtConnect with fn plan(
  self,
  start,
  goal,
  budget,
) {
  let from_start = Tree::new(start)
  let from_goal = Tree::new(goal)
  for iteration in 0.. ()
      Some(index) =>
        match
          connect_towards(
            from_goal,
            self.space,
            self.validity,
            from_start.nodes[index],
            self.step,
          ) {
          None => ()
          Some(other) =>
            if self.space.distance(
                from_start.nodes[index],
                from_goal.nodes[other],
              ) <
              1.0e-9 {
              let head = from_start.path(index)
              let tail = from_goal.path(other)
              let joined = Array::make(head.length() + tail.length() - 1, start)
              for i in 0.. InformedRrtStar raise @linear.Refused {
  if step <= 0.0 {
    raise @linear.Refused::Gain(got=step)
  }
  if neighbourhood <= 0.0 {
    raise @linear.Refused::Gain(got=neighbourhood)
  }
  {
    space,
    validity,
    step,
    goal_bias,
    neighbourhood,
    rng: @rand.seeded(seed_bytes(seed)),
  }
}

///|
pub impl @contract.Planner for InformedRrtStar with fn plan(
  self,
  start,
  goal,
  budget,
) {
  let tree = Tree::new(start)
  let mut best : Array[@linear.Vec]? = None
  let mut best_cost = 1.0e30
  for _ in 0.. self.neighbourhood {
        continue
      }
      let through = tree.costs[i] +
        self.space.distance(tree.nodes[i], candidate)
      if through < cost &&
        self.validity.motion(self.space, tree.nodes[i], candidate) {
        parent = i
        cost = through
      }
    }
    let index = tree.push(candidate, parent, cost)
    rewire(
      self.space,
      self.validity,
      tree,
      candidate,
      index,
      parent,
      cost,
      self.neighbourhood,
    )
    let to_goal = self.space.distance(candidate, goal)
    if cost + to_goal < best_cost &&
      self.validity.motion(self.space, candidate, goal) {
      let finish = tree.push(goal, index, cost + to_goal)
      best = Some(tree.path(finish))
      best_cost = cost + to_goal
    }
  }
  match best {
    Some(path) => path
    None =>
      raise @linear.Refused::Exceeded(
        limit=budget.iterations,
        got=budget.iterations,
      )
  }
}

///|
pub extend InformedRrtStar with @contract.Planner::{plan}