cp2k-rs 0.2.3

//! 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};

// K-Point configuration types
// ─── CP2K input specification ────────────────────────────────────────────────

/// How the CP2K worker should obtain the `.inp` file.
///
/// This is the "source of truth" input mechanism: either the user supplies a
/// path (`FromFile`) or a typed settings model is rendered into an input file
/// (`Generated`).
///
/// Note: The worker ultimately needs a file path for CP2K. For `Generated`,
/// the worker is expected to render settings to a temporary `.inp` file.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[allow(clippy::large_enum_variant)]
pub enum Cp2kInputSpec {
    /// Use an existing CP2K input file on disk.
    FromFile { path: String },
    /// Generate an input file from typed settings.
    Generated { settings: Cp2kSettings },
}

/// A minimal, typed CP2K settings model.
///
/// This is intentionally scoped to the most common Quickstep/DFT parameters we
/// want to support without requiring a handwritten `.inp` file.
///
/// You can always extend this model over time.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kSettings {
    /// Settings for the `&GLOBAL` section.
    pub global: Option<Cp2kGlobalSettings>,
    /// Settings for `&FORCE_EVAL` (only the common `METHOD Quickstep` path is
    /// addressed initially).
    pub force_eval: Cp2kForceEvalSettings,
}

/// Settings for `&GLOBAL`.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kGlobalSettings {
    /// `PROJECT` label.
    pub project: Option<String>,
    /// `PRINT_LEVEL`
    pub print_level: Option<String>,
    /// `RUN_TYPE`
    pub run_type: Option<String>,
    /// `WALLTIME` (CP2K string, e.g. "00:59:00")
    pub walltime: Option<String>,
    /// `SEED` values. CP2K accepts either one integer or multiple integers.
    pub seed: Option<Vec<i32>>,
    /// `ECHO_INPUT`
    pub echo_input: Option<bool>,
}

/// Settings for a `&FORCE_EVAL` section, focusing on Quickstep/DFT.
///
/// `&SUBSYS` geometry (`&CELL` and `&COORD`) is always injected from the frontend
/// geometry and is therefore not part of this settings struct.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Cp2kForceEvalSettings {
    /// `METHOD` (commonly "Quickstep").
    pub method: String,
    /// Optional stress tensor request (e.g. "ANALYTICAL").
    pub stress_tensor: Option<String>,
    /// DFT settings (for `METHOD Quickstep`).
    pub dft: Option<Cp2kDftSettings>,
    /// Subsystem settings (e.g. `&KIND` blocks).
    pub subsys: Option<Cp2kSubsysSettings>,
}

impl Default for Cp2kForceEvalSettings {
    fn default() -> Self {
        Self {
            method: "Quickstep".to_string(),
            stress_tensor: None,
            dft: Some(Cp2kDftSettings::default()),
            subsys: None,
        }
    }
}

/// Settings for `&SUBSYS` (excluding geometry injection).
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kSubsysSettings {
    /// `&KIND` blocks for element-specific basis/potential.
    pub kinds: Vec<Cp2kKindSettings>,
    /// Extra subsections to be appended under `&SUBSYS`.
    pub extra_sections: Vec<Cp2kSection>,
    /// Extra keywords to be appended under `&SUBSYS`.
    pub extra_keywords: Vec<Cp2kKeyword>,
}

/// Settings for a `&KIND` block (e.g. `&KIND Si`).
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kKindSettings {
    /// Element symbol, used as the `&KIND` parameter.
    pub element: String,
    /// `BASIS_SET` keyword.
    pub basis_set: Option<String>,
    /// `POTENTIAL` keyword.
    pub potential: Option<String>,
    /// Extra keywords inside the `&KIND` block.
    pub extra_keywords: Vec<Cp2kKeyword>,
    /// Extra subsections inside the `&KIND` block.
    pub extra_sections: Vec<Cp2kSection>,
}

/// Settings for `&DFT` (Quickstep).
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kDftSettings {
    pub basis_set_file_name: Option<String>,
    pub potential_file_name: Option<String>,

    pub qs: Option<Cp2kQsSettings>,
    pub mgrid: Option<Cp2kMgridSettings>,
    pub scf: Option<Cp2kScfSettings>,
    pub xc: Option<Cp2kXcSettings>,

    /// DFT-level settings that are often needed but not yet modelled.
    ///
    /// This is a structured escape hatch that can be extended later to a proper
    /// typed "overlay" model if needed.
    pub extra_keywords: Vec<Cp2kKeyword>,
    /// Extra subsections to be appended under `&DFT`.
    pub extra_sections: Vec<Cp2kSection>,
}

