import json

import pygame
from acados_template import AcadosSimSolver

from sovler_updates import *
from settings_acados_poly_track import *
from acados_settings_dev import acados_settings
from helpers.read_track import read_track
from helpers.race_visualization import RaceVisualization

from helpers.processing import *

track_file = "tracks/NXPCup2025.txt"
track = read_track(track_file)
[Sref, xRef, yRef, psiRef, kappa] = track

Tf = 1.2  # prediction horizon(s)
steps_per_s = 15
ref_velocity = 3.3  #m/s

degree = 5

N = (int)(Tf * steps_per_s)  # number of discretization steps
sref_N = Tf * ref_velocity  # reference for final reference progress

max_predict = 1.2  #maximum distance for which we consider curvatures
poly_estimate_dist = 1.5

# load model which will track polynom
constraint, model, poly_solver = setting_poly_track("bicycle", Tf, N, poly_degree=degree, max_predict=max_predict)

constraint, model, orig_solver, ocp = acados_settings(Tf, N, track)

nx = 6  #State dimensions
nu = 2  #Control dimensions

sim_solver = AcadosSimSolver(ocp, json_file="sim_ocp.json")

track_len = Sref[-1]
print(f"tracklen: {track_len}")
print(f"simulator time increment: {sim_solver.T}")

visualization = RaceVisualization(2560 * 1.6 , 1440 * 1.6, track, track_width=0.55, timestep=1 / steps_per_s)
visualization.flip();

for i in range(5,0, -1):
    time.sleep(1)
    print(f"Start in: {i}")

times = []
times_o = []
for offset in np.arange(-0.17, 0.17 + 0.01, 0.01):

    print(f"Running with offset: {offset} in closed loop iteration.")

    finished = False
    finished_old = False

    curr_t = 0

    # progress, lateralpos, orientation, velocity, torque, steering
    real_x = np.array([0.20858, 0, 0.0, 0.0, 0.0, 0.0])
    real_x_o = np.array([0.20858, 0, 0.0, 0.0, 0.0, 0.0])

    real_trajector_o = []
    real_trajector_o.append(real_x)
    real_trajector = []
    real_trajector.append(real_x)

    for j in range(N):
        poly_solver.set(j, "x", np.array([sref_N * j / N, 0, 0, 0, 0, 0]))

    for j in range(N):
         orig_solver.set(j, "x", np.array([sref_N * j / N, 0, 0, 0, 0, 0]))

    # simulate
    while True:
        # Get position in global space and relevant track points - simulating detection

        x_poly, y_poly = poly_solver_update(real_x=real_x,
                                            poly_solver=poly_solver,
                                            track=track,
                                            poly_estimate_dist=poly_estimate_dist,
                                            degree=degree,
                                            N=N,
                                            sref_N=sref_N,
                                            )

        original_solver_update(real_x_o, orig_solver, sref_N, N)

        x_plan, u_plan, elapsed_time, _ = run_solver(poly_solver, N, nx, nu)
        x_plan_o, u_plan_o, elapsed_time_o,_ = run_solver(orig_solver, N,nx,nu)

        visualization.draw_track_aproximation((x_poly, y_poly))
        visualization.draw_plan( x_plan, (x_poly, y_poly))
        visualization.draw_car(real_x)
        visualization.draw_car(real_x_o, (0,255,0))

        visualization.flip(clear=True)

        real_x = sim_solver.simulate(real_x, u_plan[0]);
        real_x_o = sim_solver.simulate(real_x_o, u_plan_o[0]);

        curr_t += 1.0/steps_per_s

        if real_x[0] < track_len and not finished:
            real_trajector.append(real_x)
        elif not finished:
            finished = True
            finish_time = curr_t
        if real_x_o[0] < track_len and not finished_old:
            real_trajector_o.append(real_x_o)
        elif not finished_old:
            finished_old = True
            finish_time_o = curr_t

        if finished and finished_old:
            break;

        real_x[0] = real_x[0] % track_len
        real_x_o[0] = real_x_o[0] % track_len



    print(f"Finish times: orig: {finish_time_o} poly planner: {finish_time}")

    times.append(finish_time)
    times_o.append(finish_time_o)

    visualization.draw_trajectory(real_trajector, (255,0,0))
    visualization.draw_trajectory(real_trajector_o)
    visualization.flip(clear=True)

    for i in range(10, 0, -1):
        time.sleep(1)
        print(f"Next turn starts in: {i}")

data = {
    "poly_planner": times,
    "orig_planner": times_o,
}

# Save to JSON file
with open("race_times1.json", "w") as f:
    json.dump(data, f)


waiting = True
while waiting:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            waiting = False
        elif event.type == pygame.KEYDOWN:
            waiting = False

pygame.quit()
sys.exit()