text:
# This file is part of the bapsflib package, a Python toolkit for the
# BaPSF group at UCLA.
#
# http://plasma.physics.ucla.edu/
#
# Copyright 2017-2018 Erik T. Everson and contributors
#
# License: Standard 3-clause BSD; see "LICENSES/LICENSE.txt" for full
#   license terms and contributor agreement.
#
#
"""
Module containing the main `~bapsflib._hdf.utils.hdfreaddata.HDFReadData` class.
"""
__all__ = ["HDFReadData"]

import astropy.units as u
import copy
import numpy as np
import os
import time

from typing import Union
from warnings import warn

from bapsflib.plasma import core

from .file import File
from .hdfreadcontrols import HDFReadControls
from .helpers import (
    build_sndr_for_simple_dset,
    condition_controls,
    condition_shotnum,
    do_shotnum_intersection,
)

# noinspection PyInitNewSignature
class HDFReadData(np.ndarray):
    """
    Reads digitizer and control device data from the HDF5 file. Control
    device data is extracted using
    :class:`~.hdfreadcontrols.HDFReadControls` and combined with the
    digitizer data.

    This class constructs and returns a structured numpy array.  The
    data in the array is grouped into three categories:

    #. shot numbers which are contained in the :code:`'shotnum'` field
    #. digitizer data which is contained in the :code:`'signal'` field
    #. control device data which is represented by the remaining fields
       in the numpy array.  These field names are polymorphic and are
       defined by the control device mapping class. (see
       :class:`~.hdfreadcontrols.HDFReadControls` for more detail)

    Data that is not shot number specific is stored in the :attr:`info`
    attribute.

    .. note::

        * Evey returned numpy array will have the :code:`'xyz'` field,
          which is reserved for probe position data.  If a control
          device specifies this field, then field will be filled with
          the control device data;otherwise, the field will be filled
          with :code:`numpy.nan` values.
    """

    __example_doc__ = """
    :Example: Here data is extracted from the digitizer
        :code:`'SIS crate'` and position data is mated from the
        control device :code:`'6K Compumotor'`.

        >>> # open HDF5 file
        >>> f = bapsflib.lapd.File('test.hdf5')
        >>>
        >>> # read digitizer data from board 1, channel 1,
        >>> # - this is equivalent to
        >>> #   f.read_data(1, 1)
        >>> data = HDFReadData(f, 1, 1)
        >>> data.dtype
        dtype([('shotnum', '<u4'), ('signal', '<f4', (100,)),
              ('xyz', '<f4', (3,))])
        >>>
        >>> # display shot numbers
        >>> data['shotnum']
        array([  1,  2, ..., 98, 99], dtype=uint32)
        >>>
        >>> # show 'signal' values for shot number 1
        >>> data['signal'][0]
        array([-0.41381955, -0.4134333 , -0.4118886 , ..., -0.41127062,
               -0.4105754 , -0.41119337], dtype=float32)
        >>>
        >>> # show 'xyz' values for shot number 1
        >>> data['xyz'][0]
        array([nan, nan, nan], dtype=float32)
        >>>
        >>> # read digitizer data while adding '6K Compumotor' data
        >>> # from receptacle (configuration) 3
        >>> data = HDFReadData(f, 1, 1,
                               add_controls=[('6K Compumotor', 3)])
        >>> data.dtype
        dtype([('shotnum', '<u4'), ('signal', '<f4', (100,)),
               ('xyz', '<f4', (3,)), ('ptip_rot_theta', '<f8'),
               ('ptip_rot_phi', '<f8')])
        >>>
        >>> # show 'xyz' values for shot number 1
        >>> data['xyz'][0]
        array([ -32. ,   15. , 1022.4], dtype=float32)

    """

    def __new__(
        cls,
        hdf_file: File,
        board: int,
        channel: int,
        index=slice(None),
        shotnum=slice(None),
        digitizer=None,
        config_name=None,
        adc=None,
        keep_bits=False,
        add_controls=None,
        intersection_set=True,
        **kwargs,
    ):
        """
        :param hdf_file: HDF5 file object
        :param board: analog-digital-converter board number
        :param channel: analog-digital-converter channel number
        :param index: dataset row indices to be sliced (overridden
            by :code:`shotnum`)
        :type index: Union[int, List[int], slice, numpy.ndarray]
        :param shotnum: HDF5 file shot number(s) indicating data
            entries to be extracted (overrides :code:`index`)
        :type shotnum: Union[int, List[int], slice, numpy.ndarray]
        :param str digitizer: digitizer name
        :param str adc: name of analog-digital-converter
        :param str config_name: name of the digitizer configuration
        :param bool keep_bits: set :code:`True` to keep data in bits,
            :code:`False` (DEFAULT) to convert data to voltage
        :param add_controls: a list indicating the desired control
            device names and their configuration name (if more than one
            configuration exists)
        :type controls: Union[str, Iterable[str, Tuple[str, Any]]]
        :param bool intersection_set: :code:`True` (DEFAULT) will force
            the returned shot numbers to be the intersection of
            :data:`shotnum` and the shot numbers contained in each
            control device and digitizer dataset. :code:`False` will
            return the union of shot numbers.

        Behavior of :data:`index`, :data:`shotnum` and
        :data:`intersection_set`:

        .. note::

            * The :data:`shotnum` keyword will always override the
              :data:`index` keyword, but, due to extra overhead
              required for identifying shot number locations in the
              digitizer dataset, the :data:`index` keyword will always
              execute quicker than the :data:`shotnum` keyword.
        """
        # initialize timing
        tt = []
        if "timeit" in kwargs:  # pragma: no cover
            timeit = kwargs["timeit"]
            if timeit:
                tt.append(time.time())
            else:
                timeit = False
        else:
            timeit = False

        # ---- Condition hdf_file                                   ----
        # - `hdf_file` is a lapd.File object
        #
        if not isinstance(hdf_file, File):
            raise TypeError(
                f"`hdf_file` is NOT type `{File.__module__}.{File.__qualname__}`"
            )

        # print execution timing
        if timeit:  # pragma: no cover
            tt.append(time.time())
            print(f"tt - `hdf_file` conditioning: {(tt[-1] - tt[-2]) * 1.0e3} ms")

        # ---- Examine file map object                              ----
        # grab instance of `HDFMap`
        _fmap = hdf_file.file_map

        # ---- Condition `add_controls`                             ----
        # Check for non-empty controls
        if bool(add_controls) and not bool(_fmap.controls):
            raise ValueError("There are no control devices in the HDF5 file.")

        # condition controls
        if bool(add_controls):
            controls = condition_controls(hdf_file, add_controls)
        else:
            controls = []

        # print execution timing
        if timeit:  # pragma: no cover
            tt.append(time.time())
            print(f"tt - `add_controls` conditioning: {(tt[-1] - tt[-2]) * 1.0e3} ms")

        # ---- Condition `digitizer` keyword                        ----
        if not bool(_fmap.digitizers):
            raise ValueError("There are no digitizers in the HDF5 file.")
        elif digitizer is None:
            if not bool(_fmap.main_digitizer):
                raise ValueError(
                    "No main digitizer is identified..."
                    "need to specify `digitizer` kwarg"
                )

            why = (
                f"Digitizer not specified so assuming the 'main_digitizer' "
                f"({_fmap.main_digitizer.device_name}) defined in the mappings."
            )
            warn(why)
            _dmap = _fmap.main_digitizer
        else:
            try:
                _dmap = _fmap.digitizers[digitizer]
            except KeyError:
                raise ValueError(
                    f"Specified Digitizer '{digitizer}' is not among known "
                    f"digitizers ({list(_fmap.digitizers)})"
                )

        # ---- Gather Digi Dataset Info                             ----
        #
        # Note: _dmap.construct_dataset_name has conditioning for
        #       board, channel, adc, and
        #
        # dname      - digitizer dataset name
        # dhname     - digitizer header dataset name
        # dpath      - full path to digitizer group
        # dset       - digitizer h5py.Dataset object
        # dheader    - dset associated header dataset
        # shotnumkey - field name for shot number column in dheader
        #
        # Build kwargs for construct_dataset_name()
        kwargs = {"return_info": True}
        if config_name is not None:
            kwargs["config_name"] = config_name
        if adc is not None:
            kwargs["adc"] = adc

        # Get datasets
        dname, d_info = _dmap.construct_dataset_name(board, channel, **kwargs)
        dhname = _dmap.construct_header_dataset_name(board, channel, **kwargs)
        dpath = f"{_dmap.info['group path']}/"
        dset = hdf_file.get(dpath + dname)
        dheader = hdf_file.get(dpath + dhname)

        # define `config_name`
        if config_name is None:
            config_name = _dmap.active_configs[0]

        # define `shotnumkey`
        shotnumkey = _dmap.configs[config_name]["shotnum"]["dset field"][0]

        # print execution timing
        if timeit:  # pragma: no cover
            tt.append(time.time())
            print(f"tt - get dset and dheader: {(tt[-1] - tt[-2]) * 1.0e3} ms")

        # ---- Condition shots, index, and shotnum ----
        # index   -- row index of digitizer dataset
        #            ~ indexed at 0
        #            ~ supersedes any other indexing keywords
        # shotnum -- global HDF5 file shot number
        #            ~ this is the index used to link values between
        #              datasets
        #            ~ overridden by `index`
        #
        # Through conditioning the following are (re-)defined
        # index   -- row index of digitizer dataset (dset)
        #            ~ numpy.ndarray
        #            ~ dtype = np.integer
        #            ~ shape = (num_of_indices,)
        #
        # shotnum -- global HDF5 shot numbers
        #            ~ index at 1
        #            ~ will be a filtered version of input kwarg shotnum
        #              based on intersection_set
        #            ~ numpy.ndarray
        #            ~ dtype = np.uint32
        #            ~ shape = (sn_size, )
        #
        # sni     -- bool array for providing a one-to-one mapping
        #            between shotnum and index
        #            ~ shotnum[sni] = dheader[index, shotnumkey]
        #            ~ data['signal'][sni, ...] = dset[index, ...]
        #            ~ data['singal'][np.logical_not(sni), ...] = np.nan
        #            ~ numpy.ndarray
        #            ~ dtype = np.bool
        #            ~ shape = (sn_size, )
        #            ~ np.count_nonzero(arr[0,...]) = num_of_indices
        #
        # - Indexing behavior: (depends on intersection_set)
        #
        #   ~ intersection_set = True (DEFAULT)
        #     * the returned array will only contain shot numbers that
        #       are in the intersection of shotnum, the digitizer
        #       dataset, and all the specified control device datasets
        #
        #   ~ intersection_set = False
        #     * the returned array will contain all shot numbers
        #       specified by shotnum (>= 1)
        #     * if a dataset does not included a shot number contained
        #       in shotnum, then its entry in the returned array will
        #       be given a NULL value depending on the dtype
        #
        # Determine if indexing w.r.t. `index` or `shotnum`
        index_with = "index"
        if isinstance(index, slice):
            if index == slice(None):
                if not isinstance(shotnum, slice):
                    index_with = "shotnum"
                elif shotnum != slice(None):
                    index_with = "shotnum"

        # Condition `index` and `shotnum` keywords
        # - Valid indexing types are: int, list(int), slice(), and
        #   np.ndarray
        #
        if index_with == "index":
            # Condition `index` keyword
            #
            # Note: I'm letting the slicing of dset[index, shotnumkey]
            #       throw the appropriate errors
            #
            # Define `shotnum`
            # - Note: h5py datasets can NOT be sliced using numpy arrays
            #
            # convert `index` to np.ndarray
            sn_size = dheader.size
            if isinstance(index, int):
                index = np.array([index], dtype=np.int32)
            elif isinstance(index, list):
                index = np.array(index, dtype=np.int32)
            elif isinstance(index, slice):
                start, stop, step = index.indices(sn_size)
                index = np.arange(start, stop, step, dtype=np.int32)
            elif isinstance(index, type(Ellipsis)):
                index = np.arange(0, sn_size, 1, dtype=np.int32)
            elif isinstance(index, np.ndarray):
                pass
            else:
                raise TypeError("Valid `index` type not passed.")

            # convert (VALID) negative indices to positive
            neg_index_mask = np.where((index < 0) & (index >= -sn_size), True, False)
            if np.any(neg_index_mask):
                adj_ii = index[neg_index_mask] % sn_size
                index[neg_index_mask] = adj_ii
            index = np.unique(index)

            # define `shotnum`
            shotnum = dheader[index.tolist(), shotnumkey]

            # define sni
            sni = np.ones(shotnum.shape[0], dtype=np.bool)

            # print execution timing
            if timeit:  # pragma: no cover
                tt.append(time.time())
                print(f"tt - condition index: {(tt[-1] - tt[-2]) * 1.0e3} ms")
        else:
            # Condition `shotnum` keyword
            #
            # convert `shotnum` to np.ndarray
            """
            if isinstance(shotnum, slice):
                # determine largest possible shot number
                last_sn = dheader[-1, shotnumkey]
                if shotnum.stop is not None:
                    stop_sn = max(shotnum.stop, last_sn + 1)
                else:
                    stop_sn = last_sn + 1

                # get the start, stop, and step for the shot number
                # array
                start, stop, step = shotnum.indices(stop_sn)

                # determine smallest possible shot number
                # - intersection_set = True
                #   * start = max of first_sn and shotnum.start
                # - intersection_set = False
                #   * start = min of first_sn and shotnum.start
                first_sn = [dheader[0, shotnumkey]]
                if shotnum.start is not None:
                    # ensure shot numbers are >= 1
                    if start <= 0:
                        start = 1
                else:
                    # start wasn't specified in slice object
                    start = min(first_sn)

                # adjust start for intersection_set
                if intersection_set:
                    first_sn.append(start)
                    start = max(first_sn)

                # re-define shotnum as a list
                shotnum = np.arange(start, stop, step).tolist()
            elif isinstance(shotnum, int):
                shotnum = [shotnum]
            elif isinstance(shotnum, list):
                # ensure all elements are int
                if not all(isinstance(sn, int) for sn in shotnum):
                    raise ValueError('Valid `shotnum` not passed')
            else:
                raise ValueError('Valid `shotnum` not passed')
            """
            # perform `shotnum` conditioning
            # - `shotnum` is returned as a numpy array
            shotnum = condition_shotnum(shotnum, {"digi": dheader}, {"digi": shotnumkey})

            # Calc. the corresponding `index` and `sni`
            # - `shotnum` will be converted from list to np.array
            # - `index` and `sni` will be np.array's
            """
            index, shotnum, sni = \
                condition_shotnum(shotnum, dheader, shotnumkey,
                                  intersection_set)
            """
            index, sni = build_sndr_for_simple_dset(shotnum, dheader, shotnumkey)

            # perform intersection
            if intersection_set:
                shotnum, sni_dict, index_dict = do_shotnum_intersection(
                    shotnum, {"digi": sni}, {"digi": index}
                )
                sni = sni_dict["digi"]
                index = index_dict["digi"]

            # print execution timing
            if timeit:  # pragma: no cover
                tt.append(time.time())
                print(f"tt - condition shotnum: {(tt[-1] - tt[-2]) * 1.0e3} ms")

        # ---- Retrieve Control Data                                ----
        # 1. retrieve the numpy array for control data
        # 2. re-filter shotnum if intersection_set=True s.t. only
        #    shotnum's w/ control data are returned
        #
        # grab control device dataset
        #
        # - this will ensure cdata.shape == data.shape all the time
        # - shotnum should always be a ndarray at this point
        #
        if len(controls) != 0:
            cdata = HDFReadControls(
                hdf_file,
                controls,
                assume_controls_conditioned=True,
                shotnum=shotnum,
                intersection_set=intersection_set,
            )

            # print execution timing
            if timeit:  # pragma: no cover
                tt.append(time.time())
                print(
                    f"tt - read in cdata (control data): {(tt[-1] - tt[-2]) * 1.0e3} ms"
                )

            # re-filter index, shotnum, and sni
            # - only need to be filtered if intersection_set=True
            # - for intersection_set=True, shotnum and index are
            #   one-to-one
            #
            if intersection_set:
                new_sn_mask = np.isin(shotnum, cdata["shotnum"])
                shotnum = shotnum[new_sn_mask]
                index = index[new_sn_mask]
                sni = np.ones(shotnum.shape[0], dtype=bool)
        else:
            cdata = None

        # ---- Build `obj`                                          ----
        # Define dtype and shape
        # - 1st column of the digi data header contains the global HDF5
        #   file shot number
        # - shotkey = is the field name/key of the dheader shot number
        #   column
        sigtype = np.float32 if not keep_bits else dset.dtype
        shape = shotnum.shape
        dtype = [
            ("shotnum", np.uint32, 1),
            ("signal", sigtype, dset.shape[1]),
            ("xyz", np.float32, 3),
        ]
        if len(controls) != 0:
            for subdtype in cdata.dtype.descr:
                if subdtype[0] not in [d[0] for d in dtype]:
                    dtype.append(subdtype)

        # print execution timing
        if timeit:  # pragma: no cover
            tt.append(time.time())
            print(f"tt - define dtype: {(tt[-1] - tt[-2]) * 1.0e3} ms")

        # Initialize data array
        data = np.empty(shape, dtype=dtype)

        # print execution timing
        if timeit:  # pragma: no cover
            tt.append(time.time())
            print(f"tt - initialize data np.ndarray: {(tt[-1] - tt[-2]) * 1.0e3} ms")

        # fill 'shotnum' field of data array
        data["shotnum"] = shotnum

        # fill 'signal' fields of data array
        index = index.tolist()
        if intersection_set:
            # fill signal
            data["signal"] = dset[index, ...]
        else:
            # fill signal
            data["signal"][sni] = dset[index, ...]
            if np.issubdtype(data["signal"].dtype, np.integer):
                data["signal"][np.logical_not(sni)] = 0
            else:
                # dtype is np.floating
                data["signal"][np.logical_not(sni)] = np.nan

        # fill fields related to controls
        if len(controls) != 0:
            # Note: shot numbers of cdata and data are one-to-one
            #       by this point so intersection_set is irrelevant
            #
            if not np.array_equal(data["shotnum"], cdata["shotnum"]):  # pragma: no cover
                # this should never happen
                raise ValueError("data['shotnum'] and cdata['shotnum'] are not equal")

            # fill xyz
            if "xyz" in cdata.dtype.names:
                data["xyz"] = cdata["xyz"]
            else:
                data["xyz"] = np.nan

            # fill remaining controls
            for field in cdata.dtype.names:
                if field not in ("shotnum", "xyz"):
                    data[field] = cdata[field]
        else:
            # fill xyz
            data["xyz"] = np.nan

        # print execution timing
        if timeit:  # pragma: no cover
            tt.append(time.time())
            print(f"tt - fill data array: {(tt[-1] - tt[-2]) * 1.0e3} ms")

        # Define obj to be returned
        obj = data.view(cls)

        # get voltage offset
        try:
            voffset = dheader[0, "Offset"] * u.volt
        except ValueError:
            warn("Digitizer header dataset is missing the voltage 'Offset' field. ")
            voffset = None

        # assign dataset meta-info
        obj._info = {
            "source file": os.path.abspath(hdf_file.filename),
            "device group path": _dmap.info["group path"],
            "device dataset path": dpath + dname,
            "digitizer": d_info["digitizer"],
            "configuration name": d_info["configuration name"],
            "adc": d_info["adc"],
            "bit": d_info["bit"],
            "clock rate": d_info["clock rate"],
            "sample average": d_info["sample average (hardware)"],
            "shot average": d_info["shot average (software)"],
            "board": board,
            "channel": channel,
            "voltage offset": voffset,
            "probe name": None,
            "port": (None, None),
            "signal units": u.bit,
        }
        if cdata is not None:
            obj._info["controls"] = copy.deepcopy(cdata.info["controls"])
        else:
            obj._info["controls"] = {}

        # plasma parameter dict
        obj._plasma = {
            "Bo": None,
            "kT": None,
            "kTe": None,
            "kTi": None,
            "gamma": core.FloatUnit(1.0, "arb"),
            "m_e": core.ME,
            "m_i": None,
            "n": None,
            "n_e": None,
            "n_i": None,
            "Z": None,
        }  # pragma: no cover

        # convert to voltage
        # - 'signal' dtype is assigned based on keep_bit
        #
        # obj['signal'] = obj['signal'].astype(np.float32, copy=False)
        #
        if not keep_bits:
            if obj.dv is None:
                warn("Unable to calculated voltage step size...'signal' remains as bits")
            else:
                # define offset
                offset = abs(obj.info["voltage offset"].value)

                # calc voltage
                obj["signal"] = (obj.dv.value * obj["signal"]) - offset

                # update 'signal units'
                obj._info["signal units"] = u.volt

        # print execution timing
        if timeit:  # pragma: no cover
            tt.append(time.time())
            print(f"tt - execution time: {(tt[-1] - tt[-2]) * 1.0e3} ms")

        # return obj
        return obj

    def __array_finalize__(self, obj):
        # This should only be True during explicit construction
        # if obj is None:
        if obj is None or obj.__class__ is np.ndarray:
            return

        # Define _info attribute
        self._info = getattr(
            obj,
            "_info",
            {
                "source file": None,
                "device group path": None,
                "device dataset path": None,
                "configuration name": None,
                "adc": None,
                "bit": None,
                "clock rate": None,
                "sample average": None,
                "shot average": None,
                "board": None,
                "channel": None,
                "voltage offset": None,
                "probe name": None,
                "port": (None, None),
                "signal units": None,
                "controls": {},
            },
        )

        # Define plasma attribute
        self._plasma = getattr(
            obj,
            "_plasma",
            {
                "Bo": None,
                "kT": None,
                "kTe": None,
                "kTi": None,
                "gamma": core.FloatUnit(1.0, "arb"),
                "m_e": core.ME,
                "m_i": None,
                "n": None,
                "n_e": None,
                "n_i": None,
                "Z": None,
            },
        )  # pragma: no cover

    def convert_signal(self, to_volt=False, to_bits=False, force=False):
        """converts signal from volts (bits) to bits (volts)"""
        #
        # 1. investigate if 'signal' values are np.integer or
        #    np.floating
        #    - np.integer => 'signal' is in bits
        #    - np.floating => 'signal' is in volts
        # 2. only convert if requested conversion is not current state
        # 3. update 'signal units' in self._info
        #
        raise NotImplementedError

    @property
    def info(self):
        """
        A dictionary of metadata for the extracted data. The
        dict() keys are:

        .. list-table::
            :widths: 5 3 11

            * - :const:`hdf file`
              - `str`
              - HDF5 file name the data was retrieved from
            * - :const:`dataset name`
              - `str`
              - name of the dataset
            * - :const:`dataset path`
              - `str`
              - path to said dataset in the HDF5 file
            * - :const:`digitizer`
              - `str`
              - digitizer group name
            * - :const:`configuration name`
              - `str`
              - name of data configuration
            * - :const:`adc`
              - `str`
              - analog-digital converter in which the data was recorded
                on
            * - :const:`bit`
              - `int`
              - bit resolution for the adc
            * - :const:`clock rate`
              - (`int`, `float`)
              - tuple containing clock rate, e.g. (100.0, 'MHz')
            * - :const:`clock rate`
              - `int`
              - (hardware sampling) number of data sample average
                together
            * - :const:`shot average`
              - `int`
              - (software averaging) number of shot sequences averaged
                together
            * - :const:`board`
              - `int`
              - board that the probe was connected to
            * - :const:`channel`
              - `int`
              - channel of the board that the probe was connected to
            * - :const:`voltage offset`
              - `float`
              - voltage offset of the digitized signal
            * - :const:`probe name`
              - `str`
              - name of deployed probe...empty for user to use at
                his/her discretion
            * - :const:`port`
              - (`int`, `str`)
              - 2-element tuple indicating which port the probe was
                deployed on, eg. (19, 'W')

        .. 'port' -- 2-element tuple indicating which port the probe was
                     deployed on. e.g. (19, 'W') => deployed on port 19
                     on the west side of the machine. Second elements
                     descriptors should follow:
                     'T'  = top
                     'TW' = top-west
                     'W'  = west
                     'BW' = bottom-west
                     'B'  = bottom
                     'BE' = bottome-east
                     'E'  = east
                     'TE' = top-east
        """
        return self._info

    @property
    def dt(self) -> Union[u.Quantity, None]:
        """
        Temporal step size (in sec) calculated from the
        :code:`'clock rate'` and :code:`'sample average'` items in
        :attr:`info`.  Returns :code:`None` if step size can not be
        calculated.
        """
        if not isinstance(self.info["clock rate"], u.Quantity):
            return

        # calc base dt
        dt = 1.0 / self.info["clock rate"]
        dt = dt.to("s")

        # adjust for hardware averaging
        if self.info["sample average"] is not None:
            dt = dt * float(self.info["sample average"])

        return dt

    @property
    def dv(self) -> Union[u.Quantity, None]:
        """
        Voltage step size (in volts) calculated from the :code:`'bit'`
        and :code:`'voltage offset'` items in :attr:`info`.  Returns
        :code:`None` if step size can not be calculated.
        """
        if self.info["voltage offset"] is None:
            return
        elif self.info["bit"] is None:
            return

        dv = 2.0 * abs(self.info["voltage offset"]) / (2.0 ** self.info["bit"] - 1.0)
        return dv

    @property
    def plasma(self):  # pragma: no cover
        """
        Dictionary of plasma parameters. (All quantities are in cgs
        units except temperature is in eV)

        +----------------+---------------------------------------------+
        | Base Values                                                  |
        +================+=============================================+
        | :const:`Bo`    | magnetic field                              |
        +----------------+---------------------------------------------+
        | :const:`kT`    | temperature (generic)                       |
        +----------------+---------------------------------------------+
        | :const:`kTe`   | electron temperature                        |
        +----------------+---------------------------------------------+
        | :const:`kTi`   | ion temperature                             |
        +----------------+---------------------------------------------+
        | :const:`gamma` | adiabatic index                             |
        +----------------+---------------------------------------------+
        | :const:`m_e`   | electron mass                               |
        +----------------+---------------------------------------------+
        | :const:`m_i`   | ion mass                                    |
        +----------------+---------------------------------------------+
        | :const:`n`     | plasma number density                       |
        +----------------+---------------------------------------------+
        | :const:`n_e`   | electron number density                     |
        +----------------+---------------------------------------------+
        | :const:`n_i`   | ion number density                          |
        +----------------+---------------------------------------------+
        | :const:`Z`     | ion charge number                           |
        +----------------+---------------------------------------------+
        | Calculated Values                                            |
        +----------------+---------------------------------------------+
        | :const:`fce`   | electron cyclotron frequency                |
        +----------------+---------------------------------------------+
        | :const:`fci`   | ion cyclotron frequency                     |
        +----------------+---------------------------------------------+
        | :const:`fpe`   | electron plasma frequency                   |
        +----------------+---------------------------------------------+
        | :const:`fpi`   | ion plasma frequency                        |
        +----------------+---------------------------------------------+
        | :const:`fUH`   | Upper-Hybrid Resonance frequency            |
        +----------------+---------------------------------------------+
        | :const:`lD`    | Debye Length                                |
        +----------------+---------------------------------------------+
        | :const:`lpe`   | electron-inertial length                    |
        +----------------+---------------------------------------------+
        | :const:`lpi`   | ion-inertial length                         |
        +----------------+---------------------------------------------+
        | :const:`rce`   | electron gyroradius                         |
        +----------------+---------------------------------------------+
        | :const:`rci`   | ion gyroradius                              |
        +----------------+---------------------------------------------+
        | :const:`cs`    | ion sound speed                             |
        +----------------+---------------------------------------------+
        | :const:`VA`    | Alfven speed                                |
        +----------------+---------------------------------------------+
        | :const:`vTe`   | electron thermal velocity                   |
        +----------------+---------------------------------------------+
        | :const:`vTi`   | ion thermal velocity                        |
        +----------------+---------------------------------------------+
        """
        return self._plasma

    def set_plasma(
        self, Bo, kTe, kTi, m_i, n_e, Z, gamma=None, **kwargs
    ):  # pragma: no cover
        """
        Set :attr:`plasma` and add key frequency, length, and velocity
        parameters. (all quantities in cgs except temperature is in eV)

        :param float Bo: magnetic field (in Gauss)
        :param float kTe: electron temperature (in eV)
        :param float kTi: ion temperature (in eV)
        :param float m_i: ion mass (in g)
        :param float n_e: electron number density (in cm^-3)
        :param int Z: ion charge number
        :param float gamma: adiabatic index (arb.)
        """
        # define base values
        self._plasma["Bo"] = core.FloatUnit(Bo, "G")
        self._plasma["kTe"] = core.FloatUnit(kTe, "eV")
        self._plasma["kTi"] = core.FloatUnit(kTi, "eV")
        self._plasma["m_i"] = core.FloatUnit(m_i, "g")
        self._plasma["n_e"] = core.FloatUnit(n_e, "cm^-3")
        self._plasma["Z"] = core.IntUnit(Z, "arb")

        # define ion number density
        self._plasma["n_i"] = core.FloatUnit(
            self._plasma["n_e"] / self._plasma["Z"], "cm^-3"
        )

        # define gamma (adiabatic index)
        # - default = 1.0
        if gamma is not None:
            self._plasma["gamma"] = core.FloatUnit(gamma, "arb")

        # define plasma temperature
        # - if omitted then assumed kTe
        # TODO: double check assumption
        if "kT" in kwargs:
            self._plasma["kT"] = core.FloatUnit(kwargs["kT"], "eV")
        else:
            self._plasma["kT"] = core.FloatUnit(kTe, "eV")

        # define plasma number density
        # - if omitted then assumed n_e
        if "n" in kwargs:
            self._plasma["n"] = core.FloatUnit(kwargs["n"], "cm^-3")
        else:
            self._plasma["n"] = core.FloatUnit(n_e, "cm^-3")

        # add key plasma constants
        self._update_plasma_constants()

    def set_plasma_value(self, key, value):  # pragma: no cover
        """
        Re-define one of the base plasma values (Bo, gamma, kT, kTe,
        kTi, m_i, n, n_e, or Z) in the :attr:`plasma` dictionary.

        :param str key: one of the base plasma values
        :param value: value for key
        """
        # set plasma value
        if key == "Bo":
            self._plasma["Bo"] = core.FloatUnit(value, "G")
        elif key == "gamma":
            self._plasma["gamma"] = core.FloatUnit(value, "arb")
        elif key in ["kT", "kTe", "kTi"]:
            self._plasma[key] = core.FloatUnit(value, "eV")

            if key == "kTe" and self._plasma["kt"] is None:
                self._plasma["kT"] = self._plasma[key]
        elif key == "m_i":
            self._plasma[key] = core.FloatUnit(value, "g")
        elif key in ["n", "n_e"]:
            self._plasma[key] = core.FloatUnit(value, "cm^-3")

            # re-calc n_i and n
            if key == "n_e":
                self._plasma["n_i"] = core.FloatUnit(
                    self._plasma["n_e"] / self._plasma["Z"], "cm^-3"
                )

                if self._plasma["n"] is None:
                    self._plasma["n"] = self._plasma["n_e"]
        elif key == "Z":
            self._plasma[key] = core.IntUnit(value, "arb")

            # re-calc n_i
            self._plasma["n_i"] = core.FloatUnit(
                self._plasma["n_e"] / self._plasma["Z"], "cm^-3"
            )

        # update key plasma constants
        self._update_plasma_constants()

    def _update_plasma_constants(self):  # pragma: no cover
        """
        Updates the calculated plasma constants (fci, fce, fpe, etc.) in
        :attr:`plasma`.
        """
        # add key frequencies
        self._plasma["fce"] = core.fce(**self._plasma)
        self._plasma["fci"] = core.fci(**self._plasma)
        self._plasma["fpe"] = core.fpe(**self._plasma)
        self._plasma["fpi"] = core.fpi(**self._plasma)
        self._plasma["fUH"] = core.fUH(**self._plasma)
        self._plasma["fLH"] = core.fLH(**self._plasma)

        # add key lengths
        self._plasma["lD"] = core.lD(**self._plasma)
        self._plasma["lpe"] = core.lpe(**self._plasma)
        self._plasma["lpi"] = core.lpi(**self._plasma)
        self._plasma["rce"] = core.rce(**self._plasma)
        self._plasma["rci"] = core.rci(**self._plasma)

        # add key velocities
        self._plasma["cs"] = core.cs(**self._plasma)
        self._plasma["VA"] = core.VA(**self._plasma)
        self._plasma["vTe"] = core.vTe(**self._plasma)
        self._plasma["vTi"] = core.vTi(**self._plasma)

