Documentation for Controllers

class AIF(Controller):
    """
    A controller that implements the All-ion-fragmentation (AIF) DIA fragmentation strategy.
    Should be used in conjunction with MS-DIAL for deconvolution.
    """

    def __init__(self, ms1_source_cid_energy, advanced_params=None):
        """
        Initialise an AIF controller
        Args:
            ms1_source_cid_energy: source CID energy for MS1 scan
            advanced_params: an [vimms.Controller.base.AdvancedParams][] object that contains
                             advanced parameters to control the mass spec. If left to None,
                             default values will be used.
        """
        super().__init__(advanced_params=advanced_params)
        self.scan_number = self.initial_scan_id
        self.ms1_source_cid_energy = ms1_source_cid_energy

    def write_msdial_experiment_file(self, filename):
        """
        Generates a file that can be read by MS-DIAL to perform deconvolution

        Args:
            filename: path to experiment file in MS-DIAL format

        Returns: None

        """
        heads = ["ID", "MS Type", "Start m/z", "End m/z", "Name", "CE", "DecTarget(1:Yes, 0:No)"]
        start = self.advanced_params.default_ms1_scan_window[0]
        stop = self.advanced_params.default_ms1_scan_window[1]
        ce = self.ms1_source_cid_energy
        ms1_row = ["0", "SCAN", start, stop, "0eV", 0, 0]
        aif_row = ["1", "ALL", start, stop, "{}eV".format(ce), ce, 1]

        out_dir = os.path.dirname(filename)
        create_if_not_exist(out_dir)

        with open(filename, "w", newline="") as f:
            writer = csv.writer(f, delimiter="\t", dialect="excel")
            writer.writerow(heads)
            writer.writerow(ms1_row)
            writer.writerow(aif_row)

    def update_state_after_scan(self, last_scan):
        pass

    def _process_scan(self, scan):
        """
        This method is called when a scan arrives that requires action.
        Normally means that we should schedule some more, but in DIA we don't need
        to actually look at the peaks in the scan, so just schedule the next block

        Args:
            scan: A new [vimms.MassSpec.Scan][] to process.

        Returns: newly generated scans

        """

        # For all ions fragmentation, when we receive the last scan of
        # the previous block, we make a new block. Each block is an MS1 scan
        # followed by an MS2 scan where the MS2 fragmens everything
        scans = []

        if self.scan_to_process is not None:
            # make the MS1 scan with source cid energy applied
            aif_scan = self.get_ms1_scan_params()
            aif_scan.set(ScanParameters.SOURCE_CID_ENERGY, self.ms1_source_cid_energy)
            self._check_scan(aif_scan)

            scans.append(aif_scan)
            self.scan_number += 1  # increase every time we make a scan

            # make the MS1 scan with no energy applied
            ms1_scan = self.get_ms1_scan_params()
            self._check_scan(ms1_scan)

            scans.append(ms1_scan)
            self.scan_number += 1
            self.next_processed_scan_id = self.scan_number

            # set this ms1 scan as has been processed
            self.scan_to_process = None

        return scans

