rexafs 0.2.6

use rexafs::{io, BackgroundMethod, NormalizationMethod, PrePostEdge, Spectrum, XrayFFTF, AUTOBK};

fn sample() -> Spectrum {
    io::read_qas_transmission(format!(
        "{}/tests/testfiles/Ru_QAS.dat",
        env!("CARGO_MANIFEST_DIR")
    ))
    .unwrap()
}

#[test]
fn terminal_stage_matches_explicit_stages() {
    let mut automatic = sample();
    automatic.fft().unwrap();
    let mut explicit = sample();
    explicit
        .normalize()
        .unwrap()
        .calc_background()
        .unwrap()
        .fft()
        .unwrap();
    {
        let spectrum = &explicit;
        assert_eq!(automatic.e0(), spectrum.e0());
        assert_eq!(automatic.k(), spectrum.k());
        assert_eq!(automatic.chi(), spectrum.chi());
        assert_eq!(automatic.r(), spectrum.r());
        assert_eq!(automatic.chir_mag(), spectrum.chir_mag());
    }
    let mag = automatic.chir_mag().unwrap();
    let re = automatic.chir_real().unwrap();
    let im = automatic.chir_imag().unwrap();
    for ((mag, re), im) in mag.iter().zip(re.iter()).zip(im.iter()) {
        assert!((mag - re.hypot(*im)).abs() < 1e-10);
    }
}

#[test]
fn invalid_arrays_return_errors_before_processing() {
    for (energy, mu) in [
        (vec![], vec![]),
        (vec![2., 1., 3.], vec![1.]),
        (vec![1., 2., 2.], vec![1., 2., 3.]),
        (vec![1., f64::NAN, 3.], vec![1., 2., 3.]),
        (vec![1., 2., 3.], vec![1., f64::INFINITY, 3.]),
    ] {
        assert!(Spectrum::from_arrays(&energy, &mu).is_err());
    }
}

#[test]
fn edge_changes_invalidate_results_and_are_used() {
    let mut spectrum = sample();
    spectrum.fft().unwrap();
    let edge = spectrum.e0().unwrap() + 0.25;
    spectrum.set_e0(edge);
    assert!(spectrum.norm().is_none());
    assert!(spectrum.chi().is_none());
    assert!(spectrum.r().is_none());
    spectrum.fft().unwrap();
    assert_eq!(spectrum.e0(), Some(edge));
    if let Some(BackgroundMethod::AUTOBK(a)) = &spectrum.background {
        assert_eq!(a.ek0, Some(edge));
    }
    for e0 in [f64::NAN, f64::INFINITY, 0., 1e9] {
        spectrum.set_e0(e0);
        assert!(spectrum.fft().is_err());
        assert!(spectrum.r().is_none());
    }
}

#[test]
fn stage_configuration_is_respected_and_only_dependents_are_invalidated() {
    let mut spectrum = sample();
    let mut norm = PrePostEdge::new();
    norm.pre_edge_start = Some(-200.);
    norm.pre_edge_end = Some(-65.);
    spectrum
        .set_normalization_method(Some(NormalizationMethod::PrePostEdge(norm)))
        .unwrap();
    let mut background = AUTOBK::new();
    background.rbkg = Some(1.2);
    spectrum
        .set_background_method(Some(BackgroundMethod::AUTOBK(background)))
        .unwrap();
    let mut transform = XrayFFTF::new();
    transform.kweight = Some(1.);
    spectrum.set_fft(transform.clone()).fft().unwrap();
    let norm = spectrum.norm().unwrap();
    let chi_address = spectrum.chi().unwrap().as_ptr();
    let chi = spectrum.chi().unwrap().to_vec();
    let old_r = spectrum.chir_mag().unwrap();
    transform.kweight = Some(3.);
    spectrum.set_fft(transform);
    assert!(spectrum.r().is_none());
    spectrum.fft().unwrap();
    assert_eq!(spectrum.chi().unwrap().as_ptr(), chi_address);
    assert_eq!(spectrum.chi().unwrap(), chi);
    assert_eq!(spectrum.norm().unwrap(), norm);
    assert_ne!(spectrum.chir_mag().unwrap(), old_r);
    assert_eq!(spectrum.kweight(), Some(&3.));
    let mut background = AUTOBK::new();
    background.rbkg = Some(1.5);
    spectrum
        .set_background_method(Some(BackgroundMethod::AUTOBK(background)))
        .unwrap();
    assert_eq!(spectrum.norm().unwrap(), norm);
    assert!(spectrum.chi().is_none());
    assert!(spectrum.r().is_none());
    spectrum.fft().unwrap();
    assert_ne!(spectrum.chi().unwrap(), chi);
}