/// Settings for `&QS`.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kQsSettings {
    pub eps_default: Option<f64>,
}

/// Settings for `&MGRID`.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kMgridSettings {
    pub cutoff: Option<f64>,
    pub rel_cutoff: Option<f64>,
}

/// Settings for `&SCF`.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kScfSettings {
    pub scf_guess: Option<String>,
    pub max_scf: Option<i32>,
    pub eps_scf: Option<f64>,
    pub added_mos: Option<i32>,

    pub smear: Option<Cp2kSmearSettings>,
    pub diagonalization: Option<Cp2kDiagonalizationSettings>,
    pub mixing: Option<Cp2kMixingSettings>,

    pub extra_keywords: Vec<Cp2kKeyword>,
    pub extra_sections: Vec<Cp2kSection>,
}

/// Settings for `&SMEAR`.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kSmearSettings {
    pub enabled: bool,
    pub method: Option<String>,
    pub electronic_temperature_k: Option<f64>,
}

/// Settings for `&DIAGONALIZATION`.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kDiagonalizationSettings {
    pub enabled: bool,
    pub algorithm: Option<String>,
}

/// Settings for `&MIXING`.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kMixingSettings {
    pub method: Option<String>,
    pub alpha: Option<f64>,
    pub nbroyden: Option<i32>,
}

/// Settings for `&XC`.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kXcSettings {
    /// E.g. "PBE"
    pub functional: Option<String>,

    pub extra_keywords: Vec<Cp2kKeyword>,
    pub extra_sections: Vec<Cp2kSection>,
}

/// A generic keyword representation for future structured extensions.
///
/// This is not intended to be the primary user-facing interface; it provides a
/// controlled escape hatch while keeping `Cp2kSettings` the source of truth.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Cp2kKeyword {
    pub name: String,
    pub value: String,
    /// Optional unit token, e.g. "K" or "Ry". The renderer is responsible for
    /// formatting it as `NAME [unit] value` where appropriate.
    pub unit: Option<String>,
}

/// A generic section representation for future structured extensions.
///
/// This is a controlled escape hatch; note that CP2K allows repeated sections
/// and section parameters (`&KIND Si`).
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Cp2kSection {
    /// Section name, e.g. "KIND" or "PRINT".
    pub name: String,
    /// Optional section parameter, e.g. `Si` in `&KIND Si`.
    pub parameter: Option<String>,
    pub keywords: Vec<Cp2kKeyword>,
    pub sections: Vec<Cp2kSection>,
}

// ─── K-Point configuration types ─────────────────────────────────────────────

/// Specifies how k-points are generated
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum KPointScheme {
    /// No k-points (isolated molecule)
    None,
    /// Gamma point only
    Gamma,
    /// Monkhorst-Pack grid with specified divisions in each direction
    /// (n1, n2, n3)
    MonkhorstPack(u32, u32, u32),
    /// Explicit list of k-points: each entry is [kx, ky, kz, weight]
    ExplicitList(Vec<[f64; 4]>),
}

/// Configuration for k-points and related settings
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KPointConfig {
    /// The k-point scheme
    pub scheme: KPointScheme,
    /// Whether to use symmetry reduction of k-points
    pub use_symmetry: bool,
    /// Verbosity for k-point information
    pub verbose: bool,
    /// Whether to use full grid (no symmetry reduction)
    pub full_grid: bool,
}

/// Configuration for DOS calculations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DosConfig {
    /// Whether to compute electronic density of states
    pub enabled: bool,
    /// Whether to compute partial DOS (per atom/orbital)
    pub partial_dos: bool,
    /// Energy range (min, max) in eV
    pub erange_ev: Option<(f64, f64)>,
    /// Number of energy points for DOS grid
    pub npoints: Option<usize>,
}