class AdvancedParams:
    """
    An object that stores advanced parameters used to control the mass spec
    e.g. AGC target, collision energy, orbitrap resolution etc.

    When ViMMS is connected to an actual mass spec instrument (Orbitrap Fusion in our case)
    via IAPI, most of these values are directly passed to the mass spec as they are.
    Generally you can leave these settings to their default.

    In simulated use, most of these values are not used and therefore they won't affect simulated
    results.
    """

    def __init__(
        self,
        default_ms1_scan_window=DEFAULT_MS1_SCAN_WINDOW,
        ms1_agc_target=DEFAULT_MS1_AGC_TARGET,
        ms1_max_it=DEFAULT_MS1_MAXIT,
        ms1_collision_energy=DEFAULT_MS1_COLLISION_ENERGY,
        ms1_orbitrap_resolution=DEFAULT_MS1_ORBITRAP_RESOLUTION,
        ms1_activation_type=DEFAULT_MS1_ACTIVATION_TYPE,
        ms1_mass_analyser=DEFAULT_MS1_MASS_ANALYSER,
        ms1_isolation_mode=DEFAULT_MS1_ISOLATION_MODE,
        ms1_source_cid_energy=DEFAULT_SOURCE_CID_ENERGY,
        ms2_agc_target=DEFAULT_MS2_AGC_TARGET,
        ms2_max_it=DEFAULT_MS2_MAXIT,
        ms2_collision_energy=DEFAULT_MS2_COLLISION_ENERGY,
        ms2_orbitrap_resolution=DEFAULT_MS2_ORBITRAP_RESOLUTION,
        ms2_activation_type=DEFAULT_MS2_ACTIVATION_TYPE,
        ms2_mass_analyser=DEFAULT_MS2_MASS_ANALYSER,
        ms2_isolation_mode=DEFAULT_MS2_ISOLATION_MODE,
        ms2_source_cid_energy=DEFAULT_SOURCE_CID_ENERGY,
    ):
        """
        Create an advanced parameter object

        Args:
            default_ms1_scan_window: the m/z window to perform MS1 scan
            ms1_agc_target: automatic gain control target for MS1 scan
            ms1_max_it: maximum time to acquire ions for MS1 scan
            ms1_collision_energy: the collision energy used for MS1 scan
            ms1_orbitrap_resolution: the Orbitrap resolution used for MS1 scan
            ms1_activation_type: the activation type for MS1 scan, either CID or HCD
            ms1_mass_analyser: the mass analyser to use for MS1 scan, either IonTrap or Orbitrap
            ms1_isolation_mode: the isolation mode for MS1 scan, either None, Quadrupole, IonTrap
            ms1_source_cid_energy: source CID energy
            ms2_agc_target: automatic gain control target for MS2 scan
            ms2_max_it: maximum time to acquire ions for MS2 scan
            ms2_collision_energy: the collision energy used for MS2 scan
            ms2_orbitrap_resolution: the Orbitrap resolution used for MS2 scan
            ms2_activation_type: the activation type for MS2 scan, either CID or HCD
            ms2_mass_analyser: the mass analyser to use for MS2 scan, either IonTrap or Orbitrap
            ms2_isolation_mode: the isolation mode for MS2 scan, either None, Quadrupole, IonTrap
            ms2_source_cid_energy: source CID energy
        """
        self.default_ms1_scan_window = default_ms1_scan_window

        self.ms1_agc_target = ms1_agc_target
        self.ms1_max_it = ms1_max_it
        self.ms1_collision_energy = ms1_collision_energy
        self.ms1_orbitrap_resolution = ms1_orbitrap_resolution
        self.ms1_activation_type = ms1_activation_type
        self.ms1_mass_analyser = ms1_mass_analyser
        self.ms1_isolation_mode = ms1_isolation_mode
        self.ms1_source_cid_energy = ms1_source_cid_energy

        self.ms2_agc_target = ms2_agc_target
        self.ms2_max_it = ms2_max_it
        self.ms2_collision_energy = ms2_collision_energy
        self.ms2_orbitrap_resolution = ms2_orbitrap_resolution
        self.ms2_activation_type = ms2_activation_type
        self.ms2_mass_analyser = ms2_mass_analyser
        self.ms2_isolation_mode = ms2_isolation_mode
        self.ms2_source_cid_energy = ms2_source_cid_energy

class AgentBasedController(Controller):
    """
    A class that implements an agent-based controller.
    """

    def __init__(self, agent, advanced_params=None):
        """Initialises an agent-based controller.

        Arguments:
            agent: an instance of the [vimms.Agent.AbstractAgent][] class.
            advanced_params: optional advanced parameters for the mass spec.
        """
        super().__init__(advanced_params=advanced_params)
        self.agent = agent

    def _process_scan(self, scan):
        new_tasks = []
        if self.scan_to_process is not None:
            new_tasks, self.current_task_id, self.next_processed_scan_id = self.agent.next_tasks(
                self.scan_to_process, self, self.current_task_id
            )
            self.scan_to_process = None  # has been processed
        return new_tasks

    def update_state_after_scan(self, last_scan):
        self.agent.update(last_scan, self)

class CaseControlNonOverlapController(TopNEXtController):
    """
    Case-control non-overlap controller
    """

    def __init__(
        self,
        ionisation_mode,
        isolation_width,
        N,
        mz_tol,
        rt_tol,
        min_ms1_intensity,
        roi_params,
        grid,
        smartroi_params=None,
        min_roi_length_for_fragmentation=1,
        ms1_shift=0,
        advanced_params=None,
        register_all_roi=False,
        scoring_params=GRID_CONTROLLER_SCORING_PARAMS,
        exclusion_method=ROI_EXCLUSION_DEW,
        exclusion_t_0=None,
    ):
        super().__init__(
            ionisation_mode,
            isolation_width,
            N,
            mz_tol,
            rt_tol,
            min_ms1_intensity,
            roi_params,
            grid,
            smartroi_params=smartroi_params,
            min_roi_length_for_fragmentation=min_roi_length_for_fragmentation,
            ms1_shift=ms1_shift,
            advanced_params=advanced_params,
            register_all_roi=register_all_roi,
            scoring_params=scoring_params,
            exclusion_method=exclusion_method,
            exclusion_t_0=exclusion_t_0,
        )
        self.scoring_params = scoring_params
        if self.scoring_params["theta3"] != 0 and self.register_all_roi is False:
            print("Warning: register_all_roi should be set to True id theta3 is not 0")

    def _get_scores(self):
        scores = [
            self.grid.case_control_non_overlap(
                r, self.current_roi_intensities[i], self.scoring_params
            )
            for i, r in enumerate(self.live_roi)
        ]
        return self._get_top_N_scores(scores * self._score_filters())

