holos-tda 0.9.0

use std::collections::BTreeMap;

use crate::coverage_synthesis::evaluate_coverage_plan_states_prevalidated;
use crate::monotone_search::{SearchLimits, SearchResult, SearchStatus, minimize_antitone};
use crate::{
    CoverageAction, CoverageComponent, CoverageSpecification, CoverageSynthesisLimits, Error,
    Result,
};

use super::{
    CoverageComponentFrontier, CoverageComposition, CoverageCompositionStatus,
    CoverageFrontierEntry,
};

/// Build exact component frontiers and compose a minimum-cost global plan.
///
/// Each local search runs at every activation count up to `max_activations`.
/// The dynamic program then picks one frontier entry per component under that
/// limit. A producer work limit returns
/// [`CoverageCompositionStatus::SearchIncomplete`].
pub fn compose_coverage_frontiers(
    specification: &CoverageSpecification,
    actions: &[CoverageAction],
    max_activations: usize,
    limits: CoverageSynthesisLimits,
) -> Result<CoverageComposition> {
    if limits.max_oracle_calls == 0 || limits.max_search_nodes == 0 {
        return Err(Error::InvalidInput(
            "coverage composition search limits must be positive".into(),
        ));
    }
    let components = specification.components(actions)?;
    let mut work = CompositionWork::default();
    let mut frontiers = Vec::with_capacity(components.len());
    for component in components {
        let Some(frontier) = build_frontier(
            specification,
            actions,
            component,
            max_activations.min(actions.len()),
            limits,
            &mut work,
        )?
        else {
            return Ok(CoverageComposition {
                status: CoverageCompositionStatus::SearchIncomplete,
                frontiers,
                selected: Vec::new(),
                cost: None,
                oracle_calls: work.oracle_calls,
                search_nodes: work.search_nodes,
                cache_hits: work.cache_hits,
            });
        };
        frontiers.push(frontier);
    }
    let Some((cost, selected)) = compose(&frontiers, max_activations.min(actions.len()))? else {
        return Ok(CoverageComposition {
            status: CoverageCompositionStatus::Infeasible,
            frontiers,
            selected: Vec::new(),
            cost: None,
            oracle_calls: work.oracle_calls,
            search_nodes: work.search_nodes,
            cache_hits: work.cache_hits,
        });
    };
    Ok(CoverageComposition {
        status: CoverageCompositionStatus::Optimal,
        frontiers,
        selected,
        cost: Some(cost),
        oracle_calls: work.oracle_calls,
        search_nodes: work.search_nodes,
        cache_hits: work.cache_hits,
    })
}

#[derive(Debug, Default)]
struct CompositionWork {
    oracle_calls: usize,
    search_nodes: usize,
    cache_hits: usize,
}

fn build_frontier(
    specification: &CoverageSpecification,
    actions: &[CoverageAction],
    component: CoverageComponent,
    global_limit: usize,
    limits: CoverageSynthesisLimits,
    work: &mut CompositionWork,
) -> Result<Option<CoverageComponentFrontier>> {
    let local_costs = component
        .actions()
        .iter()
        .map(|action| actions[*action].cost)
        .collect::<Vec<_>>();
    let largest_limit = global_limit.min(local_costs.len());
    let mut entries = Vec::new();
    let Some(maximum) = search_component(
        specification,
        actions,
        &component,
        &local_costs,
        largest_limit,
        limits,
        work,
    )?
    else {
        return Ok(None);
    };
    let smallest_relevant_limit = match accept_maximum(&mut entries, &component, maximum)? {
        MaximumResult::Feasible(limit) => limit,
        MaximumResult::Infeasible => {
            return Ok(Some(CoverageComponentFrontier { component, entries }));
        }
        MaximumResult::Incomplete => return Ok(None),
    };
    for local_limit in 0..smallest_relevant_limit {
        let Some(result) = search_component(
            specification,
            actions,
            &component,
            &local_costs,
            local_limit,
            limits,
            work,
        )?
        else {
            return Ok(None);
        };
        if !accept_frontier_result(&mut entries, &component, result)? {
            return Ok(None);
        }
    }
    Ok(Some(CoverageComponentFrontier { component, entries }))
}

enum MaximumResult {
    Feasible(usize),
    Infeasible,
    Incomplete,
}

fn accept_maximum(
    entries: &mut Vec<CoverageFrontierEntry>,
    component: &CoverageComponent,
    result: SearchResult,
) -> Result<MaximumResult> {
    match result.status {
        SearchStatus::Optimal => Ok(MaximumResult::Feasible(insert_result(
            entries, component, result,
        )?)),
        SearchStatus::Infeasible => Ok(MaximumResult::Infeasible),
        SearchStatus::Incomplete => Ok(MaximumResult::Incomplete),
    }
}

fn accept_frontier_result(
    entries: &mut Vec<CoverageFrontierEntry>,
    component: &CoverageComponent,
    result: SearchResult,
) -> Result<bool> {
    match result.status {
        SearchStatus::Optimal => {
            insert_result(entries, component, result)?;
            Ok(true)
        }
        SearchStatus::Infeasible => Ok(true),
        SearchStatus::Incomplete => Ok(false),
    }
}

