rexafs 0.2.5

//! American Mineralogist Crystal Structure Database, in the SQLite layout
//! shipped by Larch/larixite (`amcsd_cif1.db` trimmed, `amcsd_cif2.db`
//! full): tables `cif`, `minerals`, `spacegroups` (H-M symbol + JSON list
//! of `x,y,z` operations), `cif_elements`, `publications`, `authors`.
//! Coordinates are stored as base64-encoded little-endian integers: a decoded
//! value is the stored int32 divided by 4·10⁶. This is a database packaging
//! convention, not an experimental precision or uncertainty estimate.

use std::collections::BTreeMap;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};

use base64::Engine;
use rusqlite::{Connection, OpenFlags};

use super::{StructureHit, StructureQuery, StructureSource};
use crate::xafs::structure::cif::structure_from_cif;
use crate::xafs::structure::model::Structure;
use crate::xafs::structure::StructureError;

/// File name of the full database.
pub const AMCSD_FULL: &str = "amcsd_cif2.db";
/// File name of the trimmed database bundled with larixite.
pub const AMCSD_TRIM: &str = "amcsd_cif1.db";
/// Configured mirrors tried in order by [`download_amcsd`], following the
/// SQLite distribution used by larixite (Larch's structure toolkit).
/// The final entry is a direct figshare file URL; [`mirror_url`] appends the
/// database file name to the preceding base URLs. Availability is checked at
/// download time, not when these constants are accessed.
pub const SOURCE_URLS: [&str; 3] = [
    "https://docs.xrayabsorption.org/databases",
    "https://millenia.cars.aps.anl.gov/xraylarch/downloads",
    "https://figshare.com/ndownloader/files/54545639",
];
/// Home of the original database (Mineralogical Society of America and the
/// Mineralogical Association of Canada, NSF-funded).
pub const HOME_URL: &str = "https://rruff.geo.arizona.edu/AMS/amcsd.php";
/// Citation requested by the database's authors.
pub const CITATION: &str = "Downs, R.T. and Hall-Wallace, M. (2003) The American Mineralogist \
    Crystal Structure Database. American Mineralogist 88, 247-250";
/// Terms as stated on the AMCSD site: a public resource; a citation is
/// requested when a publication uses the data. The SQLite packaging is
/// larixite's (Matt Newville et al., xraypy).
pub const TERMS: &str = "AMCSD is a public database maintained by the Mineralogical Society of \
    America and the Mineralogical Association of Canada (NSF EAR-0112782, EAR-0622371). Please \
    cite Downs & Hall-Wallace (2003) when the data are used in a publication. SQLite packaging \
    by larixite (xraypy).";
/// The download `url` built for one mirror.
pub fn mirror_url(base: &str) -> String {
    if base.contains("figshare") {
        base.to_string()
    } else {
        format!("{base}/{AMCSD_FULL}")
    }
}

const FARRAY_SCALE: f64 = 4.0e6;

fn db_err<E: std::fmt::Display>(e: E) -> StructureError {
    StructureError::Database {
        reason: e.to_string(),
    }
}

/// Decode larixite's packed float arrays into dimensionless values or missing entries.
/// `'0'`, empty text, or invalid base64 returns an empty vector. Complete four-byte
/// little-endian integers are divided by 4e6; decoded sentinel values near 2 or 3
/// become None. Trailing incomplete bytes are ignored. This helper does not
/// validate that the returned array has the same length as a site's label array.
pub fn decode_farray(text: &str) -> Vec<Option<f64>> {
    let text = text.trim();
    if text.is_empty() || text == "0" {
        return Vec::new();
    }
    let bytes = match base64::engine::general_purpose::STANDARD.decode(text) {
        Ok(b) => b,
        Err(_) => return Vec::new(),
    };
    bytes
        .as_chunks::<4>()
        .0
        .iter()
        .map(|c| {
            let v = i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as f64 / FARRAY_SCALE;
            if (v - 2.0).abs() < 1e-5 || (v - 3.0).abs() < 1e-5 {
                None
            } else {
                Some(v)
            }
        })
        .collect()
}

/// Read-only connection to an AMCSD database file.
pub struct Amcsd {
    conn: Connection,
    path: PathBuf,
}

impl std::fmt::Debug for Amcsd {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Amcsd").field("path", &self.path).finish()
    }
}