# add example to __new__ docstring
HDFReadData.__new__.__doc__ += "\n"
for line in HDFReadData.__example_doc__.splitlines():
    HDFReadData.__new__.__doc__ += f"    {line}\n"

'''
def condition_shotnum(shotnum, dheader, shotnumkey,
                      intersection_set):
    """
    Conditions **shotnum** against the digitizer header dataset.

    :param shotnum: desired HDF5 shot number(s)
    :type shotnum: list(int)
    :param dheader: digitizer header dataset
    :param dheader: :class:`h5py.Dataset`
    :param str shotnumkey: field name in **dheader** that contains the
        shot numbers
    :param bool intersection_set: Set :code:`True` to intersect
        **shotnum** with the shot numbers in :code:`dheader[shotnumkey]`
    :return: index, shotnum, sni

    .. note::

        The returned :class:`numpy.ndarray`'s (:const:`index`,
        :const:`shotnum`, and :const:`sni`) follow the rule::

            shotnum[sni] = dheader[index, shotnumkey]
    """
    # Inputs:
    # shotnum           list(int)     - the desired shot number(s)
    # cheader           h5py.Dataset  - the digi header dataset
    # shotnumkey        str           - field name for the shot number
    #                                   column in dheader
    # intersection_set  bool          - intersect shotnum with
    #                                   dheader[shotnumkey]
    #
    # Returns:
    # index    np.array(dtype=uint32) - cdset row index for the
    #                                   specified shotnum
    # shotnum  np.array(dtype=uint32) - shot numbers
    # sni      np.array(dtype=bool)   - shotnum mask such that:
    #            shotnum[sni] = cdset[index, shotnumkey]
    #
    # remove shot numbers less-than or equal to 0
    shotnum.sort()
    shotnum = list(set(shotnum))
    shotnum.sort()
    if min(shotnum) <= 0:
        # remove values less-than or equal to 0
        new_sn = [sn for sn in shotnum if sn > 0]
        shotnum = new_sn

    # remove shot numbers greater-than largest shot number
    # in dataset
    if intersection_set:
        last_sn = dheader[-1, shotnumkey]
        if max(shotnum) > last_sn:
            new_sn = [sn for sn in shotnum if sn <= last_sn]
            shotnum = new_sn

    # ensure shotnum is not empty
    if len(shotnum) == 0:
        raise ValueError('Valid shotnum not passed.')

    # convert shotnum to np.array
    # - shotnum is always a list up to this point
    shotnum = np.array(shotnum).view()

    # Calc. corresponding `index` and `sni` for shotnum
    # - intersection will after initial calculation
    #
    if dheader.shape[0] == 1:
        # only one possible shot number
        only_sn = dheader[0, shotnumkey]
        sni = np.where(shotnum == only_sn, True, False)
        index = np.array([0]) \
            if True in sni else np.empty(shape=0, dtype=np.uint32)
    else:
        # get 1st and last shot number for further
        # conditioning
        first_sn, last_sn = dheader[[-1, 0], shotnumkey]

        if last_sn - first_sn + 1 == dheader.shape[0]:
            # shot numbers are sequential
            index = shotnum - first_sn

            # build sni
            # - this will also mask index s.t. index has
            #   no values outside dheader.shape
            sni = np.where(index < dheader.shape[0], True, False)
            index = index[sni]
        else:
            # shot numbers are NOT sequential
            # TODO: check for more efficient read in methods
            #
            step_front_read = shotnum[-1] - first_sn
            step_end_read = last_sn - shotnum[0]

            if dheader.shape[0] <= 1 + min(step_front_read,
                                           step_end_read):
                # dheader.shape is smaller than the
                # theoretical sequential reads from either
                # end of the array
                dset_sn = dheader[shotnumkey].view()
                sni = np.isin(shotnum, dset_sn)

                # define index
                index = np.where(np.isin(dset_sn, shotnum))[0]
            elif step_front_read <= step_end_read:
                # extracting from the beginning of the array is the
                # smallest
                some_dset_sn = dheader[0:step_front_read + 1,
                                       shotnumkey]
                sni = np.isin(shotnum, some_dset_sn)

                # define index
                index = np.where(np.isin(some_dset_sn, shotnum))[0]
            else:
                # extracting from the end of the array is the smallest
                start, stop, step = \
                    slice(-step_end_read - 1,
                          None,
                          None).indices(dheader.shape[0])
                some_dset_sn = dheader[start::, shotnumkey]
                sni = np.isin(shotnum, some_dset_sn)

                # define index
                # NOTE: if index is empty (i.e. index.shape[0] == 0)
                #       then adding an int still returns an empty array
                index = np.where(np.isin(some_dset_sn, shotnum))[0]
                index += start

    # filter shotnum and ensure obj will not be empty
    if intersection_set:
        # check for empty shotnum
        if True not in sni:
            raise ValueError(
                'Input shotnum would result in a null array')

        # filter
        shotnum = shotnum[sni]
        sni = np.ones(shotnum.shape, dtype=bool)
    else:
        # empty shotnum
        if shotnum.shape[0] == 0:
            raise ValueError(
                'Input shotnum would result in a null array')

    # return calculated arrays
    return index.view(), shotnum.view(), sni.view()
'''

