# -*- coding: utf-8 -*-
"""特徴量CSVから train/validation/test の固定分割を作る。"""
import argparse
from pathlib import Path

import numpy as np
import pandas as pd


ROOT = Path(__file__).resolve().parents[2]
DEFAULT_INPUT = ROOT / "検証" / "data" / "features.csv"
DEFAULT_OUTDIR = ROOT / "検証" / "data"


def _grade(score):
    if score >= 80:
        return "A"
    if score >= 60:
        return "B"
    if score >= 21:
        return "C"
    return "D"


def make_splits(df, seed=42, val=0.1, test=0.2):
    rng = np.random.default_rng(seed)
    rows = []
    for _, group in df.groupby(df["target"].apply(_grade), sort=True):
        idx = group.index.to_numpy().copy()
        rng.shuffle(idx)
        n = len(idx)
        n_test = int(round(n * test))
        n_val = int(round(n * val))
        for i in idx[:n_test]:
            rows.append((i, "test"))
        for i in idx[n_test:n_test + n_val]:
            rows.append((i, "val"))
        for i in idx[n_test + n_val:]:
            rows.append((i, "train"))
    split = pd.DataFrame(rows, columns=["row_index", "split"]).sort_values("row_index")
    return df[["sample_id"]].join(split.set_index("row_index"), how="left")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, default=DEFAULT_INPUT)
    parser.add_argument("--outdir", type=Path, default=DEFAULT_OUTDIR)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    df = pd.read_csv(args.input)
    if "sample_id" not in df.columns or "target" not in df.columns:
        raise ValueError("features.csv には sample_id と target が必要です。")

    split = make_splits(df, seed=args.seed)
    args.outdir.mkdir(parents=True, exist_ok=True)
    for name in ["train", "val", "test"]:
        part = split[split["split"] == name][["sample_id"]]
        part.to_csv(args.outdir / f"splits_{name}.csv", index=False, encoding="utf-8-sig")
        print(f"{name}: {len(part)}")
    split.to_csv(args.outdir / "splits.csv", index=False, encoding="utf-8-sig")
    print(f"saved: {args.outdir / 'splits.csv'}")


if __name__ == "__main__":
    main()