import numpy as np
from numpy.typing import NDArray


def truncate_nans(data):
    """
    Truncate trailing time columns that contain NaN values.

    The function assumes that the last axis of the array corresponds to
    time and removes all columns from the first occurrence of a NaN
    (in any row) to the end of the array.

    This guarantees that the returned array contains **no NaN values**,
    even if this means discarding some valid data points that appear
    after a NaN in another row.

    Parameters
    ----------
    data : array_like, shape (N, T)
        Input array where rows represent independent signals (e.g. q-points)
        and the last axis represents time.

    Returns
    -------
    ndarray
        Array truncated along the time axis such that no NaN values remain.
        The returned array has the same number of rows as the input but may
        have fewer time columns.

    Notes
    -----
    Example::

        [[1.0, 0.8, 0.7, 0.6, 0.5, 0.3, nan],
         [1.0, 0.8, 0.7, 0.6, 0.5, nan, nan]]

    becomes::

        [[1.0, 0.8, 0.7, 0.6, 0.5],
         [1.0, 0.8, 0.7, 0.6, 0.5]]

    because the first NaN appears in column 5.
    """

    arr = np.asarray(data)
    assert arr.ndim == 2

    # Determine which time columns contain any NaN values
    column_has_nan = np.any(np.isnan(arr), axis=0)

    # Find the first column containing a NaN
    nan_indices = np.where(column_has_nan)[0]

    # No NaNs present, return array unchanged
    if len(nan_indices) == 0:
        return arr

    # NaNs were found, Keep only the column without NaNs
    first_nan_column = nan_indices[0]
    return arr[:, :first_nan_column]


class TimeAverager:
    """Naive special purpose averager class used in dynasor to collect and time-average arrays
    obtained from sliding time-window averaging.

    It assists with keeping track of how many data samples have been added to each slot.

    It will time-average arrays of shape ``(Nq, time_window)`` where ``Ǹq`` is the
    number of q-points and ``time_window`` is the size of the time window.

    Parameters
    ----------
    time_window
        Size of the time window in which the time-average happens.
    array_length
        Length of the array to be averaged for each time-lag, i.e., number of q-points.
    """

    def __init__(self, time_window: int, array_length: int):
        assert time_window >= 1
        self._time_window = time_window
        self._array_length = array_length

        self._counts = np.zeros(time_window, dtype=int)
        self._arrays = [np.zeros(array_length) for _ in range(time_window)]

    def add_sample(self, time_lag: int, sample: np.ndarray):
        assert len(sample) == self._array_length
        self._counts[time_lag] += 1
        self._arrays[time_lag] += sample

    def get_average_at_timelag(self, time_lag: int):
        if self._counts[time_lag] == 0:
            array = np.full((self._array_length, ), np.nan)
            return array
        return self._arrays[time_lag] / self._counts[time_lag]

    def get_average_all(self) -> NDArray[float]:
        """
        Returns an averaged array of shape ``(array_length, time_window)``.
        """
        return np.array([self.get_average_at_timelag(t) for t in range(self._time_window)]).T