max_stars_repo_name: BaPSF/bapsflib

max_issues_repo_name: BaPSF/bapsflib

reason: import numpy,import astropy


text:
from itertools import product
import empowerment_maximization
import numpy as np
from functools import lru_cache
import random
import global_configuration
import env_util

EPSILON = 1e-5

# Hashes the numpy array of observations
def hash_obs(obs, player_pos):
    if player_pos is None:
        return hash(b'a')
    return hash(obs.tobytes() + bytes(player_pos))

@lru_cache(maxsize=2000)
def find_player_pos(wrapped_env, player_id):
    agent_actions = wrapped_env.get_available_actions(player_id)
    if not agent_actions:
        return None
    return list(wrapped_env.get_available_actions(player_id))[0]

def find_player_pos_vanilla(env, player_id):
    return list(env.game.get_available_actions(player_id))[0]

def find_player_health(env_state, player_id):
    player_variables = [o['Variables'] for o in env_state['Objects'] if o['Name'] == 'avatar' and o['PlayerId'] == player_id]
    if not player_variables:
        return 0
    # If health is not an object variable, return None
    if 'health' not in player_variables[0]:
        return None
    return player_variables[0]['health']

@lru_cache(maxsize=2000)
def find_alive_players(wrapped_env):
    return [player_id for player_id in range(1, wrapped_env.player_count + 1) if wrapped_env.get_available_actions(player_id)]