class Controller(ABC):
    """
    Abtract base class for controllers.
    """

    def __init__(self, advanced_params=None):
        """
        Initialise a base Controller class.

        Args:
            advanced_params: an [vimms.Controller.base.AdvancedParams][] object that contains
                             advanced parameters to control the mass spec. If left to None,
                             default values will be used.
        """
        if advanced_params is None:
            self.advanced_params = AdvancedParams()
        else:
            self.advanced_params = advanced_params

        self.scans = defaultdict(list)  # key: ms level, value: list of scans for that level
        self.scan_to_process = None
        self.environment = None
        self.next_processed_scan_id = INITIAL_SCAN_ID
        self.initial_scan_id = INITIAL_SCAN_ID
        self.current_task_id = self.initial_scan_id
        self.processing_times = []
        self.last_ms1_rt = 0.0

    def __repr__(self):
        return f"{type(self)}({','.join(f'{k}={v}' for k, v in self.__dict__.items())})"

    def get_ms1_scan_params(self, metadata=None):
        """
        Generate a default scan parameter for MS1 scan. The generated scan parameter object
        is typically passed to the mass spec (whether real or simulated) that produces
        the actual MS1 scan.
        Args:
            metadata: any additional metadata to include

        Returns: a [vimms.Common.ScanParameters][] object that describes the MS1 scan to generate.

        """
        task = get_default_scan_params(
            polarity=self.environment.mass_spec.ionisation_mode,
            default_ms1_scan_window=self.advanced_params.default_ms1_scan_window,
            agc_target=self.advanced_params.ms1_agc_target,
            max_it=self.advanced_params.ms1_max_it,
            collision_energy=self.advanced_params.ms1_collision_energy,
            source_cid_energy=self.advanced_params.ms1_source_cid_energy,
            orbitrap_resolution=self.advanced_params.ms1_orbitrap_resolution,
            activation_type=self.advanced_params.ms1_activation_type,
            mass_analyser=self.advanced_params.ms1_mass_analyser,
            isolation_mode=self.advanced_params.ms1_isolation_mode,
            metadata=metadata,
        )
        return task

    def get_ms2_scan_params(
        self, mz, intensity, precursor_scan_id, isolation_width, mz_tol, rt_tol, metadata=None
    ):
        """
        Generate a default scan parameter for MS2 scan. The generated scan parameter object
        is typically passed to the mass spec (whether real or simulated) that produces
        the actual MS2 scan.

        Args:
            mz: the m/z of the precursor ion to fragment
            intensity: the intensity of the precursor ion to fragment
            precursor_scan_id: the associated MS1 scan ID that contains this precursor ion
            isolation_width: isolation width, in Dalton
            mz_tol: m/z tolerance for dynamic exclusion (TODO: this shouldn't be here)
            rt_tol: RT tolerance for dynamic exclusion (TODO: this shouldn't be here)
            metadata: any additional metadata to include

        Returns: a [vimms.Common.ScanParameters][] object that describes the MS2 scan to generate.

        """
        task = get_dda_scan_param(
            mz,
            intensity,
            precursor_scan_id,
            isolation_width,
            mz_tol,
            rt_tol,
            agc_target=self.advanced_params.ms2_agc_target,
            max_it=self.advanced_params.ms2_max_it,
            collision_energy=self.advanced_params.ms2_collision_energy,
            source_cid_energy=self.advanced_params.ms2_source_cid_energy,
            orbitrap_resolution=self.advanced_params.ms2_orbitrap_resolution,
            activation_type=self.advanced_params.ms2_activation_type,
            mass_analyser=self.advanced_params.ms2_mass_analyser,
            isolation_mode=self.advanced_params.ms2_isolation_mode,
            polarity=self.environment.mass_spec.ionisation_mode,
            metadata=metadata,
        )
        return task

    def get_initial_tasks(self):
        """
        Gets the initial tasks to load immediately into the mass spec
        (before acquisition starts)

        Returns: an empty list of tasks, unless overridden by subclass

        """
        return []

    def get_initial_scan_params(self):
        """
        Gets the initial scan parameters to send to the mass spec that
        starts the whole process. Will default to sending an MS1 scan with
        whatever parameters passed in self.params. Subclasses can override
        this to return different types of scans.

        Returns: a [vimms.Common.ScanParameters][] object describing the initial scan to make.

        """
        return self.get_ms1_scan_params()

    def set_environment(self, env):
        """
        Set the environment used to run this controller

        Args:
            env: an [vimms.Environment.Environment] object or its subclasses.

        Returns: None

        """
        self.environment = env

    def handle_scan(self, scan, current_size, pending_size):
        """
        Basic codes to handle an incoming scan, which is generally the same for all controllers.

        Args:
            scan: A new [vimms.MassSpec.Scan][] to process.
            current_size: current size of task buffer
            pending_size: pending size of task buffer

        Returns: a list of new [vimms.Common.ScanParameters][] describing what to do next.

        """
        # record every scan that we've received
        self.scans[scan.ms_level].append(scan)

        # update ms1 time (used for ROI matching)
        if scan.ms_level == 1:
            self.last_ms1_rt = scan.rt
            self.last_ms1_scan = scan

        # we get an ms1 scan and it has some peaks AND all the pending tasks
        # have been sent and processed AND this ms1 scan is a custom scan
        # we'd sent before (not a method scan) then store it for
        # fragmentation next time

        if scan.scan_id == self.next_processed_scan_id:
            self.scan_to_process = scan
        else:
            self.scan_to_process = None

        # implemented by subclass
        if self.scan_to_process is not None:
            # track how long each scan takes to process
            start = time.time()
            new_tasks = self._process_scan(scan)
            elapsed = time.time() - start
            self.processing_times.append(elapsed)
        else:
            # this scan is not the one we want to process, but here we
            # pass it to _process_scan anyway in case the subclass wants
            # to do something with it
            new_tasks = self._process_scan(scan)
        return new_tasks

    @abstractmethod
    def update_state_after_scan(self, last_scan):
        """Update internal state after a scan has been processed.

        Arguments:
            last_scan ([vimms.MassSpec.Scan][]): the last-processed object.
        """
        pass

    @abstractmethod
    def _process_scan(self, scan):
        """Process incoming scan

        Arguments:
            scan: A new [vimms.MassSpec.Scan][] to process.
        """
        pass

    def dump_scans(self, output_method):
        """
        Dump all scans to the output format.
        Useful for debugging.

        Args:
            output_method: a function that accepts scan information as a CSV string from pandas

        Returns: None

        """
        all_scans = self.scans[1] + self.scans[2]
        all_scans.sort(key=lambda x: x.scan_id)  # sort by scan_id
        out_list = []
        for scan in all_scans:
            # ignore any scan that we didn't send (no scan_params)
            if scan.scan_params is not None:
                out = {
                    "scan_id": scan.scan_id,
                    "num_peaks": scan.num_peaks,
                    "rt": scan.rt,
                    "ms_level": scan.ms_level,
                }
                # add all the scan params to out
                out.update(scan.scan_params.get_all())
                out_list.append(out)

        # dump to csv
        df = pd.DataFrame(out_list)
        output_method(df.to_csv(index=False, line_terminator="\n"))

    def _check_scan(self, params):
        """
        Checks that the conditions that are checked in
        vimms-fusion MS class pass here. This is done to ensure that the values are
        all in a consistent state.

        Args:
            params: a [vimms.Common.ScanParameters][] object to check.

        Returns:

        """

        collision_energy = params.get(ScanParameters.COLLISION_ENERGY)
        orbitrap_resolution = params.get(ScanParameters.ORBITRAP_RESOLUTION)
        activation_type = params.get(ScanParameters.ACTIVATION_TYPE)
        mass_analyser = params.get(ScanParameters.MASS_ANALYSER)
        isolation_mode = params.get(ScanParameters.ISOLATION_MODE)
        agc_target = params.get(ScanParameters.AGC_TARGET)
        max_it = params.get(ScanParameters.MAX_IT)
        # source_cid_energy = params.get(ScanParameters.SOURCE_CID_ENERGY)
        polarity = params.get(ScanParameters.POLARITY)
        first_mass = params.get(ScanParameters.FIRST_MASS)
        last_mass = params.get(ScanParameters.LAST_MASS)

        assert collision_energy is not None
        assert orbitrap_resolution is not None
        assert activation_type is not None
        assert mass_analyser is not None
        assert isolation_mode is not None
        assert agc_target is not None
        assert max_it is not None
        assert polarity is not None
        assert first_mass is not None
        assert last_mass is not None

    def after_injection_cleanup(self):
        """
        Clean-up method at the end of each injection.

        Returns: None

        """
        pass

