# Copyright (c) 2020 LightOn, All Rights Reserved.
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.

"""
This module contains the OPU class
"""
import time
from math import sqrt
import pkg_resources

from lightonml.encoding.base import NoEncoding, NoDecoding
import warnings
from typing import Optional, Union, Tuple, TYPE_CHECKING

import numpy as np
from contextlib import ExitStack
import attr
import inspect
import lightonml
from lightonml.internal.config import get_host_option, opu_version
from lightonml.internal import config, output_roi, utils, types
from lightonml.internal.user_input import OpuUserInput, InputTraits
from lightonml.internal.simulated_device import SimulatedOpuDevice
from lightonml.context import ContextArray
from lightonml.internal.settings import OpuSettings, TransformSettings
from lightonml.internal.runner import TransformRunner, FitTransformRunner
from lightonml.internal.types import InputRoiStrategy, IntOrTuple, TransformOutput, AcqState
from lightonml.types import OutputRescaling
# Import lightonopu only for typechecking, as it's an optional module and may not be present
if TYPE_CHECKING:
    from lightonopu.internal.device import OpuDevice


# noinspection PyPep8Naming
class OPU:
    """Interface to the OPU.

    .. math:: \\mathbf{y} = \\lvert \\mathbf{R} \\mathbf{x} \\rvert^2 \\mbox{ (non-linear transform, the default)}

    .. math:: \\mathbf{y} = \\mathbf{R}\\mathbf{x} \\mbox{ (linear transform)}

    Main methods are `transform`, `linear_transform`, `fit1d` and `fit2d`,
    and accept NumPy arrays or PyTorch tensors.

    The non-linear transform (`transform`) is a native operation for the OPU, and performs at a higher
    speed than `linear_transform`.

    Acquiring/releasing hardware device resources is done by open/close and a
    context-manager interface.

    Unless `open_at_init=False`, these resources are acquired automatically at init.
    If another process or kernel has not released the resources, an error will be
    raised, call `close()` or shutdown the kernel on the OPU object to release it.

    Parameters
    ----------
    n_components : int,
        dimensionality of the target projection space.
    opu_device : OpuDevice or SimulatedOpuDevice, optional
        optical processing unit instance linked to a physical or simulated device.
        If not provided, a device is properly instantiated.
        If opu_device is of type SimulatedOpuDevice, the random matrix is generated
        at __init__, using max_n_features and n_components
    max_n_features: int, optional
        maximum number of binary features that the OPU will transform
        used only if opu_device is a SimulatedOpuDevice,
        in order to initiate the random matrix
    config_file : str, optional
        path to the configuration file (for dev purpose)
    config_override: dict, optional
        for override of the config_file (for dev purpose)
    verbose_level: int, optional
        deprecated, use lightonml.set_verbose_level() instead
        .. seealso:: `lightonml.set_verbose_level`
    input_roi_strategy: types.InputRoiStrategy, optional
        describes how to display the features on the input device
        .. seealso:: `lightonml.internal.types.InputRoiStrategy`
    open_at_init: bool, optional
        forces the setting of acquiring hardware resource at init. If
        not provided, follow system's setting (usually True)
    disable_pbar: bool, optional
        disable display of the progress bar when verbose_level is set to 1
    simulated: bool, optional
        performs the random projection using CPU, in case no OPU is available on your machine
        the random matrix is then generated at __init__, using max_n_features and n_components
    rescale: types.OutputRescaling, optional,
        output rescaling method for `linear_transform`.
        Ignored by `transform`.
        .. seealso:: `lightonml.types.OutputRescaling`

    Attributes
    ----------
    n_components: int
        dimensionality of the target projection space.
    rescale: types.OutputRescaling,
        output rescaling method for `linear_transform`.
        Ignored by `transform`.
    max_n_features: int
        maximum number of binary features that the OPU will transform
        writeable only if opu_device is a SimulatedOpuDevice,
        in order to initiate or resize the random matrix
    device: OpuDevice or SimulatedOpuDevice
        underlying hardware that performs transformation (read-only)
    input_roi_strategy: types.InputRoiStrategy, optional
        describes how to display the features on the input device
    """

    def __init__(self, n_components: int = 200000,
                 opu_device: Optional[Union["OpuDevice", SimulatedOpuDevice]] = None,
                 max_n_features: int = 1000, config_file: str = "",
                 config_override: dict = None, verbose_level: int = -1,
                 input_roi_strategy: types.InputRoiStrategy = types.InputRoiStrategy.full,
                 open_at_init: bool = None, disable_pbar=False, simulated=False,
                 rescale: Union[OutputRescaling, str] = OutputRescaling.variance):

        self.__opu_config = None
        self.__config_file = config_file
        self.__config_override = config_override
        self._max_n_features = max_n_features
        self.disable_pbar = disable_pbar
        self.rescale = rescale

        # Get trace and print functions
        if verbose_level != -1:
            warnings.warn("Verbose level arg will removed in 1.3, "
                          "Use lightonml.set_verbose_level instead",
                          DeprecationWarning)
            lightonml.set_verbose_level(verbose_level)
        else:
            verbose_level = lightonml.get_verbose_level()
        self._debug = lightonml.get_debug_fn()
        self._trace = lightonml.get_trace_fn()
        self._print = lightonml.get_print_fn()
        no_config_msg = "No configuration files for the OPU was found on this machine.\n" \
                        "You may want to run the OPU in a simulated manner, by passing the " \
                        "simulated argument to True at init.\n" \
                        "See https://docs.lighton.ai/notes/get_started.html#Simulating-an-OPU " \
                        "for more details.\n" \
                        "See also https://lighton.ai/products for getting access to our technology."

        if simulated and opu_device is not None:
            raise ValueError("simulated and opu_device arguments are conflicting")

        # Device init, or take the one passed as input
        if opu_device:
            if type(opu_device).__name__ not in ["SimulatedOpuDevice", "OpuDevice"]:
                raise TypeError("opu_device must be of type SimulatedOpuDevice or OpuDevice")
            self.device = opu_device
        elif simulated:
            self.device = SimulatedOpuDevice()
        else:
            # Instantiate device directly
            from lightonopu.internal.device import OpuDevice
            if not self.__config_file and not config.host_has_opu_config():
                # Looks like there's no OPU on this host as we didn't find configuration files
                raise RuntimeError(no_config_msg)
            opu_type = self.config["type"]
            frametime_us = self.config["input"]["frametime_us"]
            exposure_us = self.config["output"]["exposure_us"]
            seq_nb_prelim = self.config.get("sequence_nb_prelim", 0)
            name = self.config["name"]
            self.device = OpuDevice(opu_type, frametime_us, exposure_us, seq_nb_prelim,
                                    None, verbose_level, name)

        self._base_frametime_us = self.device.frametime_us
        self._base_exposure_us = self.device.exposure_us

        if self._s.simulated:
            # build the random matrix if not done already
            self._resize_rnd_matrix(max_n_features, n_components)
        else:
            # Make sure lightonopu is at 1.4.1 or later, needed for linear_reconstruction
            pkg_resources.require("lightonopu>=1.4.1")
            # initialize linear_reconstruction library
            from lightonopu import linear_reconstruction
            linear_reconstruction.init(np.prod(self.device.input_shape))

            self._output_roi = output_roi.OutputRoi(self.device.output_shape_max,
                                                    self.device.output_roi_strategy,
                                                    self._s.allowed_roi, self._s.min_n_components)
        # This also sets the output ROI
        self.n_components = n_components
        self.input_roi_strategy = input_roi_strategy
        # Runner initialized when entering fit
        self._runner = None  # type: Optional[TransformRunner]
        # ExitStack for device acquisition, initialized when entering fit
        self._acq_stack = ExitStack()
        self._trace("OPU initialized")

        # Open at init, unless relevant host.json option is False
        if open_at_init is None:
            open_at_init = get_host_option("lightonml_open_at_init", True)
        if open_at_init:
            self.open()

    def _tr_settings(self, no_input=False, **override) -> TransformSettings:
        """Returns transform settings for feeding to TransformRunner"""
        init = TransformSettings(self.input_roi_strategy, self.n_components)
        settings = attr.evolve(init, **override)
        if no_input and self.input_roi_strategy is InputRoiStrategy.auto:
            # If no input_roi, replace auto by full strategy
            settings.input_roi_strategy = InputRoiStrategy.full
            assert settings.input_roi is None
        return settings

    def fit1d(self, X=None, n_features: int = None,
              packed: bool = False, online=False, **override):
        """
        Configure OPU transform for 1d vectors

        The function can be either called with input vector, for fitting OPU
        parameters to it, or just vector dimensions, with ``n_features``.

        When input is bit-packed the packed flag must be set to True.

        When input vectors must be transformed one by one, performance will
        be improved with the online flag set to True.

        Parameters
        ----------
        X: np.ndarray or torch.Tensor
            Fit will be made on this vector to optimize transform parameters
        n_features: int
            Number of features for the input, necessary if X parameter isn't provided
        packed: bool
            Set to true if the input vectors will be already bit-packed
        online: bool, optional
            Set to true if the transforms will be made one vector after the other
            defaults to False
        override: dict, optional
            keyword args for overriding transform settings (advanced parameters)
        """
        return self.__fit(X, n_features, packed, online, False, **override)

    def fit2d(self, X=None, n_features: Tuple[int, int] = None,
              packed: bool = False, online=False, **override):
        """
        Configure OPU transform for 2d vectors

        The function can be either called with input vector, for fitting OPU
        parameters to it, or just vector dimensions, with `n_features`.

        When input is bit-packed the packed flag must be set to True.
        Number of features must be then provided with `n_features`

        When input vectors must be transformed one by one, performance will
        be improved with the online flag set to True.

        Parameters
        ----------
        X: np.ndarray or torch.Tensor
            a 2d input vector, or batch of 2d input_vectors, binary encoded, packed or not
        n_features: tuple(int)
            Number of features for the input, necessary if X parameter isn't provided, or
            if input is bit-packed
        packed: bool, optional
            whether the input data is in bit-packed representation
            if True, each input vector is assumed to be a 1d array, and the "real" number
            of features must be provided as n_features
            defaults to False
        online: bool, optional
            Set to true if the transforms will be made one vector after the other
            defaults to False
        override: dict, optional
            keyword args for overriding transform settings (advanced parameters)
        """
        return self.__fit(X, n_features, packed, online, True, **override)

    def transform(self, X, encoder_cls=NoEncoding, decoder_cls=NoDecoding) -> TransformOutput:
        """
        Performs the nonlinear random projections of one or several input vectors.

        The `fit1d` or `fit2d` method must be called before, for setting vector dimensions
        or online option.
        If you need to transform one vector after each other, add `online=True` in the fit function.

        Parameters
        ----------
        X:  np.ndarray or torch.Tensor
            input vector, or batch of input vectors.
            Each vector must have the same dimensions as the one given in `fit1d` or `fit2d`.
        encoder_cls: encoder.base.BaseTransformer, optional
            class or instance of class that transform the input into binary vectors to be processed by the opu.
        decoder_cls: encoder.base.BaseTransformer, optional
            class or instance of class that transforms the output of the opu back into the appropriate format.

        Returns
        -------
        Y: np.ndarray or torch.Tensor
             complete array of nonlinear random projections of X,
             of size self.n_components
             If input is an ndarray, type is actually ContextArray,
             with a context attribute to add metadata
        """
        assert self._runner, "Call fit1d or fit2d before transform"
        assert self.device.active, "OPU device isn't active, use opu.open() or \"with opu:\""

        if inspect.isclass(encoder_cls):
            encoder = encoder_cls()
        else:
            encoder = encoder_cls

        X_enc = encoder.transform(X)

        user_input = OpuUserInput.from_traits(X_enc, self._runner.traits)
        self._debug(str(user_input))

        if user_input.is_batch and not self._s.simulated:
            # With batch input start acquisition first
            assert self.device.acq_state.value != AcqState.online.value, \
                "Can't transform a batch of vectors when acquisition is" \
                " in online mode, only single vectors"
            with self.device.acquiring(n_images=self._s.n_samples_by_pass):
                out = self._runner.transform(user_input)
        else:
            out = self._runner.transform(user_input)
        return self._post_transform(out, user_input, encoder, decoder_cls)

    def linear_transform(self, X, encoder_cls=NoEncoding, decoder_cls=NoDecoding) -> TransformOutput:
        """
        Do a linear transform of X, for Nitro (non-linear) photonic cores.

        Parameters
        ----------
        X:  np.ndarray or torch.Tensor
            input vector, or batch of input vectors.
            Each vector must have the same dimensions as the one given in `fit1d` or `fit2d`.
        encoder_cls: encoding.base.BaseTransformer, optional
            class or instance of class that transform the input into binary vectors to be processed by the opu.
        decoder_cls: encoding.base.BaseTransformer, optional
            class or instance of class that transforms the output of the opu back into the appropriate format.

        Returns
        -------
        Y: np.ndarray or torch.Tensor
             complete array of nonlinear random projections of X,
             of size self.n_components
             If input is an ndarray, type is actually ContextArray,
             with a context attribute to add metadata
        """
        assert self._runner, "Call fit1d or fit2d before linear_transform"
        traits = self._runner.traits

        if traits.packed:
            # TODO implement for packed
            raise RuntimeError("Linear transform isn't yet implemented for packed input :/")

        if inspect.isclass(encoder_cls):
            encoder = encoder_cls()
        else:
            encoder = encoder_cls

        X_enc = encoder.transform(X)

        user_input = OpuUserInput.from_traits(X_enc, traits)
        _, result_ctx = self._raw_linear_transform(X_enc, traits, user_input)
        # Decoding, add context, and optional convert back to torch if needed
        output = self._post_transform(result_ctx, user_input, encoder, decoder_cls)
        # Rescale the output, intentionally after the decoding step
        if self.rescale is OutputRescaling.variance:
            n_features = user_input.n_features_s
            output = output / (self._s.stdev * sqrt(n_features))
        elif self.rescale is OutputRescaling.norm:
            output = output / (self._s.stdev * sqrt(self.n_components))
        return output

    def transform1d(self, *args, **kwargs):
        raise RuntimeError("transform1d is deprecated, you must now use fit1d and transform")

    def transform2d(self, *args, **kwargs):
        raise RuntimeError("transform2d is deprecated, you must now use fit2d and transform")

    def fit_transform1d(self, X, packed: bool = False,
                        **override) -> ContextArray:
        """Performs the nonlinear random projections of 1d input vector(s).

        This function is the one-liner equivalent of `fit1d` and `transform` calls.

        .. warning:: when making several transform calls, prefer calling `fit1d`
            and then `transform`, or you might encounter an inconsistency in the
            transformation matrix.

        The input data can be bit-packed, where ``n_features = 8*X.shape[-1]``
        Otherwise ``n_features = X.shape[-1]``

        If tqdm module is available, it is used for progress display

        Parameters
        ----------
        X: np.ndarray or torch.Tensor
            a 1d input vector, or batch of 1d input_vectors, binary encoded, packed or not
            batch can be 1d or 2d. In all cases ``output.shape[:-1] = X.shape[:-1]``
        packed: bool, optional
            whether the input data is in bit-packed representation
            defaults to False
        override: keyword args for overriding transform settings (advanced parameters)

        Returns
        -------
        Y: np.ndarray or torch.Tensor
             complete array of nonlinear random projections of X,
             of size self.n_components
             If input is an ndarray, type is actually ContextArray,
             with a context attribute to add metadata

        """
        self.fit1d(X, None, packed, False, **override)
        return self.transform(X)

    def fit_transform2d(self, X, packed: bool = False, n_2d_features=None,
                        **override) -> ContextArray:
        """Performs the nonlinear random projections of 2d input vector(s).

        This function is the one-liner equivalent of `fit2d` and `transform` calls.

        .. warning:: when making several transform calls, prefer calling `fit2d`
            and then `transform`, or you might encounter an inconsistency in the
            transformation matrix.

        If tqdm module is available, it is used for progress display

        Parameters
        ----------
        X: np.ndarray or torch.Tensor
            a 2d input vector, or batch of 2d input_vectors, binary encoded, packed or not
        packed: bool, optional
            whether the input data is in bit-packed representation
            if True, each input vector is assumed to be a 1d array, and the "real" number
            of features must be provided as n_2d_features
            defaults to False
        n_2d_features: list, tuple or np.ndarray of length 2
            If the input is bit-packed, specifies the shape of each input vector.
            Not needed if the input isn't bit-packed.
        override: keyword args for overriding transform settings (advanced parameters)

        Returns
        -------
        Y: np.ndarray or torch.Tensor
             complete array of nonlinear random projections of X,
             of size self.n_components
             If input is an ndarray, type is actually ContextArray,
             with a context attribute to add metadata
        """
        self.fit2d(X, n_2d_features, packed, False, **override)
        return self.transform(X)

    def __fit(self, X, n_features: IntOrTuple,
              packed: bool, online: bool, is_2d_features: bool,
              **override):
        """Internal working of the fitXd calls

        Instantiates a TransformRunner, and start online acq if needs be.
        """
        if X is not None:
            # Input is provided, do the fit with user input
            user_input = OpuUserInput.from_input(X, packed, is_2d_features, n_features)
            tr_settings = self._tr_settings(no_input=False, **override)
            self._runner = FitTransformRunner(self._s, tr_settings, user_input,
                                              device=self.device,
                                              disable_pbar=self.disable_pbar)
        else:
            # Only dimensions are provided, no fitting happens on input
            assert n_features, "either input vector or n_features must be specified"
            # tr_settings has no input_roi, since it uses X to compute it
            tr_settings = self._tr_settings(no_input=True, **override)
            traits = InputTraits(n_features, packed)
            self._runner = TransformRunner(self._s, tr_settings, traits,
                                           device=self.device,
                                           disable_pbar=self.disable_pbar)

        self._acq_stack.close()
        if online:
            if self._s.no_single_transform:
                raise RuntimeError("Online transform isn't available with this OPU")
            # Start acquisition only if online. Batch transform start their own.
            self._acq_stack.enter_context(self.device.acquiring(online=True))

    @staticmethod
    def _post_transform(output, user_input, encoder, decoder_cls):
        """Final steps after transform
           1. reshape
           2. decode the output
           3. convert to tensor if user input was tensor
           """
        output = user_input.reshape_output(output)
        # If encoder has get_params method, it's for transmitting it to decoder init
        if inspect.isclass(decoder_cls):
            if hasattr(encoder, "get_params"):
                decoder = decoder_cls(**encoder.get_params())
            else:
                decoder = decoder_cls()
        else:
            decoder = decoder_cls

        output = decoder.transform(output)

        if user_input.is_tensor:
            # noinspection PyPackageRequirements,PyUnresolvedReferences
            import torch
            return torch.from_numpy(output)
        else:
            return output

    def _raw_linear_transform(self, X, traits=None, user_input=None):
        """
        Do linear_transform of X, and return both raw OPU output and decoded output in a tuple
        """

        if traits is None:
            assert self._runner, "Call fit1d or fit2d before linear_transform"
            traits = self._runner.traits
        if user_input is None:
            user_input = OpuUserInput.from_traits(X, traits)

        if self._s.simulated:
            prepared_X = X
        else:
            assert self.device.acq_state.value != AcqState.online.value, \
                "Can't do linear transform when acquisition is" \
                " in online mode, only single vectors"
            assert self._runner.t.input_roi_strategy == InputRoiStrategy.full, \
                "ROI strategy must be full for linear_transform to be correct.\n" \
                "Set input_roi_strategy attribute to InputRoiStrategy.full."

            # X2 is now numpy 2D, whatever the initial shape and the type (torch or numpy)
            X2 = user_input.reshape_input(raveled_features=True, leave_single_dim=True)

            try:
                import lightonopu.linear_reconstruction as reconstruction
            except ImportError:
                raise RuntimeError("Need a lightonopu version with linear_reconstruction module")

            start = time.time()
            prepared_X = reconstruction.encode_batch(X2)
            self._trace(f"Encoding time {time.time() - start} s")
            # Restore the dimension after batch encoding to something suitable for formatting
            prepared_X = user_input.unravel_features(prepared_X)
        # Run the OPU transform
        prepared_input = OpuUserInput.from_traits(prepared_X, traits)
        start = time.time()
        with self.device.acquiring(n_images=self._s.n_samples_by_pass):
            rp_opu = self._runner.transform(prepared_input, linear=True)
        self._trace(f"Transform time {time.time() - start} s")

        if self._s.simulated:
            result_ctx = rp_opu
        else:
            # Decoding forgets about the context, re-add it to result afterwards
            start = time.time()
            result = reconstruction.decode_batch(rp_opu)
            self._trace(f"Decoding time {time.time() - start} s")

            result_ctx = ContextArray(result, rp_opu.context)
        return rp_opu, result_ctx

    def __enter__(self):
        """Context manager interface that acquires hardware resources
        used by the OPU device."""
        self.__active_before_enter = self.device.active
        self.open()
        return self

    def __exit__(self, *args):
        # Don't close if OPU was already active
        if not self.__active_before_enter:
            self.close()

    def open(self):
        """Acquires hardware resources used by the OPU device

        .. seealso:: `close()` or use the context manager interface for
            closing at the end af an indent block
        """
        if self.device.active:
            return
        self.device.open()
        # initial reservation for giving batch transforms a buffer ready to use
        self.device.reserve(self._s.n_samples_by_pass)
        if self._s.detect_trigger:
            # Detect trigger issue, and take action if needed
            issue = utils.detect_trigger_issue(self.device)
            if issue:
                # noinspection PyProtectedMember,PyUnresolvedReferences
                self.device._OpuDevice__opu.nb_prelim = 1
                self._debug("trigger issue detected, workaround applied")
            else:
                self._debug("trigger issue not detected")

        self._debug("OPU opened")

    def close(self):
        """Releases hardware resources used by the OPU device"""
        self._acq_stack.close()
        self.device.close()
        self._debug("OPU closed")

    @property
    def config(self):
        """Returns the internal configuration object"""
        # Load it when asked first time
        if not self.__opu_config:
            self.__opu_config = config.load_config(self.__config_file, self._trace)
            if self.__config_override is not None:
                utils.recurse_update(self.__opu_config, self.__config_override)

        return self.__opu_config

    @property
    def rescale(self):
        return self._rescale

    @rescale.setter
    def rescale(self, value):
        # If str it's the enum value
        if isinstance(value, str):
            self._rescale = OutputRescaling[value.lower()]
        else:
            assert isinstance(value, OutputRescaling)
            self._rescale = value

    @property
    def max_n_components(self):
        return self._output_roi.max_components

    @property
    def n_components(self) -> int:
        return self._n_components

    @n_components.setter
    def n_components(self, value: int):
        if self._s.simulated:
            self._resize_rnd_matrix(self.max_n_features, value)
        else:
            self.device.output_roi = self._output_roi.compute_roi(value)
        # We used to call device.reserve here, but moved to device.acquiring()
        self._n_components = value

    @property
    def max_n_features(self) -> int:
        return self._s.max_n_features

    @max_n_features.setter
    def max_n_features(self, value: int):
        if not self._s.simulated:
            raise AttributeError("max_n_feature can't be set if device is real")
        self._resize_rnd_matrix(value, self._n_components)
        self._max_n_features = value

    @property
    def _s(self) -> OpuSettings:
        """Returns immutable settings associated with the OPU
        Settings are immutable (attrs frozen), so generate it at
        each call. Performance impact is negligible"""
        # Get default value
        pass_default = attr.fields(OpuSettings).n_samples_by_pass.default

        # Common settings to both simulated and base
        kwargs = {"input_shape": self.device.input_shape,
                  "output_max_shape": self.device.output_shape_max,
                  "frametime_us": self._base_frametime_us,
                  "exposure_us": self._base_exposure_us}

        if isinstance(self.device, SimulatedOpuDevice):
            # Notice we never query self.config here, in order not to
            # need a configuration file for simulated device
            return OpuSettings(max_n_features=self._max_n_features,
                               n_samples_by_pass=pass_default,
                               simulated=True, **kwargs
                               )

        return OpuSettings(
            max_n_features=int(np.prod(self.device.input_shape)),
            # Will use defaults of OpuSettings if not found
            n_samples_by_pass=self.config.get("n_samples_by_pass", pass_default),
            min_batch_size=self.config["input"].get("minimum_batch_size", 0),
            allowed_roi=self.config["output"].get("allowed_roi"),
            # min_n_components is linked to the minimum output size
            min_n_components=self.config["output"].get("minimum_output_size", 0),
            ones_range=self.config["ones_range"],
            n_tries=self.config.get("n_transform_tries", 5),
            detect_trigger=self.config.get("detect_trigger_issue", False),
            no_single_transform=self.config.get("no_single_transform", False),
            stdev=self.config["output"].get("stdev", 1.),
            **kwargs)

    def _resize_rnd_matrix(self, n_features: int, n_components: int):
        """Resize device's random matrix"""
        assert isinstance(self.device, SimulatedOpuDevice)
        rnd_mat = self.device.random_matrix
        if rnd_mat is None or rnd_mat.shape != (n_features, n_components):
            self._print("OPU: computing the random matrix... ", end='', flush=True)
            self.device.build_random_matrix(n_features, n_components)
            self._print("OK")

    def version(self, devices=False):
        """Returns a multi-line string containing name and versions of the OPU"""
        version = []

        # Build OPU name
        if not self._s.simulated:
            version.append(opu_version(self.__opu_config))

        # module version
        version.append(f"lightonml version {lightonml.__version__}")
        try:
            # noinspection PyUnresolvedReferences
            import lightonopu
            version.append(f"lightonopu version {lightonopu.__version__}")
        except ImportError:
            pass
        if devices:
            version.append(self.device.versions())
        return '\n'.join(version)

    def __getstate__(self):
        state = self.__dict__.copy()
        # Remove logging functions, they can't be pickled
        state.pop("_debug")
        state.pop("_trace")
        state.pop("_print")
        # acq stack can't be pickled, will be restored
        state.pop("_acq_stack")
        # If acquisition is ongoing, close it
        if not self._s.simulated:
            state["__online_acq"] = self.device.acq_state.value == AcqState.online.value
        self._acq_stack.close()
        # Device itself is closed on pickling
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        # Restore logging functions removed at getstate
        self._debug = lightonml.get_debug_fn()
        self._trace = lightonml.get_trace_fn()
        self._print = lightonml.get_print_fn()
        self._acq_stack = ExitStack()
        # Restore online acquisition if it was the case
        if state.get("__online_acq", False):
            self._acq_stack.enter_context(self.device.acquiring(online=True))

