"""Merging functions for experiment lists and reflection tables."""

from __future__ import annotations

import logging
import time
from contextlib import contextmanager
from io import StringIO

import numpy as np
import pandas as pd

import iotbx.mtz
from cctbx import miller, r_free_utils, sgtbx, uctbx
from dxtbx import flumpy
from dxtbx.model import ExperimentList
from iotbx import mtz, phil
from iotbx.reflection_file_editor import is_rfree_array
from iotbx.reflection_file_utils import get_r_free_flags_scores
from libtbx import env
from mmtbx.scaling import data_statistics

from dials.util.version import dials_version

try:
    import gemmi
except ModuleNotFoundError:
    gemmi = None

from dials.algorithms.merging.reporting import (
    MergeJSONCollector,
    MergingStatisticsData,
    make_dano_table,
)
from dials.algorithms.scaling.Ih_table import (
    _reflection_table_to_iobs,
    map_indices_to_asu,
)
from dials.algorithms.scaling.scaling_library import (
    merging_stats_from_scaled_array,
    scaled_data_as_miller_array,
)
from dials.algorithms.scaling.scaling_utilities import DialsMergingStatisticsError
from dials.algorithms.symmetry.absences.run_absences_checks import (
    run_systematic_absences_checks,
)
from dials.array_family import flex
from dials.util.export_mtz import MADMergedMTZWriter, MergedMTZWriter
from dials.util.filter_reflections import filter_reflection_table
from dials.util.resolution_analysis import resolution_cc_half

from .french_wilson import french_wilson

logger = logging.getLogger("dials")


r_free_flags_phil_scope = phil.parse(
    """
r_free_flags{
    generate = True
        .type = bool
        .short_caption = Generate R-free flags if not already present
        .style = bold noauto
    include scope cctbx.r_free_utils.generate_r_free_params_str
    extend = None
        .type = bool
        .short_caption = Extend existing R-free array(s) to full resolution range
        .style = bold noauto
    old_test_flag_value = None
      .type = int
      .short_caption = Original test flag value
      .help = Overrides automatic guess of test flag value from existing set. \
        This will usually be 1 for reflection files generated by Phenix, and \
        0 for reflection files from CCP4.  Do not change unless you're sure \
        you know what flag to use!
      .expert_level = 2
    d_eps = 0.0001
        .type = float
        .short_caption = Resolution buffer
        .expert_level = 2
    d_max = None
        .type = float
        .short_caption = The maximum resolution to which to generate R-free flags
        .expert_level = 2
    d_min = None
        .type = float
        .short_caption = The minimum resolution to which to generate R-free flags
        .expert_level = 2
    relative_to_complete_set = False
        .type = bool
        .short_caption = Generate R-free flags relative to complete set
    reference = None
        .type = path
}
"""
)


@contextmanager
def collect_html_data_from_merge():
    try:
        html_collector = MergeJSONCollector()
        yield html_collector
    finally:
        html_collector.reset()


