import numpy as np
import pyqtgraph as pg
from numpy import pi
from scipy.signal import butter, sosfilt

from PyQt5 import QtCore

from acconeer.exptool import configs, utils
from acconeer.exptool.clients import SocketClient, SPIClient, UARTClient
from acconeer.exptool.pg_process import PGProccessDiedException, PGProcess
from acconeer.exptool.structs import configbase

def main():
    args = utils.ExampleArgumentParser(num_sens=1).parse_args()
    utils.config_logging(args)

    if args.socket_addr:
        client = SocketClient(args.socket_addr)
    elif args.spi:
        client = SPIClient()
    else:
        port = args.serial_port or utils.autodetect_serial_port()
        client = UARTClient(port)

    sensor_config = get_sensor_config()
    processing_config = get_processing_config()
    sensor_config.sensor = args.sensors

    session_info = client.setup_session(sensor_config)

    pg_updater = PGUpdater(sensor_config, processing_config, session_info)
    pg_process = PGProcess(pg_updater)
    pg_process.start()

    client.start_session()

    interrupt_handler = utils.ExampleInterruptHandler()
    print("Press Ctrl-C to end session")

    processor = BreathingProcessor(sensor_config, processing_config, session_info)

    while not interrupt_handler.got_signal:
        info, data = client.get_next()
        plot_data = processor.process(data)

        if plot_data is not None:
            try:
                pg_process.put_data(plot_data)
            except PGProccessDiedException:
                break

    print("Disconnecting...")
    pg_process.close()
    client.disconnect()

def get_sensor_config():
    config = configs.IQServiceConfig()
    config.range_interval = [0.3, 0.8]
    config.update_rate = 80
    config.gain = 0.5
    config.repetition_mode = configs.IQServiceConfig.RepetitionMode.SENSOR_DRIVEN
    return config

class ProcessingConfiguration(configbase.ProcessingConfig):
    VERSION = 1

    hist_plot_len = configbase.FloatParameter(
        label="Plot length",
        unit="s",
        default_value=10,
        limits=(1, 30),
        decimals=0,
    )

get_processing_config = ProcessingConfiguration

