"""
GENESIS Urban Layout Simulation Viewer
Loads OBJ models and runs Traffic, Rain, and Wind simulations using PyVista
"""

import numpy as np
import pyvista as pv
import json
import argparse
from pathlib import Path
from dataclasses import dataclass, field
from typing import List, Tuple, Dict, Optional, Any
from enum import Enum
import time


# ============================================================================
# SIMULATION MODE ENUM
# ============================================================================

class SimulationMode(Enum):
    NONE = 0
    TRAFFIC = 1
    RAIN = 2
    WIND = 3


# ============================================================================
# TRAFFIC SIMULATION
# ============================================================================

@dataclass
class TrafficAgent:
    """Single traffic agent (vehicle/pedestrian)"""
    position: np.ndarray
    velocity: np.ndarray
    road_index: int
    progress: float  # 0.0 to 1.0 along road
    speed: float
    agent_type: str  # 'car', 'bike', 'pedestrian'
    
    
class TrafficSimulation:
    """Traffic simulation with agents moving along roads"""
    
    def __init__(self, roads_data: List[Dict], num_agents: int = 50):
        self.roads = roads_data
        self.num_agents = num_agents
        self.agents: List[TrafficAgent] = []
        self.agent_colors = {
            'car': [1.0, 0.2, 0.2],      # Red
            'bike': [0.2, 1.0, 0.2],     # Green
            'pedestrian': [0.2, 0.2, 1.0] # Blue
        }
        self.paused = False
        self._initialize_agents()
        
    def _initialize_agents(self):
        """Create initial agents on roads"""
        self.agents = []
        
        if not self.roads:
            return
            
        agent_types = ['car', 'bike', 'pedestrian']
        type_speeds = {'car': 15.0, 'bike': 8.0, 'pedestrian': 2.0}
        type_weights = [0.5, 0.3, 0.2]
        
        for i in range(self.num_agents):
            # Random road selection weighted by length
            road_lengths = [r.get('length', 10) for r in self.roads]
            total_length = sum(road_lengths)
            if total_length == 0:
                continue
                
            probs = [l / total_length for l in road_lengths]
            road_idx = np.random.choice(len(self.roads), p=probs)
            road = self.roads[road_idx]
            
            # Random position along road
            progress = np.random.random()
            start = np.array(road['start_point'] + [0.5])  # Slight elevation
            end = np.array(road['end_point'] + [0.5])
            position = start + progress * (end - start)
            
            # Direction along road (randomly forward or backward)
            direction = 1 if np.random.random() > 0.5 else -1
            road_vec = end - start
            road_length = np.linalg.norm(road_vec)
            if road_length > 0:
                velocity = direction * road_vec / road_length
            else:
                velocity = np.array([1.0, 0.0, 0.0])
            
            # Agent type
            agent_type = np.random.choice(agent_types, p=type_weights)
            speed = type_speeds[agent_type] * (0.8 + 0.4 * np.random.random())
            
            agent = TrafficAgent(
                position=position,
                velocity=velocity * speed,
                road_index=road_idx,
                progress=progress,
                speed=speed,
                agent_type=agent_type
            )
            self.agents.append(agent)
    
    def update(self, dt: float = 0.016):
        """Update agent positions"""
        if self.paused or not self.roads:
            return
            
        for agent in self.agents:
            road = self.roads[agent.road_index]
            start = np.array(road['start_point'] + [0.5])
            end = np.array(road['end_point'] + [0.5])
            road_vec = end - start
            road_length = np.linalg.norm(road_vec)
            
            if road_length == 0:
                continue
            
            # Update progress along road
            direction = np.sign(np.dot(agent.velocity, road_vec))
            agent.progress += direction * agent.speed * dt / road_length
            
            # Handle road endpoints - switch to connected road or reverse
            if agent.progress > 1.0 or agent.progress < 0.0:
                # Find connected road or reverse
                current_end = end if agent.progress > 1.0 else start
                connected = self._find_connected_road(agent.road_index, current_end)
                
                if connected is not None:
                    agent.road_index = connected
                    new_road = self.roads[connected]
                    new_start = np.array(new_road['start_point'] + [0.5])
                    new_end = np.array(new_road['end_point'] + [0.5])
                    
                    # Determine entry point
                    dist_to_start = np.linalg.norm(current_end[:2] - new_start[:2])
                    dist_to_end = np.linalg.norm(current_end[:2] - new_end[:2])
                    
                    if dist_to_start < dist_to_end:
                        agent.progress = 0.0
                        new_vec = new_end - new_start
                    else:
                        agent.progress = 1.0
                        new_vec = new_start - new_end
                    
                    new_length = np.linalg.norm(new_vec)
                    if new_length > 0:
                        agent.velocity = new_vec / new_length * agent.speed
                else:
                    # Reverse direction
                    agent.velocity = -agent.velocity
                    agent.progress = np.clip(agent.progress, 0.0, 1.0)
            
            # Update actual position
            agent.position = start + agent.progress * road_vec
            agent.position[2] = 0.5  # Keep at road level
    
    def _find_connected_road(self, current_road_idx: int, 
                             endpoint: np.ndarray, 
                             tolerance: float = 5.0) -> Optional[int]:
        """Find a road connected to the given endpoint"""
        for i, road in enumerate(self.roads):
            if i == current_road_idx:
                continue
            
            start = np.array(road['start_point'])
            end = np.array(road['end_point'])
            
            if np.linalg.norm(endpoint[:2] - start) < tolerance:
                return i
            if np.linalg.norm(endpoint[:2] - end) < tolerance:
                return i
        
        return None
    
    def get_points_and_colors(self) -> Tuple[np.ndarray, np.ndarray]:
        """Get agent positions and colors for rendering"""
        if not self.agents:
            return np.array([[0, 0, 0]]), np.array([[1, 1, 1]])
        
        positions = np.array([a.position for a in self.agents])
        colors = np.array([self.agent_colors[a.agent_type] for a in self.agents])
        
        return positions, colors
    
    def reset(self):
        """Reset simulation"""
        self._initialize_agents()