def prepare_merged_reflection_table(
    experiments,
    reflection_table,
    d_min=None,
    d_max=None,
    partiality_threshold=0.99,
):
    """Filter the data and prepare a reflection table with merged data."""
    if (
        "inverse_scale_factor" in reflection_table
        and "intensity.scale.value" in reflection_table
    ):
        logger.info("Performing systematic absence checks on scaled data")
        reflections = filter_reflection_table(
            reflection_table,
            intensity_choice=["scale"],
            d_min=d_min,
            d_max=d_max,
            partiality_threshold=partiality_threshold,
        )
        reflections["intensity"] = reflections["intensity.scale.value"]
        reflections["variance"] = reflections["intensity.scale.variance"]
    elif "intensity.prf.value" in reflection_table:
        logger.info(
            "Performing systematic absence checks on unscaled profile-integrated data"
        )
        reflections = filter_reflection_table(
            reflection_table,
            intensity_choice=["profile"],
            d_min=d_min,
            d_max=d_max,
            partiality_threshold=partiality_threshold,
        )
        reflections["intensity"] = reflections["intensity.prf.value"]
        reflections["variance"] = reflections["intensity.prf.variance"]
    else:
        logger.info(
            "Performing systematic absence checks on unscaled summation-integrated data"
        )
        reflections = filter_reflection_table(
            reflection_table,
            intensity_choice=["sum"],
            d_min=d_min,
            d_max=d_max,
            partiality_threshold=partiality_threshold,
        )
        reflections["intensity"] = reflections["intensity.sum.value"]
        reflections["variance"] = reflections["intensity.sum.variance"]
    # now merge
    space_group = experiments[0].crystal.get_space_group()
    reflections["asu_miller_index"] = map_indices_to_asu(
        reflections["miller_index"], space_group
    )
    reflections["inverse_scale_factor"] = flex.double(reflections.size(), 1.0)
    merged = (
        _reflection_table_to_iobs(
            reflections, experiments[0].crystal.get_unit_cell(), space_group
        )
        .merge_equivalents(use_internal_variance=False)
        .array()
    )
    merged_reflections = flex.reflection_table()
    merged_reflections["intensity"] = merged.data()
    merged_reflections["variance"] = flex.pow2(merged.sigmas())
    merged_reflections["miller_index"] = merged.indices()
    return merged_reflections


class MTZDataClass:
    """Container class (i.e. Python 3.7 dataclass) for per-wavelength mtz dataset."""

    def __init__(
        self,
        wavelength=0.0,
        project_name="AUTOMATIC",
        dataset_name="NATIVE",
        crystal_name="XTAL",
        merged_array=None,
        merged_anomalous_array=None,
        amplitudes=None,
        anomalous_amplitudes=None,
        dano=None,
        multiplicities=None,
        anomalous_multiplicities=None,
        merged_half_datasets=None,
    ):
        self.wavelength = wavelength
        self.project_name = project_name
        self.dataset_name = dataset_name
        self.crystal_name = crystal_name
        self.merged_array = merged_array
        self.merged_anomalous_array = merged_anomalous_array
        self.amplitudes = amplitudes
        self.anomalous_amplitudes = anomalous_amplitudes
        self.dano = dano
        self.multiplicities = multiplicities
        self.anomalous_multiplicities = anomalous_multiplicities
        self.merged_half_datasets = merged_half_datasets


