import argparse
from pathlib import Path

from editing_DiTs.benchmark.multi_job import (
    MultiJobRunner,
    get_common_argparser,
    validate_common_args,
    PROJECT_ROOT,
)
from editing_DiTs.image_generators.load_image_generator import (
    ImageGeneratorRegistry,
)

# Constants specific to image generation.
DEFAULT_BENCHMARK_DIR = PROJECT_ROOT / "results/benchmark/mjhq/"
DEFAULT_EMBEDDING_DIR = DEFAULT_BENCHMARK_DIR / "embeddings"
DEFAULT_IMAGE_DIR = DEFAULT_BENCHMARK_DIR / "images"
DEFAULT_REF_DIR = Path("~/MJHQ-30K/")


def collect_schedules(schedule_dir: Path) -> list[Path]:
    """Recursively collect all JSON schedule files."""
    if schedule_dir.is_file():
        return [schedule_dir]
    return list(schedule_dir.glob("**/*.json"))


def collect_categories(embedding_dir: Path, ref_dir: Path) -> list[Path]:
    ref_dir_categories = [d.name for d in ref_dir.iterdir() if d.is_dir()]

    # if len(ref_dir_categories) != 10:
    #     raise ValueError(
    #         f"Expected 10 categories in reference directory {ref_dir}, found {len(ref_dir_categories)}."
    #     )

    embedding_dir_categories = [
        d for d in embedding_dir.iterdir() if d.is_dir()
    ]

    if set((d.name for d in embedding_dir_categories)) != set(
        ref_dir_categories
    ):
        raise ValueError(
            f"Expected categories in embedding directory {embedding_dir} to match reference categories."
        )

    return embedding_dir_categories


def collect_categories_batched(
    embedding_dir: Path, ref_dir: Path
) -> list[Path]:
    ref_dir_categories = [d.name for d in ref_dir.iterdir() if d.is_dir()]

    # if len(ref_dir_categories) != 10:
    #     raise ValueError(
    #         f"Expected 10 categories in reference directory {ref_dir}, found {len(ref_dir_categories)}."
    #     )

    embedding_dir_categories = [
        d for d in embedding_dir.iterdir() if d.is_dir()
    ]

    if set((d.name for d in embedding_dir_categories)) != set(
        ref_dir_categories
    ):
        raise ValueError(
            f"Expected categories in embedding directory {embedding_dir} to match reference categories."
        )

    all_batches = []
    for cat_dir in embedding_dir_categories:
        batch_dirs = (d for d in cat_dir.iterdir() if d.is_dir())
        all_batches.extend(batch_dirs)

    return all_batches


def generate_benchmark_commands(
    image_generator: str,
    schedule_dir: Path,
    name_rel_to: Path | None,
    python_path: Path,
    seed: int,
    seed_step: int,
    embedding_dir: Path,
    ref_dir: Path,
    image_dir: Path,
    num_images_per_prompt: int,
    batch_size: int,
    regen_not_n_imgs: int | None,
    include_seed: bool = False,
    batched: bool = False,
) -> list[str]:
    """
    Generates a flat list of command strings for each schedule file.
    """
    schedule_jsons = collect_schedules(schedule_dir)
    print(f"Found {len(schedule_jsons)} schedules.")

    if not batched:
        embedding_categories = collect_categories(embedding_dir, ref_dir)
        print(f"Found {len(embedding_categories)} categories.")
    else:
        embedding_categories = collect_categories_batched(
            embedding_dir, ref_dir
        )
        print(f"Found {len(embedding_categories)} batches of categories.")

    commands = []
    for schedule_json in schedule_jsons:
        for cat_dir in embedding_categories:
            subdir = (
                schedule_json.parent.relative_to(name_rel_to)
                if name_rel_to is not None
                else ""
            )
            cat_name = cat_dir.name
            output_subdir = image_dir / subdir / cat_name

            regen_str = (
                f"--regen-if-not-n-images {regen_not_n_imgs}"
                if regen_not_n_imgs
                else ""
            )
            dont_include_seed_str = (
                "--dont-include-seed" if not include_seed else ""
            )
            command = (
                f"{python_path} {PROJECT_ROOT}/editing_DiTs/benchmark/generate_images.py"
                f" --image-generator {image_generator}"
                f" --seed {seed}"
                f" --seed-step {seed_step}"
                f" --input-dir {cat_dir}"
                f" --output-dir {output_subdir}"
                f" --schedule-dir {schedule_json}"
                f" --num-images-per-prompt {num_images_per_prompt}"
                f" --batch-size {batch_size}"
                f" {regen_str}"
                f" {dont_include_seed_str}"
            )
            commands.append(command)

    print(f"Generated {len(commands)} commands.")
    return commands


