#!/usr/bin/env python3
"""Convert GEO batch combined CSVs into the leaderboard dataset layout.

The input is the combined GEO export produced by scripts/prepare_geo_batch_datasets.R:
  label,batch,sample_id,<numeric features...>

The output matches the existing prepared datasets under data/datasets/<dataset>/:
  <dataset>_train.csv
  <dataset>_test.csv
  <dataset>_inference.csv
  <dataset>_predictions.csv
  provenance.json
"""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
from typing import Iterable

import numpy as np
import pandas as pd
from sklearn.model_selection import GroupShuffleSplit

ROOT = Path(__file__).resolve().parent.parent
META_COLUMNS = {
    "name", "names",
    "label", "labels", "group",
    "batch", "batches",
    "sample_id", "sample", "samples",
}


def parse_args(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source-root", type=Path, default=ROOT / "data" / "geo_batch",
                        help="Directory containing the combined GEO CSVs.")
    parser.add_argument("--output-root", type=Path, default=ROOT / "data" / "datasets",
                        help="Prepared-dataset root, typically data/datasets.")
    parser.add_argument("datasets", nargs="*", default=None,
                        help="Optional dataset basenames to export. Defaults to every combined CSV in --source-root.")
    parser.add_argument("--test-fraction", type=float, default=0.2,
                        help="Fraction of samples to hold out as test, split by batch group.")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--split-trials", type=int, default=256,
                        help="Number of randomized group splits to evaluate for balance.")
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args(argv)


def sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def infer_dataset_name(path: Path) -> str:
    stem = path.stem
    for suffix in ("_inference", "_train", "_test", "_predictions", "_metadata", "_expression"):
        if stem.endswith(suffix):
            return stem[: -len(suffix)]
    return stem


def discover_source_files(source_root: Path, datasets: list[str] | None) -> list[Path]:
    if datasets:
        paths = [source_root / f"{dataset}.csv" for dataset in datasets]
    else:
        paths = []
        for path in sorted(source_root.glob("*.csv")):
            if path.name.endswith(("_expression.csv", "_metadata.csv", "_predictions.csv", "_test.csv")):
                continue
            paths.append(path)
    if not paths:
        raise ValueError(f"No source CSVs found in {source_root}")
    return paths


def load_combined_csv(path: Path) -> dict:
    frame = pd.read_csv(path)
    lower_map = {column.lower(): column for column in frame.columns}
    label_col = next((lower_map[name] for name in ("label", "labels", "group") if name in lower_map), None)
    batch_col = next((lower_map[name] for name in ("batch", "batches") if name in lower_map), None)
    sample_col = next((lower_map[name] for name in ("sample_id", "sample", "samples", "name", "names") if name in lower_map), None)
    if label_col is None:
        raise ValueError(f"{path} lacks a label column")
    if batch_col is None:
        raise ValueError(f"{path} lacks a batch column")
    if sample_col is None:
        raise ValueError(f"{path} lacks a sample/name column")

    feature_cols = [column for column in frame.columns if column not in META_COLUMNS]
    if not feature_cols:
        raise ValueError(f"{path} has no feature columns")
    features = frame[feature_cols].apply(pd.to_numeric, errors="coerce").replace([np.inf, -np.inf], np.nan).fillna(0.0)
    return {
        "dataset": infer_dataset_name(path),
        "source_csv": path,
        "frame": frame,
        "label_col": label_col,
        "batch_col": batch_col,
        "sample_col": sample_col,
        "feature_cols": feature_cols,
        "features": features,
    }


def _distribution(series: pd.Series, labels: Iterable[str]) -> pd.Series:
    return series.astype(str).value_counts(normalize=True).reindex(list(labels), fill_value=0.0)


def choose_split(frame: pd.DataFrame, label_col: str, batch_col: str, test_fraction: float, seed: int, trials: int):
    labels = sorted(frame[label_col].astype(str).unique())
    overall = _distribution(frame[label_col], labels)
    target_test = max(1, int(round(len(frame) * test_fraction)))
    batches = frame[batch_col].astype(str).to_numpy()
    y = frame[label_col].astype(str).to_numpy()
    X = np.zeros((len(frame), 1), dtype=np.float32)

    best = None
    best_score = float("inf")
    for trial in range(max(1, trials)):
        splitter = GroupShuffleSplit(n_splits=1, test_size=test_fraction, random_state=seed + trial)
        train_idx, test_idx = next(splitter.split(X, y, batches))
        train_labels = pd.Series(y[train_idx])
        test_labels = pd.Series(y[test_idx])
        if train_labels.nunique() < 2 or test_labels.nunique() < 2:
            score = 10.0
        else:
            train_dist = _distribution(train_labels, labels)
            test_dist = _distribution(test_labels, labels)
            size_penalty = abs(len(test_idx) - target_test) / max(1, target_test)
            label_penalty = float(np.abs(train_dist - overall).sum() + np.abs(test_dist - overall).sum())
            batch_penalty = 0.0
            if train_labels.nunique() < len(labels):
                batch_penalty += 0.5
            if test_labels.nunique() < len(labels):
                batch_penalty += 0.5
            score = 2.0 * size_penalty + label_penalty + batch_penalty
        if score < best_score:
            best_score = score
            best = (train_idx, test_idx)
    if best is None:
        raise RuntimeError("Could not derive a valid train/test split")
    return best[0], best[1], best_score