class MergedMTZCreator:
    """Creates a gemmi.Mtz object for merged data."""

    def __init__(
        self, space_group: sgtbx.space_group, unit_cell: uctbx.unit_cell
    ) -> None:
        """
        Initializes the MergedMTZCreator with the provided space group and unit cell.

        Args:
            space_group: The sgtbx.space_group object.
            unit_cell: The uctbx.unit_cell object with cell parameters.

        Returns:
            None
        """

        self.space_group = space_group
        self.unit_cell = unit_cell

        mtz = gemmi.Mtz(with_base=True)
        mtz.title = f"From {env.dispatcher_name}"
        date_str = time.strftime("%Y-%m-%d at %H:%M:%S %Z")
        if time.strftime("%Z") != "GMT":
            date_str += time.strftime("  (%Y-%m-%d at %H:%M:%S %Z)", time.gmtime())
        mtz.history += [
            f"From {dials_version()}, run on {date_str}",
        ]
        hall = space_group.type().hall_symbol()
        ops = gemmi.symops_from_hall(hall)
        mtz.spacegroup = gemmi.find_spacegroup_by_ops(ops)

        if unit_cell:
            mtz.set_cell_for_all(gemmi.UnitCell(*unit_cell.parameters()))

        self.type_to_sig_type = {"J": "Q", "F": "Q", "G": "L", "K": "M", "D": "Q"}

        self.mtz = mtz

    def add_data(
        self, mtz_datasets: list[MTZDataClass], r_free_array: miller.array = None
    ) -> None:
        """
        Adds data to the MTZ object based on the provided datasets and optional R-free array.

        Args:
            mtz_datasets (List[MTZDataClass]): The list of MTZDataClass objects containing the datasets.
            r_free_array (miller.array, optional): The optional R-free array to add.

        Returns:
            None
        """

        self._initialise_data_array(mtz_datasets)

        if r_free_array:
            self._add_column(r_free_array, label="FreeR_flag", type_char="I")

        if len(mtz_datasets) > 1:
            suffixes = [f"_WAVE{i + 1}" for i in range(len(mtz_datasets))]
        else:
            suffixes = [""]

        for i, dataset in enumerate(mtz_datasets):
            if dataset.dataset_name is None:
                dataset.dataset_name = "FROMDIALS"
            if dataset.crystal_name is None:
                dataset.crystal_name = f"crystal_{i + 2}"
            if dataset.project_name is None:
                dataset.project_name = "DIALS"
            d = self.mtz.add_dataset(dataset.dataset_name)
            d.crystal_name = dataset.crystal_name
            d.project_name = dataset.project_name
            d.wavelength = dataset.wavelength

            suffix = suffixes[i]
            if dataset.merged_array:
                self._add_column(dataset.merged_array, "IMEAN" + suffix, "J")
            if dataset.multiplicities:
                self._add_column(dataset.multiplicities, "N" + suffix, "I")
            if dataset.amplitudes:
                self._add_column(dataset.amplitudes, "F" + suffix, "F")
            if dataset.merged_anomalous_array:
                i_plus, i_minus = self._separate_anomalous(
                    dataset.merged_anomalous_array
                )
                self._add_column(i_plus, "I" + suffix + "(+)", "K")
                self._add_column(i_minus, "I" + suffix + "(-)", "K")
            if dataset.anomalous_multiplicities:
                n_plus, n_minus = self._separate_anomalous(
                    dataset.anomalous_multiplicities
                )
                self._add_column(n_plus, "N" + suffix + "(+)", "I")
                self._add_column(n_minus, "N" + suffix + "(-)", "I")
            if dataset.anomalous_amplitudes:
                f_plus, f_minus = self._separate_anomalous(dataset.anomalous_amplitudes)
                self._add_column(f_plus, "F" + suffix + "(+)", "G")
                self._add_column(f_minus, "F" + suffix + "(-)", "G")
            if dataset.dano:
                self._add_column(dataset.dano, "DANO" + suffix, "D")
            if dataset.merged_half_datasets:
                self._add_column(
                    dataset.merged_half_datasets.data1, "IHALF1" + suffix, "J"
                )
                self._add_column(
                    dataset.merged_half_datasets.data2, "IHALF2" + suffix, "J"
                )
                self._add_column(
                    dataset.merged_half_datasets.multiplicity1, "NHALF1" + suffix, "R"
                )
                self._add_column(
                    dataset.merged_half_datasets.multiplicity2, "NHALF2" + suffix, "R"
                )
        self.mtz.set_data(self.mtz_data.to_numpy())

    def _initialise_data_array(self, mtz_datasets):
        """
        Sets the Miller indices for the data array based on the given MTZ datasets.

        Args:
            mtz_datasets (List[MTZDataClass]): A list of MTZDataClass objects

        Returns:
            None

        This function creates a merged miller set by combining the indices of all the datasets.
        The indices are kept for use by the _add_column() function. A pandas DataFrame `mtz_data`
        is created with the indices of the merged set as rows and the columns "H", "K", and "L".
        Other columns will be added to the DataFrame by the _add_column() function.
        """

        miller_set = miller.set(
            crystal_symmetry=mtz_datasets[0].merged_array.crystal_symmetry(),
            indices=mtz_datasets[0].merged_array.indices().deep_copy(),
            anomalous_flag=False,
        )
        for dataset in mtz_datasets[1:]:
            indices = dataset.merged_array.indices()
            missing_isel = miller.match_indices(miller_set.indices(), indices).singles(
                1
            )
            miller_set.indices().extend(indices.select(missing_isel))

        self.indices = miller_set.indices()

        self.mtz_data = pd.DataFrame(
            flumpy.to_numpy(self.indices).astype("float32"),
            columns=["H", "K", "L"],
        )

    def _add_column(self, column: miller.array, label: str, type_char: str):
        """
        Add a data column (and associated sigmas if available) to the MTZ object.

        Args:
            column (miller.array): The column to add to the MTZ object.
            label (str): The label for the column.
            type_char (str): The type character for the column.

        This function adds column data to the MTZ object. It first adds the column
        to the most recently added dataset of the MTZ object. Then, it matches the
        Miller indices of the column data to those in MTZ object to find the pairs
        of matching indices. It selects the data from the column based on the pairs
        of matching indices and inserts them into the `mtz_data` DataFrame. If the
        column has sigmas the steps are repeated for them too.

        Returns:
            None
        """
        self.mtz.add_column(label, type_char)

        indices = column.indices()

        matches = miller.match_indices(self.indices, indices)
        pairs = matches.pairs()
        isel_i = pairs.column(0)
        isel_j = pairs.column(1)

        data = flex.double(len(self.indices), float("nan"))
        data.set_selected(isel_i, column.data().as_double().select(isel_j))

        self.mtz_data.insert(
            len(self.mtz_data.columns),
            label,
            flumpy.to_numpy(data).astype("float32"),
        )

        if column.sigmas() is None:
            return

        type_char = self.type_to_sig_type[type_char]
        self.mtz.add_column("SIG" + label, type_char)

        sigmas = flex.double(len(self.indices), float("nan"))
        sigmas.set_selected(isel_i, column.sigmas().as_double().select(isel_j))

        self.mtz_data.insert(
            len(self.mtz_data.columns),
            "SIG" + label,
            flumpy.to_numpy(sigmas).astype("float32"),
        )

    def _separate_anomalous(
        self, miller_array: miller.array
    ) -> tuple[miller.array, miller.array]:
        """
        Separates the anomalous pairs from a given Miller array to produce
        two arrays: one for the positive and one for the negative hemisphere

        Args:
            miller_array (miller.array): The input Miller array.

        Returns:
            Tuple[miller.array, miller.array]: A tuple containing the positive
            and negative anomalous arrays.
        """
        asu, matches = miller_array.match_bijvoet_mates()

        sel = matches.pairs_hemisphere_selection("+")
        sel.extend(matches.singles_hemisphere_selection("+"))

        indices = asu.indices().select(sel)
        data = asu.data().select(sel)
        sigmas = None
        if asu.sigmas() is not None:
            sigmas = asu.sigmas().select(sel)

        miller_set = miller.set(miller_array.crystal_symmetry(), indices)
        plus_array = miller.array(miller_set, data, sigmas)

        sel = matches.pairs_hemisphere_selection("-")
        sel.extend(matches.singles_hemisphere_selection("-"))

        indices = -asu.indices().select(sel)
        data = asu.data().select(sel)
        sigmas = None
        if asu.sigmas() is not None:
            sigmas = asu.sigmas().select(sel)

        miller_set = miller.set(miller_array.crystal_symmetry(), indices)
        minus_array = miller.array(miller_set, data, sigmas)

        return plus_array, minus_array


