import numpy as np
import os
import sys
from pathlib import Path
import matplotlib
matplotlib.use('Agg') # Headless
import matplotlib.pyplot as plt

# Add project paths
project_root = Path(__file__).parent.parent
sys.path.append(str(project_root / 'scripts' / 'ISOLATE'))

try:
    from model_wrapper import ShenronRadarModel
except ImportError as e:
    print(f"Error: {e}")
    sys.exit(1)

def run_sensitivity_sweep(frame_idx=207):
    lidar_path = project_root / 'Shenron_debug' / 'logs' / 'lidar' / f"frame_{frame_idx:06d}.npy"
    if not lidar_path.exists():
        print(f"[ERROR] Frame {lidar_path} not found.")
        return

    data = np.load(lidar_path)
    if data.shape[1] == 6:
        padded = np.zeros((data.shape[0], 7), dtype=np.float32)
        padded[:, 0:3] = data[:, 0:3]
        padded[:, 4:7] = data[:, 3:6]
        data = padded

    model = ShenronRadarModel(radar_type='awrl1432')
    
    # Thresholds to sweep for Frame 207
    thresholds = [0.0, 1e-5, 5e-5, 1e-4, 2e-4, 5e-4, 0.001]
    
    results = []
    artifacts_dir = Path(r"C:\Users\rakadu1\.gemini\antigravity\brain\67913a3c-cbc2-4fba-87e3-88fbea20f043\artifacts")

    print(f"\n🧪 Starting Sensitivity Sweep (Frame {frame_idx})...")
    
    for i, thresh in enumerate(thresholds):
        print(f"  -> Testing Threshold: {thresh}")
        model.radar_obj.sensitivity_floor = thresh
        
        # Process frame
        rich_pcd = model.process(data)
        
        count = rich_pcd.shape[0]
        mean_mag = np.mean(rich_pcd[:, 4]) if count > 0 else 0
        
        results.append({
            'threshold': thresh,
            'count': count,
            'mean_mag': mean_mag
        })
        
        # Plotting
        plt.figure(figsize=(10, 8))
        if count > 0:
            plt.scatter(rich_pcd[:, 1], rich_pcd[:, 0], c=rich_pcd[:, 4], cmap='viridis', s=12, vmin=0, vmax=100)
            plt.colorbar(label='Magnitude')
        
        plt.title(f"Threshold: {thresh} | Pts: {count} | Avg Mag: {mean_mag:.2f}")
        plt.xlabel("Side (Y)")
        plt.ylabel("Forward (X)")
        plt.xlim([-30, 30])
        plt.ylim([0, 100])
        plt.grid(True, alpha=0.3)
        
        save_path = artifacts_dir / f"sweep_207_thresh_{i}.png"
        plt.savefig(save_path)
        plt.close()
        print(f"     [DONE] Result counts: {count} pts")

if __name__ == "__main__":
    run_sensitivity_sweep(207)