# MPC com MQTT, simulação e visualização da trajetória
import numpy as np
import paho.mqtt.client as mqtt
import json
import matplotlib.pyplot as plt

# Parâmetros do MPC
angulo_max_graus = 30.0
distancia_entre_eixos = 0.92
limiar_entrada = 0.7
horizonte = 20

# Estado compartilhado
trajetoria_latlon = []
trajetoria_xy = []
visitados_execucao = []
lat0, lon0 = None, None

# Visualização
plt.ion()
fig, ax = plt.subplots()

# Funções matemáticas auxiliares
def latlon_to_xy(lat, lon, lat0, lon0):
    R = 6371000
    dlat = np.radians(lat - lat0)
    dlon = np.radians(lon - lon0)
    x = R * dlon * np.cos(np.radians((lat + lat0) / 2))
    y = R * dlat
    return x, y

def normalizar_angulo(angulo):
    if angulo < 0:
        angulo += 360
    angulo %= 360
    return angulo
        
def calcular_orientacao(p1, p2):
    dx = p2[0] - p1[0]
    dy = p2[1] - p1[1]
    orient = np.arctan2(dy, dx)
    orient -= np.radians(90)  # Corrige a base angular para bater com o seu theta ajustado
    return (orient + np.pi) % (2 * np.pi) - np.pi  # Normaliza entre [-pi, pi]

def calcular_orientacao_gps(p1_lat, p1_lon, p2_lat, p2_lon):
    lat1 = np.radians(p1_lat)
    lon1 = np.radians(p1_lon)
    lat2 = np.radians(p2_lat)
    lon2 = np.radians(p2_lon)

    delta_lon = lon2 - lon1

    y = np.sin(delta_lon) * np.cos(lat2)
    x = np.cos(lat1) * np.sin(lat2) - np.sin(lat1) * np.cos(lat2) * np.cos(delta_lon)

    theta = np.arctan2(y, x)
    bearing = (np.degrees(theta) + 360) % 360

    return bearing

def calcular_omega(v, angulo_rad, tipo):
    tan_delta = np.tan(angulo_rad)
    if tipo == "rodasFrontais":
        return (v * tan_delta) / distancia_entre_eixos
    elif tipo == "movimentoArco":
        return (2 * v * tan_delta) / distancia_entre_eixos
    return 0.0

def corrigir_pontos_visitados(x, y, limite_max_avanço=5):
    global visitados_execucao

    for idx, ponto in enumerate(trajetoria_xy):
        if visitados_execucao[idx]:
            continue

        dist = np.linalg.norm([x - ponto[0], y - ponto[1]])

        # Dentro do limiar: marca esse e os anteriores como visitados
        if dist < limiar_entrada:
            for i in range(idx + 1):
                visitados_execucao[i] = True
            return idx

        # Se já passou do limite de avanço, para
        if idx > 0 and not visitados_execucao[idx - 1] and idx >= limite_max_avanço:
            break

    # Se não marcou nenhum ponto, volta o próximo não visitado
    return proximo_nao_visitado(visitados_execucao)

def proximo_nao_visitado(visitados):
    for i, v in enumerate(visitados):
        if not v:
            return i
    return len(visitados) - 1

def calcular_melhor_candidato(x, y, theta, visitados_sim):
    idx_alvo = corrigir_pontos_visitados(x, y)
    if idx_alvo >= len(trajetoria_xy):
        return 0.0, "rodasFrontais"

    ponto_alvo = trajetoria_xy[idx_alvo]

    # Verifica se já alcançou o ponto
    if np.linalg.norm([x - ponto_alvo[0], y - ponto_alvo[1]]) < limiar_entrada:
        visitados_sim[idx_alvo] = True
        idx_alvo = corrigir_pontos_visitados(x, y)
        if idx_alvo >= len(trajetoria_xy):
            return 0.0, "rodasFrontais"
        ponto_alvo = trajetoria_xy[idx_alvo]

    orient = calcular_orientacao((x, y), ponto_alvo)
    erro_ori = abs((orient - theta + np.pi) % (2 * np.pi) - np.pi)
    #delta_theta = (orient - theta + np.pi) % (2 * np.pi) - np.pi
    delta_theta = ((-orient) - theta + np.pi) % (2 * np.pi) - np.pi
    angulo = np.clip(delta_theta, -np.radians(angulo_max_graus), np.radians(angulo_max_graus))

    # ✅ Se o erro for grande, prioriza curvas fechadas
    if erro_ori > np.radians(10):
        tipos_validos = ["movimentoArco"]
    else:
        tipos_validos = ["rodasFrontais", "movimentoArco"]

    melhor_custo = float('inf')
    melhor_tipo = "rodasFrontais"
    melhor_angulo = angulo

    for tipo in tipos_validos:
        omega = calcular_omega(1.0, angulo, tipo)
        x_sim = x + 1.0 * np.cos(theta) * 0.2
        y_sim = y + 1.0 * np.sin(theta) * 0.2
        theta_sim = theta + omega * 0.2

        erro_pos = np.linalg.norm([x_sim - ponto_alvo[0], y_sim - ponto_alvo[1]])
        orient_sim = calcular_orientacao((x_sim, y_sim), ponto_alvo)
        erro_ori = np.abs((orient_sim - theta_sim + np.pi) % (2 * np.pi) - np.pi)

        vetor_alvo = np.array(ponto_alvo) - np.array([x_sim, y_sim])
        vetor_alvo_norm = vetor_alvo / np.linalg.norm(vetor_alvo)
        vetor_movel = np.array([np.cos(theta_sim), np.sin(theta_sim)])
        cos_angulo = np.dot(vetor_alvo_norm, vetor_movel)

        # Se o ângulo for > 90° → cosseno < 0 → indo no sentido oposto
        penalidade_orientacao = 9999 if cos_angulo < 0 else 0

        custo = erro_pos * 2.0 + erro_ori * 1.5 + penalidade_orientacao

        if custo < melhor_custo:
            melhor_custo = custo
            melhor_tipo = tipo
            melhor_angulo = angulo

    return melhor_angulo, melhor_tipo

