from dataclasses import dataclass
from typing import List, Tuple

@dataclass
class Detection:
    """Represents a single object detection."""
    bbox: Tuple[int, int, int, int]  # x1, y1, x2, y2
    confidence: float
    class_id: int
    class_name: str

class DroneObjectDetector:
    """YOLO-style object detector for drone imagery."""

    def __init__(self, num_classes=80, input_size=640, num_anchors=3):
        self.num_classes = num_classes
        self.input_size = input_size
        self.num_anchors = num_anchors
        self.class_names = self._load_class_names()

    def _load_class_names(self):
        """Load COCO class names."""
        return [
            'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
            'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
            'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
            'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
            'backpack', 'umbrella', 'handbag', 'tie', 'suitcase',
            'frisbee', 'skis', 'snowboard', 'sports ball', 'kite',
            'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
            'tennis racket', 'bottle', 'wine glass', 'cup', 'fork',
            'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
            'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut',
            'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table',
            'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
            'cell phone', 'microwave', 'oven', 'toaster', 'sink',
            'refrigerator', 'book', 'clock', 'vase', 'scissors',
            'teddy bear', 'hair drier', 'toothbrush'
        ]

    def simulate_yolo_output(self, image_shape):
        """Simulate YOLO network output."""
        h, w = image_shape[:2]
        grid_size = self.input_size // 32  # 20x20 grid

        # Generate mock detections
        detections = []
        num_detections = np.random.randint(3, 10)

        for _ in range(num_detections):
            # Random bounding box
            x1 = np.random.randint(0, w - 100)
            y1 = np.random.randint(0, h - 100)
            bw = np.random.randint(50, min(200, w - x1))
            bh = np.random.randint(50, min(200, h - y1))

            confidence = np.random.uniform(0.3, 0.95)
            class_id = np.random.randint(0, self.num_classes)

            detections.append(Detection(
                bbox=(x1, y1, x1 + bw, y1 + bh),
                confidence=confidence,
                class_id=class_id,
                class_name=self.class_names[class_id]
            ))

        return detections

    def non_max_suppression(self, detections: List[Detection],
                           iou_threshold=0.5) -> List[Detection]:
        """Apply Non-Maximum Suppression."""
        if not detections:
            return []

        # Sort by confidence
        sorted_dets = sorted(detections, key=lambda d: d.confidence, reverse=True)

        keep = []
        while sorted_dets:
            current = sorted_dets.pop(0)
            keep.append(current)

            remaining = []
            for det in sorted_dets:
                if det.class_id == current.class_id:
                    iou = self._compute_iou(current.bbox, det.bbox)
                    if iou < iou_threshold:
                        remaining.append(det)
                else:
                    remaining.append(det)

            sorted_dets = remaining

        return keep

    def _compute_iou(self, box1, box2):
        """Compute Intersection over Union."""
        x1 = max(box1[0], box2[0])
        y1 = max(box1[1], box2[1])
        x2 = min(box1[2], box2[2])
        y2 = min(box1[3], box2[3])

        intersection = max(0, x2 - x1) * max(0, y2 - y1)

        area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
        area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
        union = area1 + area2 - intersection

        return intersection / union if union > 0 else 0

    def scale_detections(self, detections, original_size, model_size):
        """Scale detections from model space to original image space."""
        scale_x = original_size[1] / model_size[1]
        scale_y = original_size[0] / model_size[0]

        scaled = []
        for det in detections:
            x1 = int(det.bbox[0] * scale_x)
            y1 = int(det.bbox[1] * scale_y)
            x2 = int(det.bbox[2] * scale_x)
            y2 = int(det.bbox[3] * scale_y)

            scaled.append(Detection(
                bbox=(x1, y1, x2, y2),
                confidence=det.confidence,
                class_id=det.class_id,
                class_name=det.class_name
            ))

        return scaled

# Example: Run detection
np.random.seed(42)
detector = DroneObjectDetector(num_classes=80)

# Simulate drone frame
frame_shape = (720, 1280, 3)

# Get detections
detections = detector.simulate_yolo_output(frame_shape)
print(f"Raw detections: {len(detections)}")

# Apply NMS
nms_detections = detector.non_max_suppression(detections, iou_threshold=0.5)
print(f"After NMS: {len(nms_detections)}")

# Display top detections
for det in sorted(nms_detections, key=lambda d: d.confidence, reverse=True)[:5]:
    print(f"  {det.class_name}: {det.confidence:.2%} at {det.bbox}")

from collections import deque

