Surveys¶

def add_iso_coords(
    data: pd.DataFrame,
    val_range: tuple[int, int] = (1, 5),
    names: tuple[str, str] = ("ISOPleasant", "ISOEventful"),
    angles: tuple[int, ...] = EQUAL_ANGLES,
    *,
    overwrite: bool = False,
) -> pd.DataFrame:
    """
    Calculate and add ISO coordinates as new columns in the DataFrame.

    Parameters
    ----------
    data
        Input DataFrame containing PAQ data
    val_range
        (min, max) range of original PAQ responses, by default (1, 5)
    names
        Names for new coordinate columns, by default ("ISOPleasant", "ISOEventful")
    angles
        Angles for each PAQ in degrees, by default EQUAL_ANGLES
    overwrite
        Whether to overwrite existing ISO coordinate columns, by default False

    Returns
    -------
    :
        DataFrame with new ISO coordinate columns added

    Raises
    ------
    Warning
        If ISO coordinate columns already exist and overwrite is False

    Examples
    --------
    >>> import pandas as pd
    >>> df = pd.DataFrame({
    ...     'PAQ1': [4, 2], 'PAQ2': [3, 5], 'PAQ3': [2, 4], 'PAQ4': [1, 3],
    ...     'PAQ5': [5, 1], 'PAQ6': [3, 2], 'PAQ7': [4, 3], 'PAQ8': [2, 5]
    ... })
    >>> df_with_iso = add_iso_coords(df)
    >>> df_with_iso[['ISOPleasant', 'ISOEventful']].round(2)
       ISOPleasant  ISOEventful
    0        -0.03        -0.28
    1         0.47         0.18

    """
    for name in names:
        if name in data.columns:
            if overwrite:
                data = data.drop(name, axis=1)
            else:
                msg = (
                    f"{name} already in dataframe. Use `overwrite=True` to replace it."
                )
                raise Warning(msg)

    iso_pleasant, iso_eventful = calculate_iso_coords(
        data, val_range=val_range, angles=angles
    )
    data = data.assign(**{names[0]: iso_pleasant, names[1]: iso_eventful})

    logger.info(f"Added ISO coordinates to DataFrame with column names: {names}")
    return data

def calculate_iso_coords(
    results_df: pd.DataFrame,
    val_range: tuple[int, int] = (5, 1),
    angles: tuple[int, ...] = EQUAL_ANGLES,
) -> tuple[pd.Series, pd.Series]:
    """
    Calculate the projected ISOPleasant and ISOEventful coordinates.

    Parameters
    ----------
    results_df
        DataFrame containing PAQ data.
    val_range
        (max, min) range of original PAQ responses, by default (5, 1)
    angles
        Angles for each PAQ in degrees, by default EQUAL_ANGLES

    Returns
    -------
    :
        ISOPleasant and ISOEventful coordinate values

    Examples
    --------
    >>> import pandas as pd
    >>> df = pd.DataFrame({
    ...     'PAQ1': [4, 2], 'PAQ2': [3, 5], 'PAQ3': [2, 4], 'PAQ4': [1, 3],
    ...     'PAQ5': [5, 1], 'PAQ6': [3, 2], 'PAQ7': [4, 3], 'PAQ8': [2, 5]
    ... })
    >>> iso_pleasant, iso_eventful = calculate_iso_coords(df)
    >>> iso_pleasant.round(2)
    0   -0.03
    1    0.47
    dtype: float64
    >>> iso_eventful.round(2)
    0   -0.28
    1    0.18
    dtype: float64

    """
    scale = max(val_range) - min(val_range)

    paq_df = return_paqs(results_df, incl_ids=False)

    iso_pleasant = paq_df.apply(lambda row: _adj_iso_pl(row, angles, scale), axis=1)
    iso_eventful = paq_df.apply(lambda row: _adj_iso_ev(row, angles, scale), axis=1)

    logger.info(f"Calculated ISO coordinates for {len(results_df)} samples")
    return iso_pleasant, iso_eventful

