import numpy as np
import matplotlib.pyplot as plt

# Parâmetros do modelo
dt = 0.2       # intervalo de tempo (s)
v = 1.0        # velocidade linear (m/s)
horizonte = 50  # número de passos futuros

# Ponto inicial
x, y, theta = 0.0, 0.0, 0.0

# Trajetória alvo (trajeto desejado)
traj = np.array([
    [0, 0],
    [2, 0.5],
    [4, 1.5],
    [6, 3],
    [8, 5],
    [10, 7],
    [12, 9],
    [14, 10],
    [16, 10],
    [18, 9]
])

def calcular_orientacao(p1, p2):
    dx = p2[0] - p1[0]
    dy = p2[1] - p1[1]
    return np.arctan2(dy, dx)

# Trajetória gerada
traj_gerada = [(x, y)]

for step in range(horizonte):
    # Acha ponto alvo mais próximo
    distancias = np.linalg.norm(traj - np.array([x, y]), axis=1)
    idx_alvo = np.argmin(distancias)
    ponto_alvo = traj[min(idx_alvo + 1, len(traj) - 1)]

    # Orientação desejada
    orient_desejada = calcular_orientacao((x, y), ponto_alvo)
    delta_theta = orient_desejada - theta

    # Ajuste de ângulo para [-pi, pi]
    delta_theta = (delta_theta + np.pi) % (2 * np.pi) - np.pi

    # Calcular omega ideal
    omega = delta_theta / dt
    omega = np.clip(omega, -np.pi/4, np.pi/4)  # limitar a 45°/s

    # Aplicar movimento
    x += v * np.cos(theta) * dt
    y += v * np.sin(theta) * dt
    theta += omega * dt

    traj_gerada.append((x, y))

# Separar para plot
traj_gerada = np.array(traj_gerada)

# Plotar
plt.plot(traj[:,0], traj[:,1], 'ro--', label='Trajetória desejada')
plt.plot(traj_gerada[:,0], traj_gerada[:,1], 'bo-', label='Trajetória MPC (estimativa por passo)')
plt.title("MPC passo a passo com orientação estimada")
plt.xlabel("X")
plt.ylabel("Y")
plt.grid(True)
plt.legend()
plt.axis("equal")
plt.show()