# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.

import random
from multiprocessing.pool import ThreadPool
from pathlib import Path
from typing import Any

from tqdm import tqdm


def generate_run_name_suffix():
    adjectives = [
        "happy",
        "brave",
        "clever",
        "swift",
        "quiet",
        "bold",
        "calm",
        "eager",
        "fierce",
        "gentle",
        "honest",
        "jolly",
        "kind",
        "lively",
        "merry",
        "nimble",
        "proud",
        "quick",
        "smooth",
        "tender",
        "vigilant",
        "wise",
        "zealous",
    ]

    names = [
        "curie",
        "darwin",
        "tesla",
        "newton",
        "einstein",
        "feynman",
        "turing",
        "hopper",
        "lovelace",
        "pasteur",
        "fermi",
        "bohr",
        "planck",
        "hawking",
        "sagan",
        "faraday",
        "franklin",
        "maxwell",
        "dirac",
        "heisenberg",
    ]

    return f"{random.choice(adjectives)}{random.choice(names)}"


def get_cache_dir(run_name: str) -> str:
    cache_dir = Path.home() / ".llama" / "evals" / "cache" / run_name
    if not cache_dir.exists():
        cache_dir.mkdir(parents=True, exist_ok=True)
    return str(cache_dir)


def map_with_progress(f: callable, xs: list[Any], num_threads: int = 50, show_progress: bool = True):
    """
    Apply f to each element of xs, using a ThreadPool, and show progress.
    """
    if show_progress:
        with ThreadPool(min(num_threads, len(xs))) as pool:
            return list(tqdm(pool.imap(f, xs), total=len(xs)))
    else:
        with ThreadPool(min(num_threads, len(xs))) as pool:
            return list(pool.imap(f, xs))
