"""
msmapper utils
"""
import numpy as np
import h5py

from mmg_toolbox import version_info
from mmg_toolbox.nexus import nexus_writer as nw
from mmg_toolbox.utils.fitting import FitResults


MSMAPPER_VERSION = '1.9'


def create_bean(input_files, output_file, start=None, shape=None, step=None,
                output_mode=None, to_crystal=None, normalisation=None, polarisation=None,
                detector_region=None, reduce_box=None, third_axis=None,
                azi_plane_normal=None):
    """
    Create a bean file for msmapper in a temporary directory
     currently only allows a few standard inputs: hkl_start, shape and step values.
    :param input_files: list of scan file locations
    :param output_file: str location of output file
    :param start: [h, k, l] start of box (None to omit and calculate autobox)
    :param shape: [n, m, o] size of box in voxels (None to omit and calcualte autobox)
    :param step: [dh, dk, dl] step size in each direction - size of voxel in reciprocal lattice units
    :param output_mode: 'Volume_HKL' or 'Volume_Q' type of calculation
    :param to_crystal: for Volume_Q, use the crystal frame if True, or Lab frame otherwise
    :param normalisation: Monitor value to use for normalisation, e.g. 'rc'
    :param polarisation: Bool apply polarisation correction
    :param detector_region: [sx, ex, sy, ey] region of interest on detector
    :param reduce_box: Bool, reduce box to non-zero elements
    :param third_axis: [h, k, l] direction of Z-axis of voxel grid
    :param azi_plane_normal: [h, k, l] sets X-axis of voxel grid, normal to Z-axis
    :return: str file location of bean file
    """
    input_files = np.asarray(input_files, dtype=str).reshape(-1).tolist()
    # Remove empty entries
    while '' in input_files:
        input_files.remove('')

    if step is None:
        step = [0.001, 0.001, 0.001]
    else:
        step = np.asarray(step, dtype=float).reshape(-1).tolist()

    if shape is not None:
        shape = np.asarray(shape, dtype=int).reshape(-1).tolist()

    if start is not None:
        start = np.asarray(start, dtype=float).reshape(-1).tolist()

    if reduce_box is None:
        reduce_box = False

    bean = {
        "inputs": input_files,  # Filename of scan file
        "output": output_file,
        # Output filename - must be in processing directory, or somewhere you can write to
        "splitterName": "gaussian",  # one of the following strings "nearest", "gaussian", "negexp", "inverse"
        "splitterParameter": 2.0,
        # splitter's parameter is distance to half-height of the weight function.
        # If you use None or "" then it is treated as "nearest"
        "scaleFactor": 2.0,
        # the oversampling factor for each image; to ensure that are no gaps in between pixels in mapping
        "step": step,
        # a single value or list if 3 values and determines the lengths of each side of the voxels in the volume
        "start": start,  # location in HKL space of the bottom corner of the array.
        "shape": shape,  # size of the array to create for reciprocal space volume
        "reduceToNonZero": reduce_box  # True/False, if True, attempts to reduce the volume output
    }
    if output_mode:
        bean['outputMode'] = output_mode
    if MSMAPPER_VERSION > '1.8' and to_crystal is not None:
        bean['toCrystalFrame'] = to_crystal
    if normalisation:
        bean['monitorName'] = normalisation
    if polarisation:
        bean['correctPolarization'] = polarisation
    if detector_region:
        bean['region'] = detector_region
    if MSMAPPER_VERSION > '1.7' and third_axis:
        bean['thirdAxis'] = np.array(third_axis).tolist()
        bean['aziPlaneNormal'] = np.array(azi_plane_normal).tolist()
    return bean