[((132, 22, 132, 46), 'lightonml.get_debug_fn', 'lightonml.get_debug_fn', ({}, {}), '()', False, 'import lightonml\n'), ((133, 22, 133, 46), 'lightonml.get_trace_fn', 'lightonml.get_trace_fn', ({}, {}), '()', False, 'import lightonml\n'), ((134, 22, 134, 46), 'lightonml.get_print_fn', 'lightonml.get_print_fn', ({}, {}), '()', False, 'import lightonml\n'), ((188, 26, 188, 37), 'contextlib.ExitStack', 'ExitStack', ({}, {}), '()', False, 'from contextlib import ExitStack\n'), ((199, 15, 199, 76), 'lightonml.internal.settings.TransformSettings', 'TransformSettings', ({(199, 33, 199, 56): 'self.input_roi_strategy', (199, 58, 199, 75): 'self.n_components'}, {}), '(self.input_roi_strategy, self.n_components)', False, 'from lightonml.internal.settings import OpuSettings, TransformSettings\n'), ((200, 19, 200, 48), 'attr.evolve', 'attr.evolve', ({(200, 31, 200, 35): 'init'}, {}), '(init, **override)', False, 'import attr\n'), ((299, 11, 299, 39), 'inspect.isclass', 'inspect.isclass', ({(299, 27, 299, 38): 'encoder_cls'}, {}), '(encoder_cls)', False, 'import inspect\n'), ((306, 21, 306, 73), 'lightonml.internal.user_input.OpuUserInput.from_traits', 'OpuUserInput.from_traits', ({(306, 46, 306, 51): 'X_enc', (306, 53, 306, 72): 'self._runner.traits'}, {}), '(X_enc, self._runner.traits)', False, 'from lightonml.internal.user_input import OpuUserInput, InputTraits\n'), ((349, 11, 349, 39), 'inspect.isclass', 'inspect.isclass', ({(349, 27, 349, 38): 'encoder_cls'}, {}), '(encoder_cls)', False, 'import inspect\n'), ((356, 21, 356, 60), 'lightonml.internal.user_input.OpuUserInput.from_traits', 'OpuUserInput.from_traits', ({(356, 46, 356, 51): 'X_enc', (356, 53, 356, 59): 'traits'}, {}), '(X_enc, traits)', False, 'from lightonml.internal.user_input import OpuUserInput, InputTraits\n'), ((488, 11, 488, 39), 'inspect.isclass', 'inspect.isclass', ({(488, 27, 488, 38): 'decoder_cls'}, {}), '(decoder_cls)', False, 'import inspect\n'), ((540, 25, 540, 69), 'lightonml.internal.user_input.OpuUserInput.from_traits', 'OpuUserInput.from_traits', ({(540, 50, 540, 60): 'prepared_X', (540, 62, 540, 68): 'traits'}, {}), '(prepared_X, traits)', False, 'from lightonml.internal.user_input import OpuUserInput, InputTraits\n'), ((541, 16, 541, 27), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((734, 22, 734, 46), 'lightonml.get_debug_fn', 'lightonml.get_debug_fn', ({}, {}), '()', False, 'import lightonml\n'), ((735, 22, 735, 46), 'lightonml.get_trace_fn', 'lightonml.get_trace_fn', ({}, {}), '()', False, 'import lightonml\n'), ((736, 22, 736, 46), 'lightonml.get_print_fn', 'lightonml.get_print_fn', ({}, {}), '()', False, 'import lightonml\n'), ((737, 26, 737, 37), 'contextlib.ExitStack', 'ExitStack', ({}, {}), '()', False, 'from contextlib import ExitStack\n'), ((126, 12, 128, 45), 'warnings.warn', 'warnings.warn', ({(126, 26, 127, 67): '"""Verbose level arg will removed in 1.3, Use lightonml.set_verbose_level instead"""', (128, 26, 128, 44): 'DeprecationWarning'}, {}), "(\n    'Verbose level arg will removed in 1.3, Use lightonml.set_verbose_level instead'\n    , DeprecationWarning)", False, 'import warnings\n'), ((129, 12, 129, 54), 'lightonml.set_verbose_level', 'lightonml.set_verbose_level', ({(129, 40, 129, 53): 'verbose_level'}, {}), '(verbose_level)', False, 'import lightonml\n'), ((131, 28, 131, 57), 'lightonml.get_verbose_level', 'lightonml.get_verbose_level', ({}, {}), '()', False, 'import lightonml\n'), ((174, 12, 174, 54), 'pkg_resources.require', 'pkg_resources.require', ({(174, 34, 174, 53): '"""lightonopu>=1.4.1"""'}, {}), "('lightonopu>=1.4.1')", False, 'import pkg_resources\n'), ((179, 31, 181, 98), 'lightonml.internal.output_roi.OutputRoi', 'output_roi.OutputRoi', ({(179, 52, 179, 80): 'self.device.output_shape_max', (180, 52, 180, 83): 'self.device.output_roi_strategy', (181, 52, 181, 71): 'self._s.allowed_roi', (181, 73, 181, 97): 'self._s.min_n_components'}, {}), '(self.device.output_shape_max, self.device.\n    output_roi_strategy, self._s.allowed_roi, self._s.min_n_components)', False, 'from lightonml.internal import config, output_roi, utils, types\n'), ((193, 27, 193, 74), 'lightonml.internal.config.get_host_option', 'get_host_option', ({(193, 43, 193, 67): '"""lightonml_open_at_init"""', (193, 69, 193, 73): 'True'}, {}), "('lightonml_open_at_init', True)", False, 'from lightonml.internal.config import get_host_option, opu_version\n'), ((457, 25, 457, 87), 'lightonml.internal.user_input.OpuUserInput.from_input', 'OpuUserInput.from_input', ({(457, 49, 457, 50): 'X', (457, 52, 457, 58): 'packed', (457, 60, 457, 74): 'is_2d_features', (457, 76, 457, 86): 'n_features'}, {}), '(X, packed, is_2d_features, n_features)', False, 'from lightonml.internal.user_input import OpuUserInput, InputTraits\n'), ((459, 27, 461, 77), 'lightonml.internal.runner.FitTransformRunner', 'FitTransformRunner', (), '', False, 'from lightonml.internal.runner import TransformRunner, FitTransformRunner\n'), ((467, 21, 467, 52), 'lightonml.internal.user_input.InputTraits', 'InputTraits', ({(467, 33, 467, 43): 'n_features', (467, 45, 467, 51): 'packed'}, {}), '(n_features, packed)', False, 'from lightonml.internal.user_input import OpuUserInput, InputTraits\n'), ((468, 27, 470, 74), 'lightonml.internal.runner.TransformRunner', 'TransformRunner', (), '', False, 'from lightonml.internal.runner import TransformRunner, FitTransformRunner\n'), ((501, 19, 501, 43), 'torch.from_numpy', 'torch.from_numpy', ({(501, 36, 501, 42): 'output'}, {}), '(output)', False, 'import torch\n'), ((514, 25, 514, 60), 'lightonml.internal.user_input.OpuUserInput.from_traits', 'OpuUserInput.from_traits', ({(514, 50, 514, 51): 'X', (514, 53, 514, 59): 'traits'}, {}), '(X, traits)', False, 'from lightonml.internal.user_input import OpuUserInput, InputTraits\n'), ((534, 20, 534, 31), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((535, 25, 535, 56), 'lightonopu.linear_reconstruction.encode_batch', 'reconstruction.encode_batch', ({(535, 53, 535, 55): 'X2'}, {}), '(X2)', True, 'import lightonopu.linear_reconstruction as reconstruction\n'), ((550, 20, 550, 31), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((551, 21, 551, 56), 'lightonopu.linear_reconstruction.decode_batch', 'reconstruction.decode_batch', ({(551, 49, 551, 55): 'rp_opu'}, {}), '(rp_opu)', True, 'import lightonopu.linear_reconstruction as reconstruction\n'), ((554, 25, 554, 61), 'lightonml.context.ContextArray', 'ContextArray', ({(554, 38, 554, 44): 'result', (554, 46, 554, 60): 'rp_opu.context'}, {}), '(result, rp_opu.context)', False, 'from lightonml.context import ContextArray\n'), ((582, 20, 582, 59), 'lightonml.internal.utils.detect_trigger_issue', 'utils.detect_trigger_issue', ({(582, 47, 582, 58): 'self.device'}, {}), '(self.device)', False, 'from lightonml.internal import config, output_roi, utils, types\n'), ((603, 32, 603, 83), 'lightonml.internal.config.load_config', 'config.load_config', ({(603, 51, 603, 69): 'self.__config_file', (603, 71, 603, 82): 'self._trace'}, {}), '(self.__config_file, self._trace)', False, 'from lightonml.internal import config, output_roi, utils, types\n'), ((667, 19, 670, 32), 'lightonml.internal.settings.OpuSettings', 'OpuSettings', (), '', False, 'from lightonml.internal.settings import OpuSettings, TransformSettings\n'), ((151, 26, 151, 46), 'lightonml.internal.simulated_device.SimulatedOpuDevice', 'SimulatedOpuDevice', ({}, {}), '()', False, 'from lightonml.internal.simulated_device import SimulatedOpuDevice\n'), ((163, 26, 164, 62), 'lightonopu.internal.device.OpuDevice', 'OpuDevice', ({(163, 36, 163, 44): 'opu_type', (163, 46, 163, 58): 'frametime_us', (163, 60, 163, 71): 'exposure_us', (163, 73, 163, 86): 'seq_nb_prelim', (164, 36, 164, 40): 'None', (164, 42, 164, 55): 'verbose_level', (164, 57, 164, 61): 'name'}, {}), '(opu_type, frametime_us, exposure_us, seq_nb_prelim, None,\n    verbose_level, name)', False, 'from lightonopu.internal.device import OpuDevice\n'), ((177, 39, 177, 71), 'numpy.prod', 'np.prod', ({(177, 47, 177, 70): 'self.device.input_shape'}, {}), '(self.device.input_shape)', True, 'import numpy as np\n'), ((605, 16, 605, 79), 'lightonml.internal.utils.recurse_update', 'utils.recurse_update', ({(605, 37, 605, 54): 'self.__opu_config', (605, 56, 605, 78): 'self.__config_override'}, {}), '(self.__opu_config, self.__config_override)', False, 'from lightonml.internal import config, output_roi, utils, types\n'), ((656, 23, 656, 47), 'attr.fields', 'attr.fields', ({(656, 35, 656, 46): 'OpuSettings'}, {}), '(OpuSettings)', False, 'import attr\n'), ((702, 27, 702, 57), 'lightonml.internal.config.opu_version', 'opu_version', ({(702, 39, 702, 56): 'self.__opu_config'}, {}), '(self.__opu_config)', False, 'from lightonml.internal.config import get_host_option, opu_version\n'), ((363, 47, 363, 63), 'math.sqrt', 'sqrt', ({(363, 52, 363, 62): 'n_features'}, {}), '(n_features)', False, 'from math import sqrt\n'), ((673, 31, 673, 63), 'numpy.prod', 'np.prod', ({(673, 39, 673, 62): 'self.device.input_shape'}, {}), '(self.device.input_shape)', True, 'import numpy as np\n'), ((155, 46, 155, 74), 'lightonml.internal.config.host_has_opu_config', 'config.host_has_opu_config', ({}, {}), '()', False, 'from lightonml.internal import config, output_roi, utils, types\n'), ((365, 47, 365, 70), 'math.sqrt', 'sqrt', ({(365, 52, 365, 69): 'self.n_components'}, {}), '(self.n_components)', False, 'from math import sqrt\n'), ((544, 38, 544, 49), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((536, 41, 536, 52), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((552, 41, 552, 52), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n')]