# ============================================================================
# RAIN SIMULATION
# ============================================================================

class RainSimulation:
    """Rain particle simulation"""
    
    def __init__(self, bounds: Tuple[float, float, float, float], 
                 max_height: float = 100.0,
                 num_drops: int = 2000,
                 rain_speed: float = 50.0):
        self.bounds = bounds  # minx, miny, maxx, maxy
        self.max_height = max_height
        self.num_drops = num_drops
        self.rain_speed = rain_speed
        self.drops: np.ndarray = None
        self.velocities: np.ndarray = None
        self.paused = False
        self._initialize_drops()
    
    def _initialize_drops(self):
        """Initialize raindrop positions"""
        minx, miny, maxx, maxy = self.bounds
        
        # Random positions within bounds
        x = np.random.uniform(minx, maxx, self.num_drops)
        y = np.random.uniform(miny, maxy, self.num_drops)
        z = np.random.uniform(0, self.max_height, self.num_drops)
        
        self.drops = np.column_stack([x, y, z])
        
        # Velocities - mostly downward with slight variation
        vx = np.random.uniform(-2, 2, self.num_drops)
        vy = np.random.uniform(-2, 2, self.num_drops)
        vz = np.full(self.num_drops, -self.rain_speed)
        
        self.velocities = np.column_stack([vx, vy, vz])
    
    def update(self, dt: float = 0.016):
        """Update raindrop positions"""
        if self.paused:
            return
        
        # Move drops
        self.drops += self.velocities * dt
        
        # Reset drops that hit ground
        ground_mask = self.drops[:, 2] < 0
        if np.any(ground_mask):
            minx, miny, maxx, maxy = self.bounds
            num_reset = np.sum(ground_mask)
            
            self.drops[ground_mask, 0] = np.random.uniform(minx, maxx, num_reset)
            self.drops[ground_mask, 1] = np.random.uniform(miny, maxy, num_reset)
            self.drops[ground_mask, 2] = self.max_height
    
    def get_points_and_colors(self) -> Tuple[np.ndarray, np.ndarray]:
        """Get drop positions and colors"""
        # Color based on height - lighter at top, darker near ground
        heights = self.drops[:, 2] / self.max_height
        colors = np.column_stack([
            0.3 + 0.4 * heights,  # R
            0.5 + 0.3 * heights,  # G
            0.8 + 0.2 * heights   # B
        ])
        
        return self.drops.copy(), colors
    
    def get_rain_lines(self) -> pv.PolyData:
        """Get rain as short lines for better visualization"""
        line_length = 2.0
        
        # Create line endpoints
        starts = self.drops.copy()
        ends = self.drops + np.array([0, 0, line_length])
        
        # Build lines
        points = []
        lines = []
        
        for i in range(len(self.drops)):
            idx = len(points)
            points.append(starts[i])
            points.append(ends[i])
            lines.append([2, idx, idx + 1])
        
        if not points:
            return pv.PolyData()
        
        mesh = pv.PolyData(np.array(points))
        mesh.lines = np.array(lines)
        
        return mesh
    
    def reset(self):
        """Reset simulation"""
        self._initialize_drops()