class EnvHashWrapper():
    def __init__(self, env):
        # The environment is very likely going to be stepped. Thus, we don't want to use the original environment here.
        if not env._is_clone:
            raise ValueError('The environment given to EnvHashWrapper must be a clone of an original environment.')
        self._env = env
        self._hash = None
        self.finished = False
        self.winner = 0

    def __hash__(self) -> int:
        return self.get_hash()

    def __eq__(self, __o: object) -> bool:
        return isinstance(__o, EnvHashWrapper) and self.__hash__() == __o.__hash__()

    def get_state(self):
        return self._env.get_state()

    def get_hash(self):
        if self._hash is None:
            if self.finished and self.winner > 0:
                self._hash = hash(self.winner)
            else:
                self._hash = self._env.get_state()['Hash']
        return self._hash

    def clone(self):
        return EnvHashWrapper(self._env.clone())

    def set_winner(self, winner):
        self.finished = True
        self.winner = winner
        self._hash = None

    # We change the API here a bit, because we need this object to be immutable
    def step(self, actions):
        clone_env = self._env.clone()
        obs, rew, env_done, info = clone_env.step(actions)
        return EnvHashWrapper(clone_env), obs, rew, env_done, info

    def get_available_actions(self, player_id):
        return self._env.game.get_available_actions(player_id)

    def player_last_obs(self, player_id):
        if not self._env._player_last_observation:
            self._env.step([[0, 0] for _ in range(self.player_count)])
        return self._env._player_last_observation[player_id - 1]

    @property
    def player_count(self):
        return self._env.player_count

