refeff-io 0.2.0

//! Core-hole section adapters for FEFF `apot.bin`.

use ndarray::{Array1, ArrayView1};
use refeff_core::atomic::{
    AtomicFourPointCoulombPotentialInput, atomic_four_point_coulomb_potential,
};

use crate::error::{IoError, Result};

use super::types::{ApotBinData, ApotBinPayload, ApotBinRecords, ApotBinType, ApotBinValue};

/// FEFF `apot.bin` section containing the initial core-hole orbital and density.
pub const APOT_CORE_HOLE_SECTION_NUMBER: usize = 5;
/// Number of radial rows in FEFF ATOM/POT radial tables.
pub const APOT_CORE_HOLE_RADIAL_POINTS: usize = 251;
/// FEFF logarithmic radial-grid origin for section-5 core-hole data.
pub const APOT_CORE_HOLE_GRID_ORIGIN: f64 = -8.8;
/// FEFF logarithmic radial-grid step for section-5 core-hole data.
/// `ATOM/apot.f90` assigns the default-real literal `0.05` to `real*8 hx`.
pub const APOT_CORE_HOLE_GRID_STEP: f64 = 0.05_f32 as f64;
/// Relative tolerance used when deciding whether cached `dvcoul` needs rewriting.
pub const APOT_CORE_HOLE_TOLERANCE: f64 = 5.0e-8;

/// Typed FEFF `apot.bin` section-5 core-hole columns.
#[derive(Debug, Clone, PartialEq)]
pub struct ApotCoreHoleColumns {
    /// Core-hole large component, FEFF `dgc0`.
    pub large_component: Array1<f64>,
    /// Core-hole small component, FEFF `dpc0`.
    pub small_component: Array1<f64>,
    /// Core-hole density, FEFF `drho`.
    pub density: Array1<f64>,
    /// Core-hole Coulomb potential, FEFF `dvcoul`.
    pub coulomb_potential: Array1<f64>,
}

/// Extract validated FEFF `apot.bin` section-5 core-hole columns.
pub fn apot_core_hole_columns(apot: &ApotBinData) -> Result<ApotCoreHoleColumns> {
    let records = core_hole_records(apot)?;
    let large_component = real_record_column(records, 0, "dgc0")?;
    let small_component = real_record_column(records, 1, "dpc0")?;
    let density = real_record_column(records, 2, "drho")?;
    let coulomb_potential = real_record_column(records, 3, "dvcoul")?;
    validate_radial_len("dgc0", large_component.len())?;
    validate_radial_len("dpc0", small_component.len())?;
    validate_radial_len("drho", density.len())?;
    validate_radial_len("dvcoul", coulomb_potential.len())?;
    Ok(ApotCoreHoleColumns {
        large_component,
        small_component,
        density,
        coulomb_potential,
    })
}

/// Recompute and refresh the `dvcoul` column in FEFF `apot.bin` section 5.
///
/// FEFF `ATOM/apot.f90` stores section 5 as `dgc0`, `dpc0`, `drho`, `dvcoul`.
/// For core-hole calculations, `dvcoul` is regenerated by applying
/// `ATOM/potslw.f90` to `2 * drho * r^2` on the fixed ATOM radial grid. For
/// `nohole <= 0`, FEFF expects a zero core-hole density and zero Coulomb
/// potential.
pub fn refresh_apot_core_hole_coulomb_payload(apot: &mut ApotBinData, nohole: i32) -> Result<()> {
    let records = core_hole_records_mut(apot)?;
    let drho = real_record_column(records, 2, "drho")?;
    if drho.len() != APOT_CORE_HOLE_RADIAL_POINTS {
        return invalid_core_hole(format!(
            "core-hole drho has {} rows, expected {APOT_CORE_HOLE_RADIAL_POINTS}",
            drho.len()
        ));
    }

    let expected = apot_core_hole_coulomb_from_density(drho.view(), nohole)?;
    refresh_real_record_column(records, 3, "dvcoul", expected.view())
}

/// Build the expected FEFF section-5 `dvcoul` values from a `drho` column.
pub fn apot_core_hole_coulomb_from_density(
    drho: ArrayView1<'_, f64>,
    nohole: i32,
) -> Result<Array1<f64>> {
    if drho.len() != APOT_CORE_HOLE_RADIAL_POINTS {
        return invalid_core_hole(format!(
            "core-hole drho has {} rows, expected {APOT_CORE_HOLE_RADIAL_POINTS}",
            drho.len()
        ));
    }
    validate_finite_column("drho", drho)?;

    if nohole <= 0 {
        ensure_near_zero("drho", drho)?;
        return Ok(Array1::zeros(drho.len()));
    }

    let radii = apot_core_hole_radii(drho.len());
    let source_density =
        Array1::from_shape_fn(drho.len(), |row| 2.0 * drho[row] * radii[row] * radii[row]);
    atomic_four_point_coulomb_potential(AtomicFourPointCoulombPotentialInput {
        density: source_density.view(),
        radii: radii.view(),
        step: APOT_CORE_HOLE_GRID_STEP,
        active_len: drho.len(),
    })
    .map_err(|source| IoError::InvalidApotBin {
        line: 0,
        message: format!("core-hole Coulomb regeneration failed: {source}"),
    })
}

/// FEFF section-5 logarithmic radial grid.
#[must_use]
pub fn apot_core_hole_radii(count: usize) -> Array1<f64> {
    Array1::from_shape_fn(count, |row| {
        (APOT_CORE_HOLE_GRID_ORIGIN + APOT_CORE_HOLE_GRID_STEP * row as f64).exp()
    })
}