# ============================================================================
# WIND SIMULATION
# ============================================================================

class WindSimulation:
    """Wind particle flow simulation"""
    
    def __init__(self, bounds: Tuple[float, float, float, float],
                 max_height: float = 50.0,
                 num_particles: int = 1500,
                 wind_speed: float = 20.0,
                 wind_direction: float = 45.0):  # degrees
        self.bounds = bounds
        self.max_height = max_height
        self.num_particles = num_particles
        self.base_wind_speed = wind_speed
        self.wind_direction = np.radians(wind_direction)
        self.particles: np.ndarray = None
        self.velocities: np.ndarray = None
        self.lifetimes: np.ndarray = None
        self.max_lifetime = 5.0
        self.paused = False
        self.buildings: List[Dict] = []
        self._initialize_particles()
    
    def set_buildings(self, plots_data: List[Dict]):
        """Set building data for wind interaction"""
        self.buildings = plots_data
    
    def set_wind_direction(self, degrees: float):
        """Set wind direction in degrees"""
        self.wind_direction = np.radians(degrees)
        self._update_base_velocity()
    
    def _update_base_velocity(self):
        """Update base velocity based on wind direction"""
        vx = self.base_wind_speed * np.cos(self.wind_direction)
        vy = self.base_wind_speed * np.sin(self.wind_direction)
        
        # Add variation to existing particles
        if self.velocities is not None:
            noise_x = np.random.uniform(-3, 3, self.num_particles)
            noise_y = np.random.uniform(-3, 3, self.num_particles)
            noise_z = np.random.uniform(-1, 1, self.num_particles)
            
            self.velocities[:, 0] = vx + noise_x
            self.velocities[:, 1] = vy + noise_y
            self.velocities[:, 2] = noise_z
    
    def _initialize_particles(self):
        """Initialize wind particles"""
        minx, miny, maxx, maxy = self.bounds
        
        # Spawn particles from upwind edge
        wind_dx = np.cos(self.wind_direction)
        wind_dy = np.sin(self.wind_direction)
        
        # Determine spawn edge based on wind direction
        if abs(wind_dx) > abs(wind_dy):
            if wind_dx > 0:  # Wind from left
                x = np.full(self.num_particles, minx - 10)
            else:  # Wind from right
                x = np.full(self.num_particles, maxx + 10)
            y = np.random.uniform(miny, maxy, self.num_particles)
        else:
            x = np.random.uniform(minx, maxx, self.num_particles)
            if wind_dy > 0:  # Wind from bottom
                y = np.full(self.num_particles, miny - 10)
            else:  # Wind from top
                y = np.full(self.num_particles, maxy + 10)
        
        z = np.random.uniform(1, self.max_height, self.num_particles)
        
        self.particles = np.column_stack([x, y, z])
        
        # Initialize velocities
        vx = self.base_wind_speed * wind_dx + np.random.uniform(-3, 3, self.num_particles)
        vy = self.base_wind_speed * wind_dy + np.random.uniform(-3, 3, self.num_particles)
        vz = np.random.uniform(-1, 1, self.num_particles)
        
        self.velocities = np.column_stack([vx, vy, vz])
        
        # Lifetimes for particle recycling
        self.lifetimes = np.random.uniform(0, self.max_lifetime, self.num_particles)
    
    def update(self, dt: float = 0.016):
        """Update particle positions with building interaction"""
        if self.paused:
            return
        
        # Update lifetimes
        self.lifetimes -= dt
        
        # Move particles
        self.particles += self.velocities * dt
        
        # Add turbulence
        turbulence = np.random.uniform(-0.5, 0.5, self.particles.shape) * dt * 10
        self.velocities += turbulence
        
        # Clamp vertical velocity
        self.velocities[:, 2] = np.clip(self.velocities[:, 2], -5, 5)
        
        # Simple building interaction - deflect around buildings
        for building in self.buildings:
            bounds = building.get('bounds', [0, 0, 0, 0])
            height = building.get('max_building_height', 10)
            
            bminx, bminy, bmaxx, bmaxy = bounds
            
            # Check particles near building
            in_x = (self.particles[:, 0] > bminx) & (self.particles[:, 0] < bmaxx)
            in_y = (self.particles[:, 1] > bminy) & (self.particles[:, 1] < bmaxy)
            below_height = self.particles[:, 2] < height
            
            affected = in_x & in_y & below_height
            
            if np.any(affected):
                # Push particles up and around
                self.velocities[affected, 2] += 10 * dt
                
                # Deflect horizontally
                center_x = (bminx + bmaxx) / 2
                center_y = (bminy + bmaxy) / 2
                
                dx = self.particles[affected, 0] - center_x
                dy = self.particles[affected, 1] - center_y
                
                self.velocities[affected, 0] += np.sign(dx) * 5 * dt
                self.velocities[affected, 1] += np.sign(dy) * 5 * dt
        
        # Reset expired or out-of-bounds particles
        minx, miny, maxx, maxy = self.bounds
        margin = 50
        
        out_of_bounds = (
            (self.particles[:, 0] < minx - margin) |
            (self.particles[:, 0] > maxx + margin) |
            (self.particles[:, 1] < miny - margin) |
            (self.particles[:, 1] > maxy + margin) |
            (self.particles[:, 2] < 0) |
            (self.particles[:, 2] > self.max_height * 2) |
            (self.lifetimes < 0)
        )
        
        if np.any(out_of_bounds):
            self._respawn_particles(out_of_bounds)
    
    def _respawn_particles(self, mask: np.ndarray):
        """Respawn particles at upwind edge"""
        num_respawn = np.sum(mask)
        if num_respawn == 0:
            return
        
        minx, miny, maxx, maxy = self.bounds
        wind_dx = np.cos(self.wind_direction)
        wind_dy = np.sin(self.wind_direction)
        
        # Spawn at upwind edge
        if abs(wind_dx) > abs(wind_dy):
            if wind_dx > 0:
                self.particles[mask, 0] = minx - 10
            else:
                self.particles[mask, 0] = maxx + 10
            self.particles[mask, 1] = np.random.uniform(miny, maxy, num_respawn)
        else:
            self.particles[mask, 0] = np.random.uniform(minx, maxx, num_respawn)
            if wind_dy > 0:
                self.particles[mask, 1] = miny - 10
            else:
                self.particles[mask, 1] = maxy + 10
        
        self.particles[mask, 2] = np.random.uniform(1, self.max_height, num_respawn)
        
        # Reset velocities
        self.velocities[mask, 0] = self.base_wind_speed * wind_dx + np.random.uniform(-3, 3, num_respawn)
        self.velocities[mask, 1] = self.base_wind_speed * wind_dy + np.random.uniform(-3, 3, num_respawn)
        self.velocities[mask, 2] = np.random.uniform(-1, 1, num_respawn)
        
        # Reset lifetimes
        self.lifetimes[mask] = self.max_lifetime
    
    def get_points_and_colors(self) -> Tuple[np.ndarray, np.ndarray]:
        """Get particle positions and colors"""
        # Color based on velocity magnitude
        speeds = np.linalg.norm(self.velocities, axis=1)
        max_speed = self.base_wind_speed * 1.5
        normalized_speed = np.clip(speeds / max_speed, 0, 1)
        
        # Color gradient: blue (slow) -> cyan -> white (fast)
        colors = np.column_stack([
            0.5 + 0.5 * normalized_speed,  # R
            0.8 + 0.2 * normalized_speed,  # G
            1.0 * np.ones_like(normalized_speed)  # B
        ])
        
        return self.particles.copy(), colors
    
    def get_streamlines(self) -> pv.PolyData:
        """Get wind as short streamlines"""
        line_length = 3.0
        
        # Normalize velocities for direction
        speeds = np.linalg.norm(self.velocities, axis=1, keepdims=True)
        speeds = np.maximum(speeds, 0.001)
        directions = self.velocities / speeds
        
        # Create line endpoints
        starts = self.particles.copy()
        ends = self.particles + directions * line_length
        
        points = []
        lines = []
        
        for i in range(len(self.particles)):
            idx = len(points)
            points.append(starts[i])
            points.append(ends[i])
            lines.append([2, idx, idx + 1])
        
        if not points:
            return pv.PolyData()
        
        mesh = pv.PolyData(np.array(points))
        mesh.lines = np.array(lines)
        
        return mesh
    
    def reset(self):
        """Reset simulation"""
        self._initialize_particles()