class CEM():
    def __init__(self, env, game_conf, seed=None, samples=1):
        self.samples = samples
        self.player_count = env.player_count
        self.game_conf = game_conf
        self.action_spaces = env_util.build_action_spaces(env, self.agent_confs)
        self.rng = np.random.default_rng() if seed is None else np.random.default_rng(seed)

    @property
    def agent_confs(self):
        return self.game_conf.agents

    def cem_action(self, env, player_id, n_step):
        # If all empowerments are weighted to zero, all actions are equal.
        emp_pairs = self.get_emp_pairs(player_id)
        all_weights_zero = all([emp_pair.weight == 0 for emp_pair in emp_pairs])
        if all_weights_zero:
            best_actions = self.action_spaces[player_id-1]
        else:
            expected_empowerments_per_pair = self.calculate_expected_empowerments_per_pair(env, player_id, n_step)
            best_actions = self.find_best_actions(player_id, expected_empowerments_per_pair)
        return random.choice(best_actions)

    def get_emp_pairs(self, player_id):
        return self.agent_confs[player_id-1].empowerment_pairs

    def calculate_expected_empowerments_per_pair(self, env, player_id, n_step):
        '''
        Returns the expected empowerment for each empowerment_pair, for each action that can be taken from the current state.
        For example: [E[E^P]_{a_t}, E[E^T]_{a_t}, E[E^C]_{a_t}], if we were to calculate three different empowerments E^P, E^T and E^C.
        Here's an example of the structure, if there are 3 empowerment pairs (pair 0, pair 1, pair 2) and 3 actions (a0, a1, a2)
                          [ [1.0, 1.5, 2.0], [1.5, 1.2, 1.9], [3.0, 2.5, 1.0] ]  <- example data
                              |    |    |      |    |    |      |    |    |
                              a0   a1   a2     a0   a1   a2     a0   a1   a2
                            |---pair 0----|  |----pair 1---|  |---pair 2----|
        '''
        expected_empowerments_per_pair = []
        for emp_pair in self.get_emp_pairs(player_id):
            expected_empowerments = self.calculate_expected_empowerments(env, player_id, (emp_pair.actor, emp_pair.perceptor), n_step, True)
            expected_empowerments_per_pair.append(expected_empowerments)
        return expected_empowerments_per_pair

    def find_best_actions(self, player_id, expected_empowerments_per_pair):
        best_actions = []
        best_reward = -1e10
        for action in self.action_spaces[player_id-1]:
            empowerment_pair_confs = enumerate(self.get_emp_pairs(player_id))
            action_reward = self.calculate_action_reward(action, empowerment_pair_confs, expected_empowerments_per_pair)
            if action_reward >= best_reward + EPSILON:
                best_actions = [action]
            elif abs(action_reward-best_reward) < EPSILON:
                best_actions.append(action)
            best_reward = max(best_reward, action_reward)
            if global_configuration.verbose_calculation:
                print(f'Action {action} has expected reward {action_reward}')
        return best_actions

    def calculate_action_reward(self, action, empowerment_pair_confs, expected_empowerments_per_pair):
        action_reward = 0
        for emp_pair_i, emp_conf in empowerment_pair_confs:
            action_reward += emp_conf.weight * expected_empowerments_per_pair[emp_pair_i][action]
        return action_reward

    @lru_cache(maxsize=1000)
    def calc_pd_s_a(self, env, player_id, action):
        health_ratio = self.get_health_ratio(env, player_id)
        # if the player is dead, the empowerment is 0
        if health_ratio == 0:
            return {env: 1.0}

        result = {}
        for _ in range(self.samples):
            next_env, obs, rew, env_done, info = env.step(env_util.build_action(action, self.player_count, player_id))
            if (env_done):
                game_winner = env_util.find_winner(info)
                next_env.set_winner(game_winner)
            if next_env not in result:
                result[next_env] = 0
            result[next_env] += 1.0 / self.samples

        # Adjust for health-performance consistency
        if health_ratio < 1 - EPSILON:
            self.health_perf_consistency(env, health_ratio, result)
        return result

    def get_health_ratio(self, env, player_id):
        if self.game_conf.health_performance_consistency:
            player_health = find_player_health(env.get_state(), player_id)
            if player_health is not None:
                health_ratio = player_health / self.agent_confs[player_id-1].max_health
            else:
                health_ratio = 1.0
        else:
            health_ratio = 1.0
        return health_ratio

    def health_perf_consistency(self, env, health_ratio, transitions):
        # If the following env is not the current_env, adjust the transition probability with health_ratio
        for next_env in transitions:
            if next_env != env:
                transitions[next_env] *= health_ratio
        # Add the current env to the possible transitions
        if env not in transitions:
            transitions[env] = 0
        # Probability of staying still is its original probability and "leftovers" of the other states.
        transitions[env] = 1 - health_ratio + health_ratio * transitions[env]

    def find_assumed_policy(self, agent_id, wrapped_env):
        current_agent_conf = self.agent_confs[agent_id-1]
        # Find the policy of the agent in turn
        return current_agent_conf.assumed_policy(wrapped_env._env , self, agent_id, self.game_conf)

    def build_distribution(self, wrapped_env, action_seq, action_stepper, actor, current_step_agent, perceptor, return_obs=False, anticipation=False, trust_correction=False):
        '''
        Recursive function
        Builds the distribution p(S_t+n|s_t, a_t^n)
        s_t is given by wrapped_env
        a_t^n is given by action_seq

        Params:
            wrapped_env defines s_t             -> the original state, wrapped in a hasher
            action_seq defines a_t^n            -> the action sequence for THE PLAYER WE ARE INTERESTED IN
            action_stepper                      -> which action is next in turn in the action sequence
            actor                               -> player that we are interested in (for whom the action sequence applies to)
            current_step_player                 -> the player whose turn it is next, if this is the active player, then we use the action sequence. Otherwise, we build the distribution for all possible actions
            perceptor                           -> the agent until whose perception the distribution is built, after all actions have been taken
            return_obs                          -> if true, we don't return the end states but corresponding observations
            anticipation                        -> boolean value indicating whether this is an anticipation step or not
            trust_correction                    -> boolean value indicating whether we should correct for trust

        Returns:
            A dictionary, where the keys are the states(/observations) and the values are the probabilities: {state: probability}
        '''
        def is_final_state():
            return (action_stepper == len(action_seq) and current_step_agent == perceptor) or wrapped_env.finished

        def get_final_return_object():
            if return_obs:
                return { self.env_to_hashed_obs(wrapped_env, actor, perceptor): 1.0 }
            else:
                return { wrapped_env: 1.0 }

        def is_trust_applied():
            actor_conf = self.agent_confs[actor-1]
            if trust_correction and actor_conf.trust is not None:
                # Find the trust setting for the current agent
                trust_current = [trust_conf for trust_conf in actor_conf.trust if trust_conf.player_id == current_step_agent]
                if len(trust_current) == 1:
                    # Check if the trust correction should be applied at this step (anticipation or one of the following steps)
                    return (trust_current[0].anticipation and anticipation) or action_stepper in trust_current[0].steps
            return False

        def next_state_destroys_empowerment(current_env, next_state):
            follow_up_emp = self.calculate_state_empowerment(next_state, actor, actor, 1)
            if follow_up_emp < EPSILON:
                current_state_emp = self.calculate_state_empowerment(current_env, actor, actor, 1)
                return current_state_emp >= EPSILON
            return False

        if is_final_state():
            return get_final_return_object()

        next_step_agent = current_step_agent % self.player_count + 1
        # Increase the action stepper if this agent was the active one.
        next_action_step = action_stepper + 1 if actor == current_step_agent else action_stepper

        # If the current agent is not alive, go straight to the next agent
        if not current_step_agent in find_alive_players(wrapped_env):
            return self.build_distribution(wrapped_env, action_seq, next_action_step, actor, next_step_agent, perceptor, return_obs, anticipation, trust_correction)

        pd_s_nstep = {}

        # Find the probability of each action for the current agent
        assumed_policy = self.find_assumed_policy(current_step_agent, wrapped_env) if actor != current_step_agent else {self.action_spaces[current_step_agent-1][action_seq[action_stepper]]: 1.0}

        apply_trust = is_trust_applied()

        # Sum up the total probability of all possible actions so we can normalize in the end if needed
        total_prob = 0
        action_to_pd_s = {}
        for action in assumed_policy:
            if assumed_policy[action] == 0:
                continue
            # Get the probability distribution for the next step, given an action
            action_to_pd_s[action] = self.calc_pd_s_a(wrapped_env, current_step_agent, action)

        trust_filtered_actions = list(action_to_pd_s)
        if apply_trust:
            for action in action_to_pd_s:
                next_step_pd_s = action_to_pd_s[action]
                for next_state, next_state_prob in next_step_pd_s.items():
                    # If the actor trusts the current agent, we skip all actions that would reduce the actor's empowerment to zero
                    if next_state_prob > 0.01 and next_state_destroys_empowerment(wrapped_env, next_state):
                        trust_filtered_actions.remove(action)
                        break
        # If all actions led to no empowerment, we should consider them all.
        if len(trust_filtered_actions) == 0:
            trust_filtered_actions = list(action_to_pd_s)

        for action in trust_filtered_actions:
            # Recursively, build the distribution for each possbile follow-up state
            next_step_pd_s = action_to_pd_s[action]
            for next_state, next_state_prob in next_step_pd_s.items():
                child_distribution = self.build_distribution(next_state, action_seq, next_action_step, actor, next_step_agent, perceptor, return_obs, anticipation, trust_correction)
                # Add the follow-up states to the overall distribution of this state p(S_t+n|s_t, a_t^n)
                for next_env in child_distribution:
                    if next_env not in pd_s_nstep:
                        pd_s_nstep[next_env] = 0
                    adjusted_prob = child_distribution[next_env] * assumed_policy[action] * next_state_prob
                    pd_s_nstep[next_env] += adjusted_prob
                    total_prob += adjusted_prob

        # Normalize the probabilities. Could not be 1, if we skip one because of trust correction, for instance.
        if total_prob > 0 and total_prob < 1 - EPSILON:
            for state in pd_s_nstep:
                pd_s_nstep[state] /= total_prob
        return pd_s_nstep

    @lru_cache(maxsize=4000)
    def calculate_state_empowerment(self, wrapped_env, actor, perceptor, n_step, trust_correction=False):
        # If the player isn't alive anymore, we assume the empowerment to be zero
        if not wrapped_env.finished and find_player_pos(wrapped_env, actor) is None:
            return 0

        # All possible action combinations of length 'step'
        action_sequences = [tuple(combo) for combo in product(range(len(self.action_spaces[actor-1])), repeat=n_step)]
        # A list of end state probabilities, for each action combination. A list of dictionaries.
        cpd_S_A_nstep = [None] * len(action_sequences)
        for seq_idx, action_seq in enumerate(action_sequences):
            cpd_S_A_nstep[seq_idx] = self.build_distribution(wrapped_env, action_seq, 0, actor, actor, perceptor, return_obs=True, trust_correction=trust_correction)

        # Use the Blahut-Arimoto algorithm to calculate the optimal probability distribution
        pd_a_opt = empowerment_maximization.blahut_arimoto(cpd_S_A_nstep, self.rng)
        # Use the optimal distribution to calculate the mutual information (empowerement)
        empowerment = empowerment_maximization.mutual_information(pd_a_opt, cpd_S_A_nstep)
        return empowerment

    def calculate_expected_empowerments(self, gym_env, current_player, emp_pair, n_step, trust_correction):
        '''
        Calculates the expected empowerment for each action that can be taken from the current state, for one empowerment pair (actor, perceptor).
        '''
        # Find the probability of follow-up states for when the actor is in turn. p(s|s_t, a_t), s = state when actor is in turn
        cpd_s_a_anticipation = {}
        for action_idx, action in enumerate(self.action_spaces[current_player-1]):
            cpd_s_a_anticipation[action] = self.build_distribution(EnvHashWrapper(gym_env.clone()), (action_idx,), 0, current_player, current_player, emp_pair[0], return_obs=False, anticipation=True, trust_correction=trust_correction)
        # Make a flat list of all the reachable states
        reachable_states = set()
        for cpd_s in cpd_s_a_anticipation.values():
            for future_state in cpd_s:
                reachable_states.add(future_state)
        # Save the expected empowerment for each anticipated state
        reachable_state_empowerments = {}
        # Calculate the n-step empowerment for each state that was found earlier
        for state in reachable_states:
            empowerment = self.calculate_state_empowerment(state, emp_pair[0], emp_pair[1], n_step, trust_correction=trust_correction)
            reachable_state_empowerments[state] = empowerment
        # Calculate the expected empowerment for each action that can be taken from the current state
        expected_empowerments = {}
        for a in self.action_spaces[current_player-1]:
            for state, state_probability in cpd_s_a_anticipation[a].items():
                if not a in expected_empowerments:
                    expected_empowerments[a] = 0
                expected_empowerments[a] += state_probability * reachable_state_empowerments[state]

        return expected_empowerments

    def env_to_hashed_obs(self, wrapped_env, actor, perceptor):
        if wrapped_env.finished:
            # If the actor is in the winning team, all their actions lead to maximum empowerment
            if actor == wrapped_env.winner:
                hashed_obs = self.rng.integers(100, 4000000000)
                return hashed_obs
            # But if the actor loses future actions lead to a minimum empowerment
            elif wrapped_env.winner > 0:
                hashed_obs = wrapped_env.winner
                return hashed_obs
        latest_obs = wrapped_env.player_last_obs(perceptor)
        player_pos = find_player_pos(wrapped_env, perceptor)
        # If the player was not found, we assume the player is dead. In that case, the hashed observation is the player's id
        hashed_obs = hash_obs(latest_obs, player_pos) if player_pos is not None else hash(actor)
        return hashed_obs

