rexafs 0.2.5

//! Built-in library of common EXAFS standards (metals, oxides, sulfides,
//! halides) and molecular visualization examples, bundled as CIF text so the app works offline. The entries
//! come from the Crystallography Open Database (CC0); see
//! `data/builtin_cifs/catalog.json` and `scripts/fetch_builtin_cifs.py`
//! for provenance and `scripts/build_builtin_library.py` for the blob.

use std::io::Read;
use std::sync::OnceLock;

use serde::{Deserialize, Serialize};

use super::cif::structure_from_cif;
use super::db::{
    formula_matches, hit_from_structure, StructureHit, StructureQuery, StructureSource,
};
use super::model::Structure;
use super::StructureError;

/// Compressed catalog + CIF texts, generated by
/// `scripts/build_builtin_library.py`.
static BLOB: &[u8] = include_bytes!("builtin_library.json.gz");

/// One bundled structure.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct BuiltinEntry {
    /// Stable key (`ru_hcp`, `fe2o3_hematite`).
    pub key: String,
    /// Display name assigned to this bundled reference or visualization example.
    pub name: String,
    /// Hill formula as recorded by the source.
    pub formula: String,
    /// International Tables space-group number recorded in the catalog.
    pub sg_number: u16,
    /// Optional space-group symbol from the source catalog.
    #[serde(default)]
    pub sg: Option<String>,
    /// `metal` / `oxide` / `sulfide` / `molecule` / `other`.
    pub category: String,
    /// Origin of the structure record, retained for attribution.
    pub source: String,
    /// Source record id (`cod-4105040`).
    pub id: String,
    /// Optional source-record URL for checking provenance.
    #[serde(default)]
    pub url: Option<String>,
    /// Publication year text supplied by the source, when available.
    #[serde(default)]
    pub year: Option<String>,
    /// Source publication’s journal title, when available.
    #[serde(default)]
    pub journal: Option<String>,
    /// Source publication’s author text, when available.
    #[serde(default)]
    pub authors: Option<String>,
    /// Source publication’s DOI, when available; not generated from the structure.
    #[serde(default)]
    pub doi: Option<String>,
    /// Data-license statement retained from the catalog.
    pub license: String,
    /// Attribution text for this structure’s source.
    pub citation: String,
    /// Interpretation or provenance note; empty when none is recorded.
    #[serde(default)]
    pub note: String,
    /// Bundled CIF text parsed when this entry is loaded; no network fetch is needed.
    pub cif: String,
}

/// The bundled library.
#[derive(Debug)]
pub struct BuiltinLibrary {
    entries: Vec<BuiltinEntry>,
    hits: Vec<StructureHit>,
}

static LIBRARY: OnceLock<Result<BuiltinLibrary, String>> = OnceLock::new();

impl BuiltinLibrary {
    /// The process-wide library, decoded on first use.
    pub fn get() -> Result<&'static BuiltinLibrary, StructureError> {
        LIBRARY
            .get_or_init(|| Self::decode(BLOB).map_err(|e| e.to_string()))
            .as_ref()
            .map_err(|e| StructureError::Database {
                reason: format!("built-in library: {e}"),
            })
    }

    fn decode(blob: &[u8]) -> Result<Self, StructureError> {
        let mut text = String::new();
        flate2::read::GzDecoder::new(blob)
            .read_to_string(&mut text)
            .map_err(|e| StructureError::Database {
                reason: format!("built-in library is not valid gzip: {e}"),
            })?;
        let entries: Vec<BuiltinEntry> = serde_json::from_str(&text)?;
        let hits = entries
            .iter()
            .map(|e| {
                let mut hit = StructureHit {
                    id: e.key.clone(),
                    source: "builtin".into(),
                    formula: e.formula.replace(' ', ""),
                    name: Some(e.name.clone()),
                    space_group: e.sg.clone(),
                    elements: elements_of(&e.formula),
                    extra: Default::default(),
                };
                hit.extra.insert("category".into(), e.category.clone());
                hit.extra
                    .insert("origin".into(), format!("{} {}", e.source, e.id));
                if let Some(u) = &e.url {
                    hit.extra.insert("url".into(), u.clone());
                }
                hit.extra.insert("license".into(), e.license.clone());
                hit
            })
            .collect();
        Ok(Self { entries, hits })
    }

    /// Borrow all catalog records in bundled order, including citations and CIF text.
    pub fn entries(&self) -> &[BuiltinEntry] {
        &self.entries
    }

    /// Number of bundled catalog records.
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Whether the bundled catalog contains no records.
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Borrow the record with this exact stable key, or None when no key matches.
    pub fn entry(&self, key: &str) -> Option<&BuiltinEntry> {
        self.entries.iter().find(|e| e.key == key)
    }

    /// Parse and expand one entry.
    pub fn structure(&self, key: &str) -> Result<Structure, StructureError> {
        let entry = self.entry(key).ok_or_else(|| StructureError::Database {
            reason: format!("no built-in structure '{key}'"),
        })?;
        let mut s = structure_from_cif(&entry.cif)?;
        s.title = entry.name.clone();
        s.source = format!("built-in ({} {})", entry.source, entry.id);
        if s.mineral.is_none() {
            s.mineral = Some(entry.name.clone());
        }
        Ok(s)
    }
}

