// modules/congestion_controller/goog_cc/kalman_bandwidth_filter.h
#ifndef MODULES_CONGESTION_CONTROLLER_GOOG_CC_KALMAN_BANDWIDTH_FILTER_H_
#define MODULES_CONGESTION_CONTROLLER_GOOG_CC_KALMAN_BANDWIDTH_FILTER_H_

#include <array>
#include <cstdint>
#include <optional>
#include "api/units/data_rate.h"
#include "api/units/time_delta.h"
#include "api/units/timestamp.h"
#include "rtc_base/system/rtc_export.h"

namespace webrtc {

// GCC 内部状态上下文，用于自适应 Q/R 计算
struct GccInternalState {
  enum class DetectorState { kNormal, kOverusing, kUnderusing };
  DetectorState delay_detector_state = DetectorState::kNormal;
  double trendline_slope = 0.0;          // 延迟趋势斜率
  double trendline_var = 0.0;            // 趋势线方差
  bool in_probing = false;               // 是否处于探测阶段
  double link_capacity_kbps = -1.0;      // Probe 估算的链路容量
};

class RTC_EXPORT KalmanBandwidthFilter {
 public:
  // 配置结构体，支持通过 FieldTrial / RemoteConfig 动态下发
  struct Config {
    // 过程噪声基础系数 (sigma_p^2)，单位: (kbps/s)^2
    double process_noise_base = 500.0; 
    // 观测噪声基础方差，单位: kbps^2
    double measurement_noise_base = 2500.0; 
    // 过控/探测态下观测噪声放大倍数
    double noise_scale_overuse = 10.0;
    double noise_scale_probing = 8.0;
    // 协方差矩阵迹上限，防发散
    double max_covariance_trace = 1e8;
    // 马氏距离阈值 (Chi-square 0.99, DoF=1)
    double mahalanobis_threshold = 6.63; 
    // 物理带宽边界
    DataRate min_bitrate = DataRate::KilobitsPerSec(30);
    DataRate max_bitrate = DataRate::KilobitsPerSec(50000);
    // 预热轮数
    int warmup_samples = 3;
  };

  explicit KalmanBandwidthFilter(const Config& config = Config());
  ~KalmanBandwidthFilter();

  // 禁止拷贝，允许移动
  KalmanBandwidthFilter(const KalmanBandwidthFilter&) = delete;
  KalmanBandwidthFilter& operator=(const KalmanBandwidthFilter&) = delete;

  // 核心接口：输入 GCC 原始预估 + 内部状态，输出平滑后带宽
  // 返回 nullopt 表示熔断/预热期透传原始值
  std::optional<DataRate> Update(DataRate gcc_estimate, 
                                 Timestamp at_time,
                                 const GccInternalState& gcc_state);

  // 网络切换/会话恢复时重置状态
  void Reset(DataRate initial_bitrate = DataRate::Zero());

  // 序列化/反序列化：用于进程重启、后台切换状态保持
  std::string SerializeState() const;
  bool DeserializeState(const std::string& serialized);

  // 供监控/调试导出内部状态
  struct InternalState {
    DataRate filtered_bitrate;
    double trend_kbps_per_s;   // 估计的带宽变化率
    double covariance_trace;   // 协方差矩阵迹
    double mahalanobis_dist;   // 当前马氏距离
    bool is_warming_up;
    bool is_fallback;          // 是否处于熔断跳过更新状态
  };
  InternalState GetInternalState() const;

 private:
  // 矩阵运算内联实现 (2x2 矩阵 + 2x1 向量)，避免 Eigen 依赖
  using Matrix2x2 = std::array<std::array<double, 2>, 2>;
  using Vector2 = std::array<double, 2>;

  void Predict(TimeDelta dt);
  void UpdateStep(double measurement_kbps, double R);
  double ComputeAdaptiveR(const GccInternalState& state) const;
  Matrix2x2 ComputeAdaptiveQ(TimeDelta dt) const;
  void ClampState();
  bool CheckMahalanobis(double measurement, double S) const;