/// A CIF record as stored in the database (cell + sites + symmetry).
#[derive(Debug, Clone)]
pub struct AmcsdRecord {
    /// AMCSD record identifier in the local SQLite database.
    pub id: i64,
    /// Mineral name when present in the record.
    pub mineral: Option<String>,
    /// Formula text recorded by the database.
    pub formula: String,
    /// Hermann–Mauguin space-group symbol from the database.
    pub hm_symbol: String,
    /// Fractional-coordinate symmetry operations such as x,-y,z+1/2.
    pub symmetry_xyz: Vec<String>,
    /// Cell parameters [a, b, c, alpha, beta, gamma]; lengths in Å and angles in degrees.
    pub cell: [f64; 6],
    /// Site labels in database order, matching the coordinate/occupancy arrays.
    pub sites: Vec<String>,
    /// Fractional x coordinates; None represents missing source values.
    pub x: Vec<Option<f64>>,
    /// Fractional y coordinates; None represents missing source values.
    pub y: Vec<Option<f64>>,
    /// Fractional z coordinates; None represents missing source values.
    pub z: Vec<Option<f64>>,
    /// Dimensionless site occupancies; an empty array means the field was absent.
    pub occupancy: Vec<Option<f64>>,
    /// Original record URL, when present.
    pub url: Option<String>,
    /// Source publication/title text, when present.
    pub publication: Option<String>,
}

impl AmcsdRecord {
    /// Re-create a CIF text like larixite's `CifStructure.ciftext`.
    pub fn to_cif(&self) -> String {
        let mut out = String::new();
        out.push_str(&format!("data_amcsd_{:07}\n", self.id));
        if let Some(m) = &self.mineral {
            out.push_str(&format!("_chemical_name_mineral '{m}'\n"));
        }
        out.push_str(&format!("_chemical_formula_sum '{}'\n", self.formula));
        if let Some(p) = &self.publication {
            out.push_str(&format!("_publ_section_title\n;\n{p}\n;\n"));
        }
        if let Some(u) = &self.url {
            out.push_str(&format!("_database_code_amcsd {}\n# {u}\n", self.id));
        }
        let c = self.cell;
        out.push_str(&format!(
            "_cell_length_a {}\n_cell_length_b {}\n_cell_length_c {}\n_cell_angle_alpha {}\n_cell_angle_beta {}\n_cell_angle_gamma {}\n",
            c[0], c[1], c[2], c[3], c[4], c[5]
        ));
        out.push_str(&format!(
            "_symmetry_space_group_name_H-M '{}'\nloop_\n_space_group_symop_operation_xyz\n",
            self.hm_symbol
        ));
        for op in &self.symmetry_xyz {
            out.push_str(&format!("'{op}'\n"));
        }
        let has_occ = !self.occupancy.is_empty();
        out.push_str(
            "loop_\n_atom_site_label\n_atom_site_fract_x\n_atom_site_fract_y\n_atom_site_fract_z\n",
        );
        if has_occ {
            out.push_str("_atom_site_occupancy\n");
        }
        for (i, label) in self.sites.iter().enumerate() {
            let f = |v: &Vec<Option<f64>>| match v.get(i).copied().flatten() {
                Some(x) => format!("{x:.6}"),
                None => "?".into(),
            };
            out.push_str(&format!(
                "{label} {} {} {}",
                f(&self.x),
                f(&self.y),
                f(&self.z)
            ));
            if has_occ {
                out.push_str(&format!(" {}", f(&self.occupancy)));
            }
            out.push('\n');
        }
        out
    }

    /// Create and expand an owned structure through the shared CIF reader.
    ///
    /// Coordinates are formatted to six decimal places by to_cif first. Source and
    /// mineral naming are retained. Missing/invalid cell or site data return errors;
    /// inspect structure warnings for skipped source values. No database is modified.
    pub fn to_structure(&self) -> Result<Structure, StructureError> {
        let mut s = structure_from_cif(&self.to_cif())?;
        s.source = format!("amcsd:{}", self.id);
        if let Some(m) = &self.mineral {
            s.title = m.clone();
        }
        Ok(s)
    }
}

