"""Pre/post fine-tuning vertical offset correction."""

from __future__ import annotations

import numpy as np
import torch


def pre_fine_tuning_offsetter(
    ndvi_array: np.ndarray,
    group_stats_array: np.ndarray,
    train_start_mean: float,
    *,
    early_window: int = 65,
    current_dap: int | None = None,
    ignore_zeros: bool = True,
    trim: tuple[float, float] = (0.10, 0.90),
    offset_cap: float = 0.25,
    min_valid: int = 5,
    ramp_full_at: int = 20,
    verbose: bool = False,
) -> tuple[np.ndarray, np.ndarray, float]:
    """Apply a vertical shift so test early-season mean matches training mean.

    Returns (ndvi_shifted, group_stats_shifted, offset_value).
    """
    ndvi = np.clip(np.asarray(ndvi_array, dtype=np.float32), 0.0, 1.0)
    gs = np.asarray(group_stats_array, dtype=np.float32)

    if ndvi.ndim != 2:
        raise ValueError(f"ndvi_array must be 2D, got {ndvi.shape}")
    N, T = ndvi.shape

    K_eff = int(min(int(early_window), T, int(current_dap) if current_dap is not None else early_window))
    if K_eff <= 0:
        return ndvi.copy(), gs.copy(), 0.0

    early = ndvi[:, :K_eff]

    def _trimmed_mean(x, lo=0.10, hi=0.90):
        if x.size == 0:
            return np.nan
        q_lo, q_hi = np.quantile(x, [lo, hi])
        sel = (x >= q_lo) & (x <= q_hi)
        return float(np.mean(x[sel])) if np.any(sel) else np.nan

    per_grid = np.full(N, np.nan, dtype=np.float32)
    for i in range(N):
        xi = early[i]
        m = np.isfinite(xi)
        if ignore_zeros:
            m &= xi != 0.0
        valid = xi[m]
        if valid.size < min_valid:
            continue
        tm = _trimmed_mean(valid, trim[0], trim[1])
        per_grid[i] = tm if np.isfinite(tm) else (float(valid.mean()) if valid.size else np.nan)

    finite_vals = per_grid[np.isfinite(per_grid)]
    if finite_vals.size == 0:
        pooled = early[np.isfinite(early)]
        if ignore_zeros:
            pooled = pooled[pooled != 0.0]
        tm = _trimmed_mean(pooled, trim[0], trim[1]) if pooled.size else np.nan
        center = float(tm) if np.isfinite(tm) else 0.0
    else:
        vals = finite_vals
        if vals.size >= 4:
            q1, q3 = np.percentile(vals, [25.0, 75.0])
            iqr = q3 - q1
            filt = vals[(vals >= q1 - 1.5 * iqr) & (vals <= q3 + 1.5 * iqr)]
            if filt.size >= max(3, min_valid):
                vals = filt
        center = float(np.median(vals))

    raw_offset = float(np.clip(center - train_start_mean, -offset_cap, offset_cap))
    shrink = min(1.0, max(0.0, K_eff / max(1, ramp_full_at)))
    offset_value = raw_offset * shrink

    ndvi_shifted = np.clip(ndvi - offset_value, 0.0, 1.0).astype(np.float32)
    gs_shifted = gs.astype(np.float32, copy=True)
    if gs_shifted.ndim == 2 and gs_shifted.shape[1] >= 1:
        gs_shifted[:, 0] -= offset_value

    if verbose:
        print(f"[offset] K_eff={K_eff} center={center:.4f} offset={offset_value:.4f}")

    return ndvi_shifted, gs_shifted, float(offset_value)


def post_fine_tuning_offsetter(preds, offset_value: float):
    """Undo the pre-fine-tuning vertical shift. Works on tensors, arrays, scalars, or nested containers."""
    if isinstance(preds, torch.Tensor):
        out = preds + offset_value
        finite = torch.isfinite(out)
        out = torch.where(finite, out.clamp(0.0, 1.0), out)
        return out
    if isinstance(preds, np.ndarray):
        out = preds.astype(np.float64) + offset_value
        finite = np.isfinite(out)
        out[finite] = np.clip(out[finite], 0.0, 1.0)
        return out.astype(np.float32)
    if isinstance(preds, (list, tuple)):
        return type(preds)(post_fine_tuning_offsetter(x, offset_value) for x in preds)
    if isinstance(preds, (int, float)):
        v = float(preds) + offset_value
        return max(0.0, min(1.0, v)) if np.isfinite(v) else v
    return preds