def ipsatize(
    data: pd.DataFrame,
    method: Literal["grand_mean", "column_wise", "row_wise"] = "grand_mean",
    participant_col: str = "participant",
    scales: list[str] | None = None,
) -> pd.DataFrame:
    """
    Participant-level ipsatization for circumplex analysis.

    Removes systematic response biases before computing a correlation matrix.
    The choice of method depends on the study design and the type of bias
    being corrected.

    Parameters
    ----------
    data
        DataFrame containing PAQ scale columns and (for participant-level
        methods) a grouping column.
    method
        Centering strategy:

        ``"grand_mean"`` *(default)* — one scalar per participant: the mean
        across *all* PAQ values and *all* observations for that participant.
        Removes overall response-level differences between participants.
        **Matches the published SATP analysis (Aletta et al., 2024) and the
        original R implementation.**

        ``"column_wise"`` — eight scalars per participant: the per-scale mean
        across that participant's observations.  Removes scale-specific
        response biases.  This is the behaviour of the legacy
        :func:`person_center` function.

        ``"row_wise"`` — one scalar per observation: the mean across all PAQ
        scales within that observation.  Removes the general impression of
        each individual soundscape stimulus.  Equivalent to
        ``circumplex.ipsatize()``.
    participant_col
        Column used to group observations by participant.  Required for
        ``"grand_mean"`` and ``"column_wise"``; ignored for ``"row_wise"``.
    scales
        PAQ column names to centre.  Defaults to :data:`PAQ_IDS` when
        ``None``.

    Returns
    -------
    :
        DataFrame containing only the scale columns with centred values.
        The ``participant_col`` grouping column is excluded from the result.

    Raises
    ------
    KeyError
        If ``participant_col`` is not present in ``data`` when
        ``method`` is ``"grand_mean"`` or ``"column_wise"``.

    Examples
    --------
    >>> import pandas as pd
    >>> data = pd.DataFrame({
    ...     'PAQ1': [50., 60., 40., 30.], 'PAQ2': [50., 60., 40., 30.],
    ...     'PAQ3': [50., 60., 40., 30.], 'PAQ4': [50., 60., 40., 30.],
    ...     'PAQ5': [50., 60., 40., 30.], 'PAQ6': [50., 60., 40., 30.],
    ...     'PAQ7': [50., 60., 40., 30.], 'PAQ8': [50., 60., 40., 30.],
    ...     'participant': ['A', 'A', 'B', 'B'],
    ... })
    >>> result = ipsatize(data, method="grand_mean")
    >>> result['PAQ1'].tolist()
    [-5.0, 5.0, 5.0, -5.0]

    """
    _scales = scales if scales is not None else PAQ_IDS

    if method == "column_wise":
        means = data.groupby(participant_col)[_scales].transform("mean")
        return data[_scales] - means

    if method == "grand_mean":
        # Compute a single scalar per participant: mean across all PAQ values
        # and all observations for that participant.  Use nanmean so that
        # participants with partial NaN data still get a valid grand mean
        # computed from their non-NaN values; NaN rows are then removed by
        # downstream listwise deletion rather than silently expanding data loss
        # to the whole participant.
        grand_means = data.groupby(participant_col)[_scales].apply(
            lambda df: float(np.nanmean(df.values))
        )
        grand_mean_per_row = data[participant_col].map(grand_means)
        return data[_scales].subtract(grand_mean_per_row, axis=0)

    if method == "row_wise":
        row_means = data[_scales].mean(axis=1)
        return data[_scales].sub(row_means, axis=0)

    msg = f"method must be 'grand_mean', 'column_wise', or 'row_wise'; got {method!r}"
    raise ValueError(msg)

