// nanobind bindings for RAID. Arrays are passed as NumPy views (no copy for contiguous float64/int64), the GIL is
// released around every heavy call, and results come back as dicts (JSON-like) plus NumPy arrays for vectors.
#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/vector.h>
#include <nanobind/stl/optional.h>
#include <nlohmann/json.hpp>
#include "raid/core/types.hpp"
#include "raid/core/error.hpp"
#include "raid/backend/backend.hpp"
#include "raid/estimators/ols.hpp"
#include "raid/estimators/realized.hpp"
#include "raid/models/garch.hpp"
#include "raid/estimators/har.hpp"
#include "raid/statistics/bootstrap.hpp"
#include "raid/statistics/reductions.hpp"
#include "raid/data/bar_series.hpp"
#include "raid/data/calendar.hpp"
#include "raid/data/returns.hpp"
#include "raid/data/findata.hpp"
#include "raid/data/io.hpp"

namespace nb = nanobind;
using namespace raid;
using json = nlohmann::json;
using F64 = nb::ndarray<const double, nb::c_contig, nb::device::cpu>;
using I64 = nb::ndarray<const int64_t, nb::c_contig, nb::device::cpu>;

namespace {
nb::object json_to_py(const json& j) {
  switch (j.type()) {
    case json::value_t::null: return nb::none();
    case json::value_t::boolean: return nb::bool_(j.get<bool>());
    case json::value_t::number_integer: return nb::int_(j.get<long long>());
    case json::value_t::number_unsigned: return nb::int_(j.get<unsigned long long>());
    case json::value_t::number_float: return nb::float_(j.get<double>());
    case json::value_t::string: return nb::str(j.get<std::string>().c_str());
    case json::value_t::array: { nb::list l; for (auto& v : j) l.append(json_to_py(v)); return l; }
    case json::value_t::object: { nb::dict d; for (auto& [k, v] : j.items()) d[k.c_str()] = json_to_py(v); return d; }
    default: return nb::none();
  }
}
nb::ndarray<nb::numpy, double> to_np(std::vector<double> v) {
  auto* data = new std::vector<double>(std::move(v));
  nb::capsule owner(data, [](void* p) noexcept { delete static_cast<std::vector<double>*>(p); });
  return nb::ndarray<nb::numpy, double>(data->data(), {data->size()}, owner);
}
nb::ndarray<nb::numpy, int64_t> to_np_i64(std::vector<int64_t> v) {
  auto* data = new std::vector<int64_t>(std::move(v));
  nb::capsule owner(data, [](void* p) noexcept { delete static_cast<std::vector<int64_t>*>(p); });
  return nb::ndarray<nb::numpy, int64_t>(data->data(), {data->size()}, owner);
}
}  // namespace