  Config config_;
  // 状态向量 X = [Bandwidth_kbps, Trend_kbps_per_s]^T
  Vector2 x_ = {0.0, 0.0}; 
  // 协方差矩阵 P
  Matrix2x2 P_ = {{ {1e4, 0}, {0, 1e4} }};
  // 状态转移矩阵 F (dt 相关，每次 Predict 重算)
  // 观测矩阵 H = [1, 0] (常量)
  
  Timestamp last_time_;
  int sample_count_ = 0;
  bool fallback_mode_ = false; // 熔断标志
};

}  // namespace webrtc
#endif

// modules/congestion_controller/goog_cc/kalman_bandwidth_filter.cc
#include "kalman_bandwidth_filter.h"
#include <cmath>
#include <limits>
#include "rtc_base/checks.h"
#include "rtc_base/logging.h"
#include "rtc_base/strings/json.h"

namespace webrtc {

KalmanBandwidthFilter::KalmanBandwidthFilter(const Config& config) : config_(config) {}

std::optional<DataRate> KalmanBandwidthFilter::Update(DataRate gcc_estimate, 
                                                      Timestamp at_time,
                                                      const GccInternalState& gcc_state) {
  // 1. 时间步计算
  TimeDelta dt = (last_time_.IsInfinite() || at_time < last_time_) 
                 ? TimeDelta::Millis(200) // 首帧/异常时间戳兜底
                 : at_time - last_time_;
  last_time_ = at_time;
  double dt_sec = dt.seconds<double>();
  RTC_DCHECK(dt_sec > 0 && dt_sec < 10.0); // 防御异常 dt

  // 2. 预热期：直接透传，初始化状态
  if (sample_count_ < config_.warmup_samples) {
    x_[0] = gcc_estimate.kbps<double>();
    x_[1] = 0.0;
    sample_count_++;
    if (sample_count_ == config_.warmup_samples) {
      // 预热结束，初始化协方差
      P_ = {{ {config_.measurement_noise_base, 0}, {0, 1e4} }};
    }
    return gcc_estimate; // 预热期不滤波，建立基线
  }

  // 3. 自适应 Q/R 计算
  Matrix2x2 Q = ComputeAdaptiveQ(dt_sec);
  double R = ComputeAdaptiveR(gcc_state);

  // 4. 预测步
  Predict(dt_sec, Q);

  // 5. 观测合法性检查 (马氏距离)
  double z = gcc_estimate.kbps<double>();
  // S = H * P * H' + R = P[0][0] + R
  double S = P_[0][0] + R; 
  
  if (!CheckMahalanobis(z, S)) {
    // 异常观测：仅预测，不更新。标记熔断。
    fallback_mode_ = true;
    RTC_LOG(LS_WARNING) << "Kalman Filter: Mahalanobis check failed. z=" << z 
                        << " pred=" << x_[0] << " S=" << S;
    // 返回预测值而非原始值，避免将脏数据传给编码器
    ClampState();
    return DataRate::KilobitsPerSec(x_[0]);
  }
  fallback_mode_ = false;

  // 6. 更新步
  UpdateStep(z, R);

  // 7. 物理约束裁剪
  ClampState();

  return DataRate::KilobitsPerSec(x_[0]);
}

void KalmanBandwidthFilter::Predict(double dt_sec, const Matrix2x2& Q) {
  // F = [1, dt; 0, 1]
  // x = F * x
  double new_bandwidth = x_[0] + x_[1] * dt_sec;
  x_[1] = x_[1]; // 恒速模型，趋势不变
  x_[0] = new_bandwidth;

  // P = F * P * F' + Q
  // 手动展开 2x2 矩阵乘法，避免循环开销
  double p00 = P_[0][0] + 2 * dt_sec * P_[0][1] + dt_sec * dt_sec * P_[1][1] + Q[0][0];
  double p01 = P_[0][1] + dt_sec * P_[1][1] + Q[0][1];
  double p11 = P_[1][1] + Q[1][1];
  
  P_[0][0] = p00;
  P_[0][1] = P_[1][0] = p01;
  P_[1][1] = p11;
}

void KalmanBandwidthFilter::UpdateStep(double z, double R) {
  // K = P * H' / (H * P * H' + R) = [P[0][0]; P[1][0]] / (P[0][0] + R)
  double S = P_[0][0] + R;
  double k0 = P_[0][0] / S;
  double k1 = P_[1][0] / S;

  // y = z - H * x = z - x[0]
  double y = z - x_[0];

  // x = x + K * y
  x_[0] += k0 * y;
  x_[1] += k1 * y;

  // P = (I - K * H) * P
  // 约瑟夫形式数值更稳定: P = (I-KH)P(I-KH)' + KRK'
  // 此处简化版 (对称性可能轻微丢失，但 2x2 场景可接受，工程上建议用约瑟夫形式)
  double p00 = (1 - k0) * P_[0][0];
  double p01 = (1 - k0) * P_[0][1];
  double p11 = P_[1][1] - k1 * P_[0][1];
  
  P_[0][0] = p00;
  P_[0][1] = P_[1][0] = p01;
  P_[1][1] = p11;
}

// 自适应 R: 核心策略
double KalmanBandwidthFilter::ComputeAdaptiveR(const GccInternalState& state) const {
  double R = config_.measurement_noise_base;
  
  // 1. 基于趋势线方差调整 (GCC 内部卡尔曼输出的方差)
  if (state.trendline_var > 0) {
    // 归一化映射: 方差越大 -> 观测越不可信 -> R 越大
    R += state.trendline_var * 100.0; // 经验系数
  }

  // 2. 状态机加权
  if (state.delay_detector_state == GccInternalState::DetectorState::kOverusing) {
    R *= config_.noise_scale_overuse;
  }
  if (state.in_probing) {
    R *= config_.noise_scale_probing;
  }
  
  // 3. 链路容量约束：如果 Probe 给出容量上限，且 GCC 估计接近上限，增大 R 抑制超调
  if (state.link_capacity_kbps > 0 && x_[0] > state.link_capacity_kbps * 0.9) {
    R *= 2.0;
  }
  return R;
}

// 自适应 Q: 基于历史趋势变化率估计
Matrix2x2 KalmanBandwidthFilter::ComputeAdaptiveQ(double dt_sec) const {
  // sigma_p^2 可基于 x_[1] (趋势) 动态调整：趋势大 -> 网络波动大 -> Q 大
  double sigma_p_sq = config_.process_noise_base * (1.0 + std::abs(x_[1]) / 1000.0); 
  
  double dt2 = dt_sec * dt_sec;
  double dt3 = dt2 * dt_sec;
  
  return {{ 
    { dt3/3.0 * sigma_p_sq, dt2/2.0 * sigma_p_sq },
    { dt2/2.0 * sigma_p_sq, dt_sec * sigma_p_sq }
  }};
}

bool KalmanBandwidthFilter::CheckMahalanobis(double z, double S) const {
  if (S <= 0) return true; // 数值异常保护
  double diff = z - x_[0];
  double d2 = (diff * diff) / S;
  return d2 <= config_.mahalanobis_threshold;
}

void KalmanBandwidthFilter::ClampState() {
  double min_b = config_.min_bitrate.kbps<double>();
  double max_b = config_.max_bitrate.kbps<double>();
  
  if (x_[0] < min_b) x_[0] = min_b;
  if (x_[0] > max_b) x_[0] = max_b;
  
  // 协方差发散保护
  double trace = P_[0][0] + P_[1][1];
  if (trace > config_.max_covariance_trace || std::isnan(trace)) {
    RTC_LOG(LS_ERROR) << "Kalman Filter: Covariance divergence detected. Resetting P.";
    P_ = {{ {config_.measurement_noise_base, 0}, {0, 1e4} }};
  }
}

// 序列化实现 (JSON 格式，便于调试与持久化)
std::string KalmanBandwidthFilter::SerializeState() const {
  Json::Value root;
  root["x0"] = x_[0];
  root["x1"] = x_[1];
  root["p00"] = P_[0][0];
  root["p01"] = P_[0][1];
  root["p11"] = P_[1][1];
  root["sample_count"] = sample_count_;
  root["last_time_ms"] = last_time_.ms();
  return Json::WriteCompact(root);
}

bool KalmanBandwidthFilter::DeserializeState(const std::string& serialized) {
  Json::Value root;
  if (!Json::Parse(serialized, &root)) return false;
  x_[0] = root["x0"].asDouble();
  x_[1] = root["x1"].asDouble();
  P_[0][0] = root["p00"].asDouble();
  P_[0][1] = root["p10"].asDouble(); // 兼容旧字段名
  P_[1][0] = P_[0][1];
  P_[1][1] = root["p11"].asDouble();
  sample_count_ = root["sample_count"].asInt();
  last_time_ = Timestamp::Millis(root["last_time_ms"].asInt64());
  return true;
}

} // namespace webrtc

