pdp_lns 0.1.1

use crate::instance::Instance;
use crate::solution::{Solution, route_distance};
use highs::{Col, HighsModelStatus, RowProblem, Sense};
use std::collections::{HashMap, HashSet};
use std::hash::{Hash, Hasher};

/// Maximum routes in the pool.
const MAX_POOL_SIZE: usize = 100_000;

/// Vehicle cost in objective (hierarchical: minimize vehicles >> minimize distance).
/// Must dominate any realistic route distance.
const VEHICLE_COST: f64 = 1e6;

/// MIP time limit in seconds.
const MIP_TIME_LIMIT: f64 = 5.0;

/// Pool of feasible routes found during the search.
/// Used to solve a Set Covering MIP (Goeke 2019, Section 3.6, equations 32–36).
pub struct RoutePool {
    /// Stored routes with depot endpoints [0, ..., 0].
    routes: Vec<Vec<usize>>,
    /// Distance of each route.
    distances: Vec<f64>,
    /// Sorted pickup signature per route (used for fast dominance pruning).
    pickup_signatures: Vec<Vec<usize>>,
    /// Hash of sorted pickup signature per route (for fast grouping).
    pickup_signature_hashes: Vec<u64>,
    /// Deduplication: set of route interiors (customers without depots).
    seen: HashSet<Vec<usize>>,
    /// For each pickup node p, list of pool indices of routes containing p.
    pickup_routes: Vec<Vec<usize>>,
    /// Number of customer nodes (inst.n).
    n: usize,
    /// List of pickup node IDs.
    pickups: Vec<usize>,
    /// True if new routes were added since the last SC solve.
    dirty: bool,
    /// Number of newly added routes since the last SC solve.
    new_since_solve: usize,
}

impl RoutePool {
    pub fn new(inst: &Instance) -> Self {
        RoutePool {
            routes: Vec::new(),
            distances: Vec::new(),
            pickup_signatures: Vec::new(),
            pickup_signature_hashes: Vec::new(),
            seen: HashSet::new(),
            pickup_routes: vec![Vec::new(); inst.n + 1],
            n: inst.n,
            pickups: inst.pickups.clone(),
            dirty: false,
            new_since_solve: 0,
        }
    }

    /// Add all routes from a solution to the pool.
    pub fn add_solution(&mut self, inst: &Instance, sol: &Solution) {
        for route in &sol.routes {
            self.add_route(inst, route);
        }
    }

    fn add_route(&mut self, inst: &Instance, route: &[usize]) {
        if route.len() <= 2 || self.routes.len() >= MAX_POOL_SIZE {
            return;
        }

        let interior = &route[1..route.len() - 1];
        if self.seen.contains(interior) {
            return; // duplicate
        }
        self.seen.insert(interior.to_vec());

        let dist = route_distance(inst, route);
        let idx = self.routes.len();

        // Index by pickup nodes only (delivery is always in the same route)
        let mut pickup_signature: Vec<usize> = Vec::new();
        for &v in &route[1..route.len() - 1] {
            if inst.is_pickup(v) {
                self.pickup_routes[v].push(idx);
                pickup_signature.push(v);
            }
        }
        pickup_signature.sort_unstable();
        let mut hasher = std::collections::hash_map::DefaultHasher::new();
        pickup_signature.hash(&mut hasher);
        let pickup_signature_hash = hasher.finish();

        self.routes.push(route.to_vec());
        self.distances.push(dist);
        self.pickup_signatures.push(pickup_signature);
        self.pickup_signature_hashes.push(pickup_signature_hash);
        self.dirty = true;
        self.new_since_solve += 1;
    }

    pub fn len(&self) -> usize {
        self.routes.len()
    }

    pub fn is_empty(&self) -> bool {
        self.routes.is_empty()
    }

    /// Returns true if new routes were added since the last SC solve.
    pub fn has_new_routes(&self) -> bool {
        self.dirty
    }

    /// Number of routes added since the last SC solve.
    pub fn new_routes_since_solve(&self) -> usize {
        self.new_since_solve
    }

    /// Solve the set partitioning MIP with the default time limit.
    /// `best` is the current best solution, used as objective cutoff.
    pub fn solve(&mut self, inst: &Instance, best: &Solution) -> Option<Solution> {
        self.dirty = false;
        self.new_since_solve = 0;
        let cutoff = Self::sc_cost(best, inst);
        self.solve_with_limit(inst, MIP_TIME_LIMIT, cutoff)
    }

    /// Solve the set partitioning MIP with an extended time limit (for final solve).
    /// `best` is the current best solution, used as objective cutoff.
    pub fn solve_final(&mut self, inst: &Instance, best: &Solution) -> Option<Solution> {
        self.dirty = false;
        self.new_since_solve = 0;
        let cutoff = Self::sc_cost(best, inst);
        self.solve_with_limit(inst, 10.0, cutoff)
    }