def make_merged_mtz_file(mtz_datasets, r_free_array: miller.array = None):
    """
    Make an mtz file object for the data, adding the date, time and program.

    For multi-wavelength data, each wavelength is added as a new crystal.

    Args:
        mtz_datasets: A list of MTZDataClass objects, one per wavelength of the
            experiment.

    Returns:
        A gemmi.Mtz object or an iotbx mtz object, if gemmi is not available.
    """

    if gemmi:
        return make_merged_mtz_file_with_gemmi(mtz_datasets, r_free_array)

    if len(mtz_datasets) > 1:
        writer = MADMergedMTZWriter
    else:
        writer = MergedMTZWriter

    mtz_writer = writer(
        mtz_datasets[0].merged_array.space_group(),
        mtz_datasets[0].merged_array.unit_cell(),
    )

    if r_free_array:
        mtz_writer.add_crystal(crystal_name="HKL_base", project_name="HKL_base")
        mtz_writer.add_empty_dataset(wavelength=0, name="HKL_base")
        mtz_writer.add_dataset(
            r_free_array=r_free_array,
        )

    #### Add each wavelength as a new crystal.
    for i, dataset in enumerate(mtz_datasets):
        mtz_writer.add_crystal(
            crystal_name=dataset.crystal_name, project_name=dataset.project_name
        )
        mtz_writer.add_empty_dataset(dataset.wavelength, name=dataset.dataset_name)
        mtz_writer.add_dataset(
            dataset.merged_array,
            dataset.merged_anomalous_array,
            dataset.amplitudes,
            dataset.anomalous_amplitudes,
            dataset.dano,
            dataset.multiplicities,
            dataset.anomalous_multiplicities,
            half_datasets=dataset.merged_half_datasets,
            suffix=f"_WAVE{i + 1}" if len(mtz_datasets) > 1 else "",
        )

    return mtz_writer.mtz_file