def _write_csv(df: pd.DataFrame, path: Path):
    path.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(path, index=False)


def export_dataset(spec: dict, output_root: Path, test_fraction: float, seed: int, trials: int, overwrite: bool) -> dict:
    dataset_id = spec["dataset"]
    source_csv = spec["source_csv"]
    frame = spec["frame"]
    label_col = spec["label_col"]
    batch_col = spec["batch_col"]
    sample_col = spec["sample_col"]
    feature_cols = spec["feature_cols"]
    features = spec["features"].reset_index(drop=True)

    train_idx, test_idx, split_score = choose_split(frame, label_col, batch_col, test_fraction, seed, trials)
    train_frame = frame.iloc[train_idx].reset_index(drop=True)
    test_frame = frame.iloc[test_idx].reset_index(drop=True)
    train_features = features.iloc[train_idx].reset_index(drop=True)
    test_features = features.iloc[test_idx].reset_index(drop=True)

    names_train = train_frame[sample_col].astype(str).to_numpy()
    names_test = test_frame[sample_col].astype(str).to_numpy()
    batches_train = train_frame[batch_col].astype(str).to_numpy()
    batches_test = test_frame[batch_col].astype(str).to_numpy()
    labels_train = train_frame[label_col].astype(str).to_numpy()
    labels_test = test_frame[label_col].astype(str).to_numpy()

    train_out = pd.concat([
        pd.DataFrame({"name": names_train, "batch": batches_train, "label": labels_train}),
        train_features,
    ], axis=1)
    test_out = pd.concat([
        pd.DataFrame({"name": names_test, "batch": batches_test}),
        test_features,
    ], axis=1)
    inference_out = pd.concat([
        pd.DataFrame({"name": names_test, "batch": batches_test, "label": labels_test}),
        test_features,
    ], axis=1)
    predictions_out = pd.DataFrame({"name": names_test, "prediction": labels_test})

    target_dir = output_root / dataset_id
    if target_dir.exists() and not overwrite:
        raise FileExistsError(f"{target_dir} already exists; rerun with --overwrite to replace it")
    target_dir.mkdir(parents=True, exist_ok=True)

    _write_csv(train_out, target_dir / f"{dataset_id}_train.csv")
    _write_csv(test_out, target_dir / f"{dataset_id}_test.csv")
    _write_csv(inference_out, target_dir / f"{dataset_id}_inference.csv")
    _write_csv(predictions_out, target_dir / f"{dataset_id}_predictions.csv")

    provenance = {
        "dataset_id": dataset_id,
        "source_csv": str(source_csv),
        "source_sha256": sha256_file(source_csv),
        "source_rows": int(len(frame)),
        "source_features": len(feature_cols),
        "output_train_rows": int(len(train_out)),
        "output_test_rows": int(len(test_out)),
        "output_features": int(train_features.shape[1]),
        "n_batches": int(frame[batch_col].astype(str).nunique()),
        "n_labels": int(frame[label_col].astype(str).nunique()),
        "label_counts": frame[label_col].astype(str).value_counts().to_dict(),
        "batch_counts": frame[batch_col].astype(str).value_counts().to_dict(),
        "test_fraction": test_fraction,
        "split_seed": seed,
        "split_trials": trials,
        "split_score": split_score,
        "preprocessing": "batch-preserving group split; no additional transformation",
        "role": "prepared_geo_batch",
        "feature_columns_sha256": sha256_bytes("\n".join(feature_cols).encode()),
    }
    (target_dir / "provenance.json").write_text(json.dumps(provenance, indent=2) + "\n")
    print(f"wrote {dataset_id}: train={len(train_out)} test={len(test_out)} features={train_features.shape[1]} -> {target_dir}")
    return provenance


def main(argv=None) -> int:
    args = parse_args(argv)
    source_paths = discover_source_files(args.source_root, args.datasets)
    provenances = []
    for path in source_paths:
        if not path.exists():
            raise FileNotFoundError(path)
        provenances.append(export_dataset(load_combined_csv(path), args.output_root, args.test_fraction, args.seed, args.split_trials, args.overwrite))
    summary_path = args.output_root / "geo_batch_dataset_manifest.json"
    summary_path.write_text(json.dumps(provenances, indent=2) + "\n")
    print(f"wrote {summary_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())