class DEWFilter(ScoreFilter):
    """
    A class that implements dynamic exclusion filter
    """

    def __init__(self, rt_tol):
        """
        Initialises a dynamic exclusion filter based on time only
        Args:
            rt_tol: the RT tolerance (in seconds)
        """
        self.rt_tol = rt_tol

    def filter(self, current_rt, rois):
        """
        Check whether intensity values are above or below the threshold
        Args:
            current_rt: the current RT value
            rois: a list of [vimms.Roi.Roi][] objects.

        Returns: an array of indicators for the filter
        """

        last_frag_rts = [roi.last_frag_rt for roi in rois]

        # Handles None values by converting to NaN for which all
        # comparisons return 0
        return np.logical_not(current_rt - np.array(last_frag_rts, dtype=np.double) <= self.rt_tol)

class DiaController(Controller):
    """
    A class for doing tree and nested DIA methods.
    Also has a SWATH type controller, but reccommend to use SWATH class
    above. Method uses windows methods from DIA.py to create the pattern
    of windows needed to run the controllers.
    Note: the following method used multiple simultaneous isolation windows
    """

    def __init__(
        self,
        min_mz,
        max_mz,  # TODO: add scan overlap to DiaWindows
        window_type,
        kaufmann_design,
        num_windows,
        scan_overlap=0,
        extra_bins=0,
        dia_design="kaufmann",
        advanced_params=None,
    ):
        super().__init__(advanced_params=advanced_params)
        self.dia_design = dia_design
        self.window_type = window_type
        self.kaufmann_design = kaufmann_design
        self.extra_bins = extra_bins
        self.num_windows = num_windows
        self.scan_overlap = scan_overlap
        self.min_mz = min_mz  # scan from this mz
        self.max_mz = max_mz  # scan to this mz

        self.scan_number = self.initial_scan_id

    def update_state_after_scan(self, last_scan):
        pass

    def _process_scan(self, scan):
        # if there's a previous ms1 scan to process
        new_tasks = []
        if self.scan_to_process is not None:

            mz_tol = 10  # not used
            rt_tol = 15  # these are not used

            precursor_scan_id = self.scan_to_process.scan_id

            mzs = self.scan_to_process.mzs
            if len(mzs) > 0:  # check that ms1 scan is not empty
                default_range = [(self.min_mz, self.max_mz)]
                locations = DiaWindows(
                    mzs,
                    default_range,
                    self.dia_design,
                    self.window_type,
                    self.kaufmann_design,
                    self.extra_bins,
                    self.num_windows,
                ).locations
                for loc in locations:
                    mz = []
                    isolation_width = []
                    intensity = []
                    for sub_loc in loc[0]:
                        mz.append(sum(sub_loc) / 2)
                        isolation_width.append(sub_loc[1] - sub_loc[0])
                        intensity.append(0)
                    dda_scan_params = self.get_ms2_scan_params(
                        mz, intensity, precursor_scan_id, isolation_width, mz_tol, rt_tol
                    )

                    # push this dda scan to the mass spec queue
                    new_tasks.append(dda_scan_params)
            else:
                locations = []

            # make the MS1 scan
            task = self.get_ms1_scan_params()
            new_tasks.append(task)

            self.scan_number += len(locations) + 1
            self.next_processed_scan_id = self.scan_number
        return new_tasks