def make_merged_mtz_file_with_gemmi(mtz_datasets, r_free_array=None):
    # XXX This should replace the code in make_merged_mtz_file when
    # MergedMTZWriter and MADMergedMTZWriter are removed
    writer = MergedMTZCreator

    mtz_writer = writer(
        mtz_datasets[0].merged_array.space_group(),
        mtz_datasets[0].merged_array.unit_cell(),
    )

    mtz_writer.add_data(mtz_datasets, r_free_array)

    return mtz_writer.mtz


def merge_scaled_array(
    experiments,
    scaled_array,
    anomalous=True,
    use_internal_variance=False,
    assess_space_group=False,
    n_bins=20,
    show_additional_stats=False,
    applied_d_min=None,
):
    # assumes filtering already done and converted to combined scaled array

    # Note, merge_equivalents does not raise an error if data is unique.
    merged = scaled_array.merge_equivalents(use_internal_variance=use_internal_variance)
    merged_anom = None

    if anomalous:
        anomalous_scaled = scaled_array.as_anomalous_array()
        merged_anom = anomalous_scaled.merge_equivalents(
            use_internal_variance=use_internal_variance
        )

    # Before merge, do assessment of the space_group
    if assess_space_group:
        merged_reflections = flex.reflection_table()
        merged_reflections["intensity"] = merged.array().data()
        merged_reflections["variance"] = flex.pow2(merged.array().sigmas())
        merged_reflections["miller_index"] = merged.array().indices()
        logger.info("Running systematic absences check")
        run_systematic_absences_checks(experiments, merged_reflections)

    stats_data = MergingStatisticsData(experiments, scaled_array)

    try:
        stats, anom_stats = merging_stats_from_scaled_array(
            scaled_array,
            n_bins,
            use_internal_variance,
            additional_stats=show_additional_stats,
        )
    except DialsMergingStatisticsError as e:
        logger.error(e, exc_info=True)
    else:
        stats_data.merging_statistics_result = stats
        stats_data.anom_merging_statistics_result = anom_stats
        # try to fit resolution
        try:
            if applied_d_min is None:
                d_min = resolution_cc_half(stats, limit=0.3).d_min
            else:
                d_min = applied_d_min
        except RuntimeError as e:
            logger.debug(f"Resolution fit failed: {e}")
        else:
            max_current_res = stats.bins[-1].d_min
            if d_min and d_min - max_current_res > 0.005:
                try:
                    cut_stats, cut_anom_stats = merging_stats_from_scaled_array(
                        scaled_array.resolution_filter(d_min=d_min),
                        n_bins,
                        use_internal_variance,
                        additional_stats=show_additional_stats,
                    )
                except DialsMergingStatisticsError:
                    pass
                else:
                    if scaled_array.space_group().is_centric():
                        cut_anom_stats = None
                    stats_data.cut_merging_statistics_result = cut_stats
                    stats_data.cut_anom_merging_statistics_result = cut_anom_stats

    return merged, merged_anom, stats_data


