# encoding: utf-8
"""Collection emulator, calling an external program to generate data images.
Written originally for Global Phasing simcal,
but could be made to work with other systems

License:

This file is part of MXCuBE.

MXCuBE is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

MXCuBE is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with MXCuBE. If not, see <https://www.gnu.org/licenses/>.
"""

import os
import re
import subprocess
from collections import OrderedDict

import f90nml

from mxcubecore import HardwareRepository as HWR
from mxcubecore.HardwareObjects.mockup.CollectMockup import CollectMockup
from mxcubecore.TaskUtils import task
from mxcubecore.utils import conversion

__copyright__ = """ Copyright © 2017 - 2019 by Global Phasing Ltd. """
__license__ = "LGPLv3+"
__author__ = "Rasmus H Fogh"


class CollectEmulator(CollectMockup):
    def __init__(self, name):
        CollectMockup.__init__(self, name)

        self._counter = 1

    def _get_simcal_input(self, data_collect_parameters, crystal_data):
        """Get ordered dict with simcal input from available data"""

        # Set up and add crystal data
        result = OrderedDict()
        setup_data = result["setup_list"] = crystal_data
        instrument_data = HWR.beamline.gphl_workflow.instrument_data

        #
        # # update with instrument data
        #
        sweep_count = len(data_collect_parameters["oscillation_sequence"])

        # Setting parameters in order (may not be necessary, but ...)
        # Missing: *mu*
        remap = {
            "beam": "nominal_beam_dir",
            "det_coord_def": "det_org_dist",
            "cone_s_height": "cone_height",
        }
        tags = (
            "lambda_sd",
            "beam",
            "beam_sd_deg",
            "pol_plane_n",
            "pol_frac",
            "d_sensor",
            "min_zeta",
            "det_name",
            "det_x_axis",
            "det_y_axis",
            "det_qx",
            "det_qy",
            "det_nx",
            "det_ny",
            # "det_org_x",
            # "det_org_y",
            "det_coord_def",
        )
        for tag in tags:
            val = instrument_data.get(remap.get(tag, tag))
            if val is not None:
                setup_data[tag] = val

        (
            setup_data["det_org_x"],
            setup_data["det_org_y"],
        ) = HWR.beamline.detector.get_beam_position()

        ll0 = list(HWR.beamline.gphl_workflow.rotation_axes.values())
        setup_data["omega_axis"] = ll0[0]
        setup_data["kappa_axis"] = ll0[1]
        setup_data["phi_axis"] = ll0[2]
        ll0 = list(HWR.beamline.gphl_workflow.translation_axes.values())
        setup_data["trans_x_axis"] = ll0[0]
        setup_data["trans_y_axis"] = ll0[1]
        setup_data["trans_z_axis"] = ll0[2]
        tags = (
            "cone_radius",
            "cone_s_height",
            "beam_stop_radius",
            "beam_stop_s_length",
            "beam_stop_s_distance",
        )
        for tag in tags:
            val = instrument_data.get(remap.get(tag, tag))
            if val is not None:
                setup_data[tag] = val

        # Add/overwrite parameters from emulator configuration
        conv = conversion.convert_string_value
        for key, val in self.config.simcal_parameters.items():
            setup_data[key] = conv(val)

        setup_data["n_vertices"] = 0
        setup_data["n_triangles"] = 0
        setup_data["n_orients"] = 0
        setup_data["n_sweeps"] = sweep_count

        # Add segments
        segments = HWR.beamline.gphl_workflow.detector_segments
        if segments:
            if isinstance(segments, dict):
                segment_count = 1
            else:
                segment_count = len(segments)
            setup_data["n_segments"] = segment_count
            result["segment_list"] = segments

        # Adjustments
        val = instrument_data.get("beam")
        if val:
            setup_data["beam"] = val

        # update with diffractcal data
        # TODO check that this works also for updating segment list
        fp0 = HWR.beamline.gphl_workflow.file_paths.get("diffractcal_file")
        if os.path.isfile(fp0):
            diffractcal_data = f90nml.read(fp0)["sdcp_instrument_list"]
            for tag in setup_data.keys():
                val = diffractcal_data.get(tag)
                if val is not None:
                    setup_data[tag] = val
            ll0 = diffractcal_data["gonio_axis_dirs"]
            setup_data["omega_axis"] = ll0[:3]
            setup_data["kappa_axis"] = ll0[3:6]
            setup_data["phi_axis"] = ll0[6:]

        # get resolution limit and detector distance
        detector_distance = data_collect_parameters.get("detector_distance", 0.0)
        if not detector_distance:
            dd = data_collect_parameters.get("resolution")
            # resolution may not be set - if so you should take the current value
            if dd:
                resolution = dd.get("upper")
                if resolution:
                    self.set_resolution(resolution)
            detector_distance = HWR.beamline.detector.distance.get_value()
        # Add sweeps
        sweeps = []
        compress_data = False
        for osc in data_collect_parameters["oscillation_sequence"]:
            motors = data_collect_parameters["motors"]
            sweep = OrderedDict()

            energy = (
                data_collect_parameters.get("energy") or HWR.beamline.energy.get_value()
            )
            sweep["lambda"] = conversion.HC_OVER_E / energy
            sweep["res_limit"] = setup_data["res_limit_def"]
            sweep["exposure"] = osc["exposure_time"]
            ll0 = list(HWR.beamline.gphl_workflow.translation_axes)
            sweep["trans_xyz"] = list(motors.get(x) or 0.0 for x in ll0)
            sweep["det_coord"] = detector_distance
            # NBNB hardwired for omega scan TODO
            sweep["axis_no"] = 3
            sweep["omega_deg"] = osc["start"]
            # NB kappa and phi are overwritten from the motors dict, if set there
            sweep["kappa_deg"] = osc["kappaStart"]
            sweep["phi_deg"] = osc["phiStart"]
            sweep["step_deg"] = osc["range"]
            sweep["n_frames"] = osc["number_of_images"]
            sweep["image_no"] = osc["start_image_number"]
            # self.make_image_file_template(data_collect_parameters, suffix='cbf')

            # Extract format statement from template,
            # and convert to fortran format
            text_type = conversion.text_type
            template = text_type(data_collect_parameters["fileinfo"]["template"])
            ss0 = re.search("(%[0-9]+d)", template).group(0)
            template = template.replace(ss0, "?" * int(ss0[1:-1]))
            name_template = os.path.join(
                text_type(data_collect_parameters["fileinfo"]["directory"]),
                template,
                # data_collect_parameters['fileinfo']['template']
            )
            # We still use the normal name template for compressed data
            if name_template.endswith(".gz"):
                compress_data = True
                name_template = name_template[:-3]
            sweep["name_template"] = name_template

            # Overwrite kappa and phi from motors - if set
            val = motors.get("kappa")
            if val is not None:
                sweep["kappa_deg"] = val
            val = motors.get("kappa_phi")
            if val is not None:
                sweep["phi_deg"] = val

            # Skipped: spindle_deg=0.0, two_theta_deg=0.0, mu_air=-1, mu_sensor=-1

            sweeps.append(sweep)

        if len(sweeps) == 1:
            # NBNB in current code we can have only one sweep here,
            # but it will work for multiple
            result["sweep_list"] = sweep
        else:
            result["sweep_list"] = sweeps
        return result, compress_data

    @task
    def data_collection_hook(self):
        """Spawns data emulation using gphl simcal"""
        data_model = HWR.beamline.gphl_workflow._queue_entry.get_data_model()
        if data_model.skip_collection:
            return

        data_collect_parameters = self.current_dc_parameters

        if not HWR.beamline.gphl_workflow:
            raise ValueError("Emulator requires GPhL workflow installation")
        gphl_connection = HWR.beamline.gphl_connection
        if not gphl_connection:
            raise ValueError("Emulator requires GPhL connection installation")

        # Get program locations
        simcal_executive = gphl_connection.get_executable("simcal")
        simcal_licence_dir = (
            gphl_connection.get_bdg_licence_dir("simcal")
            or gphl_connection.config.software_paths["GPHL_INSTALLATION"]
        )
        # # Get environmental variables.
        envs = {"autoPROC_home": simcal_licence_dir}
        GPHL_XDS_PATH = gphl_connection.config.software_paths.get("GPHL_XDS_PATH")
        if GPHL_XDS_PATH:
            envs["GPHL_XDS_PATH"] = GPHL_XDS_PATH
        GPHL_CCP4_PATH = gphl_connection.config.software_paths.get("GPHL_CCP4_PATH")
        if GPHL_CCP4_PATH:
            envs["GPHL_CCP4_PATH"] = GPHL_CCP4_PATH
        text_type = conversion.text_type
        for tag, val in self.config.environment_variables.items():
            envs[text_type(tag)] = text_type(val)

        # get crystal data
        crystal_data, hklfile = HWR.beamline.gphl_workflow.get_emulation_crystal_data()

        input_data, compress_data = self._get_simcal_input(
            data_collect_parameters, crystal_data
        )

        # NB outfile is the echo output of the input file;
        # image files templates are set in the input file
        file_info = data_collect_parameters["fileinfo"]
        if not os.path.exists(file_info["directory"]):
            os.makedirs(file_info["directory"])
        infile = os.path.join(
            file_info["directory"], "simcal_in_%s.nml" % self._counter
        )

        f90nml.write(input_data, infile, force=True)
        outfile = os.path.join(
            file_info["directory"], "simcal_out_%s.nml" % self._counter
        )
        logfile = os.path.join(
            file_info["directory"], "simcal_log_%s.txt" % self._counter
        )
        self._counter += 1
        command_list = [
            simcal_executive,
            "--input",
            infile,
            "--output",
            outfile,
            "--hkl",
            hklfile,
        ]

        for tag, val in self.config.simcal_options.items():
            command_list.extend(conversion.command_option(tag, val, prefix="--"))
        self.log.info("Executing command: %s", " ".join(command_list))
        self.log.info("Executing environment: %s", sorted(envs.items()))

        if compress_data:
            command_list.append("--gzip-img")

        fp1 = open(logfile, "w")
        fp2 = subprocess.STDOUT
        # resource.setrlimit(resource.RLIMIT_STACK, (-1,-1))

        def set_ulimit():
            import resource

            resource.setrlimit(
                resource.RLIMIT_STACK, (resource.RLIM_INFINITY, resource.RLIM_INFINITY)
            )

        try:
            self.emit("collectStarted", (None, 1))
            running_process = subprocess.Popen(
                command_list, stdout=fp1, stderr=fp2, env=envs, preexec_fn=set_ulimit
            )
            gphl_connection.collect_emulator_process = running_process

            # # NBNB leaving this in causes simulations to be killed and missing images
            # super(CollectEmulator, self).data_collection_hook()

            self.log.info("Waiting for simcal collection emulation.")
            # NBNB TODO put in time-out, somehow
            return_code = running_process.wait()
        except Exception:
            self.log.error("Error in GPhL collection emulation")
            raise
        finally:
            fp1.close()
        process = gphl_connection.collect_emulator_process
        gphl_connection.collect_emulator_process = None
        if process == "ABORTED":
            self.log.info("Simcal collection emulation aborted")
        elif return_code:
            raise RuntimeError(
                "simcal process terminated with return code %s" % return_code
            )
        else:
            self.log.info("Simcal collection emulation successful")
        self.ready_event.set()
        return
