1use std::cmp::Ordering;
2
3use crate::attribution::{edit_op_node_id, effect_node_id};
4use crate::calibration;
5use crate::graph::{CausalGraph, CausalNode};
6use crate::types::{CausalPrediction, EditOpSignature, RiskFlag};
7use check_runner::EffectSignature;
8
9const ZERO_SHOT_MIN_EFFECTIVE_SAMPLES: f64 = 5.0;
10const RISK_MATCH_COVERAGE_FLOOR: f64 = 0.90;
11
12#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
13pub struct PredictionConfig {
14 pub risk_confidence_threshold: f64,
15 pub zero_shot_coverage_threshold: f64,
16 pub fuzzy_top_k: usize,
17 pub min_samples_per_signature: usize,
19}
20
21impl Default for PredictionConfig {
22 fn default() -> Self {
23 Self {
24 risk_confidence_threshold: 0.65,
25 zero_shot_coverage_threshold: 0.6,
26 fuzzy_top_k: 3,
27 min_samples_per_signature: calibration::MIN_SAMPLES_PER_SIGNATURE,
28 }
29 }
30}
31
32pub fn predict(
33 signatures: &[EditOpSignature],
34 graph: &CausalGraph,
35 risk_confidence_threshold: f64,
36 zero_shot_coverage_threshold: f64,
37) -> CausalPrediction {
38 let config = PredictionConfig {
39 risk_confidence_threshold,
40 zero_shot_coverage_threshold,
41 ..PredictionConfig::default()
42 };
43 predict_with_config(signatures, graph, &config)
44}
45
46pub fn predict_with_config(
47 signatures: &[EditOpSignature],
48 graph: &CausalGraph,
49 config: &PredictionConfig,
50) -> CausalPrediction {
51 if signatures.is_empty() {
52 return CausalPrediction {
53 predicted_correctness: 0.5,
54 predicted_novelty: 1.0,
55 confidence: 0.0,
56 coverage_fraction: 0.0,
57 risk_flags: Vec::new(),
58 zero_shot_eligible: false,
59 };
60 }
61
62 let mut positive = 0.0;
63 let mut negative = 0.0;
64 let mut coverage_total = 0.0;
65 let mut risk_candidates = Vec::new();
66 let mut sample_evidence = 0.0;
67 let mut signature_confidences = Vec::new();
68
69 for signature in signatures {
70 let matches = resolve_signature_matches(signature, graph, config.fuzzy_top_k);
71 let signature_coverage = matches
72 .iter()
73 .map(|candidate| candidate.coverage_weight)
74 .fold(0.0, f64::max);
75 coverage_total += signature_coverage;
76 if matches.is_empty() {
77 continue;
78 }
79
80 let mut signature_confidence = 1.0_f64;
81
82 for matched in matches {
83 for (target_index, edge) in graph.outgoing_edges(matched.node_index) {
84 let Some(CausalNode::Effect(effect_signature)) =
85 graph.graph.node_weight(target_index)
86 else {
87 continue;
88 };
89
90 let edge_observations = edge.stats.observations as f64;
91 let match_coverage = matched.coverage_weight;
92 let edge_conservative_confidence = calibration::advisory_confidence(
93 calibration::conservative_reliability(edge.stats.alpha, edge.stats.beta),
94 match_coverage,
95 edge_observations,
96 1,
97 config.min_samples_per_signature,
98 );
99 signature_confidence = signature_confidence.min(edge_conservative_confidence);
100
101 let signal = edge.weight.max(0.0) * edge.stats.mean() * match_coverage;
102 if signal <= f64::EPSILON {
103 continue;
104 }
105
106 sample_evidence += edge_observations * match_coverage;
107
108 if effect_signature.outcome == "pass" {
109 positive += signal;
110 } else {
111 negative += signal;
112 risk_candidates.push(RawRiskCandidate {
113 op_signature: signature.clone(),
114 predicted_effect: effect_signature.clone(),
115 raw_confidence: edge_conservative_confidence,
116 coverage_weight: matched.coverage_weight,
117 effective_sample_size: calibration::effective_sample_size(
118 edge_observations,
119 ),
120 historical_weight: edge.weight,
121 });
122 }
123 }
124 }
125
126 signature_confidences.push(signature_confidence);
127 }
128
129 let coverage_fraction = (coverage_total / signatures.len() as f64).clamp(0.0, 1.0);
130 let total_signal = positive + negative;
131 let modeled_correctness = if total_signal.abs() < f64::EPSILON {
132 0.5
133 } else {
134 (positive / total_signal).clamp(0.0, 1.0)
135 };
136 let blended_correctness = modeled_correctness;
137 let signature_confidence = if signature_confidences.is_empty() {
138 0.0
139 } else {
140 signature_confidences.into_iter().fold(1.0, f64::min)
141 };
142 let confidence = calibration::advisory_confidence(
143 signature_confidence,
144 1.0,
145 sample_evidence,
146 signatures.len(),
147 config.min_samples_per_signature,
148 );
149
150 let mut risk_flags = Vec::new();
151 for candidate in risk_candidates {
152 if candidate.raw_confidence >= config.risk_confidence_threshold
153 && candidate.coverage_weight >= RISK_MATCH_COVERAGE_FLOOR
154 && candidate.effective_sample_size >= ZERO_SHOT_MIN_EFFECTIVE_SAMPLES
155 {
156 risk_flags.push(RiskFlag {
157 op_signature: candidate.op_signature,
158 predicted_effect: candidate.predicted_effect,
159 confidence: candidate.raw_confidence,
160 historical_weight: candidate.historical_weight,
161 });
162 }
163 }
164
165 let zero_shot_eligible = coverage_fraction >= config.zero_shot_coverage_threshold
166 && confidence >= config.risk_confidence_threshold
167 && calibration::effective_sample_size(sample_evidence) >= ZERO_SHOT_MIN_EFFECTIVE_SAMPLES
168 && risk_flags.is_empty();
169
170 risk_flags.sort_by(|left, right| {
171 right
172 .confidence
173 .partial_cmp(&left.confidence)
174 .unwrap_or(Ordering::Equal)
175 .then_with(|| {
176 effect_node_id(&left.predicted_effect).cmp(&effect_node_id(&right.predicted_effect))
177 })
178 .then_with(|| {
179 edit_op_node_id(&left.op_signature).cmp(&edit_op_node_id(&right.op_signature))
180 })
181 });
182
183 CausalPrediction {
184 predicted_correctness: blended_correctness.clamp(0.0, 1.0),
185 predicted_novelty: (1.0 - coverage_fraction).clamp(0.0, 1.0),
186 confidence,
187 coverage_fraction,
188 risk_flags,
189 zero_shot_eligible,
190 }
191}
192
193#[derive(Debug, Clone)]
194struct MatchCandidate {
195 node_index: petgraph::graph::NodeIndex,
196 coverage_weight: f64,
197}
198
199#[derive(Debug, Clone)]
200struct RawRiskCandidate {
201 op_signature: EditOpSignature,
202 predicted_effect: EffectSignature,
203 raw_confidence: f64,
204 coverage_weight: f64,
205 effective_sample_size: f64,
206 historical_weight: f64,
207}
208
209fn resolve_signature_matches(
210 signature: &EditOpSignature,
211 graph: &CausalGraph,
212 fuzzy_top_k: usize,
213) -> Vec<MatchCandidate> {
214 let node_id = edit_op_node_id(signature);
215 if let Some(node_index) = graph.node_index_map.get(&node_id) {
216 return vec![MatchCandidate {
217 node_index: *node_index,
218 coverage_weight: 1.0,
219 }];
220 }
221
222 let mut fuzzy = graph
223 .cause_nodes()
224 .into_iter()
225 .filter_map(|(node_index, candidate)| {
226 let similarity = heuristic_similarity(signature, candidate);
227 (similarity > 0.0).then_some((node_index, similarity))
228 })
229 .collect::<Vec<_>>();
230
231 fuzzy.sort_by(|left, right| {
232 right
233 .1
234 .partial_cmp(&left.1)
235 .unwrap_or(Ordering::Equal)
236 .then_with(|| left.0.index().cmp(&right.0.index()))
237 });
238 fuzzy.truncate(fuzzy_top_k.max(1));
239
240 let total_similarity = fuzzy.iter().map(|(_, similarity)| *similarity).sum::<f64>();
241 if total_similarity <= f64::EPSILON {
242 return Vec::new();
243 }
244
245 fuzzy
246 .into_iter()
247 .map(|(node_index, similarity)| MatchCandidate {
248 node_index,
249 coverage_weight: similarity,
250 })
251 .collect()
252}
253
254fn heuristic_similarity(left: &EditOpSignature, right: &EditOpSignature) -> f64 {
255 let mut score = 0.0;
256 let mut max_score = 0.0;
257
258 max_score += 3.0;
259 if left.op_kind == right.op_kind {
260 score += 3.0;
261 }
262
263 max_score += 1.0;
264 if left.anchor_kind == right.anchor_kind {
265 score += 1.0;
266 }
267
268 max_score += 2.0;
269 if left.file_extension == right.file_extension {
270 score += 2.0;
271 }
272
273 max_score += 2.0;
274 if left.scope_tag == right.scope_tag {
275 score += 2.0;
276 }
277
278 max_score += 2.0;
279 score += shared_prefix_ratio(&left.context_hash, &right.context_hash) * 2.0;
280
281 (score / max_score).clamp(0.0, 1.0)
282}
283
284fn shared_prefix_ratio(left: &str, right: &str) -> f64 {
285 let max_len = left.len().min(right.len()).max(1) as f64;
286 let shared = left
287 .chars()
288 .zip(right.chars())
289 .take_while(|(left, right)| left == right)
290 .count() as f64;
291 shared / max_len
292}