// AZ-332 — pybind11 binding for OKVIS2 (production-default C1 VIO).
//
// Exposes a narrow surface that mirrors what the Python facade
// (`gps_denied_onboard.components.c1_vio.okvis2.Okvis2Strategy`)
// needs — NOT the full OKVIS2 estimator API. The surface is:
//
//   Okvis2Backend
//     ctor(yaml_config: str, camera_intrinsics_3x3: ndarray[float64, 3, 3])
//     add_frame(frame_id: str, ts_ns: int, image: ndarray[uint8, H, W, C]) -> bool
//     add_imu(ts_ns: int, accel: ndarray[float64, 3], gyro: ndarray[float64, 3]) -> None
//     get_latest_output() -> dict | None
//     reset(body_T_world: ndarray[float64, 4, 4], velocity: ndarray[float64, 3],
//           accel_bias: ndarray[float64, 3], gyro_bias: ndarray[float64, 3]) -> None
//     health() -> dict
//
// Frame buffers cross the FFI boundary as `py::array_t<uint8_t,
// c_style|forcecast>` so the camera-ingest path (AZ-265
// LiveCameraFrameSource) can hand off a contiguous numpy array without a
// copy — Risk-2 mitigation per the AZ-332 task spec.
//
// Exception envelope: every OKVIS2 / Eigen / std::runtime_error inside a
// binding method is caught and rethrown as one of three Python-side
// exceptions registered via `py::register_exception`. The Python facade
// then rewraps those into the VioError family.

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>

#include <Eigen/Core>
#include <Eigen/Geometry>

#include <array>
#include <cstdint>
#include <memory>
#include <mutex>
#include <optional>
#include <stdexcept>
#include <string>

// OKVIS2 estimator headers. The exact include path is determined by the
// vendored upstream's CMake export. The skeleton compiles without these
// headers because the actual ThreadedKFVio wiring lives in
// _build_estimator() / _drive_estimator(), which today STUB and surface a
// runtime error if invoked. Wiring them in is the follow-up task within
// AZ-332's tier2 deliverable bundle.
//
// #include <okvis/ThreadedKFVio.hpp>
// #include <okvis/Estimator.hpp>
// #include <okvis/VioParametersReader.hpp>

namespace py = pybind11;

namespace {

// ---------------------------------------------------------------------------
// Exception types — registered as Python-side classes via
// `py::register_exception` in PYBIND11_MODULE below. The Python facade
// catches these and rewraps into the VioError family.

class OkvisInitException : public std::runtime_error {
 public:
  using std::runtime_error::runtime_error;
};

class OkvisFatalException : public std::runtime_error {
 public:
  using std::runtime_error::runtime_error;
};

class OkvisOptimizationException : public std::runtime_error {
 public:
  using std::runtime_error::runtime_error;
};

// ---------------------------------------------------------------------------
// Pose / output struct produced by the estimator step.
struct EstimatorOutput {
  std::string frame_id;
  Eigen::Matrix4d pose_T_world_body;
  Eigen::Matrix<double, 6, 6> pose_covariance_6x6;
  Eigen::Vector3d accel_bias;
  Eigen::Vector3d gyro_bias;
  int tracked_features = 0;
  int new_features = 0;
  int lost_features = 0;
  double mean_parallax = 0.0;
  double mre_px = 0.0;
  std::int64_t emitted_at_ns = 0;
};

// ---------------------------------------------------------------------------
// Internal estimator state machine — INIT until N keyframes converge,
// TRACKING during nominal operation, DEGRADED on feature-count drop,
// LOST after consecutive failed updates.
enum class HealthState : int { Init = 0, Tracking = 1, Degraded = 2, Lost = 3 };

const char* state_to_str(HealthState s) {
  switch (s) {
    case HealthState::Init:
      return "init";
    case HealthState::Tracking:
      return "tracking";
    case HealthState::Degraded:
      return "degraded";
    case HealthState::Lost:
      return "lost";
  }
  return "init";
}

// ---------------------------------------------------------------------------
// Okvis2Backend — the C++ surface exposed to Python.
class Okvis2Backend {
 public:
  Okvis2Backend(const std::string& yaml_config,
                py::array_t<double, py::array::c_style | py::array::forcecast>
                    camera_intrinsics_3x3)
      : yaml_config_(yaml_config) {
    if (camera_intrinsics_3x3.ndim() != 2 ||
        camera_intrinsics_3x3.shape(0) != 3 ||
        camera_intrinsics_3x3.shape(1) != 3) {
      throw OkvisInitException(
          "Okvis2Backend: camera_intrinsics_3x3 must be a 3x3 float64 array");
    }
    auto buf = camera_intrinsics_3x3.unchecked<2>();
    for (py::ssize_t i = 0; i < 3; ++i) {
      for (py::ssize_t j = 0; j < 3; ++j) {
        K_(i, j) = buf(i, j);
      }
    }
    _build_estimator();
  }

