tokitai-operator 0.1.0

//! The cost model used by the planner.
//!
//! `CostModel` estimates `runtime_ns`, `memory_bytes`,
//! `precision_loss`, and the per-step `valuation` rationale. The
//! model is non-standard: it carries a `Fallback` enum that
//! records why a fallback path was chosen (e.g. "backend
//! capability missing", "p-adic valuation cutoff too high").
//!
//! The non-standard cost model is one of the 4 paper-novelty
//! axes; see `docs/nonstandard_cost_model.md` for the policy.
//!
use crate::backend::BackendCapabilities;
use crate::backend::hardware::{DeviceCapabilities, DeviceKind};
use crate::backend::memory::{TransferPlan, TransferStatus};
use crate::planner::{ExecutionPlan, PlanStepKind};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CostModelWeights {
    pub base_step_cost: u64,
    pub missing_kernel_penalty: u64,
    pub fallback_penalty: u64,
    pub transfer_penalty: u64,
    pub unsupported_transfer_penalty: u64,
    pub unsupported_layout_penalty: u64,
    pub padic_missing_capability_penalty: u64,
    pub sheaf_missing_capability_penalty: u64,
    pub valuation_skip_credit_per_term: u64,
    pub sheaf_locality_credit_per_open: u64,
}

impl Default for CostModelWeights {
    fn default() -> Self {
        Self::balanced()
    }
}

impl CostModelWeights {
    /// Balanced weights (3x base, 3x penalties, 1x credits) — the default
    /// cost model used by the planner. Suitable for most planning paths.
    pub fn balanced() -> Self {
        Self {
            base_step_cost: 100,
            missing_kernel_penalty: 500,
            fallback_penalty: 1_000,
            transfer_penalty: 200,
            unsupported_transfer_penalty: 1_500,
            unsupported_layout_penalty: 700,
            padic_missing_capability_penalty: 2_000,
            sheaf_missing_capability_penalty: 2_000,
            valuation_skip_credit_per_term: 25,
            sheaf_locality_credit_per_open: 40,
        }
    }

    /// Strict weights (3x base, 5x penalties, halved credits). Use when
    /// the planner must prefer the most conservative plan even at the
    /// cost of rejecting borderline-correct plans.
    pub fn strict() -> Self {
        Self {
            base_step_cost: 300,
            missing_kernel_penalty: 2_500,
            fallback_penalty: 5_000,
            transfer_penalty: 600,
            unsupported_transfer_penalty: 7_500,
            unsupported_layout_penalty: 3_500,
            padic_missing_capability_penalty: 10_000,
            sheaf_missing_capability_penalty: 10_000,
            valuation_skip_credit_per_term: 12,
            sheaf_locality_credit_per_open: 20,
        }
    }