class DiaWindows:
    """
    Class for creating windows for basic, tree and nested DIA methods. Method is used in
    DiaController in Controller/dia. Basic methods are approximately equal to a SWATH method
    with no overlapping windows
    """

    # flake8: noqa: C901
    def __init__(
        self,
        ms1_mzs,
        ms1_range,
        dia_design,
        window_type,
        kaufmann_design,
        extra_bins,
        num_windows=None,
        range_slack=0.01,
    ):
        ms1_range_difference = ms1_range[0][1] - ms1_range[0][0]
        # set the number of windows for kaufmann method
        if dia_design == "kaufmann":
            num_windows = 64
        # dont allow extra bins for basic method
        if dia_design == "basic" and extra_bins > 0:
            raise ValueError("Cannot have extra bins with 'basic' dia design.")
        # find bin walls and extra bin walls
        if window_type == "even":
            internal_bin_walls = [ms1_range[0][0]]
            for window_index in range(0, num_windows):
                internal_bin_walls.append(
                    ms1_range[0][0] + ((window_index + 1) / num_windows) * ms1_range_difference
                )
            internal_bin_walls[0] = internal_bin_walls[0] - range_slack * ms1_range_difference
            internal_bin_walls[-1] = internal_bin_walls[-1] + range_slack * ms1_range_difference
            internal_bin_walls_extra = None
            if extra_bins > 0:
                internal_bin_walls_extra = [ms1_range[0][0]]
                for window_index in range(0, num_windows * (2**extra_bins)):
                    internal_bin_walls_extra.append(
                        ms1_range[0][0]
                        + ((window_index + 1) / (num_windows * (2**extra_bins)))
                        * ms1_range_difference
                    )
                internal_bin_walls_extra[0] = (
                    internal_bin_walls_extra[0] - range_slack * ms1_range_difference
                )
                internal_bin_walls_extra[-1] = (
                    internal_bin_walls_extra[-1] + range_slack * ms1_range_difference
                )
        elif window_type == "percentile":
            internal_bin_walls = np.percentile(
                ms1_mzs, np.arange(0, 100 + 100 / num_windows, 100 / num_windows)
            ).tolist()
            internal_bin_walls[0] = internal_bin_walls[0] - range_slack * ms1_range_difference
            internal_bin_walls[-1] = internal_bin_walls[-1] + range_slack * ms1_range_difference
            internal_bin_walls_extra = None
            if extra_bins > 0:
                internal_bin_walls_extra = np.percentile(
                    ms1_mzs,
                    np.arange(
                        0,
                        100 + 100 / (num_windows * (2**extra_bins)),
                        100 / (num_windows * (2**extra_bins)),
                    ),
                ).tolist()
                internal_bin_walls_extra[0] = (
                    internal_bin_walls_extra[0] - range_slack * ms1_range_difference
                )
                internal_bin_walls_extra[-1] = (
                    internal_bin_walls_extra[-1] + range_slack * ms1_range_difference
                )
        else:
            raise ValueError("Incorrect window_type selected. Must be 'even' or 'percentile'.")
            # convert bin walls and extra bin walls into locations to scan
        if dia_design == "basic":
            self.locations = []
            for window_index in range(0, num_windows):
                self.locations.append(
                    [[(internal_bin_walls[window_index], internal_bin_walls[window_index + 1])]]
                )
        elif dia_design == "kaufmann":
            self.locations = KaufmannWindows(
                internal_bin_walls, internal_bin_walls_extra, kaufmann_design, extra_bins
            ).locations
        else:
            raise ValueError("Incorrect dia_design selected. Must be 'basic' or 'kaufmann'.")

