pdp_lns 0.1.1

use crate::instance::Instance;
use highs::{Col, HighsModelStatus, RowProblem, Sense};

/// Per-node K-nearest neighbors for sparse arc set.
/// Each node keeps at most K_NEAREST best outgoing and K_NEAREST best incoming arcs
/// (by LP reduced cost). Depot arcs are always included.
const K_NEAREST: usize = 30;

/// Result of the LP-based arc scoring.
pub struct ArcScoringResult {
    /// Boolean sparse arc set: `sparse[i * stride + j]` = true if arc (i,j) ∈ A'^-.
    pub sparse: Vec<bool>,
    /// Per-arc reduced costs from LP relaxation.
    /// `reduced_costs[i * stride + j]` = reduced cost of arc (i,j), or f64::MAX if infeasible.
    pub reduced_costs: Vec<f64>,
}

/// Solve the LP relaxation of a two-index PDPTW formulation to get reduced costs,
/// then build the sparse arc set A'^-.
///
/// Returns sparse arc set and per-arc reduced costs.
/// Depot arcs are always included in the sparse set.
///
/// Based on Goeke (2019), Section 3.1, step iii.
pub fn compute_sparse_arcs(inst: &Instance) -> ArcScoringResult {
    let total = inst.n + 1; // depot + n customers
    let cap = f64::from(inst.capacity);

    // Tighten time windows to reduce Big-M values.
    // Forward: late[p] = min(late[p], late[d] - t_pd - s_p)
    // Backward: early[d] = max(early[d], early[p] + s_p + t_pd)
    // Then transitive tightening over all feasible arcs (a few passes).
    let mut tw_e: Vec<f64> = (0..total).map(|i| inst.early(i)).collect();
    let mut tw_l: Vec<f64> = (0..total).map(|i| inst.late(i)).collect();

    // PD-pair tightening (safe, always valid)
    for &p in &inst.pickups {
        let d = inst.delivery_of(p);
        let t_pd = inst.dist(p, d);
        let s_p = inst.svc(p);
        tw_l[p] = tw_l[p].min(tw_l[d] - t_pd - s_p);
        tw_e[d] = tw_e[d].max(tw_e[p] + s_p + t_pd);
    }

    // Transitive tightening (Desrosiers et al.):
    // Forward: early[j] = max(early[j], min_{predecessors i} (early[i] + s_i + t_ij))
    //   Node j must be preceded by SOME i; the best predecessor gives the tightest valid bound.
    // Backward: late[i] = min(late[i], max_{successors j} (late[j] - s_i - t_ij))
    //   Node i must be followed by SOME j; the best successor gives the latest valid departure.
    for _pass in 0..3 {
        let mut changed = false;
        // Forward pass: tighten early[j] for each customer j
        for j in 1..total {
            let mut best_arrival = f64::MAX;
            for (i, &tw_e_i) in tw_e.iter().enumerate() {
                if i == j || !inst.is_arc_feasible(i, j) {
                    continue;
                }
                let arr = tw_e_i + inst.svc(i) + inst.dist(i, j);
                if arr < best_arrival {
                    best_arrival = arr;
                }
            }
            if best_arrival < f64::MAX
                && best_arrival > tw_e[j] + 1e-9
                && best_arrival <= tw_l[j] + 1e-9
            {
                tw_e[j] = best_arrival;
                changed = true;
            }
        }
        // Backward pass: tighten late[i] for each node i
        for i in 0..total {
            let s_i = inst.svc(i);
            let mut best_depart = f64::NEG_INFINITY;
            for (j, &tw_l_j) in tw_l.iter().enumerate() {
                if i == j || !inst.is_arc_feasible(i, j) {
                    continue;
                }
                let dep = tw_l_j - s_i - inst.dist(i, j);
                if dep > best_depart {
                    best_depart = dep;
                }
            }
            if best_depart > f64::NEG_INFINITY
                && best_depart < tw_l[i] - 1e-9
                && best_depart >= tw_e[i] - 1e-9
            {
                tw_l[i] = best_depart;
                changed = true;
            }
        }
        if !changed {
            break;
        }
    }

    // Log tightening effect
    {
        let mut sum_orig = 0.0;
        let mut sum_tight = 0.0;
        for i in 0..total {
            sum_orig += inst.late(i) - inst.early(i);
            sum_tight += tw_l[i] - tw_e[i];
        }
        let pct = if sum_orig > 0.0 {
            100.0 * (1.0 - sum_tight / sum_orig)
        } else {
            0.0
        };
        eprintln!(
            "  TW tightening: avg window {:.1} -> {:.1} ({:.1}% reduction)",
            sum_orig / total as f64,
            sum_tight / total as f64,
            pct,
        );
    }

    // Pre-filter arcs for LP: keep K_LP nearest outgoing/incoming per node (by distance),
    // plus depot & PD-pair arcs. Reject TW-infeasible arcs and arcs with excessive
    // structural waiting time (reduced cost fixing: wait > horizon/3).
    const K_LP: usize = 50;
    // Structural wait = max(0, early[j] - (late[i] + s_i + t_ij)):
    // even arriving at latest possible time from i, still wait before j's window opens.
    // Arcs with wait > horizon/3 are almost never used in good solutions.
    let horizon = tw_l[0] - tw_e[0]; // planning horizon (depot window)
    let wait_threshold = horizon / 3.0;
    let mut lp_arc_ok = vec![false; total * total];

    // Helper: check arc (i,j) passes TW feasibility + wait filter
    let arc_lp_ok = |i: usize, j: usize| -> bool {
        if j != 0 && tw_e[i] + inst.svc(i) + inst.dist(i, j) > tw_l[j] + 1e-9 {
            return false; // infeasible under tightened windows
        }
        if j != 0 {
            let wait = (tw_e[j] - (tw_l[i] + inst.svc(i) + inst.dist(i, j))).max(0.0);
            if wait > wait_threshold {
                return false;
            }
        }
        true
    };

    // Always include depot arcs and PD-pair arcs
    for i in 0..total {
        for j in 0..total {
            if i == j || !inst.is_arc_feasible(i, j) || !arc_lp_ok(i, j) {
                continue;
            }
            if i == 0 || j == 0 {
                lp_arc_ok[i * total + j] = true;
            }
        }
    }
    for &p in &inst.pickups {
        let d = inst.delivery_of(p);
        if inst.is_arc_feasible(p, d) && arc_lp_ok(p, d) {
            lp_arc_ok[p * total + d] = true;
        }
    }
    // K_LP nearest outgoing/incoming per node (by distance; wait filter already applied)
    {
        let mut neighbors: Vec<(f64, usize)> = Vec::with_capacity(total);
        for i in 0..total {
            neighbors.clear();
            for j in 0..total {
                if i == j || !inst.is_arc_feasible(i, j) || !arc_lp_ok(i, j) {
                    continue;
                }
                neighbors.push((inst.dist(i, j), j));
            }
            neighbors.sort_unstable_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
            for &(_, j) in neighbors.iter().take(K_LP) {
                lp_arc_ok[i * total + j] = true;
            }
        }
        for j in 1..total {
            neighbors.clear();
            for i in 0..total {
                if i == j || !inst.is_arc_feasible(i, j) || !arc_lp_ok(i, j) {
                    continue;
                }
                neighbors.push((inst.dist(i, j), i));
            }
            neighbors.sort_unstable_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
            for &(_, i) in neighbors.iter().take(K_LP) {
                lp_arc_ok[i * total + j] = true;
            }
        }
    }

    let mut pb = RowProblem::default();

    // Step 1: Create arc variables x_ij for LP-selected arcs
    let mut arc_cols: Vec<Col> = Vec::with_capacity(total * K_LP);
    let mut arc_map: Vec<(u16, u16)> = Vec::with_capacity(total * K_LP);
    let mut arc_to_col_idx: Vec<u32> = vec![u32::MAX; total * total];

    for i in 0..total {
        for j in 0..total {
            if !lp_arc_ok[i * total + j] {
                continue;
            }
            let cost = inst.dist(i, j);
            let col = pb.add_column(cost, 0.0..=1.0);
            arc_to_col_idx[i * total + j] = arc_cols.len() as u32;
            arc_cols.push(col);
            arc_map.push((i as u16, j as u16));
        }
    }

    let num_arc_vars = arc_cols.len();

    // Step 2: Create time variables T_i for each node
    let mut time_cols: Vec<Col> = Vec::with_capacity(total);
    for i in 0..total {
        let col = pb.add_column(0.0, tw_e[i]..=tw_l[i]);
        time_cols.push(col);
    }

    // Step 2b: Create load variables Q_i for each node
    let mut load_cols: Vec<Col> = Vec::with_capacity(total);
    for i in 0..total {
        let q_i = f64::from(inst.demand[i]);
        let (lb, ub) = if i == 0 {
            (0.0, 0.0) // depot: vehicles depart empty
        } else if q_i > 0.0 {
            (q_i, cap) // pickup: [demand, capacity]
        } else {
            (0.0, cap + q_i) // delivery: [0, capacity - |demand|]
        };
        let col = pb.add_column(0.0, lb..=ub);
        load_cols.push(col);
    }

    // Step 3: Flow conservation constraints

    // Outflow: Σ_j x_ij = 1 for all customers i ∈ {1..n}
    let mut coeffs: Vec<(Col, f64)> = Vec::new();
    for i in 1..total {
        coeffs.clear();
        for j in 0..total {
            if i == j {
                continue;
            }
            let idx = arc_to_col_idx[i * total + j];
            if idx != u32::MAX {
                coeffs.push((arc_cols[idx as usize], 1.0));
            }
        }
        if !coeffs.is_empty() {
            pb.add_row(1.0..=1.0, &coeffs);
        }
    }

    // Inflow: Σ_i x_ij = 1 for all customers j ∈ {1..n}
    for j in 1..total {
        coeffs.clear();
        for i in 0..total {
            if i == j {
                continue;
            }
            let idx = arc_to_col_idx[i * total + j];
            if idx != u32::MAX {
                coeffs.push((arc_cols[idx as usize], 1.0));
            }
        }
        if !coeffs.is_empty() {
            pb.add_row(1.0..=1.0, &coeffs);
        }
    }

    // Depot outflow: Σ_j x_0j <= K
    {
        coeffs.clear();
        for &idx in &arc_to_col_idx[1..total] {
            if idx != u32::MAX {
                coeffs.push((arc_cols[idx as usize], 1.0));
            }
        }
        if !coeffs.is_empty() {
            let k = inst.num_vehicles as f64;
            pb.add_row(0.0..=k, &coeffs);
        }
    }

    // Depot inflow: Σ_i x_i0 <= K
    {
        coeffs.clear();
        for i in 1..total {
            let idx = arc_to_col_idx[i * total];
            if idx != u32::MAX {
                coeffs.push((arc_cols[idx as usize], 1.0));
            }
        }
        if !coeffs.is_empty() {
            let k = inst.num_vehicles as f64;
            pb.add_row(0.0..=k, &coeffs);
        }
    }

    // Step 4: Time propagation constraints
    // T_i + s_i + t_ij - M_ij*(1 - x_ij) <= T_j  for all (i,j) ∈ A', j ≠ 0
    // Rearranged: T_i - T_j + M_ij * x_ij <= M_ij - s_i - t_ij
    // where M_ij = max(l_i + s_i + t_ij - e_j, 0)
    for (col_idx, &(i16, j16)) in arc_map.iter().enumerate() {
        let i = i16 as usize;
        let j = j16 as usize;
        if j == 0 {
            continue; // Skip arcs returning to depot
        }
        let t_ij = inst.dist(i, j);
        let s_i = inst.svc(i);
        let m_ij = f64::max(tw_l[i] + s_i + t_ij - tw_e[j], 0.0);

        if m_ij < 1e-10 {
            continue;
        }

        let rhs = m_ij - s_i - t_ij;
        pb.add_row(
            f64::NEG_INFINITY..=rhs,
            [
                (time_cols[i], 1.0),
                (time_cols[j], -1.0),
                (arc_cols[col_idx], m_ij),
            ],
        );
    }

    // Step 4b: Load propagation constraints
    // Q_i - Q_j + M_ij^Q * x_ij <= M_ij^Q - q_j  for all (i,j) ∈ A', j ≠ 0
    // where M_ij^Q = max(ub(Q_i) - lb(Q_j) + q_j, 0)
    for (col_idx, &(i16, j16)) in arc_map.iter().enumerate() {
        let i = i16 as usize;
        let j = j16 as usize;
        if j == 0 {
            continue;
        }
        let q_j = f64::from(inst.demand[j]);
        // ub(Q_i): 0 for depot, cap for pickup, cap + q_i for delivery
        let q_i = f64::from(inst.demand[i]);
        let ub_qi = if i == 0 {
            0.0
        } else if q_i > 0.0 {
            cap
        } else {
            cap + q_i
        };
        // lb(Q_j): q_j for pickup (positive), 0 for delivery
        let lb_qj = if q_j > 0.0 { q_j } else { 0.0 };
        let m_ij_q = f64::max(ub_qi - lb_qj + q_j, 0.0);

        if m_ij_q < 1e-10 {
            continue;
        }

        let rhs = m_ij_q - q_j;
        pb.add_row(
            f64::NEG_INFINITY..=rhs,
            [
                (load_cols[i], 1.0),
                (load_cols[j], -1.0),
                (arc_cols[col_idx], m_ij_q),
            ],
        );
    }

    // Step 5: Precedence constraints: T_p + s_p + t_{p,d} <= T_d for each PD pair
    for &p in &inst.pickups {
        let d = inst.delivery_of(p);
        let t_pd = inst.dist(p, d);
        let s_p = inst.svc(p);

        // T_p - T_d <= -(s_p + t_pd)
        pb.add_row(
            f64::NEG_INFINITY..=-(s_p + t_pd),
            [(time_cols[p], 1.0), (time_cols[d], -1.0)],
        );
    }

    // Step 5b: PD pair load coupling: Q_p - Q_d = demand[p]
    for &p in &inst.pickups {
        let d = inst.delivery_of(p);
        let q_p = f64::from(inst.demand[p]);
        pb.add_row(q_p..=q_p, [(load_cols[p], 1.0), (load_cols[d], -1.0)]);
    }

    // Step 6: Solve LP
    eprintln!(
        "  LP: {} variables ({} arcs + {} time + {} load), solving...",
        num_arc_vars + total + total,
        num_arc_vars,
        total,
        total,
    );

    let mut model = pb.optimise(Sense::Minimise);
    model.set_option("output_flag", false);
    model.set_option("threads", 4);
    model.set_option("time_limit", 10.0); // hard cap; fallback to distance-based if exceeded
    let t0 = std::time::Instant::now();
    let solved = model.solve();
    let lp_ms = t0.elapsed().as_millis();
    let status = solved.status();

    if status != HighsModelStatus::Optimal {
        eprintln!(
            "  LP status: {status:?} (took {lp_ms} ms) — falling back to distance-based scoring"
        );
        return build_distance_based_result(inst);
    }

    let solution = solved.get_solution();
    let dual_cols = solution.dual_columns();
    let obj = solved.objective_value();
    eprintln!("  LP solved: obj={obj:.2}, time={lp_ms} ms");

    // Build per-arc reduced cost matrix
    let mut reduced_costs = vec![f64::MAX; total * total];
    for (var_idx, &(i16, j16)) in arc_map.iter().enumerate() {
        let i = i16 as usize;
        let j = j16 as usize;
        reduced_costs[i * total + j] = dual_cols[var_idx];
    }

    // Step 8: Build sparse arc set A'^- using per-node K-nearest neighbors.
    // For each node, keep the K_NEAREST outgoing arcs (i,*) and K_NEAREST incoming
    // arcs (*,j) with lowest reduced cost. Depot arcs are always included.
    let total_arcs = total * (total - 1);

    // Group arcs by source (outgoing) and target (incoming)
    let mut outgoing: Vec<Vec<(f64, usize)>> = vec![Vec::new(); total]; // outgoing[i] = [(rc, j), ...]
    let mut incoming: Vec<Vec<(f64, usize)>> = vec![Vec::new(); total]; // incoming[j] = [(rc, i), ...]
    for (var_idx, &(i16, j16)) in arc_map.iter().enumerate() {
        let i = i16 as usize;
        let j = j16 as usize;
        let rc = dual_cols[var_idx];
        outgoing[i].push((rc, j));
        incoming[j].push((rc, i));
    }

    let mut sparse = vec![false; total * total];
    let mut count = 0;

    // Always include depot arcs
    for j in 1..total {
        if inst.is_arc_feasible(0, j) {
            sparse[j] = true;
            count += 1;
        }
        if inst.is_arc_feasible(j, 0) {
            sparse[j * total] = true;
            count += 1;
        }
    }

    // For each node, keep K_NEAREST best outgoing arcs
    for i in 1..total {
        let arcs = &mut outgoing[i];
        if arcs.len() > K_NEAREST {
            arcs.select_nth_unstable_by(K_NEAREST - 1, |a, b| a.0.partial_cmp(&b.0).unwrap());
        }
        for &(_, j) in arcs.iter().take(K_NEAREST) {
            if !sparse[i * total + j] {
                sparse[i * total + j] = true;
                count += 1;
            }
        }
    }

    // For each node, keep K_NEAREST best incoming arcs
    for j in 1..total {
        let arcs = &mut incoming[j];
        if arcs.len() > K_NEAREST {
            arcs.select_nth_unstable_by(K_NEAREST - 1, |a, b| a.0.partial_cmp(&b.0).unwrap());
        }
        for &(_, i) in arcs.iter().take(K_NEAREST) {
            if !sparse[i * total + j] {
                sparse[i * total + j] = true;
                count += 1;
            }
        }
    }

    eprintln!(
        "  Sparse arc set: {count}/{num_arc_vars} arcs ({:.1}% of feasible, {:.1}% of total, K={})",
        f64::from(count) / num_arc_vars as f64 * 100.0,
        f64::from(count) / total_arcs as f64 * 100.0,
        K_NEAREST,
    );

    ArcScoringResult {
        sparse,
        reduced_costs,
    }
}