def main() -> None:
    # Create the common argparser and add image-generation-specific arguments.
    common_parser = get_common_argparser()
    parser = argparse.ArgumentParser(parents=[common_parser])
    parser.add_argument(
        "--image-generator",
        type=str,
        required=True,
        help="The name of the image generator to use.",
        choices=list(ImageGeneratorRegistry.registry.keys()),
    )
    parser.add_argument(
        "--embedding-dir",
        type=Path,
        default=DEFAULT_EMBEDDING_DIR,
        help="Directory containing the embeddings for MJHQ FID.",
    )
    parser.add_argument(
        "--reference-dir",
        type=Path,
        default=DEFAULT_REF_DIR,
        help="Directory containing the reference images for MJHQ FID.",
    )
    parser.add_argument(
        "--image-dir",
        type=Path,
        default=DEFAULT_IMAGE_DIR,
        help="Directory where benchmark images will be saved.",
    )
    parser.add_argument(
        "--schedule-dir",
        type=Path,
        required=True,
        help="Directory containing schedules to run (recursively searched).",
    )
    parser.add_argument(
        "--name-schedule-dir-rel-to",
        type=Path,
        help="Anchor path for subdirectories found under --schedule-dir.",
    )
    parser.add_argument(
        "--num-images-per-prompt",
        type=int,
        default=1,
        help="Number of images to generate per prompt.",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=200,
        help="Batch size while generating images.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=0,
        help="Random seed to start image generation.",
    )
    parser.add_argument(
        "--seed-step",
        type=int,
        default=0,
        help="Amount to increment the seed for each subsequent generation.",
    )
    parser.add_argument(
        "--regen-if-not-n-images",
        type=int,
        default=None,
        help="Force re-generation if output dir does not contain exactly n images.",
    )
    parser.add_argument(
        "--include-seed",
        type=bool,
        default=False,
        help="Whether to include the seed in the generated images' names.",
    )
    parser.add_argument(
        "--batched",
        action="store_true",
        help="If each category has batched dirs under it.",
    )
    args = parser.parse_args()

    # Validate and resolve the common arguments.
    validate_common_args(args)
    if args.name_schedule_dir_rel_to is None and args.schedule_dir.is_dir():
        args.name_schedule_dir_rel_to = args.schedule_dir.resolve()

    args.embedding_dir = args.embedding_dir.resolve()
    args.image_dir = args.image_dir.resolve()
    args.schedule_dir = args.schedule_dir.resolve()
    if args.name_schedule_dir_rel_to is not None:
        args.name_schedule_dir_rel_to = args.name_schedule_dir_rel_to.resolve()

    if not args.embedding_dir.exists():
        print(f"WARNING: Embedding directory {args.embedding_dir} not found.")
    if not args.schedule_dir.exists():
        print(f"WARNING: Schedule directory {args.schedule_dir} not found.")

    # Generate a flat list of commands.
    commands = generate_benchmark_commands(
        args.image_generator,
        args.schedule_dir,
        args.name_schedule_dir_rel_to,
        args.python_path,
        args.seed,
        args.seed_step,
        args.embedding_dir,
        args.reference_dir,
        args.image_dir,
        args.num_images_per_prompt,
        args.batch_size,
        args.regen_if_not_n_images,
        args.include_seed,
        args.batched,
    )

    runner = MultiJobRunner(
        scripts_dir=args.scripts_dir,
        time_per_job=args.time_per_job,
        qos=args.qos,
        gpu=args.gpu,
        cpu=args.cpu,
        mem=args.mem,
        nodelist=args.nodes,
        dry_run=args.dry_run,
    )
    runner.submit_jobs(commands)  # automatically batched


if __name__ == "__main__":
    main()
