from typing import List, Tuple

class VFH:
    def __init__(self, n_sectors: int = 72, threshold: float = 0.5,
                 safe_distance: float = 2.0):
        self.n_sectors = n_sectors
        self.threshold = threshold
        self.safe_distance = safe_distance
        self.sector_angle = 2 * np.pi / n_sectors
    
    def build_polar_histogram(self, obstacles: List[dict], 
                               robot_position: np.ndarray) -> np.ndarray:
        histogram = np.zeros(self.n_sectors)
        
        for obs in obstacles:
            diff = obs['position'] - robot_position
            distance = np.linalg.norm(diff)
            
            if distance < self.safe_distance * 3:
                angle = np.arctan2(diff[1], diff[0])
                sector = int((angle + np.pi) / self.sector_angle) % self.n_sectors
                
                # Certainty decreases with distance
                certainty = max(0, 1 - distance / (self.safe_distance * 3))
                histogram[sector] += certainty * obs.get('size', 1.0)
        
        return histogram
    
    def smooth_histogram(self, histogram: np.ndarray, window: int = 5) -> np.ndarray:
        smoothed = np.zeros_like(histogram)
        for i in range(len(histogram)):
            weights = []
            for j in range(-window, window + 1):
                idx = (i + j) % len(histogram)
                weight = max(0, 1 - abs(j) / (window + 1))
                weights.append(histogram[idx] * weight)
            smoothed[i] = sum(weights)
        return smoothed
    
    def find_best_direction(self, histogram: np.ndarray, 
                            goal_direction: float) -> float:
        goal_sector = int((goal_direction + np.pi) / self.sector_angle) % self.n_sectors
        
        # Find valleys (sectors below threshold)
        valleys = []
        in_valley = False
        valley_start = 0
        
        for i in range(self.n_sectors * 2):
            idx = i % self.n_sectors
            if histogram[idx] < self.threshold:
                if not in_valley:
                    valley_start = idx
                    in_valley = True
            else:
                if in_valley:
                    valleys.append((valley_start, idx - 1))
                    in_valley = False
        
        if in_valley:
            valleys.append((valley_start, self.n_sectors - 1))
        
        # Select valley closest to goal
        best_direction = goal_direction
        best_cost = float('inf')
        
        for start, end in valleys:
            mid = (start + end) / 2 % self.n_sectors
            angle = mid * self.sector_angle - np.pi
            
            # Cost: angular difference from goal
            cost = abs(np.arctan2(np.sin(angle - goal_direction), 
                                  np.cos(angle - goal_direction)))
            
            if cost < best_cost:
                best_cost = cost
                best_direction = angle
        
        return best_direction
    
    def compute_velocity(self, obstacles: List[dict], robot_pos: np.ndarray,
                         robot_yaw: float, goal_direction: float) -> Tuple[float, float]:
        histogram = self.build_polar_histogram(obstacles, robot_pos)
        smoothed = self.smooth_histogram(histogram)
        best_dir = self.find_best_direction(smoothed, goal_direction)
        
        # Convert direction to velocity commands
        angle_error = np.arctan2(np.sin(best_dir - robot_yaw), 
                                 np.cos(best_dir - robot_yaw))
        
        linear_vel = 2.0 * max(0.2, 1 - abs(angle_error))
        angular_vel = 1.5 * angle_error
        
        return linear_vel, angular_vel

from typing import List, Tuple

