alpha-id/src/bin/alpha_id.rs

use alpha_id::model::{normalize_chapters, Config, Filter, Mode};
use alpha_id::pipeline::Pipeline;
use clap::{Parser, Subcommand};
use std::io::Read;

#[derive(Parser)]
#[command(name = "alpha-id")]
struct Cli {
    #[arg(long, default_value = "config/default.toml", global = true)]
    config: String,
    #[command(subcommand)]
    cmd: Cmd,
}

#[derive(Subcommand)]
enum Cmd {
    /// Build indexes (lexical now; embeddings added in Phase A)
    Index { #[arg(long)] sample: Option<usize>, #[arg(long)] full: bool, #[arg(long)] confirm: bool, #[arg(long)] pack: bool },
    /// Suggest ICD codes for a dictation (FILE or - for stdin)
    Suggest {
        input: String,
        #[arg(long, default_value = "lexical")] mode: String,
        #[arg(long, default_value_t = 10)] top: usize,
        #[arg(long)] billable_only: bool,
        #[arg(long)] valid_only: bool,
        #[arg(long)] exclude_exotic: bool,
        #[arg(long)] chapters: Option<String>,
        #[arg(long)] json: bool,
    },
    /// Run the eval harness
    Eval {
        #[arg(long, default_value = "lexical")] mode: String,
        #[arg(long, default_value_t = 200)] cases: usize,
        #[arg(long, default_value_t = 42)] seed: u64,
    },
}

fn run_eval(p: &Pipeline, m: Mode, cases: usize, seed: u64) -> (f64, f64) {
    use alpha_id::eval::{inject_errors, recall_at_k};
    use rand::SeedableRng;
    use rand::seq::SliceRandom;
    let all_billable: Vec<_> = p.entries.iter().enumerate()
        .filter(|(_, e)| e.valid)
        .filter(|(_, e)| alpha_id::corpus::primary_code(e)
            .and_then(|c| p.meta.get(c))
            .map(|mt| mt.para295 == "P").unwrap_or(false))
        .collect();
    let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
    let n = cases.min(all_billable.len());
    let billable: Vec<_> = all_billable
        .choose_multiple(&mut rng, n)
        .cloned()
        .collect();
    let mut sum = 0.0;
    for (k, (_, e)) in billable.iter().enumerate() {
        let code = alpha_id::corpus::primary_code(e).unwrap().to_string();
        let noisy = inject_errors(&e.text, 0.15, seed.wrapping_add(k as u64));
        let res = p.suggest(&noisy, m, &Filter::default(), 10);
        let preds: Vec<String> = res.suggestions.iter().map(|s| s.icd_code.clone()).collect();
        sum += recall_at_k(&preds, &code, 10);
    }
    let n = billable.len().max(1) as f64;
    // Both metrics are equal by construction: the loop runs only over
    // billable cases, so Recall@10 == BillableRecall@10 here.
    (sum / n, sum / n)
}

fn read_input(arg: &str) -> String {
    if let Ok(v) = std::env::var("ALPHA_ID_STDIN") { return v; }
    if arg == "-" {
        let mut s = String::new();
        std::io::stdin().read_to_string(&mut s).ok();
        s
    } else {
        std::fs::read_to_string(arg).unwrap_or_default()
    }
}

fn main() {
    let cli = Cli::parse();
    let cfg = Config::load(&cli.config).expect("config");
    match cli.cmd {
        Cmd::Index { sample, full, confirm, pack } => {
            use alpha_id::embed::{EmbeddingStore, build_corpus_embeddings, estimate};
            use alpha_id::ionos::IonosClient;

            if pack {
                use alpha_id::packed::{corpus_sha256, packed_path, write_packed, PackedHeader};
                let entries = alpha_id::corpus::load(&cfg.alpha_id_path).expect("corpus");
                let store = EmbeddingStore::open(&cfg.index_dir, &cfg.embed_model).expect("store");
                let cached = entries.iter().filter(|e| store.has(&e.text)).count();
                if cached != entries.len() {
                    eprintln!(
                        "error: store incomplete: {} of {} entries embedded; \
                         run `index --full --confirm` first",
                        cached, entries.len()
                    );
                    std::process::exit(2);
                }
                let rows: Vec<Vec<f32>> =
                    entries.iter().map(|e| store.get(&e.text).expect("cached vector")).collect();
                let dim = rows.first().map(|r| r.len()).unwrap_or(0) as u32;
                let corpus_name = std::path::Path::new(&cfg.alpha_id_path)
                    .file_name().and_then(|n| n.to_str())
                    .unwrap_or(&cfg.alpha_id_path).to_string();
                let header = PackedHeader {
                    n_rows: entries.len() as u32,
                    dim,
                    embed_model: cfg.embed_model.clone(),
                    corpus_sha256: corpus_sha256(&cfg.alpha_id_path).expect("corpus hash"),
                    corpus_name,
                };
                let path = packed_path(&cfg.index_dir, &cfg.embed_model);
                write_packed(&path, &header, &rows).expect("write packed");
                eprintln!("packed {} rows × {} dims → {}", header.n_rows, header.dim, path.display());
                return;
            }

            let p = Pipeline::load(&cfg).expect("load");
            let texts: Vec<String> = p.entries.iter().map(|e| e.text.clone()).collect();
            let client = IonosClient::new(&cfg.ionos_base_url, &cfg.token_path)
                .expect("ionos client (token?)");
            let store = EmbeddingStore::open(&cfg.index_dir, &cfg.embed_model).expect("store");

            if let Some(n) = sample {
                let subset: Vec<String> = texts.iter().take(n).cloned().collect();
                let est = estimate(&subset);
                eprintln!("SMOKE: embedding {} sample texts ({} chars)…", est.items, est.chars);
                let n_done = build_corpus_embeddings(&client, &store, &cfg.embed_model, &subset, 32)
                    .expect("sample embed failed");
                let dim = subset.first()
                    .and_then(|t| store.get(t))
                    .map(|v| v.len())
                    .unwrap_or(0);
                eprintln!("SMOKE OK: embedded {n_done}, vector dim = {dim}. \
                           Inspect, then run `index --full --confirm`.");
                return;
            }

            if full {
                let pending: Vec<String> = texts.iter().filter(|t| !store.has(t)).cloned().collect();
                let est = estimate(&pending);
                eprintln!("FULL RUN: {} of {} texts not yet embedded ({} chars). \
                           This calls IONOS and costs money.", est.items, texts.len(), est.chars);
                if !confirm {
                    eprintln!("Refusing without --confirm. Re-run: \
                               `alpha-id index --full --confirm`");
                    std::process::exit(2);
                }
                let n = build_corpus_embeddings(&client, &store, &cfg.embed_model, &texts, 32)
                    .expect("full embed failed");
                eprintln!("FULL RUN done: {n} newly embedded, {} total cached.",
                    texts.iter().filter(|t| store.has(t)).count());
                return;
            }

            eprintln!("Specify --sample N (smoke test first) or --full --confirm.");
        }
        Cmd::Suggest { input, mode, top, billable_only, valid_only, exclude_exotic, chapters, json } => {
            let p = Pipeline::load(&cfg).expect("load");
            let m: Mode = mode.parse().expect("mode");
            let chapters = chapters.map(|c| {
                let (padded, unknown) = normalize_chapters(&c);
                for tok in &unknown {
                    eprintln!("warning: --chapters token '{tok}' matches no known chapter (01..22); ignored");
                }
                padded
            });
            let filter = Filter {
                billable_only, valid_only, exclude_exotic,
                chapters,
            };
            let text = read_input(&input);
            let res = p.suggest(&text, m, &filter, top);
            if json {
                println!("{}", serde_json::to_string_pretty(&res).unwrap());
            } else {
                for s in &res.suggestions {
                    println!("{:<8} {:>5.3}  {}", s.icd_code, s.score, s.description);
                }
                eprintln!("[{} segments, mode {}, {} ms, degraded={}]",
                    res.diagnostics.segment_count, res.diagnostics.mode,
                    res.diagnostics.millis, res.diagnostics.degraded);
            }
        }
        Cmd::Eval { mode, cases, seed } => {
            let modes: Vec<Mode> = if mode == "both" {
                vec![Mode::Lexical, Mode::Hybrid]
            } else {
                vec![mode.parse().expect("mode")]
            };
            let p = Pipeline::load(&cfg).expect("load");
            for m in modes {
                let (r10, br10) = run_eval(&p, m, cases, seed);
                println!("mode={:?} cases={} Recall@10={:.3} BillableRecall@10={:.3}",
                         m, cases, r10, br10);
            }
        }
    }
}