class DsDAController(WrapperController):
    """
    A controller which allows running the DsDA  (Dataset-Dependent Acquisition)
    method.

    See the original publication for a description of DsDA:

    Broeckling, Hoyes, et al. "Comprehensive Tandem-Mass-Spectrometry coverage
    of complex samples enabled by Data-Set-Dependent acquisition."
    Analytical Chemistry. 90, 8020–8027 (2018).
    """

    def __init__(self, dsda_state, mzml_name, advanced_params=None, task_filter=None):
        """
        Initialise a new DsDAController instance.

        Args:
            dsda_state: An instance of [vimms.DsDA.DsDAState][], wrapping a
            live R process running DsDA.
            mzml_name: The name of the .mzML file to write for this injection.
            advanced_params: a [vimms.Controller.base.AdvancedParams][] object
                that contains advanced parameters to control the mass spec.
                See [vimms.Controller.base.AdvancedParams][] for defaults.
            task_filter: Object that examines the task list and adds or deletes
                tasks to ensure schedule remains in sync with the actual RT.
        """
        self.dsda_state = dsda_state
        self.mzml_name = mzml_name
        self.task_filter = task_filter

        if dsda_state.file_num == 0:
            self.controller = self.dsda_state.get_base_controller()
        else:
            schedule_params, rts = self.dsda_state.get_scan_params()
            self.controller = FixedScansController(
                schedule=schedule_params,
                advanced_params=advanced_params,
                expected_rts=rts,
                task_filter=task_filter,
            )

        print(self.controller)
        super().__init__()

    def after_injection_cleanup(self):
        self.dsda_state.register_mzml(self.mzml_name)