ext: py

max_stars_repo_name: Thaigun/cem-experiments

max_issues_repo_name: Thaigun/cem-experiments

max_forks_repo_name: Thaigun/cem-experiments

reason: import numpy


text:
import pandas as pd
import numpy as np
from collections import Counter
from sklearn.base import BaseEstimator, TransformerMixin

#@author: Mohsin
class TargetEncoderWithThresh(BaseEstimator, TransformerMixin):
    """
    A utlity class to help encode categorical variables using different methods.

    Ideas taken from this excellent kernel
    https://www.kaggle.com/vprokopev/mean-likelihood-encodings-a-comprehensive-study?utm_medium=email&utm_source=mailchimp&utm_campaign=datanotes-20181004

    Inputs:
    cols: (List or str) Can be either a string or list of strings with column names
    targetcol: (str) Target column to encode column/group of columns with
    thresh: (int) Minimum count of grouping to encode (Acts as smoothing). Currently not implemented TODO
    func: (str or callable) Function to be applied on column/ group of columns to encode.
          If str is provided, it should be a attribute of pandas Series
    cname: (str) Column name for new string
    func_kwargs: (dict) Additional arguments to be passed to function
    add_to_orig: (bool) Whether to return dataframe with added feature or just the feature as series
    use_prior: (bool) Use smoothing as suggested in kernel
    alpha: (float) Smoothing factor

    Output:
    pandas DataFrame/Series

    """

    def __init__(self, cols=None, targetcol=None, cname=None, thresh=0, func=np.mean, add_to_orig=False,
                 func_kwargs={}, use_prior=False, alpha=0.5):
        self.cols = cols  # Can be either a string or list of strings with column names
        self.targetcol = targetcol  # Target column to encode column/group of columns with
        self.thresh = thresh  # Minimum count of grouping to encode (Acts as smoothing)
        self.func = func  # Function to be applied on column/ group of columns to encode
        self.add_to_orig = add_to_orig  # Whether return a dataframe with added feature or just a series of feature
        self.cname = cname  # Column to new feature generated
        self.func_kwargs = func_kwargs  # Additional key word arguments to be applied to func
        self.alpha = alpha  # smoothing factor
        self.use_prior = use_prior

    # @numba.jit
    def fit(self, X, y=None):

        if isinstance(self.func, str):
            if hasattr(pd.Series, self.func):
                # print("here")
                vals = getattr(X.groupby(self.cols)[self.targetcol], self.func)
                self.dictmap = vals(**self.func_kwargs)
                prior = getattr(X[self.targetcol], self.func)(**self.func_kwargs)

        else:
            self.dictmap = X.groupby(self.cols)[self.targetcol].apply(lambda x: self.func(x, **self.func_kwargs))
            prior = X[[self.targetcol]].apply(lambda x: self.func(x, **self.func_kwargs)).values[0]
        self.counts = Counter(zip(*[X[col].tolist() for col in self.cols]))
        if len(self.cols) == 1:
            counts_greater_than_thresh = [k[0] for k, v in self.counts.items() if v >= self.thresh]
        else:
            counts_greater_than_thresh = [k for k, v in self.counts.items() if v >= self.thresh]

        # print(self.dictmap.head())
        # print(self.counts.most_common(10))
        self.dictmap = self.dictmap.loc[self.dictmap.index.isin(counts_greater_than_thresh)]
        if self.use_prior:
            self.dictmap = {k: ((self.counts[k] * v + prior * self.alpha) / (self.counts[k] + self.alpha))
                            for k, v in self.dictmap.items()}
            self.dictmap = pd.Series(self.dictmap)
            self.dictmap.index.names = self.cols
        # print(self.dictmap.head())

        if self.cname:
            self.dictmap.name = self.cname
        else:
            cname = ''
            cname = [cname + '_' + str(col) for col in self.cols]
            self.cname = '_'.join(cname) + "_" + str(self.func)
            self.dictmap.name = self.cname

        # print(self.cname)
        # self.dictmap = self.dictmap
        return self

    def transform(self, X, y=None):
        if isinstance(X, pd.DataFrame):
            X_transformed = X[self.cols]

            X_transformed = X_transformed.join(self.dictmap, on=self.cols, how='left')[self.cname]

            if self.add_to_orig:
                return pd.concat([X, X_transformed], axis=1, copy=False)
            else:
                return X_transformed.values

        else:
            raise TypeError("Input should be a pandas DataFrame")

max_stars_repo_name: mohsinkhn/ltfs-av

max_issues_repo_name: mohsinkhn/ltfs-av

max_forks_repo_name: mohsinkhn/ltfs-av


text:
import numpy as np
import tensorflow as tf
import pandas as pd

import app.crud as crud
# import crud
from sqlalchemy.orm import Session
from app.database import SessionLocal
# from database import SessionLocal

import app.models as models
# import models

m_id = 1
db: Session = SessionLocal()