def merge(
    experiments,
    reflections,
    d_min=None,
    d_max=None,
    combine_partials=True,
    partiality_threshold=0.4,
    best_unit_cell=None,
    anomalous=True,
    use_internal_variance=False,
    assess_space_group=False,
    n_bins=20,
    show_additional_stats=False,
):
    """
    Merge reflection table data and generate a summary of the merging statistics.

    This procedure filters the input data, merges the data (normal and optionally
    anomalous), assesses the space group symmetry and generates a summary
    of the merging statistics.

    Returns two merge_equivalents objects and a statistics summary.
    """

    logger.info("\nMerging scaled reflection data\n")
    # first filter bad reflections using dials.util.filter methods
    reflections = filter_reflection_table(
        reflections,
        intensity_choice=["scale"],
        d_min=d_min,
        d_max=d_max,
        combine_partials=combine_partials,
        partiality_threshold=partiality_threshold,
    )
    # ^ scale factor has been applied, so now set to 1.0 - okay as not
    # going to output scale factor in merged mtz.
    reflections["inverse_scale_factor"] = flex.double(reflections.size(), 1.0)
    scaled_array = scaled_data_as_miller_array(
        [reflections], experiments, best_unit_cell
    )
    return merge_scaled_array(
        experiments,
        scaled_array,
        anomalous,
        use_internal_variance,
        assess_space_group,
        n_bins,
        show_additional_stats=show_additional_stats,
        applied_d_min=d_min,
    )


def show_wilson_scaling_analysis(merged_intensities, n_residues=200):
    """
    Report the wilson statistics for a merged intensity array

    Args:
        merged_intensities: A merged miller intensity array.
        n_residues: The number of residues to use for the wilson analysis.
    """
    if not merged_intensities.space_group().is_centric():
        try:
            wilson_scaling = data_statistics.wilson_scaling(
                miller_array=merged_intensities, n_residues=n_residues
            )
        except (IndexError, RuntimeError) as e:
            logger.error(
                "\n"
                "Error encountered during Wilson statistics calculation:\n"
                "Perhaps there are too few unique reflections.\n"
                "%s",
                e,
                exc_info=True,
            )
        else:
            # Divert output through logger - do with StringIO rather than
            # info_handle else get way too much whitespace in output.
            out = StringIO()
            wilson_scaling.show(out=out)
            logger.info(out.getvalue())
            return wilson_scaling.iso_b_wilson
    return None


def truncate(
    merged_intensities,
    implementation: str = "dials",
    min_reflections: int = 200,
    fallback_to_flat_prior: bool = True,
):
    """
    Perform French-Wilson truncation procedure on merged intensities.

    Args:
        merged_intensities (miller.array): A merged miller intensity array (normal or anomalous)
        implementation (str): Choice of implementation of French & Wilson algorithm, either
                              "dials" or "cctbx"
        min_reflections (int): Minimum number of reflections to perform the French & Wilson
                               procedure
        fallback_to_flat_prior (bool): Fallback to assumption of a flat, positive prior,
                                       if the number of reflections are fewer than min_reflections,
                                       i.e. |F| = sqrt((Io+sqrt(Io**2 +2sigma**2))/2.0)

    Returns:
        (tuple): tuple containing:
            amplitudes: A normal all-positive miller amplitude array
            anom_amplitudes: An anomalous all-positive amplitude array, if the
                input array has the anomalous_flag set, else None.
            dano: The array of anomalous differences, if the input array has the
                anomalous_flag set, else None.
    """
    logger.info("\nPerforming French-Wilson treatment of scaled intensities")
    out = StringIO()
    n_refl = merged_intensities.size()
    if n_refl < min_reflections and fallback_to_flat_prior:
        logger.info(
            "Insufficient reflections for French & Wilson procedure, "
            "falling back to assumption of a flat, positive prior, i.e.: "
            "  |F| = sqrt((Io+sqrt(Io**2 +2sigma**2))/2.0)"
        )
        do_french_wilson = lambda ma: ma.enforce_positive_amplitudes()  # noqa: E731
    elif n_refl < min_reflections:
        raise ValueError(
            "Insufficient reflections for French & Wilson procedure. "
            "Either set fallback_to_flat_prior=True or truncate=False."
        )
    elif implementation == "cctbx":
        do_french_wilson = lambda ma: ma.french_wilson(log=out)  # noqa: E731
    else:
        do_french_wilson = french_wilson

    if merged_intensities.anomalous_flag():
        anom_amplitudes = do_french_wilson(merged_intensities)
        n_removed = merged_intensities.size() - anom_amplitudes.size()
        assert anom_amplitudes.is_xray_amplitude_array()
        amplitudes = anom_amplitudes.as_non_anomalous_array()
        amplitudes = amplitudes.merge_equivalents(use_internal_variance=False).array()
        dano = anom_amplitudes.anomalous_differences()
    else:
        anom_amplitudes = None
        dano = None
        amplitudes = do_french_wilson(merged_intensities)
        n_removed = merged_intensities.size() - amplitudes.size()
    logger.info("Total number of rejected intensities %s", n_removed)
    logger.debug(out.getvalue())
    return amplitudes, anom_amplitudes, dano


