"""Data loading, group statistics, and train/val splitting."""

from __future__ import annotations

import os
import numpy as np
import pandas as pd


def load_ndvi_from_csv(year: int, data_path: str, prefix: str = "NDVI") -> np.ndarray:
    """Load NDVI array (N, 127) from a CSV file.

    Expected file: ``{data_path}/{prefix}_{year}.csv`` with columns
    [FID, dap_0, dap_1, ..., dap_126].
    """
    csv_path = os.path.join(data_path, f"{prefix}_{year}.csv")
    df = pd.read_csv(csv_path)
    return df.iloc[:, 1:128].to_numpy(dtype=np.float32)


def generate_group_stats(
    ndvi_array: np.ndarray,
    year_array: np.ndarray,
    dap_index: int = 6,
) -> np.ndarray:
    """Compute per-sample group statistics (mean, std) at a specific DAP, grouped by year.

    Returns array of shape (N, 2).
    """
    assert ndvi_array.shape[0] == year_array.shape[0], "NDVI and year array length mismatch"
    group_stats = np.zeros((ndvi_array.shape[0], 2), dtype=np.float32)

    df = pd.DataFrame({"year": year_array, "dap_value": ndvi_array[:, dap_index]})
    year_stats = df.groupby("year")["dap_value"].agg(["mean", "std"]).to_dict()

    for i, year in enumerate(year_array):
        group_stats[i] = [year_stats["mean"][year], year_stats["std"][year]]
    return group_stats


def uniform_split(ndvi_array: np.ndarray, num_samples: int = 100) -> tuple[np.ndarray, np.ndarray]:
    """Uniformly sample validation rows from an NDVI array.

    Returns (train_set, val_set).
    """
    total_rows = len(ndvi_array)
    step = total_rows // num_samples
    indices = np.arange(0, total_rows, step)[:num_samples]
    val_set = ndvi_array[indices]
    train_set = np.delete(ndvi_array, indices, axis=0)
    return train_set, val_set


def save_group_stats(
    group_stats_array: np.ndarray,
    year_array: np.ndarray,
    results_path: str,
    prefix: str = "group_stats",
):
    """Save group stats as separate CSVs per year."""
    save_dir = os.path.join(results_path, "group_stats")
    os.makedirs(save_dir, exist_ok=True)

    df = pd.DataFrame(group_stats_array, columns=["group_mean", "group_std"])
    df["year"] = year_array

    for year in np.unique(year_array):
        df_year = df[df["year"] == year].drop(columns="year")
        path = os.path.join(save_dir, f"{prefix}_{year}.csv")
        df_year.to_csv(path, index=False)
