"""
Communication Module - Acoustic modem integration
"""

import numpy as np
from typing import List, Dict, Optional, Any
from dataclasses import dataclass
from enum import Enum


class MessageType(Enum):
    """Types of acoustic messages"""
    POSITION = "position"
    COMMAND = "command"
    SENSOR_DATA = "sensor_data"
    ALERT = "alert"
    ACKNOWLEDGMENT = "ack"


@dataclass
class Message:
    """Acoustic message"""
    msg_type: MessageType
    sender_id: int
    receiver_id: Optional[int]  # None for broadcast
    timestamp: float
    data: Dict[str, Any]
    priority: int  # 1 (low) to 5 (high)


class AcousticCommunication:
    """
    Acoustic communication system for underwater vehicles
    
    Handles message encoding/decoding, ranging, and network management
    Supports 1-10 kbps data rate at 12 kHz carrier frequency
    """
    
    def __init__(self,
                 frequency: float = 12000,  # Hz
                 bandwidth: float = 1000,  # Hz
                 range_max: float = 500,  # meters
                 data_rate: float = 1000):  # bps
        """
        Initialize acoustic communication
        
        Args:
            frequency: Carrier frequency (Hz)
            bandwidth: Channel bandwidth (Hz)
            range_max: Maximum communication range (m)
            data_rate: Data rate (bps)
        """
        self.frequency = frequency
        self.bandwidth = bandwidth
        self.range_max = range_max
        self.data_rate = data_rate
        
        # Sound velocity (nominal)
        self.sound_velocity = 1500.0  # m/s
        
        # Network state
        self.vehicle_id = 0
        self.neighbors = {}
        self.message_queue = []
        self.sent_messages = []
        self.received_messages = []
        
        # Statistics
        self.packets_sent = 0
        self.packets_received = 0
        self.packets_lost = 0
        self.total_bytes = 0
    
    def set_vehicle_id(self, vehicle_id: int):
        """Set local vehicle ID"""
        self.vehicle_id = vehicle_id
    
    def send(self,
            msg_type: MessageType,
            data: Dict[str, Any],
            receiver_id: Optional[int] = None,
            priority: int = 3) -> bool:
        """
        Send a message
        
        Args:
            msg_type: Message type
            data: Message data
            receiver_id: Target vehicle (None for broadcast)
            priority: Message priority (1-5)
            
        Returns:
            True if sent successfully
        """
        # Check range to receiver
        if receiver_id is not None and receiver_id in self.neighbors:
            distance = self.neighbors[receiver_id]['distance']
            if distance > self.range_max:
                self.packets_lost += 1
                return False
        
        # Create message
        msg = Message(
            msg_type=msg_type,
            sender_id=self.vehicle_id,
            receiver_id=receiver_id,
            timestamp=len(self.sent_messages),
            data=data,
            priority=priority
        )
        
        # Simulate transmission
        success = self._transmit(msg)
        
        if success:
            self.sent_messages.append(msg)
            self.packets_sent += 1
            self.total_bytes += self._estimate_size(msg)
        
        return success
    
    def _transmit(self, msg: Message) -> bool:
        """Simulate acoustic transmission"""
        # Calculate propagation loss
        if msg.receiver_id in self.neighbors:
            distance = self.neighbors[msg.receiver_id]['distance']
        else:
            distance = np.random.uniform(10, self.range_max)
        
        # Spreading loss (20 log10(r))
        spreading_loss = 20 * np.log10(distance + 1)
        
        # Absorption loss (simplified)
        absorption = 0.01 * distance  # dB/km
        
        total_loss = spreading_loss + absorption
        
        # Random packet loss based on SNR
        snr = 30 - total_loss  # dB
        if snr < 10:
            loss_prob = 0.5
        elif snr < 20:
            loss_prob = 0.2
        else:
            loss_prob = 0.05
        
        return np.random.random() > loss_prob
    
    def _estimate_size(self, msg: Message) -> int:
        """Estimate message size in bytes"""
        size = 20  # Header
        
        # Add data size (rough estimate)
        for key, value in msg.data.items():
            size += len(key)
            if isinstance(value, (int, float)):
                size += 8
            elif isinstance(value, str):
                size += len(value)
            elif isinstance(value, (list, np.ndarray)):
                size += len(value) * 8
        
        return size
    
    def receive(self,
               timeout: float = 30.0) -> List[Message]:
        """
        Receive messages
        
        Args:
            timeout: Receive timeout in seconds
            
        Returns:
            List of received messages
        """
        # Simulate receiving messages from queue
        received = []
        
        for msg in self.message_queue:
            if (msg.receiver_id is None or 
                msg.receiver_id == self.vehicle_id):
                
                received.append(msg)
                self.received_messages.append(msg)
                self.packets_received += 1
        
        # Clear queue
        self.message_queue = []
        
        return received
    
    def add_to_queue(self, msg: Message):
        """Add message to receive queue"""
        self.message_queue.append(msg)
    
    def update_neighbor(self,
                       vehicle_id: int,
                       position: np.ndarray,
                       signal_strength: float):
        """Update neighbor information"""
        self.neighbors[vehicle_id] = {
            'position': position.copy(),
            'signal_strength': signal_strength,
            'last_seen': len(self.received_messages),
            'distance': np.linalg.norm(position)
        }
    
    def calculate_range(self,
                       time_of_flight: float,
                       sound_velocity: Optional[float] = None) -> float:
        """
        Calculate range from time of flight
        
        Args:
            time_of_flight: Acoustic travel time (s)
            sound_velocity: Sound velocity (m/s)
            
        Returns:
            Range in meters
        """
        if sound_velocity is None:
            sound_velocity = self.sound_velocity
        
        return sound_velocity * time_of_flight
    
    def estimate_data_rate(self, distance: float) -> float:
        """
        Estimate achievable data rate at given distance
        
        Args:
            distance: Communication distance (m)
            
        Returns:
            Data rate in bps
        """
        if distance > self.range_max:
            return 0
        
        # Rate decreases with distance
        rate = self.data_rate * (1 - distance / self.range_max)
        
        return max(100, rate)
    
    def get_network_status(self) -> Dict:
        """Get network status"""
        return {
            'vehicle_id': self.vehicle_id,
            'neighbors': len(self.neighbors),
            'packets_sent': self.packets_sent,
            'packets_received': self.packets_received,
            'packets_lost': self.packets_lost,
            'total_bytes': self.total_bytes,
            'queue_size': len(self.message_queue),
            'data_rate_bps': self.data_rate,
            'range_max_m': self.range_max
        }
    
    def link_budget(self,
                   distance: float,
                   tx_power: float = 10,  # W
                   noise_level: float = -80) -> float:  # dB
        """
        Calculate acoustic link budget
        
        Args:
            distance: Communication distance (m)
            tx_power: Transmit power (W)
            noise_level: Noise level (dB)
            
        Returns:
            SNR in dB
        """
        # Convert transmit power to dB
        tx_dB = 10 * np.log10(tx_power * 1000)  # dB re 1 mW
        
        # Spreading loss
        spreading = 20 * np.log10(distance)
        
        # Absorption
        absorption = 0.01 * distance
        
        received_power = tx_dB - spreading - absorption
        
        snr = received_power - noise_level
        
        return snr
    
    def __repr__(self) -> str:
        status = self.get_network_status()
        return (f"AcousticCommunication(neighbors={status['neighbors']}, "
                f"sent={status['packets_sent']}, "
                f"rate={status['data_rate_bps']:.0f}bps)")