def merge_scaled_array_to_mtz_with_report_collection(
    params: phil.scope_extract,
    experiments: ExperimentList,
    scaled_array,
    wavelength: float | None = None,
    applied_d_min: int | None = None,
) -> tuple[mtz.object, dict]:
    if wavelength is None:
        wavelength = np.mean(
            np.array([expt.beam.get_wavelength() for expt in experiments], dtype=float)
        )
    with collect_html_data_from_merge() as collector:
        mtz_dataset = MTZDataClass(
            wavelength=wavelength,
            project_name=params.output.project_name,
            dataset_name=params.output.dataset_names[0],
            crystal_name=params.output.crystal_names[0],
        )

        merged, merged_anomalous, stats_summary = merge_scaled_array(
            experiments,
            scaled_array,
            anomalous=params.anomalous,
            assess_space_group=params.assess_space_group,
            n_bins=params.merging.n_bins,
            use_internal_variance=params.merging.use_internal_variance,
            show_additional_stats=params.output.additional_stats,
            applied_d_min=applied_d_min,
        )
        process_merged_data(
            params, mtz_dataset, merged, merged_anomalous, stats_summary
        )
        if params.r_free_flags.reference:
            r_free_array = r_free_flags_from_reference(params, [mtz_dataset])
        elif params.r_free_flags.generate:
            r_free_array = generate_r_free_flags(params, [mtz_dataset])
        else:
            r_free_array = None
        mtz = make_merged_mtz_file([mtz_dataset], r_free_array=r_free_array)
        json_data = collector.create_json()
    return mtz, json_data


def process_merged_data(params, mtz_dataset, merged, merged_anomalous, stats_summary):
    merged_array = merged.array()
    # Save the relevant data in the mtz_dataset dataclass
    # This will add the data for IMEAN/SIGIMEAN
    mtz_dataset.merged_array = merged_array
    if merged_anomalous:
        merged_anomalous_array = merged_anomalous.array()
        # This will add the data for I(+), I(-), SIGI(+), SIGI(-), N(+), N(-)
        mtz_dataset.merged_anomalous_array = merged_anomalous_array
        mtz_dataset.anomalous_multiplicities = merged_anomalous.redundancies()
    else:
        merged_anomalous_array = None
    # This will add the data for N
    mtz_dataset.multiplicities = merged.redundancies()

    if params.anomalous:
        merged_intensities = merged_anomalous_array
    else:
        merged_intensities = merged_array
    if (
        hasattr(stats_summary.merging_statistics_result, "merged_half_datasets")
        and stats_summary.merging_statistics_result.merged_half_datasets is not None
    ):
        mtz_dataset.merged_half_datasets = (
            stats_summary.merging_statistics_result.merged_half_datasets
        )

    anom_amplitudes = None
    if params.truncate:
        amplitudes, anom_amplitudes, dano = truncate(
            merged_intensities,
            implementation=params.french_wilson.implementation,
            min_reflections=params.french_wilson.min_reflections,
            fallback_to_flat_prior=params.french_wilson.fallback_to_flat_prior,
        )
        # This will add the data for F, SIGF
        mtz_dataset.amplitudes = amplitudes
        # This will add the data for F(+), F(-), SIGF(+), SIGF(-)
        mtz_dataset.anomalous_amplitudes = anom_amplitudes
        # This will add the data for DANO, SIGDANO
        mtz_dataset.dano = dano

    # print out analysis statistics
    try:
        B_iso = show_wilson_scaling_analysis(merged_intensities)
    except Exception as e:
        logger.info(e)
    else:
        stats_summary.Wilson_B_iso = B_iso

    if anom_amplitudes:
        logger.info(make_dano_table(anom_amplitudes))

    if stats_summary.merging_statistics_result:
        logger.info(stats_summary)

    if MergeJSONCollector.initiated:
        stats_summary.anomalous_amplitudes = anom_amplitudes
        MergeJSONCollector.data[mtz_dataset.wavelength] = stats_summary