class BreathingProcessor:
    peak_hist_len = 600

    phase_weights_alpha = 0.9
    peak_loc_alpha = 0.95
    sweep_alpha = 0.7
    env_alpha = 0.95

    def __init__(self, sensor_config, processing_config, session_info):
        self.config = sensor_config

        assert sensor_config.update_rate is not None

        self.f = sensor_config.update_rate
        self.hist_plot_len = int(round(processing_config.hist_plot_len * self.f))
        self.breath_hist_len = max(2000, self.hist_plot_len)

        self.peak_history = np.zeros(self.peak_hist_len, dtype="complex")
        self.movement_history = np.zeros(self.peak_hist_len, dtype="float")
        self.breath_history = np.zeros(self.breath_hist_len, dtype="float")
        self.pulse_history = np.zeros(self.hist_plot_len, dtype="float")

        self.breath_sos = np.concatenate(butter(2, 2 * 0.3 / self.f))
        self.breath_zi = np.zeros((1, 2))
        self.pulse_sos = np.concatenate(butter(2, 2 * np.array([5]) / self.f))
        self.pulse_zi = np.zeros((1, 2))

        self.last_lp_sweep = None
        self.lp_phase_weights = None
        self.lp_sweep = None
        self.lp_peak_loc = 0

        self.sweep_index = 0

    def process(self, sweep):
        if self.sweep_index == 0:
            self.lp_sweep = np.array(sweep)
            self.lp_env = np.abs(sweep)
            self.lp_peak_loc = np.argmax(self.lp_env)

            out_data = None
        else:
            self.lp_sweep = self.lp(sweep, self.lp_sweep, self.sweep_alpha)
            env = np.abs(self.lp_sweep)
            self.lp_env = self.lp(env, self.lp_env, self.env_alpha)
            peak_loc = np.argmax(self.lp_env)
            self.lp_peak_loc = self.lp(peak_loc, self.lp_peak_loc, self.peak_loc_alpha)

            peak_idx = int(round(self.lp_peak_loc))
            peak = np.mean(self.lp_sweep[peak_idx - 50 : peak_idx + 50])
            self.push(peak, self.peak_history)

            delta = self.lp_sweep * np.conj(self.last_lp_sweep)

            phase_weights = np.imag(delta)
            if self.lp_phase_weights is None:
                self.lp_phase_weights = phase_weights
            else:
                self.lp_phase_weights = self.lp(
                    phase_weights, self.lp_phase_weights, self.phase_weights_alpha)

            weights = np.abs(self.lp_phase_weights) * env

            delta_dist = np.dot(weights, np.angle(delta))
            delta_dist *= 2.5 / (2.0 * pi * sum(weights + 0.00001))

            y = self.movement_history[0] + delta_dist
            self.push(y, self.movement_history)

            y_breath, self.breath_zi = sosfilt(self.breath_sos, np.array([y]), zi=self.breath_zi)
            self.push(y_breath, self.breath_history)

            y_pulse, self.pulse_zi = sosfilt(self.pulse_sos, np.array([y]), zi=self.pulse_zi)
            self.push(y_pulse, self.pulse_history)

            maxs = self.find_peaks(self.breath_history, 100)
            mins = self.find_peaks(-self.breath_history, 100)
            max_idx = 0
            min_idx = 0
            inhale_time = None
            exhale_time = None
            inhale_dist = 0
            exhale_dist = 0
            exhale = None
            first_peak = None
            while not (inhale_time and exhale_time):
                if not (min_idx < mins.shape[0] and max_idx < maxs.shape[0]):
                    break

                if maxs[max_idx, 0] < mins[min_idx, 0]:
                    exhale_dist = mins[min_idx, 1] + maxs[max_idx, 1]
                    if (exhale_dist > 1 and exhale_dist < 20):
                        exhale_time = mins[min_idx, 0] - maxs[max_idx, 0]
                        if first_peak is None:
                            first_peak = maxs[max_idx, 0]
                        exhale = True
                    max_idx += 1
                else:
                    inhale_dist = mins[min_idx, 1] + maxs[max_idx, 1]
                    if (inhale_dist > 1 and inhale_dist < 20):
                        inhale_time = maxs[max_idx, 0] - mins[min_idx, 0]
                        exhale = False
                        if first_peak is None:
                            first_peak = mins[min_idx, 0]
                    min_idx += 1

            breathing = False
            if inhale_time is not None and exhale_time is not None:
                bpm = 60.0 / ((inhale_time + exhale_time) / self.f)
                symmetry = (inhale_dist - exhale_dist) / (inhale_dist + exhale_dist)
                first_peak_rel = first_peak / (inhale_time + exhale_time)
                if 3 < bpm < 30 and abs(symmetry) < 0.6 and first_peak_rel < 0.7:
                    breathing = True

            if breathing:
                bstr = "Exhaling" if exhale else "Inhaling"
                bpm_text = "{}, BPM {:0.1f}, depth {:0.1f} mm".format(bstr, bpm, inhale_dist)
            else:
                bpm_text = None

            # Make an explicit copy, otherwise flip will not return a new object
            breath_hist_plot = self.breath_history[: self.hist_plot_len]
            breath_hist_plot = np.array(np.flip(breath_hist_plot, axis=0))
            breath_hist_plot -= (np.max(breath_hist_plot) + np.min(breath_hist_plot)) * 0.5

            zoom_hist_plot = self.pulse_history[: self.hist_plot_len // 2]
            zoom_hist_plot = np.array(np.flip(zoom_hist_plot, axis=0))
            zoom_hist_plot -= (max(zoom_hist_plot) + min(zoom_hist_plot)) * 0.5

            out_data = {
                "peak_hist": self.peak_history[: 100],
                "peak_std_mm": 2.5 * np.std(np.unwrap(np.angle(self.peak_history))) / (2.0 * pi),
                "env_ampl": abs(self.lp_sweep),
                "env_delta": self.lp_phase_weights,
                "peak_idx": peak_idx,
                "breathing_history": breath_hist_plot,
                "breathing_text": bpm_text,
                "zoom_hist": zoom_hist_plot,
            }

        self.last_lp_sweep = self.lp_sweep
        self.sweep_index += 1
        return out_data

    def lp(self, new, state, alpha):
        return alpha * state + (1 - alpha) * new

    def push(self, val, arr):
        res = np.empty_like(arr)
        res[0] = val
        res[1 :] = arr[: -1]
        arr[...] = res

    def find_peaks(self, env, width):
        n = len(env)
        peaks = np.zeros((0, 2))
        for idx in range(0, n, width):
            mi = np.argmax(env[idx : min(idx + width, n)]) + idx
            mi2 = np.argmax(env[max(mi - width, 0) : min(mi + width, n)])
            mi2 += max(mi - width, 0)
            if mi == mi2 and (0 < mi < n - 1):
                peaks = np.concatenate((peaks, np.array([[mi, env[mi]]])), axis=0)
        return peaks

class PGUpdater:
    def __init__(self, sensor_config, processing_config, session_info):
        assert sensor_config.update_rate is not None

        f = sensor_config.update_rate
        self.depths = utils.get_range_depths(sensor_config, session_info)
        self.hist_plot_len_s = processing_config.hist_plot_len
        self.hist_plot_len = int(round(self.hist_plot_len_s * f))
        self.move_xs = (np.arange(-self.hist_plot_len, 0) + 1) / f
        self.smooth_max = utils.SmoothMax(f, hysteresis=0.4, tau_decay=1.5)

    def setup(self, win):
        win.setWindowTitle("Acconeer breathing example")
        win.resize(800, 600)

        self.env_plot = win.addPlot(title="Amplitude of IQ data and change")
        self.env_plot.setMenuEnabled(False)
        self.env_plot.addLegend()
        self.env_plot.showGrid(x=True, y=True)
        self.env_curve = self.env_plot.plot(
            pen=utils.pg_pen_cycler(0),
            name="Amplitude of IQ data",
        )
        self.delta_curve = self.env_plot.plot(
            pen=utils.pg_pen_cycler(1),
            name="Phase change between sweeps",
        )
        self.peak_vline = pg.InfiniteLine(pen=pg.mkPen("k", width=2.5, style=QtCore.Qt.DashLine))
        self.env_plot.addItem(self.peak_vline)

        self.peak_plot = win.addPlot(title="Phase of IQ at peak")
        self.peak_plot.setMenuEnabled(False)
        utils.pg_setup_polar_plot(self.peak_plot, 1)
        self.peak_curve = self.peak_plot.plot(pen=utils.pg_pen_cycler(0))
        self.peak_scatter = pg.ScatterPlotItem(brush=pg.mkBrush("k"), size=15)
        self.peak_plot.addItem(self.peak_scatter)
        self.peak_text_item = pg.TextItem(color=pg.mkColor("k"), anchor=(0, 1))
        self.peak_plot.addItem(self.peak_text_item)
        self.peak_text_item.setPos(-1.15, -1.15)

        win.nextRow()

        self.zoom_plot = win.addPlot(title="Relative movement")
        self.zoom_plot.setMenuEnabled(False)
        self.zoom_plot.showGrid(x=True, y=True)
        self.zoom_plot.setLabel("bottom", "Time (s)")
        self.zoom_plot.setLabel("left", "Movement (mm)")
        self.zoom_curve = self.zoom_plot.plot(pen=utils.pg_pen_cycler(0))

        self.move_plot = win.addPlot(title="Breathing movement")
        self.move_plot.setMenuEnabled(False)
        self.move_plot.showGrid(x=True, y=True)
        self.move_plot.setLabel("bottom", "Time (s)")
        self.move_plot.setLabel("left", "Movement (mm)")
        self.move_plot.setYRange(-2, 2)
        self.move_plot.setXRange(-self.hist_plot_len_s, 0)
        self.move_curve = self.move_plot.plot(pen=utils.pg_pen_cycler(0))
        self.move_text_item = pg.TextItem(color=pg.mkColor("k"), anchor=(0, 1))
        self.move_text_item.setPos(self.move_xs[0], -2)
        self.move_plot.addItem(self.move_text_item)

    def update(self, data):
        envelope = data["env_ampl"]
        m = self.smooth_max.update(envelope)
        plot_delta = data["env_delta"] * m * 2e-5 + 0.5 * m

        norm_peak_hist_re = np.real(data["peak_hist"]) / m
        norm_peak_hist_im = np.imag(data["peak_hist"]) / m
        peak_std_text = "Std: {:.3f}mm".format(data["peak_std_mm"])
        peak_x = self.depths[data["peak_idx"]]

        self.env_plot.setYRange(0, m)
        self.env_curve.setData(self.depths, envelope)
        self.delta_curve.setData(self.depths, plot_delta)

        self.peak_scatter.setData([norm_peak_hist_re[0]], [norm_peak_hist_im[0]])
        self.peak_curve.setData(norm_peak_hist_re, norm_peak_hist_im)
        self.peak_text_item.setText(peak_std_text)
        self.peak_vline.setValue(peak_x)

        m = max(2, max(np.abs(data["breathing_history"])))

        self.move_curve.setData(self.move_xs, data["breathing_history"])
        self.move_plot.setYRange(-m, m)
        self.move_text_item.setPos(self.move_xs[0], -m)
        self.zoom_curve.setData(self.move_xs[self.move_xs.size // 2 :], data["zoom_hist"])
        self.move_text_item.setText(data["breathing_text"])

if __name__ == "__main__":
    main()


# An implementation of the scaler interface for dials.scale

import logging
import math
import os
import bz2
from orderedset import OrderedSet

import libtbx
import numpy as np

from xia2.Handlers.Citations import Citations
from xia2.Handlers.Files import FileHandler
from xia2.lib.bits import auto_logfiler
from xia2.Handlers.Phil import PhilIndex
from xia2.lib.SymmetryLib import sort_lattices
from xia2.Handlers.Streams import banner
from xia2.Handlers.CIF import CIF, mmCIF
from xia2.Modules.Scaler.CommonScaler import CommonScaler as Scaler
from xia2.Wrappers.Dials.Scale import DialsScale
from xia2.Wrappers.Dials.Merge import DialsMerge
from xia2.Wrappers.CCP4.CCP4Factory import CCP4Factory
from xia2.Modules.AnalyseMyIntensities import AnalyseMyIntensities
from xia2.Modules.Scaler.CCP4ScalerHelpers import (
    SweepInformationHandler,
    mosflm_B_matrix,
)
from xia2.Wrappers.Dials.Symmetry import DialsSymmetry
from xia2.Wrappers.Dials.Reindex import Reindex as DialsReindex
from xia2.Wrappers.Dials.AssignUniqueIdentifiers import DialsAssignIdentifiers
from xia2.Wrappers.Dials.SplitExperiments import SplitExperiments
from xia2.Wrappers.Dials.ExportMtz import ExportMtz
from xia2.Wrappers.Dials.ExportMMCIF import ExportMMCIF
from xia2.Wrappers.Dials.TwoThetaRefine import TwoThetaRefine
from xia2.Handlers.Syminfo import Syminfo
from dxtbx.serialize import load
from dials.util.batch_handling import calculate_batch_offsets
from dials.util.export_mtz import match_wavelengths
from dials.algorithms.scaling.plots import plot_absorption_plots
from dials.array_family import flex
import dials.util.version
from cctbx.sgtbx import lattice_symmetry_group
from iotbx import mtz
import iotbx.cif
from iotbx.scalepack import no_merge_original_index
from iotbx.scalepack.merge import write as merge_scalepack_write

logger = logging.getLogger("xia2.Modules.Scaler.DialsScaler")

class DialsScaler(Scaler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self._scalr_scaled_refl_files = {}
        self._scalr_statistics = {}
        self._factory = CCP4Factory()  # allows lots of post-scaling calculations
        self._helper = DialsScalerHelper()
        self._scaler = None
        self._scaled_experiments = None
        self._scaled_reflections = None
        self._no_times_scaled = 0
        self._scaler_symmetry_check_count = 0
        self.sweep_infos = []

    def set_working_directory(self, working_directory):
        self._working_directory = working_directory
        self._factory.set_working_directory(working_directory)
        self._helper.set_working_directory(working_directory)

    def _updated_dials_scaler(self):
        # Sets the relevant parameters from the PhilIndex

        resolution = PhilIndex.params.xia2.settings.resolution
        self._scaler.set_resolution(d_min=resolution.d_min, d_max=resolution.d_max)

        self._scaler.set_intensities(PhilIndex.params.dials.scale.intensity_choice)
        full_matrix = PhilIndex.params.dials.scale.full_matrix
        if full_matrix in (libtbx.Auto, "auto", None):
            if len(self.sweep_infos) > 4:
                full_matrix = False
            else:
                full_matrix = True
        self._scaler.set_full_matrix(full_matrix)
        self._scaler.set_outlier_rejection(
            PhilIndex.params.dials.scale.outlier_rejection
        )
        self._scaler.set_outlier_zmax(PhilIndex.params.dials.scale.outlier_zmax)
        self._scaler.set_error_model(PhilIndex.params.dials.scale.error_model)
        self._scaler.set_error_model_grouping_method(
            PhilIndex.params.dials.scale.error_model_grouping
        )
        if PhilIndex.params.dials.scale.error_model_group:
            self._scaler.set_error_model_groups(
                PhilIndex.params.dials.scale.error_model_group
            )
        self._scaler.set_partiality_cutoff(
            PhilIndex.params.dials.scale.partiality_threshold
        )

        if self.get_scaler_anomalous():
            self._scaler.set_anomalous()

        exp = load.experiment_list(self.sweep_infos[0].get_experiments())[0]
        scale_interval, decay_interval = scaling_model_auto_rules(exp)

        # Model handling
        autos = (None, "auto", libtbx.Auto)
        if PhilIndex.params.dials.scale.model in autos:
            PhilIndex.params.dials.scale.model = "physical"
        self._scaler.set_model(PhilIndex.params.dials.scale.model)

        if PhilIndex.params.dials.scale.rotation_spacing:
            scale_interval = PhilIndex.params.dials.scale.rotation_spacing

        if PhilIndex.params.dials.scale.model == "physical":
            if PhilIndex.params.dials.scale.physical_model.Bfactor_spacing:
                decay_interval = (
                    PhilIndex.params.dials.scale.physical_model.Bfactor_spacing
                )
            self._scaler.set_spacing(scale_interval)
            if PhilIndex.params.dials.scale.Bfactor:
                self._scaler.set_bfactor(True, decay_interval)
            else:
                self._scaler.set_bfactor(False)
            if PhilIndex.params.dials.scale.absorption:
                self._scaler.set_absorption_correction(True)
                if PhilIndex.params.dials.scale.physical_model.absorption_level:
                    self._scaler.set_absorption_level(
                        PhilIndex.params.dials.scale.physical_model.absorption_level
                    )
                if PhilIndex.params.dials.scale.physical_model.lmax not in autos:
                    self._scaler.set_lmax(
                        PhilIndex.params.dials.scale.physical_model.lmax
                    )
                if (
                    PhilIndex.params.dials.scale.physical_model.surface_weight
                    not in autos
                ):
                    self._scaler.set_surface_weight(
                        PhilIndex.params.dials.scale.physical_model.surface_weight
                    )
                if PhilIndex.params.dials.scale.physical_model.share.absorption:
                    self._scaler.set_shared_absorption(True)
            else:
                self._scaler.set_absorption_correction(False)
        elif PhilIndex.params.dials.scale.model == "dose_decay":
            self._scaler.set_spacing(scale_interval)
            if PhilIndex.params.dials.scale.absorption:
                self._scaler.set_absorption_correction(True)
                self._scaler.set_lmax(PhilIndex.params.dials.scale.decay_model.lmax)
            else:
                self._scaler.set_absorption_correction(False)
            if PhilIndex.params.dials.scale.dose_decay_model.share.decay is not None:
                self._scale.set_shared_decay(
                    PhilIndex.params.dials.scale.dose_decay_model.share.decay
                )
            if PhilIndex.params.dials.scale.dose_decay_model.resolution_dependence:
                self._scale.set_resolution_dependence(
                    PhilIndex.dials.scale.dose_decay.resolution_dependence
                )
        elif PhilIndex.params.dials.scale.model == "KB":
            # For KB model, want both Bfactor and scale terms
            self._scaler.set_bfactor(True)
        elif PhilIndex.params.dials.scale.model == "array":
            if PhilIndex.params.dials.scale.Bfactor:
                self._scaler.set_bfactor(True, scale_interval)
                self._scaler.set_decay_bins(
                    PhilIndex.params.dials.scale.array_model.resolution_bins
                )
            else:
                self._scaler.set_bfactor(False)
            if PhilIndex.params.dials.scale.absorption:
                self._scaler.set_absorption_correction(True)
                self._scaler.set_array_absorption_bins(
                    PhilIndex.params.dials.scale.array_model.absorption_bins
                )
            else:
                self._scaler.set_absorption_correction(False)

        return self._scaler

    def _do_prescale_kb(self, experiments, reflections):
        # Pre-scale the data with KB scaling to ensure all experiments are on
        # the same scale prior to running dials.symmetry
        self._scaler = DialsScale()
        self._scaler.set_model("KB")
        self._scaler.set_full_matrix(False)
        self._scaler.set_error_model(None)
        self._scaler.set_intensities("profile")

        for expts, refl in zip(experiments, reflections):
            self._scaler.add_experiments_json(expts)
            self._scaler.add_reflections_file(refl)

        self._scaler.set_working_directory(self.get_working_directory())
        auto_logfiler(self._scaler)
        self._scaler.scale()
        prescaled_experiments = self._scaler.get_scaled_experiments()
        prescaled_reflections = self._scaler.get_scaled_reflections()
        self._scaler = None
        return prescaled_experiments, prescaled_reflections

    def _do_multisweep_symmetry_analysis(self):
        refiners = []
        experiments = []
        reflections = []

        for epoch in self._sweep_handler.get_epochs():
            si = self._sweep_handler.get_sweep_information(epoch)
            integrater = si.get_integrater()
            experiments.append(integrater.get_integrated_experiments())
            reflections.append(integrater.get_integrated_reflections())
            refiners.append(integrater.get_integrater_refiner())

        prescaled_experiments, prescaled_reflections = self._do_prescale_kb(
            experiments, reflections
        )

        logger.debug("Running multisweep dials.symmetry for %d sweeps", len(refiners))
        (
            pointgroup,
            reindex_op,
            ntr,
            pt,
            reind_refl,
            reind_exp,
            reindex_initial,
        ) = self._helper.dials_symmetry_indexer_jiffy(
            [prescaled_experiments], [prescaled_reflections], refiners, multisweep=True
        )

        FileHandler.record_temporary_file(reind_refl)
        FileHandler.record_temporary_file(reind_exp)
        return pointgroup, reindex_op, ntr, pt, reind_refl, reind_exp, reindex_initial

    def _multi_sweep_scale_prepare(self):
        need_to_return = False

        (
            pointgroup,
            reindex_op,
            ntr,
            _,
            reind_refl,
            reind_exp,
            reindex_initial,
        ) = self._do_multisweep_symmetry_analysis()
        if ntr:
            for epoch in self._sweep_handler.get_epochs():
                si = self._sweep_handler.get_sweep_information(epoch)
                si.get_integrater().integrater_reset_reindex_operator()
            self.set_scaler_done(False)
            self.set_scaler_prepare_done(False)
            need_to_return = True
            return need_to_return
        else:
            self._scalr_likely_spacegroups = [pointgroup]
            if reindex_initial:
                for epoch in self._sweep_handler.get_epochs():
                    si = self._sweep_handler.get_sweep_information(epoch)
                    self._helper.reindex_jiffy(si, pointgroup, reindex_op=reindex_op)
                # integrater reset reindex op and update in si.
            else:
                self._sweep_handler = self._helper.split_experiments(
                    reind_exp, reind_refl, self._sweep_handler
                )

        return need_to_return

    def _input_pointgroup_scale_prepare(self):
        pointgroup = self._scalr_input_pointgroup
        if self._scalr_input_spacegroup:
            self._scalr_likely_spacegroups = [self._scalr_input_spacegroup]
            pointgroup = self._scalr_input_spacegroup
        logger.debug("Using input pointgroup: %s", pointgroup)
        for epoch in self._sweep_handler.get_epochs():
            si = self._sweep_handler.get_sweep_information(epoch)
            self._helper.reindex_jiffy(si, pointgroup, "h,k,l")

    def _standard_scale_prepare(self):
        pointgroups = {}
        reindex_ops = {}
        probably_twinned = False
        need_to_return = False

        lattices = []
        # First check for the existence of multiple lattices. If only one
        # epoch, then this gives the necessary data for proceeding straight
        # to the point group check.
        for epoch in self._sweep_handler.get_epochs():
            si = self._sweep_handler.get_sweep_information(epoch)
            intgr = si.get_integrater()
            experiment = intgr.get_integrated_experiments()
            reflections = intgr.get_integrated_reflections()
            refiner = intgr.get_integrater_refiner()

            (
                pointgroup,
                reindex_op,
                ntr,
                pt,
                _,
                __,
                ___,
            ) = self._helper.dials_symmetry_indexer_jiffy(
                [experiment], [reflections], [refiner]
            )

            lattice = Syminfo.get_lattice(pointgroup)
            if lattice not in lattices:
                lattices.append(lattice)
            if ntr:
                si.get_integrater().integrater_reset_reindex_operator()
                need_to_return = True
            if pt:
                probably_twinned = True
            pointgroups[epoch] = pointgroup
            reindex_ops[epoch] = reindex_op
            logger.debug("Pointgroup: %s (%s)", pointgroup, reindex_op)

        if len(lattices) > 1:
            # Check consistency of lattices if more than one. If not, then
            # can proceed to straight to checking point group consistency
            # using the cached results.
            correct_lattice = sort_lattices(lattices)[0]
            logger.info("Correct lattice asserted to be %s", correct_lattice)

            # transfer this information back to the indexers
            for epoch in self._sweep_handler.get_epochs():
                si = self._sweep_handler.get_sweep_information(epoch)
                refiner = si.get_integrater().get_integrater_refiner()
                _tup = (correct_lattice, si.get_sweep_name())

                state = refiner.set_refiner_asserted_lattice(correct_lattice)

                if state == refiner.LATTICE_CORRECT:
                    logger.info("Lattice %s ok for sweep %s" % _tup)
                elif state == refiner.LATTICE_IMPOSSIBLE:
                    raise RuntimeError("Lattice %s impossible for %s" % _tup)
                elif state == refiner.LATTICE_POSSIBLE:
                    logger.info("Lattice %s assigned for sweep %s" % _tup)
                    need_to_return = True

        if need_to_return:
            return need_to_return

        need_to_return = False

        pointgroup_set = {pointgroups[e] for e in pointgroups}

        if len(pointgroup_set) > 1 and not probably_twinned:
            raise RuntimeError(
                "non uniform pointgroups: %s" % str(list(pointgroup_set))
            )

        if len(pointgroup_set) > 1:
            logger.debug(
                "Probably twinned, pointgroups: %s",
                " ".join(p.replace(" ", "") for p in pointgroup_set),
            )
            numbers = [Syminfo.spacegroup_name_to_number(s) for s in pointgroup_set]
            overall_pointgroup = Syminfo.spacegroup_number_to_name(min(numbers))
            self._scalr_input_pointgroup = overall_pointgroup

            logger.info("Twinning detected, assume pointgroup %s", overall_pointgroup)
            need_to_return = True
        else:
            overall_pointgroup = pointgroup_set.pop()
        self._scalr_likely_spacegroups = [overall_pointgroup]
        for epoch in self._sweep_handler.get_epochs():
            si = self._sweep_handler.get_sweep_information(epoch)
            self._helper.reindex_jiffy(si, overall_pointgroup, reindex_ops[epoch])
        return need_to_return

    def _scale_prepare(self):
        """Perform all of the preparation required to deliver the scaled
        data. This should sort together the reflection files, ensure that
        they are correctly indexed (via dials.symmetry) and generally tidy
        things up."""

        Citations.cite("dials.scale")
        # AIM discover symmetry and reindex with dials.symmetry, and set the correct
        # reflections in si.reflections, si.experiments

        self._helper.set_working_directory(self.get_working_directory())
        self._factory.set_working_directory(self.get_working_directory())

        self._sweep_handler = SweepInformationHandler(self._scalr_integraters)

        p, x = self._sweep_handler.get_project_info()
        self._scalr_pname = p
        self._scalr_xname = x

        self._helper.set_pname_xname(p, x)

        # First do stuff to work out if excluding any data
        # Note - does this actually work? I couldn't seem to get it to work
        # in either this pipeline or the standard dials pipeline
        for epoch in self._sweep_handler.get_epochs():
            si = self._sweep_handler.get_sweep_information(epoch)
            _, xname, dname = si.get_project_info()
            sname = si.get_sweep_name()

            exclude_sweep = False

            for sweep in PhilIndex.params.xia2.settings.sweep:
                if sweep.id == sname and sweep.exclude:
                    exclude_sweep = True
                    break

            if exclude_sweep:
                self._sweep_handler.remove_epoch(epoch)
                logger.debug("Excluding sweep %s", sname)
            else:
                logger.debug("%-30s %s/%s/%s", "adding data from:", xname, dname, sname)

        # If multiple files, want to run symmetry to check for consistent indexing
        # also

        # try to reproduce what CCP4ScalerA is doing

        # first assign identifiers to avoid dataset-id collisions
        # Idea is that this should be called anytime you get data anew from the
        # integrater, to intercept and assign unique ids, then set in the
        # sweep_information (si) and always use si.set_reflections/
        # si.get_reflections as we process.

        # self._sweep_handler = self._helper.assign_and_return_datasets(
        #    self._sweep_handler
        # ) symmetry now sorts out identifiers.

        need_to_return = False

        if self._scalr_input_pointgroup:
            self._input_pointgroup_scale_prepare()
        elif (
            len(self._sweep_handler.get_epochs()) > 1
            and PhilIndex.params.xia2.settings.multi_sweep_indexing
        ):
            need_to_return = self._multi_sweep_scale_prepare()
        else:
            need_to_return = self._standard_scale_prepare()

        if need_to_return:
            self.set_scaler_done(False)
            self.set_scaler_prepare_done(False)
            return

        ### After this point, point group is good and only need to
        ### reindex to consistent setting. Don't need to call back to the
        ### integator, just use the data in the sweep info.

        # First work out if we're going to reindex against external reference
        param = PhilIndex.params.xia2.settings.scale
        using_external_references = False
        reference_refl = None
        reference_expt = None
        if param.reference_reflection_file:
            if not param.reference_experiment_file:
                logger.info(
                    """
No DIALS reference experiments file provided, reference reflection file will
not be used. Reference mtz files for reindexing not currently supported for
pipeline=dials (supported for pipeline=dials-aimless).
"""
                )
            else:
                reference_refl = param.reference_reflection_file
                reference_expt = param.reference_experiment_file
                using_external_references = True
                logger.debug("Using reference reflections %s", reference_refl)
                logger.debug("Using reference experiments %s", reference_expt)

        if len(self._sweep_handler.get_epochs()) > 1:
            if PhilIndex.params.xia2.settings.unify_setting:
                self.unify_setting()

            if PhilIndex.params.xia2.settings.use_brehm_diederichs:
                self.brehm_diederichs_reindexing()
            # If not using Brehm-deidrichs reindexing, set reference as first
            # sweep, unless using external reference.
            elif not using_external_references:
                logger.debug("First sweep will be used as reference for reindexing")
                first = self._sweep_handler.get_epochs()[0]
                si = self._sweep_handler.get_sweep_information(first)
                reference_expt = si.get_experiments()
                reference_refl = si.get_reflections()

        # Now reindex to be consistent with first dataset - run reindex on each
        # dataset with reference (unless did brehm diederichs and didn't supply
        # a reference file)

        if reference_refl and reference_expt:
            exp = load.experiment_list(reference_expt)
            reference_cell = exp[0].crystal.get_unit_cell().parameters()

            # ---------- REINDEX TO CORRECT (REFERENCE) SETTING ----------
            logger.info("Reindexing all datasets to common reference")

            if using_external_references:
                epochs = self._sweep_handler.get_epochs()
            else:
                epochs = self._sweep_handler.get_epochs()[1:]
            for epoch in epochs:
                # if we are working with unified UB matrix then this should not
                # be a problem here (note, *if*; *should*)

                # what about e.g. alternative P1 settings?
                # see JIRA MXSW-904
                if PhilIndex.params.xia2.settings.unify_setting:
                    continue

                reindexer = DialsReindex()
                reindexer.set_working_directory(self.get_working_directory())
                auto_logfiler(reindexer)

                si = self._sweep_handler.get_sweep_information(epoch)
                reindexer.set_reference_filename(reference_expt)
                reindexer.set_reference_reflections(reference_refl)
                reindexer.set_indexed_filename(si.get_reflections())
                reindexer.set_experiments_filename(si.get_experiments())
                reindexer.run()

                # At this point, CCP4ScalerA would reset in integrator so that
                # the integrater calls reindex, no need to do that here as
                # have access to the files and will never need to reintegrate.

                si.set_reflections(reindexer.get_reindexed_reflections_filename())
                si.set_experiments(reindexer.get_reindexed_experiments_filename())

                # FIXME how to get some indication of the reindexing used?

                exp = load.experiment_list(
                    reindexer.get_reindexed_experiments_filename()
                )
                cell = exp[0].crystal.get_unit_cell().parameters()

                # Note - no lattice check as this will already be caught by reindex
                logger.debug("Cell: %.2f %.2f %.2f %.2f %.2f %.2f" % cell)
                logger.debug("Ref:  %.2f %.2f %.2f %.2f %.2f %.2f" % reference_cell)

                for j in range(6):
                    if (
                        math.fabs((cell[j] - reference_cell[j]) / reference_cell[j])
                        > 0.1
                    ):
                        raise RuntimeError(
                            "unit cell parameters differ in %s and %s"
                            % (reference_expt, si.get_reflections())
                        )

        # Now make sure all batches ok before finish preparing
        # This should be made safer, currently after dials.scale there is no
        # concept of 'batch', dials.export uses the calculate_batch_offsets
        # to assign batches, giving the same result as below.

        experiments_to_rebatch = []
        for epoch in self._sweep_handler.get_epochs():
            si = self._sweep_handler.get_sweep_information(epoch)
            experiment = si.get_experiments()
            experiments_to_rebatch.append(load.experiment_list(experiment)[0])
        offsets = calculate_batch_offsets(experiments_to_rebatch)

        for i, epoch in enumerate(self._sweep_handler.get_epochs()):
            si = self._sweep_handler.get_sweep_information(epoch)
            r = si.get_batch_range()
            si.set_batch_offset(offsets[i])
            si.set_batches([r[0] + offsets[i], r[1] + offsets[i]])

    def _scale(self):
        """Perform all of the operations required to deliver the scaled
        data."""
        self.sweep_infos = [
            self._sweep_handler.get_sweep_information(e)
            for e in self._sweep_handler.get_epochs()
        ]

        ### Set the parameters and datafiles for dials.scale

        self._scaler = DialsScale()
        self._scaler = self._updated_dials_scaler()

        if self._scaled_experiments and self._scaled_reflections:
            # going to continue-where-left-off
            self._scaler.add_experiments_json(self._scaled_experiments)
            self._scaler.add_reflections_file(self._scaled_reflections)
        else:
            for si in self.sweep_infos:
                self._scaler.add_experiments_json(si.get_experiments())
                self._scaler.add_reflections_file(si.get_reflections())
            # ensure we start with a clean slate in case we pre-scaled the data
            # before running dials.symmetry
            self._scaler.set_overwrite_existing_models(True)

        self._scalr_scaled_reflection_files = {"mtz_unmerged": {}, "mtz": {}}

        ### Set the resolution limit if applicable

        user_resolution_limits = {}
        highest_resolution = 100.0
        for si in self.sweep_infos:
            dname = si.get_project_info()[2]
            sname = si.get_sweep_name()
            intgr = si.get_integrater()

            if intgr.get_integrater_user_resolution():
                # record user resolution here but don't use it until later - why?
                dmin = intgr.get_integrater_high_resolution()

                if (dname, sname) not in user_resolution_limits:
                    user_resolution_limits[(dname, sname)] = dmin
                elif dmin < user_resolution_limits[(dname, sname)]:
                    user_resolution_limits[(dname, sname)] = dmin

            if (dname, sname) in self._scalr_resolution_limits:
                d_min, _ = self._scalr_resolution_limits[(dname, sname)]
                if d_min < highest_resolution:
                    highest_resolution = d_min
        if highest_resolution < 99.9:
            self._scaler.set_resolution(d_min=highest_resolution)

        ### Setup final job details and run scale

        self._scaler.set_working_directory(self.get_working_directory())
        auto_logfiler(self._scaler)
        FileHandler.record_log_file(
            f"{self._scalr_pname} {self._scalr_xname} SCALE",
            self._scaler.get_log_file(),
        )
        self._scaler.scale()
        FileHandler.record_html_file(
            f"{self._scalr_pname} {self._scalr_xname} SCALE",
            self._scaler.get_html(),
        )
        self._scaled_experiments = self._scaler.get_scaled_experiments()
        self._scaled_reflections = self._scaler.get_scaled_reflections()

        # make it so that only scaled.expt and scaled.refl are
        # the files that dials.scale knows about, so that if scale is called again,
        # scaling resumes from where it left off.
        self._scaler.clear_datafiles()

        ### Calculate the resolution limit and set done False if applicable

        highest_suggested_resolution = self.assess_resolution_limits(
            hklin=None,
            user_resolution_limits=user_resolution_limits,
            reflections=self._scaled_reflections,
            experiments=self._scaled_experiments,
        )

        if not self.get_scaler_done():
            # reset for when resolution limit applied
            logger.debug("Returning as scaling not finished...")
            return

        ### Want to do space group check after scaling. So run dials.symmetry
        ### with absences only before exporting merged and unmerged files
        ### again in correct s.g.
        if (
            not PhilIndex.params.xia2.settings.small_molecule
            and not self._scalr_input_spacegroup
        ):
            logger.notice(banner("Systematic absences check"))
            symmetry = DialsSymmetry()
            symmetry.set_experiments_filename(self._scaled_experiments)
            symmetry.set_reflections_filename(self._scaled_reflections)
            symmetry.set_working_directory(self.get_working_directory())
            symmetry.set_mode_absences_only()
            auto_logfiler(symmetry)
            symmetry.decide_pointgroup()  # bad name - actually running absences here

            self._scaled_experiments = symmetry.get_output_experiments_filename()

            sg = load.experiment_list(self._scaled_experiments)[
                0
            ].crystal.get_space_group()
            logger.info("Most likely space group: %s", sg.info())
            self._scalr_likely_spacegroups = [sg.type().lookup_symbol()]

        FileHandler.record_more_data_file(
            f"{self._scalr_pname} {self._scalr_xname} scaled",
            self._scaled_experiments,
        )
        FileHandler.record_more_data_file(
            f"{self._scalr_pname} {self._scalr_xname} scaled",
            self._scaled_reflections,
        )

        # Run twotheta refine
        self._update_scaled_unit_cell_from_scaled_data()

        ### Now export and merge so that mtz files in correct space group.

        ### For MAD case, need to generate individual merged and unmerged mtz
        ### files. First split experiments on wavelength, then run dials.export
        ### and dials.merge on each

        # Find number of dnames (i.e. number of wavelengths)
        dnames_set = OrderedSet()
        experiments = load.experiment_list(self._scaled_experiments)
        wavelengths = flex.double(
            match_wavelengths(experiments)
        )  # in experiments order
        for si in self.sweep_infos:
            dnames_set.add(
                si.get_project_info()[2]
            )  # sweep info in same order as experiments
        assert len(wavelengths) == len(dnames_set)

        scaled_unmerged_mtz_path = os.path.join(
            self.get_working_directory(),
            f"{self._scalr_pname}_{self._scalr_xname}_scaled_unmerged.mtz",
        )

        if len(dnames_set) > 1:
            self._scalr_scaled_refl_files = {}
            logger.debug("Splitting experiments by wavelength")
            # first split by wavelength
            splitter = SplitExperiments()
            splitter.add_experiments(self._scaled_experiments)
            splitter.add_reflections(self._scaled_reflections)
            splitter.set_by_wavelength(True)
            splitter.set_working_directory(self.get_working_directory())
            auto_logfiler(splitter)
            splitter.run()

            nn = len(dnames_set)
            fmt = "%%0%dd" % (math.log10(nn) + 1)

            wl_sort = flex.sort_permutation(wavelengths)
            sorted_dnames_by_wl = [dnames_set[i] for i in wl_sort]

            for i, dname in enumerate(sorted_dnames_by_wl):
                # need to sort by wavelength from low to high
                nums = fmt % i
                exporter = ExportMtz()
                exporter.set_working_directory(self.get_working_directory())
                expt_name = os.path.join(
                    self.get_working_directory(), "split_%s.expt" % nums
                )
                refl_name = os.path.join(
                    self.get_working_directory(), "split_%s.refl" % nums
                )
                FileHandler.record_temporary_file(expt_name)
                FileHandler.record_temporary_file(refl_name)
                exporter.crystal_name = self._scalr_xname
                exporter.project_name = self._scalr_pname
                exporter.set_experiments_filename(expt_name)
                exporter.set_reflections_filename(refl_name)
                exporter.set_intensity_choice("scale")
                exporter.set_partiality_threshold(
                    PhilIndex.params.dials.scale.partiality_threshold
                )  # 0.4 default
                auto_logfiler(exporter)
                mtz_filename = os.path.join(
                    self.get_working_directory(),
                    scaled_unmerged_mtz_path.rstrip(".mtz") + "_%s.mtz" % dname,
                )
                exporter.set_mtz_filename(mtz_filename)
                self._scalr_scaled_reflection_files["mtz_unmerged"][
                    dname
                ] = mtz_filename

                logger.debug("Exporting %s", mtz_filename)
                exporter.run()
                FileHandler.record_data_file(mtz_filename)

                # Export an mmCIF file using dials.export.
                exporter = ExportMMCIF()
                exporter.set_working_directory(self.get_working_directory())
                exporter.set_experiments_filename(expt_name)
                exporter.set_reflections_filename(refl_name)
                exporter.set_compression("bz2")
                exporter.set_pdb_version(
                    PhilIndex.params.xia2.settings.output.mmcif.pdb_version
                )
                exporter.set_partiality_threshold(
                    PhilIndex.params.dials.scale.partiality_threshold
                )  # 0.4 default
                mmcif_path = mtz_filename.rstrip(".mtz") + ".mmcif"
                exporter.set_filename(mmcif_path)
                auto_logfiler(exporter)
                logger.debug("Exporting %s", mmcif_path)
                exporter.run()

                FileHandler.record_temporary_file(mmcif_path)

                # now convert to .sca format
                convert_unmerged_mtz_to_sca(mtz_filename)

                merger = DialsMerge()  # merge but don't truncate
                merger.set_working_directory(self.get_working_directory())
                merger.set_experiments_filename(expt_name)
                merger.set_reflections_filename(refl_name)
                merger.set_project_name(self._scalr_pname)
                merger.set_crystal_names(self._scalr_xname)
                merger.set_dataset_names(dname)
                merger.set_partiality_threshold(
                    PhilIndex.params.dials.scale.partiality_threshold
                )
                auto_logfiler(merger)
                mtz_filename = os.path.join(
                    self.get_working_directory(),
                    "%s_%s_scaled_%s.mtz"
                    % (self._scalr_pname, self._scalr_xname, dname),
                )
                self._scalr_scaled_refl_files[dname] = mtz_filename
                self._scalr_scaled_reflection_files["mtz"][dname] = mtz_filename
                merger.set_mtz_filename(mtz_filename)

                logger.debug("Merging %s", mtz_filename)
                merger.run()
                FileHandler.record_data_file(mtz_filename)

                # now convert to .sca format
                convert_merged_mtz_to_sca(mtz_filename)

        ### For non-MAD case, run dials.export and dials.merge on scaled data.
        else:
            exporter = ExportMtz()
            exporter.crystal_name = self._scalr_xname
            exporter.project_name = self._scalr_pname
            exporter.set_working_directory(self.get_working_directory())
            exporter.set_experiments_filename(self._scaled_experiments)
            exporter.set_reflections_filename(self._scaled_reflections)
            exporter.set_intensity_choice("scale")
            exporter.set_partiality_threshold(
                PhilIndex.params.dials.scale.partiality_threshold
            )  # 0.4 default
            auto_logfiler(exporter)
            exporter.set_mtz_filename(scaled_unmerged_mtz_path)

            logger.debug("Exporting %s", scaled_unmerged_mtz_path)
            exporter.run()

            self._scalr_scaled_reflection_files["mtz_unmerged"] = {
                dnames_set[0]: scaled_unmerged_mtz_path
            }

            FileHandler.record_data_file(scaled_unmerged_mtz_path)

            # now convert to .sca format
            convert_unmerged_mtz_to_sca(scaled_unmerged_mtz_path)

            merger = DialsMerge()
            merger.set_working_directory(self.get_working_directory())
            merger.set_experiments_filename(self._scaled_experiments)
            merger.set_reflections_filename(self._scaled_reflections)
            merger.set_project_name(self._scalr_pname)
            merger.set_crystal_names(self._scalr_xname)
            merger.set_dataset_names(dnames_set[0])
            merger.set_partiality_threshold(
                PhilIndex.params.dials.scale.partiality_threshold
            )
            auto_logfiler(merger)
            mtz_filename = os.path.join(
                self.get_working_directory(),
                f"{self._scalr_pname}_{self._scalr_xname}_scaled.mtz",
            )
            self._scalr_scaled_refl_files[dnames_set[0]] = mtz_filename
            self._scalr_scaled_reflection_files["mtz"][dnames_set[0]] = mtz_filename
            merger.set_mtz_filename(mtz_filename)

            logger.debug("Merging %s", mtz_filename)
            merger.run()
            FileHandler.record_data_file(mtz_filename)

            # now export to sca format
            convert_merged_mtz_to_sca(mtz_filename)

            # Export an mmCIF file using dials.export.
            exporter = ExportMMCIF()
            exporter.set_working_directory(self.get_working_directory())
            exporter.set_experiments_filename(self._scaled_experiments)
            exporter.set_reflections_filename(self._scaled_reflections)
            exporter.set_compression("bz2")
            exporter.set_pdb_version(
                PhilIndex.params.xia2.settings.output.mmcif.pdb_version
            )
            exporter.set_partiality_threshold(
                PhilIndex.params.dials.scale.partiality_threshold
            )  # 0.4 default
            mmcif_path = scaled_unmerged_mtz_path.rstrip(".mtz") + ".mmcif"
            exporter.set_filename(mmcif_path)
            auto_logfiler(exporter)
            logger.debug("Exporting %s", mmcif_path)
            exporter.run()

            FileHandler.record_temporary_file(mmcif_path)

        # Also export just integrated data.
        for si in self.sweep_infos:
            exporter = ExportMtz()
            exporter.crystal_name = self._scalr_xname
            exporter.project_name = self._scalr_pname
            exporter.set_reflections_filename(si.get_reflections())
            exporter.set_experiments_filename(si.get_experiments())
            exporter.set_intensity_choice("profile+sum")
            pname, xname, dname = si.get_project_info()
            sweep = si.get_integrater().get_integrater_sweep_name()
            tag = f"{pname} {xname} {dname} {sweep} INTEGRATE"
            mtz_filename = os.path.join(
                self.get_working_directory(), "%s_integrated.mtz" % sweep
            )
            exporter.set_mtz_filename(mtz_filename)
            exporter.run()
            FileHandler.record_more_data_file(tag, mtz_filename)

        if PhilIndex.params.xia2.settings.merging_statistics.source == "cctbx":
            for key in self._scalr_scaled_refl_files:
                stats = self._compute_scaler_statistics(
                    self._scalr_scaled_reflection_files["mtz_unmerged"][key],
                    selected_band=(highest_suggested_resolution, None),
                    wave=key,
                )
                self._scalr_statistics[
                    (self._scalr_pname, self._scalr_xname, key)
                ] = stats

        # add CIF data
        expts = load.experiment_list(self._scaled_experiments)
        overall_absmin = 1.0
        for expt in expts:
            if (expt.scaling_model.id_ == "physical") and (
                "absorption" in expt.scaling_model.components
            ):
                plots = plot_absorption_plots(expt.scaling_model)
                correction = np.array(plots["absorption_surface"]["data"][0]["z"])
                # correction is a 2D numpy array
                absmin = np.min(correction) / np.max(correction)
                if absmin > 0:  # hope should always happen!
                    overall_absmin = min(absmin, overall_absmin)

        dials_version = dials.util.version.dials_version()
        block = CIF.get_block("xia2")
        mmblock = mmCIF.get_block("xia2")
        mmblock["_exptl.entry_id"] = "xia2"
        mmblock["_exptl.method"] = "X-RAY DIFFRACTION"
        block["_exptl_absorpt_correction_T_min"] = mmblock[
            "_exptl.absorpt_correction_T_min"
        ] = overall_absmin  # = scaled relative to 1
        block["_exptl_absorpt_correction_T_max"] = mmblock[
            "_exptl.absorpt_correction_T_max"
        ] = 1.0  #
        block["_exptl_absorpt_correction_type"] = mmblock[
            "_exptl.absorpt_correction_type"
        ] = "empirical"
        block["_exptl_absorpt_process_details"] = mmblock[
            "_exptl.absorpt_process_details"
        ] = (
            """
%s
Scaling & analysis of unmerged intensities, absorption correction using spherical harmonics
"""
            % dials_version
        )

    def _write_mmcif_output(self):
        """Migrate mmcif data generated by dials.export"""

        def _add_to_block(blockname, mmcif_path):
            mmcif_file_object = bz2.open(mmcif_path + ".bz2")
            mmblock_dials = iotbx.cif.reader(file_object=mmcif_file_object).model()
            # give the entry a consistent unique name in all fields
            if PhilIndex.params.xia2.settings.output.mmcif.pdb_version == "v5_next":
                mmblock_dials["dials"]["_pdbx_diffrn_data_section.id"] = blockname
            mmblock_dials["dials"]["_entry.id"] = blockname
            for key in mmblock_dials["dials"].keys():
                if key.endswith("entry_id"):
                    mmblock_dials["dials"][key] = blockname
            mmCIF.get_block(blockname).update(mmblock_dials["dials"])

        dnames = OrderedSet([si.get_project_info()[2] for si in self.sweep_infos])
        for dname in dnames:
            mtz_path = self._scalr_scaled_reflection_files["mtz_unmerged"][dname]
            mmcif_path = mtz_path.rstrip(".mtz") + ".mmcif"
            blockname = f"{self._scalr_pname}_{self._scalr_xname}_{dname}"
            _add_to_block(blockname, mmcif_path)

    def _update_scaled_unit_cell_from_scaled_data(self):

        params = PhilIndex.params
        fast_mode = params.dials.fast_mode
        if (
            params.xia2.settings.integrater == "dials"
            and not fast_mode
            and params.xia2.settings.scale.two_theta_refine
        ):

            logger.notice(banner("Unit cell refinement"))

            # Collect a list of all sweeps, grouped by project, crystal, wavelength
            groups_list = []
            groups = {}
            self._scalr_cell_dict = {}
            for epoch in self._sweep_handler.get_epochs():
                si = self._sweep_handler.get_sweep_information(epoch)
                pi = "_".join(si.get_project_info())  # pname, xname, dname
                groups_list.append(pi)

            p4p_file = os.path.join(
                self.get_working_directory(),
                f"{self._scalr_pname}_{self._scalr_xname}.p4p",
            )
            if len(set(groups_list)) > 1:
                # need to split up experiments and reflections
                self._sweep_handler = self._helper.split_experiments(
                    self._scaled_experiments,
                    self._scaled_reflections,
                    self._sweep_handler,
                )
                for epoch in self._sweep_handler.get_epochs():
                    si = self._sweep_handler.get_sweep_information(epoch)
                    pi = "_".join(si.get_project_info())  # pname, xname, dname
                    groups[pi] = groups.get(pi, []) + [
                        (si.get_experiments(), si.get_reflections())
                    ]  # if key exists, add another 2-tuple to the list.
                for pi in groups:
                    # Run twothetarefine on each group
                    tt_grouprefiner = TwoThetaRefine()
                    tt_grouprefiner.set_working_directory(self.get_working_directory())
                    auto_logfiler(tt_grouprefiner)
                    args = list(zip(*groups[pi]))
                    tt_grouprefiner.set_experiments(args[0])
                    tt_grouprefiner.set_reflection_files(args[1])
                    tt_grouprefiner.set_output_p4p(p4p_file)
                    tt_grouprefiner.run()
                    logger.info(
                        "%s: %6.2f %6.2f %6.2f %6.2f %6.2f %6.2f"
                        % tuple(
                            ["".join(pi.split("_")[2:])]
                            + list(tt_grouprefiner.get_unit_cell())
                        )
                    )
                    self._scalr_cell_dict[pi] = (
                        tt_grouprefiner.get_unit_cell(),
                        tt_grouprefiner.get_unit_cell_esd(),
                        tt_grouprefiner.import_cif(),
                        tt_grouprefiner.import_mmcif(),
                    )
                    if len(groups_list) > 1:
                        cif_in = tt_grouprefiner.import_cif()
                        cif_out = CIF.get_block(pi)
                        for key in sorted(cif_in.keys()):
                            cif_out[key] = cif_in[key]
                        mmcif_in = tt_grouprefiner.import_mmcif()
                        mmcif_out = mmCIF.get_block(pi)
                        # reset the entry id to be the name of the new block
                        mmcif_out["_entry.id"] = pi
                        for key in sorted(mmcif_in.keys()):
                            if key.endswith("entry_id"):
                                mmcif_out[key] = pi
                            else:
                                mmcif_out[key] = mmcif_in[key]

            # now do two theta refine on combined scaled data.
            tt_refiner = TwoThetaRefine()
            tt_refiner.set_working_directory(self.get_working_directory())
            auto_logfiler(tt_refiner)
            tt_refiner.set_experiments([self._scaled_experiments])
            tt_refiner.set_reflection_files([self._scaled_reflections])  # needs a list
            tt_refiner.set_output_p4p(p4p_file)
            tt_refiner.run()

            self._scaled_experiments = tt_refiner.get_output_experiments()
            FileHandler.record_more_data_file(
                f"{self._scalr_pname} {self._scalr_xname} scaled",
                self._scaled_experiments,
            )

            self._scalr_cell = tt_refiner.get_unit_cell()
            logger.info(
                "Overall: %6.2f %6.2f %6.2f %6.2f %6.2f %6.2f"
                % tt_refiner.get_unit_cell()
            )
            self._scalr_cell_esd = tt_refiner.get_unit_cell_esd()
            cif_in = tt_refiner.import_cif()
            mmcif_in = tt_refiner.import_mmcif()

            if params.xia2.settings.small_molecule:
                FileHandler.record_data_file(p4p_file)

            cif_out = CIF.get_block("xia2")
            mmcif_out = mmCIF.get_block("xia2")
            for key in sorted(cif_in.keys()):
                cif_out[key] = cif_in[key]
            for key in sorted(mmcif_in.keys()):
                if key.endswith("entry_id"):
                    mmcif_out[key] = "xia2"
                else:
                    mmcif_out[key] = mmcif_in[key]

            logger.debug("Unit cell obtained by two-theta refinement")

        else:
            ami = AnalyseMyIntensities()
            ami.set_working_directory(self.get_working_directory())

            average_unit_cell, _ = ami.compute_average_cell(
                [
                    self._scalr_scaled_refl_files[key]
                    for key in self._scalr_scaled_refl_files
                ]
            )

            logger.debug("Computed average unit cell (will use in all files)")
            self._scalr_cell = average_unit_cell
            self._scalr_cell_esd = None

            # Write average unit cell to .cif
            cif_out = CIF.get_block("xia2")
            cif_out[  # pylint: disable=E1137
                "_computing_cell_refinement"
            ] = "AIMLESS averaged unit cell"
            for cell, cifname in zip(
                self._scalr_cell,
                [
                    "length_a",
                    "length_b",
                    "length_c",
                    "angle_alpha",
                    "angle_beta",
                    "angle_gamma",
                ],
            ):
                cif_out["_cell_%s" % cifname] = cell  # pylint: disable=E1137

        logger.debug("%7.3f %7.3f %7.3f %7.3f %7.3f %7.3f" % self._scalr_cell)

    def apply_reindex_operator_to_sweep_info(self, si, reindex_operator, reason):
        """Use a reindex operator to reindex the data.

        Take the data from the sweep info, reindex using
        dials.reindex, and set the new data into the si.
        """
        reindexer = DialsReindex()
        reindexer.set_working_directory(self.get_working_directory())
        auto_logfiler(reindexer)

        reindexer.set_indexed_filename(si.get_reflections())
        reindexer.set_experiments_filename(si.get_experiments())
        reindexer.set_cb_op(reindex_operator)

        reindexer.run()

        si.set_reflections(reindexer.get_reindexed_reflections_filename())
        si.set_experiments(reindexer.get_reindexed_experiments_filename())

        logger.debug(
            "Reindexed with operator %s, reason is %s", reindex_operator, reason
        )

    def get_UBlattsymm_from_sweep_info(self, sweep_info):
        """Calculate U, B and lattice symmetry from experiment."""
        expt = load.experiment_list(sweep_info.get_experiments())[0]
        uc = expt.crystal.get_unit_cell()
        umatrix = expt.crystal.get_U()
        lattice_symm = lattice_symmetry_group(uc, max_delta=0.0)
        return tuple(umatrix), mosflm_B_matrix(uc), lattice_symm

    def get_mtz_data_from_sweep_info(self, sweep_info):
        """Get the data in mtz form.

        Need to run dials.export to convert the data from experiment list
        and reflection table to mtz form."""
        return self.export_to_mtz(sweep_info)

    def export_to_mtz(self, sweep_info):
        """Export to mtz, using dials.integrate phil params"""
        params = PhilIndex.params.dials.integrate
        export = ExportMtz()
        export.crystal_name = self._scalr_xname
        export.project_name = self._scalr_pname
        export.set_working_directory(self.get_working_directory())
        export.set_experiments_filename(sweep_info.get_experiments())
        export.set_reflections_filename(sweep_info.get_reflections())
        export.set_combine_partials(params.combine_partials)
        export.set_partiality_threshold(params.partiality_threshold)
        if len(sweep_info.get_batches()) == 1:
            export.set_partiality_threshold(0.1)
        if (
            len(sweep_info.get_batches()) == 1
            or PhilIndex.params.dials.fast_mode
            or not PhilIndex.params.xia2.settings.integration.profile_fitting
        ):
            # With no profiles available have to rely on summation alone
            export.set_intensity_choice("sum")

        auto_logfiler(export, "EXPORTMTZ")
        mtz_filename = os.path.join(
            self.get_working_directory(), "%s.mtz" % sweep_info.get_sweep_name()
        )
        export.set_mtz_filename(mtz_filename)
        export.run()
        return mtz_filename

class DialsScalerHelper:
    """A class to help the DIALS Scaler along a little."""

    def __init__(self):
        self._working_directory = os.getcwd()
        self._scalr_xname = None
        self._scalr_pname = None

    def set_pname_xname(self, pname, xname):
        self._scalr_xname = xname
        self._scalr_pname = pname

    def set_working_directory(self, working_directory):
        self._working_directory = working_directory

    def get_working_directory(self):
        return self._working_directory

    def assign_dataset_identifiers(self, experiments, reflections):
        """Assign unique identifiers to the datasets"""
        assigner = DialsAssignIdentifiers()
        assigner.set_working_directory(self.get_working_directory())
        auto_logfiler(assigner)
        for (exp, refl) in zip(experiments, reflections):
            assigner.add_experiments(exp)
            assigner.add_reflections(refl)
        assigner.assign_identifiers()
        return assigner

    def split_experiments(self, experiment, reflection, sweep_handler):
        """Split a multi-experiment dataset into individual datasets and set in the
        sweep handler."""
        splitter = SplitExperiments()
        splitter.add_experiments(experiment)
        splitter.add_reflections(reflection)
        splitter.set_working_directory(self.get_working_directory())
        auto_logfiler(splitter)
        splitter.run()

        nn = len(sweep_handler.get_epochs()) - 1
        fmt = "%%0%dd" % len(str(nn))

        for i, epoch in enumerate(sweep_handler.get_epochs()):
            si = sweep_handler.get_sweep_information(epoch)
            nums = fmt % i
            si.set_reflections(
                os.path.join(self.get_working_directory(), "split_%s.refl" % nums)
            )
            si.set_experiments(
                os.path.join(self.get_working_directory(), "split_%s.expt" % nums)
            )
            FileHandler.record_temporary_file(
                os.path.join(self.get_working_directory(), "split_%s.refl" % nums)
            )
            FileHandler.record_temporary_file(
                os.path.join(self.get_working_directory(), "split_%s.expt" % nums)
            )
        return sweep_handler

    def assign_and_return_datasets(self, sweep_handler):
        """Assign unique identifiers to all integrated experiments & reflections,
        and set these in the sweep_information for each epoch."""
        experiments = []
        reflections = []
        for epoch in sweep_handler.get_epochs():
            si = sweep_handler.get_sweep_information(epoch)
            integrater = si.get_integrater()
            experiments.append(integrater.get_integrated_experiments())
            reflections.append(integrater.get_integrated_reflections())
        assigner = self.assign_dataset_identifiers(experiments, reflections)
        sweep_handler = self.split_experiments(
            assigner.get_output_experiments_filename(),
            assigner.get_output_reflections_filename(),
            sweep_handler,
        )
        return sweep_handler

    def dials_symmetry_indexer_jiffy(
        self, experiments, reflections, refiners, multisweep=False
    ):
        """A jiffy to centralise the interactions between dials.symmetry
        and the Indexer, multisweep edition."""
        # First check format of input against expected input
        assert len(experiments) == len(
            reflections
        ), """
Unequal number of experiments/reflections passed to dials_symmetry_indexer_jiffy"""
        if len(experiments) > 1:
            assert multisweep, """
Passing multple datasets to indexer_jiffy but not set multisweep=True"""

        reindex_initial = False

        symmetry_analyser = self.dials_symmetry_decide_pointgroup(
            experiments, reflections
        )

        possible = symmetry_analyser.get_possible_lattices()

        logger.debug("Possible lattices (dials.symmetry):")
        logger.debug(" ".join(possible))

        # all refiners contain the same indexer link, so any good here.
        (
            correct_lattice,
            rerun_symmetry,
            need_to_return,
        ) = decide_correct_lattice_using_refiner(possible, refiners[0])

        if need_to_return and multisweep:
            if (
                PhilIndex.params.xia2.settings.integrate_p1
                and not PhilIndex.params.xia2.settings.reintegrate_correct_lattice
            ):
                need_to_return = False
                rerun_symmetry = True
            else:
                for refiner in refiners[1:]:
                    refiner.refiner_reset()

        if rerun_symmetry:
            # don't actually need to rerun, just set correct solution - this
            # call updates the relevant info in the Wrapper - but will need to reindex later
            symmetry_analyser.set_correct_lattice(correct_lattice)
            reindex_initial = True
            # rather than reindexing here, just set the reindex_inital and let the
            # scaler manage this as necessary

        logger.debug(
            "Symmetry analysis of %s", " ".join(experiments) + " ".join(reflections)
        )

        pointgroup = symmetry_analyser.get_pointgroup()
        reindex_op = symmetry_analyser.get_reindex_operator()
        probably_twinned = symmetry_analyser.get_probably_twinned()

        reindexed_reflections = symmetry_analyser.get_output_reflections_filename()
        reindexed_experiments = symmetry_analyser.get_output_experiments_filename()

        logger.debug("Pointgroup: %s (%s)", pointgroup, reindex_op)

        return (
            pointgroup,
            reindex_op,
            need_to_return,
            probably_twinned,
            reindexed_reflections,
            reindexed_experiments,
            reindex_initial,
        )

    def dials_symmetry_decide_pointgroup(self, experiments, reflections):
        """Run the symmetry analyser and return it for later inspection."""
        symmetry_analyser = DialsSymmetry()
        symmetry_analyser.set_working_directory(self.get_working_directory())
        auto_logfiler(symmetry_analyser)

        FileHandler.record_log_file(
            f"{self._scalr_pname} {self._scalr_xname} SYMMETRY",
            symmetry_analyser.get_log_file(),
        )

        for (exp, refl) in zip(experiments, reflections):
            symmetry_analyser.add_experiments(exp)
            symmetry_analyser.add_reflections(refl)
        symmetry_analyser.decide_pointgroup()

        return symmetry_analyser

    @staticmethod
    def reindex_jiffy(si, pointgroup, reindex_op):
        """Add data from si and reindex, setting back in si"""
        integrater = si.get_integrater()
        integrater.set_integrater_spacegroup_number(
            Syminfo.spacegroup_name_to_number(pointgroup)
        )
        integrater.set_integrater_reindex_operator(
            reindex_op, reason="setting point group"
        )
        integrater.set_output_format("pickle")
        integrater.get_integrater_intensities()
        # ^ This will give us the reflections in the correct point group
        si.set_reflections(integrater.get_integrated_reflections())
        si.set_experiments(integrater.get_integrated_experiments())

def decide_correct_lattice_using_refiner(possible_lattices, refiner):
    """Use the refiner to determine which of the possible lattices is the
    correct one."""
    correct_lattice, rerun_symmetry, need_to_return = (None, False, False)
    for lattice in possible_lattices:
        state = refiner.set_refiner_asserted_lattice(lattice)
        if state == refiner.LATTICE_CORRECT:
            logger.debug("Agreed lattice %s", lattice)
            correct_lattice = lattice
            break

        elif state == refiner.LATTICE_IMPOSSIBLE:
            logger.debug("Rejected lattice %s", lattice)
            rerun_symmetry = True
            continue

        elif state == refiner.LATTICE_POSSIBLE:
            logger.debug("Accepted lattice %s, will reprocess", lattice)
            need_to_return = True
            correct_lattice = lattice
            break

    if correct_lattice is None:
        correct_lattice = refiner.get_refiner_lattice()
        rerun_symmetry = True

        logger.debug("No solution found: assuming lattice from refiner")

    return correct_lattice, rerun_symmetry, need_to_return

def convert_unmerged_mtz_to_sca(mtz_filename):
    """Convert an mtz files to .sca format and write."""
    sca_filename = mtz_filename.replace("mtz", "sca")
    m = mtz.object(mtz_filename)
    for ma in m.as_miller_arrays(merge_equivalents=False, anomalous=False):
        if ma.info().labels == ["I", "SIGI"]:
            no_merge_original_index.writer(ma, file_name=sca_filename)
            FileHandler.record_data_file(sca_filename)
            break
    else:
        raise KeyError("Intensity column labels not found in MTZ file")

def convert_merged_mtz_to_sca(mtz_filename):
    """Convert an mtz files to .sca format and write."""
    # merged sca format contains 7 columns: h,k,l, I+, sigI+, I-, sigI-. We
    # always run dials.merge with anomalous=True (whether or not anomalous is
    # set in xia2), so data is always separated into I+/I- in the merged mtz.
    sca_filename = mtz_filename.replace("mtz", "sca")
    m = mtz.object(mtz_filename)
    for ma in m.as_miller_arrays(merge_equivalents=False, anomalous=True):
        if ma.info().labels == ["I(+)", "SIGI(+)", "I(-)", "SIGI(-)"]:
            merge_scalepack_write(miller_array=ma, file_name=sca_filename)
            FileHandler.record_data_file(sca_filename)
            break
    else:
        raise KeyError("Intensity column labels not found in MTZ file")

def scaling_model_auto_rules(experiment):
    """Use dials.scale rules for determining suitable parameters."""
    osc_range = experiment.scan.get_oscillation_range()
    scan_width = osc_range[1] - osc_range[0]

    if scan_width < 5.0:
        scale_interval, decay_interval = (1.0, 1.5)
    elif scan_width < 10.0:
        scale_interval, decay_interval = (2.0, 3.0)
    elif scan_width < 25.0:
        scale_interval, decay_interval = (4.0, 5.0)
    elif scan_width < 90.0:
        scale_interval, decay_interval = (8.0, 10.0)
    else:
        scale_interval, decay_interval = (15.0, 20.0)
    return scale_interval, decay_interval


import os, sys
import argparse
import cv2
import time
import dlib
import glob
import mxnet as mx
import numpy as np

sys.path.append("..")

from models import LffdDet, FeatureEncoder
from utils import FaceAligner, enlarge_bboxes, FaceCompare, face_track

det_sym_path = "model_zoo/detector/lffd_facedet.json"
det_para_path = "model_zoo/detector/lffd_facedet.params"
extor_sym_path = "model_zoo/extractor/mobilefacenet-symbol-1111.json"
extor_para_path = "model_zoo/extractor/mobilefacenet-1111.params"

HEIGHT = 1080
WIDTH = 1920
SCALE = 8
FACE_LIB = './face_lib'
IOU_THRES = 0.7
SCORE_THRES = 0.98
FACE_TRACK = True
AREA_THRES = 900

def parse_args():
    parser = argparse.ArgumentParser(description="Demo of face detection and verification")
    parser.add_argument('--context', type=str, default='cpu',
                        help='The context (cpu or gpu) of model,')
    parser.add_argument('--face-lib', type=str, default=None,
                        help='The face lib for verifivation.')
    parser.add_argument('--shape', type=tuple, default=(1080, 1920),
                        help='The shape of video/image.')
    parser.add_argument('--scale', type=int, default=8,
                        help='The scale for image resize.')
    parser.add_argument('--HEIGTH', type=int, default=1080,
                        help='The IOU for NMS in face tracking.')
    args = parser.parse_args()
    return args

def build_model(ctx):
    assert os.path.isfile(det_sym_path), "{} is not a file.".format(det_sym_path)
    assert os.path.isfile(det_para_path), "{} is not a file.".format(det_para_path)
    assert os.path.isfile(extor_sym_path), "{} is not a file.".format(extor_sym_path)
    assert os.path.isfile(extor_para_path), "{} is not a file.".format(extor_para_path)

    face_detector = LffdDet(symbol_file_path=det_sym_path,
                            model_file_path=det_para_path,
                            ctx=ctx)
    face_encoder = FeatureEncoder(symbol_path=extor_sym_path,
                                  params_path=extor_para_path,
                                  ctx=ctx)
    face_compare = FaceCompare(lib_path=FACE_LIB,
                               face_encoder=face_encoder,
                               thresh=SCORE_THRES)
    face_aligner = FaceAligner()
    return face_detector, face_compare, face_aligner

def area(bbox):
    return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])

def get_face_lib(lib_path, feats_extector):
    assert os.path.isdir(lib_path), "{} is not face-lib dir.".format(lib_path)
    folds = os.listdir(lib_path)
    face_id = [f for f in folds if f[0] is not '.']
    assert len(face_id) > 0, "There is no face under {}.".format(lib_path)
    face_ind2id = []
    face_ecoding = []
    for id in face_id:
        id_root = os.path.join(lib_path, id)
        imgs = os.listdir(id, root)
        imgs = [im for im in imgs if im.endswith('.jpg')]
        for img in imgs:
            pic = cv2.cvtColor(cv2.imread(os.path.join(id_root, img)), cv2.COLOR_BGR2RGB)
            face_ecoding.append(feats_extector(pic))
            face_ind2id.append(id)
    face_ecoding = np.concatenate(face_ecoding, axis=0)

    return face_ecoding, face_ind2id

def main():
    face_detector, face_compare, face_aligner = build_model(mx.gpu())
    # build cam
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, WIDTH)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, HEIGHT)

    pre_bboxes = []
    pre_result = []
    while(cap.isOpened()):
        ret, flame = cap.read()
        if ret  == True:
            flame = cv2.flip(flame, 1)
            h, w, _ = flame.shape
            flame_ = cv2.resize(flame,(w//SCALE, h//SCALE))
            bboxes = face_detector.predict(flame_, resize_scale=1, score_threshold=0.9, top_k=10000, \
                                                        NMS_threshold=0.4, NMS_flag=True, skip_scale_branch_list=[])

            if len(pre_bboxes) > 0:
                pre_frame_bboxes, pre_frame_result, unknown_bboxes = face_track(pre_bboxes, pre_result, bboxes)
            else:
                pre_frame_bboxes = []
                pre_frame_result = []
                unknown_bboxes = bboxes
            #to_encoding_bboxes = [enlarge_bboxes(bbox, flame_.shape) for bbox in unknown_bboxes]
            to_encoding_bboxes = unknown_bboxes
            to_encoding_faces = [face_aligner.align(flame_, bbox) for bbox in to_encoding_bboxes]
            # result = []
            final_bboxes = pre_frame_bboxes
            final_result = pre_frame_result
            for face, bbox in zip(to_encoding_faces, unknown_bboxes):
                if area(bbox) < AREA_THRES:
                    indx = -1
                    score = 0.0
                else:
                    indx, score = face_compare.compare(face)
                    if score < SCORE_THRES:
                        indx = -1
                bbox[-1] = indx
                bbox[-2] = score
                final_bboxes.append(bbox)
                final_result.append(face_compare.face_id[indx])

            # final_bboxes = pre_frame_bboxes + unknown_bboxes
            # final_result = pre_frame_result + result
            if FACE_TRACK:
                pre_bboxes = final_bboxes
                pre_result = final_result
            else:
                pre_bboxes = []
                pre_result = []

            for bbox, name in zip(final_bboxes, final_result):
                score = bbox[-2]
                bbox = [int(bb*SCALE+SCALE//2) for bb in bbox]
                if name == 'UNKNOWN':
                    color = (0, 0, 255)
                    cv2.rectangle(flame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 1)
                else:
                    color = (255, 255, 0)
                    cv2.rectangle(flame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 1)
                    cv2.rectangle(flame, (bbox[0], bbox[1]-30), (bbox[2], bbox[1]), color, -1)
                    cv2.putText(flame, ' '+name+'  {:.2f}'.format(score), (bbox[0], bbox[1]-5), cv2.FONT_HERSHEY_DUPLEX, 0.9, (255, 255, 255), 2)
            cv2.imshow('video', flame)
            if cv2.waitKey(1)&0xFF==ord('q'):
                break
        else:
            break

if __name__ == "__main__":
    main()


# Licensed under the BSD license.  See full license in LICENSE file.
# http://www.lightshowpi.com/
#
# Author: Todd Giles (todd@lightshowpi.com)
# Author: Chris Usey (chris.usey@gmail.com)
# Author: Ryan Jennings
# Author: Paul Dunn (dunnsept@gmail.com)
# Author: Tom Enos (tomslick.ca@gmail.com)

"""Play any audio file and synchronize lights to the music

When executed, this script will play an audio file, as well as turn on
and off N channels of lights to the music (by default the first 8 GPIO
channels on the Rasberry Pi), based upon music it is playing. Many
types of audio files are supported, but it has
only been tested with wav and mp3 at the time of this writing.

The timing of the lights turning on and off is based upon the frequency
response of the music being played.  A short segment of the music is
analyzed via FFT to get the frequency response across each defined
channel in the audio range.  Each light channel is then faded in and
out based upon the amplitude of the frequency response in the
corresponding audio channel.  Fading is accomplished with a software
PWM output.  Each channel can also be configured to simply turn on and
off as the frequency response in the corresponding channel crosses a
threshold.

FFT calculation can be CPU intensive and in some cases can adversely
affect playback of songs (especially if attempting to decode the song
as well, as is the case for an mp3).  For this reason, the FFT
cacluations are cached after the first time a new song is played.
The values are cached in a gzip'd text file in the same location as the
song itself.  Subsequent requests to play the same song will use the
cached information and not recompute the FFT, thus reducing CPU
utilization dramatically and allowing for clear music playback of all
audio file types.

Recent optimizations have improved this dramatically and most users are
no longer reporting adverse playback of songs even on the first
playback.

Sample usage:
To play an entire list -
sudo python synchronized_lights.py --playlist=/home/pi/music/.playlist

To play a specific song -
sudo python synchronized_lights.py --file=/home/pi/music/jingle_bells.mp3

Third party dependencies:

alsaaudio: for audio input/output
    http://pyalsaaudio.sourceforge.net/

numpy: for FFT calcuation
    http://www.numpy.org/
"""

import argparse
import atexit
import audioop
import contextlib
import csv
import json
import logging
import os
import random
import subprocess
import sys
import threading
import time

import alsaaudio as aa
import numpy as np

import audio_decoder
import audio_input
import audio_output
import fft
import configuration_manager as cm
import hardware_controller as hc
import playlist
import prepostshow
import running_stats

# TODO(mdietz): as many of these should have defaults as possible
# Configurations - TODO(todd): Move more of this into configuration manager
_CONFIG = cm.CONFIG
_MODE = cm.lightshow()['mode']
_MIN_FREQUENCY = _CONFIG.getfloat('audio_processing', 'min_frequency')
_MAX_FREQUENCY = _CONFIG.getfloat('audio_processing', 'max_frequency')
_RANDOMIZE_PLAYLIST = _CONFIG.getboolean('lightshow', 'randomize_playlist')

try:
    _CUSTOM_CHANNEL_MAPPING = [
        int(channel)
        for channel in _CONFIG.get('audio_processing',
                                   'custom_channel_mapping').split(',')]
except ValueError:
    _CUSTOM_CHANNEL_MAPPING = 0

try:
    _CUSTOM_CHANNEL_FREQUENCIES = [
        int(channel) for channel in
        _CONFIG.get('audio_processing',
                    'custom_channel_frequencies').split(',')]
except ValueError:
    _CUSTOM_CHANNEL_FREQUENCIES = 0

_PLAYLIST_PATH = cm.lightshow()['playlist_path'].replace(
    '$SYNCHRONIZED_LIGHTS_HOME', cm.HOME_DIR)

# TODO(mdietz): Changing this would necessitate rebuilding the cache. Not sure
#               the code knows that
CHUNK_SIZE = _CONFIG.getint("audio_processing", "chunk_size")

def end_early():
    """atexit function"""
    logging.critical("Atexit triggered with CLEAN_EXIT %s", CLEAN_EXIT)
    if not CLEAN_EXIT:
        hc.clean_up()

atexit.register(end_early)

def update_lights(matrix, mean, std):
    """Update the state of all the lights

    Update the state of all the lights based upon the current
    frequency response matrix

    :param matrix: row of data from cache matrix
    :type matrix: list

    :param mean: standard mean of fft values
    :type mean: list

    :param std: standard deviation of fft values
    :type std: list
    """
    for pin in xrange(0, hc.GPIOLEN):
        # Calculate output pwm, where off is at some portion of the std below
        # the mean and full on is at some portion of the std above the mean.
        brightness = matrix[pin] - mean[pin] + 0.5 * std[pin]
        brightness /= 1.25 * std[pin]
        if brightness > 1.0:
            brightness = 1.0

        if brightness < 0:
            brightness = 0

        if not hc.is_pin_pwm[pin]:
            # If pin is on / off mode we'll turn on at 1/2 brightness
            # TODO(mdietz): Configurable per channel threshold!
            if brightness > 0.5:
                hc.turn_on_light(pin, True)
            else:
                hc.turn_off_light(pin, True)
        else:
            hc.turn_on_light(pin, True, brightness)

def audio_in():
    """Control the lightshow from audio coming in from a USB audio card"""
    sample_rate = cm.lightshow()['audio_in_sample_rate']
    input_channels = cm.lightshow()['audio_in_channels']

    # Open the input stream from default input device
    audio_in_card = cm.lightshow()['audio_in_card']
    stream = aa.PCM(aa.PCM_CAPTURE, aa.PCM_NORMAL, audio_in_card)
    stream.setchannels(input_channels)
    stream.setformat(aa.PCM_FORMAT_S16_LE)  # Expose in config if needed
    stream.setrate(sample_rate)
    stream.setperiodsize(CHUNK_SIZE)

    logging.debug("Running in audio-in mode - will run until Ctrl+C "
                  "is pressed")
    print "Running in audio-in mode, use Ctrl+C to stop"

    # Start with these as our initial guesses - will calculate a rolling
    # mean / std as we get input data.
    mean = np.array([12.0] * hc.GPIOLEN, dtype='float64')
    std = np.array([1.5] * hc.GPIOLEN, dtype='float64')
    count = 2

    running_stats = running_stats.Stats(hc.GPIOLEN)

    # preload running_stats to avoid errors, and give us a show that looks
    # good right from the start
    running_stats.preload(mean, std, count)

    try:
        hc.initialize()
        fft_calc = fft.FFT(CHUNK_SIZE,
                           sample_rate,
                           hc.GPIOLEN,
                           _MIN_FREQUENCY,
                           _MAX_FREQUENCY,
                           _CUSTOM_CHANNEL_MAPPING,
                           _CUSTOM_CHANNEL_FREQUENCIES,
                           input_channels)

        # Listen on the audio input device until CTRL-C is pressed
        while True:
            length, data = stream.read()
            if length > 0:
                # if the maximum of the absolute value of all samples in
                # data is below a threshold we will disreguard it
                audio_max = audioop.max(data, 2)
                if audio_max < 250:
                    # we will fill the matrix with zeros and turn the
                    # lights off
                    matrix = np.zeros(hc.GPIOLEN, dtype="float64")
                    logging.debug("below threshold: '" + str(
                        audio_max) + "', turning the lights off")
                else:
                    matrix = fft_calc.calculate_levels(data)
                    running_stats.push(matrix)
                    mean = running_stats.mean()
                    std = running_stats.std()

                update_lights(matrix, mean, std)

    except KeyboardInterrupt:
        pass

    finally:
        print "\nStopping"
        hc.clean_up()

def load_cached_fft(fft_calc, cache_filename):
     # Read in cached fft
    mean, std, cache_matrix = None, None, None
    try:
        # load cache from file using numpy loadtxt
        cache_matrix = np.loadtxt(cache_filename)

        # compare configuration of cache file to current configuration
        cache_found = fft_calc.compare_config(cache_filename)
        if not cache_found:
            # create empty array for the cache_matrix
            cache_matrix = np.empty(shape=[0, hc.GPIOLEN])
            raise IOError()

        # get std from matrix / located at index 0
        std = np.array(cache_matrix[0])

        # get mean from matrix / located at index 1
        mean = np.array(cache_matrix[1])

        # delete mean and std from the array
        cache_matrix = np.delete(cache_matrix, 0, axis=0)
        cache_matrix = np.delete(cache_matrix, 0, axis=0)

        logging.debug("std: " + str(std) + ", mean: " + str(mean))
    except IOError:
        cache_found = False
        logging.warn("Cached sync data song_filename not found: '"
                     + cache_filename
                     + "'.  One will be generated.")

    return cache_found, mean, std, cache_matrix

# TODO(mdietz): cache dir should be configurable
def get_song_cache(song_filename, chunk_size):
    music_file = audio_decoder.open(song_filename)
    sample_rate = music_file.getframerate()
    num_channels = music_file.getnchannels()
    logging.info("Sample rate: %s" % sample_rate)
    logging.info("Channels: %s" % num_channels)
    logging.info("Frame size: %s" % music_file.getsampwidth())

    fft_calc = fft.FFT(chunk_size,
                       sample_rate,
                       hc.GPIOLEN,
                       _MIN_FREQUENCY,
                       _MAX_FREQUENCY,
                       _CUSTOM_CHANNEL_MAPPING,
                       _CUSTOM_CHANNEL_FREQUENCIES)

    # Init cache matrix
    cache_matrix = np.empty(shape=[0, hc.GPIOLEN])
    cache_filename = os.path.dirname(song_filename) + "/." + os.path.basename(
        song_filename) + ".sync"
    config_found = fft_calc.compare_config(cache_filename)

    cache_found = False
    if config_found:
        cache_found, mean, std, cache_matrix = load_cached_fft(fft_calc, cache_filename)

    if not cache_found:
        # Reset the cache matrix
        cache_matrix = np.empty(shape=[0, hc.GPIOLEN])

        # The values 12 and 1.5 are good estimates for first time playing back
        # (i.e. before we have the actual mean and standard deviations
        # calculated for each channel).
        mean = [12.0] * hc.GPIOLEN
        std = [1.5] * hc.GPIOLEN
        total = 0
        while True:
            data = music_file.readframes(chunk_size)
            if not data:
                break
            total += len(data)

            matrix = fft_calc.calculate_levels(data)
            # Add the matrix to the end of the cache
            cache_matrix = np.vstack([cache_matrix, matrix])

        for i in range(0, hc.GPIOLEN):
            std[i] = np.std([item for item in cache_matrix[:, i]
                             if item > 0])
            mean[i] = np.mean([item for item in cache_matrix[:, i]
                               if item > 0])

        # Add mean and std to the top of the cache
        cache_matrix = np.vstack([mean, cache_matrix])
        cache_matrix = np.vstack([std, cache_matrix])

        # Save the cache using numpy savetxt
        np.savetxt(cache_filename, cache_matrix)

        # Save fft config
        fft_calc.save_config()

        logging.info("Cached sync data written to '." + cache_filename
                     + "' [" + str(len(cache_matrix)) + " rows]")

        logging.info("Cached config data written to '." +
                     fft_calc.config_filename)
    music_file.close()
    return mean, std, cache_matrix

# TODO(mdietz): the project assumes that this is simply a python bin
#               that is called by a bash script. I think we should
#               re-work that to be a long running thing, all in Python
def play_song(num_songs):
    """Play the next song from the play list (or --file argument)."""
    play_now = int(cm.get_state('play_now', "0"))

    # Make sure one of --playlist or --file was specified
    if args.file is None and args.playlist is None:
        print "One of --playlist or --file must be specified"
        sys.exit()

    current_playlist = playlist.Playlist(args.playlist, num_songs)

    # Initialize Lights
    hc.initialize()

    for (song_title,
         song_filename,
         chunk_size) in current_playlist.get_song():

        # Handle the pre/post show
        if not play_now:
            result = prepostshow.PrePostShow('preshow', hc).execute()

            if result == prepostshow.PrePostShow.play_now_interrupt:
                play_now = int(cm.get_state('play_now', "0"))

        # TODO(mdietz): What the hell is this play_now variable really for?
        # Ensure play_now is reset before beginning playback
        if play_now:
            cm.update_state('play_now', "0")
            play_now = 0

        mean, std, cache_matrix = get_song_cache(song_filename, chunk_size)

        # NOTE(mdietz): Adapt this to a standard radio, not an SDR. The SDR
        #               has a clear extra amount of delay
        light_show_delay = _CONFIG.getfloat("lightshow", "light_delay")
        logging.info("Delaying light show by %f seconds" % light_show_delay)

        audio_in_stream = audio_input.get_audio_input_handler(song_filename,
                                                              chunk_size)
        audio_out_stream = audio_output.get_audio_output_handler(
            audio_in_stream.num_channels, audio_in_stream.sample_rate,
            song_title, chunk_size)

        try:
            # Process audio
            row = 0
            start_time = time.time()
            while True:
                data = audio_in_stream.next_chunk()
                if not data or play_now:
                    break

                audio_out_stream.write(data)
                # TODO(mdietz): This actually pretty much works, but it would
                #               be nice to figure out what the actual delay
                #               time is, and also make it a config value
                # TODO(mdietz): I may be able to time the popen to first stdout
                #               from the fm proc for a dynamic delay
                if time.time() - start_time < light_show_delay:
                    continue

                matrix = cache_matrix[row]
                update_lights(matrix, mean, std)

                # Read next chunk of data from music

                # Load new application state in case we've been interrupted
                # TODO(mdietz): not the way to do this. Read from a db,
                #               accept a signal or some other OOB proc
                cm.load_state()
                play_now = int(cm.get_state('play_now', "0"))
                row += 1

            # Cleanup the fm process if there is one
        except Exception:
            logging.exception("Error in playback")
        finally:
            audio_out_stream.cleanup()

        # check for postshow
        prepostshow.PrePostShow('postshow', hc).execute()

if __name__ == "__main__":
    # Arguments
    global CLEAN_EXIT
    CLEAN_EXIT = False
    parser = argparse.ArgumentParser()
    parser.add_argument('--log', default='INFO',
                        help='Set the logging level. levels:INFO, DEBUG,'
                             'WARNING, ERROR, CRITICAL')
    parser.add_argument("--num_songs", default=4,
                        help="Number of songs to play from the playlist")

    filegroup = parser.add_mutually_exclusive_group()
    filegroup.add_argument('--playlist', default=_PLAYLIST_PATH,
                           help='Playlist to choose song from.')
    filegroup.add_argument('--file', help='path to the song to play '
                                          '(required if no playlist is '
                                          'designated)')
    args = parser.parse_args()

    logging.basicConfig(filename=cm.LOG_DIR + '/music_and_lights.play.dbg',
                        format='[%(asctime)s] %(levelname)s '
                               '{%(pathname)s:%(lineno)d} - %(message)s',
                        level=logging.INFO)

    # logging levels
    levels = {'DEBUG': logging.DEBUG,
              'INFO': logging.INFO,
              'WARNING': logging.WARNING,
              'ERROR': logging.ERROR,
              'CRITICAL': logging.CRITICAL}

    level = levels.get(parser.parse_args().log.upper())
    logging.getLogger().setLevel(level)

    try:
        if cm.lightshow()['mode'] == 'audio-in':
            audio_in()
        else:
            play_song(args.num_songs)
            CLEAN_EXIT = True
    except Exception:
        logging.exception("Something blew up")
    finally:
        hc.turn_on_all_lights()