/// Fallback: build sparse arc set using distance-based K-nearest neighbors.
/// Uses distances as pseudo reduced costs.
fn build_distance_based_result(inst: &Instance) -> ArcScoringResult {
    let total = inst.n + 1;

    // Group feasible arcs by source and target
    let mut outgoing: Vec<Vec<(f64, usize)>> = vec![Vec::new(); total];
    let mut incoming: Vec<Vec<(f64, usize)>> = vec![Vec::new(); total];
    for (i, out_arcs) in outgoing.iter_mut().enumerate() {
        for (j, in_arcs) in incoming.iter_mut().enumerate() {
            if i == j || !inst.is_arc_feasible(i, j) {
                continue;
            }
            let d = inst.dist(i, j);
            out_arcs.push((d, j));
            in_arcs.push((d, i));
        }
    }

    let mut sparse = vec![false; total * total];

    // Always include depot arcs
    for j in 1..total {
        if inst.is_arc_feasible(0, j) {
            sparse[j] = true;
        }
        if inst.is_arc_feasible(j, 0) {
            sparse[j * total] = true;
        }
    }

    for i in 1..total {
        let arcs = &mut outgoing[i];
        if arcs.len() > K_NEAREST {
            arcs.select_nth_unstable_by(K_NEAREST - 1, |a, b| a.0.partial_cmp(&b.0).unwrap());
        }
        for &(_, j) in arcs.iter().take(K_NEAREST) {
            if !sparse[i * total + j] {
                sparse[i * total + j] = true;
            }
        }
    }

    for j in 1..total {
        let arcs = &mut incoming[j];
        if arcs.len() > K_NEAREST {
            arcs.select_nth_unstable_by(K_NEAREST - 1, |a, b| a.0.partial_cmp(&b.0).unwrap());
        }
        for &(_, i) in arcs.iter().take(K_NEAREST) {
            if !sparse[i * total + j] {
                sparse[i * total + j] = true;
            }
        }
    }

    // Use distances as pseudo reduced costs for fallback
    let mut reduced_costs = vec![f64::MAX; total * total];
    for i in 0..total {
        for j in 0..total {
            if i != j && inst.is_arc_feasible(i, j) {
                reduced_costs[i * total + j] = inst.dist(i, j);
            }
        }
    }

    ArcScoringResult {
        sparse,
        reduced_costs,
    }
}