# -*- coding: utf-8 -*-
"""Hybrid Kalman Filter pkg

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1w4b8ek2jRx6pgadV-9OgInnxOHEmyUUv
"""

import numpy as np
from scipy.linalg import expm
from math import *
import random
import matplotlib.pyplot as plt


class HybridKalmanFilter:
    def __init__(
        self,
        dt,
        noise_std,
        noise_covariance_factor,
        n,
        max_x,
        max_y,
        motion_type=None,
        verbose=False,
        state_transition_factor=1,
        process_noise_covariance_factor=0.01,
        u=0,
        P_init=np.matrix(
            [
                [1, 0, 0, 0, 0, 0],
                [0, 1, 0, 0, 0, 0],
                [0, 0, 5, 0, 0, 0],
                [0, 0, 0, 5, 0, 0],
                [0, 0, 0, 0, 13, 0],
                [0, 0, 0, 0, 0, 13],
            ]
        ),
        B=0,
        linear_timesteps=True,
        timestep_var=0.1,
        linear_motion_x_factor=2,
        linear_motion_y_factor=5,
        nonlinear_motion_x_factor=6,
        nonlinear_motion_y_factor=6,
        seed=42,
    ):
        self.x = {}
        self.P = {}
        self.dt = dt
        self.noise_std = noise_std  # Noise/variance for observation gaussians
        self.noise_covariance_factor = noise_covariance_factor
        self.process_noise_covariance_factor = process_noise_covariance_factor
        self.u = u  # Control
        self.P_init = P_init  # Initial uncertainty weights
        self.max_x = (
            max_x  # Maximum x coordinate for P-RNG of objects for state estimation
        )
        self.max_y = (
            max_y  # Maximum y coordinate for P-RNG of objects for state estimation
        )
        self.n = n  # Number of observations/objects to generate
        self.state_transition_factor = state_transition_factor
        self.seed = seed
        self.verbose = verbose
        self.true_positions = None
        self.noisy_observations = None
        self.motion_type = motion_type
        self.linear_motion_x_factor = linear_motion_x_factor
        self.linear_motion_y_factor = linear_motion_y_factor
        self.nonlinear_motion_x_factor = nonlinear_motion_x_factor
        self.nonlinear_motion_y_factor = nonlinear_motion_y_factor
        self.linear_timesteps = linear_timesteps
        self.timestep_var = timestep_var

        # Continuous-time state transition matrix
        self.A = (
            np.matrix(
                [
                    [0, 0, 1, 0, 0, 0],
                    [0, 0, 0, 1, 0, 0],
                    [0, 0, 0, 0, 1, 0],
                    [0, 0, 0, 0, 0, 1],
                    [0, 0, 0, 0, 0, 0],
                    [0, 0, 0, 0, 0, 0],
                ]
            )
            * self.state_transition_factor
        )

        # Process noise covariance matrix
        self.Q_c = np.matrix(np.eye(6)) * self.process_noise_covariance_factor

        # Discretize the state transition matrix
        self.F = expm(self.A * self.dt)

        # Discretize input control matrix
        self.B = B * expm(self.A * self.dt) * self.u

        # Discretize the process noise covariance matrix
        self.Q_d = self.Q_c * self.dt  # Simplified approximation

        # Measurement mapping matrix
        self.H = np.matrix([[1, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0]])

        # measurement noise covariance matrix
        self.R = (
            np.matrix([[1, 0], [0, 1]]) * self.noise_covariance_factor
        )  # 0.008 in my implementation

    # Function to generate random true_positions
    def generate_true_positions(self) -> dict:
        true_positions = {}
        for i in range(1, self.n + 1):
            true_positions[i] = (
                random.randrange(0, self.max_x),
                random.randrange(0, self.max_y),
            )
        self.true_positions = true_positions
        return self.true_positions

    # Function to simulate noisy observations
    def generate_noisy_observations(self) -> dict:
        if self.true_positions is None:
            true_positions = self.generate_true_positions()
        else:
            true_positions = self.true_positions
        noisy_observations = {}
        for id, (x_true, y_true) in true_positions.items():
            x_noisy = x_true + np.random.normal(0, self.noise_std)
            y_noisy = y_true + np.random.normal(0, self.noise_std)
            noisy_observations[id] = (x_noisy, y_noisy)
        self.noisy_observations = noisy_observations
        return self.noisy_observations

    def predict(self, id):
        # using discretized equations from continuous representation (x_x|x+1 = e^At + B* integral from 0 to dT of (e^A*tau * dtau))
        self.x[id] = self.F * self.x[id] + self.B * expm(self.F * self.dt) * self.u
        self.P[id] = (self.F * self.P[id] * np.transpose(self.F)) + self.Q_d

    def measurement_update(self, id, coords):
        I = np.identity(n=6)
        x_meas, y_meas = coords

        z = np.matrix([[x_meas], [y_meas]])

        K = (
            self.P[id]
            * np.transpose(self.H)
            * np.linalg.inv((self.H * self.P[id] * np.transpose(self.H) + self.R))
        )
        self.x[id] = self.x[id] + K * (z - self.H * self.x[id])
        self.P[id] = (I - K * self.H) * self.P[id]

    def execute_hybrid_kf(self, time_steps, visualize=True):
        self.true_positions = self.generate_true_positions()
        self.noisy_observations = self.generate_noisy_observations()
        t = 0
        true_positions_history = {id: [] for id in self.true_positions}
        estimated_positions_history = {id: [] for id in self.true_positions}

        while t < time_steps:
            if self.verbose:
                print("True observations:", self.true_positions)
                print("Noisy observations:", self.noisy_observations)

            # Update true positions
            for id, (x, y) in self.true_positions.items():
                if self.motion_type == "linear":
                    self.true_positions[id] = (
                        x + self.linear_motion_x_factor,
                        y + self.linear_motion_y_factor,
                    )
                elif self.motion_type == "nonlinear":
                    self.true_positions[id] = (
                        x * self.nonlinear_motion_x_factor,
                        y * self.nonlinear_motion_y_factor,
                    )

                true_positions_history[id].append(
                    (self.true_positions[id],)
                )  # Tuple for consistency

            # Update noisy observations
            for id, (x, y) in self.noisy_observations.items():
                if self.motion_type == "linear":
                    self.noisy_observations[id] = (
                        x + self.linear_motion_x_factor,
                        y + self.linear_motion_y_factor,
                    )
                elif self.motion_type == "nonlinear":
                    self.noisy_observations[id] = (
                        x * self.nonlinear_motion_x_factor,
                        y * self.nonlinear_motion_y_factor,
                    )

            for id, coords in self.noisy_observations.items():
                if id not in self.x:
                    # Initialize state estimates with reasonable values
                    self.x[id] = np.matrix([[0], [0], [0], [0], [0], [0]])
                if id not in self.P:
                    self.P[id] = self.P_init

                self.predict(id)
                self.measurement_update(id, coords)
                position_estimate = self.x[id][:2].T
                if self.verbose == True:
                    print(
                        f"Position estimate for object {id} at time {t}: {position_estimate[0]}"
                    )
                estimated_positions_history[id].append(position_estimate.tolist())

            # Update t in non-linear fashion
            if not self.linear_timesteps:
                t = t + np.random.normal(
                    self.dt, self.timestep_var
                )  # timestep mean and variance
            else:
                t += self.dt

        # Plotting
        if visualize:
            fig, axs = plt.subplots(
                len(self.true_positions), 1, figsize=(10, 6 * len(self.true_positions))
            )
            for idx, (id, true_positions) in enumerate(true_positions_history.items()):
                estimated_positions = estimated_positions_history.get(id, [])
                max_length = max(len(true_positions), len(estimated_positions))
                true_positions = true_positions[:max_length]
                estimated_positions = estimated_positions[:max_length]

                axs[idx].plot(
                    range(1, max_length + 1),
                    [pos[0][0] for pos in true_positions],
                    label=f"True Position {id}",
                    linestyle="--",
                    marker="o",
                )
                axs[idx].plot(
                    range(1, max_length + 1),
                    [pos[0][0] for pos in estimated_positions],
                    label=f"Estimated Position {id}",
                    linestyle="-",
                    marker="x",
                )
                axs[idx].set_xlabel("Time Steps")
                axs[idx].set_ylabel("Position")
                axs[idx].set_title(f"True and Estimated Positions for Object {id}")
                axs[idx].legend()
