#!/usr/bin/env python3
"""
Generate deterministic level seeds for policy training, OOD training, validation, and evaluation.

This script creates reproducible seed sets that can be used with the procgen-afh
fork's ordered level seeds feature (level_seeds parameter).

Usage:
    python generate_level_seeds.py --policy-train 10000 --ood-train 5000 --validation 500 --ood-eval 1000 -o seeds.json
    python generate_level_seeds.py --paper --ood-train 5000 --validation 500 --ood-eval 1000 -o seeds/paper.json

The generated seeds are:
- Deterministic: Same base seed always produces same seed sets
- Non-overlapping: Each phase gets unique seeds
- Shuffled within phase: Seeds are randomly ordered but reproducibly so (unless --paper is used)

Paper mode (--paper):
- policy_train: sequential seeds [0, 100,000) without shuffling
- ood_train/validation/ood_eval: randomly shuffled from [100,000, ...) using --base-seed
"""

import argparse
import json
from pathlib import Path
from typing import Dict, List

import numpy as np


# Paper mode constants
PAPER_POLICY_TRAIN_SEEDS = 100_000  # Sequential seeds [0, 100000) for policy training


def generate_seeds(
    num_policy_train: int,
    num_ood_train: int,
    num_validation: int,
    num_ood_eval: int,
    base_seed: int = 0,
    paper_mode: bool = False,
) -> Dict[str, List[int]]:
    """
    Generate deterministic, non-overlapping seed sets for each training phase.

    The seeds are generated by:
    1. Creating a large pool of sequential integers
    2. Shuffling them deterministically using the base_seed (unless paper_mode for policy_train)
    3. Splitting into non-overlapping chunks for each phase

    Args:
        num_policy_train: Number of seeds for policy training (ignored in paper_mode)
        num_ood_train: Number of seeds for OOD detector training
        num_validation: Number of seeds for validation
        num_ood_eval: Number of seeds for OOD evaluation
        base_seed: Base random seed for reproducible generation
        paper_mode: If True, use sequential seeds [0, 100000) for policy_train only

    Returns:
        Dictionary with keys 'policy_train', 'ood_train', 'validation', 'ood_eval' containing seed lists
    """
    if paper_mode:
        # Paper mode: sequential seeds [0, 100000) for policy_train
        # Other seeds are randomly sampled from a non-overlapping range
        policy_train_seeds = list(range(PAPER_POLICY_TRAIN_SEEDS))

        # Generate other seeds from range starting after policy_train
        # to ensure no overlap
        rng = np.random.default_rng(base_seed)
        other_pool_size = num_ood_train + num_validation + num_ood_eval
        other_pool = np.arange(
            PAPER_POLICY_TRAIN_SEEDS, PAPER_POLICY_TRAIN_SEEDS + other_pool_size
        )
        rng.shuffle(other_pool)

        ood_train_end = num_ood_train
        validation_end = ood_train_end + num_validation

        seeds = {
            "policy_train": policy_train_seeds,
            "ood_train": other_pool[:ood_train_end].tolist(),
            "validation": other_pool[ood_train_end:validation_end].tolist(),
            "ood_eval": other_pool[validation_end:].tolist(),
        }
    else:
        total_seeds = num_policy_train + num_ood_train + num_validation + num_ood_eval

        # Create RNG with base seed for reproducibility
        rng = np.random.default_rng(base_seed)

        # Generate a pool of unique seeds and shuffle them
        # We use a large range to avoid any potential collisions
        seed_pool = np.arange(total_seeds)
        rng.shuffle(seed_pool)

        # Split into non-overlapping sets
        policy_end = num_policy_train
        ood_train_end = policy_end + num_ood_train
        validation_end = ood_train_end + num_validation

        seeds = {
            "policy_train": seed_pool[:policy_end].tolist(),
            "ood_train": seed_pool[policy_end:ood_train_end].tolist(),
            "validation": seed_pool[ood_train_end:validation_end].tolist(),
            "ood_eval": seed_pool[validation_end:].tolist(),
        }

    return seeds


