import numpy as np
import matplotlib.pyplot as plt
import json
from pathlib import Path

def analyze_session(session_path):
    session_path = Path(session_path)
    radar_types = ['awrl1432', 'radarbook']
    
    # 1. Load Metrics
    stats = {}
    for r in radar_types:
        metrics_file = session_path / r / "metrology" / "metrics.jsonl"
        if metrics_file.exists():
            with open(metrics_file, 'r') as f:
                stats[r] = [json.loads(line) for line in f]
    
    # 2. Pick a frame to analyze (e.g., middle of braking)
    # Let's target frame 85 (where relative velocity is near zero)
    frame_idx = 85
    frame_name = f"frame_{frame_idx:06d}"
    
    plt.figure(figsize=(15, 10))
    plt.suptitle(f"Metrology Diagnostic: {session_path.name} (Frame {frame_idx})", fontsize=16)
    
    for i, r in enumerate(radar_types):
        rd_path = session_path / r / "metrology" / "rd" / f"{frame_name}.npy"
        cfar_path = session_path / r / "metrology" / "cfar" / f"{frame_name}.npy"
        
        if not rd_path.exists():
            print(f"Warning: {rd_path} not found.")
            continue
            
        rd = np.load(rd_path)
        cfar = np.load(cfar_path)
        
        # Log scaling for visualization
        rd_log = 10 * np.log10(rd + 1e-9)
        cfar_log = 10 * np.log10(cfar + 1e-9)
        
        # Plot Range-Doppler Heatmap
        plt.subplot(2, 2, i+1)
        im = plt.imshow(np.flipud(rd_log), aspect='auto', cmap='viridis')
        plt.title(f"{r.upper()} - Range-Doppler Energy")
        plt.colorbar(im, label='dB')
        plt.ylabel("Range Bins")
        plt.xlabel("Doppler Bins (Center=0)")
        
        # Plot SNR (Signal / Threshold)
        # This shows what the CFAR "sees" as a potential target
        plt.subplot(2, 2, i+3)
        snr_map = rd_log - cfar_log
        im = plt.imshow(np.flipud(snr_map), aspect='auto', cmap='RdYlGn', vmin=-10, vmax=20)
        plt.title(f"{r.upper()} - SNR over CFAR Threshold")
        plt.colorbar(im, label='dB')
        plt.ylabel("Range Bins")
        plt.xlabel("Doppler Bins")

    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    
    # Save the analysis plot
    output_img = session_path / "metrology_comparison.png"
    plt.savefig(output_img)
    print(f"\n[SUCCESS] Analysis image saved to: {output_img}")
    
    # Print SNR stats comparison
    print("\n--- Signal Statistics (Frame 85) ---")
    for r in radar_types:
        if r in stats and len(stats[r]) > frame_idx:
            m = stats[r][frame_idx]
            print(f"[{r.upper()}] Peak SNR: {m.get('peak_snr_db', 0):.2f} dB | Active Bins: {m.get('active_bins', 0)}")

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