def return_paqs(
    df: pd.DataFrame, other_cols: list[str] | None = None, *, incl_ids: bool = True
) -> pd.DataFrame:
    """
    Return only the PAQ columns from a DataFrame.

    Parameters
    ----------
    df
        Input DataFrame containing PAQ data.
    other_cols
        Other columns to include in the output, by default None.
    incl_ids
        Whether to include ID columns (RecordID, GroupID, etc.), by default True.

    Returns
    -------
    :
        DataFrame containing only the PAQ columns and optionally ID and other specified
        columns.

    Examples
    --------
    >>> import pandas as pd
    >>> df = pd.DataFrame({
    ...     'RecordID': [1, 2],
    ...     'PAQ1': [4, 3],
    ...     'PAQ2': [2, 5],
    ...     'PAQ3': [1, 2],
    ...     'PAQ4': [3, 4],
    ...     'PAQ5': [5, 1],
    ...     'PAQ6': [2, 3],
    ...     'PAQ7': [4, 5],
    ...     'PAQ8': [1, 2],
    ...     'OtherCol': ['A', 'B']
    ... })
    >>> return_paqs(df)
       RecordID  PAQ1  PAQ2  PAQ3  PAQ4  PAQ5  PAQ6  PAQ7  PAQ8
    0         1     4     2     1     3     5     2     4     1
    1         2     3     5     2     4     1     3     5     2
    >>> return_paqs(df, incl_ids=False, other_cols=['OtherCol'])
       PAQ1  PAQ2  PAQ3  PAQ4  PAQ5  PAQ6  PAQ7  PAQ8 OtherCol
    0     4     2     1     3     5     2     4     1        A
    1     3     5     2     4     1     3     5     2        B

    """
    cols = PAQ_IDS.copy()

    if incl_ids:
        id_cols = [
            name
            for name in ["RecordID", "GroupID", "SessionID", "LocationID"]
            if name in df.columns
        ]
        cols = id_cols + cols

    if other_cols:
        cols.extend(other_cols)

    logger.debug(f"Returning PAQ columns: {cols}")
    return df[cols]

def simulation(
    n: int = 3000,
    val_range: tuple[int, int] = (1, 5),
    *,
    seed: int | None = None,
    incl_iso_coords: bool = False,
    **coord_kwargs: Unpack[_AddISOCoordsKwargs],
) -> pd.DataFrame:
    """
    Generate random PAQ responses for simulation purposes.

    Parameters
    ----------
    n
        Number of samples to simulate, by default 3000
    val_range
        Range of values for PAQ responses, by default (1, 5)
    seed
        Optional random seed for deterministic output, by default None
    incl_iso_coords
        Whether to add calculated ISO coordinates, by default False
    **coord_kwargs
        Optional keyword arguments passed directly to the `add_iso_coords` function
        if `incl_iso_coords` is True. These can include:

        - `names` (tuple[str, str]): Names for the new ISO coordinate columns.
        - `angles` (tuple[int, ...]): Angles for each PAQ used in calculation.
        - `overwrite` (bool): Whether to overwrite existing ISO coordinate columns.

    Returns
    -------
    :
        DataFrame of randomly generated PAQ responses

    Examples
    --------
    >>> data = simulation(n=5,incl_iso_coords=True)
    >>> data.shape
    (5, 10)
    >>> list(data.columns)
    ['PAQ1', 'PAQ2', 'PAQ3', 'PAQ4', 'PAQ5', 'PAQ6', 'PAQ7', 'PAQ8', 'ISOPleasant', 'ISOEventful']

    """  # noqa: E501
    data = pd.DataFrame(
        np.random.default_rng(seed).integers(
            min(val_range), max(val_range) + 1, size=(n, 8)
        ),
        columns=PAQ_IDS,
    )

    if incl_iso_coords:
        data = add_iso_coords(data, val_range=val_range, **coord_kwargs)

    logger.info(f"Generated simulated PAQ data with {n} samples")
    return data

def rename_paqs(
    df: pd.DataFrame, paq_aliases: list | tuple | dict | None = None
) -> pd.DataFrame:
    """
    Rename the PAQ columns in a DataFrame to standard PAQ IDs.

    Parameters
    ----------
    df
        Input DataFrame containing PAQ data.
    paq_aliases
        Specify which PAQs are to be renamed. If None, will check if the column names
        are in pre-defined options. If a tuple, the order must match PAQ_IDS.
        If a dict, keys are current names and values are desired PAQ IDs.

    Returns
    -------
    :
        DataFrame with renamed PAQ columns.

    Raises
    ------
    ValueError
        If paq_aliases is not a tuple, list, or dictionary.

    Examples
    --------
    >>> import pandas as pd
    >>> df = pd.DataFrame({
    ...     'pleasant': [4, 3],
    ...     'vibrant': [2, 5],
    ...     'other_col': [1, 2]
    ... })
    >>> rename_paqs(df)
       PAQ1  PAQ2  other_col
    0     4     2          1
    1     3     5          2
    >>> df_custom = pd.DataFrame({
    ...     'pl': [4, 3],
    ...     'vb': [2, 5],
    ... })
    >>> rename_paqs(df_custom, paq_aliases={'pl': 'PAQ1', 'vb': 'PAQ2'})
       PAQ1  PAQ2
    0     4     2
    1     3     5

    """
    if paq_aliases is None:
        if any(paq_id in df.columns for paq_id in PAQ_IDS):
            logger.info("PAQs already correctly named.")
            return df
        if any(paq_name in df.columns for paq_name in PAQ_LABELS):
            paq_aliases = PAQ_LABELS

    if isinstance(paq_aliases, list | tuple):
        rename_dict = dict(zip(paq_aliases, PAQ_IDS, strict=False))
    elif isinstance(paq_aliases, dict):
        rename_dict = paq_aliases
    else:
        msg = "paq_aliases must be a tuple, list, or dictionary."
        raise TypeError(msg)

    logger.debug(f"Renaming PAQs with the following mapping: {rename_dict}")
    return df.rename(columns=rename_dict)