class FixedScansController(Controller):
    """
    A controller which takes a schedule of scans, and converts them into
    tasks in a queue.

    The base class for pre-scheduled controllers like
    [vimms.Controller.misc.DsDAController][] and
    [vimms.Controller.misc.MatchingController][].
    """

    def __init__(self, schedule=None, advanced_params=None, expected_rts=None, task_filter=None):
        """
        Create a FixedScansController.

        Args:
            schedule: List of [vimms.Common.ScanParameter][] objects.
            advanced_params: Instance of [vimms.Controller.base.AdvancedParams][].
            expected_rts: List of expected RTs for tasks with indices corresponding
                to schedule. Only needed if e.g. resynchronising RTs using
                task_filter.
            task_filter: Object that dynamically returns tasks on request. For
                example [vimms.Controller.misc.TaskFilter][] can be used to
                resynchronise unexpected RTs with the expected ones in the schedule.
        """
        super().__init__(advanced_params=advanced_params)
        self.tasks = None
        self.initial_task = None
        self.task_idx = 1  # First scan of schedule is always run
        self.expected_rts = expected_rts
        self.task_filter = task_filter

        if schedule is not None and len(schedule) > 0:
            # if schedule is provided, set it
            self.set_tasks(schedule)
            self.scan_id = self.initial_task.get(ScanParameters.SCAN_ID) + 1
            self.precursor_id = (
                self.scan_id - 1 if self.initial_task.get(ScanParameters.MS_LEVEL) == 1 else None
            )

    def get_initial_tasks(self):
        """
        Returns all initial tasks for the mass spec queue.

        Returns: List of tasks.
        """
        # the remaining scan parameters in the schedule must have been set
        assert self.tasks is not None
        if self.task_filter is None:
            return self.tasks[1:]
        else:
            return []

    def get_initial_scan_params(self):
        """
        Returns the initial scan parameter object to send when
        acquisition starts

        Returns: The initial task.
        """
        # the first scan parameters in the schedule must have been set
        assert self.initial_task is not None
        return self.initial_task

    def set_tasks(self, schedule):
        """
        Set a new schedule for this controller.

        Args:
            schedule: A list of [vimms.Common.ScanParameter][].
        """
        assert isinstance(schedule, list)
        self.initial_task = schedule[0]  # used for sending the first scan
        self.tasks = schedule  # used for sending all the other scans

    def handle_scan(self, scan, current_size, pending_size):
        # simply record every scan that we've received, but return no new tasks
        logger.debug("Time %f Received %s" % (scan.rt, scan))
        self.scans[scan.ms_level].append(scan)
        return self._process_scan(scan)

    def update_state_after_scan(self, last_scan):
        pass

    def _process_scan(self, scan):
        if self.task_filter is None:
            return []
        else:
            self.task_idx, new_task = self.task_filter.get_task(
                scan, self.scan_id, self.precursor_id, self.task_idx, self.expected_rts, self.tasks
            )
            if new_task.get(ScanParameters.MS_LEVEL) == 1:
                self.precursor_id = self.scan_id
            self.scan_id += 1
            return [new_task]

class FlexibleNonOverlapController(TopNEXtController):
    """
    TODO: this class can probably be removed.
    """

    def __init__(
        self,
        ionisation_mode,
        isolation_width,
        N,
        mz_tol,
        rt_tol,
        min_ms1_intensity,
        roi_params,
        grid,
        smartroi_params=None,
        min_roi_length_for_fragmentation=1,
        ms1_shift=0,
        advanced_params=None,
        register_all_roi=False,
        scoring_params=GRID_CONTROLLER_SCORING_PARAMS,
        exclusion_method=ROI_EXCLUSION_DEW,
        exclusion_t_0=None,
    ):  # weighted dew parameters
        super().__init__(
            ionisation_mode,
            isolation_width,
            N,
            mz_tol,
            rt_tol,
            min_ms1_intensity,
            roi_params,
            grid,
            smartroi_params=smartroi_params,
            min_roi_length_for_fragmentation=min_roi_length_for_fragmentation,
            ms1_shift=ms1_shift,
            advanced_params=advanced_params,
            register_all_roi=register_all_roi,
            scoring_params=scoring_params,
            exclusion_method=exclusion_method,
            exclusion_t_0=exclusion_t_0,
        )
        self.scoring_params = scoring_params
        if self.scoring_params["theta3"] != 0 and self.register_all_roi is False:
            print("Warning: register_all_roi should be set to True id theta3 is not 0")

    def _overlap_scores(self):
        scores = [
            self.grid.flexible_non_overlap(
                r, self.roi_builder.current_roi_intensities[i], self.scoring_params
            )
            for i, r in enumerate(self.roi_builder.live_roi)
        ]
        return scores

