"""Internal support module.

Conceptually, this module provides the narrow NWB access layer needed by
session-aware analyses and PSTH-building workflows.

It exists as a separate unit so NWB table access and behavioral-timeseries
lookups stay isolated from generic coercion code and higher-level analysis
logic.

It connects NWB objects to the broader loading and session helper stack.
"""

from __future__ import annotations

from typing import Any, Iterable, List

import numpy as np


__all__ = [
    "extract_fiber_location_from_trials",
    "get_nwb_trials_dataframe",
    "load_nwb_behavior_timeseries_data",
    "load_nwb_behavior_timeseries_timestamps",
    "matlab_like_trial_columns",
]


def get_nwb_trials_dataframe(nwb: Any) -> Any:
    """Return the NWB trials table as a pandas-like DataFrame.

    Parameters
    ----------
    nwb:
        NWB object exposing either `trials` or `intervals["trials"]`.

    Returns
    -------
    Any
        Trials table converted to a DataFrame-like object.
    """

    try:
        return nwb.trials.to_dataframe()  # type: ignore[attr-defined]
    except Exception:
        # Some NWB exports expose trials through the generic intervals mapping
        # instead of the top-level `trials` attribute.
        return nwb.intervals["trials"].to_dataframe()  # type: ignore[index]



def matlab_like_trial_columns(
    columns: Iterable[str],
    drop_first: int = 1,
    drop_last: int = 3,
) -> List[str]:
    """Return a sliced view of trial columns using fixed boundary trimming.

    Parameters
    ----------
    columns:
        Ordered trial-column names.
    drop_first:
        Number of columns to drop from the left boundary.
    drop_last:
        Number of columns to drop from the right boundary.

    Returns
    -------
    list[str]
        Trimmed column list. If the input is too short for the requested trim,
        an empty list is returned.
    """

    cols = list(columns)
    min_required = drop_first + drop_last + 1
    if len(cols) < min_required:
        return []
    return cols[drop_first : len(cols) - drop_last]



def extract_fiber_location_from_trials(
    trials_df: Any,
    opto_area_column: str = "opto_area",
) -> Any:
    """Extract unique opto-area labels from a trials table.

    Parameters
    ----------
    trials_df:
        Trials DataFrame-like object.
    opto_area_column:
        Column name containing opto target labels.

    Returns
    -------
    str or np.ndarray
        Empty string when no valid labels exist, a single string for one valid
        label, or an object array when several distinct labels are present.
    """

    import pandas as pd  # local import to keep base module lightweight

    if opto_area_column not in trials_df.columns:
        return ""

    fiber_values: List[str] = []
    vals = pd.unique(trials_df[opto_area_column])
    # Normalize away empty strings and stringified NaNs so downstream plotting
    # code only sees meaningful target labels.
    for value in vals:
        if pd.isna(value):
            continue
        s = str(value).strip()
        if not s or s.lower() == "nan":
            continue
        fiber_values.append(s)

    if len(fiber_values) == 0:
        return ""
    if len(fiber_values) == 1:
        return fiber_values[0]
    return np.asarray(fiber_values, dtype=object)



def load_nwb_behavior_timeseries_data(
    nwb_obj: Any,
    interface: str,
    ts_name: str,
) -> np.ndarray:
    """Load one behavioral time-series data array from an NWB object.

    Parameters
    ----------
    nwb_obj:
        NWB object containing a `processing["behavior"]` group.
    interface:
        Name of the behavioral interface.
    ts_name:
        Time-series name inside the interface.

    Returns
    -------
    np.ndarray
        Raw time-series data converted to a NumPy array.
    """

    ts = nwb_obj.processing["behavior"].data_interfaces[interface].time_series[ts_name]
    return np.asarray(ts.data[:])



def load_nwb_behavior_timeseries_timestamps(
    nwb_obj: Any,
    interface: str,
    ts_name: str,
) -> np.ndarray:
    """Load one behavioral time-series timestamp vector in double precision.

    Parameters
    ----------
    nwb_obj:
        NWB object containing a `processing["behavior"]` group.
    interface:
        Name of the behavioral interface.
    ts_name:
        Time-series name inside the interface.

    Returns
    -------
    np.ndarray
        Timestamp vector as ``float64``. These are large absolute session
        timestamps (up to ~10^4 s); Code_M loads them as double via
        ``.timestamps.load``. Truncating to float32 here (~0.5 ms) would shift
        the nearest-frame trial alignment relative to Code_M, so keep float64.
    """

    ts = nwb_obj.processing["behavior"].data_interfaces[interface].time_series[ts_name]
    return np.asarray(ts.timestamps[:], dtype=np.float64)