fn elements_of(formula: &str) -> Vec<String> {
    formula
        .split_whitespace()
        .map(|tok| {
            tok.chars()
                .take_while(|c| c.is_ascii_alphabetic())
                .collect::<String>()
        })
        .filter(|s| !s.is_empty())
        .collect()
}

impl StructureSource for BuiltinLibrary {
    fn name(&self) -> &str {
        "built-in"
    }

    fn search(&self, query: &StructureQuery) -> Result<Vec<StructureHit>, StructureError> {
        let text = query.text.as_deref().map(str::trim).unwrap_or("");
        let elements_only = StructureQuery {
            text: None,
            ..query.clone()
        };
        let mut hits: Vec<StructureHit> = self
            .hits
            .iter()
            .filter(|h| elements_only.matches(h) && (query.matches(h) || name_matches(h, text)))
            .cloned()
            .collect();
        // Exact key / formula match first, then alphabetical by name.
        let lower = text.to_ascii_lowercase();
        let exact = |h: &StructureHit| {
            !lower.is_empty() && (h.id == lower || formula_matches(&h.formula, text))
        };
        hits.sort_by(|a, b| exact(b).cmp(&exact(a)).then_with(|| a.name.cmp(&b.name)));
        if query.limit > 0 {
            hits.truncate(query.limit);
        }
        Ok(hits)
    }

    fn fetch(&self, hit: &StructureHit) -> Result<Structure, StructureError> {
        let mut s = self.structure(&hit.id)?;
        // Keep the hit's naming (the catalog name) on the structure.
        let h = hit_from_structure(&s, "builtin", &hit.id);
        if s.formula_sum.is_none() {
            s.formula_sum = Some(h.formula);
        }
        Ok(s)
    }
}

fn name_matches(hit: &StructureHit, text: &str) -> bool {
    if text.is_empty() {
        return true;
    }
    let t = text.to_ascii_lowercase();
    hit.name
        .as_deref()
        .is_some_and(|n| n.to_ascii_lowercase().contains(&t))
        || hit.id.contains(&t)
        || hit
            .extra
            .get("category")
            .is_some_and(|c| c.eq_ignore_ascii_case(&t))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::xafs::structure::cluster::{
        absorber_sites, build_cluster, AbsorberSelection, ClusterOptions,
    };

    #[test]
    fn library_decodes_and_every_entry_expands_and_clusters() {
        let lib = BuiltinLibrary::get().expect("bundled library");
        assert!(lib.len() >= 40, "only {} entries", lib.len());
        for e in lib.entries() {
            let s = lib
                .structure(&e.key)
                .unwrap_or_else(|err| panic!("{}: {err}", e.key));
            assert!(!s.sites.is_empty(), "{} has no sites", e.key);
            assert_eq!(
                s.space_group.number,
                Some(e.sg_number),
                "{}: space group {:?} vs catalog {}",
                e.key,
                s.space_group.number,
                e.sg_number
            );
            let first = s.elements()[0].clone();
            let sites = absorber_sites(&s, &first);
            assert!(!sites.is_empty(), "{}: no {first} site", e.key);
            let cluster = build_cluster(
                &s,
                &AbsorberSelection::SiteIndex(sites[0]),
                &ClusterOptions {
                    radius: 8.0,
                    ..Default::default()
                },
            )
            .unwrap_or_else(|err| panic!("{}: cluster {err}", e.key));
            assert!(
                cluster.atoms.len() > 10,
                "{}: only {} atoms",
                e.key,
                cluster.atoms.len()
            );
            assert!(cluster.shells(0.01).len() >= 2, "{}: too few shells", e.key);
        }
    }

    #[test]
    fn search_by_name_formula_and_element() {
        let lib = BuiltinLibrary::get().unwrap();
        let hits = lib.search(&StructureQuery::text("ruthenium")).unwrap();
        assert!(hits.iter().any(|h| h.id == "ru_hcp"), "{hits:?}");
        let hits = lib.search(&StructureQuery::text("RuO2")).unwrap();
        assert_eq!(hits.first().map(|h| h.id.as_str()), Some("ruo2"));
        let hits = lib
            .search(&StructureQuery::text("").with_elements(["Fe", "S"]))
            .unwrap();
        assert!(hits.iter().any(|h| h.id == "fes2_pyrite"));
        let all = lib.search(&StructureQuery::text("")).unwrap();
        assert_eq!(all.len(), lib.len());
        let s = lib.fetch(&hits[0]).unwrap();
        assert!(
            s.elements().iter().any(|e| e == "Fe"),
            "{} -> {:?} {:?}",
            hits[0].id,
            s.title,
            s.elements()
        );
    }
}