"""Shared internal helpers used across modules."""

from __future__ import annotations

import numpy as np
import torch

try:
    from scipy.signal import savgol_filter
except ImportError:
    savgol_filter = None


def to_numpy(x) -> np.ndarray:
    """Convert tensor or array-like to numpy."""
    if isinstance(x, torch.Tensor):
        return x.detach().cpu().numpy()
    return np.asarray(x)


def clip01(a: np.ndarray) -> np.ndarray:
    """Clip array to [0, 1]."""
    return np.clip(a, 0.0, 1.0)


def gather_last(x: torch.Tensor, lengths: torch.Tensor, feat_idx: int) -> torch.Tensor:
    """Gather the last valid timestep for a given feature. x: (B, L, F), lengths: (B,)."""
    B = x.size(0)
    idx = (lengths.to(x.device) - 1).clamp(min=0)
    return x[torch.arange(B, device=x.device), idx, feat_idx]


def last_ndvi(input_seq: torch.Tensor, input_lengths: torch.Tensor) -> torch.Tensor:
    """Extract last NDVI value per sample. Returns (B, 1)."""
    return gather_last(input_seq, input_lengths, feat_idx=0).unsqueeze(1)


def last_time_and_step_raw(
    input_seq: torch.Tensor,
    input_lengths: torch.Tensor,
    fallback_dap_step: float,
    min_step: float,
    max_step: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return last RAW DAP (B, 1) and estimated step size (B, 1)."""
    device = input_seq.device
    B = input_seq.size(0)
    idx_last = (input_lengths.to(device) - 1).clamp(min=0)
    bidx = torch.arange(B, device=device)
    t_last = input_seq[bidx, idx_last, 1]

    has_prev = (input_lengths > 1).to(device)
    idx_prev = (idx_last - 1).clamp(min=0)
    t_prev = input_seq[bidx, idx_prev, 1]
    step_est = t_last - t_prev

    fb = torch.full_like(step_est, fallback_dap_step)
    step = torch.where(has_prev.bool(), step_est, fb)
    step = torch.clamp(step, min=min_step, max=max_step)
    return t_last.unsqueeze(1), step.unsqueeze(1)


def maybe_jitter_step(
    step: torch.Tensor,
    training: bool,
    step_jitter_std: float,
    min_step: float,
    max_step: float,
) -> torch.Tensor:
    """Optionally add noise to step size during training."""
    if training and step_jitter_std > 0.0:
        noise = torch.randn_like(step) * (step_jitter_std * step.abs().clamp_min(1e-6))
        step = torch.clamp(step + noise, min=min_step, max=max_step)
    return step


def smooth1d(y: np.ndarray, kind: str = "ma", w: int = 7, poly: int = 2) -> np.ndarray:
    """1-D smoothing: 'ma' (moving average) or 'sg' (Savitzky-Golay)."""
    y = np.asarray(y)
    if y.size < 3 or kind in (None, "none"):
        return y
    if kind == "ma":
        w = max(1, int(w))
        if w == 1:
            return y
        pad = w // 2
        ypad = np.pad(y, (pad, pad), mode="edge")
        return np.convolve(ypad, np.ones(w, dtype=float) / w, mode="valid")
    if kind == "sg" and savgol_filter is not None:
        w = int(w) if int(w) % 2 == 1 else int(w) + 1
        poly = min(int(poly), w - 1)
        return savgol_filter(y, window_length=max(3, w), polyorder=max(1, poly), mode="interp")
    return y


def safe_r2(y, yhat) -> float:
    """R-squared robust to NaN/Inf."""
    y = to_numpy(y).astype(np.float32).ravel()
    yhat = to_numpy(yhat).astype(np.float32).ravel()
    m = np.isfinite(y) & np.isfinite(yhat)
    if m.sum() < 2:
        return np.nan
    y, yhat = y[m], yhat[m]
    ss_res = float(np.sum((y - yhat) ** 2))
    ss_tot = float(np.sum((y - np.mean(y)) ** 2))
    if ss_tot == 0.0:
        return np.nan
    return 1.0 - ss_res / ss_tot


def safe_rmse(y, yhat) -> float:
    """RMSE robust to NaN/Inf."""
    y = to_numpy(y).astype(np.float32).ravel()
    yhat = to_numpy(yhat).astype(np.float32).ravel()
    m = np.isfinite(y) & np.isfinite(yhat)
    if m.sum() == 0:
        return np.nan
    return float(np.sqrt(np.mean((y[m] - yhat[m]) ** 2)))