def calculate_iso_coords(
    results_df: pd.DataFrame,
    val_range: tuple[int, int] = (5, 1),
    angles: tuple[int, ...] = EQUAL_ANGLES,
) -> tuple[pd.Series, pd.Series]:
    """
    Calculate the projected ISOPleasant and ISOEventful coordinates.

    Parameters
    ----------
    results_df
        DataFrame containing PAQ data.
    val_range
        (max, min) range of original PAQ responses, by default (5, 1)
    angles
        Angles for each PAQ in degrees, by default EQUAL_ANGLES

    Returns
    -------
    :
        ISOPleasant and ISOEventful coordinate values

    Examples
    --------
    >>> import pandas as pd
    >>> df = pd.DataFrame({
    ...     'PAQ1': [4, 2], 'PAQ2': [3, 5], 'PAQ3': [2, 4], 'PAQ4': [1, 3],
    ...     'PAQ5': [5, 1], 'PAQ6': [3, 2], 'PAQ7': [4, 3], 'PAQ8': [2, 5]
    ... })
    >>> iso_pleasant, iso_eventful = calculate_iso_coords(df)
    >>> iso_pleasant.round(2)
    0   -0.03
    1    0.47
    dtype: float64
    >>> iso_eventful.round(2)
    0   -0.28
    1    0.18
    dtype: float64

    """
    scale = max(val_range) - min(val_range)

    paq_df = return_paqs(results_df, incl_ids=False)

    iso_pleasant = paq_df.apply(lambda row: _adj_iso_pl(row, angles, scale), axis=1)
    iso_eventful = paq_df.apply(lambda row: _adj_iso_ev(row, angles, scale), axis=1)

    logger.info(f"Calculated ISO coordinates for {len(results_df)} samples")
    return iso_pleasant, iso_eventful

def add_iso_coords(
    data: pd.DataFrame,
    val_range: tuple[int, int] = (1, 5),
    names: tuple[str, str] = ("ISOPleasant", "ISOEventful"),
    angles: tuple[int, ...] = EQUAL_ANGLES,
    *,
    overwrite: bool = False,
) -> pd.DataFrame:
    """
    Calculate and add ISO coordinates as new columns in the DataFrame.

    Parameters
    ----------
    data
        Input DataFrame containing PAQ data
    val_range
        (min, max) range of original PAQ responses, by default (1, 5)
    names
        Names for new coordinate columns, by default ("ISOPleasant", "ISOEventful")
    angles
        Angles for each PAQ in degrees, by default EQUAL_ANGLES
    overwrite
        Whether to overwrite existing ISO coordinate columns, by default False

    Returns
    -------
    :
        DataFrame with new ISO coordinate columns added

    Raises
    ------
    Warning
        If ISO coordinate columns already exist and overwrite is False

    Examples
    --------
    >>> import pandas as pd
    >>> df = pd.DataFrame({
    ...     'PAQ1': [4, 2], 'PAQ2': [3, 5], 'PAQ3': [2, 4], 'PAQ4': [1, 3],
    ...     'PAQ5': [5, 1], 'PAQ6': [3, 2], 'PAQ7': [4, 3], 'PAQ8': [2, 5]
    ... })
    >>> df_with_iso = add_iso_coords(df)
    >>> df_with_iso[['ISOPleasant', 'ISOEventful']].round(2)
       ISOPleasant  ISOEventful
    0        -0.03        -0.28
    1         0.47         0.18

    """
    for name in names:
        if name in data.columns:
            if overwrite:
                data = data.drop(name, axis=1)
            else:
                msg = (
                    f"{name} already in dataframe. Use `overwrite=True` to replace it."
                )
                raise Warning(msg)

    iso_pleasant, iso_eventful = calculate_iso_coords(
        data, val_range=val_range, angles=angles
    )
    data = data.assign(**{names[0]: iso_pleasant, names[1]: iso_eventful})

    logger.info(f"Added ISO coordinates to DataFrame with column names: {names}")
    return data

