f1_nodes_functions.py

import json
import sys
import argparse
import re
from pathlib import Path
from typing import List, Dict, Tuple, cast
from difflib import SequenceMatcher
import numpy as np



try:
    from rapidfuzz import fuzz
except ImportError:
    class fuzz:
        @staticmethod
        def ratio(a, b):
            return SequenceMatcher(None, a, b).ratio() * 100

try:
    from scipy.optimize import linear_sum_assignment
    HAS_SCIPY = True
except ImportError:
    HAS_SCIPY = False

# ========================= ANSI COLORS =========================
class C:
    HDR = "\033[95m"
    OKB = "\033[94m"
    OKG = "\033[92m"
    W   = "\033[93m"
    R   = "\033[91m"
    B   = "\033[1m"
    DIM = "\033[2m"
    END = "\033[0m"
    HEADER = "\033[95m"
    OKCYAN = "\033[96m"
    FAIL = "\033[91m"
    WARNING = "\033[93m"
    ENDC = "\033[0m"



LOG_FILE = None

def custom_print(message: str = "", end: str = "\n"):

    global LOG_FILE
    full_message = message + end
    
    if LOG_FILE:
        LOG_FILE.write(full_message)
        LOG_FILE.flush() 
        
    sys.stdout.write(full_message)
    sys.stdout.flush()

# ========================= FUNCTIONS=========================

def _normalize_name_nodes(name: str) -> str:
    """Standardize node names for comparison."""
    return str(name or "").replace('-', '').replace('_', '').upper()

def rf_score(a: str, b: str) -> float:
    """Calculate the similarity score between two node names."""
    a_norm = _normalize_name_nodes(a)
    b_norm = _normalize_name_nodes(b)
    if hasattr(fuzz, "token_set_ratio"):
        return float(fuzz.token_set_ratio(a_norm, b_norm))
    else:
        return 100.0 * SequenceMatcher(None, a_norm, b_norm).ratio()


def load_nodes(data: Dict) -> List[Dict]:
    nodes = []
    
    devices = data["result"]["network_topology"].get("devices", [])
    for device in devices:

        nodes.append({
            "device_name": device.get("device_name"),
            "network_layer": str(device.get("network_layer", "access")).lower(),
            "device_type": device.get("device_type", "router"),
            "location_site": str(device.get("location_site")).lower(),
        })
        
    return nodes


def hungarian_match_nodes(sim: np.ndarray):
    """Apply the Hungarian algorithm."""
    if not HAS_SCIPY:
        return []
    
    m, n = sim.shape
    cost = 100.0 - sim
    row_ind, col_ind = linear_sum_assignment(cost) 
    pairs = [(r, c, float(sim[r, c])) for r, c in zip(row_ind, col_ind) if r < m and c < n]
    return pairs



