text: # Online Bayesian linear regression in 1d using Kalman Filter
  # Based on: https://github.com/probml/pmtk3/blob/master/demos/linregOnlineDemoKalman.m
  
  # The latent state corresponds to the current estimate of the regression weights w.
  # The observation model has the form
  # p(y(t) |  w(t), x(t)) = Gauss( C(t) * w(t), R(t))
  # where C(t) = X(t,:) is the observation matrix for step t.
  # The dynamics model has the form
  # p(w(t) | w(t-1)) = Gauss(A * w(t-1), Q)
  # where Q>0 allows for parameter  drift.
  # We show that the result is equivalent to batch (offline) Bayesian inference.
  
  import jax.numpy as jnp
  
  import matplotlib.pyplot as plt
  from numpy.linalg import inv
  
  from jsl.lds.kalman_filter import LDS, kalman_filter
  
  
  def kf_linreg(X, y, R, mu0, Sigma0, F, Q):
      """
      Online estimation of a linear regression
      using Kalman Filters
      Parameters
      ----------
      X: array(n_obs, dimension)
          Matrix of features
      y: array(n_obs,)
          Array of observations
      Q: float
          Known variance
      mu0: array(dimension)
          Prior mean
      Sigma0: array(dimesion, dimension)
          Prior covariance matrix
      Returns
      -------
      * array(n_obs, dimension)
          Online estimation of parameters
      * array(n_obs, dimension, dimension)
          Online estimation of uncertainty
      """
      C = lambda t: X[t][None, ...]
      lds = LDS(F, C, Q, R, mu0, Sigma0)
  
      mu_hist, Sigma_hist, _, _ = kalman_filter(lds, y)
      return mu_hist, Sigma_hist
  
  
  def posterior_lreg(X, y, R, mu0, Sigma0):
      """
      Compute mean and covariance matrix of a
      Bayesian Linear regression
      Parameters
      ----------
      X: array(n_obs, dimension)
          Matrix of features
      y: array(n_obs,)
          Array of observations
      R: float
          Known variance
      mu0: array(dimension)
          Prior mean
      Sigma0: array(dimesion, dimension)
          Prior covariance matrix
      Returns
      -------
      * array(dimension)
          Posterior mean
      * array(n_obs, dimension, dimension)
          Posterior covariance matrix
      """
      Sn_bayes_inv = inv(Sigma0) + X.T @ X / R
      Sn_bayes = inv(Sn_bayes_inv)
      mn_bayes = Sn_bayes @ (inv(Sigma0) @ mu0 + X.T @ y / R)
  
      return mn_bayes, Sn_bayes
  
  def main():
      n_obs = 21
      timesteps = jnp.arange(n_obs)
      x = jnp.linspace(0, 20, n_obs)
      X = jnp.c_[jnp.ones(n_obs), x]
      F = jnp.eye(2)
      mu0 = jnp.zeros(2)
      Sigma0 = jnp.eye(2) * 10.
  
      Q, R = 0, 1
      # Data from original matlab example
      y = jnp.array([2.4865, -0.3033, -4.0531, -4.3359, -6.1742, -5.604, -3.5069, -2.3257, -4.6377,
                     -0.2327, -1.9858, 1.0284, -2.264, -0.4508, 1.1672, 6.6524, 4.1452, 5.2677, 6.3403, 9.6264, 14.7842])
  
      # Online estimation
      mu_hist, Sigma_hist = kf_linreg(X, y, R, mu0, Sigma0, F, Q)
      kf_var = Sigma_hist[-1, [0, 1], [0, 1]]
      w0_hist, w1_hist = mu_hist.T
      w0_err, w1_err = jnp.sqrt(Sigma_hist[:, [0, 1], [0, 1]].T)
  
      # Offline estimation
      (w0_post, w1_post), Sigma_post = posterior_lreg(X, y, R, mu0, Sigma0)
      w0_std, w1_std = jnp.sqrt(Sigma_post[[0, 1], [0, 1]])
  
      dict_figures = {}
  
      fig, ax = plt.subplots()
      ax.errorbar(timesteps, w0_hist, w0_err, fmt="-o", label="$w_0$", color="black", fillstyle="none")
      ax.errorbar(timesteps, w1_hist, w1_err, fmt="-o", label="$w_1$", color="tab:red")
  
      ax.axhline(y=w0_post, c="black", label="$w_0$ batch")
      ax.axhline(y=w1_post, c="tab:red", linestyle="--", label="$w_1$ batch")
  
      ax.fill_between(timesteps, w0_post - w0_std, w0_post + w0_std, color="black", alpha=0.4)
      ax.fill_between(timesteps, w1_post - w1_std, w1_post + w1_std, color="tab:red", alpha=0.4)
  
      plt.legend()
      ax.set_xlabel("time")
      ax.set_ylabel("weights")
      ax.set_ylim(-8, 4)
      ax.set_xlim(-0.5, n_obs)
      dict_figures["linreg_online_kalman"] = fig
      plt.savefig("prev.png")
      return dict_figures
  
  
  if __name__ == "__main__":
      from jsl.demos.plot_utils import savefig
      plt.rcParams["axes.spines.right"] = False
      plt.rcParams["axes.spines.top"] = False
      dict_figures = main()
      savefig(dict_figures)
      plt.show()