fn core_hole_records(apot: &ApotBinData) -> Result<&ApotBinRecords> {
    let section = apot
        .sections
        .iter()
        .find(|section| section.section_number == APOT_CORE_HOLE_SECTION_NUMBER)
        .ok_or_else(|| invalid_core_hole_error("missing section 5 core-hole payload"))?;
    validate_core_hole_section_labels(&section.column_labels)?;
    match &section.payload {
        ApotBinPayload::Records(records) => {
            validate_core_hole_column_types(records)?;
            Ok(records)
        }
        ApotBinPayload::HeadersOnly | ApotBinPayload::Matrix(_) => {
            invalid_core_hole("section 5 is not a row-record payload")
        }
    }
}

fn core_hole_records_mut(apot: &mut ApotBinData) -> Result<&mut ApotBinRecords> {
    let section = apot
        .sections
        .iter_mut()
        .find(|section| section.section_number == APOT_CORE_HOLE_SECTION_NUMBER)
        .ok_or_else(|| invalid_core_hole_error("missing section 5 core-hole payload"))?;
    validate_core_hole_section_labels(&section.column_labels)?;
    match &mut section.payload {
        ApotBinPayload::Records(records) => {
            validate_core_hole_column_types(records)?;
            Ok(records)
        }
        ApotBinPayload::HeadersOnly | ApotBinPayload::Matrix(_) => {
            invalid_core_hole("section 5 is not a row-record payload")
        }
    }
}

fn validate_core_hole_section_labels(column_labels: &[String]) -> Result<()> {
    if column_labels.len() < 4
        || column_labels[0] != "dgc0"
        || column_labels[1] != "dpc0"
        || column_labels[2] != "drho"
        || column_labels[3] != "dvcoul"
    {
        return invalid_core_hole(format!(
            "section 5 has unexpected column labels {column_labels:?}"
        ));
    }
    Ok(())
}

fn validate_core_hole_column_types(records: &ApotBinRecords) -> Result<()> {
    if records.column_types.len() < 4 || records.column_types[..4] != [ApotBinType::Double; 4] {
        return invalid_core_hole(format!(
            "section 5 has unexpected column types {:?}",
            records.column_types
        ));
    }
    Ok(())
}

fn real_record_column(
    records: &ApotBinRecords,
    column: usize,
    name: &'static str,
) -> Result<Array1<f64>> {
    let mut values = Vec::with_capacity(records.rows.len());
    for (row_index, row) in records.rows.iter().enumerate() {
        let value = row.get(column).ok_or_else(|| {
            invalid_core_hole_error(format!("core-hole row {row_index} missing {name}"))
        })?;
        match value {
            ApotBinValue::Real(value) if value.is_finite() => values.push(*value),
            ApotBinValue::Real(value) => {
                return invalid_core_hole(format!(
                    "core-hole {name} row {row_index} is non-finite: {value}"
                ));
            }
            ApotBinValue::Int(_) | ApotBinValue::Complex(_) | ApotBinValue::Text(_) => {
                return invalid_core_hole(format!(
                    "core-hole {name} row {row_index} is not real-valued"
                ));
            }
        }
    }
    Ok(Array1::from_vec(values))
}

fn refresh_real_record_column(
    records: &mut ApotBinRecords,
    column: usize,
    name: &'static str,
    expected: ArrayView1<'_, f64>,
) -> Result<()> {
    if records.rows.len() != expected.len() {
        return invalid_core_hole(format!(
            "core-hole {name} has {} rows, expected {}",
            records.rows.len(),
            expected.len()
        ));
    }
    for (row_index, (row, &expected)) in records.rows.iter_mut().zip(expected.iter()).enumerate() {
        let value = row.get_mut(column).ok_or_else(|| {
            invalid_core_hole_error(format!("core-hole row {row_index} missing {name}"))
        })?;
        let ApotBinValue::Real(actual) = value else {
            return invalid_core_hole(format!(
                "core-hole {name} row {row_index} is not real-valued"
            ));
        };
        let allowed = APOT_CORE_HOLE_TOLERANCE * expected.abs().max(1.0);
        let difference = (*actual - expected).abs();
        if !actual.is_finite() || difference > allowed {
            *actual = expected;
        }
    }
    Ok(())
}

fn validate_finite_column(name: &'static str, values: ArrayView1<'_, f64>) -> Result<()> {
    for (row, &value) in values.iter().enumerate() {
        if !value.is_finite() {
            return invalid_core_hole(format!("core-hole {name} row {row} is non-finite: {value}"));
        }
    }
    Ok(())
}

fn validate_radial_len(name: &'static str, len: usize) -> Result<()> {
    if len != APOT_CORE_HOLE_RADIAL_POINTS {
        return invalid_core_hole(format!(
            "core-hole {name} has {len} rows, expected {APOT_CORE_HOLE_RADIAL_POINTS}"
        ));
    }
    Ok(())
}

fn ensure_near_zero(name: &'static str, values: ArrayView1<'_, f64>) -> Result<()> {
    for (row, &value) in values.iter().enumerate() {
        if value.abs() > APOT_CORE_HOLE_TOLERANCE {
            return invalid_core_hole(format!(
                "core-hole {name} row {row} expected zero for nohole<=0, got {value:e}"
            ));
        }
    }
    Ok(())
}

fn invalid_core_hole<T>(message: impl Into<String>) -> Result<T> {
    Err(invalid_core_hole_error(message))
}

fn invalid_core_hole_error(message: impl Into<String>) -> IoError {
    IoError::InvalidApotBin {
        line: 0,
        message: message.into(),
    }
}