alpha-id/src/pipeline.rs

use crate::claml;
use crate::corpus::{self};
use crate::embed::EmbeddingStore;
use crate::fusion::{cross_segment_dedupe, rrf_merge};
use crate::ionos::IonosClient;
use crate::lexical::LexicalIndex;
use crate::model::*;
use crate::rerank;
use crate::segment::split_segments;
use crate::tags::tags_for;
use crate::vector::VectorIndex;
use std::collections::HashMap;
use std::time::Instant;

pub struct Pipeline {
    pub entries: Vec<AlphaIdEntry>,
    pub meta: HashMap<String, IcdMeta>,
    pub lexical: LexicalIndex,
    pub pool_size: usize,
    pub valid_by_code: HashMap<String, bool>,
    cfg: Config,
    vector: Option<VectorIndex>,
}

impl Pipeline {
    pub fn load(cfg: &Config) -> Result<Self, AppError> {
        let entries = corpus::load(&cfg.alpha_id_path)?;
        let meta = claml::load(&cfg.claml_path)?;
        let lexical = LexicalIndex::build_in_ram(&entries)?;
        let valid_by_code = {
            let mut m: HashMap<String, bool> = HashMap::new();
            for e in &entries {
                if let Some(code) = corpus::primary_code(e) {
                    let v = m.entry(code.to_string()).or_insert(false);
                    *v = *v || e.valid;
                }
            }
            m
        };
        // Build the semantic index only if EVERY corpus entry already has a
        // cached embedding; otherwise leave it `None` so Mode Hybrid degrades to
        // Lexical rather than mispairing partial vectors.
        let vector = match EmbeddingStore::open(&cfg.index_dir, &cfg.embed_model) {
            Ok(store) if entries.iter().all(|e| store.has(&e.text)) => {
                let vecs: Option<Vec<Vec<f32>>> =
                    entries.iter().map(|e| store.get(&e.text)).collect();
                vecs.map(VectorIndex::from_vectors)
            }
            _ => None,
        };
        Ok(Self {
            entries,
            meta,
            lexical,
            pool_size: cfg.pool_size,
            valid_by_code,
            cfg: cfg.clone(),
            vector,
        })
    }