def save_seeds(
    seeds: Dict[str, List[int]],
    output_path: Path,
    base_seed: int,
    paper_mode: bool = False,
) -> None:
    """
    Save seeds to a JSON file with metadata.

    Args:
        seeds: Dictionary of seed lists
        output_path: Path to save the JSON file
        base_seed: The base seed used for generation (for documentation)
        paper_mode: Whether seeds were generated in paper mode
    """
    output_data = {
        "metadata": {
            "base_seed": base_seed,
            "paper_mode": paper_mode,
            "num_policy_train": len(seeds["policy_train"]),
            "num_ood_train": len(seeds["ood_train"]),
            "num_validation": len(seeds["validation"]),
            "num_ood_eval": len(seeds["ood_eval"]),
            "description": (
                "Deterministic level seeds for procgen environments. "
                "Use with level_seeds parameter in ProcgenGym3Env."
                + (
                    f" Paper mode: policy_train uses sequential seeds [0, {PAPER_POLICY_TRAIN_SEEDS})."
                    if paper_mode
                    else ""
                )
            ),
        },
        "seeds": seeds,
    }

    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w") as f:
        json.dump(output_data, f, indent=2)

    print(f"Saved seeds to {output_path}")
    if paper_mode:
        print(f"  - Mode: paper (policy_train: sequential [0, {PAPER_POLICY_TRAIN_SEEDS}))")
    print(f"  - Policy training: {len(seeds['policy_train'])} seeds")
    print(f"  - OOD training:    {len(seeds['ood_train'])} seeds")
    print(f"  - Validation:      {len(seeds['validation'])} seeds")
    print(f"  - OOD evaluation:  {len(seeds['ood_eval'])} seeds")


def load_seeds(input_path: Path) -> Dict[str, List[int]]:
    """
    Load seeds from a JSON file.

    Args:
        input_path: Path to the JSON file

    Returns:
        Dictionary with seed lists for each phase
    """
    with open(input_path) as f:
        data = json.load(f)

    return data["seeds"]


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate deterministic level seeds for policy training, OOD training, and evaluation.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Generate seeds with custom counts
  python generate_level_seeds.py --policy-train 50000 --ood-train 10000 --validation 500 --ood-eval 5000 -o seeds/custom.json

  # Use a specific base seed for reproducibility
  python generate_level_seeds.py --policy-train 10000 --ood-train 5000 --validation 500 --ood-eval 1000 --base-seed 42 -o seeds/seed42.json

  # Paper mode: sequential seeds [0, 100000) for policy_train
  python generate_level_seeds.py --paper --ood-train 5000 --validation 500 --ood-eval 1000 -o seeds/paper.json

The generated JSON can be loaded and used with procgen:

  from generate_level_seeds import load_seeds
  from procgen import ProcgenGym3Env

  seeds = load_seeds("seeds.json")
  env = ProcgenGym3Env(num=4, env_name="coinrun", level_seeds=seeds["ood_eval"])
        """,
    )

    parser.add_argument(
        "--policy-train",
        type=int,
        default=None,
        help="Number of seeds for policy training (required unless --paper)",
    )
    parser.add_argument(
        "--ood-train",
        type=int,
        default=None,
        help="Number of seeds for OOD detector training (required)",
    )
    parser.add_argument(
        "--validation",
        type=int,
        default=None,
        help="Number of seeds for validation (required)",
    )
    parser.add_argument(
        "--ood-eval",
        type=int,
        default=None,
        help="Number of seeds for OOD evaluation (required)",
    )
    parser.add_argument(
        "--base-seed",
        type=int,
        default=0,
        help="Base random seed for deterministic generation (default: 0)",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=Path("seeds/level_seeds.json"),
        help="Output file path (default: seeds/level_seeds.json)",
    )
    parser.add_argument(
        "--paper",
        action="store_true",
        help=(
            f"Paper mode: use sequential seeds [0, {PAPER_POLICY_TRAIN_SEEDS}) for policy_train. "
            "Other seeds use shuffled seeds from a non-overlapping range. "
            "Ignores --policy-train but respects --ood-train, --validation, --ood-eval, --base-seed."
        ),
    )

    args = parser.parse_args()

    # Validate arguments
    if args.ood_train is None or args.validation is None or args.ood_eval is None:
        parser.error("--ood-train, --validation, and --ood-eval are required")

    if args.paper:
        # In paper mode, only validate ood_train, validation, and ood_eval
        if args.ood_train < 0 or args.validation < 0 or args.ood_eval < 0:
            parser.error("Seed counts must be non-negative")
    else:
        if args.policy_train is None:
            parser.error("--policy-train is required (or use --paper mode)")

        if (
            args.policy_train < 0
            or args.ood_train < 0
            or args.validation < 0
            or args.ood_eval < 0
        ):
            parser.error("Seed counts must be non-negative")

        if args.policy_train + args.ood_train + args.validation + args.ood_eval == 0:
            parser.error("At least one phase must have seeds")

    # Generate and save seeds
    seeds = generate_seeds(
        num_policy_train=args.policy_train,
        num_ood_train=args.ood_train,
        num_validation=args.validation,
        num_ood_eval=args.ood_eval,
        base_seed=args.base_seed,
        paper_mode=args.paper,
    )

    save_seeds(seeds, args.output, args.base_seed, paper_mode=args.paper)


if __name__ == "__main__":
    main()

