#!/usr/bin/env python
"""Generate deterministic Dirichlet partition manifests for all conditions.

This does NOT train anything. It loads each dataset once (labels only), computes
the frozen stratified split, generates the Dirichlet partition for every
``(alpha, seed)`` and persists a manifest plus characterization/rare-unseen
metadata. Training later loads these manifests; it never regenerates a partition.

Usage::

    python scripts/generate_partitions.py --datasets xiiotid 5gnidd \
        --alphas 0.75 0.5 0.25 --seeds 0 1 2 --clients 10
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

from protean_research.data import loading, manifest as manifest_mod, partition as partition_mod  # noqa: E402
from protean_research.data import rare_unseen, splits as splits_mod  # noqa: E402
from protean_research.data.characterization import characterize_support  # noqa: E402
from protean_research.data.pipeline import (  # noqa: E402
    TEST_FRACTION,
    TRAIN_FRACTION,
    VALIDATION_FRACTION,
)
from protean_research.utils.config import expand_path, load_datasets_config, load_yaml  # noqa: E402


def write_csv(path: Path, rows: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if not rows:
        path.write_text("", encoding="utf-8")
        return
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--datasets", nargs="+", default=["xiiotid", "5gnidd"])
    parser.add_argument("--alphas", nargs="+", type=float, default=[0.75, 0.5, 0.25])
    parser.add_argument("--seeds", nargs="+", type=int, default=[0, 1, 2])
    parser.add_argument("--clients", type=int, default=10)
    parser.add_argument("--partitions-dir", default=None)
    parser.add_argument("--analysis-dir", default=None)
    args = parser.parse_args(argv)

    datasets_cfg = load_datasets_config()
    protocol = load_yaml(REPO_ROOT / "configs" / "experiments" / "experiment_protocol.yaml")
    artifacts_cfg = protocol.get("artifacts", {})
    partitions_dir = expand_path(args.partitions_dir or artifacts_cfg.get("partition_manifest_dir", "artifacts/partitions"))
    analysis_dir = expand_path(args.analysis_dir or artifacts_cfg.get("analysis_dir", "artifacts/analysis"))

    created = []
    for dataset_key in args.datasets:
        spec = loading.load_spec(dataset_key, datasets_cfg)
        cleaned = loading.clean_frame(loading.read_raw(spec), spec)
        for alpha in args.alphas:
            for seed in args.seeds:
                split = splits_mod.stratified_split(
                    cleaned.labels, dataset_key, seed,
                    train_fraction=TRAIN_FRACTION,
                    validation_fraction=VALIDATION_FRACTION,
                    test_fraction=TEST_FRACTION,
                )
                labels_train = cleaned.labels[split.train]
                class_indices = partition_mod.class_index_map(labels_train, len(cleaned.class_names))
                rng_seed = partition_mod.partition_rng_seed(dataset_key, alpha, seed)
                assignments = partition_mod.dirichlet_partition(
                    class_indices, args.clients, alpha, rng_seed
                )
                manifest = manifest_mod.build_manifest(
                    dataset_key=dataset_key,
                    dataset_name=spec.name,
                    alpha=alpha,
                    seed=seed,
                    num_clients=args.clients,
                    class_names=cleaned.class_names,
                    benign_class=spec.benign_label,
                    labels=cleaned.labels,
                    split=split,
                    assignments=assignments,
                    rng_seed=rng_seed,
                    train_fraction=TRAIN_FRACTION,
                    validation_fraction=VALIDATION_FRACTION,
                    test_fraction=TEST_FRACTION,
                )
                path = manifest_mod.manifest_path(partitions_dir, dataset_key, alpha, seed)
                manifest_mod.save_manifest(path, manifest)

                # Verify exact reload before declaring success.
                manifest_mod.reproduce_train_assignments(manifest, cleaned.labels)

                support = manifest["support"]
                import numpy as np

                support_arr = np.asarray(support, dtype=np.int64)
                tag = f"{dataset_key}__alpha{alpha:g}__seed{seed}"
                write_csv(analysis_dir / f"{tag}__characterization.csv",
                          characterize_support(support_arr, cleaned.class_names))
                write_csv(analysis_dir / f"{tag}__rare_unseen.csv",
                          rare_unseen.rare_unseen_rows(
                              support_arr, cleaned.class_names, cleaned.benign_index,
                              dataset_key, alpha, seed))
                created.append({"dataset": dataset_key, "alpha": alpha, "seed": seed,
                                "path": str(path), "partition_hash": manifest["partition_hash"]})
                print(f"created {path.name}")

    print(json.dumps({"created": len(created)}, indent=2))
    return 0


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