"""
A collection of wrappers around the the euxfel_bunch_pattern pkg

The euxfel_bunch_pattern package provides generic methods to extract
information from the bunch pattern tables. To ease its use from within
the toolbox some of its methods are wrapped. Like this they show up in
the users namespace in a self-explanatory way.
"""
import logging

import euxfel_bunch_pattern as ebp

__all__ = [
    'is_sase_3',
    'is_sase_1',
    'is_ppl',
    'is_pulse_at',
]

PPL_SCS = ebp.LASER_SEED6
log = logging.getLogger(__name__)


def _convert_data(bpt_dec):
    bpt_conv = bpt_dec

    if type(bpt_dec).__module__ == 'xarray.core.dataarray':
        bpt_conv = bpt_dec.where(bpt_dec.values == True, other=0)
    elif type(bpt_dec).__module__ == 'numpy':
        bpt_conv = bpt_dec.astype(int)
    else:
        dtype = type(bpt_dec).__module__
        log.warning(f"Could not convert data type {dtype}."
                    "Return raw euxfel_bp table.")

    return bpt_conv


def is_pulse_at(bpt, loc):
    """
    Check for prescence of a pulse at the location provided.

    Parameters
    ----------
    bpt : numpy array, xarray DataArray
        The bunch pattern data.
    loc : str
        The location where to check: {'sase1', 'sase3', 'scs_ppl'}

    Returns
    -------
    boolean : numpy array, xarray DataArray
      true if a pulse is present at *loc*.
    """
    if loc == 'sase3':
        bpt_dec = ebp.is_sase(bpt, 3)
    elif loc == 'sase1':
        bpt_dec = ebp.is_sase(bpt, 1)
    elif loc == 'scs_ppl':
        bpt_dec = ebp.is_laser(bpt, laser=PPL_SCS)
    else:
        raise ValueError(f'loc argument is {loc}, expected "sase1", ' +
                         '"sase3" or "scs_ppl"')

    return _convert_data(bpt_dec)


def is_sase_3(bpt):
    """
    Check for prescence of a SASE3 pulse.

    Parameters
    ----------
    bpt : numpy array, xarray DataArray
        The bunch pattern data.

    Returns
    -------
    boolean : numpy array, xarray DataArray
      true if SASE3 pulse is present.
    """
    bpt_dec = ebp.is_sase(bpt, 3)
    return _convert_data(bpt_dec)


def is_sase_1(bpt):
    """
    Check for prescence of a SASE1 pulse.

    Parameters
    ----------
    bpt : numpy array, xarray DataArray
        The bunch pattern data.

    Returns
    -------
    boolean : numpy array, xarray DataArray
      true if SASE1 pulse is present.
    """
    bpt_dec = ebp.is_sase(bpt, 1)
    return _convert_data(bpt_dec)


def is_ppl(bpt):
    """
    Check for prescence of pp-laser pulse.

    Parameters
    ----------
    bpt : numpy array, xarray DataArray
        The bunch pattern data.

    Returns
    -------
    boolean : numpy array, xarray DataArray
      true if pp-laser pulse is present.
    """
    bpt_dec = ebp.is_laser(bpt, laser=PPL_SCS)
    return _convert_data(bpt_dec)