# ============================================================================
# MAIN SIMULATION VIEWER
# ============================================================================

class UrbanSimulationViewer:
    """Main simulation viewer combining all simulations"""
    
    def __init__(self, obj_path: str, json_path: str = None):
        self.obj_path = Path(obj_path)
        self.json_path = Path(json_path) if json_path else None
        
        # Load data
        self.mesh = None
        self.layout_data = None
        self.bounds = None
        
        # Simulations
        self.traffic_sim = None
        self.rain_sim = None
        self.wind_sim = None
        
        # Visualization
        self.plotter = None
        self.current_mode = SimulationMode.NONE
        self.simulation_actors = {}
        self.is_paused = False
        self.last_update_time = time.time()
        
        # Load resources
        self._load_resources()
        self._initialize_simulations()
    
    def _load_resources(self):
        """Load OBJ mesh and JSON data"""
        print("=" * 60)
        print("LOADING RESOURCES")
        print("=" * 60)
        
        # Load OBJ
        if self.obj_path.exists():
            print(f"Loading OBJ: {self.obj_path}")
            self.mesh = pv.read(str(self.obj_path))
            self.bounds = self.mesh.bounds
            print(f"  Vertices: {self.mesh.n_points}")
            print(f"  Faces: {self.mesh.n_cells}")
            print(f"  Bounds: {self.bounds}")
        else:
            raise FileNotFoundError(f"OBJ file not found: {self.obj_path}")
        
        # Load JSON if provided
        if self.json_path and self.json_path.exists():
            print(f"Loading JSON: {self.json_path}")
            with open(self.json_path, 'r') as f:
                self.layout_data = json.load(f)
            print(f"  Roads: {len(self.layout_data.get('roads', []))}")
            print(f"  Plots: {len(self.layout_data.get('plots', []))}")
        else:
            print("No JSON data - using mesh bounds for simulation")
            # Create minimal data from mesh bounds
            minx, maxx, miny, maxy, minz, maxz = self.bounds
            self.layout_data = {
                'roads': [],
                'plots': [],
                'site_info': {
                    'bounds': [minx, miny, maxx, maxy]
                }
            }
    
    def _initialize_simulations(self):
        """Initialize all simulation systems"""
        print("\n" + "=" * 60)
        print("INITIALIZING SIMULATIONS")
        print("=" * 60)
        
        # Get bounds
        if self.layout_data and 'site_info' in self.layout_data:
            bounds = self.layout_data['site_info']['bounds']
            sim_bounds = (bounds[0], bounds[1], bounds[2], bounds[3])
        else:
            minx, maxx, miny, maxy, _, _ = self.bounds
            sim_bounds = (minx, miny, maxx, maxy)
        
        # Calculate max building height
        max_height = 50.0
        if self.layout_data and 'plots' in self.layout_data:
            heights = [p.get('max_building_height', 10) for p in self.layout_data['plots']]
            if heights:
                max_height = max(heights) + 20
        
        # Traffic simulation
        roads = self.layout_data.get('roads', []) if self.layout_data else []
        if roads:
            self.traffic_sim = TrafficSimulation(roads, num_agents=100)
            print(f"Traffic: {len(self.traffic_sim.agents)} agents on {len(roads)} roads")
        else:
            # Generate fake roads from bounds for demo
            fake_roads = self._generate_fake_roads(sim_bounds)
            self.traffic_sim = TrafficSimulation(fake_roads, num_agents=50)
            print(f"Traffic: {len(self.traffic_sim.agents)} agents (generated roads)")
        
        # Rain simulation
        self.rain_sim = RainSimulation(
            bounds=sim_bounds,
            max_height=max_height + 50,
            num_drops=3000,
            rain_speed=60.0
        )
        print(f"Rain: {self.rain_sim.num_drops} drops")
        
        # Wind simulation
        self.wind_sim = WindSimulation(
            bounds=sim_bounds,
            max_height=max_height,
            num_particles=2000,
            wind_speed=25.0,
            wind_direction=45.0
        )
        
        # Set buildings for wind interaction
        if self.layout_data and 'plots' in self.layout_data:
            self.wind_sim.set_buildings(self.layout_data['plots'])
        
        print(f"Wind: {self.wind_sim.num_particles} particles")
    
    def _generate_fake_roads(self, bounds: Tuple) -> List[Dict]:
        """Generate fake roads if JSON not provided"""
        minx, miny, maxx, maxy = bounds
        roads = []
        
        # Grid of roads
        num_h = 5
        num_v = 5
        
        for i in range(num_h):
            y = miny + (i + 1) * (maxy - miny) / (num_h + 1)
            roads.append({
                'start_point': [minx, y],
                'end_point': [maxx, y],
                'width': 12.0,
                'length': maxx - minx
            })
        
        for i in range(num_v):
            x = minx + (i + 1) * (maxx - minx) / (num_v + 1)
            roads.append({
                'start_point': [x, miny],
                'end_point': [x, maxy],
                'width': 12.0,
                'length': maxy - miny
            })
        
        return roads
    
    def _setup_plotter(self):
        """Setup PyVista plotter with controls"""
        self.plotter = pv.Plotter(title="GENESIS Urban Simulation Viewer")
        self.plotter.set_background('black')
        
        # Add mesh
        self.plotter.add_mesh(
            self.mesh, 
            color='lightgray',
            opacity=0.9,
            show_edges=False,
            name='layout_mesh'
        )
        
        # Set top-down camera
        self._set_top_down_view()
        
        # Add text overlay
        self._add_ui_text()
        
        # Setup key bindings
        self.plotter.add_key_event('space', self._toggle_pause)
        self.plotter.add_key_event('r', self._reset_simulations)
        self.plotter.add_key_event('0', lambda: self._set_mode(SimulationMode.NONE))
        self.plotter.add_key_event('1', lambda: self._set_mode(SimulationMode.TRAFFIC))
        self.plotter.add_key_event('2', lambda: self._set_mode(SimulationMode.RAIN))
        self.plotter.add_key_event('3', lambda: self._set_mode(SimulationMode.WIND))
        self.plotter.add_key_event('Up', lambda: self._adjust_wind_direction(15))
        self.plotter.add_key_event('Down', lambda: self._adjust_wind_direction(-15))
        self.plotter.add_key_event('t', self._toggle_top_view)
        self.plotter.add_key_event('p', self._toggle_perspective)
    
    def _set_top_down_view(self):
        """Set camera to top-down view"""
        if self.bounds:
            minx, maxx, miny, maxy, minz, maxz = self.bounds
            cx = (minx + maxx) / 2
            cy = (miny + maxy) / 2
            
            # Calculate camera height based on scene size
            width = maxx - minx
            height = maxy - miny
            cam_height = max(width, height) * 1.5
            
            self.plotter.camera_position = [
                (cx, cy, cam_height),  # Camera position
                (cx, cy, 0),           # Focal point
                (0, 1, 0)              # Up vector
            ]
    
    def _toggle_top_view(self):
        """Toggle between top-down and isometric view"""
        if self.bounds:
            minx, maxx, miny, maxy, minz, maxz = self.bounds
            cx = (minx + maxx) / 2
            cy = (miny + maxy) / 2
            
            width = maxx - minx
            height = maxy - miny
            cam_dist = max(width, height) * 1.2
            
            self.plotter.camera_position = [
                (cx, cy, cam_dist),
                (cx, cy, 0),
                (0, 1, 0)
            ]
    
    def _toggle_perspective(self):
        """Toggle to perspective/isometric view"""
        if self.bounds:
            minx, maxx, miny, maxy, minz, maxz = self.bounds
            cx = (minx + maxx) / 2
            cy = (miny + maxy) / 2
            
            width = maxx - minx
            height = maxy - miny
            cam_dist = max(width, height) * 0.8
            
            self.plotter.camera_position = [
                (cx - cam_dist, cy - cam_dist, cam_dist * 0.7),
                (cx, cy, 0),
                (0, 0, 1)
            ]
    
    def _add_ui_text(self):
        """Add UI text overlay"""
        instructions = (
            "Controls:\n"
            "SPACE - Pause/Resume\n"
            "R - Reset\n"
            "0 - No overlay\n"
            "1 - Traffic\n"
            "2 - Rain\n"
            "3 - Wind\n"
            "T - Top view\n"
            "P - Perspective\n"
            "↑/↓ - Wind direction"
        )
        
        self.plotter.add_text(
            instructions,
            position='upper_left',
            font_size=10,
            color='white',
            name='instructions'
        )
        
        self.status_text = self.plotter.add_text(
            "Mode: None | Status: Running",
            position='upper_right',
            font_size=12,
            color='cyan',
            name='status'
        )
    
    def _update_status_text(self):
        """Update status text"""
        mode_names = {
            SimulationMode.NONE: "None",
            SimulationMode.TRAFFIC: "Traffic",
            SimulationMode.RAIN: "Rain",
            SimulationMode.WIND: "Wind"
        }
        
        status = "Paused" if self.is_paused else "Running"
        mode = mode_names.get(self.current_mode, "Unknown")
        
        extra = ""
        if self.current_mode == SimulationMode.WIND:
            direction = np.degrees(self.wind_sim.wind_direction)
            extra = f" | Direction: {direction:.0f}°"
        
        text = f"Mode: {mode} | Status: {status}{extra}"
        
        self.plotter.add_text(
            text,
            position='upper_right',
            font_size=12,
            color='cyan',
            name='status'
        )
    
    def _toggle_pause(self):
        """Toggle pause state"""
        self.is_paused = not self.is_paused
        
        if self.traffic_sim:
            self.traffic_sim.paused = self.is_paused
        if self.rain_sim:
            self.rain_sim.paused = self.is_paused
        if self.wind_sim:
            self.wind_sim.paused = self.is_paused
        
        self._update_status_text()
        print(f"{'Paused' if self.is_paused else 'Resumed'}")
    
    def _reset_simulations(self):
        """Reset all simulations"""
        if self.traffic_sim:
            self.traffic_sim.reset()
        if self.rain_sim:
            self.rain_sim.reset()
        if self.wind_sim:
            self.wind_sim.reset()
        
        print("Simulations reset")
    
    def _set_mode(self, mode: SimulationMode):
        """Set simulation display mode"""
        self.current_mode = mode
        self._clear_simulation_actors()
        self._update_status_text()
        
        mode_names = {
            SimulationMode.NONE: "None",
            SimulationMode.TRAFFIC: "Traffic",
            SimulationMode.RAIN: "Rain",
            SimulationMode.WIND: "Wind"
        }
        print(f"Mode: {mode_names.get(mode, 'Unknown')}")
    
    def _adjust_wind_direction(self, delta: float):
        """Adjust wind direction"""
        if self.wind_sim:
            new_dir = np.degrees(self.wind_sim.wind_direction) + delta
            self.wind_sim.set_wind_direction(new_dir)
            self._update_status_text()
            print(f"Wind direction: {new_dir:.0f}°")
    
    def _clear_simulation_actors(self):
        """Clear simulation visualization actors"""
        for name in list(self.simulation_actors.keys()):
            try:
                self.plotter.remove_actor(name)
            except:
                pass
        self.simulation_actors = {}
    
    def _update_callback(self):
        """Main update callback for animation"""
        current_time = time.time()
        dt = min(current_time - self.last_update_time, 0.1)  # Cap dt
        self.last_update_time = current_time
        
        # Update simulations
        if self.current_mode == SimulationMode.TRAFFIC and self.traffic_sim:
            self.traffic_sim.update(dt)
            self._render_traffic()
        
        elif self.current_mode == SimulationMode.RAIN and self.rain_sim:
            self.rain_sim.update(dt)
            self._render_rain()
        
        elif self.current_mode == SimulationMode.WIND and self.wind_sim:
            self.wind_sim.update(dt)
            self._render_wind()
    
    def _render_traffic(self):
        """Render traffic agents"""
        if not self.traffic_sim:
            return
        
        points, colors = self.traffic_sim.get_points_and_colors()
        
        if len(points) > 0:
            point_cloud = pv.PolyData(points)
            point_cloud['colors'] = (colors * 255).astype(np.uint8)
            
            # Remove old actor
            if 'traffic_points' in self.simulation_actors:
                try:
                    self.plotter.remove_actor('traffic_points')
                except:
                    pass
            
            actor = self.plotter.add_mesh(
                point_cloud,
                scalars='colors',
                rgb=True,
                point_size=15,
                render_points_as_spheres=True,
                name='traffic_points'
            )
            self.simulation_actors['traffic_points'] = actor
    
    def _render_rain(self):
        """Render rain particles"""
        if not self.rain_sim:
            return
        
        # Use lines for rain
        rain_mesh = self.rain_sim.get_rain_lines()
        
        if rain_mesh.n_points > 0:
            if 'rain_lines' in self.simulation_actors:
                try:
                    self.plotter.remove_actor('rain_lines')
                except:
                    pass
            
            actor = self.plotter.add_mesh(
                rain_mesh,
                color='lightblue',
                line_width=1,
                opacity=0.6,
                name='rain_lines'
            )
            self.simulation_actors['rain_lines'] = actor
    
    def _render_wind(self):
        """Render wind particles"""
        if not self.wind_sim:
            return
        
        # Use streamlines for wind
        wind_mesh = self.wind_sim.get_streamlines()
        
        if wind_mesh.n_points > 0:
            if 'wind_lines' in self.simulation_actors:
                try:
                    self.plotter.remove_actor('wind_lines')
                except:
                    pass
            
            actor = self.plotter.add_mesh(
                wind_mesh,
                color='white',
                line_width=2,
                opacity=0.7,
                name='wind_lines'
            )
            self.simulation_actors['wind_lines'] = actor
    
    def run(self):
        """Run the simulation viewer"""
        print("\n" + "=" * 60)
        print("STARTING SIMULATION VIEWER")
        print("=" * 60)
        print("\nControls:")
        print("  SPACE  - Pause/Resume")
        print("  R      - Reset simulations")
        print("  0      - No overlay")
        print("  1      - Traffic simulation")
        print("  2      - Rain simulation")
        print("  3      - Wind simulation")
        print("  T      - Top-down view")
        print("  P      - Perspective view")
        print("  ↑/↓    - Change wind direction")
        print("\nClose window to exit.")
        
        self._setup_plotter()
        
        # Start with traffic mode
        self._set_mode(SimulationMode.TRAFFIC)
        
        # Timer-driven animation (works on older PyVista)
        self.plotter.add_timer_event(max_steps=10**9, duration=33, callback=self._update_callback)

        # Show plotter
        self.plotter.show()