impl Amcsd {
    /// Open an existing database read-only.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self, StructureError> {
        let path = path.as_ref().to_path_buf();
        let conn =
            Connection::open_with_flags(&path, OpenFlags::SQLITE_OPEN_READ_ONLY).map_err(db_err)?;
        let ok: i64 = conn
            .query_row(
                "select count(*) from sqlite_master where type='table' and name in ('cif','minerals','spacegroups')",
                [],
                |r| r.get(0),
            )
            .map_err(db_err)?;
        if ok != 3 {
            return Err(StructureError::Database {
                reason: format!("{} is not an AMCSD database", path.display()),
            });
        }
        Ok(Self { conn, path })
    }

    /// Borrow the path of the read-only SQLite database.
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Query the number of CIF records; database/query failures return an error.
    pub fn len(&self) -> Result<usize, StructureError> {
        self.conn
            .query_row("select count(*) from cif", [], |r| r.get::<_, i64>(0))
            .map(|n| n as usize)
            .map_err(db_err)
    }

    /// Query whether the database has no CIF records; propagates database errors.
    pub fn is_empty(&self) -> Result<bool, StructureError> {
        Ok(self.len()? == 0)
    }

    /// All mineral names.
    pub fn minerals(&self) -> Result<Vec<String>, StructureError> {
        let mut st = self
            .conn
            .prepare("select name from minerals order by name")
            .map_err(db_err)?;
        let rows = st
            .query_map([], |r| r.get::<_, String>(0))
            .map_err(db_err)?;
        rows.collect::<Result<Vec<_>, _>>().map_err(db_err)
    }

    fn elements_of(&self, cif_id: i64) -> Result<Vec<String>, StructureError> {
        let mut st = self
            .conn
            .prepare("select element from cif_elements where cif_id = ?1")
            .map_err(db_err)?;
        let rows = st
            .query_map([cif_id.to_string()], |r| r.get::<_, String>(0))
            .map_err(db_err)?;
        let mut out: Vec<String> = rows.collect::<Result<Vec<_>, _>>().map_err(db_err)?;
        out.sort();
        out.dedup();
        Ok(out)
    }

    /// Fetch one record by AMCSD id.
    pub fn record(&self, id: i64) -> Result<AmcsdRecord, StructureError> {
        let row = self
            .conn
            .query_row(
                "select c.id, m.name, c.formula, s.hm_notation, s.symmetry_xyz, c.a, c.b, c.c, c.alpha, c.beta, c.gamma, c.atoms_sites, c.atoms_x, c.atoms_y, c.atoms_z, c.atoms_occupancy, c.amcsd_url, c.pub_title \
                 from cif c left join minerals m on m.id = c.mineral_id left join spacegroups s on s.id = c.spacegroup_id where c.id = ?1",
                [id],
                |r| {
                    Ok((
                        r.get::<_, i64>(0)?,
                        r.get::<_, Option<String>>(1)?,
                        r.get::<_, Option<String>>(2)?,
                        r.get::<_, Option<String>>(3)?,
                        r.get::<_, Option<String>>(4)?,
                        [
                            r.get::<_, Option<String>>(5)?,
                            r.get::<_, Option<String>>(6)?,
                            r.get::<_, Option<String>>(7)?,
                            r.get::<_, Option<String>>(8)?,
                            r.get::<_, Option<String>>(9)?,
                            r.get::<_, Option<String>>(10)?,
                        ],
                        r.get::<_, Option<String>>(11)?,
                        r.get::<_, Option<String>>(12)?,
                        r.get::<_, Option<String>>(13)?,
                        r.get::<_, Option<String>>(14)?,
                        r.get::<_, Option<String>>(15)?,
                        r.get::<_, Option<String>>(16)?,
                        r.get::<_, Option<String>>(17)?,
                    ))
                },
            )
            .map_err(|e| StructureError::Database {
                reason: format!("AMCSD id {id}: {e}"),
            })?;
        let (id, mineral, formula, hm, symxyz, cell, sites, ax, ay, az, aocc, url, pub_title) = row;
        let cell_num = |v: &Option<String>| -> Result<f64, StructureError> {
            let text = v.clone().unwrap_or_default();
            let text = text
                .split('(')
                .next()
                .unwrap_or("")
                .trim()
                .replace(',', ".");
            text.parse::<f64>().map_err(|_| StructureError::Database {
                reason: format!("AMCSD id {id}: bad cell value '{text}'"),
            })
        };
        let cell = [
            cell_num(&cell[0])?,
            cell_num(&cell[1])?,
            cell_num(&cell[2])?,
            cell_num(&cell[3])?,
            cell_num(&cell[4])?,
            cell_num(&cell[5])?,
        ];
        let sites: Vec<String> = sites
            .as_deref()
            .and_then(|s| serde_json::from_str::<Vec<String>>(s).ok())
            .unwrap_or_default();
        let symmetry_xyz: Vec<String> = symxyz
            .as_deref()
            .and_then(|s| serde_json::from_str::<Vec<String>>(s).ok())
            .unwrap_or_else(|| vec!["x,y,z".into()]);
        let mineral = mineral.filter(|m| m != "<missing>" && !m.is_empty());
        Ok(AmcsdRecord {
            id,
            mineral,
            formula: formula.unwrap_or_default(),
            hm_symbol: hm.unwrap_or_else(|| "P 1".into()),
            symmetry_xyz,
            cell,
            sites,
            x: decode_farray(ax.as_deref().unwrap_or("0")),
            y: decode_farray(ay.as_deref().unwrap_or("0")),
            z: decode_farray(az.as_deref().unwrap_or("0")),
            occupancy: decode_farray(aocc.as_deref().unwrap_or("0")),
            url: url.filter(|u| !u.is_empty()),
            publication: pub_title.filter(|p| !p.is_empty() && p != "<missing>"),
        })
    }

    fn hit_for(
        &self,
        id: i64,
        mineral: Option<String>,
        formula: String,
        hm: Option<String>,
    ) -> Result<StructureHit, StructureError> {
        let mut extra = BTreeMap::new();
        extra.insert("amcsd_id".into(), id.to_string());
        Ok(StructureHit {
            id: id.to_string(),
            source: "amcsd".into(),
            formula: formula.split_whitespace().collect::<String>(),
            name: mineral.filter(|m| m != "<missing>" && !m.is_empty()),
            space_group: hm,
            elements: self.elements_of(id)?,
            extra,
        })
    }
}