def processar_mpc(pos_lat, pos_lon, theta, velocidade, dt, angulo_controle_atual):
    print(f"Posicao recebida: Velocidade={velocidade}, dt={dt}, Orientacao={np.degrees(theta)}")
    global trajetoria_xy, visitados_execucao, lat0, lon0
    if not trajetoria_xy:
        return None

    x, y = latlon_to_xy(pos_lat, pos_lon, lat0, lon0)

    visitados_sim = visitados_execucao.copy()
    sim = [(x, y)]
    x_temp, y_temp, theta_temp = x, y, theta
    for _ in range(horizonte):
        angulo_mpc, tipo_mpc = calcular_melhor_candidato(x_temp, y_temp, theta_temp, visitados_sim)
        omega = calcular_omega(velocidade, angulo_mpc, tipo_mpc)

        # Corrige o ângulo apenas para visualização
        theta_plot = theta_temp + np.radians(90)

        # Avança a simulação com o ângulo corrigido para o traço
        x_temp += velocidade * np.cos(theta_plot) * dt
        y_temp += velocidade * np.sin(theta_plot) * dt
        theta_temp += omega * dt

        sim.append((x_temp, y_temp))

    # Atualiza visualização
    ax.clear()
    if trajetoria_xy:
        tx, ty = zip(*trajetoria_xy)
        ax.plot(tx, ty, 'r.-', label="Trajetória alvo")
    if sim:
        sx, sy = zip(*sim)
        ax.plot(sx, sy, 'g:', label="Previsão MPC")
    ax.plot(sim[0][0], sim[0][1], 'bo', label="Posição atual")
    ax.set_title("Visualização MPC (tempo real)")
    ax.set_xlabel("X (m)")
    ax.set_ylabel("Y (m)")
    ax.axis("equal")
    ax.legend()
    plt.pause(0.001)

    angulo_final, tipo_final = calcular_melhor_candidato(sim[0][0], sim[0][1], theta, visitados_execucao)
    #angulo_final = angulo_final * -1
    print(f"angulo: {np.degrees(angulo_final)}, tipo: {tipo_final}")
    return {
        "angulo": np.degrees(angulo_final),
        "tipo": 0 if tipo_final == "rodasFrontais" else 3,
        "velocidade": velocidade,
        "dt": dt,
        "orientacao": np.degrees(theta)
    }

# MQTT setup
def on_connect(client, userdata, flags, rc):
    print("Conectado ao MQTT")
    client.subscribe("mpc/rota")
    client.subscribe("mpc/posicao")

def on_message(client, userdata, msg):
    global trajetoria_latlon, trajetoria_xy, visitados_execucao, lat0, lon0
    if msg.topic == "mpc/rota":
        dados = json.loads(msg.payload.decode())
        trajetoria_latlon = dados["pontos"]
        lat0, lon0 = trajetoria_latlon[0][0], trajetoria_latlon[0][1]
        trajetoria_xy = [latlon_to_xy(lat, lon, lat0, lon0) for lat, lon in trajetoria_latlon]
        visitados_execucao = [False] * len(trajetoria_xy)
        print(f"Trajetória recebida com {len(trajetoria_xy)} pontos")

    elif msg.topic == "mpc/posicao":
        dados = json.loads(msg.payload.decode())
        lat = dados["lat"]
        lon = dados["lon"]
        #angulo = normalizar_angulo(dados["theta"] - 90)
        theta = np.radians(dados["theta"])
        velocidade = dados["velocidade"]
        dt = dados["dt"]
        angulo_controle = dados.get("anguloControle", 0.0)
        comando = processar_mpc(lat, lon, theta, velocidade, dt, angulo_controle)
        if comando:
            client.publish("mpc/comando", json.dumps(comando))

client = mqtt.Client()
client.on_connect = on_connect
client.on_message = on_message
client.connect("localhost", 1883, 60)

print("Aguardando trajetória e posição...")
client.loop_forever()