def combined_miller_set(mtz_datasets: list[MTZDataClass]) -> miller.array:
    miller_set = miller.set(
        crystal_symmetry=mtz_datasets[0].merged_array.crystal_symmetry(),
        indices=mtz_datasets[0].merged_array.indices().deep_copy(),
        anomalous_flag=False,
    )
    for dataset in mtz_datasets[1:]:
        indices = dataset.merged_array.indices()
        missing_isel = miller.match_indices(miller_set.indices(), indices).singles(1)
        miller_set.indices().extend(indices.select(missing_isel))
    return miller_set


def r_free_flags_from_reference(
    params: phil.scope_extract,
    mtz_datasets: list[MTZDataClass],
) -> miller.array:
    mtz = iotbx.mtz.object(params.r_free_flags.reference)
    r_free_arrays = []
    for ma in mtz.as_miller_arrays():
        if is_rfree_array(ma, ma.info()):
            r_free_arrays.append(ma)
    flag_scores = get_r_free_flags_scores(
        miller_arrays=r_free_arrays,
        test_flag_value=params.r_free_flags.old_test_flag_value,
    )
    i = flag_scores.scores.index(max(flag_scores.scores))
    r_free_array = r_free_arrays[i]
    logger.info(
        f"Copying FreeR flags from {params.r_free_flags.reference} with labels={r_free_array.info().labels}"
    )

    if params.r_free_flags.extend:
        miller_set = combined_miller_set(mtz_datasets)
        d_max, d_min = miller_set.d_max_min()
        if params.r_free_flags.d_max:
            d_max = max(params.r_free_flags.d_max, d_max)
        if params.r_free_flags.d_min:
            d_min = min(params.r_free_flags.d_min, d_min)
        missing_set = r_free_array.complete_set(
            d_min=d_min - params.r_free_flags.d_eps,
            d_max=d_max + params.r_free_flags.d_eps,
        ).lone_set(r_free_array.map_to_asu())

        if r_free_utils.looks_like_ccp4_flags(r_free_array):
            flag_format = "ccp4"
        else:
            flag_format = "cns"
        if missing_set.size():
            missing_flags = missing_set.generate_r_free_flags(
                fraction=params.r_free_flags.fraction,
                max_free=2000,
                lattice_symmetry_max_delta=5.0,
                use_lattice_symmetry=params.r_free_flags.use_lattice_symmetry,
                n_shells=params.r_free_flags.n_shells,
                format=flag_format,
            )
            r_free_array = r_free_array.concatenate(other=missing_flags)

    return r_free_array


def generate_r_free_flags(
    params: phil.scope_extract,
    mtz_datasets: list[MTZDataClass],
) -> miller.array:
    miller_set = combined_miller_set(mtz_datasets)

    if params.r_free_flags.relative_to_complete_set:
        d_max, d_min = miller_set.d_max_min()
        if params.r_free_flags.d_max:
            d_max = max(params.r_free_flags.d_max, d_max)
        if params.r_free_flags.d_min:
            d_min = min(params.r_free_flags.d_min, d_min)
        miller_set = miller_set.complete_set(
            d_min=d_min - params.r_free_flags.d_eps,
            d_max=d_max + params.r_free_flags.d_eps,
        )

    return miller_set.generate_r_free_flags(
        fraction=params.r_free_flags.fraction,
        max_free=params.r_free_flags.max_free,
        lattice_symmetry_max_delta=params.r_free_flags.lattice_symmetry_max_delta,
        use_lattice_symmetry=params.r_free_flags.use_lattice_symmetry,
        n_shells=params.r_free_flags.n_shells,
        format="ccp4",
    )
