001// Copyright (c) FIRST and other WPILib contributors.
002// Open Source Software; you can modify and/or share it under the terms of
003// the WPILib BSD license file in the root directory of this project.
004
005package edu.wpi.first.util;
006
007import java.util.TreeMap;
008
009/**
010 * Interpolating Tree Maps are used to get values at points that are not defined by making a guess
011 * from points that are defined. This uses linear interpolation.
012 */
013public class InterpolatingTreeMap<K extends Number, V extends Number> {
014  private final TreeMap<K, V> m_map = new TreeMap<>();
015
016  /**
017   * Inserts a key-value pair.
018   *
019   * @param key The key.
020   * @param value The value.
021   */
022  public void put(K key, V value) {
023    m_map.put(key, value);
024  }
025
026  /**
027   * Returns the value associated with a given key.
028   *
029   * <p>If there's no matching key, the value returned will be a linear interpolation between the
030   * keys before and after the provided one.
031   *
032   * @param key The key.
033   * @return The value associated with the given key.
034   */
035  public Double get(K key) {
036    V val = m_map.get(key);
037    if (val == null) {
038      K ceilingKey = m_map.ceilingKey(key);
039      K floorKey = m_map.floorKey(key);
040
041      if (ceilingKey == null && floorKey == null) {
042        return null;
043      }
044      if (ceilingKey == null) {
045        return m_map.get(floorKey).doubleValue();
046      }
047      if (floorKey == null) {
048        return m_map.get(ceilingKey).doubleValue();
049      }
050      V floor = m_map.get(floorKey);
051      V ceiling = m_map.get(ceilingKey);
052
053      return interpolate(floor, ceiling, inverseInterpolate(ceilingKey, key, floorKey));
054    } else {
055      return val.doubleValue();
056    }
057  }
058
059  /** Clears the contents. */
060  public void clear() {
061    m_map.clear();
062  }
063
064  /**
065   * Return the value interpolated between val1 and val2 by the interpolant d.
066   *
067   * @param val1 The lower part of the interpolation range.
068   * @param val2 The upper part of the interpolation range.
069   * @param d The interpolant in the range [0, 1].
070   * @return The interpolated value.
071   */
072  private double interpolate(V val1, V val2, double d) {
073    double dydx = val2.doubleValue() - val1.doubleValue();
074    return dydx * d + val1.doubleValue();
075  }
076
077  /**
078   * Return where within interpolation range [0, 1] q is between down and up.
079   *
080   * @param up Upper part of interpolation range.
081   * @param q Query.
082   * @param down Lower part of interpolation range.
083   * @return Interpolant in range [0, 1].
084   */
085  private double inverseInterpolate(K up, K q, K down) {
086    double upperToLower = up.doubleValue() - down.doubleValue();
087    if (upperToLower <= 0) {
088      return 0.0;
089    }
090    double queryToLower = q.doubleValue() - down.doubleValue();
091    if (queryToLower <= 0) {
092      return 0.0;
093    }
094    return queryToLower / upperToLower;
095  }
096}