# SPDX-License-Identifier: MIT

import unittest
from mock import MagicMock, patch
import os
import gc
from pyFAI import detectors

import numpy as np

from qtpy import QtWidgets, QtCore
from qtpy.QtTest import QTest

from ..utility import (
    QtTest,
    unittest_data_path,
    click_button,
    click_checkbox,
    delete_if_exists,
)
from ...model.DioptasModel import DioptasModel
from ...controller.CalibrationController import (
    CalibrationController,
    get_available_detectors,
)
from ...widgets.CalibrationWidget import CalibrationWidget

class TestCalibrationController(QtTest):
    def setUp(self):
        # Mocking the functions which will block the unittest for some reason...
        # Use patch to ensure proper cleanup
        self.processEvents_patcher = patch.object(QtWidgets.QApplication, 'processEvents', MagicMock())
        self.setValue_patcher = patch.object(QtWidgets.QProgressDialog, 'setValue', MagicMock())
        self.processEvents_patcher.start()
        self.setValue_patcher.start()

        self.model = DioptasModel()

        self.widget = CalibrationWidget()
        self.controller = CalibrationController(
            widget=self.widget, dioptas_model=self.model
        )

    def tearDown(self):
        # Stop the patchers to restore original behavior
        self.processEvents_patcher.stop()
        self.setValue_patcher.stop()

        delete_if_exists(os.path.join(unittest_data_path, "detector_with_spline.h5"))

        # Clean up widgets
        self.widget.close()
        self.widget.deleteLater()
        del self.widget
        del self.controller
        del self.model
        gc.collect()

    def mock_integrate_functions(self):
        self.model.calibration_model.integrate_1d = MagicMock(
            return_value=([np.linspace(0, 100), np.linspace(0, 100)])
        )
        self.model.calibration_model.integrate_2d = MagicMock()

    def test_load_detector(self):
        detector_names, detector_classes = get_available_detectors()
        det_ind = 9
        self.widget.detectors_cb.setCurrentIndex(
            det_ind + 3
        )  # +3 since there is also the custom element at 0
        # and 2 separators
        self.assertIsInstance(
            self.model.calibration_model.detector, detector_classes[det_ind]
        )

        detector_gb = (
            self.widget.calibration_control_widget.calibration_parameters_widget.detector_gb
        )
        self.assertAlmostEqual(
            float(detector_gb.pixel_width_txt.text()) * 1e-6,
            self.model.calibration_model.orig_pixel1,
        )
        self.assertAlmostEqual(
            float(detector_gb.pixel_height_txt.text()) * 1e-6,
            self.model.calibration_model.orig_pixel2,
        )

        self.assertFalse(detector_gb.pixel_width_txt.isEnabled())
        self.assertFalse(detector_gb.pixel_width_txt.isEnabled())

        self.widget.detectors_cb.setCurrentIndex(0)
        self.assertNotIsInstance(
            self.model.calibration_model.detector, detector_classes[det_ind]
        )
        self.assertAlmostEqual(
            float(detector_gb.pixel_width_txt.text()) * 1e-6,
            self.model.calibration_model.orig_pixel1,
        )
        self.assertAlmostEqual(
            float(detector_gb.pixel_height_txt.text()) * 1e-6,
            self.model.calibration_model.orig_pixel2,
        )
        self.assertTrue(detector_gb.pixel_width_txt.isEnabled())
        self.assertTrue(detector_gb.pixel_width_txt.isEnabled())

    def test_load_detector_transform_and_reset(self):
        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
        )
        QTest.mouseClick(self.widget.load_calibration_btn, QtCore.Qt.LeftButton)

        detector_gb = (
            self.widget.calibration_control_widget.calibration_parameters_widget.detector_gb
        )
        detector_gb.detector_cb.setCurrentIndex(
            detector_gb.detector_cb.findText("Pilatus CdTe 1M")
        )

        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "CeO2_Pilatus1M.tif")
        )
        QTest.mouseClick(self.widget.load_img_btn, QtCore.Qt.LeftButton)

        QTest.mouseClick(self.widget.rotate_m90_btn, QtCore.Qt.LeftButton)
        QTest.mouseClick(self.widget.rotate_m90_btn, QtCore.Qt.LeftButton)

        detector_gb.detector_cb.setCurrentIndex(
            detector_gb.detector_cb.findText("Custom")
        )

        QTest.mouseClick(self.widget.rotate_m90_btn, QtCore.Qt.LeftButton)

    def test_load_detector_from_file(self):
        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "detector.h5")
        )
        click_button(self.widget.detector_load_btn)

        self.assertAlmostEqual(self.model.calibration_model.orig_pixel1, 100e-6)

        self.assertTrue(self.widget.detector_reset_btn.isEnabled())
        self.widget.show()  # to check visibility
        self.assertFalse(self.widget.detectors_cb.isVisible())
        self.assertTrue(self.widget.detector_name_lbl.isVisible())
        self.assertEqual(self.widget.detector_name_lbl.text(), "detector.h5")
        detector_gb = (
            self.widget.calibration_control_widget.calibration_parameters_widget.detector_gb
        )
        self.assertFalse(detector_gb.pixel_width_txt.isEnabled())
        self.assertFalse(detector_gb.pixel_height_txt.isEnabled())
        self.assertAlmostEqual(self.widget.get_pixel_size()[0], 100e-6)

        click_button(self.widget.detector_reset_btn)
        self.assertTrue(self.widget.detectors_cb.isVisible())
        self.assertFalse(self.widget.detector_name_lbl.isVisible())
        self.assertEqual(self.widget.detectors_cb.currentIndex(), 0)
        self.assertTrue(detector_gb.pixel_width_txt.isEnabled())
        self.assertTrue(detector_gb.pixel_height_txt.isEnabled())

    def test_load_detector_with_distortion(self):
        # create detector and save it
        spline_detector = detectors.Detector()
        spline_detector.splinefile = os.path.join(
            unittest_data_path, "distortion", "f4mnew.spline"
        )
        spline_detector.save(
            os.path.join(unittest_data_path, "detector_with_spline.h5")
        )

        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "detector_with_spline.h5")
        )
        click_button(self.widget.detector_load_btn)

        self.assertEqual(self.widget.spline_filename_txt.text(), "from Detector")

    def test_load_detector_and_image_with_different_dimension(self):
        # set the detector
        self.widget.detectors_cb.setCurrentIndex(
            self.widget.detectors_cb.findText("Pilatus CdTe 1M")
        )

        # load calibration
        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
        )
        QTest.mouseClick(self.widget.load_calibration_btn, QtCore.Qt.LeftButton)

        # load image file with different dimension than detector, which should automatically
        # reset the detector to a custom detector
        QtWidgets.QMessageBox.critical = MagicMock()
        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "LaB6_40keV_MarCCD.tif")
        )
        click_button(self.widget.load_img_btn)

        QtWidgets.QMessageBox.critical.assert_called_once()
        self.assertEqual(self.widget.detectors_cb.currentText(), "Custom")

    def test_automatic_calibration(self):
        self.mock_integrate_functions()
        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "LaB6_40keV_MarCCD.tif")
        )
        QTest.mouseClick(self.widget.load_img_btn, QtCore.Qt.LeftButton)
        self.controller.search_peaks(1179.6, 1129.4)
        self.controller.search_peaks(1268.5, 1119.8)
        self.controller.widget.sv_wavelength_txt.setText("0.31")
        self.controller.widget.sv_distance_txt.setText("200")
        self.controller.widget.set_pixel_size(79e-6, 79e-6)
        calibrant_index = self.widget.calibrant_cb.findText("LaB6")
        self.controller.widget.calibrant_cb.setCurrentIndex(calibrant_index)
        #
        QTest.mouseClick(self.widget.calibrate_btn, QtCore.Qt.LeftButton)
        self.app.processEvents()
        self.model.calibration_model.integrate_1d.assert_called_once()
        self.model.calibration_model.integrate_2d.assert_called_once()
        # Progress dialog should be updated during calibration
        self.assertGreater(QtWidgets.QProgressDialog.setValue.call_count, 0)

        calibration_parameter = (
            self.model.calibration_model.get_calibration_parameter()[0]
        )
        self.assertAlmostEqual(calibration_parameter["dist"], 0.1968, places=3)

    def test_searching_peaks_automatic_increase(self):
        self.mock_integrate_functions()
        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "LaB6_40keV_MarCCD.tif")
        )
        QTest.mouseClick(self.widget.load_img_btn, QtCore.Qt.LeftButton)
        self.controller.search_peaks(1179.6, 1129.4)
        self.controller.search_peaks(1268.5, 1119.8)
        self.assertEqual(3, self.widget.peak_num_sb.value())

    def test_searching_peaks_manual_index_assignment(self):
        self.mock_integrate_functions()
        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "LaB6_40keV_MarCCD.tif")
        )
        QTest.mouseClick(self.widget.load_img_btn, QtCore.Qt.LeftButton)

        click_checkbox(self.widget.automatic_peak_num_inc_cb)
        self.assertFalse(self.widget.automatic_peak_num_inc_cb.isChecked())

        self.controller.search_peaks(1179.6, 1129.4)
        self.controller.search_peaks(1268.5, 1119.8)
        self.assertEqual(1, self.widget.peak_num_sb.value())

    def test_splines(self):
        self.mock_integrate_functions()
        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "distortion", "f4mnew.spline")
        )
        click_button(self.widget.load_spline_btn)

        self.assertIsNotNone(self.model.calibration_model.distortion_spline_filename)
        self.assertEqual(self.widget.spline_filename_txt.text(), "f4mnew.spline")
        self.assertTrue(self.widget.spline_reset_btn.isEnabled())
        #
        click_button(self.widget.spline_reset_btn)
        self.assertIsNone(self.model.calibration_model.distortion_spline_filename)
        self.assertEqual(self.widget.spline_filename_txt.text(), "None")

    def test_loading_and_saving_of_calibration_files(self):
        self.mock_integrate_functions()
        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "LaB6_40keV_MarCCD.poni")
        )
        QTest.mouseClick(self.widget.load_calibration_btn, QtCore.Qt.LeftButton)
        QtWidgets.QFileDialog.getSaveFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "calibration.poni")
        )
        QTest.mouseClick(self.widget.save_calibration_btn, QtCore.Qt.LeftButton)
        self.assertTrue(
            os.path.exists(os.path.join(unittest_data_path, "calibration.poni"))
        )
        os.remove(os.path.join(unittest_data_path, "calibration.poni"))

    def test_selecting_configuration_updates_parameter_display(self):
        self.mock_integrate_functions()
        calibration1 = {
            "dist": 0.2,
            "poni1": 0.08,
            "poni2": 0.081,
            "rot1": 0.0043,
            "rot2": 0.002,
            "rot3": 0.001,
            "pixel1": 7.9e-5,
            "pixel2": 7.9e-5,
            "wavelength": 0.3344,
            "polarization_factor": 0.99,
        }
        calibration2 = {
            "dist": 0.3,
            "poni1": 0.04,
            "poni2": 0.021,
            "rot1": 0.0053,
            "rot2": 0.002,
            "rot3": 0.0013,
            "pixel1": 7.4e-5,
            "pixel2": 7.6e-5,
            "wavelength": 0.31,
            "polarization_factor": 0.98,
        }

        self.model.calibration_model.set_pyFAI(calibration1)
        self.model.add_configuration()
        self.model.calibration_model.set_pyFAI(calibration2)

        self.model.select_configuration(0)

        model_calibration = self.model.configurations[
            0
        ].calibration_model.pattern_geometry.getPyFAI()
        del model_calibration["detector"]
        for key in ("splineFile", "splinefile", "max_shape", "orientation"):
            model_calibration.pop(key, None)

        current_displayed_calibration = self.widget.get_pyFAI_parameter()
        del current_displayed_calibration["polarization_factor"]
        self.assertEqual(model_calibration, current_displayed_calibration)

        self.model.select_configuration(1)
        model_calibration = self.model.configurations[
            1
        ].calibration_model.pattern_geometry.getPyFAI()
        del model_calibration["detector"]
        for key in ("splineFile", "splinefile", "max_shape", "orientation"):
            model_calibration.pop(key, None)
        current_displayed_calibration = self.widget.get_pyFAI_parameter()
        del current_displayed_calibration["polarization_factor"]

        self.assertEqual(model_calibration, current_displayed_calibration)

    @unittest.skip("Does not work for unknown reasons")
    def test_calibrant_with_small_set_of_d_spacings(self):
        self.mock_integrate_functions()
        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "LaB6_40keV_MarCCD.tif")
        )
        QTest.mouseClick(self.widget.load_img_btn, QtCore.Qt.LeftButton)
        self.controller.search_peaks(1179.6, 1129.4)
        self.controller.search_peaks(1268.5, 1119.8)
        calibrant_index = self.widget.calibrant_cb.findText("CuO")
        self.controller.widget.calibrant_cb.setCurrentIndex(calibrant_index)
        QtWidgets.QMessageBox.critical = MagicMock()
        click_button(self.widget.calibrate_btn)
        QtWidgets.QMessageBox.critical.assert_called_once()

    def test_loading_calibration_without_an_image_before(self):
        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "LaB6_40keV_MarCCD.poni")
        )
        QTest.mouseClick(self.widget.load_calibration_btn, QtCore.Qt.LeftButton)
        self.widget.get_pyFAI_parameter()  # would cause error if GUI not updated
        self.widget.get_fit2d_parameter()  # would cause error if GUI not updated

    def test_detector_rotation_does_not_emit_reset_detector_signal(self):
        QtWidgets.QFileDialog.getOpenFileName = MagicMock(
            return_value=os.path.join(unittest_data_path, "LaB6_40keV_MarCCD.tif")
        )
        click_button(self.widget.load_img_btn)
        self.model.calibration_model.detector_reset.emit = MagicMock()

        click_button(self.widget.rotate_m90_btn)
        self.model.calibration_model.detector_reset.emit.assert_not_called()


def test_click_calibrate_without_defined_points(
    calibration_controller, calibration_model
):
    QtWidgets.QMessageBox.critical = MagicMock()
    calibration_model.calibrate = MagicMock()
    click_button(calibration_controller.widget.calibrate_btn)
    assert QtWidgets.QMessageBox.critical.called
    assert not calibration_model.calibrate.called


def test_click_refine_without_defined_points(calibration_controller, calibration_model):
    QtWidgets.QMessageBox.critical = MagicMock()
    click_button(calibration_controller.widget.refine_btn)
    assert QtWidgets.QMessageBox.critical.called
