rexafs 0.2.6

use crate::plot::config::{FitPlotOptions, GroupPlotOptions, PanelKind, PanelSpec, PlotConfig};
use crate::plot::errors::PlotError;
use crate::plot::{fitting_plots, group_plots, panels, spectrum_plots};
use crate::xafs::fitting::types::FeffFitResult;
use crate::xafs::xasgroup::XASGroup;
use crate::xafs::xasspectrum::XASSpectrum;
use ruviz::core::Plot;
use std::path::Path;

pub(crate) enum PlotSource<'a> {
    Spectrum(&'a mut XASSpectrum),
    Group(&'a mut XASGroup),
    Fit(&'a mut FeffFitResult),
}

/// Configure and render panels while borrowing a spectrum, group or fit result.
///
/// Select a panel before its options; invalid combinations are reported when
/// rendering. Default output is 800 × 600 pixels at 300 DPI, with a legend and
/// no title. Missing spectrum stages may be computed during rendering; processing
/// and rendering failures are returned as [`PlotError`].
pub struct XASPlotBuilder<'a> {
    source: PlotSource<'a>,
    panels: Vec<PanelSpec>,
    config: PlotConfig,
    group: GroupPlotOptions,
    fit: FitPlotOptions,
    pending_error: Option<PlotError>,
}

impl<'a> XASPlotBuilder<'a> {
    pub(crate) fn for_spectrum(spectrum: &'a mut XASSpectrum) -> Self {
        Self {
            source: PlotSource::Spectrum(spectrum),
            panels: Vec::new(),
            config: PlotConfig::default(),
            group: GroupPlotOptions::default(),
            fit: FitPlotOptions::default(),
            pending_error: None,
        }
    }

    pub(crate) fn for_group(group: &'a mut XASGroup) -> Self {
        Self {
            source: PlotSource::Group(group),
            panels: Vec::new(),
            config: PlotConfig::default(),
            group: GroupPlotOptions::default(),
            fit: FitPlotOptions::default(),
            pending_error: None,
        }
    }

    pub(crate) fn for_fit(fit: &'a mut FeffFitResult) -> Self {
        Self {
            source: PlotSource::Fit(fit),
            panels: Vec::new(),
            config: PlotConfig::default(),
            group: GroupPlotOptions::default(),
            fit: FitPlotOptions::default(),
            pending_error: None,
        }
    }

    fn set_error(&mut self, error: PlotError) {
        if self.pending_error.is_none() {
            self.pending_error = Some(error);
        }
    }

    fn has_panels(&self) -> bool {
        !self.panels.is_empty()
    }

    fn is_group_source(&self) -> bool {
        matches!(self.source, PlotSource::Group(_))
    }

    fn is_fit_source(&self) -> bool {
        matches!(self.source, PlotSource::Fit(_))
    }

    fn last_panel_kind(&self) -> Option<PanelKind> {
        self.panels.last().map(|panel| panel.kind)
    }

    fn require_panel_for_option(&mut self, option_name: &str) -> bool {
        if self.has_panels() {
            true
        } else {
            self.set_error(PlotError::invalid_option(format!(
                "{option_name}() requires a panel selector before it"
            )));
            false
        }
    }

    fn apply_panel_option<F>(&mut self, option_name: &str, mut f: F)
    where
        F: FnMut(&mut PanelSpec) -> Result<(), PlotError>,
    {
        if !self.require_panel_for_option(option_name) {
            return;
        }

        if let Some(panel) = self.panels.last_mut() {
            if let Err(error) = f(panel) {
                self.set_error(error);
            }
        }
    }

    fn build_plots(&mut self) -> Result<Vec<Plot>, PlotError> {
        if let Some(error) = self.pending_error.take() {
            return Err(error);
        }
        if self.panels.is_empty() {
            return Err(PlotError::EmptySelection);
        }

        let panel_specs = self.panels.clone();
        let show_legend = self.config.show_legend;
        let mut plots = Vec::with_capacity(panel_specs.len());

        match &mut self.source {
            PlotSource::Spectrum(spectrum) => {
                for panel in panel_specs {
                    let data = spectrum_plots::extract_spectrum_panel_data(spectrum, &panel)?;
                    let plot = panels::render_panel(data, show_legend)?;
                    plots.push(plot);
                }
            }
            PlotSource::Group(group) => {
                for panel in panel_specs {
                    let data = group_plots::extract_group_panel_data(group, &panel, &self.group)?;
                    let plot = panels::render_panel(data, show_legend)?;
                    plots.push(plot);
                }
            }
            PlotSource::Fit(fit) => {
                for panel in panel_specs {
                    let data = fitting_plots::extract_fit_panel_data(fit, &panel, &self.fit)?;
                    let plot = panels::render_panel(data, show_legend)?;
                    plots.push(plot);
                }
            }
        }

        Ok(plots)
    }

    fn apply_single_config(&self, mut plot: Plot) -> Plot {
        plot = plot.size_px(self.config.width, self.config.height);
        plot = plot.dpi(self.config.dpi);
        if let Some(title) = &self.config.title {
            plot = plot.title(title.clone());
        }
        plot
    }

    /// Append an energy panel, preferring flattened absorption.
    /// For spectra, attempts normalization when at least eight points are available;
    /// if flattening is unavailable, displays working mu with the corresponding label.
    pub fn mu(mut self) -> Self {
        self.panels.push(PanelSpec::new(PanelKind::Mu));
        self
    }

    /// Append normalized absorption against energy in eV.
    /// Computes missing normalization; a failed prerequisite returns a rendering error.
    pub fn norm(mut self) -> Self {
        self.panels.push(PanelSpec::new(PanelKind::Norm));
        self
    }

    /// Append weighted EXAFS against wave number in Å⁻¹.
    /// Computes a missing background. Default display weight is 2 for spectra/groups
    /// and the stored fit weight for fit results; use [`Self::kweight`] to override it.
    pub fn k(mut self) -> Self {
        self.panels.push(PanelSpec::new(PanelKind::K));
        self
    }

    /// Append Fourier components against distance in Å, initially magnitude only.
    /// Computes a missing forward transform; the default displayed interval is 0–6 Å.
    /// Component units are Å^(-(w + 1)) for the forward transform's k weight `w`
    /// and dimensionless chi(k). No scattering phase correction is applied.
    pub fn r(mut self) -> Self {
        self.panels.push(PanelSpec::new(PanelKind::R));
        self
    }

    /// Set the most recently selected k panel's display weight.
    /// This multiplies the plotted chi by k raised to this weight without changing
    /// the stored processing settings. Calling it before `.k()` records an error.
    pub fn kweight(mut self, weight: f64) -> Self {
        self.apply_panel_option("kweight", |panel| {
            if panel.kind != PanelKind::K {
                return Err(PlotError::invalid_option(
                    "kweight() is only valid for k() panels",
                ));
            }
            panel.kweight = Some(weight);
            Ok(())
        });
        self
    }

    /// Set the latest R panel to magnitude, real and imaginary traces when true,
    /// or magnitude alone when false. Requires an R panel.
    pub fn components(mut self, show: bool) -> Self {
        self.apply_panel_option("components", |panel| {
            if panel.kind != PanelKind::R {
                return Err(PlotError::invalid_option(
                    "components() is only valid for r() panels",
                ));
            }
            if show {
                panel.r_mag = Some(true);
                panel.r_real = true;
                panel.r_imag = true;
            } else {
                panel.r_mag = Some(true);
                panel.r_real = false;
                panel.r_imag = false;
            }
            Ok(())
        });
        self
    }

    /// Include magnitude in the latest R panel. Magnitude is the initial default;
    /// use this explicitly to retain it together with `.real()` or `.imag()`.
    pub fn mag(mut self) -> Self {
        self.apply_panel_option("mag", |panel| {
            if panel.kind != PanelKind::R {
                return Err(PlotError::invalid_option(
                    "mag() is only valid for r() panels",
                ));
            }
            panel.r_mag = Some(true);
            Ok(())
        });
        self
    }

    /// Include the real component in the latest R panel. Hides the implicit
    /// magnitude unless `.mag()` or `.components(...)` selected it explicitly.
    pub fn real(mut self) -> Self {
        self.apply_panel_option("real", |panel| {
            if panel.kind != PanelKind::R {
                return Err(PlotError::invalid_option(
                    "real() is only valid for r() panels",
                ));
            }
            if panel.r_mag.is_none() {
                panel.r_mag = Some(false);
            }
            panel.r_real = true;
            Ok(())
        });
        self
    }

    /// Include the imaginary component in the latest R panel. Hides the implicit
    /// magnitude unless `.mag()` or `.components(...)` selected it explicitly.
    pub fn imag(mut self) -> Self {
        self.apply_panel_option("imag", |panel| {
            if panel.kind != PanelKind::R {
                return Err(PlotError::invalid_option(
                    "imag() is only valid for r() panels",
                ));
            }
            if panel.r_mag.is_none() {
                panel.r_mag = Some(false);
            }
            panel.r_imag = true;
            Ok(())
        });
        self
    }

    /// Show fitted pre/post-edge curves on the latest normalization panel.
    /// Default: false. Requires `.norm()` before this option.
    pub fn edges(mut self, show: bool) -> Self {
        self.apply_panel_option("edges", |panel| {
            if panel.kind != PanelKind::Norm {
                return Err(PlotError::invalid_option(
                    "edges() is only valid for norm() panels",
                ));
            }
            panel.edges = show;
            Ok(())
        });
        self
    }

    /// Show both the window curve and its range markers on the latest k panel.
    /// Default: false. Equivalent to setting `window_fn` and `window_box` together.
    pub fn window(mut self, show: bool) -> Self {
        self.apply_panel_option("window", |panel| {
            if panel.kind != PanelKind::K {
                return Err(PlotError::invalid_option(
                    "window() is only valid for k() panels",
                ));
            }
            panel.window_fn = show;
            panel.window_box = show;
            Ok(())
        });
        self
    }

    /// Show the window curve on the latest k panel. Default: false.
    /// This is a display overlay, not a change to transform parameters.
    pub fn window_fn(mut self, show: bool) -> Self {
        self.apply_panel_option("window_fn", |panel| {
            if panel.kind != PanelKind::K {
                return Err(PlotError::invalid_option(
                    "window_fn() is only valid for k() panels",
                ));
            }
            panel.window_fn = show;
            Ok(())
        });
        self
    }

    /// Show lower/upper range markers on the latest k panel, or on an R panel
    /// for a fit result. Default: false. Despite the name, these are two lines rather
    /// than a rectangle; other panel/source combinations return a rendering error.
    pub fn window_box(mut self, show: bool) -> Self {
        let fit_source = self.is_fit_source();
        self.apply_panel_option("window_box", |panel| match panel.kind {
            PanelKind::K => {
                panel.window_box = show;
                Ok(())
            }
            PanelKind::R if fit_source => {
                panel.window_box = show;
                Ok(())
            }
            PanelKind::R => Err(PlotError::invalid_option(
                "window_box() on r() panels is only valid for FeffFitResult sources",
            )),
            _ => Err(PlotError::invalid_option(
                "window_box() is only valid for k() panels, or r() panels on FeffFitResult",
            )),
        });
        self
    }

    /// Vertically offset successive selected group spectra by `offset` in plotted
    /// y units. Default: overlaid. Applies to the group's panels and requires at least
    /// one selected panel; spectrum/fit sources are unsupported.
    pub fn stacked(mut self, offset: f64) -> Self {
        if !self.require_panel_for_option("stacked") {
            return self;
        }
        if !self.is_group_source() {
            self.set_error(PlotError::invalid_option(
                "stacked() is only valid for XASGroup sources",
            ));
            return self;
        }
        self.group.stacked = Some(offset);
        self
    }

    /// Copy the zero-based group indices to display across its panels.
    /// Default: all members. Requires a group source and at least one panel.
    pub fn select(mut self, indices: &[usize]) -> Self {
        if !self.require_panel_for_option("select") {
            return self;
        }
        if !self.is_group_source() {
            self.set_error(PlotError::invalid_option(
                "select() is only valid for XASGroup sources",
            ));
            return self;
        }
        self.group.selected = Some(indices.to_vec());
        self
    }

    /// Select one zero-based dataset from a stored fit result for its panels.
    /// Default: the first dataset, or legacy top-level arrays if no datasets exist.
    /// Requires a fit source and at least one panel; an unavailable index produces
    /// a rendering error.
    pub fn dataset(mut self, index: usize) -> Self {
        if !self.require_panel_for_option("dataset") {
            return self;
        }
        if !self.is_fit_source() {
            self.set_error(PlotError::invalid_option(
                "dataset() is only valid for FeffFitResult sources",
            ));
            return self;
        }
        self.fit.dataset = Some(index);
        self
    }

    /// Show individual fit-path contributions on k panels. Default: false.
    /// Requires a fit source with k as its latest panel. R magnitude panels already
    /// include stored path magnitudes independently of this option.
    pub fn paths(mut self, show: bool) -> Self {
        if !self.require_panel_for_option("paths") {
            return self;
        }
        if !self.is_fit_source() {
            self.set_error(PlotError::invalid_option(
                "paths() is only valid for FeffFitResult sources",
            ));
            return self;
        }
        if self.last_panel_kind() != Some(PanelKind::K) {
            self.set_error(PlotError::invalid_option(
                "paths() is only valid when the last selected panel is k()",
            ));
            return self;
        }
        self.fit.paths = show;
        self
    }

    /// Copy a title for the output figure. Default: no title.
    pub fn title(mut self, title: &str) -> Self {
        self.config.title = Some(title.to_string());
        self
    }

    /// Set output width in pixels, clamped to at least 1. Default: 800.
    pub fn width(mut self, width: u32) -> Self {
        self.config.width = width.max(1);
        self
    }

    /// Set output height in pixels, clamped to at least 1. Default: 600.
    pub fn height(mut self, height: u32) -> Self {
        self.config.height = height.max(1);
        self
    }

    /// Show or hide legends. Default: visible.
    pub fn legend(mut self, show: bool) -> Self {
        self.config.show_legend = show;
        self
    }

    /// Render the selected panels and write a PNG file, overwriting the destination.
    /// Multiple panels are combined in one figure. Processing, configuration, rendering
    /// and file-write failures return [`PlotError`].
    pub fn save_png<P: AsRef<Path>>(mut self, path: P) -> Result<(), PlotError> {
        let mut plots = self.build_plots()?;
        if plots.len() == 1 {
            let plot = self.apply_single_config(plots.remove(0));
            plot.save(path)?;
            return Ok(());
        }

        let figure = panels::build_subplot_figure(
            plots,
            self.config.width,
            self.config.height,
            self.config.title.as_deref(),
        )?;
        figure.save_with_dpi(path, self.config.dpi as f32)?;
        Ok(())
    }

    /// Render exactly one panel into an owned ruviz plot for further customization.
    /// No selected panel or multiple panels returns an error; use [`Self::save_png`]
    /// or [`Self::to_svg_panels`] for multiple panels.
    pub fn render_plot(mut self) -> Result<Plot, PlotError> {
        let mut plots = self.build_plots()?;
        if plots.len() != 1 {
            return Err(PlotError::MultiPanelRenderUnsupported);
        }
        Ok(self.apply_single_config(plots.remove(0)))
    }

    /// Render exactly one panel into an owned SVG string without writing a file.
    /// Uses the same prerequisites and single-panel restriction as [`Self::render_plot`].
    pub fn to_svg(self) -> Result<String, PlotError> {
        let plot = self.render_plot()?;
        Ok(plot.render_to_svg()?)
    }

    /// Render one owned SVG string per selected panel, preserving panel order.
    /// Each SVG uses the configured width/height; the combined-figure title is not
    /// inserted. No file is written and no combined multi-panel SVG is produced.
    pub fn to_svg_panels(mut self) -> Result<Vec<String>, PlotError> {
        let plots = self.build_plots()?;
        let mut svg_panels = Vec::with_capacity(plots.len());

        for plot in plots {
            let panel = plot.size_px(self.config.width, self.config.height);
            svg_panels.push(panel.render_to_svg()?);
        }

        Ok(svg_panels)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::plot::traits::PlotXAS;
    use crate::xafs::fitting::types::FeffFitResult;
    use nalgebra::DVector;

    #[test]
    fn option_before_panel_is_rejected() {
        let mut spectrum = XASSpectrum::new();
        let error = spectrum
            .plot()
            .kweight(2.0)
            .mu()
            .render_plot()
            .expect_err("must fail");
        assert!(matches!(error, PlotError::InvalidOption { .. }));
    }

    #[test]
    fn invalid_panel_option_combination_is_rejected() {
        let mut spectrum = XASSpectrum::new();
        let error = spectrum
            .plot()
            .mu()
            .components(true)
            .render_plot()
            .expect_err("must fail");
        assert!(matches!(error, PlotError::InvalidOption { .. }));
    }

    #[test]
    fn window_alias_enables_fn_and_box_for_k_panel() {
        let mut spectrum = XASSpectrum::new();
        let builder = spectrum.plot().k().window(true);
        assert!(builder.pending_error.is_none());
        assert!(builder.panels[0].window_fn);
        assert!(builder.panels[0].window_box);
    }

    #[test]
    fn window_fn_is_rejected_for_r_panel() {
        let mut spectrum = XASSpectrum::new();
        let builder = spectrum.plot().r().window_fn(true);
        assert!(matches!(
            builder.pending_error,
            Some(PlotError::InvalidOption { .. })
        ));
    }

    #[test]
    fn window_box_on_r_panel_is_fit_only() {
        let mut spectrum = XASSpectrum::new();
        let spectrum_builder = spectrum.plot().r().window_box(true);
        assert!(matches!(
            spectrum_builder.pending_error,
            Some(PlotError::InvalidOption { .. })
        ));

        let mut fit = FeffFitResult::default();
        let fit_builder = fit.plot().r().window_box(true);
        assert!(fit_builder.pending_error.is_none());
        assert!(fit_builder.panels[0].window_box);
    }

    #[test]
    fn dataset_index_validation_propagates() {
        let mut fit = FeffFitResult::default();
        let error = fit
            .plot()
            .k()
            .dataset(3)
            .save_png("/tmp/ignore.png")
            .expect_err("must fail");
        assert!(matches!(error, PlotError::IndexOutOfRange { .. }));
    }

    #[test]
    fn multi_panel_single_output_methods_are_rejected() {
        let mut spectrum = XASSpectrum::new();
        spectrum.set_spectrum(
            DVector::from_vec(vec![1.0, 2.0, 3.0, 4.0]),
            DVector::from_vec(vec![1.0, 1.5, 2.0, 2.5]),
        );

        let render_error = spectrum
            .plot()
            .mu()
            .mu()
            .render_plot()
            .expect_err("multi-panel render_plot must fail");
        assert!(matches!(
            render_error,
            PlotError::MultiPanelRenderUnsupported
        ));

        let svg_error = spectrum
            .plot()
            .mu()
            .mu()
            .to_svg()
            .expect_err("multi-panel to_svg must fail");
        assert!(matches!(svg_error, PlotError::MultiPanelRenderUnsupported));

        let panels = spectrum
            .plot()
            .mu()
            .mu()
            .to_svg_panels()
            .expect("multi-panel svg panels should succeed");
        assert_eq!(panels.len(), 2);
    }
}