class DWAPlanner:
    def __init__(self, config):
        self.max_speed = config.get('max_speed', 2.0)
        self.min_speed = config.get('min_speed', 0.0)
        self.max_yaw_rate = config.get('max_yaw_rate', 40.0 * np.pi / 180)
        self.max_accel = config.get('max_accel', 0.2)
        self.max_delta_yaw_rate = config.get('max_delta_yaw_rate', 40.0 * np.pi / 180)
        self.v_resolution = config.get('v_resolution', 0.02)
        self.yaw_rate_resolution = config.get('yaw_rate_resolution', 1.0 * np.pi / 180)
        self.dt = config.get('dt', 0.1)
        self.predict_time = config.get('predict_time', 2.0)
        self.heading_cost = config.get('heading_cost', 0.15)
        self.clearance_cost = config.get('clearance_cost', 1.0)
        self.velocity_cost = config.get('velocity_cost', 0.2)
    
    def compute_dynamic_window(self, state) -> Tuple[float, float, float, float]:
        """Compute dynamic window based on current state and limits."""
        vs = (self.min_speed, self.max_speed,
              -self.max_yaw_rate, self.max_yaw_rate)
        
        vd = (state[3] - self.max_accel * self.dt,
              state[3] + self.max_accel * self.dt,
              state[4] - self.max_delta_yaw_rate * self.dt,
              state[4] + self.max_delta_yaw_rate * self.dt)
        
        return (max(vs[0], vd[0]), min(vs[1], vd[1]),
                max(vs[2], vd[2]), min(vs[3], vd[3]))
    
    def predict_trajectory(self, state, v, yaw_rate) -> np.ndarray:
        """Simulate trajectory for given velocity command."""
        trajectory = [state.copy()]
        s = state.copy()
        
        for _ in range(int(self.predict_time / self.dt)):
            s[0] += v * np.cos(s[2]) * self.dt
            s[1] += v * np.sin(s[2]) * self.dt
            s[2] += yaw_rate * self.dt
            s[3] = v
            s[4] = yaw_rate
            trajectory.append(s.copy())
        
        return np.array(trajectory)
    
    def heading_score(self, trajectory, goal) -> float:
        """Score based on alignment with goal direction."""
        final_pos = trajectory[-1, :2]
        goal_direction = np.arctan2(goal[1] - final_pos[1], goal[0] - final_pos[0])
        heading = trajectory[-1, 2]
        return 1.0 - abs(np.arctan2(np.sin(goal_direction - heading),
                                      np.cos(goal_direction - heading))) / np.pi
    
    def clearance_score(self, trajectory, obstacles) -> float:
        """Score based on minimum distance to obstacles."""
        min_dist = float('inf')
        
        for point in trajectory:
            for obs in obstacles:
                dist = np.linalg.norm(point[:2] - obs['position'][:2]) - obs['radius']
                min_dist = min(min_dist, dist)
        
        if min_dist < 0.1:
            return 0.0
        return min(1.0, min_dist / 5.0)
    
    def velocity_score(self, trajectory) -> float:
        """Score based on maintaining desired speed."""
        return trajectory[-1, 3] / self.max_speed
    
    def select_best_velocity(self, state, goal, obstacles) -> Tuple[float, float]:
        dw = self.compute_dynamic_window(state)
        best_score = -float('inf')
        best_v, best_yaw = state[3], state[4]
        
        for v in np.arange(dw[0], dw[1], self.v_resolution):
            for yaw in np.arange(dw[2], dw[3], self.yaw_rate_resolution):
                traj = self.predict_trajectory(state, v, yaw)
                
                h_score = self.heading_score(traj, goal)
                c_score = self.clearance_score(traj, obstacles)
                v_score = self.velocity_score(traj)
                
                if c_score < 0.01:
                    continue
                
                total = (self.heading_cost * h_score +
                        self.clearance_cost * c_score +
                        self.velocity_cost * v_score)
                
                if total > best_score:
                    best_score = total
                    best_v, best_yaw = v, yaw
        
        return best_v, best_yaw

