import os
import sys
import numpy as np
import tqdm
from pathlib import Path

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

# Import the model wrapper
try:
    from scripts.ISOLATE.model_wrapper import ShenronRadarModel
except ImportError as e:
    print(f"Error: Failed to import ShenronRadarModel. Ensure scripts/ISOLATE/model_wrapper.py exists. ({e})")
    sys.exit(1)

def process_session(session_path):
    print(f"\n>>> Processing session: {session_path.name}")
    
    lidar_dir = session_path / "lidar"
    if not lidar_dir.exists():
        print(f"    [SKIP] No 'lidar' folder found.")
        return
        
    # Find all .npy files in lidar/
    lidar_files = sorted(list(lidar_dir.glob("*.npy")))
    if not lidar_files:
        print(f"    [SKIP] No .npy files in 'lidar' folder.")
        return

    output_dir = session_path / "shenron_radar"
    output_dir.mkdir(exist_ok=True)
    
    # Initialize the model once per session
    print(f"    Initializing ShenronRadarModel...")
    model = ShenronRadarModel(radar_type='radarbook')
    
    print(f"    Generating Shenron Radar data for {len(lidar_files)} frames...")
    
    for lidar_file in tqdm.tqdm(lidar_files, desc="      Simulating Radar", unit="frame"):
        try:
            # 1. Load Semantic LiDAR data
            # Expected raw: [x, y, z, cos, obj, tag] (6 cols)
            # Expected Shenron input: [x, y, z, intensity, cos, obj, tag] (7 cols)
            data = np.load(lidar_file)
            
            if data.shape[1] == 6:
                # Pad with a dummy intensity column at index 3
                # This aligns 'tag' to index 6 as expected by our lidar.py mapping
                padded_data = np.zeros((data.shape[0], 7))
                padded_data[:, 0:3] = data[:, 0:3] # x, y, z
                padded_data[:, 4:7] = data[:, 3:6] # cos, obj, tag
                data = padded_data
            
            # 2. Process through the physics-based model
            # returns rich PCD: [M, 5] (x, y, z, velocity, magnitude)
            rich_pcd = model.process(data)
            
            # 3. Save to disk
            output_file = output_dir / lidar_file.name
            np.save(output_file, rich_pcd)
            
        except Exception as e:
            print(f"\n    [ERROR] Failed to process {lidar_file.name}: {e}")

def main():
    data_root = project_root / "data"
    if not data_root.exists():
        print(f"Error: {data_root} not found.")
        return
        
    # Get all session folders
    sessions = sorted([d for d in data_root.iterdir() if d.is_dir()])
    
    if not sessions:
        print("No simulation sessions found in data/.")
        return
        
    print(f"Found {len(sessions)} sessions.")
    
    for session in sessions:
        # Check if the session has frames.jsonl to confirm it's a valid data folder
        if (session / "frames.jsonl").exists():
            process_session(session)
        else:
            print(f"Skipping {session.name} (no frames.jsonl found).")
            
    print("\n" + "="*50)
    print("SHENRON BATCH PROCESSING COMPLETE!")
    print("="*50)

if __name__ == "__main__":
    main()