// test/kalman_bandwidth_filter_test.cc
#include "modules/congestion_controller/goog_cc/kalman_bandwidth_filter.h"
#include "test/gtest.h"
#include "api/units/data_rate.h"
#include "api/units/time_delta.h"

namespace webrtc {
namespace test {

TEST(KalmanBandwidthFilterTest, SteadyStateConvergence) {
  KalmanBandwidthFilter filter;
  GccInternalState state;
  Timestamp now = Timestamp::Millis(0);
  
  // 预热
  for (int i = 0; i < 5; ++i) {
    filter.Update(DataRate::KilobitsPerSec(1000), now + TimeDelta::Millis(i*200), state);
  }
  
  // 稳态输入 1000kbps，带噪声 ±50kbps
  double sum = 0;
  for (int i = 5; i < 105; ++i) {
    double noise = (i % 2 == 0) ? 50 : -50;
    auto out = filter.Update(DataRate::KilobitsPerSec(1000 + noise), 
                             now + TimeDelta::Millis(i*200), state);
    ASSERT_TRUE(out.has_value());
    sum += out->kbps<double>();
  }
  double avg = sum / 100;
  EXPECT_NEAR(avg, 1000, 5.0); // 稳态误差 < 5kbps
  // 验证抖动抑制：输出方差应远小于输入方差
}

TEST(KalmanBandwidthFilterTest, MahalanobisRejection) {
  KalmanBandwidthFilter filter;
  GccInternalState state;
  Timestamp now = Timestamp::Millis(0);
  
  // 正常建立状态
  for (int i=0; i<10; ++i) filter.Update(DataRate::KilobitsPerSec(1000), now + TimeDelta::Millis(i*200), state);
  
  // 注入异常脏数据：带宽突变至 10000kbps (物理不可能)
  auto out = filter.Update(DataRate::KilobitsPerSec(10000), now + TimeDelta::Millis(2200), state);
  ASSERT_TRUE(out.has_value());
  // 应被熔断，输出应接近预测值(1000kbps)，而非 10000kbps
  EXPECT_LT(out->kbps<double>(), 2000); 
  EXPECT_GT(out->kbps<double>(), 500);
  
  // 下一帧恢复正常，应快速跟回
  out = filter.Update(DataRate::KilobitsPerSec(1000), now + TimeDelta::Millis(2400), state);
  EXPECT_NEAR(out->kbps<double>(), 1000, 50);
}

TEST(KalmanBandwidthFilterTest, SerializationRoundTrip) {
  KalmanBandwidthFilter filter1;
  GccInternalState state;
  Timestamp now = Timestamp::Millis(10000);
  filter1.Update(DataRate::KilobitsPerSec(2000), now, state);
  
  std::string serialized = filter1.SerializeState();
  KalmanBandwidthFilter filter2;
  ASSERT_TRUE(filter2.DeserializeState(serialized));
  
  auto s1 = filter1.GetInternalState();
  auto s2 = filter2.GetInternalState();
  EXPECT_DOUBLE_EQ(s1.filtered_bitrate.kbps<double>(), s2.filtered_bitrate.kbps<double>());
  EXPECT_DOUBLE_EQ(s1.trend_kbps_per_s, s2.trend_kbps_per_s);
}

} // namespace test
} // namespace webrtc

// 网络线程安全的配置热更
void KalmanBandwidthFilter::OnConfigUpdated(const Config& new_config) {
  // 1. 校验合法性
  if (new_config.measurement_noise_base <= 0) return;
  
  // 2. 平滑过渡：协方差矩阵 P 不重置，仅更新 config_ 副本
  //    避免配置变更瞬间引入状态突变
  config_ = new_config; 
  
  // 3. 上报配置变更事件 (含版本号) 供监控关联
  RtcEventLog::LogConfigChange("KalmanFilter", new_config.ToString());
}