#[test]
fn direct_settings_match_enum_and_optional_enum_configuration() {
    let mut normalization = PrePostEdge::new();
    normalization.pre_edge_start = Some(-200.0);
    normalization.pre_edge_end = Some(-65.0);
    let mut background = AUTOBK::new();
    background.rbkg = Some(1.2);

    let mut direct = sample();
    direct
        .set_normalization_method(normalization.clone())
        .unwrap()
        .set_background_method(background.clone())
        .unwrap()
        .fft()
        .unwrap();

    let mut enum_form = sample();
    enum_form
        .set_normalization_method(NormalizationMethod::from(normalization.clone()))
        .unwrap()
        .set_background_method(BackgroundMethod::from(background.clone()))
        .unwrap()
        .fft()
        .unwrap();

    let mut optional_form = sample();
    optional_form
        .set_normalization_method(Some(NormalizationMethod::PrePostEdge(normalization)))
        .unwrap()
        .set_background_method(Some(BackgroundMethod::AUTOBK(background)))
        .unwrap()
        .fft()
        .unwrap();

    for other in [&enum_form, &optional_form] {
        assert_eq!(direct.norm(), other.norm());
        assert_eq!(direct.chi(), other.chi());
        assert_eq!(direct.chir_mag(), other.chir_mag());
    }

    let norm = direct.norm().unwrap();
    direct.set_background_method(AUTOBK::new()).unwrap();
    assert_eq!(direct.norm().unwrap(), norm);
    assert!(direct.chi().is_none());
    assert!(direct.chir().is_none());
    direct.set_normalization_method(PrePostEdge::new()).unwrap();
    assert!(direct.norm().is_none());

    // These unannotated calls must continue to infer the original enum types.
    direct.set_normalization_method(None).unwrap();
    direct.set_background_method(None).unwrap();
    direct.fft().unwrap();
    let mut defaults = sample();
    defaults.fft().unwrap();
    assert_eq!(direct.norm(), defaults.norm());
    assert_eq!(direct.chi(), defaults.chi());
    assert_eq!(direct.chir_mag(), defaults.chir_mag());
}

#[test]
fn unsupported_selected_algorithms_never_fall_back_to_defaults() {
    let mut spectrum = sample();
    spectrum.fft().unwrap();
    spectrum
        .set_normalization_method(Some(NormalizationMethod::new_mback()))
        .unwrap();
    assert!(spectrum.fft().unwrap_err().to_string().contains("MBack"));
    assert!(matches!(
        spectrum.normalization,
        Some(NormalizationMethod::MBack(_))
    ));
    spectrum.set_normalization_method(None).unwrap();
    spectrum
        .set_background_method(Some(BackgroundMethod::new_ilpbkg()))
        .unwrap();
    assert!(spectrum.fft().unwrap_err().to_string().contains("ILPBkg"));
    assert!(matches!(
        spectrum.background,
        Some(BackgroundMethod::ILPBkg(_))
    ));
}

#[test]
fn replacing_data_invalidates_previous_stages() {
    let mut spectrum = sample();
    spectrum.fft().unwrap();
    let energy = spectrum.energy.clone().unwrap();
    let mu = spectrum.mu.clone().unwrap();
    spectrum.set_spectrum(energy, mu);
    assert!(spectrum.e0().is_none());
    assert!(spectrum.chi().is_none());
    assert!(spectrum.r().is_none());
    spectrum.fft().unwrap();
    assert!(spectrum.r().is_some());
}

#[test]
fn explicit_edge_step_survives_invalidation_and_serialization() {
    let mut spectrum = sample();
    let mut norm = PrePostEdge::new();
    norm.edge_step = Some(2.5);
    spectrum
        .set_normalization_method(Some(NormalizationMethod::PrePostEdge(norm)))
        .unwrap();
    spectrum.fft().unwrap();
    let json = serde_json::to_string(&spectrum).unwrap();
    let mut restored: Spectrum = serde_json::from_str(&json).unwrap();
    for spectrum in [&mut spectrum, &mut restored] {
        spectrum
            .set_e0(spectrum.e0().unwrap() + 0.25)
            .fft()
            .unwrap();
        assert_eq!(
            spectrum.normalization.as_ref().unwrap().get_edge_step(),
            Some(2.5)
        );
    }
}

#[test]
fn public_edge_step_edits_survive_invalidation_and_serialization() {
    for initial in [None, Some(2.5)] {
        let mut spectrum = sample();
        let mut norm = PrePostEdge::new();
        norm.edge_step = initial;
        spectrum
            .set_normalization_method(Some(NormalizationMethod::PrePostEdge(norm)))
            .unwrap()
            .fft()
            .unwrap();

        // Test both editing a restored result and restoring an already edited one.
        for serialize_before_edit in [false, true] {
            let mut edited = spectrum.clone();
            if serialize_before_edit {
                edited = serde_json::from_str(&serde_json::to_string(&edited).unwrap()).unwrap();
            }
            let Some(NormalizationMethod::PrePostEdge(norm)) = &mut edited.normalization else {
                panic!("expected pre/post-edge normalization");
            };
            norm.edge_step = Some(4.25);
            if !serialize_before_edit {
                edited = serde_json::from_str(&serde_json::to_string(&edited).unwrap()).unwrap();
            }
            edited.invalidate_derived().fft().unwrap();
            assert_eq!(
                edited.normalization.as_ref().unwrap().get_edge_step(),
                Some(4.25)
            );
            edited.set_e0(edited.e0().unwrap() + 0.25).fft().unwrap();
            assert_eq!(
                edited.normalization.as_ref().unwrap().get_edge_step(),
                Some(4.25)
            );

            // Clearing an explicit scale restores automatic scale estimation.
            let Some(NormalizationMethod::PrePostEdge(norm)) = &mut edited.normalization else {
                unreachable!();
            };
            norm.edge_step = None;
            let mut automatic = sample();
            automatic
                .set_normalization_method(edited.normalization.clone())
                .unwrap()
                .normalize()
                .unwrap();
            edited.invalidate_derived().normalize().unwrap();
            approx::assert_abs_diff_eq!(
                edited.norm().unwrap(),
                automatic.norm().unwrap(),
                epsilon = 1e-10
            );
        }
    }
}