# Original methods from DL Notebook
# method run input as size (4097x??)
def testAug1(testSet):
    Divide1 = []
    for m in range(5):
        for k in range(4):
            divide1 = testSet.iloc[:, m][k * 1024:1024 + k * 1024].values.flatten()
            Divide1.append(divide1)
    CUT1 = []
    for r in Divide1:
        for p in range(3):
            chop1 = r[p * 256:512 + 256 * p]
            CUT1.append(chop1)
    return pd.DataFrame(tf.keras.utils.normalize(np.array(CUT1).T))

def load_huy_model(exists: bool, input, p_id: int, s_id: int):
    print("---------DL---------")
    if exists:
        print("Returning last result")
        print(input)
        return input  # crud.get_last_result(db, p_id, m_id)
    else:
        print("Running Model")
        testInstance = testAug1(input)
        Xtest = np.expand_dims(testInstance.T, 2)
        model = tf.keras.models.load_model("./")
        res = model.predict(Xtest)
        res = np.where(res > 0.5, 1, 0).flatten()
        counts = np.bincount(res)
        final = np.argmax(counts)

        res = models.Result(session_id=int(s_id), patient_id=int(p_id), result=int(final), model_id=m_id)
        return crud.create_patient_result(db, res)

# Get table from database ##
def get_session_data(p_id: int):
    session = crud.get_latest_session_table_by_id(db, p_id)
    s_id = session.split('_')[3]
    exists = crud.is_result_present(db, p_id, s_id, m_id)
    print("Getting Session Data")
    print("Exists: ", exists)
    if exists:
        return exists, [], s_id
    else:  # print(crud.get_table_data(db, latest_session))
        return exists, crud.get_table_data(db, session), s_id

# Picking only 4096 rows from data to insert into DL Model
def edit_data(p_id: int):
    exists, raw_session_table, s_id = get_session_data(p_id)
    print("----Editing Data----")
    if exists:
        print("Returning last result")
        return True, crud.get_last_result(db, p_id, m_id), p_id, s_id
    else:
        # Read column names from file
        data = []
        count = 0
        for row in raw_session_table:
            row_without_timestamp = list()
            for i in range(len(row)):
                if i > 0:
                    row_without_timestamp.append(row[i])
            if count <= 4096:
                data.append(row_without_timestamp)
                count += 1
                continue
            break

        data = (np.array(data))
        data = pd.DataFrame(data)
        print(data)
        print(data.shape)

        print("Returning edited data")
        return False, data, p_id, s_id

def start(pat_id: int):
    exists, data, p_id, s_id = edit_data(pat_id)
    return load_huy_model(exists, data, p_id, s_id)

if __name__ == "__main__":
    # Test Run
    exists, data, p_id, s_id = edit_data(17)
    # load_huy_model(edit_data(17))

max_stars_repo_name: AayushKucheria/digital-health

max_issues_repo_name: AayushKucheria/digital-health

max_forks_repo_name: AayushKucheria/digital-health


text:
# --------------
# Import packages
import numpy as np
import pandas as pd
from scipy.stats import mode

bank = pd.read_csv(path)

# Let's check which variable is categorical and which one is numerical so that you will get a basic idea about the features of the bank dataset.
categorical_var = bank.select_dtypes(include = 'object')
print(categorical_var)

numerical_var = bank.select_dtypes(include = 'number')
print(numerical_var)

# --------------
## Filling the missing values
from scipy.stats import mode
banks = bank.drop(columns = ['Loan_ID'])

print(banks.isnull().sum())

bank_mode = banks.mode()
print(bank_mode)

for x in banks.columns.values:
  banks[x]=banks[x].fillna(value=bank_mode[x].iloc[0])

print(banks.isnull().sum())

# --------------
## For basic idea of the average loan amount of a person - Loan Amount vs Gender

avg_loan_amount = pd.pivot_table(banks,index=['Gender','Married','Self_Employed'] , values='LoanAmount', aggfunc = np.mean)

print(avg_loan_amount)

# --------------
# The percentage of loan approved based on a person's employment type.
# Loan Approval vs Employment

loan_approved_se = banks[(banks['Self_Employed']=='Yes') & (banks['Loan_Status']=='Y')].shape[0]
loan_approved_nse = banks[(banks['Self_Employed']=='No') & (banks['Loan_Status']=='Y')].shape[0]

percentage_se = (loan_approved_se / 614) * 100
percentage_nse = (loan_approved_nse / 614) * 100

# --------------
# Transform the loan tenure from months to years
# company wants to find out those applicants with long loan amount term

loan_term = banks['Loan_Amount_Term'].apply(lambda m : m / 12)

big_loan_term = loan_term[loan_term>=25].shape[0]

# --------------
# Income/ Credit History vs Loan Amount
# the average income of an applicant and the average loan given to a person based on their income

loan_groupby = banks.groupby('Loan_Status')

loan_groupby = loan_groupby[['ApplicantIncome', 'Credit_History']]

mean_values = loan_groupby.mean()

max_stars_repo_name: daizybhadresha/ga-learner-dsmp-repo

max_issues_repo_name: daizybhadresha/ga-learner-dsmp-repo

max_forks_repo_name: daizybhadresha/ga-learner-dsmp-repo

reason: import numpy,from scipy


text:
import tensorflow as tf
from feeder import VarFeeder
import os
import argparse
import pandas as pd
import numpy
import logging
import yaml
import datetime
import math
import numpy as np
from typing import Tuple
import pickle

log = logging.getLogger()

SHUFFLE_BUFFER = 100

def __data_parser(serialized_example):

    features = tf.parse_single_example(serialized_example,
                                       features={'uph': tf.FixedLenFeature([], tf.string),
                                                 'price_cat_0_n': tf.FixedLenFeature([], tf.float32),
                                                 'price_cat_1_n': tf.FixedLenFeature([], tf.float32),
                                                 'price_cat_2_n': tf.FixedLenFeature([], tf.float32),
                                                 'price_cat_3_n': tf.FixedLenFeature([], tf.float32),
                                                 'hour': tf.FixedLenFeature([], tf.int64),
                                                 'g_g_m_n': tf.FixedLenFeature([], tf.float32),
                                                 'g_g_f_n': tf.FixedLenFeature([], tf.float32),
                                                 'g_g_x_n': tf.FixedLenFeature([], tf.float32),
                                                 'a_1_n': tf.FixedLenFeature([], tf.float32),
                                                 'a_2_n': tf.FixedLenFeature([], tf.float32),
                                                 'a_3_n': tf.FixedLenFeature([], tf.float32),
                                                 'a_4_n': tf.FixedLenFeature([], tf.float32),
                                                 't_3G_n': tf.FixedLenFeature([], tf.float32),
                                                 't_4G_n': tf.FixedLenFeature([], tf.float32),
                                                 'si_1_n': tf.FixedLenFeature([], tf.float32),
                                                 'si_2_n': tf.FixedLenFeature([], tf.float32),
                                                 'si_3_n': tf.FixedLenFeature([], tf.float32),
                                                 'dow_sin': tf.FixedLenSequenceFeature([], tf.float32, allow_missing=True),
                                                 'dow_cos': tf.FixedLenSequenceFeature([], tf.float32, allow_missing=True),
                                                 'ts': tf.FixedLenSequenceFeature([], tf.float32, allow_missing=True)
                                                 })

    uckey = tf.cast(features['uph'], tf.string)
    price_cat_0 = tf.cast(features['price_cat_0_n'], tf.float32)
    price_cat_1 = tf.cast(features['price_cat_1_n'], tf.float32)
    price_cat_2 = tf.cast(features['price_cat_2_n'], tf.float32)
    price_cat_3 = tf.cast(features['price_cat_3_n'], tf.float32)
    hour = tf.cast(features['hour'], tf.float32)
    gender_x = tf.cast(features['g_g_x_n'], tf.float32)
    gender_f = tf.cast(features['g_g_f_n'], tf.float32)
    gender_m = tf.cast(features['g_g_m_n'], tf.float32)
    age_1 = tf.cast(features['a_1_n'], tf.float32)
    age_2 = tf.cast(features['a_2_n'], tf.float32)
    age_3 = tf.cast(features['a_3_n'], tf.float32)
    age_4 = tf.cast(features['a_4_n'], tf.float32)
    t_3G = tf.cast(features['t_3G_n'], tf.float32)
    t_4G = tf.cast(features['t_4G_n'], tf.float32)
    si_1 = tf.cast(features['si_1_n'], tf.float32)
    si_2 = tf.cast(features['si_2_n'], tf.float32)
    si_3 = tf.cast(features['si_3_n'], tf.float32)
    hits = tf.cast(features['ts'], tf.float32)
    dow_sin = tf.cast(features['dow_sin'], tf.float32)
    dow_cos = tf.cast(features['dow_cos'], tf.float32)
    hits = tf.cast(features['ts'], tf.float32)

    return hour, gender_f, gender_m, gender_x, hits, price_cat_0, price_cat_1, price_cat_2, \
        price_cat_3, age_1, age_2, age_3, age_4, t_3G, \
        t_4G, si_1, si_2, si_3, uckey, dow_cos, dow_sin

