"""
Upload the bathroom-toilet dataset with train/val split to HuggingFace.

Replicates the exact split used in training (see experiment/data/datasets.py):
    sklearn.model_selection.train_test_split(image_ids, test_size=0.2, random_state=42)

Usage:
    python -m experiment.scripts.upload_split_to_hf \
        --csv_path CC3M-Dataset/bathroom_filter/bathroom_toilet_labels.csv \
        --image_dir CC3M-Dataset/cc3m_images/train \
        --repo_id <your-hf-username>/bathroom-toilet-cc3m
"""

import argparse
import csv
import os

from datasets import Dataset, DatasetDict, Features, Value, Image, ClassLabel
from sklearn.model_selection import train_test_split
from huggingface_hub import login

# Must match experiment/data/datasets.py
SPLIT_SEED = 42
SPLIT_TEST_SIZE = 0.2


def load_rows(csv_path: str, image_dir: str):
    """Load all rows from the labels CSV, filtering out non-bathroom and missing images."""
    rows = []
    missing = 0
    skipped_non_bathroom = 0
    with open(csv_path, "r") as f:
        reader = csv.DictReader(f)
        for row in reader:
            bathroom = int(row["bathroom"])
            toilet = int(row["toilet"])
            if bathroom == 0 and toilet == 0:
                skipped_non_bathroom += 1
                continue
            image_path = os.path.join(image_dir, f"{row['image_id']}.jpg")
            if not os.path.exists(image_path):
                missing += 1
                continue
            rows.append({
                "image_id": row["image_id"],
                "image": image_path,
                "bathroom": bathroom,
                "toilet": toilet,
            })
    print(f"Loaded {len(rows)} rows "
          f"({skipped_non_bathroom} non-bathroom skipped, {missing} images not found)")
    return rows


def split_rows(rows: list[dict]) -> tuple[list[dict], list[dict]]:
    """Replicate the exact train/val split from datasets.py."""
    all_ids = [r["image_id"] for r in rows]
    train_ids, val_ids = train_test_split(
        all_ids, test_size=SPLIT_TEST_SIZE, random_state=SPLIT_SEED,
    )
    train_set = set(train_ids)
    val_set = set(val_ids)

    train_rows = [r for r in rows if r["image_id"] in train_set]
    val_rows = [r for r in rows if r["image_id"] in val_set]

    print(f"Split: {len(train_rows)} train, {len(val_rows)} val "
          f"(seed={SPLIT_SEED}, test_size={SPLIT_TEST_SIZE})")
    return train_rows, val_rows


def rows_to_dataset(rows: list[dict]) -> Dataset:
    """Convert list of row dicts to a HuggingFace Dataset."""
    return Dataset.from_dict(
        {
            "image_id": [r["image_id"] for r in rows],
            "image": [r["image"] for r in rows],
            "bathroom": [r["bathroom"] for r in rows],
            "toilet": [r["toilet"] for r in rows],
        },
        features=Features({
            "image_id": Value("string"),
            "image": Image(),
            "bathroom": ClassLabel(names=["no", "yes"]),
            "toilet": ClassLabel(names=["no", "yes"]),
        }),
    )


def print_report(train_rows, val_rows):
    """Print category breakdown matching the training script style."""
    for name, rows in [("train", train_rows), ("val", val_rows)]:
        cats: dict[str, int] = {}
        for r in rows:
            key = f"bathroom={r['bathroom']},toilet={r['toilet']}"
            cats[key] = cats.get(key, 0) + 1
        print(f"  {name}: {cats}")


def main():
    parser = argparse.ArgumentParser(
        description="Upload bathroom-toilet dataset splits to HuggingFace"
    )
    parser.add_argument(
        "--csv_path",
        type=str,
        default="CC3M-Dataset/bathroom_filter/bathroom_toilet_labels.csv",
    )
    parser.add_argument(
        "--image_dir",
        type=str,
        default="CC3M-Dataset/cc3m_images/train",
    )
    parser.add_argument(
        "--repo_id",
        type=str,
        required=True,
        help="HuggingFace repo id, e.g. your-username/bathroom-toilet-cc3m",
    )
    parser.add_argument(
        "--private",
        action="store_true",
        help="Make the dataset private on HuggingFace",
    )
    args = parser.parse_args()

    # Auth
    hf_token = os.environ.get("HUGGING_FACE_API_KEY")
    if hf_token:
        login(token=hf_token)
    else:
        print("No HUGGING_FACE_API_KEY in env, using cached HF credentials")

    # Load and split
    rows = load_rows(args.csv_path, args.image_dir)
    if len(rows) == 0:
        print("ERROR: No data loaded. Check csv_path and image_dir.")
        return

    train_rows, val_rows = split_rows(rows)
    print_report(train_rows, val_rows)

    # Build HF datasets
    print("\nBuilding HuggingFace datasets...")
    ds = DatasetDict({
        "train": rows_to_dataset(train_rows),
        "val": rows_to_dataset(val_rows),
    })
    print(ds)

    # Upload
    print(f"\nPushing to {args.repo_id}...")
    ds.push_to_hub(args.repo_id, private=args.private)
    print(f"Done! Dataset available at: https://huggingface.co/datasets/{args.repo_id}")


if __name__ == "__main__":
    main()