import matplotlib.pyplot as plt
import numpy as np

from hybrid_tools import SKF, HybridSimulator
from hybrid_tools.basic_hybrid_systems import bouncing_ball

""" Define dynamics and resets. """
hybrid_system = bouncing_ball()

""" Initialize states and covariance. """
# Start ball at height 2.5m with zero velocity
n_states = 2
mean_init_state = np.array([2.5, 0])
mean_init_cov = 0.1 * np.eye(n_states)
init_mode = "I"  # Modes are {I, J}

""" Define timesteps. """
dt = 0.005

""" Define parameters. """
# Parameters: [e (coefficient of restitution), g (gravity)]
# e = 0.8 means ball retains 80% of velocity after bounce
# g = 9.81 m/s^2 (standard gravity)
parameters = np.array([0.8, 9.81])

""" Initialize filter. """
skf = SKF(
    init_state=mean_init_state,
    init_mode=init_mode,
    init_cov=mean_init_cov,
    dt=dt,
    parameters=parameters,
    hybrid_system=hybrid_system,
)

""" Initialize simulator. """
actual_init_state = np.random.multivariate_normal(mean_init_state, mean_init_cov)
hybrid_simulator = HybridSimulator(
    init_state=actual_init_state,
    init_mode=init_mode,
    dt=dt,
    parameters=parameters,
    hybrid_system=hybrid_system,
)

n_simulate_timesteps = 500
timesteps = np.arange(0.0, n_simulate_timesteps * dt, dt)
measurements = np.zeros((n_simulate_timesteps - 1, n_states))
actual_states = np.zeros((n_simulate_timesteps, n_states))
filtered_states = np.zeros((n_simulate_timesteps, n_states))

actual_states[0, :] = hybrid_simulator.get_state()
filtered_states[0, :] = mean_init_state

mode_list = []
actual_mode_list = []

zero_input = np.array([0.0])
# Larger steps between posterior updates lead to less mode mismatches.
# TODO: we can apply a rule where if the residual is large
# - If we recently transitioned, we should check what happened if we didn't.
# - If we haven't transitioned, but are in some neighborhood of guards, we should check what happens if we simulate more in time until the guard is hit.
posterior_update_frequency = 10
for time_idx in range(1, n_simulate_timesteps):
    hybrid_simulator.simulate_timestep(0, np.array([0]))
    actual_states[time_idx, :] = hybrid_simulator.get_state()
    measurements[time_idx - 1, :] = hybrid_simulator.get_measurement(measurement_noise_flag=True)
    skf.predict(timesteps[time_idx], zero_input)
    if time_idx % posterior_update_frequency == 0:
        _, _ = skf.update(timesteps[time_idx], zero_input, measurements[time_idx - 1, :])
    filtered_states[time_idx, :] = skf.current_state
    mode_list.append(hybrid_simulator.current_mode)

# Phase plot
plt.figure(figsize=(10, 6))
plt.plot(actual_states[:, 0], actual_states[:, 1], "k-", label="Actual states")
plt.plot(measurements[:, 0], measurements[:, 1], "r.", label="Measurements", markersize=2)
plt.plot(filtered_states[:, 0], filtered_states[:, 1], "b--", label="Filtered states")
plt.xlabel("Position (m)")
plt.ylabel("Velocity (m/s)")
plt.title("Bouncing Ball Phase Plot")
plt.legend()
plt.grid(True, alpha=0.3)
plt.savefig("bouncing_ball.png")
plt.close()

# Time series plots with mode highlighting
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), sharex=True)

# Highlight mode regions
mode_colors = {"I": "lightblue", "J": "lightcoral"}
mode_labels = {"I": "Mode I", "J": "Mode J"}
current_mode = mode_list[0] if mode_list else init_mode
start_idx = 1
labeled_modes = set()  # Track which modes have been labeled

for idx in range(1, len(mode_list)):
    if mode_list[idx] != current_mode or idx == len(mode_list) - 1:
        end_idx = idx if mode_list[idx] != current_mode else idx + 1
        # Add label only if this mode hasn't been labeled yet
        label = mode_labels[current_mode] if current_mode not in labeled_modes else None
        if label:
            labeled_modes.add(current_mode)
        ax1.axvspan(
            timesteps[start_idx],
            timesteps[end_idx],
            alpha=0.3,
            color=mode_colors[current_mode],
            label=label,
        )
        ax2.axvspan(
            timesteps[start_idx],
            timesteps[end_idx],
            alpha=0.3,
            color=mode_colors[current_mode],
            label=label,
        )
        current_mode = mode_list[idx]
        start_idx = idx

# Plot position
ax1.plot(timesteps, actual_states[:, 0], "k-", label="Actual", linewidth=1.5)
ax1.plot(timesteps[:-1], measurements[:, 0], "r.", label="Measurements", markersize=2)
ax1.plot(timesteps, filtered_states[:, 0], "b--", label="Filtered", linewidth=1.5)
ax1.set_ylabel("Position (m)", fontsize=12)
ax1.set_title("Bouncing Ball Time Series", fontsize=14, fontweight="bold")
ax1.legend(loc="upper right")
ax1.grid(True, alpha=0.3)

# Plot velocity
ax2.plot(timesteps, actual_states[:, 1], "k-", label="Actual", linewidth=1.5)
ax2.plot(timesteps[:-1], measurements[:, 1], "r.", label="Measurements", markersize=2)
ax2.plot(timesteps, filtered_states[:, 1], "b--", label="Filtered", linewidth=1.5)
ax2.set_xlabel("Time (s)", fontsize=12)
ax2.set_ylabel("Velocity (m/s)", fontsize=12)
ax2.legend(loc="upper right")
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig("bouncing_ball_timeseries.png", dpi=150)
plt.close()
# plt.show()