# ============================================================================
# CLI INTERFACE
# ============================================================================

def main():
    parser = argparse.ArgumentParser(
        description='GENESIS Urban Layout Simulation Viewer',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python simulation_viewer.py layout.obj
  python simulation_viewer.py layout.obj --json layout.json
  python simulation_viewer.py layout.obj -j layout.json --wind-dir 90
        """
    )
    
    parser.add_argument('obj_file', type=str, help='Path to OBJ file')
    parser.add_argument('--json', '-j', type=str, help='Path to JSON layout file')
    parser.add_argument('--wind-dir', type=float, default=45.0,
                        help='Initial wind direction in degrees (default: 45)')
    parser.add_argument('--num-traffic', type=int, default=100,
                        help='Number of traffic agents (default: 100)')
    parser.add_argument('--num-rain', type=int, default=3000,
                        help='Number of rain drops (default: 3000)')
    parser.add_argument('--num-wind', type=int, default=2000,
                        help='Number of wind particles (default: 2000)')
    
    args = parser.parse_args()
    
    # Validate files
    obj_path = Path(args.obj_file)
    if not obj_path.exists():
        print(f"Error: OBJ file not found: {obj_path}")
        return 1
    
    json_path = args.json
    if json_path and not Path(json_path).exists():
        print(f"Warning: JSON file not found: {json_path}")
        print("Proceeding without layout data...")
        json_path = None
    
    # Create and run viewer
    try:
        viewer = UrbanSimulationViewer(str(obj_path), json_path)
        
        # Apply custom settings
        if viewer.wind_sim:
            viewer.wind_sim.set_wind_direction(args.wind_dir)
        
        viewer.run()
        
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        return 1
    
    return 0


if __name__ == "__main__":
    exit(main())