from deps.sim.agent import Car
from deps.sim.simulation import Sim
from deps.sim.sensors import Camera, Lidar, IMU
from deps.algorithms.Controller import KeyboardController
from deps.tools.timings import Timings
from deps.tools.utils import load_configs

import platform
import numpy as np
import matplotlib
system = platform.system() 
if system == "Linux":
    matplotlib.use("TkAgg")
elif system == "Darwin": # macOS
    pass
elif system == "Windows":
    matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
import cv2
from deps.utils import logodds_to_prob, lidar_noise, remove_out_of_range
from deps.utils import normalize_grid_logodds, normalize_grid_probability, draw_robot

############### PASTE FUNCTIONS INTO functions.py AND IMPORT THEM HERE ###############
from functions import bresenham_line_array, compute_lidar_hitpoints, OGM, prob_to_logodds


# Declare user-specific paths to files.
ENV_PATH = "configs/env/simple_env.yaml"
CAR_PATH = "configs/car/car_config.yaml"
CAR_URDF_PATH = "configs/resources/f10_racecar/racecar_ackermann.urdf"
MAP_PATH = "configs/map/map_config.yaml"

# FPS constants.
CTRL_FPS = 100 # Perform control at 100Hz
LIDAR_FPS = 15 # Simulate lidar at 15Hz (15Hz is for easier visualization, increase to 30Hz for better mapping)
CAMERA_FPS = 30 # Simulate camera at 30Hz
PRINT_FPS = 5 # Update map plot at 5Hz 

# Declare sensors.
SENSORS = [Camera, Lidar, IMU]

# Declare controller
controller = KeyboardController(ctrl_fps=CTRL_FPS)

