import pygame
import numpy as np

# ==============================
#  SETUP
# ==============================
pygame.init()
WIDTH, HEIGHT = 900, 650
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("🚗 Kalman Filter Car Simulation")

font_small = pygame.font.SysFont("Consolas", 14)
font_large = pygame.font.SysFont("Consolas", 20, bold=True)
clock = pygame.time.Clock()

# Colors
WHITE = (255, 255, 255)
BLACK = (10, 10, 20)
GREEN = (80, 255, 80)
BLUE = (80, 180, 255)
RED = (255, 80, 80)
YELLOW = (255, 255, 100)
PANEL_BG = (30, 30, 50)
GRAY = (150, 150, 180)

# ==============================
#  KALMAN FILTER SETUP
# ==============================
dt = 0.1

A = np.array([[1, 0, dt, 0],
              [0, 1, 0, dt],
              [0, 0, 1, 0],
              [0, 0, 0, 1]])

H = np.array([[1, 0, 0, 0],
              [0, 1, 0, 0]])

Q = np.eye(4) * 0.01
R = np.eye(2) * 25

x_est = np.array([[WIDTH/2], [HEIGHT/2], [0], [0]])
x_true = np.copy(x_est)
P = np.eye(4) * 500

trail_true, trail_meas, trail_est = [], [], []

# ==============================
#  FUNCTIONS
# ==============================
def predict(x, P):
    x_pred = A @ x
    P_pred = A @ P @ A.T + Q
    return x_pred, P_pred

def update(x_pred, P_pred, z):
    S = H @ P_pred @ H.T + R
    K = P_pred @ H.T @ np.linalg.inv(S)
    y = z - H @ x_pred
    x_upd = x_pred + K @ y
    P_upd = (np.eye(4) - K @ H) @ P_pred
    return x_upd, P_upd

def draw_text(text, pos, color=WHITE, font=font_small):
    surf = font.render(text, True, color)
    screen.blit(surf, pos)

def draw_trail(trail, color):
    if len(trail) > 2:
        pygame.draw.lines(screen, color, False, trail, 2)

# ==============================
#  MAIN LOOP
# ==============================
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    keys = pygame.key.get_pressed()
    if keys[pygame.K_ESCAPE]:
        running = False

    # Control car
    if keys[pygame.K_UP]:
        x_true[3] -= 2
    if keys[pygame.K_DOWN]:
        x_true[3] += 2
    if keys[pygame.K_LEFT]:
        x_true[2] -= 2
    if keys[pygame.K_RIGHT]:
        x_true[2] += 2

    # True motion
    x_true = A @ x_true

    # Noisy measurement
    z = H @ x_true + np.random.multivariate_normal([0, 0], R).reshape(2, 1)

    # Prediction
    x_pred, P_pred = predict(x_est, P)

    # Update
    x_est, P = update(x_pred, P_pred, z)

    # Store trails
    trail_true.append((int(x_true[0, 0]), int(x_true[1, 0])))
    trail_meas.append((int(z[0, 0]), int(z[1, 0])))
    trail_est.append((int(x_est[0, 0]), int(x_est[1, 0])))

    if len(trail_true) > 150:
        trail_true.pop(0)
        trail_meas.pop(0)
        trail_est.pop(0)

    # ==============================
    #  DRAW
    # ==============================
    screen.fill(BLACK)

    # Draw trails
    draw_trail(trail_true, (0, 180, 0))
    draw_trail(trail_meas, (180, 0, 0))
    draw_trail(trail_est, (0, 100, 255))

    # === UNCERTAINTY CIRCLE (based on covariance P) ===
    P_pos = P[0:2, 0:2]
    uncertainty = float(np.sqrt(np.trace(P_pos))) * 0.2  # smooth scaling
    uncertainty = np.clip(uncertainty, 5, 120)

    # Draw uncertainty circle for current estimate
    pygame.draw.circle(screen, BLUE,
                       (int(x_est[0, 0]), int(x_est[1, 0])),
                       int(uncertainty), 2)

    # Draw car positions
    pygame.draw.circle(screen, GREEN, (int(x_true[0, 0]), int(x_true[1, 0])), 7)  # True
    pygame.draw.circle(screen, RED, (int(z[0, 0]), int(z[1, 0])), 5)              # Measured
    pygame.draw.circle(screen, BLUE, (int(x_est[0, 0]), int(x_est[1, 0])), 7, 2)  # Estimated

    # Panel
    pygame.draw.rect(screen, PANEL_BG, (10, 10, 310, 260), border_radius=10)
    pygame.draw.rect(screen, GRAY, (10, 10, 310, 260), 2, border_radius=10)

    draw_text("KALMAN FILTER STATUS", (30, 20), YELLOW, font_large)
    draw_text(f"True Pos:      ({x_true[0,0]:7.2f}, {x_true[1,0]:7.2f})", (30, 60), GREEN)
    draw_text(f"Measured Pos:  ({z[0,0]:7.2f}, {z[1,0]:7.2f})", (30, 80), RED)
    draw_text(f"Estimated Pos: ({x_est[0,0]:7.2f}, {x_est[1,0]:7.2f})", (30, 100), BLUE)

    error = np.linalg.norm(x_true[:2] - x_est[:2])
    draw_text(f"Prediction Error: {error:7.2f} px", (30, 130), YELLOW)
    draw_text(f"Uncertainty: {uncertainty:7.2f} px", (30, 150), YELLOW)

    draw_text("Controls: ⬆️ ⬇️ ⬅️ ➡️ | ESC to quit", (30, 200), (180, 180, 255))

    pygame.display.flip()
    clock.tick(30)

pygame.quit()