"""Evaluation utilities: per-DAP result saving and metrics."""

import os
import numpy as np
import pandas as pd

from ndvi_forecast._utils import to_numpy, safe_rmse, safe_r2


def save_dap_results(
    base_dir: str,
    dap: int,
    ndvi_inputs: np.ndarray,
    preds,
    targets,
    grid_ids=None,
):
    """Save per-DAP inputs, predictions, ground truth, and metrics to CSVs.

    Creates ``{base_dir}/dap_{dap:03d}/`` with: inputs.csv, preds.csv,
    ground_truth.csv, metrics.csv, diagnostics.csv, overall.txt.
    """
    dap_dir = os.path.join(base_dir, f"dap_{dap:03d}")
    os.makedirs(dap_dir, exist_ok=True)

    preds_np = to_numpy(preds).astype(np.float32)
    targets_np = to_numpy(targets).astype(np.float32)
    inputs_np = to_numpy(ndvi_inputs).astype(np.float32)

    if preds_np.ndim == 1:
        preds_np = preds_np[np.newaxis, :]
    if targets_np.ndim == 1:
        targets_np = targets_np[np.newaxis, :]

    N = preds_np.shape[0]
    L = preds_np.shape[1]

    if grid_ids is None:
        grid_ids = np.arange(N)

    # Save raw arrays
    pd.DataFrame(inputs_np).to_csv(os.path.join(dap_dir, "inputs.csv"), index=False)
    pd.DataFrame(preds_np).to_csv(os.path.join(dap_dir, "preds.csv"), index=False)
    pd.DataFrame(targets_np).to_csv(os.path.join(dap_dir, "ground_truth.csv"), index=False)

    # Per-grid metrics
    rows = []
    for i in range(N):
        p, t = preds_np[i], targets_np[i]
        m = np.isfinite(p) & np.isfinite(t)
        if m.sum() < 2:
            rows.append({"grid_id": grid_ids[i], "rmse": np.nan, "r2": np.nan})
            continue
        rmse = safe_rmse(t[m], p[m])
        r2 = safe_r2(t[m], p[m])
        rows.append({"grid_id": grid_ids[i], "rmse": rmse, "r2": r2})

    pd.DataFrame(rows).to_csv(os.path.join(dap_dir, "metrics.csv"), index=False)

    # Diagnostics
    diag_rows = []
    for i in range(N):
        p, t = preds_np[i], targets_np[i]
        pred_slope = np.diff(p) if L > 1 else np.array([0.0])
        tgt_slope = np.diff(t) if L > 1 else np.array([0.0])
        slope_err = float(np.nanmean(np.abs(pred_slope - tgt_slope)))
        pred_peak = int(np.argmax(p)) if L > 0 else -1
        tgt_peak = int(np.argmax(t)) if L > 0 else -1
        boundary_err = float(np.abs(p[0] - t[0])) if L > 0 else np.nan
        diag_rows.append({
            "grid_id": grid_ids[i],
            "slope_mae": slope_err,
            "pred_peak_idx": pred_peak,
            "tgt_peak_idx": tgt_peak,
            "peak_offset": pred_peak - tgt_peak,
            "boundary_err": boundary_err,
        })
    pd.DataFrame(diag_rows).to_csv(os.path.join(dap_dir, "diagnostics.csv"), index=False)

    # Overall summary
    all_rmse = safe_rmse(targets_np, preds_np)
    all_r2 = safe_r2(targets_np, preds_np)

    diffs = preds_np[:, np.newaxis, :] - preds_np[np.newaxis, :, :]
    diversity = float(np.sqrt(np.nanmean(diffs ** 2) + 1e-12)) if N > 1 else 0.0

    with open(os.path.join(dap_dir, "overall.txt"), "w") as f:
        f.write(f"DAP: {dap}\n")
        f.write(f"N_grids: {N}\n")
        f.write(f"Horizon: {L}\n")
        f.write(f"RMSE: {all_rmse:.6f}\n")
        f.write(f"R2: {all_r2:.6f}\n")
        f.write(f"Diversity: {diversity:.6f}\n")
