#include "obstacle_scanner/tracking.hpp"

#include <algorithm>
#include <cmath>
#include <stdexcept>
#include <vector>

namespace obstacle_scanner
{
namespace
{

struct AssociationCandidate
{
  std::size_t observation_index{0};
  std::size_t track_index{0};
  Eigen::Vector2d center{Eigen::Vector2d::Zero()};
  double mahalanobis_squared{0.0};
  double euclidean_distance{0.0};
  bool geometry_valid{true};
};

Eigen::Matrix<double, 2, 4> position_matrix()
{
  Eigen::Matrix<double, 2, 4> matrix = Eigen::Matrix<double, 2, 4>::Zero();
  matrix(0, 0) = 1.0;
  matrix(1, 1) = 1.0;
  return matrix;
}

bool innovation_for_track(
  const Eigen::Vector4d & state,
  const Eigen::Matrix4d & covariance,
  const Eigen::Vector2d & center,
  double position_stddev,
  double & mahalanobis_squared,
  double & euclidean_distance)
{
  const auto h = position_matrix();
  const Eigen::Vector2d innovation = center - h * state;
  const Eigen::Matrix2d measurement_noise =
    position_stddev * position_stddev * Eigen::Matrix2d::Identity();
  const Eigen::Matrix2d innovation_covariance =
    h * covariance * h.transpose() + measurement_noise;
  Eigen::LDLT<Eigen::Matrix2d> ldlt(innovation_covariance);
  if (ldlt.info() != Eigen::Success) {
    return false;
  }
  const Eigen::Vector2d solved = ldlt.solve(innovation);
  if (!solved.allFinite()) {
    return false;
  }
  mahalanobis_squared = std::max(0.0, innovation.dot(solved));
  euclidean_distance = innovation.norm();
  return std::isfinite(mahalanobis_squared) && std::isfinite(euclidean_distance);
}

void validate_config(const TrackerConfig & config)
{
  const auto finite_positive = [](double value) {
      return std::isfinite(value) && value > 0.0;
    };
  const auto finite_nonnegative = [](double value) {
      return std::isfinite(value) && value >= 0.0;
    };

  const bool valid =
    config.track_confirm_fit_hits > 0 &&
    config.track_confirm_total_hits > 0 &&
    config.track_publish_misses >= 0 &&
    config.track_delete_misses > config.track_publish_misses &&
    finite_nonnegative(config.process_accel_noise) &&
    finite_positive(config.initial_velocity_stddev) &&
    finite_positive(config.fit_position_stddev) &&
    finite_positive(config.chord_position_stddev) &&
    finite_positive(config.mahalanobis_gate) &&
    finite_positive(config.max_association_distance) &&
    finite_positive(config.max_position_stddev) &&
    finite_nonnegative(config.chord_tolerance) &&
    finite_positive(config.nominal_dt) &&
    finite_positive(config.min_dt) &&
    finite_positive(config.max_dt) &&
    config.min_dt <= config.nominal_dt &&
    config.nominal_dt <= config.max_dt;
  if (!valid) {
    throw std::invalid_argument("invalid obstacle tracker configuration");
  }
}

}  // namespace

ObstacleTracker::ObstacleTracker(const TrackerConfig & config)
: config_(config)
{
  validate_config(config_);
}

FrameDebug ObstacleTracker::update(
  const std::vector<Observation> & observations,
  double stamp_seconds)
{
  FrameDebug debug;
  debug.fit_observations = static_cast<int>(std::count_if(
      observations.begin(), observations.end(),
      [](const Observation & observation) {
        return observation.type == ObservationType::CircleFit;
      }));
  debug.chord_observations = static_cast<int>(std::count_if(
      observations.begin(), observations.end(),
      [](const Observation & observation) {
        return observation.type == ObservationType::TwoPointChord;
      }));

  if (!std::isfinite(stamp_seconds)) {
    stamp_seconds = std::isfinite(last_stamp_seconds_) ?
      last_stamp_seconds_ + config_.nominal_dt : 0.0;
  }
  double dt = config_.nominal_dt;
  if (std::isfinite(last_stamp_seconds_)) {
    const double raw_dt = stamp_seconds - last_stamp_seconds_;
    if (std::isfinite(raw_dt) && raw_dt > 0.0) {
      dt = std::clamp(raw_dt, config_.min_dt, config_.max_dt);
    }
  }
  last_stamp_seconds_ = stamp_seconds;

  for (auto & track : tracks_) {
    track.matched = false;
  }
  predict_tracks(dt);

  double innovation_sum = 0.0;
  int innovation_count = 0;
  std::vector<bool> observation_matched(observations.size(), false);

  std::vector<AssociationCandidate> fit_candidates;
  fit_candidates.reserve(observations.size() * tracks_.size());
  for (std::size_t observation_index = 0;
    observation_index < observations.size(); ++observation_index)
  {
    const auto & observation = observations[observation_index];
    if (observation.type != ObservationType::CircleFit) {
      continue;
    }
    for (std::size_t track_index = 0; track_index < tracks_.size(); ++track_index) {
      const auto & track = tracks_[track_index];
      double mahalanobis_squared = 0.0;
      double euclidean_distance = 0.0;
      if (innovation_for_track(
          track.state, track.covariance, observation.center,
          config_.fit_position_stddev, mahalanobis_squared, euclidean_distance))
      {
        fit_candidates.push_back({
          observation_index, track_index, observation.center,
          mahalanobis_squared, euclidean_distance, true});
      }
    }
  }
  std::sort(
    fit_candidates.begin(), fit_candidates.end(),
    [](const AssociationCandidate & lhs, const AssociationCandidate & rhs) {
      return lhs.mahalanobis_squared < rhs.mahalanobis_squared;
    });

  for (const auto & candidate : fit_candidates) {
    if (observation_matched[candidate.observation_index] ||
      tracks_[candidate.track_index].matched)
    {
      continue;
    }
    if (candidate.mahalanobis_squared > config_.mahalanobis_gate ||
      candidate.euclidean_distance > config_.max_association_distance)
    {
      continue;
    }
    auto & track = tracks_[candidate.track_index];
    const auto & observation = observations[candidate.observation_index];
    track.fit_hits++;
    track.total_hits++;
    update_track(
      track, candidate.center, config_.fit_position_stddev, UpdateSource::Fit,
      debug, innovation_sum, innovation_count);
    track.radius = observation.radius;
    observation_matched[candidate.observation_index] = true;
    debug.fit_updates++;
  }

  for (std::size_t observation_index = 0;
    observation_index < observations.size(); ++observation_index)
  {
    const auto & observation = observations[observation_index];
    if (observation.type != ObservationType::CircleFit ||
      observation_matched[observation_index])
    {
      continue;
    }
    if (!tracks_.empty()) {
      debug.association_rejections++;
    }
    Track track;
    track.id = next_id_++;
    track.state.head<2>() = observation.center;
    track.covariance.diagonal() <<
      config_.fit_position_stddev * config_.fit_position_stddev,
      config_.fit_position_stddev * config_.fit_position_stddev,
      config_.initial_velocity_stddev * config_.initial_velocity_stddev,
      config_.initial_velocity_stddev * config_.initial_velocity_stddev;
    track.radius = observation.radius;
    track.fit_hits = 1;
    track.total_hits = 1;
    track.matched = true;
    track.last_source = UpdateSource::Fit;
    track.confirmed =
      track.fit_hits >= config_.track_confirm_fit_hits ||
      track.total_hits >= config_.track_confirm_total_hits;
    tracks_.push_back(track);
    observation_matched[observation_index] = true;
    debug.tracks_created++;
    debug.fit_updates++;
    if (track.confirmed) {
      debug.tracks_confirmed++;
    }
  }

  std::vector<AssociationCandidate> chord_candidates;
  chord_candidates.reserve(observations.size() * tracks_.size());
  std::vector<bool> chord_geometry_valid(observations.size(), false);
  for (std::size_t observation_index = 0;
    observation_index < observations.size(); ++observation_index)
  {
    const auto & observation = observations[observation_index];
    if (observation.type != ObservationType::TwoPointChord) {
      continue;
    }
    for (std::size_t track_index = 0; track_index < tracks_.size(); ++track_index) {
      const auto & track = tracks_[track_index];
      if (track.matched) {
        continue;
      }
      Eigen::Vector2d center = Eigen::Vector2d::Zero();
      if (!chord_center_for_track(observation, track, center)) {
        continue;
      }
      chord_geometry_valid[observation_index] = true;
      double mahalanobis_squared = 0.0;
      double euclidean_distance = 0.0;
      if (innovation_for_track(
          track.state, track.covariance, center,
          config_.chord_position_stddev, mahalanobis_squared, euclidean_distance))
      {
        chord_candidates.push_back({
          observation_index, track_index, center,
          mahalanobis_squared, euclidean_distance, true});
      }
    }
  }
  std::sort(
    chord_candidates.begin(), chord_candidates.end(),
    [](const AssociationCandidate & lhs, const AssociationCandidate & rhs) {
      return lhs.mahalanobis_squared < rhs.mahalanobis_squared;
    });

  std::vector<bool> chord_matched(observations.size(), false);
  for (const auto & candidate : chord_candidates) {
    if (chord_matched[candidate.observation_index] ||
      tracks_[candidate.track_index].matched)
    {
      continue;
    }
    if (candidate.mahalanobis_squared > config_.mahalanobis_gate ||
      candidate.euclidean_distance > config_.max_association_distance)
    {
      continue;
    }
    auto & track = tracks_[candidate.track_index];
    track.total_hits++;
    update_track(
      track, candidate.center, config_.chord_position_stddev, UpdateSource::Chord,
      debug, innovation_sum, innovation_count);
    chord_matched[candidate.observation_index] = true;
    debug.chord_updates++;
  }

  for (std::size_t observation_index = 0;
    observation_index < observations.size(); ++observation_index)
  {
    const auto & observation = observations[observation_index];
    if (observation.type != ObservationType::TwoPointChord ||
      chord_matched[observation_index])
    {
      continue;
    }
    bool has_in_gate_candidate = false;
    for (const auto & candidate : chord_candidates) {
      if (candidate.observation_index == observation_index &&
        candidate.mahalanobis_squared <= config_.mahalanobis_gate &&
        candidate.euclidean_distance <= config_.max_association_distance)
      {
        has_in_gate_candidate = true;
        break;
      }
    }
    if (chord_geometry_valid[observation_index] || has_in_gate_candidate) {
      debug.association_rejections++;
    } else if (!tracks_.empty()) {
      debug.chord_rejections++;
    }
  }

  for (auto & track : tracks_) {
    if (track.matched) {
      track.missed_count = 0;
      continue;
    }
    track.missed_count++;
    track.last_source = UpdateSource::Predict;
  }

  const auto old_size = tracks_.size();
  tracks_.erase(
    std::remove_if(
      tracks_.begin(), tracks_.end(),
      [&](const Track & track) {
        return (!track.confirmed && track.missed_count >= 1) ||
               (track.confirmed && track.missed_count >= config_.track_delete_misses);
      }),
    tracks_.end());
  debug.tracks_deleted = static_cast<int>(old_size - tracks_.size());

  refresh_debug_outputs();
  debug.active_tracks = static_cast<int>(tracks_.size());
  debug.published_tracks = static_cast<int>(confirmed_tracks().size());
  debug.predicted_published = static_cast<int>(std::count_if(
      tracks_.begin(), tracks_.end(),
      [&](const Track & track) {
        return !track.matched && is_publishable(track);
      }));
  if (innovation_count > 0) {
    debug.mean_innovation = innovation_sum / static_cast<double>(innovation_count);
  }
  return debug;
}

void ObstacleTracker::predict_tracks(double dt)
{
  Eigen::Matrix4d transition = Eigen::Matrix4d::Identity();
  transition(0, 2) = dt;
  transition(1, 3) = dt;
  Eigen::Matrix<double, 4, 2> gain;
  gain << 0.5 * dt * dt, 0.0,
    0.0, 0.5 * dt * dt,
    dt, 0.0,
    0.0, dt;
  const Eigen::Matrix4d process_noise =
    std::pow(config_.process_accel_noise, 2) * gain * gain.transpose();

  for (auto & track : tracks_) {
    track.state = transition * track.state;
    track.covariance =
      transition * track.covariance * transition.transpose() + process_noise;
    track.covariance =
      0.5 * (track.covariance + track.covariance.transpose());
  }
}

void ObstacleTracker::update_track(
  Track & track,
  const Eigen::Vector2d & center,
  double position_stddev,
  UpdateSource source,
  FrameDebug & debug,
  double & innovation_sum,
  int & innovation_count)
{
  const auto h = position_matrix();
  const Eigen::Matrix2d measurement_noise =
    position_stddev * position_stddev * Eigen::Matrix2d::Identity();
  const Eigen::Vector2d innovation = center - h * track.state;
  const Eigen::Matrix2d innovation_covariance =
    h * track.covariance * h.transpose() + measurement_noise;
  Eigen::LDLT<Eigen::Matrix2d> ldlt(innovation_covariance);
  if (ldlt.info() != Eigen::Success) {
    return;
  }
  const Eigen::Matrix<double, 4, 2> kalman_gain =
    track.covariance * h.transpose() * ldlt.solve(Eigen::Matrix2d::Identity());
  track.state += kalman_gain * innovation;
  const Eigen::Matrix4d identity = Eigen::Matrix4d::Identity();
  const Eigen::Matrix4d residual_projection = identity - kalman_gain * h;
  track.covariance =
    residual_projection * track.covariance * residual_projection.transpose() +
    kalman_gain * measurement_noise * kalman_gain.transpose();
  track.covariance =
    0.5 * (track.covariance + track.covariance.transpose());
  track.missed_count = 0;
  track.last_source = source;
  track.matched = true;
  innovation_sum += innovation.norm();
  innovation_count++;
  debug.max_innovation = std::max(debug.max_innovation, innovation.norm());
  if (
    !track.confirmed &&
    (track.fit_hits >= config_.track_confirm_fit_hits ||
    track.total_hits >= config_.track_confirm_total_hits))
  {
    track.confirmed = true;
    debug.tracks_confirmed++;
  }
}

std::vector<TrackOutput> ObstacleTracker::confirmed_tracks() const
{
  std::vector<TrackOutput> outputs;
  for (const auto & track : tracks_) {
    if (is_publishable(track)) {
      outputs.push_back(to_output(track));
    }
  }
  return outputs;
}

const std::vector<TrackOutput> & ObstacleTracker::tracks_for_debug() const
{
  return debug_outputs_;
}

bool ObstacleTracker::chord_center_for_track(
  const Observation & observation,
  const Track & track,
  Eigen::Vector2d & center) const
{
  const Eigen::Vector2d chord = observation.p2 - observation.p1;
  const double chord_length = chord.norm();
  if (chord_length <= 1e-9) {
    return false;
  }
  if (chord_length > 2.0 * track.radius + config_.chord_tolerance) {
    return false;
  }

  const double half_chord = 0.5 * chord_length;
  const double h_squared = track.radius * track.radius - half_chord * half_chord;
  if (h_squared < 0.0) {
    return false;
  }

  const Eigen::Vector2d middle = 0.5 * (observation.p1 + observation.p2);
  const Eigen::Vector2d direction = chord / chord_length;
  const Eigen::Vector2d normal(-direction.y(), direction.x());
  const double h = std::sqrt(h_squared);
  const Eigen::Vector2d candidate_a = middle + h * normal;
  const Eigen::Vector2d candidate_b = middle - h * normal;
  center = candidate_a.norm() >= candidate_b.norm() ? candidate_a : candidate_b;
  return true;
}

bool ObstacleTracker::is_publishable(const Track & track) const
{
  const double position_stddev = std::sqrt(std::max(
      track.covariance(0, 0), track.covariance(1, 1)));
  return track.confirmed &&
         track.missed_count <= config_.track_publish_misses &&
         std::isfinite(position_stddev) &&
         position_stddev <= config_.max_position_stddev;
}

TrackOutput ObstacleTracker::to_output(const Track & track)
{
  TrackOutput output;
  output.id = track.id;
  output.center_x = track.state.x();
  output.center_y = track.state.y();
  output.radius = track.radius;
  output.hit_count = track.total_hits;
  output.missed_count = track.missed_count;
  output.last_source = track.last_source;
  return output;
}

void ObstacleTracker::refresh_debug_outputs()
{
  debug_outputs_.clear();
  debug_outputs_.reserve(tracks_.size());
  for (const auto & track : tracks_) {
    debug_outputs_.push_back(to_output(track));
  }
}

std::string update_source_name(UpdateSource source)
{
  switch (source) {
    case UpdateSource::Fit:
      return "fit";
    case UpdateSource::Chord:
      return "chord";
    case UpdateSource::Predict:
      return "predict";
  }
  return "unknown";
}

}  // namespace obstacle_scanner