    /// Permissive weights (half base, half penalties, 2x credits). Use
    /// when the planner is allowed to take more risks in exchange for
    /// cheaper plans, e.g. in synthesis-style code where fallbacks are
    /// acceptable.
    pub fn permissive() -> Self {
        Self {
            base_step_cost: 50,
            missing_kernel_penalty: 250,
            fallback_penalty: 500,
            transfer_penalty: 100,
            unsupported_transfer_penalty: 750,
            unsupported_layout_penalty: 350,
            padic_missing_capability_penalty: 1_000,
            sheaf_missing_capability_penalty: 1_000,
            valuation_skip_credit_per_term: 50,
            sheaf_locality_credit_per_open: 80,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BackendCostModel {
    weights: CostModelWeights,
}

impl Default for BackendCostModel {
    fn default() -> Self {
        Self::new()
    }
}

impl BackendCostModel {
    pub fn new() -> Self {
        Self {
            weights: CostModelWeights::default(),
        }
    }

    pub fn with_weights(weights: CostModelWeights) -> Self {
        Self { weights }
    }

    pub fn score_plan(
        &self,
        name: impl Into<String>,
        plan: &ExecutionPlan,
        device: &DeviceCapabilities,
        transfers: &[TransferPlan],
    ) -> PlanCandidate {
        let mut score = self
            .weights
            .base_step_cost
            .saturating_mul(plan.steps.len() as u64);
        let mut evidence = vec![format!(
            "candidate scored for backend {} on {:?}",
            device.backend.name, device.target.kind
        )];
        let mut reasons = Vec::new();

        let missing_kernel_steps = plan
            .steps
            .iter()
            .filter(|step| step.lowering_rule_id.is_none())
            .count();
        if missing_kernel_steps > 0 {
            let penalty = self
                .weights
                .missing_kernel_penalty
                .saturating_mul(missing_kernel_steps as u64);
            score = score.saturating_add(penalty);
            reasons.push(CostReason::MissingKernel {
                steps: missing_kernel_steps,
                penalty,
            });
            evidence.push(format!(
                "missing lowering/kernel penalty: steps={}, penalty={}",
                missing_kernel_steps, penalty
            ));
        }

        if let Some(fallback) = &plan.fallback {
            score = score.saturating_add(self.weights.fallback_penalty);
            reasons.push(CostReason::Fallback {
                backend: fallback.backend.clone(),
                reason: fallback.reason.clone(),
                penalty: self.weights.fallback_penalty,
            });
            evidence.push(format!(
                "fallback penalty: backend={}, reason={}",
                fallback.backend, fallback.reason
            ));
        }

        score =
            self.score_backend_support(plan, &device.backend, score, &mut reasons, &mut evidence);
        score = self.score_transfers(transfers, score, &mut reasons, &mut evidence);
        score = self.apply_semantic_opportunity_credits(plan, score, &mut reasons, &mut evidence);

        PlanCandidate {
            name: name.into(),
            backend: device.backend.name.clone(),
            device_kind: device.target.kind,
            capability_fingerprint: device.fingerprint(),
            score,
            selected: false,
            reasons,
            evidence,
        }
    }

    pub fn choose_best(&self, mut candidates: Vec<PlanCandidate>) -> Option<PlanSelection> {
        candidates.sort_by(|lhs, rhs| {
            lhs.score
                .cmp(&rhs.score)
                .then_with(|| lhs.backend.cmp(&rhs.backend))
                .then_with(|| lhs.name.cmp(&rhs.name))
        });
        let mut candidates = candidates;
        let selected = candidates.first_mut()?;
        selected.selected = true;
        let selected_name = selected.name.clone();
        let selected_backend = selected.backend.clone();
        let selected_score = selected.score;
        let rationale = format!(
            "selected candidate {} on backend {} with score {}",
            selected_name, selected_backend, selected_score
        );
        Some(PlanSelection {
            selected_name,
            selected_backend,
            selected_score,
            rationale,
            candidates,
        })
    }

    fn score_backend_support(
        &self,
        plan: &ExecutionPlan,
        backend: &BackendCapabilities,
        mut score: u64,
        reasons: &mut Vec<CostReason>,
        evidence: &mut Vec<String>,
    ) -> u64 {
        let has_padic = plan
            .steps
            .iter()
            .any(|step| step.domain.starts_with("Q_") || step.domain.contains("padic"));
        if has_padic && !supports_domain(backend, "padic:fixed_precision") {
            score = score.saturating_add(self.weights.padic_missing_capability_penalty);
            reasons.push(CostReason::PadicMissingCapability {
                backend: backend.name.clone(),
                penalty: self.weights.padic_missing_capability_penalty,
            });
            evidence.push(format!(
                "backend {} lacks padic:fixed_precision for p-adic plan",
                backend.name
            ));
        }

        let has_sheaf = plan.steps.iter().any(|step| {
            matches!(step.kind, PlanStepKind::CoverGlueCheck { .. })
                || step.domain.starts_with("cover:")
        });
        if has_sheaf && !supports_domain(backend, "sheaf:finite_site") {
            score = score.saturating_add(self.weights.sheaf_missing_capability_penalty);
            reasons.push(CostReason::SheafMissingCapability {
                backend: backend.name.clone(),
                penalty: self.weights.sheaf_missing_capability_penalty,
            });
            evidence.push(format!(
                "backend {} lacks sheaf:finite_site for cover-local plan",
                backend.name
            ));
        }

        score
    }

    fn score_transfers(
        &self,
        transfers: &[TransferPlan],
        mut score: u64,
        reasons: &mut Vec<CostReason>,
        evidence: &mut Vec<String>,
    ) -> u64 {
        for transfer in transfers {
            match &transfer.status {
                TransferStatus::NoOp => {
                    evidence.push(format!(
                        "transfer {} -> {} is no-op",
                        transfer.source.id, transfer.destination.id
                    ));
                }
                TransferStatus::Supported => {
                    score = score.saturating_add(self.weights.transfer_penalty);
                    reasons.push(CostReason::Transfer {
                        direction: format!("{:?}", transfer.direction),
                        penalty: self.weights.transfer_penalty,
                    });
                    evidence.push(format!(
                        "transfer {} -> {} adds movement cost",
                        transfer.source.id, transfer.destination.id
                    ));
                }
                TransferStatus::Unsupported(reason) => {
                    let mut penalty = self.weights.unsupported_transfer_penalty;
                    if format!("{reason:?}").contains("UnsupportedLayout") {
                        penalty = penalty.saturating_add(self.weights.unsupported_layout_penalty);
                    }
                    score = score.saturating_add(penalty);
                    reasons.push(CostReason::UnsupportedTransfer {
                        reason: format!("{reason:?}"),
                        penalty,
                    });
                    evidence.push(format!(
                        "transfer {} -> {} unsupported: {reason:?}",
                        transfer.source.id, transfer.destination.id
                    ));
                }
            }
        }
        score
    }

    fn apply_semantic_opportunity_credits(
        &self,
        plan: &ExecutionPlan,
        mut score: u64,
        reasons: &mut Vec<CostReason>,
        evidence: &mut Vec<String>,
    ) -> u64 {
        if let Some(valuation) = &plan.cost.semantic.valuation {
            let skipped = valuation.estimated_skipped_terms as u64;
            let credit = skipped.saturating_mul(self.weights.valuation_skip_credit_per_term);
            if credit > 0 {
                score = score.saturating_sub(credit);
                reasons.push(CostReason::PadicValuationOpportunity {
                    skipped_terms: valuation.estimated_skipped_terms,
                    credit,
                });
                evidence.push(format!(
                    "p-adic valuation opportunity credit: skipped_terms={}, credit={}, rationale={}",
                    valuation.estimated_skipped_terms, credit, valuation.rationale
                ));
            } else {
                evidence.push(format!(
                    "p-adic valuation opportunity recorded without static skip credit: {}; certificate_policy={}; valuation_bucket_fingerprint_policy={}; precision_margin_floor={}",
                    valuation.rationale,
                    valuation.certificate_policy.as_deref().unwrap_or("none"),
                    valuation
                        .valuation_bucket_fingerprint_policy
                        .as_deref()
                        .unwrap_or("none"),
                    valuation
                        .precision_margin_floor
                        .map(|value| value.to_string())
                        .unwrap_or_else(|| "none".to_string())
                ));
            }
        }

        for step in &plan.steps {
            if let PlanStepKind::CoverGlueCheck { opens, .. } = &step.kind {
                let credit = (opens.len() as u64)
                    .saturating_mul(self.weights.sheaf_locality_credit_per_open);
                score = score.saturating_sub(credit);
                reasons.push(CostReason::SheafLocalityOpportunity {
                    opens: opens.len(),
                    credit,
                });
                evidence.push(format!(
                    "sheaf locality opportunity credit: opens={}, credit={}",
                    opens.len(),
                    credit
                ));
            }
        }

        score
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PlanCandidate {
    pub name: String,
    pub backend: String,
    pub device_kind: DeviceKind,
    pub capability_fingerprint: String,
    pub score: u64,
    pub selected: bool,
    pub reasons: Vec<CostReason>,
    pub evidence: Vec<String>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PlanSelection {
    pub selected_name: String,
    pub selected_backend: String,
    pub selected_score: u64,
    pub rationale: String,
    pub candidates: Vec<PlanCandidate>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CostReason {
    MissingKernel {
        steps: usize,
        penalty: u64,
    },
    Fallback {
        backend: String,
        reason: String,
        penalty: u64,
    },
    Transfer {
        direction: String,
        penalty: u64,
    },
    UnsupportedTransfer {
        reason: String,
        penalty: u64,
    },
    PadicMissingCapability {
        backend: String,
        penalty: u64,
    },
    SheafMissingCapability {
        backend: String,
        penalty: u64,
    },
    PadicValuationOpportunity {
        skipped_terms: usize,
        credit: u64,
    },
    SheafLocalityOpportunity {
        opens: usize,
        credit: u64,
    },
}

fn supports_domain(backend: &BackendCapabilities, domain: &str) -> bool {
    backend.supported_domains.iter().any(|item| item == domain)
}