NB_MODULE(_core, m) {
  m.doc() = "RAID compiled core (C++20, CPU + optional Metal)";
  m.attr("__version__") = RAID_VERSION;
  nb::exception<RaidError>(m, "RaidError");

  m.def("devices", []() { json out = json::array(); for (auto& d : backend::list_devices()) out.push_back(d.to_json()); return json_to_py(out); });

  m.def("ols", [](F64 X, F64 y, const std::string& cov, int hac_lags, std::optional<I64> clusters, double ridge, std::vector<usize> unpen, std::optional<F64> weights, std::vector<std::string> names) {
    RAID_REQUIRE(X.ndim() == 2 && y.ndim() == 1 && X.shape(0) == y.shape(0), invalid_argument, "X (n,k), y (n,)");
    const Eigen::Index n = static_cast<Eigen::Index>(X.shape(0)), k = static_cast<Eigen::Index>(X.shape(1));
    Eigen::Map<const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> Xm(X.data(), n, k);
    Eigen::Map<const Eigen::VectorXd> ym(y.data(), n);
    est::OlsOptions o; o.cov.type = stats::covtype_from_string(cov); o.cov.hac_lags = hac_lags; o.ridge_lambda = ridge; o.unpenalized = unpen; o.names = names;
    std::vector<i64> cl; if (clusters) { cl.assign(clusters->data(), clusters->data() + clusters->shape(0)); o.cov.clusters = cl; }
    std::vector<f64> w; if (weights) { w.assign(weights->data(), weights->data() + weights->shape(0)); o.weights = w; }
    Eigen::MatrixXd Xd = Xm; json out; { nb::gil_scoped_release rel; auto r = est::ols(Xd, ym, o); out = r.estimate.to_json(); }
    return json_to_py(out);
  }, nb::arg("X"), nb::arg("y"), nb::arg("cov") = "classical", nb::arg("hac_lags") = -1, nb::arg("clusters") = nb::none(), nb::arg("ridge_lambda") = 0.0, nb::arg("unpenalized") = std::vector<usize>{}, nb::arg("weights") = nb::none(), nb::arg("names") = std::vector<std::string>{});

  m.def("realized_measures", [](F64 r, I64 off, const std::string& backend_s, bool bpv_fs, bool deterministic) {
    est::RealizedOptions o; o.plan.backend = backend_from_string(backend_s); o.bpv_finite_sample = bpv_fs; o.plan.deterministic = deterministic;
    backend::Batch b{std::span<const f64>(r.data(), r.shape(0)), std::span<const i64>(off.data(), off.shape(0))};
    est::RealizedMeasures mres; { nb::gil_scoped_release rel; mres = est::realized_measures_batch(b, o); }
    nb::dict d; d["rv"] = to_np(mres.rv); d["bpv"] = to_np(mres.bpv); d["rsv_neg"] = to_np(mres.rsv_neg); d["rsv_pos"] = to_np(mres.rsv_pos); d["rq"] = to_np(mres.rq); d["jump"] = to_np(mres.jump); d["tripv"] = to_np(mres.tripv); d["ret"] = to_np(mres.daily_return);
    d["backend"] = std::string(to_string(mres.stats.backend)); d["precision"] = std::string(to_string(mres.stats.precision)); d["wall_seconds"] = mres.stats.wall_seconds; d["device_seconds"] = mres.stats.device_seconds; d["note"] = mres.stats.note; d["warnings"] = mres.warnings;
    return d;
  }, nb::arg("returns"), nb::arg("offsets"), nb::arg("backend") = "auto", nb::arg("bpv_finite_sample") = true, nb::arg("deterministic") = true);

  m.def("garch_fit", [](F64 r, const std::string& model, const std::string& dist, int horizon, const std::string& cov, int max_iter, double tol) {
    models::GarchFitOptions o; o.model = model == "gjr" ? models::GarchModel::gjr : (model == "egarch" ? models::GarchModel::egarch : models::GarchModel::garch);
    o.dist = dist == "student_t" || dist == "t" ? models::GarchDist::student_t : models::GarchDist::gaussian; o.forecast_horizon = horizon; o.cov = cov; o.max_iter = max_iter; o.tol = tol;
    std::span<const f64> rs(r.data(), r.shape(0)); json out; { nb::gil_scoped_release rel; out = models::garch_fit(rs, o).to_json(); }
    return json_to_py(out);
  }, nb::arg("returns"), nb::arg("model") = "garch", nb::arg("dist") = "gaussian", nb::arg("horizon") = 0, nb::arg("cov") = "qmle", nb::arg("max_iter") = 500, nb::arg("tol") = 1e-8);

  m.def("validate_bars", [](I64 t, F64 o, F64 h, F64 l, F64 c, F64 v, const std::string& symbol, const std::string& interval, const std::string& session) {
    BarSeries b; b.symbol = symbol; b.interval = interval_from_string(interval); const usize n = t.shape(0); b.reserve(n);
    for (usize i = 0; i < n; ++i) b.push_back(Timestamp(t.data()[i]), o.data()[i], h.data()[i], l.data()[i], c.data()[i], v.data()[i]);
    json out; { nb::gil_scoped_release rel; auto st = b.normalize(); MarketCalendar cal; auto q = b.validate(cal, session == "extended" ? SessionPolicy::extended : SessionPolicy::regular); q.duplicates = st.duplicates_removed; q.unsorted = st.reordered; out = q.to_json(); }
    return json_to_py(out);
  });

  m.def("calendar_sessions", [](const std::string& a, const std::string& b) { MarketCalendar cal; std::vector<std::string> out; for (auto d : cal.sessions(date_from_iso(a), date_from_iso(b))) out.push_back(d.iso()); return out; });

  // ---- HAR family ----
  m.def("har_fit", [](F64 rv, std::vector<usize> lags, int horizon, bool log_model, bool harq, std::optional<F64> rq, const std::string& cov, int hac_lags, const std::string& target) {
    est::HarSpec s; s.lags = lags; s.horizon = horizon; s.log = log_model; s.harq = harq; s.cov.type = stats::covtype_from_string(cov); s.cov.hac_lags = hac_lags; s.target = target == "single" ? est::HarTarget::single : est::HarTarget::average;
    std::span<const f64> rvs(rv.data(), rv.shape(0)); std::span<const f64> rqs; if (rq) rqs = std::span<const f64>(rq->data(), rq->shape(0));
    json out; { nb::gil_scoped_release rel; auto e = est::har_fit(rvs, s, rqs); out = e.to_json(); out["forecast_next"] = est::har_forecast(e, rvs, s, rqs); }
    return json_to_py(out);
  }, nb::arg("rv"), nb::arg("lags") = std::vector<usize>{1, 5, 22}, nb::arg("horizon") = 1, nb::arg("log") = false, nb::arg("harq") = false, nb::arg("rq") = nb::none(), nb::arg("cov") = "hac", nb::arg("hac_lags") = -1, nb::arg("target") = "average");
  m.def("har_rolling", [](F64 rv, std::vector<usize> lags, int horizon, bool log_model, usize window, usize step, bool harq, std::optional<F64> rq) {
    est::HarSpec s; s.lags = lags; s.horizon = horizon; s.log = log_model; s.harq = harq; std::span<const f64> rvs(rv.data(), rv.shape(0)); std::span<const f64> rqs; if (rq) rqs = std::span<const f64>(rq->data(), rq->shape(0));
    est::HarRolling r; { nb::gil_scoped_release rel; r = est::har_rolling(rvs, s, window, step, rqs); }
    nb::dict d; std::vector<int64_t> o(r.origin.begin(), r.origin.end()); d["origin"] = to_np_i64(std::move(o)); d["forecast"] = to_np(r.forecast); d["actual"] = to_np(r.actual); d["n_fits"] = r.n_fits; d["seconds"] = r.seconds;
    auto l = est::forecast_loss(r.actual, r.forecast); d["mse"] = l.mse; d["qlike"] = l.qlike; d["n_loss"] = l.n; return d;
  }, nb::arg("rv"), nb::arg("lags") = std::vector<usize>{1, 5, 22}, nb::arg("horizon") = 1, nb::arg("log") = false, nb::arg("window") = 1000, nb::arg("step") = 1, nb::arg("harq") = false, nb::arg("rq") = nb::none());
  m.def("har_batch", [](F64 rv, I64 off, std::vector<usize> lags, int horizon, bool log_model, const std::string& backend_s) {
    est::HarSpec s; s.lags = lags; s.horizon = horizon; s.log = log_model; ExecutionPlan p; p.backend = backend_from_string(backend_s);
    est::HarBatchResult r; { nb::gil_scoped_release rel; r = est::har_batch(std::span<const f64>(rv.data(), rv.shape(0)), std::span<const i64>(off.data(), off.shape(0)), s, p); }
    nb::list fits; for (auto& e : r.fits) fits.append(json_to_py(e.to_json())); nb::dict d; d["fits"] = fits; d["backend"] = std::string(to_string(r.stats.backend)); d["precision"] = std::string(to_string(r.stats.precision)); d["wall_seconds"] = r.stats.wall_seconds; d["device_seconds"] = r.stats.device_seconds; return d;
  }, nb::arg("rv"), nb::arg("offsets"), nb::arg("lags") = std::vector<usize>{1, 5, 22}, nb::arg("horizon") = 1, nb::arg("log") = false, nb::arg("backend") = "auto");
  // ---- bootstrap ----
  m.def("bootstrap", [](F64 x, const std::string& stat, const std::string& scheme, double mean_block, usize replicas, uint64_t seed, double level) {
    stats::BootstrapOptions o; o.scheme = scheme == "circular" ? stats::BootstrapScheme::circular_block : (scheme == "iid" ? stats::BootstrapScheme::iid : stats::BootstrapScheme::stationary); o.mean_block = mean_block; o.replicas = replicas; o.seed = seed; o.ci_level = level;
    std::function<f64(std::span<const f64>)> f; if (stat == "mean") f = [](std::span<const f64> v) { return stats::mean<f64>(v); }; else if (stat == "rv" || stat == "sum_squares") f = [](std::span<const f64> v) { return stats::sum_squares<f64>(v); }; else if (stat == "variance") f = [](std::span<const f64> v) { return stats::variance<f64>(v); }; else throw RaidError(ErrorCode::invalid_argument, "stat must be mean|rv|variance");
    stats::BootstrapResult r; { nb::gil_scoped_release rel; r = stats::bootstrap(std::span<const f64>(x.data(), x.shape(0)), f, o); }
    nb::dict d; d["point"] = r.point; d["se"] = r.se; d["ci_lo"] = r.ci_lo; d["ci_hi"] = r.ci_hi; d["bias"] = r.bias; d["failures"] = r.failures; d["block_length"] = r.block_length; d["statistics"] = to_np(r.statistics); d["seconds"] = r.seconds; return d;
  }, nb::arg("x"), nb::arg("stat") = "mean", nb::arg("scheme") = "stationary", nb::arg("mean_block") = 0.0, nb::arg("replicas") = 999, nb::arg("seed") = 20261102, nb::arg("level") = 0.95);

  // ---- FinData client (keys from FINDATA_API_KEY / FINDATA_API_KEY_FILE; never passed through Python) ----
  m.def("findata_bars", [](const std::string& symbol, const std::string& interval, const std::string& start, const std::string& end, const std::string& adjustment, const std::string& cache_dir, bool offline) {
    findata::Config cfg = findata::Config::from_env(); if (!cache_dir.empty()) cfg.cache_dir = cache_dir; cfg.offline = offline; findata::Client cl(cfg);
    findata::BarsRequest r; r.symbol = symbol; r.interval = interval_from_string(interval); r.start = date_from_iso(start); r.end = date_from_iso(end); r.adjustment = adjustment;
    findata::BarsResult res; { nb::gil_scoped_release rel; res = cl.bars(r); }
    nb::dict d; std::vector<int64_t> t(res.bars.t_open.size()); for (usize i = 0; i < t.size(); ++i) t[i] = res.bars.t_open[i].us;
    d["timestamp_us"] = to_np_i64(std::move(t)); d["open"] = to_np(res.bars.open); d["high"] = to_np(res.bars.high); d["low"] = to_np(res.bars.low); d["close"] = to_np(res.bars.close); d["volume"] = to_np(res.bars.volume);
    std::vector<int64_t> so(res.bars.session_offsets.begin(), res.bars.session_offsets.end()); d["session_offsets"] = to_np_i64(std::move(so)); std::vector<std::string> sd; for (auto x : res.bars.session_dates) sd.push_back(x.iso()); d["sessions"] = sd;
    d["symbol"] = res.bars.symbol; d["instrument_id"] = res.bars.instrument_id; d["dataset_id"] = res.dataset_id; d["dataset_version"] = res.dataset_version; d["from_cache"] = res.from_cache; d["cache_path"] = res.cache_path; d["duplicates_dropped"] = res.duplicates_dropped;
    std::vector<std::string> rid; for (auto& e : res.exchanges) rid.push_back(e.request_id); d["request_ids"] = rid; d["key_fingerprint"] = cl.key_fingerprint();
    return d;
  }, nb::arg("symbol"), nb::arg("interval") = "1m", nb::arg("start"), nb::arg("end"), nb::arg("adjustment") = "raw", nb::arg("cache_dir") = "", nb::arg("offline") = false);
  m.def("findata_export", [](std::vector<std::string> symbols, const std::string& dataset_id, const std::string& start, const std::string& end, const std::string& adjustment, const std::string& out_path, const std::string& cache_dir) {
    findata::Config cfg = findata::Config::from_env(); if (!cache_dir.empty()) cfg.cache_dir = cache_dir; findata::Client cl(cfg);
    findata::ExportRequest r; r.symbols = symbols; r.dataset_id = dataset_id; r.start = date_from_iso(start); r.end = date_from_iso(end); r.adjustment = adjustment;
    findata::ExportResult res; { nb::gil_scoped_release rel; res = cl.export_bars(r, out_path); }
    nb::dict d; d["export_id"] = res.export_id; d["path"] = res.path; d["sha256"] = res.sha256; d["dataset_version"] = res.dataset_version; d["rows"] = res.rows; d["bytes"] = res.bytes; d["from_cache"] = res.from_cache; d["seconds"] = res.seconds_total; d["key_fingerprint"] = cl.key_fingerprint();
    return d;
  }, nb::arg("symbols"), nb::arg("dataset_id") = "market_bars.stock.1m.raw", nb::arg("start"), nb::arg("end"), nb::arg("adjustment") = "raw", nb::arg("out_path") = "", nb::arg("cache_dir") = "");
  m.def("findata_json", [](const std::string& what, const std::string& a, const std::string& b, const std::string& c, int n) {
    findata::Client cl; json out; { nb::gil_scoped_release rel;
      if (what == "health") out = cl.health(); else if (what == "catalog") out = cl.catalog_datasets(); else if (what == "dataset") out = cl.dataset(a);
      else if (what == "sessions") out = cl.sessions(a, interval_from_string(b.empty() ? "1m" : b), date_from_iso(c.substr(0, 10)), date_from_iso(c.substr(11, 10)));
      else if (what == "coverage") out = cl.coverage_symbols_all(a.empty() ? "market_bars.stock.1m.raw" : a, n, 100000);
      else if (what == "fundamentals") { std::vector<std::string> m; if (!c.empty()) { usize p0 = 0; for (usize i = 0; i <= c.size(); ++i) if (i == c.size() || c[i] == ',') { m.push_back(c.substr(p0, i - p0)); p0 = i + 1; } } out = cl.fundamentals(a, b, m); }
      else if (what == "query") out = cl.query(a, static_cast<usize>(n)); else throw RaidError(ErrorCode::invalid_argument, "unknown findata_json kind " + what); }
    return json_to_py(out);
  }, nb::arg("what"), nb::arg("a") = "", nb::arg("b") = "", nb::arg("c") = "", nb::arg("n") = 1000);
#ifdef RAID_WITH_ARROW
  m.def("read_bars_parquet", [](const std::string& path, std::vector<std::string> symbols) {
    data::ParquetBars pb; { nb::gil_scoped_release rel; pb = data::read_bars_parquet(path, symbols); }
    nb::dict out; for (auto& [sym, b] : pb.series) { nb::dict d; std::vector<int64_t> t(b.t_open.size()); for (usize i = 0; i < t.size(); ++i) t[i] = b.t_open[i].us; d["timestamp_us"] = to_np_i64(std::move(t)); d["open"] = to_np(b.open); d["high"] = to_np(b.high); d["low"] = to_np(b.low); d["close"] = to_np(b.close); d["volume"] = to_np(b.volume); std::vector<int64_t> so(b.session_offsets.begin(), b.session_offsets.end()); d["session_offsets"] = to_np_i64(std::move(so)); std::vector<std::string> sd; for (auto x : b.session_dates) sd.push_back(x.iso()); d["sessions"] = sd; d["instrument_id"] = b.instrument_id; out[sym.c_str()] = d; }
    nb::dict meta; for (auto& [k, v] : pb.metadata) meta[k.c_str()] = v; out["__meta__"] = meta; out["__rows__"] = pb.rows_read; out["__duplicates_dropped__"] = pb.duplicates_dropped; return out;
  }, nb::arg("path"), nb::arg("symbols") = std::vector<std::string>{});
#endif

  m.def("intraday_returns", [](I64 t, F64 c, const std::string& kind) {
    BarSeries b; b.interval = Interval::m1; const usize n = t.shape(0); b.reserve(n);
    for (usize i = 0; i < n; ++i) b.push_back(Timestamp(t.data()[i]), c.data()[i], c.data()[i], c.data()[i], c.data()[i], 0.0);
    data::IntradayReturns ir; { nb::gil_scoped_release rel; b.normalize(); MarketCalendar cal; auto reg = b.filter_session(cal, SessionPolicy::regular); ir = data::intraday_returns(reg, kind == "arithmetic" ? data::ReturnKind::arithmetic : data::ReturnKind::log); }
    nb::dict d; d["r"] = to_np(ir.r); d["offsets"] = to_np_i64(ir.session_offsets); std::vector<std::string> s; for (auto dd : ir.session_dates) s.push_back(dd.iso()); d["sessions"] = s; d["overnight"] = to_np(ir.overnight); d["close_to_close"] = to_np(ir.close_to_close); d["open_to_close"] = to_np(ir.open_to_close);
    return d;
  });
}