cp2k-rs 0.1.0

//! IPC protocol types for the CP2K worker process.
//!
//! Serialization uses bincode for compact binary encoding.
//! All messages are length-prefixed: [u32 length][payload bytes].

use serde::{Deserialize, Serialize};

/// A command sent from the Python frontend to the worker.
#[derive(Debug, Serialize, Deserialize)]
pub enum Command {
    /// Initialize a new force environment.
    InitForceEnv { input: String, output: String },
    /// Calculate energy and forces for the current geometry.
    CalcEnergyForce,
    /// Calculate energy only.
    CalcEnergy,
    /// Query the number of atoms.
    GetNatom,
    /// Query the number of particles.
    GetNparticle,
    /// Retrieve particle positions as a flat f64 array.
    GetPositions,
    /// Retrieve forces as a flat f64 array.
    GetForces,
    /// Retrieve the potential energy.
    GetPotentialEnergy,
    /// Retrieve the 3x3 simulation cell (row-major flat f64).
    GetCell,
    /// Retrieve the 3x3 QMMM cell.
    GetQmmmCell,
    /// Set particle positions.
    SetPositions { data: Vec<f64> },
    /// Set particle velocities.
    SetVelocities { data: Vec<f64> },
    /// Set simulation cell (9 floats, row-major).
    SetCell { data: Vec<f64> },
    /// Get number of MOs in the active space.
    GetMoCount,
    // ----- extended interface -----
    #[cfg(feature = "extended")]
    IsQuickstep,
    #[cfg(feature = "extended")]
    GetStressTensor,
    #[cfg(feature = "extended")]
    GetVirialTensor,
    #[cfg(feature = "extended")]
    GetNmo { spin: i32 },
    #[cfg(feature = "extended")]
    GetEigenvalues { spin: i32 },
    #[cfg(feature = "extended")]
    GetOccupationNumbers { spin: i32 },
    #[cfg(feature = "extended")]
    GetHomoLumo { spin: i32 },
    #[cfg(feature = "extended")]
    GetMullikenCharges,
    #[cfg(feature = "extended")]
    GetHirshfeldCharges,
    #[cfg(feature = "extended")]
    GetDipoleMoment,
    #[cfg(feature = "extended")]
    GetScfInfo,
    #[cfg(feature = "extended")]
    GetEnergyComponents,
    #[cfg(feature = "extended")]
    GetNelectron,
    #[cfg(feature = "extended")]
    GetFermiEnergy,
    #[cfg(feature = "extended")]
    GetTotalSpin,
    /// Ask the worker to shut down cleanly.
    Shutdown,
    /// Get CP2K version string.
    GetVersion,
}

/// A request wrapper carrying a unique ID and a command.
#[derive(Debug, Serialize, Deserialize)]
pub struct Request {
    pub request_id: u64,
    pub command: Command,
}

/// Status of a response.
#[derive(Debug, Serialize, Deserialize)]
pub enum Status {
    Ok,
    Error(String),
}

/// The payload of a successful response.
#[derive(Debug, Serialize, Deserialize)]
pub enum Payload {
    /// No data to return (void operations).
    Empty,
    /// A single integer value.
    Int(i64),
    /// A single unsigned integer value.
    UInt(u64),
    /// A single float value.
    Float(f64),
    /// A boolean value.
    Bool(bool),
    /// A 1-D f64 array (flat).
    Array1(Vec<f64>),
    /// A 2-D f64 array (shape[0], shape[1], flat row-major data).
    Array2 {
        rows: usize,
        cols: usize,
        data: Vec<f64>,
    },
    /// A string value.
    String(String),
    /// Four values: two floats and two ints (HOMO/LUMO).
    HomoLumo {
        homo: f64,
        lumo: f64,
        homo_idx: i32,
        lumo_idx: i32,
    },
    /// Five floats (energy components).
    EnergyComponents {
        e_kin: f64,
        e_hartree: f64,
        e_xc: f64,
        e_core: f64,
        e_total: f64,
    },
    /// SCF info: steps, converged, energy_change.
    ScfInfo {
        nsteps: i32,
        converged: bool,
        energy_change: f64,
    },
}

/// A response from the worker back to the Python frontend.
#[derive(Debug, Serialize, Deserialize)]
pub struct Response {
    pub request_id: u64,
    pub status: Status,
    pub payload: Payload,
}

impl Response {
    pub fn ok(request_id: u64, payload: Payload) -> Self {
        Response {
            request_id,
            status: Status::Ok,
            payload,
        }
    }

    pub fn error(request_id: u64, msg: impl Into<String>) -> Self {
        Response {
            request_id,
            status: Status::Error(msg.into()),
            payload: Payload::Empty,
        }
    }
}