class ObstacleAvoidanceNet(nn.Module):
    def __init__(self, lidar_size=360, state_size=4, hidden_size=256):
        super().__init__()
        
        # LiDAR encoder
        self.lidar_conv = nn.Sequential(
            nn.Conv1d(1, 32, kernel_size=5, stride=2, padding=2),
            nn.ReLU(),
            nn.Conv1d(32, 64, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool1d(16)
        )
        
        # State encoder
        self.state_fc = nn.Sequential(
            nn.Linear(state_size, 64),
            nn.ReLU(),
            nn.Linear(64, 64)
        )
        
        # Fusion and decision
        self.fusion = nn.Sequential(
            nn.Linear(64 * 16 + 64, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU()
        )
        
        # Policy head
        self.policy_mean = nn.Linear(hidden_size, 2)  # v, omega
        self.policy_log_std = nn.Linear(hidden_size, 2)
        
        # Value head
        self.value = nn.Linear(hidden_size, 1)
    
    def forward(self, lidar, state):
        # Encode lidar
        lidar_features = self.lidar_conv(lidar.unsqueeze(1))
        lidar_features = lidar_features.view(lidar_features.size(0), -1)
        
        # Encode state
        state_features = self.state_fc(state)
        
        # Fuse features
        combined = torch.cat([lidar_features, state_features], dim=1)
        features = self.fusion(combined)
        
        # Policy
        mean = self.policy_mean(features)
        log_std = self.policy_log_std(features).clamp(-20, 2)
        std = log_std.exp()
        
        # Value
        value = self.value(features)
        
        return mean, std, value
    
    def select_action(self, lidar, state):
        mean, std, value = self.forward(lidar, state)
        normal = torch.distributions.Normal(mean, std)
        action = normal.sample()
        log_prob = normal.log_prob(action).sum(dim=-1)
        return action, log_prob, value


class PPOTrainer:
    def __init__(self, model, lr=3e-4, gamma=0.99, clip_eps=0.2):
        self.model = model
        self.optimizer = torch.optim.Adam(model.parameters(), lr=lr)
        self.gamma = gamma
        self.clip_eps = clip_eps
    
    def compute_gae(self, rewards, values, dones):
        advantages = []
        gae = 0
        for t in reversed(range(len(rewards))):
            if t == len(rewards) - 1:
                next_value = 0
            else:
                next_value = values[t + 1]
            delta = rewards[t] + self.gamma * next_value * (1 - dones[t]) - values[t]
            gae = delta + self.gamma * 0.95 * (1 - dones[t]) * gae
            advantages.insert(0, gae)
        return torch.tensor(advantages, dtype=torch.float32)
    
    def update(self, batch):
        old_actions = torch.tensor(batch['actions'], dtype=torch.float32)
        old_log_probs = torch.tensor(batch['log_probs'], dtype=torch.float32)
        returns = self.compute_gae(batch['rewards'], batch['values'], batch['dones'])
        
        for _ in range(10):
            mean, std, values = self.model(
                torch.tensor(batch['lidar']),
                torch.tensor(batch['states'])
            )
            
            normal = torch.distributions.Normal(mean, std)
            new_log_probs = normal.log_prob(old_actions).sum(dim=-1)
            
            ratio = (new_log_probs - old_log_probs).exp()
            clipped = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps)
            
            policy_loss = -torch.min(ratio * returns, clipped * returns).mean()
            value_loss = F.mse_loss(values.squeeze(), returns)
            
            loss = policy_loss + 0.5 * value_loss
            
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

from collections import deque

class MultiSensorAvoidance:
    def __init__(self):
        self.lidar_buffer = deque(maxlen=10)
        self.camera_buffer = deque(maxlen=5)
        self ultrasonic_buffer = deque(maxlen=5)
        
        # Thresholds
        self.critical_distance = 1.0
        self.warning_distance = 3.0
        self.safe_distance = 5.0
        
        # Weights for sensor fusion
        self.weights = {'lidar': 0.5, 'camera': 0.3, 'ultrasonic': 0.2}
    
    def process_lidar(self, scan, angles):
        """Process LiDAR scan data."""
        obstacles = []
        for r, theta in zip(scan, angles):
            if r < self.safe_distance:
                x = r * np.cos(theta)
                y = r * np.sin(theta)
                obstacles.append({
                    'position': np.array([x, y]),
                    'distance': r,
                    'confidence': 1.0,
                    'sensor': 'lidar'
                })
        return obstacles
    
    def process_depth_camera(self, depth_image, fov_h, fov_v):
        """Process depth camera data."""
        h, w = depth_image.shape
        obstacles = []
        
        for i in range(0, h, 10):
            for j in range(0, w, 10):
                depth = depth_image[i, j]
                if 0.5 < depth < self.safe_distance:
                    theta_h = (j / w - 0.5) * fov_h
                    theta_v = (i / h - 0.5) * fov_v
                    x = depth * np.sin(theta_h)
                    y = depth * np.cos(theta_h)
                    z = depth * np.sin(theta_v)
                    
                    obstacles.append({
                        'position': np.array([x, y, z]),
                        'distance': depth,
                        'confidence': 0.8,
                        'sensor': 'camera'
                    })
        return obstacles
    
    def fuse_obstacles(self, all_obstacles):
        """Fuse obstacles from multiple sensors."""
        fused = []
        
        for obs in all_obstacles:
            # Check if similar obstacle exists in fused list
            merged = False
            for f in fused:
                if np.linalg.norm(obs['position'] - f['position']) < 1.0:
                    # Update with higher confidence
                    if obs['confidence'] > f['confidence']:
                        f['position'] = obs['position']
                        f['confidence'] = obs['confidence']
                    merged = True
                    break
            
            if not merged:
                fused.append(obs.copy())
        
        return fused
    
    def compute_avoidance_vector(self, obstacles, goal_direction):
        """Compute avoidance velocity vector."""
        repulsive = np.zeros(2)
        
        for obs in obstacles:
            dist = obs['distance']
            if dist < self.warning_distance:
                direction = -obs['position'][:2] / (dist + 0.01)
                strength = (self.warning_distance - dist) / self.warning_distance
                repulsive += direction * strength * obs['confidence']
        
        # Attractive toward goal
        attractive = np.array([np.cos(goal_direction), np.sin(goal_direction)])
        
        # Combine
        alpha = 1.0
        beta = 2.0
        avoidance = alpha * attractive + beta * repulsive
        
        # Normalize and scale
        norm = np.linalg.norm(avoidance)
        if norm > 0:
            avoidance = avoidance / norm
        
        return avoidance
    
    def safety_check(self, obstacles):
        """Check for critical safety situations."""
        critical = False
        for obs in obstacles:
            if obs['distance'] < self.critical_distance:
                critical = True
                break
        
        if critical:
            # Emergency stop and hover
            return {'action': 'hover', 'velocity': np.zeros(3)}
        
        return None
    
    def run_avoidance_loop(self, lidar_scan, depth_image, goal_direction):
        """Main avoidance loop."""
        # Process sensors
        lidar_obs = self.process_lidar(lidar_scan, np.linspace(-np.pi, np.pi, len(lidar_scan)))
        camera_obs = self.process_depth_camera(depth_image, 1.2, 0.9)
        
        # Safety check
        all_obs = lidar_obs + camera_obs
        safety = self.safety_check(all_obs)
        if safety:
            return safety
        
        # Fuse obstacles
        fused_obs = self.fuse_obstacles(all_obs)
        
        # Compute avoidance
        avoidance_vec = self.compute_avoidance_vector(fused_obs, goal_direction)
        
        # Convert to velocity command
        linear_vel = 1.5
        angular_vel = np.arctan2(avoidance_vec[1], avoidance_vec[0]) - goal_direction
        angular_vel = np.arctan2(np.sin(angular_vel), np.cos(angular_vel))
        
        return {
            'action': 'move',
            'linear_velocity': linear_vel,
            'angular_velocity': angular_vel * 2.0,
            'obstacle_count': len(fused_obs),
            'min_distance': min(o['distance'] for o in fused_obs) if fused_obs else float('inf')
        }