def run(cfg):

    with open(cfg['tf_statistics_path'], 'rb') as f:
        tf_stat = pickle.load(f)

    names = []
    tfrecord_location = cfg['tfrecords_local_path']
    for file in os.listdir(tfrecord_location):
        if file.startswith("part"):
            names.append(file)
    file_paths = [os.path.join(tfrecord_location, name) for name in names]

    # read and make the dataset from tfrecord
    dataset = tf.data.TFRecordDataset(file_paths)
    dataset = dataset.map(__data_parser)

    batch_size = cfg['batch_size']
    duration = cfg['duration']

    dataset = dataset.batch(batch_size).shuffle(SHUFFLE_BUFFER)
    iterator = dataset.make_one_shot_iterator()
    next_el = iterator.get_next()

    holidays = numpy.ones((duration,), dtype=float)

    lagged_ix = numpy.ones((duration, 4), dtype=float)
    quarter_autocorr = numpy.ones((batch_size,), dtype=float)
    page_popularity = numpy.ones((batch_size,), dtype=float)

    date_list = tf_stat['days']

    dow =[(a,b) for a,b in zip(tf_stat['dow_sin'],tf_stat['dow_cos'])]

    with tf.Session() as sess:
        x = sess.run(next_el)
        l1 = x[19][0]
        l2 = x[20][0]
        m = [[v] for v in l1]
        [m[i].append(l2[i]) for i in range(0, len(l1))]
        page_indx = list(x[18])

        tensors = dict(
            hits=pd.DataFrame(x[4], index=page_indx, columns=date_list),
            lagged_ix=lagged_ix,
            page_ix=page_indx,
            pf_age=pd.DataFrame(x[9:13], columns=page_indx).T,
            pf_si=pd.DataFrame(x[15:18], columns=page_indx).T,
            pf_network=pd.DataFrame(x[13:15], columns=page_indx).T,
            pf_price_cat=pd.DataFrame(x[5:9], columns=page_indx).T,
            pf_gender=pd.DataFrame(x[1:4], columns=page_indx).T,
            page_popularity=page_popularity,
            quarter_autocorr=quarter_autocorr,
            dow=dow,
            holiday=holidays)

        data_len = tensors['hits'].shape[1]
        plain = dict(
            data_days=data_len - cfg['add_days'],
            features_days=data_len,
            data_start=date_list[0],
            data_end=date_list[-1],
            features_end=date_list[-1],
            n_pages=batch_size
        )
        VarFeeder(cfg['data_dir'], tensors, plain)

if __name__ == '__main__':

    # This is the format for the shell arguments
    # python trainer/tfrecord_reader.py --add_days=63 factdata-12192019.tfrecord trainer/data/vars

    # parser = argparse.ArgumentParser(description='Prepare data')
    # parser.add_argument('config_file')
    # parser.add_argument('data_dir')
    # parser.add_argument('--valid_threshold', default=0.0, type=float, help="Series minimal length threshold (pct of data length)")
    # parser.add_argument('--add_days', default=64, type=int, help="Add N days in a future for prediction")
    # parser.add_argument('--start', help="Effective start date. Data before the start is dropped")
    # parser.add_argument('--end', help="Effective end date. Data past the end is dropped")
    # parser.add_argument('--corr_backoffset', default=0, type=int, help='Offset for correlation calculation')
    # args = parser.parse_args()

    parser = argparse.ArgumentParser(description='Prepare data')
    parser.add_argument('config_file')
    args = parser.parse_args()

    with open(args.config_file, 'r') as ymlfile:
        cfg = yaml.load(ymlfile)

    cfg = cfg['tfrecorder_reader']

    run(cfg)


text:
import numpy as np

from pyautofinance.common.metrics.metric import Metric

class MedianProfit(Metric):
    name = 'MedianProfit'

    def __init__(self, profits):
        self.profits = profits
        super().__init__()

    def _get_metric_value(self):
        array = np.array(self.profits)
        return np.median(array)

    def __gt__(self, other):
        return self.value > other.value


text:
import os
import time
from collections import defaultdict
import logging

import pandas as pd
import numpy as np

from joblib import load, dump
from tqdm import tqdm

from models.primitive_models import PrimitiveModels
from models.ease import EASE

from utils.helpers import get_movies_by_ids, get_user_profiles, get_control_items
from utils.metrics import serendipity
from utils.logging import init_logging
from utils.parser import parse_args
from utils.helpers import get_movies_by_profile, get_movies_by_ids

np.random.seed(42)
args = parse_args()

init_logging('ease_gs.log', path=args.save_path)

def get_recommendations():
    ease = EASE()
    ease.fit(ratings, implicit=False, lambda_=param)

    pred = ease.predict(ratings, list(control_items.keys()), np.unique(ratings['movieId']), 10)

    recommendations_ = defaultdict(list)

    for user_id, df in pred.groupby('userId'):
        recommendations_[user_id].extend(df['movieId'].values)

    return recommendations_

def evaluate():
    hit = 0  # used for recall calculation
    total_recommendations = 0
    all_recommendations = []  # used for coverage calculation
    serendipity_results = []

    recommendations = get_recommendations()
    users = control_items.keys()

    for user_id in tqdm(users, total=len(users)):
        user_profile = user_embeddings[user_id]
        prediction = np.array(recommendations[user_id])
        primitive_predictions = primitive_recommendations[user_id]

        if prediction is None or prediction.ndim == 0:
            continue

        if control_items[user_id] in prediction:  # if prediction contains control item increase hit counter
            hit += 1

        all_recommendations.extend(list(prediction))
        total_recommendations += 1

        recommended_items = get_movies_by_ids(movies, prediction)

        if len(user_profile) == 0:  # if user profile is empty
            continue

        primitive_recs = np.array(list(primitive_recommendations[user_id].values()))

        rel.append(relevance(recommended_items, user_embeddings[user_id]))
        unexp.append(unexpectedness(prediction, primitive_recs))

    if total_recommendations > 0:
        recall = hit / total_recommendations
    else:
        recall = 0

    coverage = np.unique(all_recommendations).shape[0] / np.unique(np.unique(ratings['movieId'])).shape[0]
    return recall, coverage, np.nanmean(serendipity_results)

# Datasets
DATASET = '1m'
DATA_PATH = 'data/movielens/' + DATASET + '/clean/'

ratings = pd.read_csv(DATA_PATH + 'ratings.csv')
movies = pd.read_csv(os.path.join(DATA_PATH, 'movies.csv'))
item_ratings = load(DATA_PATH + 'item_sum_dif_rating.pickle')

train_data = pd.read_csv(DATA_PATH + 'train_data.csv', index_col='userId')
test_df = pd.read_csv(DATA_PATH + 'test_data.csv', index_col='userId')
test_data, _ = get_control_items(ratings, user_profiles=test_df)
ratings, control_items = get_control_items(ratings, user_ids=test_df.index.values)

# user_embeddings = get_user_profiles(ratings)

user_embeddings = {}
for user_id, user_profile in test_data.iterrows():
    user_embeddings[user_id] = get_movies_by_profile(movies, user_profile)

n = 10  # number of recommendations for each user

primitive_recommendations = None

if not os.path.isfile('cache/primitive_recommendations.joblib'):
    primitive_models = PrimitiveModels(train_data, ratings, item_ratings)
    primitive_recommendations = primitive_models.make_recommendations(test_data, n)
    dump(primitive_recommendations, 'cache/primitive_recommendations.joblib')
else:
    primitive_recommendations = load('cache/primitive_recommendations.joblib')

params = [1.0, 10.0, 50.0, 100.0, 250.0, 500.0, 750.0, 1000.0, 2000.0, 5000.0, 10000.0]
param = None
history = []

if __name__ == '__main__':
    for p in params:
        param = p
        start_time = time.time()
        logging.debug('Started: lambda={}'.format(p))
        recall, coverage, ser = evaluate()
        logging.debug('Recall: {}, coverage: {}, serendipity: {}'.format(recall, coverage, ser))
        logging.debug('Finished: lambda={}'.format(p))
        time_elapsed = time.time() - start_time
        history.append((p, recall, coverage, ser, time_elapsed))

dump(history, os.path.join(args.save_path, 'ease.joblib'))


text:
"""INDIVIDUAL ET PREPROCESSING AND ANALYSIS OF CHILD EXPERIMENT"""

import os
import itertools
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib import cm

def convert_gaze_coords(datafile):
    datafile['BPOGX']=(datafile.loc[:,'BPOGX']-0.5)*1920
    datafile['BPOGY']=-1*(datafile.loc[:,'BPOGY']-0.5)*1080
    datafile['FPOGX']=(datafile.loc[:,'FPOGX']-0.5)*1920
    datafile['FPOGY']=-1*(datafile.loc[:,'FPOGY']-0.5)*1080
    datafile['RPOGX']=(datafile.loc[:,'RPOGX']-0.5)*1920
    datafile['RPOGY']=-1*(datafile.loc[:,'RPOGY']-0.5)*1080
    datafile['LPOGX']=(datafile.loc[:,'LPOGX']-0.5)*1920
    datafile['LPOGY']=-1*(datafile.loc[:,'LPOGY']-0.5)*1080

    return datafile

file=r'C:/Users/presi/Desktop/PhD/Children experiment/Tasks/Task1/ET_Data/L_513_2021_Nov_14_1513.tsv'
IMG_PATH=r'C:/Users/presi/Desktop/PhD/Children experiment/Tasks/Images'

df=pd.read_csv(file,sep="\t")

df=convert_gaze_coords(df)

df2=df[df['BPOGV']==1]

for j,k in itertools.product(os.listdir(IMG_PATH),range(0,4)):
    dataset=df2.query("USER.str.startswith('"+str(j).split('.')[0]+".JPG_HIGH_"+str(k)+"_START').values")
    x=dataset.groupby('FPOGID')['FPOGX'].mean()
    y=dataset.groupby('FPOGID')['FPOGY'].mean()
    fix_dur=dataset.groupby('FPOGID')['FPOGD'].max()
    fig, ax = plt.subplots(figsize=(20, 11))
    ax.scatter(x,
               y,
               zorder=1,
               marker='o',
               s=fix_dur*10000,
               color='lime',
               alpha=0.5)
    ax.plot(x,
            y,
            '-o',
            linewidth=3,
            color='blue')
    img = plt.imread(IMG_PATH+ "/" + str(j))
    plt.imshow(img,
               zorder=0,
               extent=[-960, 960, -540, 540],
               aspect='auto')
    for i in range(len(fix_dur)):
        ax.annotate(str(i+1),
                    xy=(fix_dur.iloc[i],
                        fix_dur.iloc[i]),
                    xytext=(x.iloc[i],
                            y.iloc[i]),
                    fontsize=25,
                    color='black',
                    ha='center',
                    va='center')
    plt.xlabel('X coordinates (in pixels)', size=20)
    plt.ylabel('Y coordinates (in pixels)', size=20)
    plt.title('Scanpath for Subject 13 / Image '+str(j)+' / Block '+str(k+1), size=30)
    plt.show()