"""Load and subsample the RAID dataset using direct CSV download.

Streaming RAID gives heavily biased samples because the data is sorted on
disk by domain. This version downloads the train split CSV directly via the
Hugging Face datasets library, loads it into pandas, shuffles properly, then
applies stratified subsampling.

The label is derived from the `model` field: "human" -> 0, anything else -> 1.
"""
from __future__ import annotations

import argparse
import random
from pathlib import Path

import pandas as pd
from datasets import load_dataset
from tqdm import tqdm


def load_raid_subsample(
    target_n: int = 2000,
    seed: int = 42,
    include_adversarial: bool = False,
    min_text_chars: int = 200,
    max_text_chars: int = 4000,
) -> pd.DataFrame:
    """Download RAID train split and return a stratified subsample."""

    print("Loading RAID train split (this downloads ~800MB on first run, cached after)...")
    ds = load_dataset("liamdugan/raid", split="train")
    df = ds.to_pandas()
    print(f"Full dataset loaded: {len(df):,} rows.")

    # Filter adversarial unless requested.
    if not include_adversarial:
        df = df[df["attack"] == "none"].copy()
        print(f"After removing adversarial: {len(df):,} rows.")

    # Filter by text length.
    df["text_len"] = df["generation"].str.len()
    df = df[(df["text_len"] >= min_text_chars) & (df["text_len"] <= max_text_chars)].copy()
    print(f"After length filter: {len(df):,} rows.")

    # Derive binary label.
    df["label"] = (df["model"] != "human").astype(int)
    df = df.rename(columns={"model": "generator", "generation": "text"})

    # Shuffle before sampling.
    df = df.sample(frac=1, random_state=seed).reset_index(drop=True)

    # Stratified sample: equal numbers per (domain, label) bucket.
    domains = df["domain"].unique()
    labels = [0, 1]
    n_buckets = len(domains) * len(labels)
    per_bucket = max(1, target_n // n_buckets)

    chunks = []
    for domain in domains:
        for label in labels:
            bucket = df[(df["domain"] == domain) & (df["label"] == label)]
            n_take = min(per_bucket, len(bucket))
            if n_take > 0:
                chunks.append(bucket.head(n_take))

    result = pd.concat(chunks, ignore_index=True)
    result = result.sample(frac=1, random_state=seed).reset_index(drop=True)

    # Keep only what we need.
    result = result[["id", "text", "domain", "generator", "label", "attack"]].copy()
    print(f"Final subsample: {len(result):,} rows.")
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--target-n", type=int, default=2000)
    parser.add_argument("--output", type=str, default="data/raid_subsample.parquet")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--include-adversarial", action="store_true")
    args = parser.parse_args()

    df = load_raid_subsample(
        target_n=args.target_n,
        seed=args.seed,
        include_adversarial=args.include_adversarial,
    )
    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(out_path, index=False)
    print(f"Wrote {len(df)} rows to {out_path}")
    print("\nLabel distribution:")
    print(df["label"].value_counts())
    print("\nDomain distribution:")
    print(df["domain"].value_counts())
    print("\nGenerator distribution:")
    print(df["generator"].value_counts())


if __name__ == "__main__":
    main()