import numpy as np
import matplotlib.pyplot as plt

def initial_state(num_positions):
    return np.ones(num_positions) / num_positions

def predict_state(prior_prob, velocity):
    num_positions = len(prior_prob)
    predicted_prob = np.zeros(num_positions)
    for i in range(num_positions):
        predicted_prob[i] = prior_prob[(i - velocity) % num_positions]
    return predicted_prob

def update_state(predicted_prob, observation, observation_noise):
    likelihood = np.exp(-0.5 * ((np.arange(len(predicted_prob)) - observation) / observation_noise)**2)
    posterior_prob = likelihood * predicted_prob
    posterior_prob /= np.sum(posterior_prob)  # Normalize
    return posterior_prob

# Simulation parameters
num_positions = 100
velocity = 1
observation_noise = 10
num_steps = 50

# Initialize
true_positions = np.zeros(num_steps)
observations = np.zeros(num_steps)
prior_prob = initial_state(num_positions)

# Test initial state
init_state_2 = np.zeros(num_positions)
init_state_2[int(num_positions/2)] = 1
print(init_state_2)
p_plus = predict_state(init_state_2, 1)
print(p_plus)

# Initialize
true_positions = np.zeros(num_steps)
observations = np.zeros(num_steps)
estimated_positions = np.zeros(num_steps)  # Store estimated positions

# Run the Bayesian update cycle
for step in range(num_steps):
    # True target position update (constant velocity)
    true_positions[step] = step * velocity

    # Generate noisy observation
    observations[step] = true_positions[step] + np.random.normal(0, observation_noise)

    # Bayesian prediction
    predicted_prob = predict_state(prior_prob, velocity)

    # Bayesian update
    prior_prob = update_state(predicted_prob, observations[step], observation_noise)

# Calculate Metrics
mse = np.mean((true_positions - estimated_positions) ** 2)

# Symmetric Mean Absolute Percentage Error (SMAPE)
smape_numerator = 2 * np.abs(estimated_positions - true_positions)
smape_denominator = np.abs(estimated_positions) + np.abs(true_positions)
smape = np.mean(np.nan_to_num(smape_numerator / smape_denominator))

# Mean Absolute Percentage Error (MAPE)
mape_numerator = np.abs(estimated_positions - true_positions)
mape_denominator = np.abs(true_positions)
mape = np.mean(np.nan_to_num(mape_numerator / mape_denominator)) * 100

r2 = 1 - np.sum((true_positions - estimated_positions) ** 2) / np.sum((true_positions - np.mean(true_positions)) ** 2)
mae = np.mean(np.abs(true_positions - estimated_positions))

# Print Metrics
print("Mean Squared Error (MSE):", mse)
print("Symmetric Mean Absolute Percentage Error (SMAPE):", smape)
print("Mean Absolute Percentage Error (MAPE):", mape)
print("R-squared (R2):", r2)
print("Mean Absolute Error (MAE):", mae)

# Plot the results
plt.figure(figsize=(12, 6))
plt.plot(np.arange(num_positions), prior_prob, label='Posterior Distribution')
plt.scatter(true_positions, np.ones_like(true_positions) * 0.01, color='r', marker='x', label='True Positions')
plt.scatter(observations, np.ones_like(observations) * 0.01, color='g', marker='o', label='Noisy Observations')
plt.title('Bayesian Multi-Target Tracking')
plt.xlabel('Position')
plt.ylabel('Probability')
plt.legend()
plt.show()