  // Push a nav-camera frame into the estimator.
  // Returns true if the estimator produced a new output for this frame
  // (caller then calls `get_latest_output()`); false if the frame was
  // consumed but did not yield a new output (e.g. dropped as non-keyframe).
  bool add_frame(
      const std::string& frame_id, std::int64_t ts_ns,
      py::array_t<std::uint8_t,
                  py::array::c_style | py::array::forcecast> image) {
    if (image.ndim() < 2 || image.ndim() > 3) {
      throw OkvisOptimizationException(
          "Okvis2Backend.add_frame: image must be 2-D (grayscale) or 3-D (HxWxC)");
    }
    pending_frame_id_ = frame_id;
    pending_ts_ns_ = ts_ns;
    return _drive_estimator(image);
  }

  void add_imu(std::int64_t ts_ns,
               py::array_t<double,
                           py::array::c_style | py::array::forcecast> accel,
               py::array_t<double,
                           py::array::c_style | py::array::forcecast> gyro) {
    if (accel.size() != 3 || gyro.size() != 3) {
      throw OkvisOptimizationException(
          "Okvis2Backend.add_imu: accel and gyro must be length-3 float64 arrays");
    }
    if (ts_ns <= last_imu_ts_ns_) {
      throw OkvisOptimizationException(
          "Okvis2Backend.add_imu: ts_ns must be strict-monotonic");
    }
    last_imu_ts_ns_ = ts_ns;
    // Real OKVIS2 IMU push lands here once the estimator is wired in.
    // For the skeleton we just record the most recent sample — the
    // estimator's IMU integration is performed inside ThreadedKFVio.
    auto a = accel.unchecked<1>();
    auto g = gyro.unchecked<1>();
    last_accel_ = Eigen::Vector3d(a(0), a(1), a(2));
    last_gyro_ = Eigen::Vector3d(g(0), g(1), g(2));
  }

  std::optional<py::dict> get_latest_output() const {
    std::lock_guard<std::mutex> lk(output_mtx_);
    if (!latest_output_.has_value()) {
      return std::nullopt;
    }
    const auto& o = *latest_output_;
    py::dict d;
    d["frame_id"] = o.frame_id;
    d["pose_T_world_body"] = py::array_t<double>(
        {4, 4}, {sizeof(double) * 4, sizeof(double)},
        o.pose_T_world_body.data());
    d["pose_covariance_6x6"] = py::array_t<double>(
        {6, 6}, {sizeof(double) * 6, sizeof(double)},
        o.pose_covariance_6x6.data());
    d["accel_bias"] = py::array_t<double>(
        {3}, {sizeof(double)}, o.accel_bias.data());
    d["gyro_bias"] = py::array_t<double>(
        {3}, {sizeof(double)}, o.gyro_bias.data());
    d["tracked_features"] = o.tracked_features;
    d["new_features"] = o.new_features;
    d["lost_features"] = o.lost_features;
    d["mean_parallax"] = o.mean_parallax;
    d["mre_px"] = o.mre_px;
    d["emitted_at_ns"] = o.emitted_at_ns;
    return d;
  }

  void reset(py::array_t<double,
                         py::array::c_style | py::array::forcecast> body_T_world,
             py::array_t<double,
                         py::array::c_style | py::array::forcecast> velocity,
             py::array_t<double,
                         py::array::c_style | py::array::forcecast> accel_bias,
             py::array_t<double,
                         py::array::c_style | py::array::forcecast> gyro_bias) {
    if (body_T_world.ndim() != 2 || body_T_world.shape(0) != 4 ||
        body_T_world.shape(1) != 4) {
      throw OkvisInitException(
          "Okvis2Backend.reset: body_T_world must be a 4x4 float64 array");
    }
    if (velocity.size() != 3 || accel_bias.size() != 3 || gyro_bias.size() != 3) {
      throw OkvisInitException(
          "Okvis2Backend.reset: velocity / *_bias must be length-3 float64 arrays");
    }
    auto T = body_T_world.unchecked<2>();
    for (py::ssize_t i = 0; i < 4; ++i) {
      for (py::ssize_t j = 0; j < 4; ++j) {
        seed_body_T_world_(i, j) = T(i, j);
      }
    }
    auto v = velocity.unchecked<1>();
    auto ab = accel_bias.unchecked<1>();
    auto gb = gyro_bias.unchecked<1>();
    seed_velocity_ = Eigen::Vector3d(v(0), v(1), v(2));
    seed_accel_bias_ = Eigen::Vector3d(ab(0), ab(1), ab(2));
    seed_gyro_bias_ = Eigen::Vector3d(gb(0), gb(1), gb(2));

    state_ = HealthState::Init;
    consecutive_lost_ = 0;
    {
      std::lock_guard<std::mutex> lk(output_mtx_);
      latest_output_.reset();
    }
    _build_estimator();
  }