def likert_data_quality(
    df: pd.DataFrame, val_range: tuple[int, int] = (1, 5), *, allow_na: bool = False
) -> list[int] | None:
    """
    Perform basic quality checks on PAQ (Likert scale) data.

    Parameters
    ----------
    df
        DataFrame containing PAQ data
    allow_na
        Whether to allow NaN values in PAQ data, by default False
    val_range
        Valid range for PAQ values, by default (1, 5)

    Returns
    -------
    :

    Examples
    --------
    >>> import pandas as pd
    >>> import numpy as np
    >>> df = pd.DataFrame({
    ...     'PAQ1': [np.nan, 2, 3, 3], 'PAQ2': [3, 2, 6, 3], 'PAQ3': [2, 2, 3, 3],
    ...     'PAQ4': [1, 2, 3, 3], 'PAQ5': [5, 2, 3, 3], 'PAQ6': [3, 2, 3, 3],
    ...     'PAQ7': [4, 2, 3, 3], 'PAQ8': [2, 2, 3, 3]
    ... })
    >>> likert_data_quality(df)
    [0, 1, 2]
    >>> likert_data_quality(df,allow_na=True)
    [1, 2]

    """
    paqs = return_paqs(df, incl_ids=False)
    invalid_indices = []

    for idx, row in paqs.iterrows():
        # Convert the index to int to ensure type compatibility
        row_idx = int(idx) if isinstance(idx, str) else idx
        row_array = row.to_numpy()
        is_constant = row_array.shape[0] > 0 and (row_array[0] == row_array).all()

        if (not allow_na and row.isna().any()) or (
            row.notna().all()
            and (
                row.min() < min(val_range)
                or row.max() > max(val_range)
                or (is_constant and row.iloc[0] != np.mean(val_range))
            )
        ):
            invalid_indices.append(row_idx)

    if invalid_indices:
        logger.info(f"Found {len(invalid_indices)} samples with data quality issues")
        return invalid_indices

    logger.info("PAQ data quality check passed")
    return None

def simulation(
    n: int = 3000,
    val_range: tuple[int, int] = (1, 5),
    *,
    seed: int | None = None,
    incl_iso_coords: bool = False,
    **coord_kwargs: Unpack[_AddISOCoordsKwargs],
) -> pd.DataFrame:
    """
    Generate random PAQ responses for simulation purposes.

    Parameters
    ----------
    n
        Number of samples to simulate, by default 3000
    val_range
        Range of values for PAQ responses, by default (1, 5)
    seed
        Optional random seed for deterministic output, by default None
    incl_iso_coords
        Whether to add calculated ISO coordinates, by default False
    **coord_kwargs
        Optional keyword arguments passed directly to the `add_iso_coords` function
        if `incl_iso_coords` is True. These can include:

        - `names` (tuple[str, str]): Names for the new ISO coordinate columns.
        - `angles` (tuple[int, ...]): Angles for each PAQ used in calculation.
        - `overwrite` (bool): Whether to overwrite existing ISO coordinate columns.

    Returns
    -------
    :
        DataFrame of randomly generated PAQ responses

    Examples
    --------
    >>> data = simulation(n=5,incl_iso_coords=True)
    >>> data.shape
    (5, 10)
    >>> list(data.columns)
    ['PAQ1', 'PAQ2', 'PAQ3', 'PAQ4', 'PAQ5', 'PAQ6', 'PAQ7', 'PAQ8', 'ISOPleasant', 'ISOEventful']

    """  # noqa: E501
    data = pd.DataFrame(
        np.random.default_rng(seed).integers(
            min(val_range), max(val_range) + 1, size=(n, 8)
        ),
        columns=PAQ_IDS,
    )

    if incl_iso_coords:
        data = add_iso_coords(data, val_range=val_range, **coord_kwargs)

    logger.info(f"Generated simulated PAQ data with {n} samples")
    return data