import numpy as np
import json
from pathlib import Path

def get_peak_tracking(session_path):
    session_path = Path(session_path)
    radar_types = ['awrl1432', 'radarbook']
    
    results = {}
    for r in radar_types:
        rd_dir = session_path / r / "metrology" / "rd"
        if not rd_dir.exists(): continue
        
        frames = sorted(list(rd_dir.glob("*.npy")))
        track = []
        for f_path in frames:
            rd = np.load(f_path)
            # Find coordinates of the max energy
            idx = np.unravel_index(np.argmax(rd), rd.shape)
            val = 10 * np.log10(np.max(rd) + 1e-9)
            track.append({"frame": f_path.stem, "range_bin": int(idx[0]), "doppler_bin": int(idx[1]), "db": float(val)})
        results[r] = track

    print("\n--- Peak Tracking Comparison (First 20 frames) ---")
    header = f"{'Frame':<15} | {'1432 R-Bin':<12} | {'RB R-Bin':<12} | {'1432 dB':<10} | {'RB dB':<10}"
    print(header)
    print("-" * len(header))
    
    for i in range(min(20, len(results.get('awrl1432', [])))):
        f14 = results['awrl1432'][i]
        frb = results['radarbook'][i]
        print(f"{f14['frame']:<15} | {f14['range_bin']:<12} | {frb['range_bin']:<12} | {f14['db']:<10.1f} | {frb['db']:<10.1f}")

if __name__ == "__main__":
    target_session = Path("data/braking_20260414_142359")
    if target_session.exists():
        get_peak_tracking(target_session)
    else:
        print(f"Session {target_session} not found.")