import os
import json
import numpy as np
import matplotlib.pyplot as plt

DATA_PATH = "data"


# ---------------- LOAD FRAME METADATA ----------------
def load_frame(frame_id):
    with open(os.path.join(DATA_PATH, "frames.jsonl")) as f:
        for line in f:
            data = json.loads(line)
            if data["frame_id"] == frame_id:
                return data
    return None


# ---------------- RADAR ----------------
def inspect_radar(frame_id):
    path = os.path.join(DATA_PATH, "radar", f"frame_{frame_id:06d}.npy")
    data = np.load(path)

    print("\n--- RADAR DATA ---")
    print("Shape:", data.shape)

    if len(data) > 0:
        print("\nFirst 5 detections:")
        for row in data[:5]:
            print(f"Depth={row[0]:.2f} m | Azimuth={row[1]:.3f} rad | Vel={row[3]:.2f} m/s")
    else:
        print("No radar detections")

    return data


# ---------------- LIDAR ----------------
def inspect_lidar(frame_id):
    path = os.path.join(DATA_PATH, "lidar", f"frame_{frame_id:06d}.npy")
    data = np.load(path)

    print("\n--- LIDAR DATA ---")
    print("Shape:", data.shape)

    print("\nFirst 5 points:")
    for row in data[:5]:
        print(f"x={row[0]:.2f}, y={row[1]:.2f}, z={row[2]:.2f}")

    return data


# ---------------- GROUND TRUTH ----------------
def inspect_gt(frame_id):
    frame = load_frame(frame_id)

    print("\n--- GROUND TRUTH ---")

    gt = frame["ground_truth"]

    print(f"Total vehicles: {len(gt)}")

    for v in gt[:5]:
        print(f"ID={v['id']} | x={v['x']:.2f}, y={v['y']:.2f} | speed={v['speed']:.2f}")


# ---------------- PLOT RADAR ----------------
def plot_radar(radar_data):
    if len(radar_data) == 0:
        print("No radar data to plot")
        return

    r = radar_data[:, 0]
    az = radar_data[:, 1]

    x = r * np.cos(az)
    y = r * np.sin(az)

    plt.figure()
    plt.scatter(x, y, s=10)
    plt.title("Radar XY")
    plt.xlabel("X (m)")
    plt.ylabel("Y (m)")
    plt.axis("equal")
    plt.grid()
    plt.show()


# ---------------- PLOT LIDAR ----------------
def plot_lidar(lidar_data):
    x = lidar_data[:, 0]
    y = lidar_data[:, 1]

    plt.figure()
    plt.scatter(x, y, s=1)
    plt.title("Lidar Top View")
    plt.axis("equal")
    plt.grid()
    plt.show()


# ---------------- MAIN ----------------
if __name__ == "__main__":
    frame_id = int(input("Enter frame number: "))

    frame = load_frame(frame_id)

    if frame is None:
        print("Frame not found")
        exit()

    print("\n========== FRAME INFO ==========")
    print("Frame ID:", frame["frame_id"])
    print("Timestamp:", frame["timestamp"])

    inspect_gt(frame_id)

    radar = inspect_radar(frame_id)
    lidar = inspect_lidar(frame_id)

    # Optional plots
    plot_choice = input("\nPlot data? (y/n): ").lower()

    if plot_choice == "y":
        plot_radar(radar)
        plot_lidar(lidar)