/// A command sent from the Python frontend to the worker.
#[derive(Debug, Serialize, Deserialize)]
#[allow(clippy::large_enum_variant)]
pub enum Command {
    /// Initialize a new force environment.
    ///
    /// This is the legacy entry point that takes a path to a `.inp` file.
    /// Prefer [`InitForceEnvWithGeometry`], which supports geometry injection.
    InitForceEnv { input: String, output: String },

    /// Initialise a new force environment and supply the starting geometry.
    ///
    /// The worker constructs a temporary input by taking the input specified by
    /// `input_spec` and injecting `&CELL` and `&COORD` sections into `&SUBSYS`
    /// before passing the (temporary) file to CP2K.
    ///
    /// For `input_spec = FromFile`, the file is read and patched.
    /// For `input_spec = Generated`, the settings are rendered and then patched
    /// (patching is still used to guarantee geometry correctness).
    ///
    /// Fields
    /// ------
    /// input_spec      Input source (file or generated settings).
    /// output          Path for CP2K output.
    /// symbols         Element symbol for each atom, e.g. `["W", "Al", "O"]`.
    /// positions_angstrom Flat interleaved array [x0,y0,z0, x1,y1,z1, …] in Å.
    /// cell_angstrom       9-element row-major cell matrix in Å.
    ///                     cell_angstrom[i*3 + j] = j-th component of i-th lattice vector.
    /// periodic            Periodicity string passed to CP2K's PERIODIC keyword.
    ///                     Use `"XYZ"` for fully periodic systems (the common case),
    ///                     `"NONE"` for isolated molecules. Any valid CP2K PERIODIC
    ///                     token is accepted.
    InitForceEnvWithGeometry {
        input: Cp2kInputSpec,
        output: String,
        symbols: Vec<String>,
        positions_angstrom: Vec<f64>,
        cell_angstrom: Vec<f64>,
        periodic: String,
        /// Optional project label for the &GLOBAL block
        project_label: Option<String>,
        /// Optional print level for the &GLOBAL block
        print_level: Option<String>,
        /// Optional RUN_TYPE for the &GLOBAL block
        run_type: Option<String>,
        /// Optional CHARGE for the &DFT block
        charge: Option<i32>,
        /// Optional MULTIPLICITY for the &DFT block
        multiplicity: Option<i32>,
        /// Optional WALLTIME for the &GLOBAL block
        walltime: Option<String>,
        /// Optional SEED for the &GLOBAL block
        seed: Option<Vec<i32>>,
        /// Optional ECHO_INPUT for the &GLOBAL block
        echo_input: Option<bool>,
        /// Optional UKS (Unrestricted Kohn-Sham) for the &DFT block
        uks: Option<bool>,
        /// Optional ROKS (Restricted Open Kohn-Sham) for the &DFT block
        roks: Option<bool>,
        /// Optional WFN_RESTART_FILE_NAME for the &DFT block
        wfn_restart_file: Option<String>,
        /// Optional k-point configuration for periodic systems
        kpoint_config: Option<KPointConfig>,
        /// Optional DOS configuration
        dos_config: Option<DosConfig>,
    },
    /// 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,
    #[cfg(feature = "extended")]
    GetGridInfo { spin: i32 },
    #[cfg(feature = "extended")]
    GetElectronDensity { spin: i32 },
    #[cfg(feature = "extended")]
    GetMoCoeffInfo { spin: i32 },
    #[cfg(feature = "extended")]
    GetMoCoefficients { spin: i32 },
    /// Get the number of k-points (0 = Gamma-point-only).
    #[cfg(feature = "extended")]
    GetNkpoints,
    /// Get Kohn-Sham eigenvalues for one k-point and one spin channel.
    #[cfg(feature = "extended")]
    GetKpointEigenvalues { kpt_idx: i32, spin: i32 },
    /// 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,
    },
    /// Grid metadata for the electron density.
    GridInfo {
        npts: [i32; 3],
        origin: [f64; 3],
        dh: [[f64; 3]; 3],
    },
    /// Large array data stored in POSIX shared memory (3D).
    SharedArray3 {
        shm_name: String,
        dims: [usize; 3],
        byte_size: usize,
    },
    /// Large array data stored in POSIX shared memory (2D).
    SharedArray2 {
        shm_name: String,
        rows: usize,
        cols: usize,
        byte_size: usize,
    },
    /// MO coefficient dimensions.
    MoCoeffInfo { nao: usize, nmo: usize },
}

/// 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,
        }
    }
}