core.py

import networkx as nx
import numpy as np
import scipy.sparse as sparse
import osqp
import time
import warnings
import random
import concurrent.futures
import os, sys
_root = os.path.dirname(os.path.abspath(__file__))
_pr = os.path.dirname(_root)
if _pr not in sys.path:
    sys.path.append(_pr)

from .cycles import get_triangles, get_shortest_hop_count_cycle_bfs
from .dijkstra import parallel_worker_F_adapted
from .utils import calculate_overlap

# Suppress warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)
warnings.filterwarnings("ignore", category=sparse.SparseEfficiencyWarning)

def calculate_loop_modulus_rho_preprocessed(
    graph: nx.Graph,
    tolerance=1e-4,
    max_iterations=None,
    initial_set_target_size=None,
    overlap_penalty_alpha=0.1,
    osqp_eps_abs=1e-5,
    osqp_eps_rel=1e-5,
    osqp_max_iter=10000,
    cycles_to_add_per_iter=5,
    parallel_seeds=4,
    verbose_iterations=False,
    print_interval=20
    ):
    """
    Calculate the loop modulus of a graph using a parallel adapted Dijkstra approach.
    
    Args:
        graph: NetworkX graph
        tolerance: Convergence tolerance
        max_iterations: Maximum number of iterations
        initial_set_target_size: Target size for initial cycle set
        overlap_penalty_alpha: Penalty factor for cycle overlap during preprocessing
        osqp_eps_abs: Absolute tolerance for OSQP solver
        osqp_eps_rel: Relative tolerance for OSQP solver
        osqp_max_iter: Maximum iterations for OSQP solver
        cycles_to_add_per_iter: Number of violated cycles to add per iteration
        parallel_seeds: Number of parallel Dijkstra searches
        verbose_iterations: Whether to print verbose output for every iteration
        print_interval: Print interval for iterations
        
    Returns:
        final_rho_star_map: Dictionary mapping edges to rho* values
        edge_list: List of edges in the graph
        added_cycles: List of all cycles added as constraints
        final_modulus: Final modulus value
        qp_solves_count: Number of QP solves performed
        total_time: Total computation time
    """
    num_nodes = graph.number_of_nodes() 
    num_edges = graph.number_of_edges()
    run_start_time = time.time()
    
    # --- Input Validation & Edge List Setup ---
    if not isinstance(graph, nx.Graph): 
        raise TypeError("Graph must be NetworkX Graph.")
    if num_edges == 0: 
        return {}, [], [], 0.0, 0, 0.0
        
    try: 
        edge_list = [tuple(sorted(e)) for e in graph.edges()]
        edge_to_idx = {e: i for i, e in enumerate(edge_list)}
        num_edges = len(edge_list)
    except: 
        print("Error creating edge list")
        return {}, [], [], 0.0, 0, 0.0

    # --- Preprocessing Step ---
    print("--- Preprocessing: Finding initial cycles (Triangles Only) ---")
    preprocessing_start_time = time.time()
    try: 
        candidate_cycles_list = get_triangles(graph)
    except Exception as e: 
        print(f"Error get_triangles: {e}")
        candidate_cycles_list = []
    print(f"Found {len(candidate_cycles_list)} triangles.")

    initial_cycles_list = []
    shortest_cycle = None

    if candidate_cycles_list:
        # Score cycles by edge centrality (using degree product)
        degrees = dict(graph.degree())
        edge_centrality = {e: degrees.get(e[0], 0) * degrees.get(e[1], 0) for e in edge_list}
        valid_cands = [
            {'cycle': c, 'score': sum(edge_centrality.get(e, 0) for e in c)} 
            for c in candidate_cycles_list 
            if all(e in edge_centrality for e in c)
        ]
        valid_cands.sort(key=lambda x: x['score'], reverse=True)
        print(f"Scored {len(valid_cands)} valid candidate cycles.")
        
        if valid_cands:
            print("Selecting initial set...")
            initial_set = set()
            covered = set()
            sel_data = []
            
            if initial_set_target_size is None: 
                target = max(10, min(len(valid_cands), num_edges // 4))
            else: 
                target = min(initial_set_target_size, len(valid_cands))
            print(f"Target size: {target}")
            
            # Seed with best cycle if target > 0 and candidates exist
            if target > 0:
                first = valid_cands.pop(0)
                initial_cycles_list.append(first['cycle'])
                first_fset = frozenset(first['cycle'])
                initial_set.add(first_fset)
                sel_data.append({'fset': first_fset, 'cycle': first['cycle']})
                covered.update(first['cycle'])
                
            rem = valid_cands  # Remaining candidates
            
            # Greedy selection loop
            while len(initial_cycles_list) < target and rem:
                best_c_data = None
                best_met = -float('inf')
                best_idx = -1
                
                for idx, c_data in enumerate(rem):
                    c_cyc = c_data['cycle']
                    c_fset = frozenset(c_cyc)
                    
                    if c_fset in initial_set: 
                        continue
                        
                    new_cov = len(set(c_cyc) - covered)
                    
                    if new_cov > 0:
                        # Calculate metric: coverage minus penalty for overlap
                        add_ov = sum(calculate_overlap(c_cyc, s['cycle']) for s in sel_data)
                        pen = overlap_penalty_alpha * add_ov
                        met = new_cov - pen
                    else:
                        met = -float('inf')
                        
                    if met > best_met:
                        best_met = met
                        best_c_data = c_data
                        best_idx = idx
                        
                if best_c_data:
                    best_cyc = best_c_data['cycle']
                    best_fset = frozenset(best_cyc)
                    initial_cycles_list.append(best_cyc)
                    initial_set.add(best_fset)
                    sel_data.append({'fset': best_fset, 'cycle': best_cyc})
                    covered.update(best_cyc)
                    rem.pop(best_idx)
                else:
                    break

    # Fallback: Run BFS if initial_cycles_list is empty
    if not initial_cycles_list:
        print("Preprocessing failed or yielded empty set. Using shortest hop cycle from BFS.")
        try:
            shortest_cycle, _ = get_shortest_hop_count_cycle_bfs(graph)
            if shortest_cycle:  # Check if BFS found a cycle
                initial_cycles_list = [shortest_cycle]
            else:  # BFS also failed
                print("Graph appears to be acyclic.")
                return {}, edge_list, [], 0.0, 0, time.time() - run_start_time
        except Exception as e:
            print(f"Error in fallback BFS: {e}")
            return {}, edge_list, [], 0.0, 0, time.time() - run_start_time

    preprocessing_end_time = time.time()
    print(f"Preprocessing finished ({preprocessing_end_time - preprocessing_start_time:.2f}s). Using {len(initial_cycles_list)} cycles.")
    
    # --- Initial Constraint Matrix & Variables ---
    added_cycles = list(initial_cycles_list)
    N_constraints_list = []
    
    for c_edges in initial_cycles_list:
        c_row = np.zeros(num_edges, dtype=float)
        valid = all(edge in edge_to_idx for edge in c_edges)
        
        if valid:
            for e in c_edges:
                c_row[edge_to_idx[e]] = 1.0
            N_constraints_list.append(c_row)
        else: 
            print(f"Warning: Skipping invalid initial cycle.")

    if not N_constraints_list: 
        print("Error: No valid constraints.")
        return {}, edge_list, [], 0.0, 0, time.time() - run_start_time
        
    rho = np.zeros(num_edges)
    initial_solve_time = 0.0
    qp_solves_count = 0
    current_modulus = 0.0
    
    # OSQP solver setup
    osqp_settings = {
        'eps_abs': osqp_eps_abs,
        'eps_rel': osqp_eps_rel,
        'max_iter': osqp_max_iter,
        'verbose': False,
        'check_termination': 10,
        'warm_start': True
    }
    P_osqp = sparse.csc_matrix(2.0 * sparse.identity(num_edges))
    q_osqp = np.zeros(num_edges)
    solver_osqp = osqp.OSQP()
    warm_start_x = None
    warm_start_y = None

    # --- Initial QP Solve ---
    if N_constraints_list:
        print(f"\n--- Solving Initial QP with OSQP ({len(N_constraints_list)}) ---")
        qp_start_time = time.time()
        n_con_init = len(N_constraints_list)
        
        try:
            N_mat = sparse.csc_matrix(np.array(N_constraints_list))
            I_mat = sparse.identity(num_edges, format='csc')
            A = sparse.vstack([N_mat, I_mat], format='csc')
            l = np.hstack([np.ones(n_con_init), np.zeros(num_edges)])
            u = np.full(n_con_init + num_edges, np.inf)
        except Exception as e:
            print(f"Sparse matrix error: {e}")
            return {}, edge_list, added_cycles, 0.0, 0, time.time() - run_start_time
            
        try:
            solver_osqp.setup(P=P_osqp, q=q_osqp, A=A, l=l, u=u, **osqp_settings)
            results = solver_osqp.solve()
            qp_solves_count = 1
            
            if results.info.status != 'solved':
                print(f"Warning for Initial OSQP: {results.info.status}.")
                
            if results.x is not None:
                rho = results.x
                rho[rho < 0] = 0
                warm_start_x = results.x
                
            if results.y is not None and len(results.y) == len(l):
                warm_start_y = results.y
                
            current_modulus = rho @ rho
            qp_end_time = time.time()
            initial_solve_time = qp_end_time - qp_start_time
            print(f"Initial QP solved in {initial_solve_time:.2f}s. Modulus={current_modulus:.6f}")
        except Exception as e:
            print(f"OSQP Error: {e}")

    # --- Main Iteration Loop ---
    if max_iterations is None:
        max_iterations = num_edges * 2
        
    print(f"\nStarting Main Loop. Tolerance={tolerance}, Max Iter={max_iterations}, Parallel Seeds={parallel_seeds}, Cycles/Iter={cycles_to_add_per_iter}")
    loop_start_time = time.time()
    additional_constraints_added_total = 0
    converged = False
    added_cycles_frozensets = {frozenset(c) for c in added_cycles}
    N_constraints_dynamic = list(N_constraints_list)
    available_nodes = list(graph.nodes())
    
    if not available_nodes:
        print("Error: No nodes in graph.")
        return {}, edge_list, added_cycles, current_modulus, qp_solves_count, time.time() - run_start_time
        
    last_qp_solve_time = 0.0
    stop_iteration = False

    for iteration in range(max_iterations):
        current_rho_map = {edge: rho[edge_to_idx[edge]] for edge in edge_list}
        violation_threshold = 1.0 - tolerance

        # --- Parallel Cycle Finding ---
        start_find_time = time.time()
        violated_cycles_found_map = {}
        seeds_used = 0
        
        num_seeds_to_run = min(parallel_seeds, len(available_nodes))
        if num_seeds_to_run > 0:
            seed_nodes = random.sample(available_nodes, num_seeds_to_run)
            seeds_used = len(seed_nodes)
            from .utils_shortest import shortest_cycle
            result = shortest_cycle(graph, bound=violation_threshold/2, seeds=set(seed_nodes))
            if result:
                edges, cyc_len = result
                if cyc_len < violation_threshold:
                    violated_cycles_found_map[frozenset(edges)] = cyc_len
                            
        violated_cycles_data_sorted = sorted(violated_cycles_found_map.items(), key=lambda item: item[1])
        end_find_time = time.time()
        find_duration = end_find_time - start_find_time

        if not violated_cycles_data_sorted:
            converged = True
            break

        cycles_to_add_data = []
        num_added_this_iter = 0
        
        for cycle_fset, cycle_len in violated_cycles_data_sorted:
            if cycle_fset not in added_cycles_frozensets:
                cycles_to_add_data.append((list(cycle_fset), cycle_len))
                num_added_this_iter += 1
                if num_added_this_iter >= cycles_to_add_per_iter:
                    break

        print_this_iter = verbose_iterations or (iteration % print_interval == 0) or (num_added_this_iter == 0)
        if print_this_iter:
            print(f"\nIter {iteration + 1} ({find_duration:.2f}s): Found {len(violated_cycles_data_sorted)} candidates ({seeds_used} seeds). Adding {num_added_this_iter} new.")
            if cycles_to_add_data:
                print(f"  Best l_rho added: {cycles_to_add_data[0][1]:.6f}")

        if num_added_this_iter == 0:
            converged = True
            break

        # --- Add new constraints ---
        new_constraints_rows = []
        stop_iteration = False  # Reset flag for this iteration
        
        for cycle_edges, cycle_len in cycles_to_add_data:
            added_cycles.append(cycle_edges)
            added_cycles_frozensets.add(frozenset(cycle_edges))
            constraint_row = np.zeros(num_edges)
            valid = True
            
            for edge in cycle_edges:
                idx = edge_to_idx.get(edge)
                if idx is not None:
                    constraint_row[idx] = 1.0
                else:
                    print(f"Map Error edge {edge}")
                    valid = False
                    break
                    
            if not valid:
                added_cycles.pop()
                added_cycles_frozensets.remove(frozenset(cycle_edges))
                stop_iteration = True
                break
                
            new_constraints_rows.append(constraint_row)
            
        if stop_iteration:
            converged = False
            break  # Exit outer loop if error adding constraints

        N_constraints_dynamic.extend(new_constraints_rows)
        additional_constraints_added_total += num_added_this_iter

        # --- Solve QP with OSQP ---
        if print_this_iter:
            print(f"Solving OSQP ({len(N_constraints_dynamic)})")
            
        qp_iter_start_time = time.time()
        num_constraints_current = len(N_constraints_dynamic)
        
        try:
            N_mat = sparse.csc_matrix(np.array(N_constraints_dynamic))
            I_mat = sparse.identity(num_edges, format='csc')
            A_up = sparse.vstack([N_mat, I_mat], format='csc')
            l_up = np.hstack([np.ones(num_constraints_current), np.zeros(num_edges)])
            u_up = np.full(num_constraints_current + num_edges, np.inf)
        except:
            print(f"Sparse error iter {iteration+1}")
            converged = False
            break
            
        try:
            solver_osqp.setup(P=P_osqp, q=q_osqp, A=A_up, l=l_up, u=u_up, **osqp_settings)
            
            if warm_start_x is not None:
                if warm_start_y is not None and len(warm_start_y) == len(l_up):
                    solver_osqp.warm_start(x=warm_start_x, y=warm_start_y)
                else:
                    solver_osqp.warm_start(x=warm_start_x)
                    warm_start_y = None
                    
            results = solver_osqp.solve()
            qp_solves_count += 1
            status_ok = False
            
            if results.info.status == 'solved':
                status_ok = True
            elif results.info.status in ['solved inaccurate', 'max_iter_reached']:
                print(f" OSQP status {results.info.status} (iter {iteration+1}). Accept.")
                status_ok = True
            else:
                print(f"Warning: OSQP status {results.info.status} iter {iteration+1}. Stop.")
                converged = False
                break
                
            if status_ok and results.x is not None:
                rho = results.x
                rho[rho < 0] = 0
                warm_start_x = results.x
                
            if status_ok and results.y is not None:
                warm_start_y = results.y if len(results.y) == len(l_up) else None
                
            current_modulus = rho @ rho
            qp_iter_end_time = time.time()
            last_qp_solve_time = qp_iter_end_time - qp_iter_start_time
            
            if print_this_iter:
                print(f"QP solved ({last_qp_solve_time:.2f}s). Mod={current_modulus:.6f}")
                
        except Exception as e:
            print(f"OSQP Error iter {iteration+1}: {e}")
            converged = False
            break

    else:
        converged = False
        print(f"\nWarning: Reached max iterations ({max_iterations}).")

    loop_end_time = time.time()
    main_loop_time = loop_end_time - loop_start_time
    total_time = initial_solve_time + main_loop_time

    print(f"\nFinished Modulus Calculation main loop.")
    # Final state print
    print(f"--- Final State Summary ---")
    final_mod_val = rho @ rho
    if converged:
        print(f" Converged. Last check found no new violated cycles or length >= tolerance.")
    elif stop_iteration:
        print(f" Stopped due to error during constraint processing.")
    else:
        print(f" Stopped for other reason (e.g., max_iter, QP error).")
        
    print(f" Final Modulus approx: {final_mod_val:.6f}")
    
    if qp_solves_count > 1:  # Avoid printing if only initial solve ran
         print(f" Last QP solve time: {last_qp_solve_time:.2f}s")

    print(f"Constraints added (after initial): {additional_constraints_added_total}, Total constraints: {len(added_cycles)}")
    if not converged:
        print("Warning: Algorithm may not have converged fully.")

    final_rho_star_map = {edge: rho[edge_to_idx[edge]] for edge in edge_list}
    final_modulus = rho @ rho
    print(f"Final Modulus approx Mod2(L'): {final_modulus:.6f}")
    
    return final_rho_star_map, edge_list, added_cycles, final_modulus, qp_solves_count, total_time