def update_msmapper_nexus(filename: str, hkl_slice: tuple[np.ndarray, np.ndarray, np.ndarray],
                          orthogonal_axes: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None,
                          average_axes: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None,
                          fit_options: dict | None = None, h_result: FitResults | None = None,
                          k_result: FitResults | None = None, l_result: FitResults | None = None,
                          q_result: FitResults | None = None, tth_result: FitResults | None = None):

    """
    Update the NeXus file generated by msmapper with additional analysis data.

    This function opens an existing NeXus file produced by **msmapper** and appends
    analysis-related datasets and metadata. Typical additions include HKL slices,
    optional orthogonal/averaged axis data, fitting results (e.g., peak fits along
    H, K, L, Q, and 2θ), and any options/configuration used during fitting.

    The function is non-destructive with respect to the original raw datasets
    written by msmapper; new fields are appended under appropriate NXgroups.

    :param filename: Path to the NeXus (.nxs/.h5) file to be updated in-place.
    :type filename: str

    :param hkl_slice: slices of the reciprocal space volume in each of h, k, l directions
    :type hkl_slice: tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]

    :param orthogonal_axes: Optional tuple of three arrays representing an
        orthogonal coordinate system (e.g., lab-frame axes) mapped to the same points
        as the HKL slice. If provided, these are stored alongside HKL for reference.
        Order is expected to be (axis_x, axis_y, axis_z) and shape must match volume.
    :type orthogonal_axes: tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray] | None

    :param average_axes: Optional tuple of three arrays representing averaged axes
        (e.g., binned or smoothed directions) corresponding to the HKL slice. Order
        is (wavevector==|Q|, two-theta [Deg], averaged intensities).
    :type average_axes: tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray] | None

    :param fit_options: Optional dictionary of fitting options used to produce the
        FitResults (e.g., model type, bounds, initial guesses, weighting). This is
        serialized into the NeXus file for provenance.
    :type fit_options: dict | None

    :param h_result: Peak Fit results for the H direction (e.g., peak positions, widths,
        amplitudes, residuals). If provided, these are written under the appropriate
        group in the NeXus file.
    :type h_result: FitResults | None

    :param k_result: Fit results for the K direction.
    :type k_result: FitResults | None

    :param l_result: Fit results for the L direction.
    :type l_result: FitResults | None

    :param q_result: Fit results for Q (|Q| or wavevector magnitude).
    :type q_result: FitResults | None

    :param tth_result: Fit results for two-theta (2θ).
    :type tth_result: FitResults | None

    :returns: None. The function updates the specified NeXus file on disk.
    :rtype: None

    :raises FileNotFoundError: If ``filename`` does not exist.
    :raises ValueError: If array lengths in ``hkl_slice`` (or provided axes) do not match,
        or if the input shapes are incompatible with the expected structure.
    :raises IOError: If the file cannot be opened for read/write.
    :raises RuntimeError: If writing the analysis groups/datasets fails due to
        backend/format issues.

    **Notes**
    - The function assumes the target file already follows the msmapper NeXus layout
      and will append analysis under a dedicated NXgroup (e.g., ``/entry/analysis``).

    **Example**
    ::
        import numpy as np

        H = np.linspace(-1.0, 1.0, 201)
        K = np.zeros_like(H)
        L = np.zeros_like(H)

        # Optional axes (example placeholders)
        ox = H.copy()
        oy = K.copy()
        oz = L.copy()

        update_msmapper_nexus(
            filename="scan_001.nxs",
            hkl_slice=(H, K, L),
            orthogonal_axes=(ox, oy, oz),
            average_axes=None,
            fit_options={"model": "gaussian", "max_iter": 500},
            h_result=h_fit,  # instances of FitResults
            k_result=None,
            l_result=None,
            q_result=q_fit,
            tth_result=None
        )
    """

    h_slice, k_slice, l_slice = hkl_slice
    qx, qy, qz = orthogonal_axes
    wavevector, tth, mag_average = average_axes

    print(f"Updating NeXus file: {filename}")

    with h5py.File(filename, 'a') as hdf:
        entry = nw.add_nxentry(hdf, 'analysis', definition=None, default=True)
        nw.add_nxprocess(entry, 'process', program='mmg_toolbox', version=version_info())

        # hkl
        h = hdf['/processed/reciprocal_space/h-axis'][:]
        k = hdf['/processed/reciprocal_space/k-axis'][:]
        l = hdf['/processed/reciprocal_space/l-axis'][:]
        h_axis = nw.add_nxdata(entry, 'h_axis', ['h'], 'intensity')
        h_axis['h'] = h5py.SoftLink('/processed/reciprocal_space/h-axis')
        nw.add_nxfield(h_axis, 'intensity', h_slice)

        k_axis = nw.add_nxdata(entry, 'k_axis', ['k'], 'intensity')
        k_axis['k'] = h5py.SoftLink('/processed/reciprocal_space/k-axis')
        nw.add_nxfield(k_axis, 'intensity', k_slice)

        l_axis = nw.add_nxdata(entry, 'l_axis', ['l'], 'intensity', default=True)
        l_axis['l'] = h5py.SoftLink('/processed/reciprocal_space/l-axis')
        nw.add_nxfield(l_axis, 'intensity', l_slice)

        # Q
        orthog = nw.add_nxdata(entry, 'orthogonal_axes', ['Qx', 'Qy', 'Qz'], 'volume', 'weight')
        nw.add_nxfield(orthog, 'Qx', qx, units='1/angstrom')
        nw.add_nxfield(orthog, 'Qy', qy, units='1/angstrom')
        nw.add_nxfield(orthog, 'Qz', qz, units='1/angstrom')
        orthog['volume'] = h5py.SoftLink('/processed/reciprocal_space/volume')
        orthog['weight'] = h5py.SoftLink('/processed/reciprocal_space/weight')
        nw.add_attr(orthog, Qx_indices=1, Qy_indices=2, Qz_indices=0, default_slice=len(l) // 2)

        # volume averaged at each magnitude
        qvsi = nw.add_nxdata(entry, 'wavevector', ['q'], 'intensity')
        nw.add_nxfield(qvsi, 'q', wavevector, units='1/angstrom')
        nw.add_nxfield(qvsi, 'intensity', mag_average)

        tthvsi = nw.add_nxdata(entry, 'twotheta', ['tth'], 'intensity')
        nw.add_nxfield(tthvsi, 'tth', tth, units='deg')
        nw.add_nxfield(tthvsi, 'intensity', mag_average)

        if fit_options:
            fit_proc = nw.add_nxprocess(entry, 'peak_fit',
                                        program='mmg_toolbox.utils.fitting.multipeakfit',
                                        **fit_options)
            nw.add_nxnote(fit_proc, 'fit_results_h', 'fit results for h-axis', data=str(h_result))
            nw.add_nxparameters(fit_proc, 'fit_data_h', **h_result.results())
            nw.add_nxnote(fit_proc, 'fit_results_k', 'fit results for k-axis', data=str(k_result))
            nw.add_nxparameters(fit_proc, 'fit_data_k', **k_result.results())
            nw.add_nxnote(fit_proc, 'fit_results_l', 'fit results for l-axis', data=str(l_result))
            nw.add_nxparameters(fit_proc, 'fit_data_l', **l_result.results())
            nw.add_nxnote(fit_proc, 'fit_results_q', 'fit results for wavevector magnitude |Q|', data=str(q_result))
            nw.add_nxparameters(fit_proc, 'fit_data_q', **q_result.results())
            nw.add_nxnote(fit_proc, 'fit_results_tth', 'fit results for remapped two-theta', data=str(tth_result))
            nw.add_nxparameters(fit_proc, 'fit_data_tth', **tth_result.results())

            nw.add_nxfield(h_axis, 'fit', h_result.fit_data(h, ntimes=1)[1], add_to_signal=True)
            nw.add_nxfield(k_axis, 'fit', k_result.fit_data(k, ntimes=1)[1], add_to_signal=True)
            nw.add_nxfield(l_axis, 'fit', l_result.fit_data(l, ntimes=1)[1], add_to_signal=True)
            nw.add_nxfield(qvsi, 'fit', q_result.fit_data(wavevector, ntimes=1)[1], add_to_signal=True)
            nw.add_nxfield(tthvsi, 'fit', tth_result.fit_data(tth, ntimes=1)[1], add_to_signal=True)

