"""Validation Against Mantid
============================

Here we check that using the appropriate `abinslib` features we can obtain
very close agreement with output from the Mantid `Abins v1` algorithm.
"""

# %%
# Setup and Data Fetching
# -----------------------
# First, we fetch the necessary reference data. This includes
# the input phonon modes and corresponding JSON spectra generated by
# simulations in Mantid. To see the exact settings in Mantid, examine the
# validation workflows at
# https://github.com/ISISNeutronMuon/abINS_lib/tree/main/dev/validation .

from euphonic import QpointPhononModes, Quantity, Spectrum1D
from euphonic.plot import plot_1d_to_axis
import matplotlib.pyplot as plt
import numpy as np

from abinslib.almost_isotropic_incoherent import (
    calculate_almost_isotropic_incoherent_spectra,
    mantid_like_combination_spectra,
)
from abinslib.data import get_validation_data
from abinslib.displacements import Displacements
from abinslib.isotropic_incoherent import q_scaling_isotropic_incoherent_spectra
from abinslib.util import apply_weights, calculate_indirect_q2

# Get QpointPhononModes object from validation dataset
input_file = get_validation_data("ethanol_qpoint_phonon_modes.json")
modes = QpointPhononModes.from_json_file(input_file)

# Set some calculation parameters
temperature = Quantity(10, "kelvin")

# TOSCA kinematic constraints (Backscattering, Ef = 32 cm^-1).
# The odd value of backscattering angle for validation is related to
# Mantid-Abins implementation history; 135° is fine for normal calculations.
angle = 134.98885653282196 * np.pi / 180
final_energy = Quantity(32.0, "cm_1").to("hartree")

# Prepare displacement data
displacements = Displacements.from_modes(modes, temperature=temperature)
dw = displacements.to_atomic_displacements()


def plot_validation_comparison(
    abins_spectrum: Spectrum1D,
    mantid_spectrum: Spectrum1D,
    abins_label: str,
    mantid_label: str,
    title: str,
    xlim: float,
    width: Quantity = Quantity(1, "meV"),
) -> None:
    """Plot two spectra for comparison.

    Fixed-width Gaussian broadening is added as a visual aid and approximation
    to instrumental resolution.
    """
    fig, ax = plt.subplots()
    plot_1d_to_axis(abins_spectrum.broaden(width), ax, label=abins_label)
    plot_1d_to_axis(
        mantid_spectrum.broaden(width), ax, linestyle="--", label=mantid_label
    )

    ax.legend()
    ax.set_title(title + f"\n{width:~P} broadening")
    ax.set_xlim(0, xlim)
    ax.set_xlabel("Energy transfer / meV")
    ax.set_ylabel("$S(\\omega)$ / barn meV$^{-1}$")
    ax.set_yscale("log")

    fig.tight_layout()
    plt.show()


# %%
# Isotropic Fundamentals
# ----------------------
# For the fully isotropic fundamental spectrum, we apply Q^2 and Debye-Waller
# terms to the binned data, so provide Q^2 values for the bin-centre energies.

mantid_iso_file = get_validation_data("ethanol_mantid_isotropic_fundamentals.json")
mantid_iso = Spectrum1D.from_json_file(mantid_iso_file)
bins = mantid_iso.get_bin_edges(restrict_range=False)

q2_iso = calculate_indirect_q2(
    mantid_iso.get_bin_centres(),
    angle=angle,
    final_energy=final_energy,
)

spectra_iso = q_scaling_isotropic_incoherent_spectra(
    modes, displacements, dw, q2_iso, bins
)
abinslib_iso = apply_weights(spectra_iso).sum()

plot_validation_comparison(
    abins_spectrum=abinslib_iso,
    mantid_spectrum=mantid_iso,
    abins_label="q_scaling_isotropic_incoherent_spectra",
    mantid_label="Mantid-Abins isotropic fundamentals",
    title=(
        "Ethanol isotropic fundamentals:\n"
        r"$134.99^\circ E_f = 32 \mathrm{cm}^{-1}$"
    ),
    xlim=200,
)

# %%
# Almost-Isotropic Fundamentals
# -----------------------------
# For the almost-isotropic fundamentals, the $Q^2$ scaling relies on the
# individual mode frequencies rather than the binned energy transfer centres.

mantid_almost_iso_file = get_validation_data(
    "ethanol_mantid_almost_isotropic_fundamentals.json"
)
mantid_almost_iso = Spectrum1D.from_json_file(mantid_almost_iso_file)
bins = mantid_almost_iso.get_bin_edges(restrict_range=False)

q2_almost_iso = calculate_indirect_q2(
    modes.frequencies,
    angle=angle,
    final_energy=final_energy,
)

spectra_almost_iso = calculate_almost_isotropic_incoherent_spectra(
    modes, displacements, dw, q2_almost_iso, bins
)
abinslib_almost_iso = apply_weights(spectra_almost_iso).sum()

plot_validation_comparison(
    abins_spectrum=abinslib_almost_iso,
    mantid_spectrum=mantid_almost_iso,
    abins_label="calculate_almost_isotropic_incoherent_spectra",
    mantid_label="Mantid-Abins almost-isotropic fundamentals",
    title=(
        "Ethanol almost-isotropic fundamentals:\n"
        r"$134.99^\circ E_f = 32 \mathrm{cm}^{-1}$"
    ),
    xlim=200,
)

# %%
# Second-Order Spectra
# --------------------
# Second-order incoherent-approximation spectra with Mantid-like implementation.
# Q^4 scaling and Debye-Waller factor are applied to the data after binning,
# so we need to provide Q^2 values for the bins rather than combination mdoes.

mantid_second_order_file = get_validation_data("ethanol_mantid_second_order.json")
mantid_second_order = Spectrum1D.from_json_file(mantid_second_order_file)
bins = mantid_second_order.get_bin_edges(restrict_range=False)

q2_second_order = calculate_indirect_q2(
    mantid_second_order.get_bin_centres(),
    angle=angle,
    final_energy=final_energy,
)

spectra_second_order = mantid_like_combination_spectra(
    modes, displacements, dw, q2_second_order, bins
)
abinslib_second_order = apply_weights(spectra_second_order).sum()

plot_validation_comparison(
    abins_spectrum=abinslib_second_order,
    mantid_spectrum=mantid_second_order,
    abins_label="mantid_like_combination_spectra",
    mantid_label="Mantid-Abins almost-isotropic second-order",
    title=(
        "Ethanol second order:\n"
        r"$134.99^\circ E_f = 32 \mathrm{cm}^{-1}$"
    ),
    xlim=420,
)