from __future__ import annotations

import numpy as np
import pandas as pd


def _interpolate_series(values: np.ndarray, timestamps: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    series = pd.Series(values.astype(np.float32), index=pd.Index(timestamps.astype(np.float32)))
    if series.isna().all():
        empty = np.zeros_like(values, dtype=np.float32)
        valid_mask = np.zeros_like(values, dtype=np.float32)
        return empty, valid_mask

    if not np.all(np.diff(series.index.to_numpy()) >= 0):
        raise ValueError("timestamps must be non-decreasing for PCHIP interpolation.")

    if pd.isna(series.iloc[0]):
        series.iloc[0] = series.dropna().iloc[0]
    if pd.isna(series.iloc[-1]):
        series.iloc[-1] = series.dropna().iloc[-1]

    out = series.interpolate(method="pchip", limit_direction="both")
    out = out.ffill().bfill().fillna(0.0)
    valid_mask = np.ones_like(values, dtype=np.float32)
    return out.to_numpy(dtype=np.float32), valid_mask


def interpolate_batch(
    X_with_nan: np.ndarray,
    timestamps: np.ndarray | None = None,
    return_valid_mask: bool = False,
) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
    X = np.asarray(X_with_nan, dtype=np.float32)
    if X.ndim != 3:
        raise ValueError(f"Expected 3D array, got {X.shape}.")
    if timestamps is None:
        timestamps = np.broadcast_to(
            np.arange(X.shape[1], dtype=np.float32)[None, :, None],
            X.shape,
        ).copy()
    else:
        timestamps = np.asarray(timestamps, dtype=np.float32)
        if timestamps.ndim == 2:
            if timestamps.shape != X.shape[:2]:
                raise ValueError(
                    "Expected timestamps with shape [n_samples, n_steps] when using 2D timestamps, "
                    f"got {timestamps.shape} for X shape {X.shape}."
                )
            timestamps = np.repeat(timestamps[:, :, None], X.shape[2], axis=2)
        elif timestamps.ndim == 3:
            if timestamps.shape[0] != X.shape[0] or timestamps.shape[1] != X.shape[1]:
                raise ValueError(
                    "Expected timestamps to match X on sample and step dimensions, "
                    f"got timestamps shape {timestamps.shape} for X shape {X.shape}."
                )
            if timestamps.shape[2] == 1:
                timestamps = np.repeat(timestamps, X.shape[2], axis=2)
            elif timestamps.shape[2] != X.shape[2]:
                raise ValueError(
                    "Expected timestamps feature dimension to be 1 or match X, "
                    f"got timestamps shape {timestamps.shape} for X shape {X.shape}."
                )
        else:
            raise ValueError(
                "Expected timestamps with 2 or 3 dimensions, "
                f"got {timestamps.shape}."
            )

    interpolated = np.zeros_like(X, dtype=np.float32)
    valid_mask = np.zeros_like(X, dtype=np.float32)
    for sample_idx in range(X.shape[0]):
        for feature_idx in range(X.shape[2]):
            series_interp, series_valid = _interpolate_series(
                X[sample_idx, :, feature_idx],
                timestamps[sample_idx, :, feature_idx],
            )
            interpolated[sample_idx, :, feature_idx] = series_interp
            valid_mask[sample_idx, :, feature_idx] = series_valid
    if return_valid_mask:
        return interpolated, valid_mask
    return interpolated
