use js_sys::{Object, Reflect};
use std::str::FromStr;
use wasm_bindgen::prelude::*;
use crate::annotator::Annotator;
use crate::numbering::segment;
use crate::types::{Chain, Scheme};
#[wasm_bindgen(typescript_custom_section)]
const TS_TYPES: &str = r#"
/** Numbered residues keyed by IMGT/Kabat position string (e.g. `"112A"`). */
export type Numbering = Record<string, string>;
/** Result returned by {@link Annotator.number}. */
export interface NumberingResult {
/** Detected chain type: `"H"`, `"K"`, `"L"`, `"A"`, `"B"`, `"G"`, or `"D"`. */
chain: string;
/** Numbering scheme used: `"imgt"` or `"kabat"`. */
scheme: string;
/** Alignment confidence score between 0 and 1. */
confidence: number;
/** Position-to-residue mapping for the aligned region. */
numbering: Numbering;
}
/** FR/CDR segments returned by {@link Annotator.segment}. */
export interface SegmentationResult {
fr1: string;
cdr1: string;
fr2: string;
cdr2: string;
fr3: string;
cdr3: string;
fr4: string;
/** Residues before FR1 (non-canonical N-terminal extension). */
prefix: string;
/** Residues after FR4 (non-canonical C-terminal extension). */
postfix: string;
}
/**
* Annotates antibody and T-cell receptor sequences with IMGT or Kabat position numbers.
*
* @param chains - Chain types to consider during auto-detection. Each entry is a
* case-insensitive string. Accepted values:
* - Antibody heavy chain: `"IGH"` / `"H"` / `"heavy"`
* - Antibody kappa chain: `"IGK"` / `"K"` / `"kappa"`
* - Antibody lambda chain: `"IGL"` / `"L"` / `"lambda"`
* - TCR alpha chain: `"TRA"` / `"A"` / `"alpha"`
* - TCR beta chain: `"TRB"` / `"B"` / `"beta"`
* - TCR gamma chain: `"TRG"` / `"G"` / `"gamma"`
* - TCR delta chain: `"TRD"` / `"D"` / `"delta"`
*
* Pass all chains you want to consider; the annotator scores each and picks the
* best-matching one. To consider every supported chain pass all seven values.
*
* @param scheme - Numbering scheme to use for output positions. Accepted values
* (case-insensitive):
* - `"IMGT"` / `"i"` — IMGT numbering (recommended; used internally)
* - `"Kabat"` / `"k"` — Kabat numbering (derived from IMGT)
*
* @param min_confidence - Optional minimum alignment confidence threshold in the
* range `[0, 1]`. Sequences scoring below this value are rejected with an error.
* Defaults to `0.5` when `null` or omitted.
*/
export class Annotator {
free(): void;
[Symbol.dispose](): void;
constructor(chains: string[], scheme: string, min_confidence?: number | null);
number(sequence: string): NumberingResult;
segment(sequence: string): SegmentationResult;
}
"#;
#[wasm_bindgen]
impl Annotator {
#[wasm_bindgen(constructor, js_name = "new", skip_typescript)]
pub fn wasm_new(
chains: Vec<String>,
scheme: String,
min_confidence: Option<f32>,
) -> Result<Annotator, JsValue> {
let parsed_chains = chains
.iter()
.map(|chain| {
Chain::from_str(chain)
.map_err(|_| JsValue::from_str(&format!("Invalid chain: {}", chain)))
})
.collect::<Result<Vec<_>, _>>()?;
let parsed_scheme = Scheme::from_str(&scheme)
.map_err(|_| JsValue::from_str(&format!("Invalid scheme: {}", scheme)))?;
Annotator::new(&parsed_chains, parsed_scheme, min_confidence)
.map_err(|e| JsValue::from_str(&e.to_string()))
}
#[wasm_bindgen(js_name = "number", skip_typescript)]
pub fn wasm_number(&self, sequence: &str) -> Result<JsValue, JsValue> {
let result = self
.number(sequence)
.map_err(|e| JsValue::from_str(&e.to_string()))?;
let aligned_seq = &sequence[result.query_start..=result.query_end];
let numbering = Object::new();
for (pos, ch) in result.positions.iter().zip(aligned_seq.chars()) {
Reflect::set(
&numbering,
&JsValue::from_str(&pos.to_string()),
&JsValue::from_str(&ch.to_string()),
)
.unwrap();
}
let dict = Object::new();
Reflect::set(&dict, &"chain".into(), &result.chain.to_string().into()).unwrap();
Reflect::set(&dict, &"scheme".into(), &result.scheme.to_string().into()).unwrap();
Reflect::set(&dict, &"confidence".into(), &result.confidence.into()).unwrap();
Reflect::set(&dict, &"numbering".into(), &numbering.into()).unwrap();
Ok(dict.into())
}
#[wasm_bindgen(js_name = "segment", skip_typescript)]
pub fn wasm_segment(&self, sequence: &str) -> Result<JsValue, JsValue> {
let result = self
.number(sequence)
.map_err(|e| JsValue::from_str(&e.to_string()))?;
let aligned_seq = &sequence[result.query_start..=result.query_end];
let segments = segment(&result.positions, aligned_seq, result.scheme);
let dict = Object::new();
for (region, seq) in &segments {
Reflect::set(&dict, &JsValue::from_str(region), &JsValue::from_str(seq)).unwrap();
}
Ok(dict.into())
}
}
#[cfg(test)]
mod tests {
use super::*;
use wasm_bindgen_test::wasm_bindgen_test;
const IGH_SEQ: &str =
"QVQLVQSGAEVKRPGSSVTVSCKASGGSFSTYALSWVRQAPGRGLEWMGGVIPLLTITNYAPRFQGRITITADRSTSTAYLELNSLRPEDTAVYYCAREGTTGKPIGAFAHWGQGTLVTVSS";
#[wasm_bindgen_test]
fn test_number_igh() {
let ann = Annotator::wasm_new(
vec!["IGH".to_string(), "IGK".to_string(), "IGL".to_string()],
"IMGT".to_string(),
None,
)
.unwrap();
let result = ann.wasm_number(IGH_SEQ).unwrap();
let chain = Reflect::get(&result, &"chain".into()).unwrap();
assert_eq!(chain.as_string().unwrap(), "H");
let confidence = Reflect::get(&result, &"confidence".into()).unwrap();
assert!(confidence.as_f64().unwrap() > 0.5);
}
#[wasm_bindgen_test]
fn test_segment_igh() {
let ann = Annotator::wasm_new(
vec!["IGH".to_string(), "IGK".to_string(), "IGL".to_string()],
"IMGT".to_string(),
None,
)
.unwrap();
let result = ann.wasm_segment(IGH_SEQ).unwrap();
let fr1 = Reflect::get(&result, &"fr1".into()).unwrap();
assert!(!fr1.as_string().unwrap().is_empty());
}
#[wasm_bindgen_test]
fn test_invalid_chain_errors() {
let err = Annotator::wasm_new(vec!["INVALID".to_string()], "IMGT".to_string(), None);
assert!(err.is_err());
}
#[wasm_bindgen_test]
fn test_invalid_scheme_errors() {
let err = Annotator::wasm_new(vec!["IGH".to_string()], "NOTASCHEME".to_string(), None);
assert!(err.is_err());
}
}