from __future__ import annotations

import json

import procrunner
import pytest

import iotbx.cif
from cctbx import sgtbx, uctbx
from dxtbx.model.experiment_list import ExperimentListFactory
from dxtbx.serialize import load
from iotbx import mtz

from dials.array_family import flex
from dials.command_line.slice_sequence import slice_experiments, slice_reflections
from dials.util.multi_dataset_handling import assign_unique_identifiers


def run_export(export_format, dials_data, tmp_path):
    result = procrunner.run(
        [
            "dials.export",
            "format=" + export_format,
            dials_data("centroid_test_data", pathlib=True) / "experiments.json",
            dials_data("centroid_test_data", pathlib=True) / "integrated.pickle",
        ],
        working_directory=tmp_path,
    )
    assert not result.returncode and not result.stderr
    assert (tmp_path / f"integrated.{export_format}").is_file()


def test_nxs(dials_data, tmp_path):
    run_export("nxs", dials_data, tmp_path)


def test_mtz(dials_data, tmp_path):
    result = procrunner.run(
        [
            "dials.export",
            "format=mtz",
            "project_name=ham",
            "crystal_name=spam",
            dials_data("centroid_test_data", pathlib=True) / "experiments.json",
            dials_data("centroid_test_data", pathlib=True) / "integrated.pickle",
        ],
        working_directory=tmp_path,
    )
    assert not result.returncode and not result.stderr
    assert (tmp_path / "integrated.mtz").is_file()
    mtz_obj = mtz.object(str(tmp_path / "integrated.mtz"))
    assert mtz_obj.crystals()[1].name() == "spam"
    assert mtz_obj.crystals()[1].project_name() == "ham"


def test_mtz_recalculated_cell(dials_data, tmp_path):
    # First run dials.two_theta_refine to ensure that the crystals have
    # recalculated_unit_cell set
    scaled_expt = (
        dials_data("x4wide_processed", pathlib=True) / "AUTOMATIC_DEFAULT_scaled.expt"
    )
    scaled_refl = (
        dials_data("x4wide_processed", pathlib=True) / "AUTOMATIC_DEFAULT_scaled.refl"
    )
    result = procrunner.run(
        ["dials.two_theta_refine", scaled_expt, scaled_refl],
        working_directory=tmp_path,
    )
    assert (tmp_path / "refined_cell.expt").is_file()
    refined_expt = load.experiment_list(
        tmp_path / "refined_cell.expt", check_format=False
    )
    ttr_cell = refined_expt.crystals()[0].get_recalculated_unit_cell()

    d_min = 1.3