"""
Power Budget Module - Mission energy planning
"""

import numpy as np
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass


@dataclass
class SensorSpec:
    """Sensor power specifications"""
    name: str
    power_consumption: float  # W
    sampling_rate: float  # Hz
    duty_cycle: float  # 0-1
    critical: bool  # Essential for mission


@dataclass
class MissionPhase:
    """Mission phase power requirements"""
    name: str
    duration: float  # hours
    base_power: float  # W
    active_sensors: List[str]
    propulsion_power: float  # W


class PowerBudget:
    """
    Power budget calculator for hadal missions
    
    Calculates energy requirements and ensures mission feasibility
    """
    
    def __init__(self,
                 mission_duration: float = 72,  # hours
                 battery_capacity: float = 180):  # Wh
        """
        Initialize power budget
        
        Args:
            mission_duration: Planned mission duration
            battery_capacity: Total battery capacity (Wh)
        """
        self.mission_duration = mission_duration
        self.battery_capacity = battery_capacity
        
        # Define standard sensors
        self.available_sensors = {
            'oxygen': SensorSpec('oxygen', 0.5, 0.1, 0.2, True),
            'pressure': SensorSpec('pressure', 0.3, 1.0, 1.0, True),
            'temperature': SensorSpec('temperature', 0.2, 1.0, 1.0, True),
            'acoustic': SensorSpec('acoustic', 2.0, 0.033, 0.1, True),  # 30 sec updates
            'bioluminescence': SensorSpec('bioluminescence', 0.8, 10, 0.5, False),
            'altimeter': SensorSpec('altimeter', 0.4, 1.0, 1.0, True),
            'camera': SensorSpec('camera', 5.0, 0.1, 0.05, False),
            'ctd': SensorSpec('ctd', 1.0, 0.1, 0.1, True)
        }
        
        # Define mission phases
        self.phases = [
            MissionPhase('descent', 6, 1.0, 
                        ['pressure', 'temperature', 'acoustic'], 0.5),
            MissionPhase('survey', 48, 2.5,
                        ['oxygen', 'pressure', 'temperature', 'acoustic', 
                         'bioluminescence', 'altimeter', 'ctd'], 0.3),
            MissionPhase('ascent', 8, 1.0,
                        ['pressure', 'temperature', 'acoustic'], 0.3),
            MissionPhase('standby', 10, 0.5,
                        ['pressure', 'temperature'], 0)
        ]
        
        # Harvesting parameters
        self.harvesting_rate = 0.0028  # W (2.8 mW from validation)
        self.harvesting_efficiency = 0.87
        
        # Results
        self.budget = {}
    
    def calculate_requirements(self,
                              sensors: List[str],
                              sampling_rate: Optional[Dict[str, float]] = None) -> float:
        """
        Calculate power requirements for sensor suite
        
        Args:
            sensors: List of sensor names
            sampling_rate: Optional custom sampling rates
            
        Returns:
            Total power consumption in W
        """
        total_power = 0.0
        
        for sensor_name in sensors:
            if sensor_name not in self.available_sensors:
                continue
            
            sensor = self.available_sensors[sensor_name]
            
            # Use custom sampling rate if provided
            if sampling_rate and sensor_name in sampling_rate:
                rate_factor = sampling_rate[sensor_name] / sensor.sampling_rate
            else:
                rate_factor = 1.0
            
            # Calculate average power (accounting for duty cycle)
            avg_power = (sensor.power_consumption * 
                        sensor.duty_cycle * 
                        rate_factor)
            
            total_power += avg_power
        
        return total_power
    
    def calculate_phase_power(self, phase: MissionPhase) -> float:
        """Calculate power for a mission phase"""
        sensor_power = self.calculate_requirements(phase.active_sensors)
        total_power = sensor_power + phase.base_power + phase.propulsion_power
        return total_power
    
    def calculate_mission_energy(self) -> Dict[str, float]:
        """
        Calculate total mission energy requirements
        
        Returns:
            Dictionary with phase energies and total
        """
        phase_energies = {}
        total_energy = 0.0
        
        for phase in self.phases:
            power = self.calculate_phase_power(phase)
            energy = power * phase.duration
            phase_energies[phase.name] = energy
            total_energy += energy
        
        phase_energies['total'] = total_energy
        self.budget = phase_energies
        
        return phase_energies
    
    def add_harvesting(self,
                      harvested_power: float,
                      harvesting_duty_cycle: float = 0.8) -> Dict[str, float]:
        """
        Add energy harvesting contribution
        
        Args:
            harvested_power: Average harvested power (W)
            harvesting_duty_cycle: Fraction of time harvesting
            
        Returns:
            Updated energy budget
        """
        if not self.budget:
            self.calculate_mission_energy()
        
        total_harvested = (harvested_power * 
                          harvesting_duty_cycle * 
                          self.mission_duration)
        
        net_energy = self.budget['total'] - total_harvested
        
        return {
            'mission_energy_wh': self.budget['total'],
            'harvested_energy_wh': total_harvested,
            'net_energy_wh': net_energy,
            'battery_capacity_wh': self.battery_capacity,
            'margin_wh': self.battery_capacity - net_energy,
            'margin_percent': (self.battery_capacity - net_energy) / self.battery_capacity * 100
        }
    
    def check_feasibility(self,
                         reserve_percent: float = 20) -> Tuple[bool, Dict]:
        """
        Check if mission is feasible with current power budget
        
        Args:
            reserve_percent: Required reserve capacity
            
        Returns:
            (feasible, diagnostics)
        """
        if not self.budget:
            self.calculate_mission_energy()
        
        required_energy = self.budget['total']
        available_energy = self.battery_capacity * (1 - reserve_percent / 100)
        
        feasible = required_energy <= available_energy
        
        # Add harvesting if available
        harvesting_energy = self.harvesting_rate * self.mission_duration
        with_harvesting = required_energy <= (available_energy + harvesting_energy)
        
        # Phase-by-phase check
        phase_check = {}
        cumulative_energy = 0
        
        for phase in self.phases:
            phase_power = self.calculate_phase_power(phase)
            phase_energy = phase_power * phase.duration
            cumulative_energy += phase_energy
            
            phase_check[phase.name] = {
                'energy_wh': phase_energy,
                'cumulative_wh': cumulative_energy,
                'feasible': cumulative_energy <= available_energy
            }
        
        return feasible, {
            'required_wh': required_energy,
            'available_wh': available_energy,
            'harvesting_wh': harvesting_energy,
            'feasible': feasible,
            'feasible_with_harvesting': with_harvesting,
            'phase_check': phase_check,
            'margin_wh': available_energy - required_energy,
            'margin_percent': (available_energy - required_energy) / available_energy * 100
        }
    
    def optimize_sensors(self,
                        critical_only: bool = False,
                        max_power: Optional[float] = None) -> List[str]:
        """
        Optimize sensor selection for power constraints
        
        Args:
            critical_only: Only use critical sensors
            max_power: Maximum allowed power consumption
            
        Returns:
            Optimized sensor list
        """
        if critical_only:
            return [s for s, spec in self.available_sensors.items() if spec.critical]
        
        if max_power is None:
            return list(self.available_sensors.keys())
        
        # Greedy selection based on power/utility ratio
        sensors_with_value = []
        for name, spec in self.available_sensors.items():
            # Simple utility heuristic
            if spec.critical:
                utility = 10
            elif name == 'bioluminescence':
                utility = 8
            elif name == 'camera':
                utility = 5
            else:
                utility = 3
            
            avg_power = spec.power_consumption * spec.duty_cycle
            value_ratio = utility / avg_power if avg_power > 0 else float('inf')
            
            sensors_with_value.append((name, avg_power, value_ratio))
        
        # Sort by value ratio
        sensors_with_value.sort(key=lambda x: x[2], reverse=True)
        
        # Select until power budget exceeded
        selected = []
        total_power = 0
        
        for name, power, _ in sensors_with_value:
            if total_power + power <= max_power:
                selected.append(name)
                total_power += power
        
        return selected
    
    def get_power_profile(self) -> Dict[str, List[float]]:
        """Get power consumption profile over mission"""
        time_points = []
        power_points = []
        
        current_time = 0
        
        for phase in self.phases:
            phase_power = self.calculate_phase_power(phase)
            
            for _ in range(int(phase.duration)):
                time_points.append(current_time)
                power_points.append(phase_power)
                current_time += 1
        
        return {
            'time_hours': time_points,
            'power_watts': power_points
        }
    
    def __repr__(self) -> str:
        if not self.budget:
            self.calculate_mission_energy()
        
        feasible, _ = self.check_feasibility()
        status = "FEASIBLE" if feasible else "INFEASIBLE"
        
        return (f"PowerBudget(duration={self.mission_duration}h, "
                f"energy={self.budget['total']:.1f}Wh, "
                f"status={status})")
