#include "SlowYoGenerator.h"

#include <algorithm>
#include <cmath>
#include <functional>
#include <iostream>
#include <numeric>
#include <vector>


// based on
// https://stackoverflow.com/a/27030598/9686600
std::vector<double> linspace(const double & start,
                             const double & end,
                             const std::size_t & num)
{
  std::vector<double> linspaced;

  if (num == 0U) {
    return linspaced;
  }

  if (num == 1U) {
    linspaced.push_back(start);
    return linspaced;
  }

  const double delta = (end - start) / (num - 1U);

  for (std::size_t idx=0U; idx < num - 1U; ++idx) {
    linspaced.push_back(start + delta * idx);
  }
  linspaced.push_back(end); // I want to ensure that start and end
                            // are exactly the same as the input
  return linspaced;
}


std::vector<double> arange(const double & start,
  const double & stop,
  const double & step)
{
  std::vector<double> v;
  for(std::size_t idx=0U; idx*step<stop; ++idx) {
    v.push_back(start + idx * step);
  }
  return v;
}


void print_vector(const std::vector<double> & v)
{
  std::cout << "size: " << v.size() << std::endl;
  for (std::size_t idx=0U; idx<v.size(); ++idx) {
    std::cout << v[idx] << " ";
  }
  std::cout << std::endl;
}


void minmax(const std::vector<double> & v,
  double & mn, std::size_t & argmn,
  double & mx, std::size_t & argmx)
{
  mn = 1e8;
  argmn = 0U;
  mx = -1e8;
  argmx = 0U;
  for (std::size_t idx=0U; idx<v.size(); ++idx) {
    if (v[idx] < mn) {
      mn = v[idx];
      argmn = idx;
    }
    if (v[idx] > mx) {
      mx = v[idx];
      argmx = idx;
    }
  }
}


void scale(std::vector<double> &v, const double & s)
{
  for(std::size_t idx=0U; idx < v.size(); ++idx) {
    v[idx] *= s;
  }
}


std::vector<double> vector_add(const std::vector<double> & v1,
  const std::vector<double> & v2)
{
  std::vector<double> v; v.reserve(v1.size());
  for (std::size_t idx=0U; idx<v1.size(); ++idx) {
    v.push_back(v1[idx] + v2[idx]);
  }
  return v;
}


void bias(std::vector<double> & v, const double & c)
{
  for (std::size_t idx=0U; idx<v.size(); ++idx) {
    v[idx] += c;
  }
}


std::vector<double> diff(const std::vector<double> & v)
{
  std::vector<double> diffed; diffed.reserve(v.size() - 1U);
  for (std::size_t idx=0U; idx<v.size()-1U; ++idx) {
    diffed.push_back(v[idx+1] - v[idx]);
  }
  return diffed;
}


std::vector<double> cumulative_sum(const std::vector<double> & v)
{
  std::vector<double> c_sum; c_sum.resize(v.size());
  std::partial_sum(v.begin(), v.end(), c_sum.begin());
  return c_sum;
}


struct Pt
{
  double x;
  double y;
  Pt(const double & x, const double & y) : x(x), y(y) {}

  Pt operator-(const Pt & that) const
  {
    return Pt(this->x - that.x, this->y - that.y);
  }

  Pt operator+(const Pt & that) const
  {
    return Pt(this->x + that.x, this->y + that.y);
  }

  Pt operator*(const double & s) const
  {
    return Pt(s * this->x, s * this->y);
  }
};