class GenericBox(Box):
    """Makes no particular assumptions about bounding boxes."""

    def __repr__(self):
        return "Generic{}".format(super().__repr__())

    def contains_point(self, pt):
        return (
            self.pt1.x <= pt.x and self.pt1.y <= pt.y and self.pt2.x >= pt.x and self.pt2.y >= pt.y
        )

    def interval_contains(self, inv):
        if inv.is_vertical():
            return (
                self.pt1.x <= inv.pt1.x
                and self.pt2.x >= inv.pt1.x
                and self.pt1.y >= inv.pt1.y
                and self.pt2.y <= inv.pt2.y
            )
        else:
            return (
                self.pt1.x >= inv.pt1.x
                and self.pt2.x <= inv.pt2.x
                and self.pt1.y <= inv.pt1.y
                and self.pt2.y >= inv.pt1.y
            )

    def overlaps_with_box(self, other_box):
        return (
            self.pt1.x < other_box.pt2.x
            and self.pt2.x > other_box.pt1.x
            and self.pt1.y < other_box.pt2.y
            and self.pt2.y > other_box.pt1.y
        )

    def contains_box(self, other_box):
        return (
            self.pt1.x <= other_box.pt1.x
            and self.pt1.y <= other_box.pt1.y
            and self.pt2.x >= other_box.pt2.x
            and self.pt2.y >= other_box.pt2.y
        )

    def combine_max(self, other_box):
        return GenericBox(
            min(self.pt1.x, other_box.pt1.x),
            max(self.pt2.x, other_box.pt2.x),
            min(self.pt1.y, other_box.pt1.y),
            max(self.pt2.y, other_box.pt2.y),
            parents=(self.parents + other_box.parents),
            intensity=max(self.intensity, other_box.intensity),
        )

    def apply_min_box_ppm(self, xwidth=None, ywidth=None, round_digits=8):
        x1, y1 = self.pt1
        x2, y2 = self.pt2

        if xwidth is not None:
            mid = (x1 + x2) / 2
            dist = (xwidth / 1e6) * mid
            if dist > x2 - x1:
                x1 = mid - dist / 2
                x2 = mid + dist / 2

        if ywidth is not None:
            mid = (y1 + y2) / 2
            dist = (ywidth / 1e6) * mid
            if dist > y2 - y1:
                y1 = mid - dist / 2
                y2 = mid + dist / 2

        new_box = self.copy()
        new_box.pt1.x, new_box.pt1.y = x1, y1
        new_box.pt2.x, new_box.pt2.y = x2, y2
        new_box.round(ndigits=round_digits)
        return new_box

    def overlap_raw(self, other_box):
        if not self.overlaps_with_box(other_box):
            return 0.0
        return (min(self.pt2.x, other_box.pt2.x) - max(self.pt1.x, other_box.pt1.x)) * (
            min(self.pt2.y, other_box.pt2.y) - max(self.pt1.y, other_box.pt1.y)
        )

    def overlap_2(self, other_box):
        if not self.overlaps_with_box(other_box):
            return 0.0
        b = type(self)(
            max(self.pt1.x, other_box.pt1.x),
            min(self.pt2.x, other_box.pt2.x),
            max(self.pt1.y, other_box.pt1.y),
            min(self.pt2.y, other_box.pt2.y),
        )
        return b.area() / (self.area() + other_box.area() - b.area())

    def overlap_3(self, other_box):
        if not self.overlaps_with_box(other_box):
            return 0.0
        b = type(self)(
            max(self.pt1.x, other_box.pt1.x),
            min(self.pt2.x, other_box.pt2.x),
            max(self.pt1.y, other_box.pt1.y),
            min(self.pt2.y, other_box.pt2.y),
        )
        return b.area() / self.area()

    def non_overlap_split(self, other_box):
        """Finds 1 to 4 boxes describing the polygon of area of this box
        not overlapped by other_box. If one box is found, crops this box to
        dimensions of that box, and returns None. Otherwise, returns list of
        2 to 4 boxes. Number of boxes found is equal to number of edges
        overlapping area does NOT share with this box."""
        if not self.overlaps_with_box(other_box):
            return None
        x1, x2, y1, y2 = self.pt1.x, self.pt2.x, self.pt1.y, self.pt2.y
        split_boxes = []
        if other_box.pt1.x > self.pt1.x:
            x1 = other_box.pt1.x
            split_boxes.append(
                type(self)(self.pt1.x, x1, y1, y2, parents=self.parents, intensity=self.intensity)
            )
        if other_box.pt2.x < self.pt2.x:
            x2 = other_box.pt2.x
            split_boxes.append(
                type(self)(x2, self.pt2.x, y1, y2, parents=self.parents, intensity=self.intensity)
            )
        if other_box.pt1.y > self.pt1.y:
            y1 = other_box.pt1.y
            split_boxes.append(
                type(self)(x1, x2, self.pt1.y, y1, parents=self.parents, intensity=self.intensity)
            )
        if other_box.pt2.y < self.pt2.y:
            y2 = other_box.pt2.y
            split_boxes.append(
                type(self)(x1, x2, y2, self.pt2.y, parents=self.parents, intensity=self.intensity)
            )
        return split_boxes

    def split_all(self, other_box):
        if not self.overlaps_with_box(other_box):
            return None, None, None
        both_parents = self.parents + other_box.parents
        both_box = type(self)(
            max(self.pt1.x, other_box.pt1.x),
            min(self.pt2.x, other_box.pt2.x),
            max(self.pt1.y, other_box.pt1.y),
            min(self.pt2.y, other_box.pt2.y),
            parents=both_parents,
            intensity=max(self.intensity, other_box.intensity),
        )
        b1_boxes = self.non_overlap_split(other_box)
        b2_boxes = other_box.non_overlap_split(self)
        return b1_boxes, b2_boxes, both_box