"""
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos
"""

from collections.abc import Callable
from dataclasses import dataclass
import math

import numpy as np


Array = np.ndarray
Prediction = Callable[[Array, float, dict[str, float]], Array]
Measurement = Callable[[Array, dict[str, float]], Array]
Injection = Callable[[Array, Array], Array]
NativePrediction = Callable[[Array, Array, float, Array, dict[str, float]], None]
NativeUpdate = Callable[[Array, Array, int, Array, Array], Array]


@dataclass(frozen=True)
class Observation:
  kind: int
  values: Array
  noise: Array


@dataclass(frozen=True)
class ModelDefinition:
  state_size: int
  error_size: int
  transition: Prediction
  measurements: dict[int, Measurement]
  process_noise: Array
  observation_noise: dict[int, Array]
  inject_error: Injection | None = None
  error_projection: Callable[[Array], Array] | None = None
  normalize: Callable[[Array], Array] | None = None
  native_predict: NativePrediction | None = None
  native_update: NativeUpdate | None = None


@dataclass
class _Snapshot:
  time: float
  state: Array
  covariance: Array


@dataclass
class _Event:
  time: float
  observation: Observation
  order: int


class StateEstimator:
  def __init__(self, model: ModelDefinition, initial_state: Array, initial_covariance: Array,
               max_rewind_age: float = 0.0):
    self.model = model
    self.parameters: dict[str, float] = {}
    self.max_rewind_age = max_rewind_age
    self._order = 0
    self.init_state(initial_state, initial_covariance, None)

  @property
  def x(self) -> Array:
    return self._state.copy()

  @property
  def P(self) -> Array:
    return self._covariance.copy()

  @property
  def t(self) -> float:
    return self._time

  def set_global(self, name: str, value: float) -> None:
    self.parameters[name] = float(value)

  def init_state(self, state: Array, covs: Array, filter_time: float | None) -> None:
    state = np.asarray(state, dtype=np.float64).reshape(-1)
    covariance = np.asarray(covs, dtype=np.float64)
    self._validate_state(state, covariance)
    self._state = self._normalize(state.copy())
    self._covariance = self._stabilize(covariance.copy())
    self._time = math.nan if filter_time is None else float(filter_time)
    self._events: list[_Event] = []
    self._snapshots = [_Snapshot(self._time, self._state.copy(), self._covariance.copy())]

  def set_filter_time(self, filter_time: float | None) -> None:
    self._time = math.nan if filter_time is None else float(filter_time)

  def reset_rewind(self) -> None:
    self._events.clear()
    self._snapshots = [_Snapshot(self._time, self._state.copy(), self._covariance.copy())]

  def predict(self, time: float) -> None:
    time = float(time)
    if math.isnan(self._time):
      self._time = time
      return
    if time < self._time:
      raise ValueError("prediction time precedes estimator time")
    dt = time - self._time
    if dt == 0.0:
      return
    if self.model.native_predict is not None:
      self.model.native_predict(self._state, self._covariance, dt, self.model.process_noise, self.parameters)
      self._time = time
      return
    previous = self._state.copy()
    transition_jacobian = self._jacobian(lambda value: self.model.transition(value, dt, self.parameters), previous)
    predicted = self.model.transition(previous, dt, self.parameters)
    projection = self._error_projection(previous)
    if self.model.error_size == self.model.state_size:
      error_transition = transition_jacobian
    else:
      error_transition = np.linalg.pinv(self._error_projection(predicted)) @ transition_jacobian @ projection
    self._state = self._normalize(predicted)
    self._covariance = self._stabilize(error_transition @ self._covariance @ error_transition.T + dt * self.model.process_noise)
    self._time = time

  def predict_and_observe(self, time: float, kind: int, measurements: Array, noise: Array | None = None):
    values = self._measurement_batch(kind, measurements)
    noises = self._noise_batch(kind, len(values), noise)
    event = _Event(float(time), Observation(kind, values, noises), self._order)
    self._order += 1
    if not math.isnan(self._time) and event.time < self._time:
      if self.max_rewind_age <= 0.0 or self._time - event.time > self.max_rewind_age:
        return None
      return self._rewind(event)
    result = self._apply_event(event)
    self._events.append(event)
    self._snapshots.append(_Snapshot(self._time, self._state.copy(), self._covariance.copy()))
    self._trim_history()
    return result

  def _apply_event(self, event: _Event):
    self.predict(event.time)
    prior_state = self._state.copy()
    prior_covariance = self._covariance.copy()
    innovations = []
    for measurement, noise in zip(event.observation.values, event.observation.noise, strict=True):
      innovations.append(self._update(event.observation.kind, measurement, noise))
    return (event.time, self.x, prior_state, self.P, prior_covariance, event.observation.kind,
            tuple(innovations), event.observation.values.copy(), event.observation.noise.copy())

  def _update(self, kind: int, measurement: Array, noise: Array) -> Array:
    measurement_function = self.model.measurements.get(kind)
    if measurement_function is None:
      raise KeyError(f"unknown observation kind {kind}")
    measurement = np.asarray(measurement, dtype=np.float64).reshape(-1)
    if noise.shape != (measurement.size, measurement.size):
      raise ValueError("observation noise dimension mismatch")
    if self.model.native_update is not None:
      innovation = self.model.native_update(self._state, self._covariance, kind, measurement, noise)
      return innovation
    predicted = np.asarray(measurement_function(self._state, self.parameters), dtype=np.float64).reshape(-1)
    if predicted.shape != measurement.shape:
      raise ValueError("measurement dimension mismatch")
    innovation = measurement - predicted
    state_jacobian = self._jacobian(lambda value: measurement_function(value, self.parameters), self._state)
    observation_jacobian = state_jacobian @ self._error_projection(self._state)
    innovation_covariance = observation_jacobian @ self._covariance @ observation_jacobian.T + noise
    gain = np.linalg.solve(innovation_covariance, observation_jacobian @ self._covariance).T
    delta = gain @ innovation
    self._state = self._normalize(self._inject(self._state, delta))
    identity = np.eye(self.model.error_size)
    residual = identity - gain @ observation_jacobian
    self._covariance = self._stabilize(residual @ self._covariance @ residual.T + gain @ noise @ gain.T)
    self._require_finite()
    return innovation

  def _rewind(self, new_event: _Event):
    events = sorted(self._events + [new_event], key=lambda event: (event.time, event.order))
    base_index = max(i for i, snapshot in enumerate(self._snapshots) if math.isnan(snapshot.time) or snapshot.time <= new_event.time)
    base = self._snapshots[base_index]
    retained = self._events[:base_index]
    retained_orders = {event.order for event in retained}
    replay = [event for event in events if event.order not in retained_orders]
    self._state = base.state.copy()
    self._covariance = base.covariance.copy()
    self._time = base.time
    self._events = retained.copy()
    self._snapshots = self._snapshots[:base_index + 1]
    result = None
    for event in replay:
      current = self._apply_event(event)
      self._events.append(event)
      self._snapshots.append(_Snapshot(self._time, self._state.copy(), self._covariance.copy()))
      if event is new_event:
        result = current
    self._trim_history()
    return result

  def _trim_history(self) -> None:
    if self.max_rewind_age <= 0.0 or math.isnan(self._time):
      return
    cutoff = self._time - self.max_rewind_age
    remove = 0
    while remove < len(self._events) and self._events[remove].time < cutoff:
      remove += 1
    if remove:
      self._events = self._events[remove:]
      self._snapshots = self._snapshots[remove:]

  def _measurement_batch(self, kind: int, measurements: Array) -> Array:
    measurement_function = self.model.measurements.get(kind)
    if measurement_function is None:
      raise KeyError(f"unknown observation kind {kind}")
    if self.model.native_update is not None and kind in self.model.observation_noise:
      expected = self.model.observation_noise[kind].shape[0]
    else:
      expected = np.asarray(measurement_function(self._state, self.parameters)).size
    values = np.asarray(measurements, dtype=np.float64)
    if values.ndim == 1:
      values = values.reshape(1, -1)
    elif values.ndim != 2:
      raise ValueError("measurements must be one or two dimensional")
    if values.shape[1] != expected:
      raise ValueError("measurement dimension mismatch")
    return values

  def _noise_batch(self, kind: int, count: int, noise: Array | None) -> Array:
    if noise is None:
      base = self.model.observation_noise.get(kind)
      if base is None:
        raise KeyError(f"missing observation noise for kind {kind}")
      return np.repeat(np.asarray(base, dtype=np.float64)[None, :, :], count, axis=0)
    noises = np.asarray(noise, dtype=np.float64)
    if noises.ndim == 2:
      noises = noises[None, :, :]
    if noises.shape[0] == 1 and count > 1:
      noises = np.repeat(noises, count, axis=0)
    if noises.shape[0] != count:
      raise ValueError("observation noise batch mismatch")
    return noises

  def _jacobian(self, function: Callable[[Array], Array], value: Array) -> Array:
    output = np.asarray(function(value), dtype=np.float64).reshape(-1)
    result = np.empty((output.size, value.size), dtype=np.float64)
    for index in range(value.size):
      step = np.cbrt(np.finfo(np.float64).eps) * max(1.0, abs(value[index]))
      upper = value.copy()
      lower = value.copy()
      upper[index] += step
      lower[index] -= step
      result[:, index] = (np.asarray(function(upper)).reshape(-1) - np.asarray(function(lower)).reshape(-1)) / (2.0 * step)
    return result

  def _inject(self, state: Array, delta: Array) -> Array:
    if self.model.inject_error is None:
      return state + delta
    return self.model.inject_error(state, delta)

  def _error_projection(self, state: Array) -> Array:
    if self.model.error_projection is None:
      return np.eye(self.model.state_size, self.model.error_size)
    return np.asarray(self.model.error_projection(state), dtype=np.float64)

  def _normalize(self, state: Array) -> Array:
    if self.model.normalize is None:
      return np.asarray(state, dtype=np.float64).reshape(-1)
    return np.asarray(self.model.normalize(state), dtype=np.float64).reshape(-1)

  def _stabilize(self, covariance: Array) -> Array:
    covariance = (covariance + covariance.T) * 0.5
    eigenvalues, eigenvectors = np.linalg.eigh(covariance)
    if eigenvalues[0] < -1e-10:
      raise FloatingPointError("covariance is not positive semidefinite")
    return (eigenvectors * np.maximum(eigenvalues, 0.0)) @ eigenvectors.T

  def _validate_state(self, state: Array, covariance: Array) -> None:
    if state.shape != (self.model.state_size,):
      raise ValueError("state dimension mismatch")
    if covariance.shape != (self.model.error_size, self.model.error_size):
      raise ValueError("covariance dimension mismatch")
    if self.model.process_noise.shape != covariance.shape:
      raise ValueError("process noise dimension mismatch")
    if not np.isfinite(state).all() or not np.isfinite(covariance).all():
      raise ValueError("state and covariance must be finite")

  def _require_finite(self) -> None:
    if not np.isfinite(self._state).all() or not np.isfinite(self._covariance).all():
      raise FloatingPointError("estimator produced non-finite values")


class EstimatorModel:
  def __init__(self, estimator: StateEstimator):
    self.filter = estimator

  @property
  def x(self) -> Array:
    return self.filter.x

  @property
  def P(self) -> Array:
    return self.filter.P

  @property
  def t(self) -> float:
    return self.filter.t

  def init_state(self, state: Array, covs: Array, filter_time: float | None) -> None:
    self.filter.init_state(state, covs, filter_time)

  def predict(self, time: float) -> None:
    self.filter.predict(time)

  def predict_and_observe(self, time: float, kind: int, measurements: Array, noise: Array | None = None):
    return self.filter.predict_and_observe(time, kind, measurements, noise)