  py::dict health() const {
    py::dict d;
    d["state"] = std::string(state_to_str(state_));
    d["consecutive_lost"] = consecutive_lost_;
    d["bias_norm"] = std::sqrt(
        seed_accel_bias_.squaredNorm() + seed_gyro_bias_.squaredNorm());
    return d;
  }

 private:
  void _build_estimator() {
    // Real wiring: instantiate okvis::ThreadedKFVio from yaml_config_,
    // attach output callback that fills latest_output_ under output_mtx_.
    //
    // The skeleton intentionally throws on any actual frame ingest so a
    // production binary that loads this binding before AZ-332's
    // estimator wiring lands cannot silently report misleading poses.
    estimator_built_ = false;
  }

  bool _drive_estimator(
      py::array_t<std::uint8_t,
                  py::array::c_style | py::array::forcecast> /*image*/) {
    if (!estimator_built_) {
      // Skeleton path — pybind11 binding compiles and loads but the
      // OKVIS2 estimator is not yet wired. Tier-2 follow-up wires it up.
      throw OkvisFatalException(
          "Okvis2Backend: OKVIS2 estimator not yet wired — this binding "
          "is the AZ-332 skeleton; tier2 follow-up wires okvis::ThreadedKFVio");
    }
    return false;
  }

  std::string yaml_config_;
  Eigen::Matrix3d K_ = Eigen::Matrix3d::Identity();
  Eigen::Matrix4d seed_body_T_world_ = Eigen::Matrix4d::Identity();
  Eigen::Vector3d seed_velocity_ = Eigen::Vector3d::Zero();
  Eigen::Vector3d seed_accel_bias_ = Eigen::Vector3d::Zero();
  Eigen::Vector3d seed_gyro_bias_ = Eigen::Vector3d::Zero();
  Eigen::Vector3d last_accel_ = Eigen::Vector3d::Zero();
  Eigen::Vector3d last_gyro_ = Eigen::Vector3d::Zero();

  HealthState state_ = HealthState::Init;
  int consecutive_lost_ = 0;
  std::int64_t last_imu_ts_ns_ = -1;
  std::string pending_frame_id_;
  std::int64_t pending_ts_ns_ = 0;
  bool estimator_built_ = false;

  mutable std::mutex output_mtx_;
  std::optional<EstimatorOutput> latest_output_;
};

}  // namespace

PYBIND11_MODULE(okvis2_binding, m) {
  m.doc() =
      "OKVIS2 pybind11 binding (AZ-332). Wraps okvis::ThreadedKFVio for the "
      "Python Okvis2Strategy facade. Tier2 follow-up wires the real estimator.";

  py::register_exception<OkvisInitException>(m, "OkvisInitException");
  py::register_exception<OkvisFatalException>(m, "OkvisFatalException");
  py::register_exception<OkvisOptimizationException>(
      m, "OkvisOptimizationException");

  py::class_<Okvis2Backend>(m, "Okvis2Backend")
      .def(py::init<const std::string&,
                    py::array_t<double, py::array::c_style | py::array::forcecast>>(),
           py::arg("yaml_config"), py::arg("camera_intrinsics_3x3"))
      .def("add_frame", &Okvis2Backend::add_frame, py::arg("frame_id"),
           py::arg("ts_ns"), py::arg("image"))
      .def("add_imu", &Okvis2Backend::add_imu, py::arg("ts_ns"),
           py::arg("accel"), py::arg("gyro"))
      .def("get_latest_output", &Okvis2Backend::get_latest_output)
      .def("reset", &Okvis2Backend::reset, py::arg("body_T_world"),
           py::arg("velocity"), py::arg("accel_bias"), py::arg("gyro_bias"))
      .def("health", &Okvis2Backend::health);
}