class KalmanFilter:
    """Kalman filter for track prediction."""

    def __init__(self):
        # State: [x, y, w, h, dx, dy, dw, dh]
        self.state = np.zeros(8)
        self.P = np.eye(8) * 10  # Covariance matrix
        self.Q = np.eye(8) * 0.1  # Process noise
        self.R = np.eye(4) * 1.0  # Measurement noise

        # State transition matrix
        self.F = np.eye(8)
        self.F[0, 4] = 1  # x += dx
        self.F[1, 5] = 1  # y += dy
        self.F[2, 6] = 1  # w += dw
        self.F[3, 7] = 1  # h += dh

        # Measurement matrix
        self.H = np.zeros((4, 8))
        self.H[:4, :4] = np.eye(4)

    def predict(self):
        self.state = self.F @ self.state
        self.P = self.F @ self.P @ self.F.T + self.Q
        return self.state[:4]

    def update(self, measurement):
        y = measurement - self.H @ self.state
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)

        self.state = self.state + K @ y
        self.P = (np.eye(8) - K @ self.H) @ self.P

class Track:
    """Single object track."""

    def __init__(self, track_id, initial_bbox):
        self.track_id = track_id
        self.bbox = initial_bbox
        self.kalman = KalmanFilter()
        self.kalman.state[:4] = np.array(initial_bbox, dtype=float)
        self.hits = 1
        self.time_since_update = 0
        self.history = deque(maxlen=30)
        self.appearance_descriptor = np.random.randn(128)

    def predict(self):
        predicted = self.kalman.predict()
        self.time_since_update += 1
        return predicted

    def update(self, bbox, descriptor=None):
        self.bbox = bbox
        self.kalman.update(np.array(bbox, dtype=float))
        self.hits += 1
        self.time_since_update = 0
        self.history.append(bbox)

        if descriptor is not None:
            # Update appearance with exponential moving average
            alpha = 0.7
            self.appearance_descriptor = (alpha * self.appearance_descriptor +
                                         (1 - alpha) * descriptor)

class DeepSORTTracker:
    """Deep SORT multi-object tracker."""

    def __init__(self, max_age=30, min_hits=3, iou_threshold=0.3):
        self.tracks = []
        self.next_id = 0
        self.max_age = max_age
        self.min_hits = min_hits
        self.iou_threshold = iou_threshold

    def update(self, detections, descriptors=None):
        """Update tracker with new detections."""
        # Predict existing tracks
        for track in self.tracks:
            track.predict()

        # Match detections to tracks
        if self.tracks and detections:
            matched, unmatched_tracks, unmatched_dets = self._match(
                detections, descriptors
            )

            # Update matched tracks
            for track_idx, det_idx in matched:
                desc = descriptors[det_idx] if descriptors is not None else None
                self.tracks[track_idx].update(detections[det_idx], desc)

            # Create new tracks for unmatched detections
            for det_idx in unmatched_dets:
                desc = descriptors[det_idx] if descriptors is not None else None
                self._create_track(detections[det_idx], desc)
        else:
            for det in detections:
                self._create_track(det)

        # Remove old tracks
        self.tracks = [t for t in self.tracks
                      if t.time_since_update <= self.max_age]

        # Return confirmed tracks
        confirmed = [t for t in self.tracks if t.hits >= self.min_hits]
        return confirmed

    def _match(self, detections, descriptors=None):
        """Match detections to tracks using IoU and appearance."""
        n_tracks = len(self.tracks)
        n_dets = len(detections)

        # Compute cost matrix
        cost_matrix = np.zeros((n_tracks, n_dets))

        for i, track in enumerate(self.tracks):
            for j, det in enumerate(detections):
                # IoU cost
                iou = self._compute_iou(track.bbox, det)
                cost_matrix[i, j] = 1 - iou  # Convert to cost

                # Appearance cost
                if descriptors is not None:
                    desc_dist = np.linalg.norm(
                        track.appearance_descriptor - descriptors[j]
                    )
                    cost_matrix[i, j] += 0.5 * min(desc_dist, 1.0)

        # Hungarian algorithm (simplified greedy matching)
        matched = []
        unmatched_tracks = list(range(n_tracks))
        unmatched_dets = list(range(n_dets))

        for _ in range(min(n_tracks, n_dets)):
            if not unmatched_tracks or not unmatched_dets:
                break

            # Find minimum cost pair
            min_cost = float('inf')
            best_match = None

            for i in unmatched_tracks:
                for j in unmatched_dets:
                    if cost_matrix[i, j] < min_cost:
                        min_cost = cost_matrix[i, j]
                        best_match = (i, j)

            if best_match and min_cost < (1 - self.iou_threshold):
                matched.append(best_match)
                unmatched_tracks.remove(best_match[0])
                unmatched_dets.remove(best_match[1])
            else:
                break

        return matched, unmatched_tracks, unmatched_dets

    def _compute_iou(self, box1, box2):
        """Compute IoU between two boxes."""
        x1 = max(box1[0], box2[0])
        y1 = max(box1[1], box2[1])
        x2 = min(box1[2], box2[2])
        y2 = min(box1[3], box2[3])

        intersection = max(0, x2 - x1) * max(0, y2 - y1)
        area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
        area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
        union = area1 + area2 - intersection

        return intersection / union if union > 0 else 0

    def _create_track(self, bbox, descriptor=None):
        """Create a new track."""
        track = Track(self.next_id, bbox)
        if descriptor is not None:
            track.appearance_descriptor = descriptor
        self.tracks.append(track)
        self.next_id += 1