impl StructureSource for Amcsd {
    fn name(&self) -> &str {
        "amcsd"
    }

    fn search(&self, query: &StructureQuery) -> Result<Vec<StructureHit>, StructureError> {
        let limit = if query.limit == 0 { 200 } else { query.limit };
        // Pre-filter in SQL by required elements and text, then apply the
        // full query on the hits.
        let mut sql = String::from(
            "select c.id, m.name, c.formula, s.hm_notation from cif c \
             left join minerals m on m.id = c.mineral_id \
             left join spacegroups s on s.id = c.spacegroup_id where 1=1",
        );
        let mut params: Vec<String> = Vec::new();
        for el in &query.elements {
            sql.push_str(" and c.id in (select cif_id from cif_elements where element = ?)");
            params.push(el.clone());
        }
        for el in &query.exclude {
            sql.push_str(" and c.id not in (select cif_id from cif_elements where element = ?)");
            params.push(el.clone());
        }
        if let Some(text) = query
            .text
            .as_deref()
            .map(str::trim)
            .filter(|t| !t.is_empty())
        {
            let compact: String = text.split_whitespace().collect();
            sql.push_str(" and (lower(m.name) like ? or lower(replace(c.formula,' ','')) like ? or c.id = ?)");
            params.push(format!("%{}%", text.to_ascii_lowercase()));
            params.push(format!("%{}%", compact.to_ascii_lowercase()));
            params.push(text.to_string());
        }
        sql.push_str(" order by m.name, c.id limit ?");
        params.push((limit * 4).to_string());
        let mut st = self.conn.prepare(&sql).map_err(db_err)?;
        let rows = st
            .query_map(rusqlite::params_from_iter(params.iter()), |r| {
                Ok((
                    r.get::<_, i64>(0)?,
                    r.get::<_, Option<String>>(1)?,
                    r.get::<_, Option<String>>(2)?,
                    r.get::<_, Option<String>>(3)?,
                ))
            })
            .map_err(db_err)?;
        let mut hits = Vec::new();
        for row in rows {
            let (id, mineral, formula, hm) = row.map_err(db_err)?;
            let hit = self.hit_for(id, mineral, formula.unwrap_or_default(), hm)?;
            // Text already matched in SQL (name/formula/id); re-check element
            // constraints only.
            let mut q = query.clone();
            q.text = None;
            if q.matches(&hit) {
                hits.push(hit);
                if hits.len() >= limit {
                    break;
                }
            }
        }
        Ok(hits)
    }

    fn fetch(&self, hit: &StructureHit) -> Result<Structure, StructureError> {
        let id: i64 = hit.id.parse().map_err(|_| StructureError::Database {
            reason: format!("bad AMCSD id '{}'", hit.id),
        })?;
        self.record(id)?.to_structure()
    }
}