def normalize_gen_to_ref_topology(ref_nodes, gen_nodes, min_sim: float = 9.0):
    """Normalize GEN sites and types to match REF using the Hungarian algorithm."""
    if not gen_nodes or not ref_nodes:
        return gen_nodes
    
    ref_sites = list(set(n.get('location_site') for n in ref_nodes))
    ref_types = list(set(n.get('device_type') for n in ref_nodes))
    
    gen_sites = list(set(n.get('location_site') for n in gen_nodes))
    gen_types = list(set(n.get('device_type') for n in gen_nodes))
    
    
    site_mapping = {}
    type_mapping = {}
    

    if ref_sites and gen_sites:
        site_sim = build_similarity_matrix_nodes(ref_sites, gen_sites)
        
        for j, gen_site in enumerate(gen_sites):
            best_ref_idx = max(range(len(ref_sites)), key=lambda i: site_sim[i][j])
            best_score = site_sim[best_ref_idx][j]
            
            if best_score >= min_sim:
                site_mapping[gen_site] = ref_sites[best_ref_idx]
                print(f"  GEN site: {gen_site} → REF site (max similarity): {ref_sites[best_ref_idx]} (score: {best_score:.2f})")
            else:
                print(f"  Warning: GEN site '{gen_site}' has no reliable match (best score: {best_score:.2f})")
        
        print(f"  Applying site normalization:")
        for node in gen_nodes:
            original_site = node.get('location_site')
            if original_site in site_mapping:
                new_site = site_mapping[original_site]
                node['location_site'] = new_site
                print(f"    Node {node['device_name']}: site {original_site} → {new_site}")
    
    print(f"  Pre-normalizing types containing 'server':")
    for node in gen_nodes:
        original_type = node.get('device_type', '')
        if original_type and any(keyword in original_type.lower() for keyword in ['server', 'videodelivery', 'haproxy']):
            node['device_type'] = 'Server'
            print(f"    Node {node['device_name']}: type {original_type} → Server (contains 'server')")
    
    if ref_types and gen_types:
        print(f"  Normalizing types by layer:")
        
        ref_nodes_by_layer = {}
        gen_nodes_by_layer = {}
        
        for node in ref_nodes:
            layer = node.get('network_layer')
            if layer not in ref_nodes_by_layer:
                ref_nodes_by_layer[layer] = []
            ref_nodes_by_layer[layer].append(node)
        
        for node in gen_nodes:
            layer = node.get('network_layer')
            if layer not in gen_nodes_by_layer:
                gen_nodes_by_layer[layer] = []
            gen_nodes_by_layer[layer].append(node)
        
        all_layers = set(ref_nodes_by_layer.keys()) | set(gen_nodes_by_layer.keys())
        
        for layer in sorted(all_layers):
            ref_layer_nodes = ref_nodes_by_layer.get(layer, [])
            gen_layer_nodes = gen_nodes_by_layer.get(layer, [])
            
            if not ref_layer_nodes or not gen_layer_nodes:
                continue
            
            ref_layer_types = list(set(n.get('device_type') for n in ref_layer_nodes))
            gen_layer_types = list(set(n.get('device_type') for n in gen_layer_nodes))
            
            print(f"    Layer {layer}: {len(ref_layer_types)} types REF vs {len(gen_layer_types)} types GEN")
            
            type_sim = build_similarity_matrix_nodes(ref_layer_types, gen_layer_types)
            type_pairs = hungarian_match_nodes(type_sim)
            
            layer_type_mapping = {}
            for i, j, s in type_pairs:
                if s >= min_sim:
                    layer_type_mapping[gen_layer_types[j]] = ref_layer_types[i]
                    print(f"      Type GEN: {gen_layer_types[j]} → Type REF: {ref_layer_types[i]} (score: {s:.2f})")
            
            type_mapping.update(layer_type_mapping)
        
        print(f"  Applying type normalization:")
        for node in gen_nodes:
            original_type = node.get('device_type')
            if original_type in type_mapping:
                new_type = type_mapping[original_type]
                node['device_type'] = new_type
                print(f"    Node {node['device_name']}: type {original_type} → {new_type}")
    
    print(f"Normalization complete: {len(site_mapping)} sites and {len(type_mapping)} types normalized")
    return gen_nodes, site_mapping, type_mapping