# Example: Track objects across frames
np.random.seed(42)
tracker = DeepSORTTracker(max_age=15, min_hits=3)

print("Simulating drone tracking sequence...")
for frame_id in range(20):
    # Simulate detections (objects moving)
    detections = []
    for obj_id in range(3):
        x = 100 + frame_id * 10 + obj_id * 200 + np.random.randint(-5, 5)
        y = 150 + np.random.randint(-10, 10)
        detections.append((x, y, x + 60, y + 80))

    # Update tracker
    tracked = tracker.update(detections)

    if (frame_id + 1) % 5 == 0:
        print(f"Frame {frame_id+1}: Tracking {len(tracked)} objects")
        for track in tracked:
            print(f"  ID {track.track_id}: {track.bbox}")

print(f"\nFinal active tracks: {len(tracker.tracks)}")
print(f"Total IDs assigned: {tracker.next_id}")

class DroneDetectionAugmentor:
    """Augmentation strategies for drone detection."""

    @staticmethod
    def random_scale(image, bboxes, scale_range=(0.5, 1.5)):
        """Random scale augmentation."""
        scale = np.random.uniform(*scale_range)
        h, w = image.shape[:2]
        new_h, new_w = int(h * scale), int(w * scale)

        # Scale image (simplified)
        scaled_image = np.random.randint(0, 255, (new_h, new_w, 3), dtype=np.uint8)

        # Scale bboxes
        scaled_bboxes = []
        for bbox in bboxes:
            x1, y1, x2, y2 = bbox
            scaled_bboxes.append((
                int(x1 * scale), int(y1 * scale),
                int(x2 * scale), int(y2 * scale)
            ))

        return scaled_image, scaled_bboxes

    @staticmethod
    def altitude_simulation(image, target_altitude=50):
        """Simulate different altitudes."""
        # Higher altitude = smaller objects, different blur
        blur_level = max(1, target_altitude // 25)

        # Simple averaging blur
        kernel_size = blur_level * 2 + 1
        h, w = image.shape[:2]

        # Simulate smaller objects at higher altitude
        scale_factor = 100 / target_altitude
        new_size = max(50, int(min(h, w) * scale_factor))

        return new_size, scale_factor

    @staticmethod
    def topdown_rotation(image, bboxes, max_angle=30):
        """Simulate slight camera angle variations."""
        angle = np.random.uniform(-max_angle, max_angle)

        # Rotate bboxes (simplified)
        rotated_bboxes = []
        for bbox in bboxes:
            cx = (bbox[0] + bbox[2]) / 2
            cy = (bbox[1] + bbox[3]) / 2
            w = bbox[2] - bbox[0]
            h = bbox[3] - bbox[1]

            # Expand bbox to account for rotation
            new_w = int(abs(w * np.cos(np.radians(angle))) +
                       abs(h * np.sin(np.radians(angle))))
            new_h = int(abs(h * np.cos(np.radians(angle))) +
                       abs(w * np.sin(np.radians(angle))))

            rotated_bboxes.append((
                int(cx - new_w/2), int(cy - new_h/2),
                int(cx + new_w/2), int(cy + new_h/2)
            ))

        return rotated_bboxes

# Example augmentation
np.random.seed(42)
augmentor = DroneDetectionAugmentor()

sample_bboxes = [(100, 100, 200, 180), (300, 150, 380, 250)]

# Simulate altitude change
for alt in [25, 50, 100, 200]:
    new_size, scale = augmentor.altitude_simulation(None, alt)
    print(f"Altitude {alt}m: Object scale factor = {scale:.2f}")

# Test rotation augmentation
rotated = augmentor.topdown_rotation(None, sample_bboxes, max_angle=15)
print(f"\nOriginal bboxes: {sample_bboxes}")
print(f"Rotated bboxes: {rotated}")

from collections import defaultdict

class DroneSurveillanceTracker:
    """Complete surveillance tracking system for drones."""

    def __init__(self, grid_size=100):
        self.grid_size = grid_size
        self.zone_density = defaultdict(int)
        self.track_trajectories = {}
        self.alerts = []

    def process_detection(self, track_id, position, timestamp):
        """Process a tracked object's position."""
        # Update trajectory
        if track_id not in self.track_trajectories:
            self.track_trajectories[track_id] = []
        self.track_trajectories[track_id].append({
            'position': position,
            'timestamp': timestamp
        })

        # Update zone density
        zone = self._get_zone(position)
        self.zone_density[zone] += 1

        # Check for anomalies
        anomalies = self._check_anomalies(track_id, position)

        return anomalies

    def _get_zone(self, position):
        """Get grid zone from position."""
        x, y = position
        return (x // self.grid_size, y // self.grid_size)

    def _check_anomalies(self, track_id, position):
        """Check for anomalous behavior."""
        anomalies = []
        trajectory = self.track_trajectories.get(track_id, [])

        if len(trajectory) < 5:
            return anomalies

        # Check loitering (staying in small area)
        recent = trajectory[-10:]
        positions = [p['position'] for p in recent]
        center = np.mean(positions, axis=0)
        distances = [np.linalg.norm(np.array(p) - center) for p in positions]

        if np.mean(distances) < 30 and len(recent) >= 8:
            anomalies.append({
                'type': 'loitering',
                'track_id': track_id,
                'severity': 'medium'
            })

        # Check speed (sudden stops or acceleration)
        if len(trajectory) >= 2:
            speeds = []
            for i in range(1, len(trajectory)):
                dt = trajectory[i]['timestamp'] - trajectory[i-1]['timestamp']
                if dt > 0:
                    dist = np.linalg.norm(
                        np.array(trajectory[i]['position']) -
                        np.array(trajectory[i-1]['position'])
                    )
                    speeds.append(dist / dt)

            if len(speeds) >= 3:
                speed_var = np.std(speeds) / (np.mean(speeds) + 1e-6)
                if speed_var > 2.0:
                    anomalies.append({
                        'type': 'erratic_movement',
                        'track_id': track_id,
                        'severity': 'low'
                    })

        return anomalies

    def get_heatmap(self):
        """Generate density heatmap."""
        if not self.zone_density:
            return np.zeros((10, 10))

        max_density = max(self.zone_density.values())
        heatmap = np.zeros((10, 10))

        for (zx, zy), count in self.zone_density.items():
            if 0 <= zx < 10 and 0 <= zy < 10:
                heatmap[zy, zx] = count / max_density

        return heatmap

    def generate_report(self):
        """Generate surveillance report."""
        report = {
            'total_tracks': len(self.track_trajectories),
            'total_alerts': len(self.alerts),
            'hot_zones': [],
            'longest_track': 0,
            'average_speed': 0
        }

        # Find hot zones
        sorted_zones = sorted(self.zone_density.items(),
                            key=lambda x: x[1], reverse=True)
        report['hot_zones'] = sorted_zones[:5]

        # Find longest track
        for track_id, trajectory in self.track_trajectories.items():
            if len(trajectory) > report['longest_track']:
                report['longest_track'] = len(trajectory)

        return report

# Example surveillance session
np.random.seed(42)
tracker = DroneSurveillanceTracker(grid_size=100)

print("Running drone surveillance...")
all_anomalies = []

# Simulate 5 tracked objects over 30 seconds
for t in range(30):
    for track_id in range(5):
        # Simulate movement
        base_x = 200 + track_id * 150
        base_y = 300 + (t * 5 if track_id < 3 else np.random.randint(-5, 5))

        position = (
            base_x + np.random.randint(-10, 10),
            base_y + np.random.randint(-10, 10)
        )

        anomalies = tracker.process_detection(track_id, position, t)
        all_anomalies.extend(anomalies)

# Generate report
report = tracker.generate_report()
print(f"\nSurveillance Report:")
print(f"  Tracked objects: {report['total_tracks']}")
print(f"  Total alerts: {len(all_anomalies)}")
print(f"  Anomaly types: {set(a['type'] for a in all_anomalies)}")

# Display heatmap
heatmap = tracker.get_heatmap()
print(f"\nDensity Heatmap (top 3 zones):")
for (zx, zy), count in sorted(tracker.zone_density.items(),
                               key=lambda x: x[1], reverse=True)[:3]:
    print(f"  Zone ({zx}, {zy}): {count} detections")

print("\nSurveillance tracking system operational!")