if __name__ == "__main__":
    # Load LiDAR and mapping config data
    configs = load_configs(CAR_PATH)
    lidar_config = configs["lidar_configs"]
    num_rays = lidar_config["num_rays"]
    lidar_angle1 = lidar_config["lidar_angle1"] - 90 # Provided angle assumes car is orientated at 90 deg
    lidar_angle2 = lidar_config["lidar_angle2"] - 90
    lidar_angle_start = np.deg2rad(lidar_angle1) # start angle, default is -45 deg (45 degrees right of forward)
    lidar_angle_end = np.deg2rad(lidar_angle2)  # end angle, default is 45 deg (45 degrees left of forward)
    max_ray_len = lidar_config["ray_len"]

    map_configs = load_configs(MAP_PATH)
    res = map_configs["resolution"]
    x_min = map_configs["x_min"] 
    y_min = map_configs["y_min"] 
    x_max = map_configs["x_max"] 
    y_max = map_configs["y_max"] 

    p_hit = map_configs["lidar_accuracy"]
    logodds_hit = prob_to_logodds(p_hit) # logodds value to add/subtract for cells
    display_logodds_or_prob = map_configs["plot_logodds_or_prob"]

    noisy_lidar = map_configs["noisy_lidar"] 
    lidar_offset = map_configs["lidar_offset"] # LiDAR Pose in Robot Frame

    # Initialize simulation.
    sim = Sim(time_step_freq=120, debug=False)
    sim.create_env(env_config=ENV_PATH) # add extra argument to load file with map

    car = Car(
            urdf_path=CAR_URDF_PATH, 
            car_config=CAR_PATH, 
            sensors=SENSORS,
            debug=False
    )

    car.place_car(sim.floor)

    # Set sim response time.
    ctrl_time = Timings(CTRL_FPS)
    camera_time = Timings(CAMERA_FPS)
    print_frequency = Timings(PRINT_FPS)
    lidar_frequency = Timings(LIDAR_FPS)

    pose = None # Car's state (in this case, it's the car's pose)

    # Computing LiDAR ray angles
    lidar_ray_angles = np.linspace(lidar_angle_end, lidar_angle_start, num_rays)

    # Initializing OGM # Map Initialization Function
    capped_cell_logodds = 25*logodds_hit
    MAP = OGM(x_min, y_min, x_max, y_max, res, logodds_hit, capped_cell_logodds)
    MAP.initialize_OGM()

    # OpenCV Map Visualization Initialization
    cv2.namedWindow("Probabilistic OGM", cv2.WINDOW_NORMAL)
    if system == "Windows":
        cv2.resizeWindow("Probabilistic OGM", 800, 800)
    else: cv2.resizeWindow("Probabilistic OGM", 800, 600)
    
    # Creating Grayscale LUT
    gray_lut = np.arange(256, dtype=np.uint8).reshape(-1, 1)
    
    # Matplotlib Colorbar
    fig, ax = plt.subplots(figsize=(2, 6))
    ax.set_visible(False)
    if display_logodds_or_prob:
        c_min, c_max = -capped_cell_logodds, capped_cell_logodds
        c_label = "Logodds"
    else:
        c_min, c_max = 0.0, 1.0
        c_label = "Probability"
    im = plt.imshow(np.array([[]]), cmap="gray_r", vmin=c_min, vmax=c_max)
    cbar = plt.colorbar(im, label=c_label)
    plt.tight_layout()
    plt.ion()
    plt.show(block=False)
    plt.pause(0.1)

    while True:
        pose = car.get_state_rpy(to_array=False,radian=True)
        x, y, roll, pitch, yaw = pose

        rays_data, dists, coords = car.get_sensor_data("lidar", common=False) 
        car.simulate_sensor("lidar", rays_data)

        if print_frequency.update_time():
            # Normalize Grid for Color Map
            if display_logodds_or_prob:
                grid_plot = normalize_grid_logodds(MAP.grid.T, capped_cell_logodds)
            else:
                grid_plot = normalize_grid_probability(MAP.grid.T)

            # Apply Color Map
            grid_gray = cv2.LUT(grid_plot, gray_lut)
            grid_colored = cv2.cvtColor(grid_gray, cv2.COLOR_GRAY2BGR)

            # Draw Robot
            grid_colored = draw_robot(grid_colored, [x, y, yaw], res, x_min, y_min)

            # Update Map
            cv2.imshow("Probabilistic OGM", grid_colored)
            cv2.waitKey(1)
            plt.pause(0.001)
            
        if lidar_frequency.update_time():
            lidar_ray_ranges = dists

            # Removing out of range rays (needs to be done before noise)
            lidar_dists, lidar_angles, coords = remove_out_of_range(lidar_ray_ranges, lidar_ray_angles, coords, lower_bound=0.05, upper_bound=max_ray_len)

            # LiDAR ray data
            if noisy_lidar:
                lidar_dists = lidar_noise(lidar_dists)

            # This simulated LiDAR has a max range of 8 meters, any measurement of 8 meters or greater has been removed.
            # The LiDAR originally has 100 rays, the variables ray_ranges and ray_angles have already had out of range values removed.
            # ray_ranges and ray_angles are aligned (each index of both corresponds to the same ray).

            # Using your colab functions, update the map
            # Given: 
                # Pose of Robot: x, y, yaw
                # Ray Data: lidar_angles, lidar_dists (1 x n arrays, n is the number of in range values)
                # Parameters from above: res, x_min, x_max, y_min, y_max, p_hit, logodds_hit
                    # all angles and lengths are in radians and meters
            
            robot_pose = [x, y, yaw]
            
            # Frame Transformation Function: Converting ray data to endpoints in global frame 
            x_endpoints, y_endpoints = compute_lidar_hitpoints(robot_pose, lidar_dists, lidar_angles, lidar_offset)
            
            # Bresenham Function: Calculating bresenham between robot pose (position and angle) and each ray endpoint
            (x_frees, y_frees), (x_hits, y_hits) = bresenham_line_array(x, y, x_endpoints, y_endpoints, res)

            # Map Update Function: Updating free squares and hit squares in the grid 
            MAP.update_OGM(x_frees, y_frees, x_hits, y_hits)

        if ctrl_time.update_time():
            yaw = np.rad2deg(yaw)
            v, s = controller.navigate(x, y, yaw)    
           
            car.act(v, s)            

            sim.step() # Advance one time step in the sim
            sim.view(x,y, yaw, "distant")

        # End Simulation (press x in map window)
        key = cv2.waitKey(1) & 0xFF
        if key == ord("x"):
            break

    cv2.destroyAllWindows()
    plt.close(fig_cbar) 
        