def get_mapping(ref_nodes, gen_nodes, min_sim: float = 9.0):
    """Hierarchical mapping: layer → site → type using the Hungarian algorithm (adapted from gen_F1_nodes.py logic)."""
    if not gen_nodes:
        return {}
        
    mapping: Dict[str, str] = {}
    
    ref_groups: Dict[str, Dict[str, List[str]]] = {} 
    gen_groups: Dict[str, Dict[str, List[str]]] = {}
    
    for n in ref_nodes:
        if n["device_name"] and isinstance(n["device_name"], str): 
            layer = n["network_layer"]
            site = n.get("location_site")
            device_type = n.get("device_type")
            
            if layer not in ref_groups:
                ref_groups[layer] = {}
            if site not in ref_groups[layer]:
                ref_groups[layer][site] = {}
            if device_type not in ref_groups[layer][site]:
                ref_groups[layer][site][device_type] = []
            
            ref_groups[layer][site][device_type].append(n["device_name"])
            


    for n in gen_nodes:
        if n["device_name"] and isinstance(n["device_name"], str):
            layer = n["network_layer"]
            site = n.get("location_site")
            device_type = n.get("device_type")
            
            if layer not in gen_groups:
                gen_groups[layer] = {}
            if site not in gen_groups[layer]:
                gen_groups[layer][site] = {}
            if device_type not in gen_groups[layer][site]:
                gen_groups[layer][site][device_type] = []
            
            gen_groups[layer][site][device_type].append(n["device_name"])
    
    all_keys = set(ref_groups.keys()) | set(gen_groups.keys())
    
    custom_print(f"\n{C.HDR}{C.B}===== Hierarchical mapping: Layer → Site → Type (Hungarian algorithm) ====={C.END}")
    
    for layer in sorted(all_keys, key=lambda x: (x != 'null', x)):
        ref_layer_dict = ref_groups.get(layer, {})
        gen_layer_dict = gen_groups.get(layer, {})
        
        display_layer = layer.upper() if layer != 'null' else layer
        custom_print(f"\n{C.OKB}--- Layer: {display_layer} ---{C.END}")
        


        all_sites = set(ref_layer_dict.keys()) | set(gen_layer_dict.keys())
        
        for site in sorted(all_sites):
            ref_site_types = ref_layer_dict.get(site, {})
            gen_site_types = gen_layer_dict.get(site, {})
            
            if not ref_site_types and not gen_site_types:
                continue
                
            custom_print(f"  Site: {site}")
            

            all_types = set(ref_site_types.keys()) | set(gen_site_types.keys())
            
            for device_type in sorted(all_types):
                ref_names = ref_site_types.get(device_type, [])
                gen_names = gen_site_types.get(device_type, [])
                
                if not ref_names or not gen_names:
                    if ref_names:
                        custom_print(f"    {C.R}Local FN {device_type}: {', '.join(ref_names)}{C.END}")
                    if gen_names:
                        custom_print(f"    {C.R}Local FP {device_type}: {', '.join(gen_names)}{C.END}")
                    continue
                sim = build_similarity_matrix_nodes(ref_names, gen_names)
                pairs = hungarian_match_nodes(sim)
                
                group_matches = {}
                matches = []
                for i, j, s in pairs:
                    r_name = ref_names[i]
                    g_name = gen_names[j]
                    if s >= min_sim and r_name and g_name:
                        group_matches[r_name] = g_name
                        mapping[r_name] = g_name
                        matches.append((r_name, g_name, s))
                
                if matches:
                    custom_print(f"    {C.OKG}TP {device_type}: {len(matches)} match(s){C.END}")
                    for r_name, g_name, score in matches:
                        custom_print(f"      - REF: {r_name:<20} <--> GEN: {g_name:<20} (Score: {score:.2f})")
                else:
                    custom_print(f"    {C.W}No matches for {device_type} (threshold: {min_sim:.1f}%){C.END}")
                

                fn_nodes = [name for name in ref_names if name not in group_matches]
                fp_nodes = [name for name in gen_names if name not in group_matches.values()]
                
                if fn_nodes:
                    custom_print(f"    {C.R}Local FN {device_type} ({len(fn_nodes)}): {', '.join(fn_nodes[:3])}{'...' if len(fn_nodes) > 3 else ''}{C.END}")
                if fp_nodes:
                    custom_print(f"    {C.R}Local FP {device_type} ({len(fp_nodes)}): {', '.join(fp_nodes[:3])}{'...' if len(fp_nodes) > 3 else ''}{C.END}")
                    
    custom_print(f"\n{C.OKG}Final mapping: {len(mapping)} correspondences{C.END}")
    return mapping

def rf_score_nodes(a: str, b: str) -> float:
    """Calculate the similarity score between two node names."""
    a_norm = _normalize_name_nodes(a)
    b_norm = _normalize_name_nodes(b)
    if hasattr(fuzz, 'token_set_ratio'):
        return float(fuzz.token_set_ratio(a_norm, b_norm))
    else:
        return fuzz.ratio(a_norm, b_norm)

def build_similarity_matrix_nodes(ref_names: List[str], gen_names: List[str]) -> np.ndarray:
    """Build the similarity matrix (based on gen_F1_nodes.py logic)."""
    M = np.zeros((len(ref_names), len(gen_names)), dtype=float)
    for i, a in enumerate(ref_names):
        for j, b in enumerate(gen_names):
            if a and b:
                M[i, j] = rf_score_nodes(a, b)
    return M

