readcon-core 0.14.0

#include "readcon-core.hpp"
#include <iostream>
#include <vector>

// A helper function to print the frame's data in detail.
void print_frame_details(int frame_number, const readcon::ConFrame &frame) {
    std::cout << "\n==================== FRAME " << frame_number
              << " ====================\n";

    // Print header information using the accessor methods
    auto pre_headers = frame.prebox_header();
    std::cout << "Pre-box Header 1: \"" << pre_headers[0] << "\"\n";
    std::cout << "Pre-box Header 2: \"" << pre_headers[1] << "\"\n";

    // Print cell information
    auto cell = frame.cell();
    auto angles = frame.angles();
    std::cout << "Cell Dimensions:  " << cell[0] << ", " << cell[1] << ", "
              << cell[2] << std::endl;
    std::cout << "Cell Angles:      " << angles[0] << ", " << angles[1] << ", "
              << angles[2] << std::endl;

    auto post_headers = frame.postbox_header();
    std::cout << "Post-box Header 1:\"" << post_headers[0] << "\"\n";
    std::cout << "Post-box Header 2:\"" << post_headers[1] << "\"\n";

    // Print velocity status
    std::cout << "Has Velocities:   " << std::boolalpha << frame.has_velocities()
              << std::endl;

    // Print atom information
    auto atoms = frame.atoms(); // Call once to avoid repeated conversions
    std::cout << "--- Atoms (" << atoms.size() << ") ---\n";

    // Print details for the first 5 atoms for brevity
    int atoms_to_print = 0;
    for (const auto &atom : atoms) {
        if (atoms_to_print >= 5) {
            std::cout << "... and " << atoms.size() - 5 << " more."
                      << std::endl;
            break;
        }
        const auto fixed = atom.fixed_mask();
        std::cout << "  ID: " << atom.atom_id << ", Z: " << atom.atomic_number
                  << ", Pos: (" << atom.x << ", " << atom.y << ", " << atom.z
                  << ")"
                  << ", Fixed: (" << fixed[0] << ", " << fixed[1] << ", "
                  << fixed[2] << ")";
        if (auto v = atom.velocity()) {
            std::cout << ", Vel: (" << (*v)[0] << ", " << (*v)[1] << ", "
                      << (*v)[2] << ")";
        }
        std::cout << std::endl;
        atoms_to_print++;
    }
}

// Section buffer / DLPack smoke (readcon-core-8cx2): no AoS required.
static int smoke_section_buffers(const readcon::ConFrame &frame) {
    using readcon::RKRStatus;
    using readcon::RKR_STATUS_SUCCESS;
    using readcon::RKR_STATUS_SECTION_ABSENT;
    const std::size_t n = frame.atom_count();
    if (n == 0) {
        std::cerr << "section smoke: atom_count == 0\n";
        return 1;
    }
    std::vector<double> pos(n * 3);
    RKRStatus st = frame.copy_positions(pos.data(), pos.size());
    if (st != RKR_STATUS_SUCCESS) {
        std::cerr << "copy_positions failed: " << static_cast<int>(st) << "\n";
        return 1;
    }
    std::vector<double> vel(n * 3);
    st = frame.copy_velocities(vel.data(), vel.size());
    if (st != RKR_STATUS_SUCCESS && st != RKR_STATUS_SECTION_ABSENT) {
        std::cerr << "copy_velocities unexpected: " << static_cast<int>(st) << "\n";
        return 1;
    }
    std::vector<double> frc(n * 3);
    st = frame.copy_forces(frc.data(), frc.size());
    if (st != RKR_STATUS_SUCCESS && st != RKR_STATUS_SECTION_ABSENT) {
        std::cerr << "copy_forces unexpected: " << static_cast<int>(st) << "\n";
        return 1;
    }
    std::vector<double> mass(n);
    st = frame.copy_masses(mass.data(), mass.size());
    if (st != RKR_STATUS_SUCCESS && st != RKR_STATUS_SECTION_ABSENT) {
        std::cerr << "copy_masses unexpected: " << static_cast<int>(st) << "\n";
        return 1;
    }
    std::cout << "section_buffers ok natoms=" << n
              << " pos0=(" << pos[0] << "," << pos[1] << "," << pos[2] << ")\n";
    return 0;
}

int main(int argc, char *argv[]) {
    if (argc < 2 || argc > 3) {
        std::cerr << "Usage: " << argv[0] << " <input.con> [output.con]"
                  << std::endl;
        return 1;
    }

    std::string input_filename = argv[1];

    try {
        // --- Read-only and Summarize Mode (Memory-Efficient) ---
        if (argc == 2) {
            std::cout << "Mode: Read-only. Iterating lazily through frames in: "
                      << input_filename << std::endl;
            readcon::ConFrameIterator frame_iterator(input_filename);

            int frame_count = 0;
            // This loop is memory-efficient. It processes one frame at a time
            // without storing them all in memory.
            for (const auto &frame : frame_iterator) {
                frame_count++;
                print_frame_details(frame_count, frame);
                if (frame_count == 1) {
                    if (smoke_section_buffers(frame) != 0) {
                        return 2;
                    }
                }
            }
            std::cout
                << "\n==================================================\n";
            std::cout << "Iteration complete. Total frames processed: "
                      << frame_count << std::endl;
        }
        // --- Read and Write Mode ---
        else { // argc == 3
            std::string output_filename = argv[2];
            std::cout << "Mode: Read-Write. Reading from '" << input_filename
                      << "' and writing to '" << output_filename << "'."
                      << std::endl;

            readcon::ConFrameIterator frame_iterator(input_filename);

            // In write mode, we must collect all frames first.
            std::vector<readcon::ConFrame> all_frames;
            // Use 'auto &&frame' to enable move semantics: ConFrameIterator
            // yields rvalue references, so moves are without copy
            for (auto &&frame : frame_iterator) {
                // We must move the frame from the iterator into our vector.
                all_frames.push_back(std::move(frame));
            }

            if (all_frames.empty()) {
                std::cout << "No valid frames found to write." << std::endl;
            } else {
                print_frame_details(all_frames.size(), all_frames.back());
                std::cout << "\nWriting " << all_frames.size()
                          << " frames...\n";

                // Use the new, ergonomic ConFrameWriter object.
                // RAII ensures the file is properly closed when the writer goes
                // out of scope.
                readcon::ConFrameWriter writer(output_filename);
                writer.extend(all_frames);

                std::cout << "Successfully wrote all frames." << std::endl;
            }
        }

    } catch (const std::exception &e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }

    return 0;
}