    /// Look up ICD metadata: exact code first, then — only for 6-char dotted codes
    /// of the form `Lxx.yz` — the 5-char parent `Lxx.y`.
    ///
    /// The ClaML map (populated by `claml::load`) now contains synthesized
    /// 5th-digit entries (e.g. `E11.90`, `I10.91`) with correct `Para295 = "P"`,
    /// so the exact hit succeeds for the vast majority of corpus codes.  The
    /// 5-char-parent fallback covers any residual gap where a 6-char code is not
    /// explicitly present but its 5-char parent is.
    ///
    /// No 3-char-root fallback is applied: that level caused ~1 193 codes to
    /// inherit the parent's `Para295 = "V"` and be wrongly dropped under
    /// `--billable-only` (Goal 1 regression).  Codes absent from ClaML entirely
    /// yield `None` (accepted, documented out-of-scope gap).
    fn meta_lookup<'a>(&'a self, code: &str) -> Option<&'a IcdMeta> {
        if let Some(m) = self.meta.get(code) {
            return Some(m);
        }
        // 5-char-parent fallback only for 6-char dotted codes: Lxx.yz → Lxx.y
        // Pattern: one letter, two digits, dot, two digits (total 6 chars).
        if code.len() == 6 {
            let bytes = code.as_bytes();
            if bytes[0].is_ascii_alphabetic()
                && bytes[1].is_ascii_digit()
                && bytes[2].is_ascii_digit()
                && bytes[3] == b'.'
                && bytes[4].is_ascii_digit()
                && bytes[5].is_ascii_digit()
            {
                if let Some(m) = self.meta.get(&code[..5]) {
                    return Some(m);
                }
            }
        }
        None
    }

    fn to_suggestion(&self, code: &str, phrase: &str, score: f32,
                     segs: Vec<usize>, alpha_valid: bool) -> Option<Suggestion> {
        let m = self.meta_lookup(code)?; // None ≈ 31 codes absent from ClaML entirely; accepted gap (see meta_lookup doc)
        Some(Suggestion {
            icd_code: code.to_string(),
            description: m.description.clone(),
            score,
            matched_phrase: phrase.to_string(),
            source_segments: segs,
            tags: tags_for(m, alpha_valid),
        })
    }

    /// Mode Lexical candidate collection: lexical pool per segment. This is the
    /// exact original behavior; it is also the degrade fallback for Mode Hybrid.
    fn collect_lexical(&self, segments: &[String]) -> Vec<Candidate> {
        let mut candidates: Vec<Candidate> = Vec::new();
        for (si, seg) in segments.iter().enumerate() {
            let lex = self.lexical.search(seg, self.pool_size).unwrap_or_default();
            for (entry_idx, score) in lex {
                let e = &self.entries[entry_idx];
                if let Some(code) = corpus::primary_code(e) {
                    candidates.push(Candidate {
                        icd_code: code.to_string(),
                        alpha_text: e.text.clone(),
                        segment_idx: si,
                        lexical: Some(score),
                        semantic: None,
                        rerank: None,
                    });
                }
            }
        }
        candidates
    }

    /// Mode Hybrid candidate collection: per segment, fuse a lexical pool with a
    /// semantic pool (IONOS embedding + in-RAM vector search) via RRF, then
    /// cross-encoder rerank the fused pool. ANY error (IONOS unreachable, no
    /// vector index, embed/rerank failure) propagates so `suggest` degrades.
    fn collect_hybrid(&self, segments: &[String], calls: &mut usize)
        -> Result<Vec<Candidate>, AppError> {
        let client = IonosClient::new(&self.cfg.ionos_base_url, &self.cfg.token_path)?;
        let vector = self.vector.as_ref()
            .ok_or_else(|| AppError::Ionos("no vector index".into()))?;
        let mut candidates: Vec<Candidate> = Vec::new();

        for (si, seg) in segments.iter().enumerate() {
            // Lexical pool (entry indices).
            let lex = self.lexical.search(seg, self.pool_size).unwrap_or_default();
            let lex_ids: Vec<usize> = lex.iter().map(|(i, _)| *i).collect();

            // Semantic pool: embed the segment, then exact cosine search.
            let body = serde_json::json!({
                "model": self.cfg.embed_model, "input": [seg]
            });
            let resp: serde_json::Value = client.post_json("/embeddings", &body)?;
            *calls += 1;
            let q: Vec<f32> = resp["data"][0]["embedding"].as_array()
                .ok_or_else(|| AppError::Ionos("embeddings: no data[0].embedding".into()))?
                .iter()
                .map(|x| x.as_f64()
                    .ok_or_else(|| AppError::Ionos("embeddings: non-numeric component".into())))
                .collect::<Result<Vec<f64>, _>>()?
                .into_iter().map(|f| f as f32).collect();
            let sem = vector.top_k(&q, self.pool_size);
            let sem_ids: Vec<usize> = sem.iter().map(|(i, _)| *i).collect();

            // Fuse the two ranked index lists; cap at pool_size.
            let mut fused = rrf_merge(&lex_ids, &sem_ids, 60);
            fused.truncate(self.pool_size);

            // Cross-encoder rerank the fused pool.
            let docs: Vec<String> = fused.iter()
                .map(|&i| self.entries[i].text.clone())
                .collect();
            let scores = rerank::rerank(&client, &self.cfg.rerank_model, seg, &docs)?;
            *calls += 1;

            for (rank, &entry_idx) in fused.iter().enumerate() {
                let e = &self.entries[entry_idx];
                if let Some(code) = corpus::primary_code(e) {
                    candidates.push(Candidate {
                        icd_code: code.to_string(),
                        alpha_text: e.text.clone(),
                        segment_idx: si,
                        lexical: None,
                        semantic: None,
                        rerank: scores.get(rank).copied(),
                    });
                }
            }
        }
        Ok(candidates)
    }

    pub fn suggest(&self, dictation: &str, mode: Mode, filter: &Filter, top_k: usize)
        -> SuggestResult {
        let t0 = Instant::now();
        let segments = split_segments(dictation);
        let mut ionos_calls = 0usize;
        let mut degraded = false;

        let candidates: Vec<Candidate> = match mode {
            Mode::Lexical => self.collect_lexical(&segments),
            Mode::Hybrid => match self.collect_hybrid(&segments, &mut ionos_calls) {
                Ok(c) => c,
                Err(_) => {
                    // Total degradation: any failure → Mode Lexical fallback,
                    // request still answered, never panics.
                    degraded = true;
                    self.collect_lexical(&segments)
                }
            },
        };

        let merged = cross_segment_dedupe(candidates);
        let mut suggestions = Vec::new();
        for (d, score) in merged {
            let alpha_valid = self.valid_by_code.get(&d.icd_code).copied().unwrap_or(false);
            if let Some(s) = self.to_suggestion(&d.icd_code, &d.best_phrase, score,
                                                d.source_segments.clone(), alpha_valid) {
                if filter.accepts(&s.tags) {
                    suggestions.push(s);
                }
            }
            if suggestions.len() >= top_k { break; }
        }

        SuggestResult {
            suggestions,
            diagnostics: Diagnostics {
                mode: format!("{:?}", mode),
                degraded,
                segment_count: segments.len(),
                ionos_calls,
                millis: t0.elapsed().as_millis(),
            },
        }
    }
}