"""Save and load data splits."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import pandas as pd

from vito_crop_classification.constants import get_data_folder


def save_splits(
    df_train: pd.DataFrame,
    df_val: pd.DataFrame,
    df_test: pd.DataFrame,
    data_f: Path | str,
    data_cfg: str,
    scale_cfg: dict[str, tuple[float, float]],
    feat_cfg: dict[str, Any] | None = None,
) -> None:
    """Save split keys to disk."""
    assert "key_ts" in df_train.columns, "key_ts column is missing from dataset"

    # create save directory
    save_dir = get_data_folder() / data_f if type(data_f) == str else data_f
    save_dir.mkdir(parents=True, exist_ok=True)

    # save path
    save_path = save_dir / f"cfg_{data_cfg}.json"
    with open(save_path, "w") as f:
        json.dump(
            {
                "scale_cfg": scale_cfg,
                "train": df_train["key_ts"].to_list(),
                "val": df_val["key_ts"].to_list(),
                "test": df_test["key_ts"].to_list(),
            },
            f,
        )

    # Save feature settings
    save_path = save_dir / "cfg_data.json"
    if feat_cfg is not None:
        with open(save_path, "w") as f:
            json.dump(feat_cfg, f)