SlowYoGenerator::SlowYoGenerator(
  const double & _min_depth,  // meters
  const double & _max_depth,  // meters
  const double & _dive_angle,  // degrees
  const double & _speed,  // meters per second
  const double & _alpha)  // smoothing factor:
                         //   max: 1.0==sinusoidal
                         //   min: 0.0==trapezoidal
  : min_depth(_min_depth),
    max_depth(_max_depth),
    dive_angle(_dive_angle),
    speed(_speed),
    alpha(std::min(std::max(0.0, _alpha), 1.0))
{
  /*
  std::cout << "min depth: [" << min_depth
            << "] -- max depth: [" << max_depth
            << "] -- dive angle: [" << dive_angle
            << "] -- speed: [" << speed
            << "] -- alpha: [" << alpha << "]" << std::endl;
  */
  const double CMD_RATE = 2.5;  // Hz
  const double RAD_PER_DEG = M_PI / 180.0;
  // sin(pitch) = opposite / hypotenuse = depth rate / path speed
  const double depth_rate = speed * sin(RAD_PER_DEG * dive_angle);  // z-rate
  const double T_tr = 2.0 * (max_depth - min_depth) / depth_rate;
  // follows from max depth rate (related to max pitch) of cosine depth profile
  // when derivative (depth rate) is max at period T/4 and solving for T
  const double T_sn = M_PI * (max_depth - min_depth) / depth_rate;
  // blend the periods with smooth factor
  slowyo.T = T_tr * (1.0 - alpha) + T_sn * alpha;
  slowyo.T *= 1.1;  // fudge for control delay
  slowyo.dt = 1.0 / CMD_RATE;
  slowyo.t = arange(0.0, slowyo.T, 1.0 / CMD_RATE);

  // use slightly longer arrays to help smooth period wrapping
  std::vector<double> tt = arange(0.0, 1.2 * slowyo.T, 1.0 / CMD_RATE);
  std::vector<double> sn; sn.reserve(tt.size());  // sinusoid
  std::vector<double> tr; tr.reserve(tt.size());  // triangle
  for (std::size_t idx=0U; idx<tt.size(); ++idx) {
    double v = sin(2.0 * M_PI * tt[idx] / slowyo.T);
    sn.push_back(v);
    tr.push_back(asin(v));
  }

  // find peaks
  double sn_min = 0.0;
  std::size_t sn_argmin = 0U;
  double sn_max = 0.0;
  std::size_t sn_argmax = 0U;
  minmax(sn, sn_min, sn_argmin, sn_max, sn_argmax);

  // find peaks
  double tr_min = 0.0;
  std::size_t tr_argmin = 0U;
  double tr_max = 0.0;
  std::size_t tr_argmax = 0U;
  minmax(tr, tr_min, tr_argmin, tr_max, tr_argmax);
  scale(tr, 1.0 / tr_max);  // normalize triangle

  // blend sinusoid + triangle with smoothing factor
  std::vector<double> tr_alpha(tr.begin(), tr.end());
  scale(tr_alpha, 1.0 - alpha);
  std::vector<double> sn_alpha(sn.begin(), sn.end());
  scale(sn_alpha, alpha);
  std::vector<double> y = vector_add(tr_alpha, sn_alpha);

  // smooth peaks of profile using bezier curve
  const double bezier_window = static_cast<std::size_t>(y.size() / 60.0);
  std::vector<std::size_t> sn_idx; sn_idx.reserve(2U);
  sn_idx.push_back(sn_argmax);
  sn_idx.push_back(sn_argmin);
  for (std::size_t idx=0; idx<2U; ++idx) {  // for each peak
    // bezier control points
    Pt P0(tt[sn_idx[idx] - bezier_window], y[sn_idx[idx] - bezier_window]);
    Pt P1(tt[sn_idx[idx]], y[sn_idx[idx]]);
    Pt P2(tt[sn_idx[idx] + bezier_window], y[sn_idx[idx] + bezier_window]);
    // use as many points as window size to replace in y
    std::vector<double> b = linspace(0.0, 1.0, 2U*bezier_window + 1U);
    std::vector<Pt> B; B.reserve(b.size());
    for (std::size_t jdx=0U; jdx<b.size(); ++jdx) {
      // points along bezier curve
      B.push_back(P1 + (P0 - P1) * (1 - b[jdx]) * (1 - b[jdx]) + (P2 - P1) * b[jdx] * b[jdx]);
    }

    std::size_t window_idx = sn_idx[idx] - bezier_window;
    // std::size_t window_stop = sn_idx[idx] + bezier_window + 1;
    for (std::size_t kdx=0U; kdx<B.size(); ++kdx) {
      y[window_idx++] = B[kdx].y;  // replace y with smooth points
    }
  }

  // find peaks
  double y_min = 0.0;
  std::size_t y_argmin = 0U;
  double y_max = 0.0;
  std::size_t y_argmax = 0U;
  minmax(y, y_min, y_argmin, y_max, y_argmax);

  // scale and offset to satisfy depth envelope
  scale(y, 0.5*(max_depth - min_depth) / y_max);
  bias(y, 0.5*(max_depth - min_depth));

  // simple derivative to get depth rate profile
  y = diff(y);
  // shift so peak of original (derivative = 0) is at t=0
  std::rotate(y.begin(), y.begin() + sn_argmax, y.end());
  // take half the period and flip it (so we start by diving instead of surfacing at t=0)
  std::vector<double> out(y.begin(),
                          y.begin() + 2U*sn_argmax + 1U);
  scale(out, -1.0);
  slowyo.depth_rate_profile = out;
  // flip the half-period back so we surface again and append to the other half-period
  // so we get a full depth rate profile
  scale(out, -1.0);
  slowyo.depth_rate_profile.insert(slowyo.depth_rate_profile.end(),
                                   out.begin(), out.end());
  // integrate to get depth profile
  slowyo.depth_profile = cumulative_sum(slowyo.depth_rate_profile);
  // start from min_depth
  bias(slowyo.depth_profile, min_depth);

  // actual depth rate in meters per second not meters per sample (@ CMD_RATE)
  scale(slowyo.depth_rate_profile, 1.0 / slowyo.dt);

  /*
  double dr_min = 0.0;
  std::size_t dr_argmin = 0U;
  double dr_max = 0.0;
  std::size_t dr_argmax = 0U;
  minmax(slowyo.depth_rate_profile, dr_min, dr_argmin, dr_max, dr_argmax);
  std::cout << "max depth rate=[" << dr_max << "]" << std::endl;
  */

  // pitch = asin(opposite / hypotenuse) = asin(depth rate / path speed)
  slowyo.pitch_profile = std::vector<double>(slowyo.depth_rate_profile.begin(),
                                             slowyo.depth_rate_profile.end());
  scale(slowyo.pitch_profile, 1.0 / speed);
  std::transform(slowyo.pitch_profile.begin(),
                 slowyo.pitch_profile.end(),
                 slowyo.pitch_profile.begin(),
                 asin);

  /*
  double p_min = 0.0;
  std::size_t p_argmin = 0U;
  double p_max = 0.0;
  std::size_t p_argmax = 0U;
  minmax(slowyo.pitch_profile, p_min, p_argmin, p_max, p_argmax);
  std::cout << "max pitch=[" << p_max / RAD_PER_DEG << "]" << std::endl;
  */
}