"""
Implémentation du protocole LEACH-C (Centralisé)
"""

import random
import math
from node import Node
from metrics import Metrics
from config import BS_POSITION


class LEACHC:
    """
    Implémentation de LEACH-C en version centralisée.
    
    La BS reçoit l'état de tous les nœuds et calcule les clusters optimaux.
    """
    
    def __init__(self, nodes, probability_ch, packet_size):
        """
        Initialise le protocole LEACH-C.
        
        Args:
            nodes (list): Liste des nœuds
            probability_ch (float): Probabilité d'activité des nœuds (non utilisée pour élection)
            packet_size (int): Taille des paquets en bits
        """
        self.nodes = nodes
        self.probability_ch = probability_ch  # Utilisé pour activité des nœuds
        self.packet_size = packet_size
        self.metrics = Metrics()
        self.ch_nodes = []
        self.round_num = 0
    
    def calculate_optimal_clusters(self):
        """
        Calcule une partition optimale des clusters.
        Heuristique : sélectionner les nœuds avec le plus d'énergie comme CHs.
        Nombre de CHs ≈ sqrt(N)/2 est une bonne heuristique.
        
        Returns:
            list: Liste des CHs calculés par la BS
        """
        alive_nodes = [n for n in self.nodes if n.is_alive]
        
        if not alive_nodes:
            return []
        
        # Nombre optimal de CHs (heuristique classique)
        num_ch_optimal = max(1, int(math.sqrt(len(alive_nodes)) / 2))
        
        # Trier les nœuds par énergie décroissante
        sorted_nodes = sorted(alive_nodes, key=lambda n: n.energy, reverse=True)
        
        # Sélectionner les CHs
        ch_candidates = sorted_nodes[:num_ch_optimal]
        
        return ch_candidates
    
    def form_clusters(self, ch_nodes):
        """
        La BS assigne les nœuds à leurs CHs respectifs.
        Heuristique : chaque nœud rejoint le CH le plus proche.
        
        Args:
            ch_nodes (list): Liste des CHs
        """
        # Reset cluster affiliation
        for node in self.nodes:
            node.cluster_id = None
        
        # Assignation des nœuds aux CHs
        for node in self.nodes:
            if not node.is_alive or node.is_cluster_head:
                continue
            
            # Trouver le CH le plus proche
            closest_ch = None
            min_distance = float('inf')
            
            for ch in ch_nodes:
                if ch.is_alive:
                    dist = node.distance_to(ch.x, ch.y)
                    if dist < min_distance:
                        min_distance = dist
                        closest_ch = ch
            
            if closest_ch:
                node.cluster_id = closest_ch.node_id
    
    def send_ch_info_to_bs(self, ch_nodes):
        """
        Les CHs envoient leurs info à la BS.
        Coût énergétique : transmission du CH vers la BS.
        
        Args:
            ch_nodes (list): Liste des CHs
        """
        for ch in ch_nodes:
            if ch.is_alive:
                # Transmettre l'info (simplification : peu d'octets)
                info_bits = 32  # Position et énergie
                distance_to_bs = ch.distance_to(BS_POSITION[0], BS_POSITION[1])
                ch.transmit(info_bits, distance_to_bs)
    
    def communication_phase(self, ch_nodes):
        """
        Phase de communication :
        - Les nœuds envoient les données à leur CH
        - Les CHs agrègent et envoient à la BS
        
        Args:
            ch_nodes (list): Liste des CHs
        
        Returns:
            tuple: (packets_to_ch, packets_to_bs)
        """
        packets_to_ch = 0
        packets_to_bs = 0
        
        # Phase 1 : Nœuds → CHs
        for node in self.nodes:
            if not node.is_alive or node.is_cluster_head:
                continue
            
            # Déterminer si le nœud doit envoyer (probabilité p)
            if random.random() < self.probability_ch:
                closest_ch = None
                min_distance = float('inf')
                
                for ch in ch_nodes:
                    if ch.is_alive:
                        dist = node.distance_to(ch.x, ch.y)
                        if dist < min_distance:
                            min_distance = dist
                            closest_ch = ch
                
                if closest_ch:
                    # Transmission du nœud au CH
                    energy_tx = node.transmit(self.packet_size, min_distance)
                    
                    # Réception au CH
                    energy_rx = closest_ch.receive(self.packet_size)
                    
                    packets_to_ch += 1
        
        # Phase 2 : CHs → BS (après agrégation)
        for ch in ch_nodes:
            if not ch.is_alive:
                continue
            
            # Compter les nœuds dans le cluster
            cluster_members = sum(1 for n in self.nodes 
                                 if n.cluster_id == ch.node_id and n.is_alive)
            
            if cluster_members > 0:
                # Agrégation des données reçues
                aggregated_bits = cluster_members * self.packet_size
                energy_agg = ch.aggregate(aggregated_bits)
                
                # Transmission à la BS
                distance_to_bs = ch.distance_to(BS_POSITION[0], BS_POSITION[1])
                energy_tx_bs = ch.transmit(aggregated_bits, distance_to_bs)
                
                packets_to_bs += 1
        
        return packets_to_ch, packets_to_bs
    
    def run_round(self):
        """
        Exécute une ronde complète du protocole LEACH-C.
        
        Returns:
            bool: True si la ronde a réussi, False si aucun CH disponible
        """
        self.round_num += 1
        
        # Reset pour la nouvelle ronde
        for node in self.nodes:
            node.reset_for_round()
        
        # Phase 1 : BS calcule les CHs optimaux
        ch_candidates = self.calculate_optimal_clusters()
        
        if not ch_candidates:
            # Aucun nœud vivant
            self.metrics.record_round(self.round_num, self.nodes, [], 0, 0, muted=True)
            self.metrics.update_dead_nodes(self.nodes)
            return False
        
        # Marquer les CHs
        self.ch_nodes = []
        for ch in ch_candidates:
            ch.is_cluster_head = True
            self.ch_nodes.append(ch)
        
        # Phase 2 : Formation des clusters par la BS
        self.form_clusters(self.ch_nodes)
        
        # Phase 3 : Envoi de l'info des CHs à la BS
        self.send_ch_info_to_bs(self.ch_nodes)
        
        # Phase 4 : Communication
        packets_to_ch, packets_to_bs = self.communication_phase(self.ch_nodes)
        
        # Phase 5 : Mobilité - déplacement aléatoire des nœuds
        for node in self.nodes:
            if node.is_alive:
                node.move()
        
        # Enregistrement des métriques
        self.metrics.record_round(self.round_num, self.nodes, self.ch_nodes, 
                                 packets_to_ch, packets_to_bs, muted=False)
        self.metrics.update_dead_nodes(self.nodes)
        
        return True
    
    def run_simulation(self, num_rounds):
        """
        Lance la simulation complète.
        
        Args:
            num_rounds (int): Nombre de rounds à simuler
        """
        for _ in range(num_rounds):
            # Vérifier s'il y a encore des nœuds vivants
            alive_count = sum(1 for n in self.nodes if n.is_alive)
            if alive_count == 0:
                break
            
            self.run_round()
    
    def get_metrics(self, total_rounds):
        """Retourne les métriques de la simulation."""
        return self.metrics.get_summary(total_rounds)
    
    def get_detailed_metrics(self):
        """Retourne les données détaillées de chaque round."""
        return self.metrics.get_rounds_data()