#[test]
fn automatic_edge_step_is_recomputed_after_input_changes() {
    let mut spectrum = sample();
    spectrum.normalize().unwrap();
    let original_step = spectrum
        .normalization
        .as_ref()
        .unwrap()
        .get_edge_step()
        .unwrap();
    *spectrum.mu.as_mut().unwrap() *= 2.0;
    spectrum.invalidate_derived().normalize().unwrap();
    let new_step = spectrum
        .normalization
        .as_ref()
        .unwrap()
        .get_edge_step()
        .unwrap();
    assert!((new_step - 2.0 * original_step).abs() < 1e-10);
}

#[test]
fn normalization_edge_override_propagates_to_background() {
    let mut spectrum = sample();
    spectrum.fft().unwrap();
    let edge = spectrum.e0().unwrap() + 0.5;
    let mut norm = PrePostEdge::new();
    norm.e0 = Some(edge);
    spectrum
        .set_normalization_method(Some(NormalizationMethod::PrePostEdge(norm)))
        .unwrap()
        .fft()
        .unwrap();
    assert_eq!(spectrum.e0(), Some(edge));
    if let Some(BackgroundMethod::AUTOBK(a)) = &spectrum.background {
        assert_eq!(a.ek0, Some(edge));
    }
}

#[test]
fn failed_transform_preserves_parameters_and_clears_results() {
    let mut spectrum = sample();
    spectrum.fft().unwrap();
    let mut params = XrayFFTF::new();
    params.nfft = Some(0);
    params.kweight = Some(3.);
    spectrum.set_fft(params);
    assert!(spectrum.fft().is_err());
    assert_eq!(spectrum.xftf.as_ref().unwrap().nfft, Some(0));
    assert_eq!(spectrum.xftf.as_ref().unwrap().kweight, Some(3.));
    assert!(spectrum.r().is_none());
    assert!(spectrum.chi().is_some());
}

#[test]
fn fft_window_axis_tracks_resampling_without_replacing_background_k() {
    let mut spectrum = sample();
    spectrum.fft().unwrap();
    let original = spectrum.k().unwrap().to_vec();
    for step in [0.025, 0.1] {
        let ft = XrayFFTF {
            grid: rexafs::FFTGrid::Larch,
            kstep: Some(step),
            ..Default::default()
        };
        spectrum.set_fft(ft).fft().unwrap();
        assert_eq!(spectrum.k().unwrap(), original.as_slice());
        let axis = spectrum.kwin_k().unwrap();
        assert_eq!(axis.len(), spectrum.kwin().unwrap().len());
        for (i, &k) in axis.iter().enumerate() {
            assert_eq!(k, i as f64 * step);
        }
    }
}

#[test]
fn inverse_configuration_clears_cached_outputs_and_preserves_forward_results() {
    let mut spectrum = sample();
    spectrum.ifft().unwrap();
    let forward = spectrum.chir_mag().unwrap();
    let background = spectrum.chi().unwrap().to_vec();
    let unfiltered = spectrum.chiq().unwrap();
    // A computed configuration can be reused without leaking another spectrum's
    // cached q/chi(q). Both Rust array backends must honor the same contract.
    let mut inverse = spectrum.xftr.clone().unwrap();
    assert!(inverse.chiq.is_some());
    inverse.rmin = Some(1.0);
    inverse.rmax = Some(3.0);
    inverse.dr = Some(0.5);
    spectrum.set_ifft(inverse);
    assert!(spectrum.q().is_none());
    assert!(spectrum.chiq().is_none());
    assert!(spectrum.xftr.as_ref().unwrap().rwin.is_none());
    assert_eq!(spectrum.chir_mag().unwrap(), forward);
    assert_eq!(spectrum.chi().unwrap(), background);
    spectrum.ifft().unwrap();
    assert_ne!(spectrum.chiq().unwrap(), unfiltered);
    assert_eq!(spectrum.q().unwrap().len(), spectrum.chiq().unwrap().len());
    assert_eq!(spectrum.chir_mag().unwrap(), forward);

    let mut invalid = rexafs::XrayFFTR::new();
    invalid.nfft = Some(0);
    spectrum.set_ifft(invalid);
    assert!(spectrum.ifft().is_err());
    assert!(spectrum.chiq().is_none());
    assert_eq!(spectrum.chir_mag().unwrap(), forward);
}