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

# 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=100):
    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
    # We want to find a floor that keeps the car but kills trees.
    # Normalization factor is ~0.0003. Initial P_inc is ~1-100.
    thresholds = [0.0, 5e-5, 1e-4, 5e-4, 1e-3, 5e-3]
    
    results = []
    artifacts_dir = Path(r"C:\Users\rakadu1\.gemini\antigravity\brain\67913a3c-cbc2-4fba-87e3-88fbea20f043\artifacts")
    artifacts_dir.mkdir(parents=True, exist_ok=True)

    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=10, vmin=0, vmax=100)
            plt.colorbar(label='Magnitude')
        
        plt.title(f"Sensitivity Floor: {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_thresh_{i}.png"
        plt.savefig(save_path)
        plt.close()
        print(f"     [DONE] Result counts: {count} pts")

    # Final summary plot
    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    plt.plot([str(t) for t in thresholds], [r['count'] for r in results], marker='o', color='blue')
    plt.title("Detected Point Count")
    plt.ylabel("Reflections")
    plt.xticks(rotation=45)

    plt.subplot(1, 2, 2)
    plt.plot([str(t) for t in thresholds], [r['mean_mag'] for r in results], marker='o', color='red')
    plt.title("Avg SNR Magnitude")
    plt.ylabel("Magnitude")
    plt.xticks(rotation=45)

    plt.tight_layout()
    plt.savefig(artifacts_dir / "sweep_summary.png")
    plt.close()

if __name__ == "__main__":
    run_sensitivity_sweep()