    /// Compute which routes are non-dominated ("active") for the MIP.
    /// Two routes with the same pickup set → keep only the shortest.
    fn compute_active_routes(&self) -> Vec<bool> {
        let num_routes = self.routes.len();
        let mut active = vec![true; num_routes];

        // Single pass: keep only the shortest route for each signature.
        // Group first by precomputed signature hash; on hash collisions compare actual signatures.
        let mut best_by_hash: HashMap<u64, Vec<usize>> = HashMap::new();
        for r in 0..num_routes {
            let sig_hash = self.pickup_signature_hashes[r];
            let sig = &self.pickup_signatures[r];
            let bucket = best_by_hash.entry(sig_hash).or_default();

            let mut matched_pos: Option<usize> = None;
            for (pos, &best_r) in bucket.iter().enumerate() {
                if self.pickup_signatures[best_r] == *sig {
                    matched_pos = Some(pos);
                    break;
                }
            }

            if let Some(pos) = matched_pos {
                let best_r = bucket[pos];
                if self.distances[r] < self.distances[best_r] {
                    active[best_r] = false;
                    bucket[pos] = r;
                } else {
                    active[r] = false;
                }
            } else {
                bucket.push(r);
            }
        }

        active
    }

    /// Solve the LP relaxation of the Set Cover problem and return per-node
    /// shadow prices (dual values).  Routes use distance-only cost (no VEHICLE_COST)
    /// so that duals give meaningful differentiation between pickups.
    ///
    /// Returns a vector of length `inst.n + 1` indexed by node ID.
    /// `result[p]` = shadow price for pickup `p`; 0.0 for non-pickup nodes.
    #[allow(clippy::cast_precision_loss)]
    pub fn solve_lp_duals(&self, inst: &Instance) -> Option<Vec<f64>> {
        if self.routes.is_empty() {
            return None;
        }
        // Check that all pickups are covered by at least one route in the pool
        for &p in &self.pickups {
            if self.pickup_routes[p].is_empty() {
                return None;
            }
        }

        let num_routes = self.routes.len();
        let active = self.compute_active_routes();

        let mut pb = RowProblem::default();

        // Map pool index → LP column (only for active routes)
        let mut pool_to_col: Vec<Option<Col>> = vec![None; num_routes];
        for r in 0..num_routes {
            if !active[r] {
                continue;
            }
            // Distance-only cost: no VEHICLE_COST so duals differentiate meaningfully
            let cost = self.distances[r];
            let col = pb.add_column(cost, 0.0..=1.0); // continuous variable
            pool_to_col[r] = Some(col);
        }

        // Covering constraints: for each pickup p, Σ_{r: p∈r} y_r >= 1
        let mut coeffs: Vec<(Col, f64)> = Vec::new();
        let mut constraint_to_pickup: Vec<usize> = Vec::new();
        for &p in &self.pickups {
            coeffs.clear();
            coeffs.extend(
                self.pickup_routes[p]
                    .iter()
                    .filter_map(|&r_idx| pool_to_col[r_idx].map(|col| (col, 1.0))),
            );
            if coeffs.is_empty() {
                continue; // should not happen given the check above
            }
            pb.add_row(1.0..=f64::INFINITY, &coeffs);
            constraint_to_pickup.push(p);
        }

        let mut model = pb.optimise(Sense::Minimise);
        model.set_option("time_limit", 2.0);
        model.set_option("threads", 4);
        model.set_option("output_flag", false);

        let Ok(solved) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| model.solve()))
        else {
            return None;
        };

        if solved.status() != HighsModelStatus::Optimal {
            return None;
        }

        let solution = solved.get_solution();
        let dual_rows = solution.dual_rows();

        // Map constraint duals to per-node vector
        let mut result = vec![0.0_f64; inst.n + 1];
        for (constraint_idx, &pickup) in constraint_to_pickup.iter().enumerate() {
            if constraint_idx < dual_rows.len() {
                // LP duals for >= constraints are non-negative in a minimization;
                // take absolute value for robustness against solver sign conventions.
                result[pickup] = dual_rows[constraint_idx].abs();
            }
        }

        Some(result)
    }

    /// Compute the SC objective value for a solution.
    #[allow(clippy::cast_precision_loss)]
    fn sc_cost(sol: &Solution, inst: &Instance) -> f64 {
        sol.num_vehicles() as f64 * VEHICLE_COST + sol.total_distance(inst)
    }

    /// Solve the set partitioning MIP and return a solution if one is found.
    ///
    /// Formulation (Goeke 2019, equations 32–36, with = instead of >=):
    ///   min  Σ (VEHICLE_COST + c_r) · y_r           (hierarchical objective)
    ///   s.t. Σ_{r: p∈r} y_r = 1   ∀ pickup p       (each PD pair covered exactly once)
    ///        y_r ∈ {0, 1}                            (binary selection)
    ///
    /// `cutoff` is an upper bound — only solutions strictly better are accepted.
    fn solve_with_limit(&self, inst: &Instance, time_limit: f64, cutoff: f64) -> Option<Solution> {
        if self.routes.is_empty() {
            return None;
        }

        // Check that all pickups are covered by at least one route in the pool
        for &p in &self.pickups {
            if self.pickup_routes[p].is_empty() {
                return None;
            }
        }

        let num_routes = self.routes.len();

        // Prune dominated routes: for routes with identical pickup sets,
        // keep only the one with minimum distance.
        let active = self.compute_active_routes();
        let active_count = active.iter().filter(|&&a| a).count();

        let mut pb = RowProblem::default();

        // Map pool index → MIP column (only for active routes)
        let mut pool_to_col: Vec<Option<Col>> = vec![None; num_routes];
        let mut col_to_pool: Vec<usize> = Vec::with_capacity(active_count);
        for r in 0..num_routes {
            if !active[r] {
                continue;
            }
            let cost = VEHICLE_COST + self.distances[r];
            let col = pb.add_integer_column(cost, 0..=1);
            pool_to_col[r] = Some(col);
            col_to_pool.push(r);
        }

        // Partitioning constraints: for each pickup p, Σ_{r: p∈r} y_r = 1
        let mut coeffs: Vec<(Col, f64)> = Vec::new();
        for &p in &self.pickups {
            coeffs.clear();
            coeffs.extend(
                self.pickup_routes[p]
                    .iter()
                    .filter_map(|&r_idx| pool_to_col[r_idx].map(|col| (col, 1.0))),
            );
            pb.add_row(1.0..=1.0, &coeffs);
        }

        // Solve MIP with time limit and objective cutoff
        let mut model = pb.optimise(Sense::Minimise);
        model.set_option("time_limit", time_limit);
        model.set_option("objective_target", cutoff - 1e-6);
        model.set_option("output_flag", false); // suppress HiGHS logs
        model.set_option("threads", 4);
        model.set_option("presolve", "on");
        model.set_option("mip_rel_gap", 0.001); // 0.1% optimality gap is good enough

        let Ok(solved) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| model.solve()))
        else {
            eprintln!("  SC: HiGHS panicked, skipping");
            return None;
        };
        let status = solved.status();

        // Accept optimal or time-limit-with-incumbent
        if status != HighsModelStatus::Optimal && status != HighsModelStatus::ReachedTimeLimit {
            return None;
        }

        let solution = solved.get_solution();
        let primal = solution.columns();

        // When ReachedTimeLimit, HiGHS may return LP relaxation values rather than
        // an integer incumbent. Verify all selected variables are binary (0 or 1).
        if status == HighsModelStatus::ReachedTimeLimit {
            let all_binary = primal.iter().all(|&v| !(0.01..=0.99).contains(&v));
            if !all_binary {
                return None; // fractional LP solution, not a valid integer incumbent
            }
        }

        // Extract selected pool indices (map MIP columns back to pool indices)
        let mut selected_pool_indices: Vec<usize> = Vec::new();
        for (col_idx, &pool_idx) in col_to_pool.iter().enumerate() {
            if primal[col_idx] > 0.5 {
                selected_pool_indices.push(pool_idx);
            }
        }

        if selected_pool_indices.is_empty() {
            return None;
        }

        // Verify all pickups are covered exactly once (partitioning).
        let mut covered_pickups = vec![false; self.n + 1];
        for &pool_idx in &selected_pool_indices {
            for &p in &self.pickup_signatures[pool_idx] {
                if covered_pickups[p] {
                    return None; // duplicate coverage — not a valid partition
                }
                covered_pickups[p] = true;
            }
        }
        for &p in &self.pickups {
            if !covered_pickups[p] {
                return None; // incomplete solution
            }
        }

        let selected_routes: Vec<Vec<usize>> = selected_pool_indices
            .iter()
            .map(|&pool_idx| self.routes[pool_idx].clone())
            .collect();

        let sc_sol = Solution {
            routes: selected_routes,
            unassigned: Vec::new(),
        };

        let nv = sc_sol.num_vehicles();
        let dist = sc_sol.total_distance(inst);
        eprintln!(
            "  SC: {active_count}/{num_routes} routes → {nv} vehicles, dist={dist:.2} ({status:?})",
        );

        Some(sc_sol)
    }
}