readcon-core 0.10.0

mod common;
use readcon_core::iterators::ConFrameIterator;
use readcon_core::types::ConFrameBuilder;
use readcon_core::writer::ConFrameWriter;
use std::collections::BTreeMap;
use std::fs;
use std::path::Path;

#[test]
fn test_writer_roundtrip() {
    let fdat_original =
        fs::read_to_string(test_case!("tiny_multi_cuh2.con")).expect("Can't find test file.");
    let parser_original = ConFrameIterator::new(&fdat_original);
    let frames_original: Vec<_> = parser_original.map(|r| r.unwrap()).collect();

    // Ensure we have something to test.
    assert!(!frames_original.is_empty());

    let mut buffer: Vec<u8> = Vec::new();

    // Introduce a scope so the writer is dropped and flushes any buffered data, releasing the mutable reference to the buffer.
    {
        let mut writer = ConFrameWriter::new(&mut buffer);
        writer
            .extend(frames_original.iter())
            .expect("Failed to write to buffer.");
    } // `writer` is dropped here, ensuring all data is flushed and the mutable reference ends.

    // Convert the buffer back to a string and re-parse.
    // This move is now valid because the borrow has ended.
    let fdat_roundtrip = String::from_utf8(buffer).expect("Buffer is not valid UTF-8.");
    let parser_roundtrip = ConFrameIterator::new(&fdat_roundtrip);
    let frames_roundtrip: Vec<_> = parser_roundtrip.map(|r| r.unwrap()).collect();

    assert_eq!(
        frames_original.len(),
        frames_roundtrip.len(),
        "Number of frames should be the same after roundtrip."
    );
    assert_eq!(
        frames_original, frames_roundtrip,
        "Frame data should be identical after a read-write-read roundtrip."
    );
}

#[test]
fn test_builder_roundtrip() {
    let mut builder =
        ConFrameBuilder::new([15.345600, 21.702000, 100.000000], [90.0, 90.0, 90.0]);
    builder
        .prebox_header("Random Number Seed")
        .postbox_header(["0 0".to_string(), "218 0 1".to_string()]);
    builder.add_atom(
        "Cu",
        0.639400000000001,
        0.904500000000000,
        6.975299999999995,
        [true, true, true],
        0,
        63.546,
    );
    builder.add_atom(
        "Cu",
        3.196999999999999,
        0.904500000000000,
        6.975299999999995,
        [true, true, true],
        1,
        63.546,
    );
    builder.add_atom(
        "H",
        8.682_3,
        9.946999999999997,
        11.732999999999993,
        [false, false, false],
        2,
        1.008,
    );
    let frame = builder.build();

    let mut buffer: Vec<u8> = Vec::new();
    {
        let mut writer = ConFrameWriter::with_precision(&mut buffer, 17);
        writer.write_frame(&frame).expect("Failed to write frame.");
    }

    let fdat = String::from_utf8(buffer).expect("Buffer is not valid UTF-8.");
    let parser = ConFrameIterator::new(&fdat);
    let frames_roundtrip: Vec<_> = parser.map(|r| r.unwrap()).collect();
    assert_eq!(frames_roundtrip.len(), 1);

    let rt = &frames_roundtrip[0];
    assert_eq!(rt.header.natm_types, 2);
    assert_eq!(rt.header.natms_per_type, vec![2, 1]);
    // Verify precision 17 roundtrip preserves coordinates
    assert_eq!(rt.atom_data[0].x, 0.639400000000001);
    assert_eq!(rt.atom_data[0].y, 0.904500000000000);
    assert_eq!(rt.atom_data[2].x, 8.682_3);
}

#[test]
fn test_writer_precision_default_vs_high() {
    let mut builder = ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
    builder.add_atom(
        "Cu",
        1.234_567_890_123_456_7,
        0.0,
        0.0,
        [false, false, false],
        0,
        63.546,
    );
    let frame = builder.build();

    // Default precision (6)
    let mut buf6: Vec<u8> = Vec::new();
    {
        let mut w = ConFrameWriter::new(&mut buf6);
        w.write_frame(&frame).unwrap();
    }
    let s6 = String::from_utf8(buf6).unwrap();
    let frames6: Vec<_> = ConFrameIterator::new(&s6).map(|r| r.unwrap()).collect();
    // 6 decimal places means ~1e-6 precision loss
    assert!((frames6[0].atom_data[0].x - 1.234568).abs() < 1e-5);

    // High precision (17)
    let mut buf17: Vec<u8> = Vec::new();
    {
        let mut w = ConFrameWriter::with_precision(&mut buf17, 17);
        w.write_frame(&frame).unwrap();
    }
    let s17 = String::from_utf8(buf17).unwrap();
    let frames17: Vec<_> = ConFrameIterator::new(&s17).map(|r| r.unwrap()).collect();
    // 17 decimal places preserves the full f64
    assert!((frames17[0].atom_data[0].x - 1.234_567_890_123_456_7).abs() < 1e-14);
}