/// Download the full AMCSD database to `dest` (a file path), trying each
/// mirror. `progress(received_bytes, total_bytes)` is called as data
/// arrives. Requires the `http` feature (implied by `materials-project` and `cod`).
#[cfg(feature = "http")]
pub fn download_amcsd<P: AsRef<Path>>(
    dest: P,
    progress: impl FnMut(u64, Option<u64>),
) -> Result<PathBuf, StructureError> {
    download_amcsd_cancellable(dest, progress, &std::sync::atomic::AtomicBool::new(false))
}

/// [`download_amcsd`] that stops early when `cancel` becomes `true`.
///
/// A cancelled download removes its partial file and returns
/// [`StructureError::Network`] with the reason `"cancelled"`. Requires the
/// `http` feature (also enabled by `materials-project` and `cod`).
#[cfg(feature = "http")]
pub fn download_amcsd_cancellable<P: AsRef<Path>>(
    dest: P,
    progress: impl FnMut(u64, Option<u64>),
    cancel: &std::sync::atomic::AtomicBool,
) -> Result<PathBuf, StructureError> {
    download_amcsd_with(dest, progress, cancel, |_| {})
}

/// [`download_amcsd_cancellable`] that also reports the mirror URL being
/// fetched through `on_source` (once per attempt), so a UI can name it.
///
/// Mirrors are tried in [`SOURCE_URLS`] order; a mirror that answers with a
/// non-200 status or with a file that is not a valid AMCSD database is
/// skipped and the next one is tried.
#[cfg(feature = "http")]
pub fn download_amcsd_with<P: AsRef<Path>>(
    dest: P,
    mut progress: impl FnMut(u64, Option<u64>),
    cancel: &std::sync::atomic::AtomicBool,
    mut on_source: impl FnMut(&str),
) -> Result<PathBuf, StructureError> {
    use std::sync::atomic::Ordering;
    let dest = dest.as_ref().to_path_buf();
    let mut errors: Vec<String> = Vec::new();
    for base in SOURCE_URLS {
        let url = mirror_url(base);
        on_source(&url);
        let agent: ureq::Agent = ureq::Agent::config_builder()
            .user_agent(concat!(
                "rexafs/",
                env!("CARGO_PKG_VERSION"),
                " (+https://github.com/Ameyanagi/rexafs)"
            ))
            .build()
            .into();
        let mut resp = match agent.get(&url).call() {
            Ok(resp) => resp,
            Err(e) => {
                errors.push(format!("{url}: {e}"));
                continue;
            }
        };
        if resp.status() != 200 {
            errors.push(format!("{url}: HTTP {}", resp.status()));
            continue;
        }
        let total = resp
            .headers()
            .get("content-length")
            .and_then(|v| v.to_str().ok())
            .and_then(|v| v.parse::<u64>().ok());
        let tmp = dest.with_extension("part");
        let mut file = std::fs::File::create(&tmp).map_err(|source| StructureError::Io {
            path: tmp.display().to_string(),
            source,
        })?;
        let mut reader = resp.body_mut().as_reader();
        let mut buf = vec![0u8; 1 << 16];
        let mut received = 0u64;
        let mut failed = None;
        loop {
            let n = match reader.read(&mut buf) {
                Ok(n) => n,
                Err(e) => {
                    failed = Some(format!("{url}: {e}"));
                    break;
                }
            };
            if n == 0 {
                break;
            }
            file.write_all(&buf[..n])
                .map_err(|source| StructureError::Io {
                    path: tmp.display().to_string(),
                    source,
                })?;
            received += n as u64;
            progress(received, total);
            if cancel.load(Ordering::Relaxed) {
                drop(file);
                let _ = std::fs::remove_file(&tmp);
                return Err(StructureError::Network {
                    reason: "cancelled".into(),
                });
            }
        }
        drop(file);
        if let Some(e) = failed {
            let _ = std::fs::remove_file(&tmp);
            errors.push(e);
            continue;
        }
        // Validate before declaring success; an HTML error page or an
        // empty body from a deferred download must not replace `dest`.
        if let Err(e) = Amcsd::open(&tmp) {
            let _ = std::fs::remove_file(&tmp);
            errors.push(format!("{url}: not a valid AMCSD database ({e})"));
            continue;
        }
        std::fs::rename(&tmp, &dest).map_err(|source| StructureError::Io {
            path: dest.display().to_string(),
            source,
        })?;
        return Ok(dest);
    }
    Err(StructureError::Network {
        reason: if errors.is_empty() {
            "no mirrors".into()
        } else {
            errors.join("; ")
        },
    })
}