import numpy as np
import json
from pathlib import Path

def track_full_state(session_path):
    session_path = Path(session_path)
    radar_types = ['awrl1432', 'radarbook']
    
    # Load axes
    axes = {}
    for r in radar_types:
        met_base = session_path / r / "metrology"
        axes[r] = {
            "range": np.load(met_base / "range_axis.npy"),
            "angle": np.load(met_base / "angle_axis.npy")
        }

    print(f"\n--- FULL STATE PEAK TRACKING: {session_path.name} ---")
    header = f"{'Frame':<15} | {'Radar':<10} | {'Range(m)':<10} | {'Angle(deg)':<10} | {'Mag(dB)':<10}"
    print(header)
    print("-" * len(header))

    for frame_idx in range(6, 16): # Core samples
        frame_name = f"frame_{frame_idx:06d}"
        
        for r in radar_types:
            rd_path = session_path / r / "metrology" / "rd" / f"{frame_name}.npy"
            ra_path = session_path / r / "metrology" / "ra" / f"{frame_name}.npy"
            
            if not rd_path.exists() or not ra_path.exists(): continue
            
            rd = np.load(rd_path)
            ra = np.load(ra_path)
            
            # Peak in RD (Range/Doppler)
            rd_idx = np.unravel_index(np.argmax(rd), rd.shape)
            r_bin = rd_idx[0]
            r_m = axes[r]["range"][r_bin]
            peak_db = 10 * np.log10(np.max(rd) + 1e-9)
            
            # Peak in RA (Range/Azimuth)
            ra_idx = np.unravel_index(np.argmax(ra), ra.shape)
            a_bin = ra_idx[1]
            a_deg = np.degrees(axes[r]["angle"][a_bin])
            
            name = "1432" if r == 'awrl1432' else "R-Book"
            print(f"{frame_name:<15} | {name:<10} | {r_m:<10.2f} | {a_deg:<10.2f} | {peak_db:<10.1f}")
        print("-" * 15)

if __name__ == "__main__":
    track_full_state("data/braking_20260414_142359")