#[test]
fn test_builder_velocity_roundtrip() {
    let mut builder = ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
    builder
        .add_atom("Cu", 1.0, 2.0, 3.0, [true, true, true], 0, 63.546)
        .with_velocity([0.1, 0.2, 0.3]);
    builder
        .add_atom("H", 4.0, 5.0, 6.0, [false, false, false], 1, 1.008)
        .with_velocity([0.4, 0.5, 0.6]);
    let frame = builder.build();

    assert!(frame.has_velocities());

    let mut buffer: Vec<u8> = Vec::new();
    {
        let mut writer = ConFrameWriter::new(&mut buffer);
        writer.write_frame(&frame).expect("Failed to write frame.");
    }

    let fdat = String::from_utf8(buffer).expect("Buffer is not valid UTF-8.");
    let parser = ConFrameIterator::new(&fdat);
    let frames: Vec<_> = parser.map(|r| r.unwrap()).collect();
    assert_eq!(frames.len(), 1);
    assert!(frames[0].has_velocities());
    assert_eq!(frames[0].atom_data[0].velocity, Some([0.1, 0.2, 0.3]));
    assert_eq!(frames[0].atom_data[1].velocity, Some([0.4, 0.5, 0.6]));
}

#[test]
fn test_writer_emits_empty_sections_when_validate_true() {
    let mut metadata = BTreeMap::new();
    metadata.insert("validate".to_string(), serde_json::json!(true));
    let mut builder = ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
    builder.metadata(metadata);
    builder.add_atom("Cu", 1.0, 2.0, 3.0, [false, false, false], 0, 63.546);
    let frame = builder.build();

    let mut buffer: Vec<u8> = Vec::new();
    {
        let mut writer = ConFrameWriter::new(&mut buffer);
        writer.write_frame(&frame).expect("Failed to write frame.");
    }

    let fdat = String::from_utf8(buffer).expect("Buffer is not valid UTF-8.");
    let metadata: serde_json::Value = serde_json::from_str(fdat.lines().nth(1).unwrap()).unwrap();

    assert_eq!(metadata.get("sections"), Some(&serde_json::json!([])));

    let frames: Vec<_> = ConFrameIterator::new(&fdat)
        .map(|result| result.unwrap())
        .collect();
    assert_eq!(frames.len(), 1);
}

/// Verifies that non-sequential atom_id values (representing the original
/// atom index before type-based reordering) survive a write-read roundtrip.
///
/// This is the key scenario from eOn commit 8b8d929: atoms like
/// C(0),C(1),C(2),O(3),C(4),C(5) get reordered to C(0),C(1),C(2),C(4),C(5),O(3)
/// in the .con file, and column 5 must preserve the original indices.
#[test]
fn test_nonsequential_atom_index_roundtrip() {
    let mut builder = ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
    // Simulate atoms added in original order: C,C,C,O,C,C
    // The builder groups by symbol, so C atoms come first, then O.
    // We assign atom_id matching the original pre-grouping index.
    builder.add_atom("C", 0.0, 0.0, 0.0, [false, false, false], 0, 12.011);
    builder.add_atom("C", 1.0, 0.0, 0.0, [false, false, false], 1, 12.011);
    builder.add_atom("C", 2.0, 0.0, 0.0, [false, false, false], 2, 12.011);
    builder.add_atom("O", 3.0, 0.0, 0.0, [false, false, false], 3, 15.999);
    builder.add_atom("C", 4.0, 0.0, 0.0, [false, false, false], 4, 12.011);
    builder.add_atom("C", 5.0, 0.0, 0.0, [false, false, false], 5, 12.011);
    let frame = builder.build();

    // After build(), atoms are grouped: C(0),C(1),C(2),C(4),C(5),O(3)
    assert_eq!(frame.header.natm_types, 2);
    assert_eq!(frame.header.natms_per_type, vec![5, 1]);
    assert_eq!(frame.atom_data[3].atom_id, 4); // C at x=4, original idx 4
    assert_eq!(frame.atom_data[4].atom_id, 5); // C at x=5, original idx 5
    assert_eq!(frame.atom_data[5].atom_id, 3); // O at x=3, original idx 3

    // Write and re-read
    let mut buffer: Vec<u8> = Vec::new();
    {
        let mut writer = ConFrameWriter::new(&mut buffer);
        writer.write_frame(&frame).expect("Failed to write frame.");
    }
    let fdat = String::from_utf8(buffer).expect("Buffer is not valid UTF-8.");
    let parser = ConFrameIterator::new(&fdat);
    let frames: Vec<_> = parser.map(|r| r.unwrap()).collect();
    assert_eq!(frames.len(), 1);

    let rt = &frames[0];
    // Non-sequential atom_id values must survive the roundtrip
    assert_eq!(rt.atom_data[3].atom_id, 4);
    assert_eq!(rt.atom_data[4].atom_id, 5);
    assert_eq!(rt.atom_data[5].atom_id, 3);
}