def hungarian_match(sim: np.ndarray):
    """Apply the Hungarian algorithm."""
    if not HAS_SCIPY:
            custom_print(f"{C.R}Error: SciPy is not installed. Required for the Hungarian algorithm.{C.END}")

    m, n = sim.shape
    cost = 100.0 - sim
    row_ind, col_ind = linear_sum_assignment(cost) 
    pairs = [(r, c, cast(float, sim[r, c])) for r, c in zip(row_ind, col_ind) if r < m and c < n]
    return pairs

def convert_nodes_format(nodes: List[Dict]) -> List[Dict]:
    converted = []
    for n in nodes:
        device_type = n.get("device_type", "router")
        
        converted.append({
            "device_name": n["name"],
            "network_layer": n["layer"],
            "device_type": device_type,
            "location_site": n.get("site", "default"),
            "connections": []  
        })
    return converted

def mapping_by_layer(ref_nodes: List[Dict], gen_nodes: List[Dict],
                     min_sim: float, test_filename: str) -> Dict[str, str]:
    """Create the initial REF->GEN mapping using hierarchical layer/site/type grouping."""
    
    if not gen_nodes:
        return {}
    
    ref_nodes_converted = convert_nodes_format(ref_nodes)
    gen_nodes_converted = convert_nodes_format(gen_nodes)
    
    mapping = get_mapping(ref_nodes_converted, gen_nodes_converted, min_sim)
    
    custom_print(f"\n{C.OKG}Final mapping size (TP_nodes):{C.END} {len(mapping)}")
    return mapping

def calculate_node_metrics(mapping: Dict[str, str], ref_nodes: List[Dict], gen_nodes: List[Dict]):
    """
    Calcul of TP, FP, FN and F1.
    """
    
    ref_count = len(ref_nodes) 
    gen_count = len(gen_nodes) 
    
    if gen_count == 0 and len(mapping) == 0:
        TP_nodes = 0
        FP_nodes = 0
        FN_nodes = ref_count 
        P_nodes = 0.0
        R_nodes = 0.0
        F1_nodes = 0.0
        percent_mapping = 0.0
    else:
        TP_nodes = len(mapping)
        FN_nodes = ref_count - TP_nodes
        FP_nodes = gen_count - TP_nodes
        
        P_nodes = TP_nodes / (TP_nodes + FP_nodes) if (TP_nodes + FP_nodes) else 0.0
        R_nodes = TP_nodes / (TP_nodes + FN_nodes) if (TP_nodes + FN_nodes) else 0.0
        F1_nodes = 2 * P_nodes * R_nodes / (P_nodes + R_nodes) if (P_nodes + R_nodes) else 0.0
        percent_mapping = R_nodes * 100.0

    # -------------------------

    custom_print(f"\n{C.HDR}{C.B}===== NODE F1-SCORE RESULTS (Final Calculation) ====={C.END}")
    custom_print(f"  • REF nodes (Total): {ref_count}")
    custom_print(f"  • GEN nodes (Total): {gen_count}") 
    custom_print("-" * 50)
    custom_print(f"  • True Positives ({C.B}TP_nodes{C.END}): {TP_nodes}")
    custom_print(f"  • False Positives ({C.B}FP_nodes{C.END}): {FP_nodes}") 
    custom_print(f"  • False Negatives ({C.B}FN_nodes{C.END}): {FN_nodes}")
    custom_print("-" * 50)
    custom_print(f"  • Precision ($P_{{nodes}}$): {P_nodes:.6f}")
    custom_print(f"  • Recall ($R_{{nodes}}$): {R_nodes:.6f}")
    custom_print(f"  • {C.B}Mapping percentage{C.END}: {percent_mapping:.2f} %")
    custom_print(f"{C.OKG}  • Node F1-score ($F1_{{nodes}}$): {F1_nodes:.6f}{C.END}")

    return {
        "TP_nodes": TP_nodes, "FP_nodes": FP_nodes, "FN_nodes": FN_nodes,
        "P_nodes": P_nodes, "R_nodes": R_nodes, "F1_nodes": F1_nodes,
        "ref_count": ref_count, "gen_count": gen_count, "percent_mapping": percent_mapping
    }