#[allow(clippy::too_many_arguments)]
fn search_component(
    specification: &CoverageSpecification,
    actions: &[CoverageAction],
    component: &CoverageComponent,
    local_costs: &[u64],
    local_limit: usize,
    limits: CoverageSynthesisLimits,
    work: &mut CompositionWork,
) -> Result<Option<SearchResult>> {
    let remaining_calls = limits.max_oracle_calls.saturating_sub(work.oracle_calls);
    let remaining_nodes = limits.max_search_nodes.saturating_sub(work.search_nodes);
    if remaining_calls == 0 || remaining_nodes == 0 {
        return Ok(None);
    }
    let result = minimize_antitone(
        local_costs,
        local_limit,
        SearchLimits {
            oracle_calls: remaining_calls,
            search_nodes: remaining_nodes,
        },
        |local_selected| {
            let selected = local_selected
                .iter()
                .map(|local| component.actions()[*local])
                .collect::<Vec<_>>();
            Ok(!evaluate_coverage_plan_states_prevalidated(
                specification,
                actions,
                &selected,
                component.states(),
                limits.coverage,
            )?
            .criterion_holds)
        },
    )?;
    work.oracle_calls = checked_sum(
        work.oracle_calls,
        result.oracle_calls,
        "coverage composition oracle call count overflows",
    )?;
    work.search_nodes = checked_sum(
        work.search_nodes,
        result.search_nodes,
        "coverage composition search node count overflows",
    )?;
    work.cache_hits = checked_sum(
        work.cache_hits,
        result.cache_hits,
        "coverage composition cache hit count overflows",
    )?;
    Ok(Some(result))
}

fn insert_result(
    entries: &mut Vec<CoverageFrontierEntry>,
    component: &CoverageComponent,
    result: SearchResult,
) -> Result<usize> {
    let selected = result
        .selected
        .iter()
        .map(|local| component.actions()[*local])
        .collect::<Vec<_>>();
    let activations = selected.len();
    let cost = result
        .upper_bound
        .ok_or_else(|| Error::InvalidInput("optimal component result has no cost".into()))?;
    insert_nondominated(
        entries,
        CoverageFrontierEntry {
            activations,
            cost,
            selected,
        },
    );
    Ok(activations)
}

fn insert_nondominated(entries: &mut Vec<CoverageFrontierEntry>, candidate: CoverageFrontierEntry) {
    if entries.iter().any(|entry| dominates(entry, &candidate)) {
        return;
    }
    entries.retain(|entry| !dominates(&candidate, entry));
    entries.push(candidate);
    entries.sort_by(|left, right| {
        (left.activations, left.cost, &left.selected).cmp(&(
            right.activations,
            right.cost,
            &right.selected,
        ))
    });
}

fn dominates(left: &CoverageFrontierEntry, right: &CoverageFrontierEntry) -> bool {
    left.activations <= right.activations
        && left.cost <= right.cost
        && (left.activations < right.activations
            || left.cost < right.cost
            || left.selected <= right.selected)
}

fn compose(
    frontiers: &[CoverageComponentFrontier],
    max_activations: usize,
) -> Result<Option<(u64, Vec<usize>)>> {
    let mut partial = BTreeMap::from([(0usize, (0u64, Vec::new()))]);
    for frontier in frontiers {
        let mut next = BTreeMap::<usize, (u64, Vec<usize>)>::new();
        for (&used, (cost, selected)) in &partial {
            for entry in &frontier.entries {
                let activations = used.checked_add(entry.activations).ok_or_else(|| {
                    Error::InvalidInput("coverage composition activation count overflows".into())
                })?;
                if activations > max_activations {
                    continue;
                }
                let combined_cost = cost.checked_add(entry.cost).ok_or_else(|| {
                    Error::InvalidInput("coverage composition cost overflows".into())
                })?;
                let combined_selected = merge(selected, &entry.selected);
                let replace = next.get(&activations).is_none_or(|(best_cost, best)| {
                    (combined_cost, &combined_selected) < (*best_cost, best)
                });
                if replace {
                    next.insert(activations, (combined_cost, combined_selected));
                }
            }
        }
        partial = next;
        if partial.is_empty() {
            return Ok(None);
        }
    }
    Ok(partial
        .into_values()
        .min_by(|left, right| (left.0, &left.1).cmp(&(right.0, &right.1))))
}

fn merge(left: &[usize], right: &[usize]) -> Vec<usize> {
    let mut output = Vec::with_capacity(left.len() + right.len());
    let (mut i, mut j) = (0, 0);
    while i < left.len() || j < right.len() {
        if j == right.len() || (i < left.len() && left[i] < right[j]) {
            output.push(left[i]);
            i += 1;
        } else {
            output.push(right[j]);
            j += 1;
        }
    }
    output
}

fn checked_sum(left: usize, right: usize, message: &str) -> Result<usize> {
    left.checked_add(right)
        .ok_or_else(|| Error::InvalidInput(message.into()))
}