            try:
                folder_date = datetime.strptime("_".join(item.name.split("_")[:2]), date_format)
                if folder_date > most_recent_time:
                    most_recent_time = folder_date
                    most_recent_folder = item
            except (ValueError, IndexError):
                continue

    return most_recent_folder

def set_demo_mode(env_args_list: list[EnvArgs]):
    """Set the demo mode for the experiments. This can be useful for generating videos for demos."""
    for env_args in env_args_list:
        env_args.viewport = {"width": 1280, "height": 720}
        env_args.record_video = True
        env_args.wait_for_user_message = False
        env_args.slow_mo = 1000

def _convert_env_args(env_args_list) -> list[EnvArgs]:
    """Return a list where every element is the *new* EnvArgs.

    For backward compatibility, we need to convert the old EnvArgs to the new one.

    Args:
        env_args_list (list): list of EnvArgs objects to convert

    Returns:
        list: list of converted EnvArgs objects

    Raises:
        TypeError: If an element in env_args_list is not of expected type.
    """
    from bgym import EnvArgs as BGymEnvArgs

    new_list = []
    for ea in env_args_list:
        # already new → keep as‑is
        if isinstance(ea, (EnvArgs, AbstractEnvArgs)):
            new_list.append(ea)
        # old → convert
        elif isinstance(ea, BGymEnvArgs):
            new_list.append(EnvArgs(**asdict(ea)))
        else:
            raise TypeError(f"Unexpected type: {type(ea)}")
    return new_list

# def _flag_sequential_exp(exp_args_list: list[ExpArgs], benchmark: Benchmark):
#     if benchmark.name.startswith("visualwebarena"):
#         sequential_subset = benchmark.subset_from_glob("requires_reset", "True")
#         sequential_subset = set(
#             [env_args.task_name for env_args in sequential_subset.env_args_list]
#         )
#         for exp_args in exp_args_list:
#             if exp_args.env_args.task_name in sequential_subset:
#                 exp_args.sequential = True

# def ablation_study(start_agent: AgentArgs, changes, benchmark: str, demo_mode=False):
#     """Ablation study of an agent.

#     Changes is a list of tuples (path_to_attribute, value) to change in the agent
#     configuration.

#     Args:


    return slugify(study_name, max_length=200, allow_unicode=True)

@dataclass
class SequentialStudies(AbstractStudy):
    """
    Sequential execution of multiple studies.

    This is required for e.g. WebArena, where a server reset is required between evaluations of each agent.
    """

    studies: list[Study]

    @property
    def name(self):
        """The name of the study."""
        agent_names = [a.agent_name for study in self.studies for a in study.agent_args]
        benchmark_names = [study.benchmark.name for study in self.studies]
        return _make_study_name(agent_names, benchmark_names, self.suffix)

    def find_incomplete(self, include_errors=True):
        for study in self.studies:
            study.find_incomplete(include_errors=include_errors)

    def run(
        self,
        n_jobs=1,
        parallel_backend="ray",
        strict_reproducibility=False,
        n_relaunch=3,
        exp_root=RESULTS_DIR,
    ):
        # This sequence of of making directories is important to make sure objects are materialized
        # properly before saving. Otherwise relaunch may not work properly.
        self.make_dir()
        for study in self.studies:
            study.make_dir(exp_root=self.dir)

        self.save(exp_root=exp_root)
        self._run(n_jobs, parallel_backend, strict_reproducibility, n_relaunch)
        _, summary_df, _ = self.get_results()
        logger.info("\n" + str(summary_df))
        logger.info(f"SequentialStudies {self.name} finished.")


run

class SequentialStudies(AbstractStudy):
    """
    Sequential execution of multiple studies.

    This is required for e.g. WebArena, where a server reset is required between evaluations of each agent.
    """

    studies: list[Study]

    @property
    def name(self):
        """The name of the study."""
        agent_names = [a.agent_name for study in self.studies for a in study.agent_args]
        benchmark_names = [study.benchmark.name for study in self.studies]
        return _make_study_name(agent_names, benchmark_names, self.suffix)

    def find_incomplete(self, include_errors=True):
        for study in self.studies:
            study.find_incomplete(include_errors=include_errors)

    def run(
        self,
        n_jobs=1,
        parallel_backend="ray",
        strict_reproducibility=False,
        n_relaunch=3,
        exp_root=RESULTS_DIR,
    ):
        # This sequence of of making directories is important to make sure objects are materialized
        # properly before saving. Otherwise relaunch may not work properly.
        self.make_dir()
        for study in self.studies:
            study.make_dir(exp_root=self.dir)

        self.save(exp_root=exp_root)
        self._run(n_jobs, parallel_backend, strict_reproducibility, n_relaunch)
        _, summary_df, _ = self.get_results()
        logger.info("\n" + str(summary_df))
        logger.info(f"SequentialStudies {self.name} finished.")

    def _run(self, n_jobs=1, parallel_backend="ray", strict_reproducibility=False, n_relaunch=3):
        for study in self.studies:
            study.run(n_jobs, parallel_backend, strict_reproducibility, n_relaunch)

    def override_max_steps(self, max_steps):
        for study in self.studies:
            study.override_max_steps(max_steps)

    def append_to_journal(self, strict_reproducibility=True):
        for study in self.studies:
            study.append_to_journal(strict_reproducibility=strict_reproducibility)

def _init_worker(server_queue: Queue):
    """Run once at the initialization of the worker in the multiprocessing.Pool.

    This is typically used to initialize different environment variables of the WebArena server for
    multiple instances in parallel.


            suffix=suffix,
            comment=comment,
            ignore_dependencies=ignore_dependencies,
        )

class AbstractStudy(ABC):
    """Abstract class for a study."""

    dir: Path = None
    suffix: str = ""

    @abstractmethod
    def find_incomplete(self, include_errors=True):
        """Prepare the study for relaunching by finding incomplete experiments"""

    @abstractmethod
    def run(self, n_jobs=1, parallel_backend="ray", strict_reproducibility=False, n_relaunch=3):
        """Run the study"""

    def make_dir(self, exp_root=RESULTS_DIR):
        """Create a directory for the study"""
        if self.dir is None:
            dir_name = f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}_{self.name}"

            self.dir = Path(exp_root) / dir_name
        self.dir.mkdir(parents=True, exist_ok=True)

    def save(self, exp_root=RESULTS_DIR):
        """Pickle the study to the directory"""
        # TODO perhaps remove exp_args_list before pickling and when loading bring them from the individual directories

        self.make_dir(exp_root=exp_root)
        with gzip.open(self.dir / "study.pkl.gz", "wb") as f:
            pickle.dump(self, f)

    def get_results(self, suffix="", also_save=True):
        """Recursively load all results from the study directory and summarize them."""
        result_df = inspect_results.load_result_df(self.dir)
        error_report = inspect_results.error_report(result_df, max_stack_trace=3, use_log=True)
        summary_df = inspect_results.summarize_study(result_df)

        if also_save:
            suffix = f"_{suffix}" if suffix else ""
            result_df.to_csv(self.dir / f"result_df{suffix}.csv")
            summary_df.to_csv(self.dir / f"summary_df{suffix}.csv")
            (self.dir / f"error_report{suffix}.md").write_text(error_report)


save

    dir: Path = None
    suffix: str = ""

    @abstractmethod
    def find_incomplete(self, include_errors=True):
        """Prepare the study for relaunching by finding incomplete experiments"""

    @abstractmethod
    def run(self, n_jobs=1, parallel_backend="ray", strict_reproducibility=False, n_relaunch=3):
        """Run the study"""

    def make_dir(self, exp_root=RESULTS_DIR):
        """Create a directory for the study"""
        if self.dir is None:
            dir_name = f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}_{self.name}"

            self.dir = Path(exp_root) / dir_name
        self.dir.mkdir(parents=True, exist_ok=True)

    def save(self, exp_root=RESULTS_DIR):
        """Pickle the study to the directory"""
        # TODO perhaps remove exp_args_list before pickling and when loading bring them from the individual directories

        self.make_dir(exp_root=exp_root)
        with gzip.open(self.dir / "study.pkl.gz", "wb") as f:
            pickle.dump(self, f)

    def get_results(self, suffix="", also_save=True):
        """Recursively load all results from the study directory and summarize them."""
        result_df = inspect_results.load_result_df(self.dir)
        error_report = inspect_results.error_report(result_df, max_stack_trace=3, use_log=True)
        summary_df = inspect_results.summarize_study(result_df)

        if also_save:
            suffix = f"_{suffix}" if suffix else ""
            result_df.to_csv(self.dir / f"result_df{suffix}.csv")
            summary_df.to_csv(self.dir / f"summary_df{suffix}.csv")
            (self.dir / f"error_report{suffix}.md").write_text(error_report)

        return result_df, summary_df, error_report

    def shuffle_exps(self):
        """Shuffle the experiments in the study."""
        self.exp_args_list = random.sample(self.exp_args_list, len(self.exp_args_list))


    @abstractmethod
    def run(self, n_jobs=1, parallel_backend="ray", strict_reproducibility=False, n_relaunch=3):
        """Run the study"""

    def make_dir(self, exp_root=RESULTS_DIR):
        """Create a directory for the study"""
        if self.dir is None:
            dir_name = f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}_{self.name}"

            self.dir = Path(exp_root) / dir_name
        self.dir.mkdir(parents=True, exist_ok=True)

    def save(self, exp_root=RESULTS_DIR):
        """Pickle the study to the directory"""
        # TODO perhaps remove exp_args_list before pickling and when loading bring them from the individual directories

        self.make_dir(exp_root=exp_root)
        with gzip.open(self.dir / "study.pkl.gz", "wb") as f:
            pickle.dump(self, f)

    def get_results(self, suffix="", also_save=True):
        """Recursively load all results from the study directory and summarize them."""
        result_df = inspect_results.load_result_df(self.dir)
        error_report = inspect_results.error_report(result_df, max_stack_trace=3, use_log=True)
        summary_df = inspect_results.summarize_study(result_df)

        if also_save:
            suffix = f"_{suffix}" if suffix else ""
            result_df.to_csv(self.dir / f"result_df{suffix}.csv")
            summary_df.to_csv(self.dir / f"summary_df{suffix}.csv")
            (self.dir / f"error_report{suffix}.md").write_text(error_report)

        return result_df, summary_df, error_report

    def shuffle_exps(self):
        """Shuffle the experiments in the study."""
        self.exp_args_list = random.sample(self.exp_args_list, len(self.exp_args_list))

@dataclass
class Study(AbstractStudy):
    """A study coresponds to one or multiple agents evaluated on a benchmark.

    This is part of the high level API to help keep experiments organized and reproducible.

    Attributes:
        agent_args: list[AgentArgs]
            The agent configuration(s) to run. *IMPORTANT*: these objects will be pickled and
            unpickled.  Make sure they are imported from a package that is accessible from
            PYTHONPATH. Otherwise, it won't load in agentlab-xray.
        benchmark: Benchmark | str
            The benchmark to run the agents on. See DEFAULT_BENCHMARKS for the main ones. You
            can also make your own by modifying an existing one.


        # TODO perhaps remove exp_args_list before pickling and when loading bring them from the individual directories

        self.make_dir(exp_root=exp_root)
        with gzip.open(self.dir / "study.pkl.gz", "wb") as f:
            pickle.dump(self, f)

    def get_results(self, suffix="", also_save=True):
        """Recursively load all results from the study directory and summarize them."""
        result_df = inspect_results.load_result_df(self.dir)
        error_report = inspect_results.error_report(result_df, max_stack_trace=3, use_log=True)
        summary_df = inspect_results.summarize_study(result_df)

        if also_save:
            suffix = f"_{suffix}" if suffix else ""
            result_df.to_csv(self.dir / f"result_df{suffix}.csv")
            summary_df.to_csv(self.dir / f"summary_df{suffix}.csv")
            (self.dir / f"error_report{suffix}.md").write_text(error_report)

        return result_df, summary_df, error_report

    def shuffle_exps(self):
        """Shuffle the experiments in the study."""
        self.exp_args_list = random.sample(self.exp_args_list, len(self.exp_args_list))

@dataclass
class Study(AbstractStudy):
    """A study coresponds to one or multiple agents evaluated on a benchmark.

    This is part of the high level API to help keep experiments organized and reproducible.

    Attributes:
        agent_args: list[AgentArgs]
            The agent configuration(s) to run. *IMPORTANT*: these objects will be pickled and
            unpickled.  Make sure they are imported from a package that is accessible from
            PYTHONPATH. Otherwise, it won't load in agentlab-xray.
        benchmark: Benchmark | str
            The benchmark to run the agents on. See DEFAULT_BENCHMARKS for the main ones. You
            can also make your own by modifying an existing one.
        dir: Path
            The directory where the study will be saved. If None, a directory will be created in
            RESULTS_DIR.
        suffix: str


            The average step timeout in seconds. This is used to stop the experiments if they are
            taking too long. The default is 60 seconds.
        demo_mode: bool
            If True, the experiments will be run in demo mode, which will record videos, and enable
            visual effects for actions.
    """

    agent_args: list[AgentArgs] = None
    benchmark: Benchmark | str = None
    dir: Path = None
    suffix: str = ""  # used for adding a personnal comment to the study name
    uuid: str = None
    reproducibility_info: dict = None
    logging_level: int = logging.DEBUG
    logging_level_stdout: int = logging.WARNING
    comment: str = None  # Extra comments from the authors of this study
    ignore_dependencies: bool = False
    avg_step_timeout: int = 60
    demo_mode: bool = False

    def __post_init__(self):
        """Initialize the study. Set the uuid, and generate the exp_args_list."""
        self.uuid = uuid.uuid4()
        if isinstance(self.benchmark, str):
            self.benchmark = DEFAULT_BENCHMARKS[self.benchmark.lower()]()

        self.benchmark.env_args_list = _convert_env_args(self.benchmark.env_args_list)

        if isinstance(self.dir, str):
            self.dir = Path(self.dir)
        self.make_exp_args_list()

    def make_exp_args_list(self):
        """Generate the exp_args_list from the agent_args and the benchmark."""
        self.exp_args_list = self.agents_on_benchmark(
            self.agent_args,
            self.benchmark,
            logging_level=self.logging_level,
            logging_level_stdout=self.logging_level_stdout,
            ignore_dependencies=self.ignore_dependencies,
            demo_mode=self.demo_mode,
        )

    def find_incomplete(self, include_errors=True):
        """Find incomplete or errored experiments in the study directory for relaunching.

        Args:
            include_errors: bool
                If True, include errored experiments in the list.

        Returns:


    reproducibility_info: dict = None
    logging_level: int = logging.DEBUG
    logging_level_stdout: int = logging.WARNING
    comment: str = None  # Extra comments from the authors of this study
    ignore_dependencies: bool = False
    avg_step_timeout: int = 60
    demo_mode: bool = False

    def __post_init__(self):
        """Initialize the study. Set the uuid, and generate the exp_args_list."""
        self.uuid = uuid.uuid4()
        if isinstance(self.benchmark, str):
            self.benchmark = DEFAULT_BENCHMARKS[self.benchmark.lower()]()

        self.benchmark.env_args_list = _convert_env_args(self.benchmark.env_args_list)

        if isinstance(self.dir, str):
            self.dir = Path(self.dir)
        self.make_exp_args_list()

    def make_exp_args_list(self):
        """Generate the exp_args_list from the agent_args and the benchmark."""
        self.exp_args_list = self.agents_on_benchmark(
            self.agent_args,
            self.benchmark,
            logging_level=self.logging_level,
            logging_level_stdout=self.logging_level_stdout,
            ignore_dependencies=self.ignore_dependencies,
            demo_mode=self.demo_mode,
        )

    def find_incomplete(self, include_errors=True):
        """Find incomplete or errored experiments in the study directory for relaunching.

        Args:
            include_errors: bool
                If True, include errored experiments in the list.

        Returns:
            list[ExpArgs]: The list of all experiments with completed ones replaced by a
                dummy exp_args to keep the task dependencies.
        """
        self.exp_args_list = find_incomplete(self.dir, include_errors=include_errors)
        n_incomplete = non_dummy_count(self.exp_args_list)
        n_error = [
            getattr(exp_args, "status", "incomplete") == "error" for exp_args in self.exp_args_list
        ].count(True)
        return n_incomplete, n_error

    def load_exp_args_list(self):


        )

    def find_incomplete(self, include_errors=True):
        """Find incomplete or errored experiments in the study directory for relaunching.

        Args:
            include_errors: bool
                If True, include errored experiments in the list.

        Returns:
            list[ExpArgs]: The list of all experiments with completed ones replaced by a
                dummy exp_args to keep the task dependencies.
        """
        self.exp_args_list = find_incomplete(self.dir, include_errors=include_errors)
        n_incomplete = non_dummy_count(self.exp_args_list)
        n_error = [
            getattr(exp_args, "status", "incomplete") == "error" for exp_args in self.exp_args_list
        ].count(True)
        return n_incomplete, n_error

    def load_exp_args_list(self):
        logger.info(f"Loading experiments from {self.dir}")
        self.exp_args_list = list(inspect_results.yield_all_exp_results(savedir_base=self.dir))

    def set_reproducibility_info(self, strict_reproducibility=False, comment=None):
        """Gather relevant information that may affect the reproducibility of the experiment

        e.g.: versions of BrowserGym, benchmark, AgentLab...

        Args:
            strict_reproducibility: bool
                If True, all modifications have to be committed before running the experiments.
                Also, if relaunching a study, it will not be possible if the code has changed.
            comment: str
                Extra comment to add to the reproducibility information.
        """
        agent_names = [a.agent_name for a in self.agent_args]
        info = repro.get_reproducibility_info(
            agent_names,
            self.benchmark,
            self.uuid,
            ignore_changes=not strict_reproducibility,
            comment=comment,


        Args:
            include_errors: bool
                If True, include errored experiments in the list.

        Returns:
            list[ExpArgs]: The list of all experiments with completed ones replaced by a
                dummy exp_args to keep the task dependencies.
        """
        self.exp_args_list = find_incomplete(self.dir, include_errors=include_errors)
        n_incomplete = non_dummy_count(self.exp_args_list)
        n_error = [
            getattr(exp_args, "status", "incomplete") == "error" for exp_args in self.exp_args_list
        ].count(True)
        return n_incomplete, n_error

    def load_exp_args_list(self):
        logger.info(f"Loading experiments from {self.dir}")
        self.exp_args_list = list(inspect_results.yield_all_exp_results(savedir_base=self.dir))

    def set_reproducibility_info(self, strict_reproducibility=False, comment=None):
        """Gather relevant information that may affect the reproducibility of the experiment

        e.g.: versions of BrowserGym, benchmark, AgentLab...

        Args:
            strict_reproducibility: bool
                If True, all modifications have to be committed before running the experiments.
                Also, if relaunching a study, it will not be possible if the code has changed.
            comment: str
                Extra comment to add to the reproducibility information.
        """
        agent_names = [a.agent_name for a in self.agent_args]
        info = repro.get_reproducibility_info(
            agent_names,
            self.benchmark,
            self.uuid,
            ignore_changes=not strict_reproducibility,
            comment=comment,
            allow_bypass_benchmark_version=not strict_reproducibility,
        )
        if self.reproducibility_info is not None:
            repro.assert_compatible(
                self.reproducibility_info, info, raise_if_incompatible=strict_reproducibility
            )
        self.reproducibility_info = info

    def run(
        self,
        n_jobs=1,
        parallel_backend="ray",
        strict_reproducibility=False,
        n_relaunch=3,
        relaunch_errors=True,
        exp_root=RESULTS_DIR,
    ):
        self.set_reproducibility_info(
            strict_reproducibility=strict_reproducibility, comment=self.comment
        )
        self.save(exp_root=exp_root)

        n_exp = len(self.exp_args_list)
        last_error_count = None

        for i in range(n_relaunch):
            logger.info(f"Launching study {self.name} - trial {i + 1} / {n_relaunch}")


            max_workers=len(parallel_servers), initializer=_init_worker, initargs=(server_queue,)
        ) as executor:
            # Create list of arguments for each study
            study_args = [
                (study, n_jobs, parallel_backend, strict_reproducibility, n_relaunch)
                for study in self.studies
            ]

            # Submit all tasks and wait for completion
            futures = [executor.submit(_run_study, *args) for args in study_args]

            # Wait for all futures to complete and raise any exceptions
            for future in futures:
                future.result()

@dataclass
class ParallelStudies_alt(SequentialStudies):
    parallel_servers: list[BaseServer] | int = None

    def _run(
        self,
        n_jobs=1,
        parallel_backend="ray",
        strict_reproducibility=False,
        n_relaunch=3,
    ):
        parallel_servers = self.parallel_servers
        if isinstance(parallel_servers, int):
            parallel_servers = [BaseServer() for _ in range(parallel_servers)]

        server_queue = Manager().Queue()
        for server in parallel_servers:
            server_queue.put(server)

        with Pool(len(parallel_servers), initializer=_init_worker, initargs=(server_queue,)) as p:
            p.starmap(
                _run_study,
                [
                    (study, n_jobs, parallel_backend, strict_reproducibility, n_relaunch)
                    for study in self.studies
                ],
            )

def get_most_recent_study(
    root_dir: Path = None, date_format: str = "%Y-%m-%d_%H-%M-%S", contains=None
):
    """Return the most recent directory based on the date in the folder name.

    Args:
        root_dir: The directory to search in
        date_format: The format of the date in the folder name
        contains: If not None, only consider folders that contains this string

    Returns:
        Path: The most recent folder satisfying the conditions
    """

    if root_dir is None:
        root_dir = RESULTS_DIR

    most_recent_folder = None


        # This sequence of of making directories is important to make sure objects are materialized
        # properly before saving. Otherwise relaunch may not work properly.
        self.make_dir()
        for study in self.studies:
            study.make_dir(exp_root=self.dir)

        self.save(exp_root=exp_root)
        self._run(n_jobs, parallel_backend, strict_reproducibility, n_relaunch)
        _, summary_df, _ = self.get_results()
        logger.info("\n" + str(summary_df))
        logger.info(f"SequentialStudies {self.name} finished.")

    def _run(self, n_jobs=1, parallel_backend="ray", strict_reproducibility=False, n_relaunch=3):
        for study in self.studies:
            study.run(n_jobs, parallel_backend, strict_reproducibility, n_relaunch)

    def override_max_steps(self, max_steps):
        for study in self.studies:
            study.override_max_steps(max_steps)

    def append_to_journal(self, strict_reproducibility=True):
        for study in self.studies:
            study.append_to_journal(strict_reproducibility=strict_reproducibility)

def _init_worker(server_queue: Queue):
    """Run once at the initialization of the worker in the multiprocessing.Pool.

    This is typically used to initialize different environment variables of the WebArena server for
    multiple instances in parallel.

    Args:
        server_queue: Queue
            A queue of object implementing BaseServer to initialize (or anything with a init
            method).
    """
    print("initializing server instance with on process", os.getpid())
    print(f"using queue {server_queue}")
    server_instance = server_queue.get()  # type: "WebArenaInstanceVars"
    logger.warning(f"Initializing server instance {server_instance} from process {os.getpid()}")
    server_instance.init()


name

        benchmark_name = benchmark_names[0]
    else:
        benchmark_name = f"{len(benchmark_names)}_benchmarks"

    study_name = f"{agent_name}_on_{benchmark_name}_{suffix if suffix else ''}"

    return slugify(study_name, max_length=200, allow_unicode=True)

@dataclass
class SequentialStudies(AbstractStudy):
    """
    Sequential execution of multiple studies.

    This is required for e.g. WebArena, where a server reset is required between evaluations of each agent.
    """

    studies: list[Study]

    @property
    def name(self):
        """The name of the study."""
        agent_names = [a.agent_name for study in self.studies for a in study.agent_args]
        benchmark_names = [study.benchmark.name for study in self.studies]
        return _make_study_name(agent_names, benchmark_names, self.suffix)

    def find_incomplete(self, include_errors=True):
        for study in self.studies:
            study.find_incomplete(include_errors=include_errors)

    def run(
        self,
        n_jobs=1,
        parallel_backend="ray",
        strict_reproducibility=False,
        n_relaunch=3,
        exp_root=RESULTS_DIR,
    ):
        # This sequence of of making directories is important to make sure objects are materialized
        # properly before saving. Otherwise relaunch may not work properly.
        self.make_dir()
        for study in self.studies:
            study.make_dir(exp_root=self.dir)

        self.save(exp_root=exp_root)


        strict_reproducibility=False,
        n_relaunch=3,
        exp_root=RESULTS_DIR,
    ):
        # This sequence of of making directories is important to make sure objects are materialized
        # properly before saving. Otherwise relaunch may not work properly.
        self.make_dir()
        for study in self.studies:
            study.make_dir(exp_root=self.dir)

        self.save(exp_root=exp_root)
        self._run(n_jobs, parallel_backend, strict_reproducibility, n_relaunch)
        _, summary_df, _ = self.get_results()
        logger.info("\n" + str(summary_df))
        logger.info(f"SequentialStudies {self.name} finished.")

    def _run(self, n_jobs=1, parallel_backend="ray", strict_reproducibility=False, n_relaunch=3):
        for study in self.studies:
            study.run(n_jobs, parallel_backend, strict_reproducibility, n_relaunch)

    def override_max_steps(self, max_steps):
        for study in self.studies:
            study.override_max_steps(max_steps)

    def append_to_journal(self, strict_reproducibility=True):
        for study in self.studies:
            study.append_to_journal(strict_reproducibility=strict_reproducibility)

def _init_worker(server_queue: Queue):
    """Run once at the initialization of the worker in the multiprocessing.Pool.

    This is typically used to initialize different environment variables of the WebArena server for
    multiple instances in parallel.

    Args:
        server_queue: Queue
            A queue of object implementing BaseServer to initialize (or anything with a init
            method).
    """
    print("initializing server instance with on process", os.getpid())
    print(f"using queue {server_queue}")
    server_instance = server_queue.get()  # type: "WebArenaInstanceVars"


load

            strict_reproducibility: bool
                If True, incomplete experiments will raise an error.
        """
        _, summary_df, _ = self.get_results()
        repro.append_to_journal(
            self.reproducibility_info,
            summary_df,
            strict_reproducibility=strict_reproducibility,
        )

    @property
    def name(self):
        agent_names = [a.agent_name for a in self.agent_args]
        return _make_study_name(agent_names, [self.benchmark.name], self.suffix)

    def override_max_steps(self, max_steps):
        for exp_args in self.exp_args_list:
            exp_args.env_args.max_steps = max_steps

    @staticmethod
    def load(dir: Path) -> "Study":
        dir = Path(dir)
        study_path = dir / "study.pkl.gz"
        if not study_path.exists() and dir.is_dir():
            # For backward compatibility
            first_result = next(
                inspect_results.yield_all_exp_results(savedir_base=dir, progress_fn=None)
            )
            benchmark_name = first_result.exp_args.env_args.task_name.split(".")[0]
            agent_args = first_result.exp_args.agent_args
            study = Study(agent_args=agent_args, benchmark=benchmark_name, dir=dir)
        else:
            with gzip.open(dir / "study.pkl.gz", "rb") as f:
                study = pickle.load(f)  # type: Study
            study.dir = dir

            # # just a check
            # for i, exp_args in enumerate(study.exp_args_list):
            #     if exp_args.order != i:
            #         logging.warning(f"The order of the experiments is not correct. {exp_args.order} != {i}")

        return study

    @staticmethod
    def load_most_recent(root_dir: Path = None, contains=None) -> "Study":
        return Study.load(get_most_recent_study(root_dir, contains=contains))

    def agents_on_benchmark(
        self,
        agents: list[AgentArgs] | AgentArgs,
        benchmark: Benchmark,
        demo_mode=False,
        logging_level: int = logging.INFO,
        logging_level_stdout: int = logging.INFO,
        ignore_dependencies=False,
    ):
        """Run one or multiple agents on a benchmark.

        Args:
            agents: list[AgentArgs] | AgentArgs
                The agent configuration(s) to run.
            benchmark: Benchmark


            # For backward compatibility
            first_result = next(
                inspect_results.yield_all_exp_results(savedir_base=dir, progress_fn=None)
            )
            benchmark_name = first_result.exp_args.env_args.task_name.split(".")[0]
            agent_args = first_result.exp_args.agent_args
            study = Study(agent_args=agent_args, benchmark=benchmark_name, dir=dir)
        else:
            with gzip.open(dir / "study.pkl.gz", "rb") as f:
                study = pickle.load(f)  # type: Study
            study.dir = dir

            # # just a check
            # for i, exp_args in enumerate(study.exp_args_list):
            #     if exp_args.order != i:
            #         logging.warning(f"The order of the experiments is not correct. {exp_args.order} != {i}")

        return study

    @staticmethod
    def load_most_recent(root_dir: Path = None, contains=None) -> "Study":
        return Study.load(get_most_recent_study(root_dir, contains=contains))

    def agents_on_benchmark(
        self,
        agents: list[AgentArgs] | AgentArgs,
        benchmark: Benchmark,
        demo_mode=False,
        logging_level: int = logging.INFO,
        logging_level_stdout: int = logging.INFO,
        ignore_dependencies=False,
    ):
        """Run one or multiple agents on a benchmark.

        Args:
            agents: list[AgentArgs] | AgentArgs
                The agent configuration(s) to run.
            benchmark: Benchmark
                The benchmark to run the agents on.
            demo_mode: bool
                If True, the experiments will be run in demo mode.
            logging_level: int


            )
            benchmark_name = first_result.exp_args.env_args.task_name.split(".")[0]
            agent_args = first_result.exp_args.agent_args
            study = Study(agent_args=agent_args, benchmark=benchmark_name, dir=dir)
        else:
            with gzip.open(dir / "study.pkl.gz", "rb") as f:
                study = pickle.load(f)  # type: Study
            study.dir = dir

            # # just a check
            # for i, exp_args in enumerate(study.exp_args_list):
            #     if exp_args.order != i:
            #         logging.warning(f"The order of the experiments is not correct. {exp_args.order} != {i}")

        return study

    @staticmethod
    def load_most_recent(root_dir: Path = None, contains=None) -> "Study":
        return Study.load(get_most_recent_study(root_dir, contains=contains))

    def agents_on_benchmark(
        self,
        agents: list[AgentArgs] | AgentArgs,
        benchmark: Benchmark,
        demo_mode=False,
        logging_level: int = logging.INFO,
        logging_level_stdout: int = logging.INFO,
        ignore_dependencies=False,
    ):
        """Run one or multiple agents on a benchmark.

        Args:
            agents: list[AgentArgs] | AgentArgs
                The agent configuration(s) to run.
            benchmark: Benchmark
                The benchmark to run the agents on.
            demo_mode: bool
                If True, the experiments will be run in demo mode.
            logging_level: int
                The logging level for individual jobs.
            logging_level_stdout: int
                The logging level for the stdout.
            ignore_dependencies: bool
                If True, the dependencies will be ignored and all experiments can be run in parallel.

        Returns:
            list[ExpArgs]: The list of experiments to run.

        Raises:
            ValueError: If multiple agents are run on a benchmark that requires manual reset.
        """

        if not isinstance(agents, (list, tuple)):
            agents = [agents]

        if benchmark.name.startswith("visualwebarena") or benchmark.name.startswith("webarena"):
            if len(agents) > 1:
                raise ValueError(
                    f"Only one agent can be run on {benchmark.name} since the instance requires manual reset after each evaluation."
                )

        for agent in agents:
            agent.set_benchmark(
                benchmark, demo_mode
            )  # the agent can adapt (lightly?) to the benchmark

        env_args_list = benchmark.env_args_list
        if demo_mode:
            set_demo_mode(env_args_list)

        exp_args_list = []

        for agent in agents:
            for env_args in env_args_list:
                exp_args = ExpArgs(
                    agent_args=agent,
                    env_args=env_args,
                    logging_level=logging_level,
                    logging_level_stdout=logging_level_stdout,
                )
                exp_args_list.append(exp_args)

        for i, exp_args in enumerate(exp_args_list):
            exp_args.order = i

        # not required with ray, but keeping around if we would need it for visualwebareana on joblib
        # _flag_sequential_exp(exp_args_list, benchmark)

        if not ignore_dependencies:
            # populate the depends_on field based on the task dependencies in the benchmark
            exp_args_list = add_dependencies(exp_args_list, benchmark.dependency_graph_over_tasks())
        else:
            logger.warning(
                f"Ignoring dependencies for benchmark {benchmark.name}. This could lead to different results."
            )

        return exp_args_list

def _make_study_name(agent_names, benchmark_names, suffix=None):
    """Make a study name from the agent and benchmark names."""
    # extract unique agent and benchmark names
    agent_names = list(set(agent_names))
    benchmark_names = list(set(benchmark_names))

    if len(agent_names) == 1:
        agent_name = agent_names[0]
    else:
        agent_name = f"{len(agent_names)}_agents"

    if len(benchmark_names) == 1:
        benchmark_name = benchmark_names[0]
    else:
        benchmark_name = f"{len(benchmark_names)}_benchmarks"

    study_name = f"{agent_name}_on_{benchmark_name}_{suffix if suffix else ''}"


"""Dirty script to visualize the dependency graph of a benchmark, e.g. webarena, vsisualwebarena,
etc. You may have to detust it to make it work for you."""

import math

import bgym
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from bgym import DEFAULT_BENCHMARKS

def clean_dict(dependency_dict: dict[str, list[str]]) -> dict[str, list[str]]:
    new_dep = {}
    for key, deps in dependency_dict.items():
        new_key = key.split(".")[-1]

        new_dep[new_key] = [dep.split(".")[-1] for dep in deps]
    return new_dep

def dict_to_networkx(dependency_dict: dict[str, list[str]]) -> nx.DiGraph:

    G = nx.DiGraph()
    i = 0
    # Add edges from each node to its dependencies
    for node, dependencies in dependency_dict.items():
        i += 1
        if i > 20:
            pass

        print(node, dependencies)
        # Add edges from the node to each of its dependencies
        for dep in dependencies:
            G.add_edge(dep, node)
    return G

def plot_graph(G, ax, title=None, node_color="lightblue", node_size=40, font_size=8):
    """
    Plot a single graph component on the given matplotlib axis.

    Args:
        G: NetworkX graph (should be a single connected component)
        ax: Matplotlib axis to plot on
        title: Optional title for the subplot
        node_color: Color for the nodes
        node_size: Size of the nodes
        font_size: Size of the node labels
    """
    # Use a simple layout for better performance
    # pos = nx.spring_layout(G, k=0.1, iterations=100)

    pos = nx.kamada_kawai_layout(G)

    # pos = nx.spectral_layout(G)

    def name_to_size(name):
        if "-" in name:
            start, end = name.split("-")

            n_nodes = int(end) - int(start) + 1
        else:
            n_nodes = 1
        size_factor = node_size / 10
        return n_nodes * size_factor

    # compute size based on name
    sizes = [name_to_size(name) for name in G.nodes]

    nx.draw(
        G,
        pos,
        ax=ax,
        with_labels=True,
        node_color=node_color,
        node_size=sizes,
        font_size=font_size,
        font_weight="normal",
        arrows=True,
        arrowsize=15,
    )

    if title:
        ax.set_title(title)
    ax.axis("off")

def plot_components_grid(
    components, max_cols=4, node_color="lightblue", node_size=2000, font_size=10
):
    """
    Plot components in a grid layout.

    Args:
        components: List of NetworkX graphs, one per component
        max_cols: Maximum number of columns in the grid
        node_color: Color for the nodes
        node_size: Size of the nodes
        font_size: Size of the node labels

    Returns:
        matplotlib figure
    """
    n_components = len(components)

    if n_components == 0:
        print("No components found")
        return None

    # Calculate grid dimensions
    ncols = min(n_components, max_cols)
    nrows = math.ceil(n_components / ncols)

    # Create figure with a reasonable size per subplot
    fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 4 * nrows))
    fig.suptitle("Dependency Graph Components", size=16)

    # Make axes iterable even if there's only one
    if n_components == 1:
        axes = np.array([[axes]])
    elif nrows == 1:
        axes = np.array([axes])
    elif ncols == 1:
        axes = axes.reshape(-1, 1)

    # Plot each component
    for idx, component in enumerate(components):
        i, j = divmod(idx, ncols)
        title = f"Component {idx+1} ({component.number_of_nodes()} nodes)"
        plot_graph(
            component,
            axes[i, j],
            title,
            node_color=node_color,
            node_size=node_size,
            font_size=font_size,
        )

    # Remove empty subplots
    for idx in range(n_components, nrows * ncols):
        i, j = divmod(idx, ncols)
        axes[i, j].remove()

    plt.tight_layout()
    return fig

def compress_sequential_chains(dep_dict: dict[str, list[str]]) -> dict[str, list[str]]:
    """
    Compress chains of sequential numbers in a dependency dictionary.
    Returns a new dictionary with compressed chains using range notation.

    Args:
        dep_dict: Dictionary mapping string numbers to list of string number dependencies

    Returns:
        Dictionary with compressed chains using range notation
    """
    # Convert to integers for easier processing
    int_dict = {int(k): [int(x) for x in v] for k, v in dep_dict.items()}

    # Find chains
    chains = []
    current_chain = []

    # Sort nodes for sequential processing
    nodes = sorted(int_dict.keys())

    i = 0
    while i < len(nodes):
        node = nodes[i]

        # Start new chain
        if not current_chain:
            current_chain = [node]
            i += 1
            continue

        # Check if this node continues the chain
        last_node = current_chain[-1]

        # Conditions for chain continuation:
        # 1. Numbers are consecutive
        # 2. Current node has exactly one dependency
        # 3. That dependency is the previous node in chain
        # 4. The previous node has exactly one successor
        is_consecutive = node == last_node + 1
        has_single_dep = len(int_dict[node]) == 1
        deps_on_last = has_single_dep and int_dict[node][0] == last_node
        last_has_single_successor = sum(1 for k, v in int_dict.items() if last_node in v) == 1

        if is_consecutive and deps_on_last and last_has_single_successor:
            current_chain.append(node)
        else:
            if len(current_chain) > 1:
                chains.append(current_chain)
            current_chain = [node]

        i += 1

    # Add last chain if it exists
    if len(current_chain) > 1:
        chains.append(current_chain)

    # Create compressed dictionary
    compressed_dict = {}
    processed_nodes = set()

    # Add compressed chains
    for chain in chains:
        chain_name = f"{chain[0]}-{chain[-1]}"
        # Find dependencies of first node in chain
        deps = int_dict[chain[0]]
        compressed_dict[chain_name] = [str(d) for d in deps]
        processed_nodes.update(chain)

    # Add remaining non-chain nodes
    for node in nodes:
        if node not in processed_nodes:
            compressed_dict[str(node)] = [str(d) for d in int_dict[node]]

    # Update dependencies to use compressed names
    for k in compressed_dict:
        deps = compressed_dict[k]
        new_deps = []
        for dep in deps:
            dep_int = int(dep)
            # Find if this dependency is part of a chain
            chain_found = False
            for chain in chains:
                if dep_int in chain:
                    new_deps.append(f"{chain[0]}-{chain[-1]}")
                    chain_found = True
                    break
            if not chain_found:
                new_deps.append(dep)
        compressed_dict[k] = new_deps

    return compressed_dict

def compress_chains(G):
    """
    Compress chains in a directed graph by merging nodes that have single parent and single child.

    Args:
        G: NetworkX directed graph

    Returns:
        NetworkX directed graph with compressed chains
    """
    G_compressed = G.copy()
    processed_nodes = set()

    while True:
        # Find nodes with exactly one parent and one child
        nodes_to_compress = []
        for node in list(
            G_compressed.nodes()
        ):  # Create a list to avoid modification during iteration
            if node in processed_nodes:
                continue

            predecessors = list(G_compressed.predecessors(node))
            successors = list(G_compressed.successors(node))

            if len(predecessors) == 1 and len(successors) == 1:
                pred = predecessors[0]
                succ = successors[0]

                # Skip if any node in the chain is already processed
                if pred in processed_nodes or succ in processed_nodes:
                    continue

                # Only compress if middle node has single parent/child
                pred_preds = list(G_compressed.predecessors(pred))
                succ_succs = list(G_compressed.successors(succ))

                if len(pred_preds) <= 1 and len(succ_succs) <= 1:
                    nodes_to_compress.append((pred, node, succ))
                    processed_nodes.update([pred, node, succ])

        if not nodes_to_compress:
            break

        # Process each chain
        for pred, mid, succ in nodes_to_compress:
            if not all(G_compressed.has_node(n) for n in [pred, mid, succ]):
                continue

            # Create new merged node name
            new_node = ",".join(str(n) for n in [pred, mid, succ])

            # Add the new node
            G_compressed.add_node(new_node)

            # Add edges from all predecessors of first node
            for p in list(G_compressed.predecessors(pred)):
                G_compressed.add_edge(p, new_node)

            # Add edges to all successors of last node
            for s in list(G_compressed.successors(succ)):
                G_compressed.add_edge(new_node, s)

            # Remove the old nodes
            G_compressed.remove_nodes_from([pred, mid, succ])

    return G_compressed

# benchmark = DEFAULT_BENCHMARKS["webarena"]()
benchmark = DEFAULT_BENCHMARKS["visualwebarena"]()

dep_graph = benchmark.dependency_graph_over_tasks()
dep_graph = clean_dict(dep_graph)

dep_graph = compress_sequential_chains(dep_graph)
graph = dict_to_networkx(dep_graph)

# graph = compress_chains(graph)

components = nx.weakly_connected_components(graph)
components = [graph.subgraph(component).copy() for component in components]
plot_components_grid(components)
plt.show()


"""Dirty script to visualize the dependency graph of a benchmark, e.g. webarena, vsisualwebarena,
etc. You may have to detust it to make it work for you."""

import math

import bgym
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from bgym import DEFAULT_BENCHMARKS

def clean_dict(dependency_dict: dict[str, list[str]]) -> dict[str, list[str]]:
    new_dep = {}
    for key, deps in dependency_dict.items():
        new_key = key.split(".")[-1]

        new_dep[new_key] = [dep.split(".")[-1] for dep in deps]
    return new_dep

def dict_to_networkx(dependency_dict: dict[str, list[str]]) -> nx.DiGraph:

    G = nx.DiGraph()
    i = 0
    # Add edges from each node to its dependencies
    for node, dependencies in dependency_dict.items():
        i += 1
        if i > 20:
            pass

        print(node, dependencies)
        # Add edges from the node to each of its dependencies
        for dep in dependencies:
            G.add_edge(dep, node)
    return G

def plot_graph(G, ax, title=None, node_color="lightblue", node_size=40, font_size=8):


etc. You may have to detust it to make it work for you."""

import math

import bgym
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from bgym import DEFAULT_BENCHMARKS

def clean_dict(dependency_dict: dict[str, list[str]]) -> dict[str, list[str]]:
    new_dep = {}
    for key, deps in dependency_dict.items():
        new_key = key.split(".")[-1]

        new_dep[new_key] = [dep.split(".")[-1] for dep in deps]
    return new_dep

def dict_to_networkx(dependency_dict: dict[str, list[str]]) -> nx.DiGraph:

    G = nx.DiGraph()
    i = 0
    # Add edges from each node to its dependencies
    for node, dependencies in dependency_dict.items():
        i += 1
        if i > 20:
            pass

        print(node, dependencies)
        # Add edges from the node to each of its dependencies
        for dep in dependencies:
            G.add_edge(dep, node)
    return G

def plot_graph(G, ax, title=None, node_color="lightblue", node_size=40, font_size=8):
    """
    Plot a single graph component on the given matplotlib axis.

    Args:
        G: NetworkX graph (should be a single connected component)
        ax: Matplotlib axis to plot on
        title: Optional title for the subplot
        node_color: Color for the nodes
        node_size: Size of the nodes
        font_size: Size of the node labels
    """
    # Use a simple layout for better performance
    # pos = nx.spring_layout(G, k=0.1, iterations=100)

    pos = nx.kamada_kawai_layout(G)

    # pos = nx.spectral_layout(G)


    return new_dep

def dict_to_networkx(dependency_dict: dict[str, list[str]]) -> nx.DiGraph:

    G = nx.DiGraph()
    i = 0
    # Add edges from each node to its dependencies
    for node, dependencies in dependency_dict.items():
        i += 1
        if i > 20:
            pass

        print(node, dependencies)
        # Add edges from the node to each of its dependencies
        for dep in dependencies:
            G.add_edge(dep, node)
    return G

def plot_graph(G, ax, title=None, node_color="lightblue", node_size=40, font_size=8):
    """
    Plot a single graph component on the given matplotlib axis.

    Args:
        G: NetworkX graph (should be a single connected component)
        ax: Matplotlib axis to plot on
        title: Optional title for the subplot
        node_color: Color for the nodes
        node_size: Size of the nodes
        font_size: Size of the node labels
    """
    # Use a simple layout for better performance
    # pos = nx.spring_layout(G, k=0.1, iterations=100)

    pos = nx.kamada_kawai_layout(G)

    # pos = nx.spectral_layout(G)

    def name_to_size(name):
        if "-" in name:
            start, end = name.split("-")

            n_nodes = int(end) - int(start) + 1
        else:
            n_nodes = 1
        size_factor = node_size / 10
        return n_nodes * size_factor

    # compute size based on name
    sizes = [name_to_size(name) for name in G.nodes]

    nx.draw(
        G,
        pos,
        ax=ax,
        with_labels=True,
        node_color=node_color,
        node_size=sizes,
        font_size=font_size,
        font_weight="normal",
        arrows=True,
        arrowsize=15,
    )

    if title:
        ax.set_title(title)
    ax.axis("off")

def plot_components_grid(
    components, max_cols=4, node_color="lightblue", node_size=2000, font_size=10
):
    """
    Plot components in a grid layout.

    Args:
        components: List of NetworkX graphs, one per component
        max_cols: Maximum number of columns in the grid
        node_color: Color for the nodes
        node_size: Size of the nodes
        font_size: Size of the node labels

    Returns:
        matplotlib figure
    """
    n_components = len(components)


    sizes = [name_to_size(name) for name in G.nodes]

    nx.draw(
        G,
        pos,
        ax=ax,
        with_labels=True,
        node_color=node_color,
        node_size=sizes,
        font_size=font_size,
        font_weight="normal",
        arrows=True,
        arrowsize=15,
    )

    if title:
        ax.set_title(title)
    ax.axis("off")

def plot_components_grid(
    components, max_cols=4, node_color="lightblue", node_size=2000, font_size=10
):
    """
    Plot components in a grid layout.

    Args:
        components: List of NetworkX graphs, one per component
        max_cols: Maximum number of columns in the grid
        node_color: Color for the nodes
        node_size: Size of the nodes
        font_size: Size of the node labels

    Returns:
        matplotlib figure
    """
    n_components = len(components)

    if n_components == 0:
        print("No components found")
        return None

    # Calculate grid dimensions
    ncols = min(n_components, max_cols)
    nrows = math.ceil(n_components / ncols)

    # Create figure with a reasonable size per subplot
    fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 4 * nrows))
    fig.suptitle("Dependency Graph Components", size=16)

    # Make axes iterable even if there's only one
    if n_components == 1:
        axes = np.array([[axes]])
    elif nrows == 1:
        axes = np.array([axes])
    elif ncols == 1:
        axes = axes.reshape(-1, 1)

    # Plot each component
    for idx, component in enumerate(components):
        i, j = divmod(idx, ncols)
        title = f"Component {idx+1} ({component.number_of_nodes()} nodes)"
        plot_graph(
            component,
            axes[i, j],
            title,
            node_color=node_color,
            node_size=node_size,
            font_size=font_size,
        )

    # Remove empty subplots
    for idx in range(n_components, nrows * ncols):
        i, j = divmod(idx, ncols)
        axes[i, j].remove()

    plt.tight_layout()
    return fig

def compress_sequential_chains(dep_dict: dict[str, list[str]]) -> dict[str, list[str]]:
    """
    Compress chains of sequential numbers in a dependency dictionary.
    Returns a new dictionary with compressed chains using range notation.

    Args:
        dep_dict: Dictionary mapping string numbers to list of string number dependencies

    Returns:
        Dictionary with compressed chains using range notation
    """
    # Convert to integers for easier processing
    int_dict = {int(k): [int(x) for x in v] for k, v in dep_dict.items()}

    # Find chains
    chains = []
    current_chain = []


        i, j = divmod(idx, ncols)
        title = f"Component {idx+1} ({component.number_of_nodes()} nodes)"
        plot_graph(
            component,
            axes[i, j],
            title,
            node_color=node_color,
            node_size=node_size,
            font_size=font_size,
        )

    # Remove empty subplots
    for idx in range(n_components, nrows * ncols):
        i, j = divmod(idx, ncols)
        axes[i, j].remove()

    plt.tight_layout()
    return fig

def compress_sequential_chains(dep_dict: dict[str, list[str]]) -> dict[str, list[str]]:
    """
    Compress chains of sequential numbers in a dependency dictionary.
    Returns a new dictionary with compressed chains using range notation.

    Args:
        dep_dict: Dictionary mapping string numbers to list of string number dependencies

    Returns:
        Dictionary with compressed chains using range notation
    """
    # Convert to integers for easier processing
    int_dict = {int(k): [int(x) for x in v] for k, v in dep_dict.items()}

    # Find chains
    chains = []
    current_chain = []

    # Sort nodes for sequential processing
    nodes = sorted(int_dict.keys())

    i = 0
    while i < len(nodes):
        node = nodes[i]

        # Start new chain
        if not current_chain:
            current_chain = [node]
            i += 1
            continue

        # Check if this node continues the chain
        last_node = current_chain[-1]

        # Conditions for chain continuation:
        # 1. Numbers are consecutive
        # 2. Current node has exactly one dependency
        # 3. That dependency is the previous node in chain
        # 4. The previous node has exactly one successor
        is_consecutive = node == last_node + 1
        has_single_dep = len(int_dict[node]) == 1
        deps_on_last = has_single_dep and int_dict[node][0] == last_node
        last_has_single_successor = sum(1 for k, v in int_dict.items() if last_node in v) == 1

        if is_consecutive and deps_on_last and last_has_single_successor:
            current_chain.append(node)
        else:
            if len(current_chain) > 1:
                chains.append(current_chain)
            current_chain = [node]

        i += 1

    # Add last chain if it exists
    if len(current_chain) > 1:
        chains.append(current_chain)

    # Create compressed dictionary
    compressed_dict = {}
    processed_nodes = set()

    # Add compressed chains
    for chain in chains:
        chain_name = f"{chain[0]}-{chain[-1]}"
        # Find dependencies of first node in chain
        deps = int_dict[chain[0]]
        compressed_dict[chain_name] = [str(d) for d in deps]
        processed_nodes.update(chain)

    # Add remaining non-chain nodes
    for node in nodes:
        if node not in processed_nodes:
            compressed_dict[str(node)] = [str(d) for d in int_dict[node]]

    # Update dependencies to use compressed names
    for k in compressed_dict:
        deps = compressed_dict[k]
        new_deps = []
        for dep in deps:
            dep_int = int(dep)
            # Find if this dependency is part of a chain
            chain_found = False
            for chain in chains:
                if dep_int in chain:
                    new_deps.append(f"{chain[0]}-{chain[-1]}")
                    chain_found = True
                    break
            if not chain_found:
                new_deps.append(dep)
        compressed_dict[k] = new_deps

    return compressed_dict

def compress_chains(G):
    """
    Compress chains in a directed graph by merging nodes that have single parent and single child.

    Args:
        G: NetworkX directed graph

    Returns:
        NetworkX directed graph with compressed chains
    """
    G_compressed = G.copy()
    processed_nodes = set()

    while True:
        # Find nodes with exactly one parent and one child
        nodes_to_compress = []
        for node in list(
            G_compressed.nodes()


    # Update dependencies to use compressed names
    for k in compressed_dict:
        deps = compressed_dict[k]
        new_deps = []
        for dep in deps:
            dep_int = int(dep)
            # Find if this dependency is part of a chain
            chain_found = False
            for chain in chains:
                if dep_int in chain:
                    new_deps.append(f"{chain[0]}-{chain[-1]}")
                    chain_found = True
                    break
            if not chain_found:
                new_deps.append(dep)
        compressed_dict[k] = new_deps

    return compressed_dict

def compress_chains(G):
    """
    Compress chains in a directed graph by merging nodes that have single parent and single child.

    Args:
        G: NetworkX directed graph

    Returns:
        NetworkX directed graph with compressed chains
    """
    G_compressed = G.copy()
    processed_nodes = set()

    while True:
        # Find nodes with exactly one parent and one child
        nodes_to_compress = []
        for node in list(
            G_compressed.nodes()
        ):  # Create a list to avoid modification during iteration
            if node in processed_nodes:
                continue

            predecessors = list(G_compressed.predecessors(node))
            successors = list(G_compressed.successors(node))

            if len(predecessors) == 1 and len(successors) == 1:
                pred = predecessors[0]
                succ = successors[0]

                # Skip if any node in the chain is already processed
                if pred in processed_nodes or succ in processed_nodes:
                    continue

                # Only compress if middle node has single parent/child
                pred_preds = list(G_compressed.predecessors(pred))
                succ_succs = list(G_compressed.successors(succ))

                if len(pred_preds) <= 1 and len(succ_succs) <= 1:
                    nodes_to_compress.append((pred, node, succ))
                    processed_nodes.update([pred, node, succ])

        if not nodes_to_compress:
            break

        # Process each chain
        for pred, mid, succ in nodes_to_compress:
            if not all(G_compressed.has_node(n) for n in [pred, mid, succ]):
                continue

            # Create new merged node name
            new_node = ",".join(str(n) for n in [pred, mid, succ])

            # Add the new node
            G_compressed.add_node(new_node)

            # Add edges from all predecessors of first node
            for p in list(G_compressed.predecessors(pred)):
                G_compressed.add_edge(p, new_node)

            # Add edges to all successors of last node
            for s in list(G_compressed.successors(succ)):
                G_compressed.add_edge(new_node, s)

            # Remove the old nodes
            G_compressed.remove_nodes_from([pred, mid, succ])

    return G_compressed

# benchmark = DEFAULT_BENCHMARKS["webarena"]()
benchmark = DEFAULT_BENCHMARKS["visualwebarena"]()

dep_graph = benchmark.dependency_graph_over_tasks()
dep_graph = clean_dict(dep_graph)

dep_graph = compress_sequential_chains(dep_graph)
graph = dict_to_networkx(dep_graph)

# graph = compress_chains(graph)

components = nx.weakly_connected_components(graph)
components = [graph.subgraph(component).copy() for component in components]
plot_components_grid(components)
plt.show()


def plot_graph(G, ax, title=None, node_color="lightblue", node_size=40, font_size=8):
    """
    Plot a single graph component on the given matplotlib axis.

    Args:
        G: NetworkX graph (should be a single connected component)
        ax: Matplotlib axis to plot on
        title: Optional title for the subplot
        node_color: Color for the nodes
        node_size: Size of the nodes
        font_size: Size of the node labels
    """
    # Use a simple layout for better performance
    # pos = nx.spring_layout(G, k=0.1, iterations=100)

    pos = nx.kamada_kawai_layout(G)

    # pos = nx.spectral_layout(G)

    def name_to_size(name):
        if "-" in name:
            start, end = name.split("-")

            n_nodes = int(end) - int(start) + 1
        else:
            n_nodes = 1
        size_factor = node_size / 10
        return n_nodes * size_factor

    # compute size based on name
    sizes = [name_to_size(name) for name in G.nodes]

    nx.draw(
        G,
        pos,
        ax=ax,
        with_labels=True,
        node_color=node_color,
        node_size=sizes,
        font_size=font_size,
        font_weight="normal",
        arrows=True,
        arrowsize=15,
    )

    if title:
        ax.set_title(title)
    ax.axis("off")


import csv
import logging
import os
import platform
from datetime import datetime
from importlib import metadata
from pathlib import Path

import bgym
import pandas as pd
from bgym import Benchmark
from git import InvalidGitRepositoryError, Repo
from git.config import GitConfigParser

import agentlab
from agentlab.experiments.exp_utils import RESULTS_DIR

def _get_repo(module):
    return Repo(Path(module.__file__).resolve().parent, search_parent_directories=True)

def _get_benchmark_version(
    benchmark: Benchmark, allow_bypass_benchmark_version: bool = False
) -> str:
    benchmark_name = benchmark.name

    if hasattr(benchmark, "get_version"):
        return benchmark.get_version()

    # in between 2 pull requests
    if benchmark_name.startswith("miniwob"):
        return metadata.distribution("browsergym.miniwob").version
    elif benchmark_name.startswith("workarena"):
        return metadata.distribution("browsergym.workarena").version
    elif benchmark_name.startswith("webarena"):
        return metadata.distribution("browsergym.webarena").version
    elif benchmark_name.startswith("visualwebarena"):
        return metadata.distribution("browsergym.visualwebarena").version
    elif benchmark_name.startswith("weblinx"):
        try:
            return metadata.distribution("weblinx_browsergym").version
        except metadata.PackageNotFoundError:
            return "0.0.1rc1"
    elif benchmark_name.startswith("assistantbench"):
        return metadata.distribution("browsergym.assistantbench").version
    else:
        if allow_bypass_benchmark_version:
            return "bypassed"
        else:
            raise ValueError(f"Unknown benchmark {benchmark_name}")

def _get_git_username(repo: Repo) -> str:
    """
    Retrieves the first available Git username from various sources.

    Note: overlycomplex designed by Claude and not fully tested.

    This function checks multiple locations for the Git username in the following order:
    1. Repository-specific configuration
    2. GitHub API (if the remote is a GitHub repository)
    3. Global Git configuration
    4. System Git configuration
    5. Environment variables (GIT_AUTHOR_NAME and GIT_COMMITTER_NAME)

    Args:
        repo (Repo): A GitPython Repo object representing the Git repository.

    Returns:
        str: The first non-None username found, or None if no username is found.
    """
    # Repository-specific configuration
    try:
        username = repo.config_reader().get_value("user", "name", None)
        if username:
            return username
    except Exception:
        pass

    try:
        # GitHub username
        remote_url = repo.remotes.origin.url
        if "github.com" in remote_url:
            import json
            import re
            import urllib.request

            match = re.search(r"github\.com[:/](.+)/(.+)\.git", remote_url)
            if match:
                owner, repo_name = match.groups()
                api_url = f"https://api.github.com/repos/{owner}/{repo_name}"
                with urllib.request.urlopen(api_url) as response:
                    data = json.loads(response.read().decode())
                    username = data["owner"]["login"]
                    if username:
                        return username
    except Exception:
        pass

    try:
        # Global configuration
        username = GitConfigParser(repo.git.config("--global", "--list"), read_only=True).get_value(
            "user", "name", None
        )
        if username:
            return username
    except Exception:
        pass

    try:
        # System configuration
        username = GitConfigParser(repo.git.config("--system", "--list"), read_only=True).get_value(
            "user", "name", None
        )
        if username:
            return username
    except Exception:
        pass

    # Environment variables
    return os.environ.get("GIT_AUTHOR_NAME") or os.environ.get("GIT_COMMITTER_NAME")

def _get_git_info(module, changes_white_list=()) -> tuple[str, list[tuple[str, Path]]]:
    """
    Retrieve comprehensive git information for the given module.

    This function attempts to find the git repository containing the specified
    module and returns the current commit hash and a comprehensive list of all
    files that contribute to the repository's state.

    Args:
        module: The Python module object to check for git information.
        changes_white_list: A list of file paths to ignore when checking for changes.

    Returns:
        tuple: A tuple containing two elements:
            - str or None: The current git commit hash, or None if not a git repo.
            - list of tuple: A list of (status, Path) tuples for all modified files.
              Empty list if not a git repo. Status can be 'M' (modified), 'A' (added),
              'D' (deleted), 'R' (renamed), 'C' (copied), 'U' (updated but unmerged),
              or '??' (untracked).
    """

    try:
        repo = _get_repo(module)

        git_hash = repo.head.object.hexsha

        modified_files = []

        # Staged changes
        staged_changes = repo.index.diff(repo.head.commit)
        for change in staged_changes:
            modified_files.append((change.change_type, Path(change.a_path)))

        # Unstaged changes
        unstaged_changes = repo.index.diff(None)
        for change in unstaged_changes:
            modified_files.append((change.change_type, Path(change.a_path)))

        # Untracked files
        untracked_files = repo.untracked_files
        for file in untracked_files:
            modified_files.append(("??", Path(file)))

        # wildcard matching from white list
        modified_files_filtered = []
        for status, file in modified_files:
            if any(file.match(pattern) for pattern in changes_white_list):
                continue
            modified_files_filtered.append((status, file))

        return git_hash, modified_files_filtered
    except InvalidGitRepositoryError:
        return None, []

def get_reproducibility_info(
    agent_names: str | list[str],
    benchmark: Benchmark,
    study_id: str = "",
    comment=None,
    changes_white_list=(  # Files that are often modified during experiments but do not affect reproducibility
        "*/reproducibility_script.py",
        "*reproducibility_journal.csv",
        "*main.py",
        "*inspect_results.ipynb",
    ),
    ignore_changes=False,
    allow_bypass_benchmark_version=False,
):
    """
    Retrieve a dict of information that could influence the reproducibility of an experiment.
    """
    from browsergym import core

    import agentlab

    if isinstance(agent_names, str):
        agent_names = [agent_names]

    try:
        repo = _get_repo(agentlab)
    except InvalidGitRepositoryError:
        repo = None

    info = {
        "git_user": _get_git_username(repo),
        "agent_names": agent_names,
        "benchmark": benchmark.name,
        "study_id": study_id,
        "comment": comment,
        "benchmark_version": _get_benchmark_version(benchmark, allow_bypass_benchmark_version),
        "date": datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
        "os": f"{platform.system()} ({platform.version()})",
        "python_version": platform.python_version(),
        "playwright_version": metadata.distribution("playwright").version,
    }

    def add_git_info(module_name, module):
        git_hash, modified_files = _get_git_info(module, changes_white_list)

        modified_files_str = "\n".join([f"  {status}: {file}" for status, file in modified_files])

        if len(modified_files) > 0:
            msg = (
                f"Module {module_name} has uncommitted changes. "
                f"Modified files:  \n{modified_files_str}\n"
            )
            if ignore_changes:
                logging.warning(
                    msg + "Ignoring changes as requested and proceeding to experiments."
                )
            else:
                raise ValueError(
                    msg + "Please commit or stash your changes before running the experiment."
                )

        info[f"{module_name}_version"] = module.__version__
        info[f"{module_name}_git_hash"] = git_hash
        info[f"{module_name}__local_modifications"] = modified_files_str

    add_git_info("agentlab", agentlab)
    add_git_info("browsergym", core)
    return info

def assert_compatible(info: dict, old_info: dict, raise_if_incompatible=True):
    """Make sure that the two info dicts are compatible."""
    # TODO may need to adapt if there are multiple agents, and the re-run on
    # error only has a subset of agents. Hence old_info.agent_name != info.agent_name
    for key in info.keys():
        if key in ("date", "avg_reward", "std_err", "n_completed", "n_err"):
            continue
        if info[key] != old_info[key]:
            _raise_or_warn(
                f"Reproducibility info already exist and is not compatible."
                f"Key {key} has changed from {old_info[key]} to {info[key]}."
                f"Set strict_reproducibility=False to bypass this error.",
                raise_error=raise_if_incompatible,
            )

def _raise_or_warn(msg, raise_error=True):
    if raise_error:
        raise ValueError(msg)
    else:
        logging.warning(msg)

def _verify_report(report_df: pd.DataFrame, agent_names=list[str], strict_reproducibility=True):

    report_df = report_df.reset_index()

    unique_agent_names = report_df["agent.agent_name"].unique()
    if set(agent_names) != set(unique_agent_names):
        raise ValueError(
            f"Agent names in the report {unique_agent_names} do not match the agent names {agent_names}."
        )
    if len(set(agent_names)) != len(agent_names):
        raise ValueError(f"Duplicate agent names {agent_names}.")

    report_df = report_df.set_index("agent.agent_name", inplace=False)

    for idx in report_df.index:
        n_err = report_df.loc[idx, "n_err"].item()
        n_completed, n_total = report_df.loc[idx, "n_completed"].split("/")
        if n_err > 0:
            _raise_or_warn(
                f"Experiment {idx} has {n_err} errors. Please rerun the study and make sure all tasks are completed.",
                raise_error=strict_reproducibility,
            )
        if n_completed != n_total:
            _raise_or_warn(
                f"Experiment {idx} has {n_completed} completed tasks out of {n_total}. "
                f"Please rerun the study and make sure all tasks are completed.",
                raise_error=strict_reproducibility,
            )
    return report_df

def _get_csv_headers(file_path: str) -> list[str]:
    with open(file_path, "r", newline="") as file:
        reader = csv.reader(file)
        try:
            headers = next(reader)
        except StopIteration:
            headers = None
    return headers

def _add_result_to_info(info: dict, report_df: pd.DataFrame):
    """Extracts the results from the report and adds them to the info dict inplace"""

    for key in ("avg_reward", "std_err", "n_err", "n_completed"):
        value = report_df.loc[info["agent_name"], key]
        if hasattr(value, "item"):
            value = value.item()
        info[key] = value

def append_to_journal(
    info, report_df: pd.DataFrame, journal_path=None, strict_reproducibility=True
):
    """Append the info and results to the reproducibility journal."""
    if journal_path is None:
        try:
            _get_repo(agentlab)  # if not based on git clone, this will raise an error
            journal_path = (
                Path(agentlab.__file__).parent.parent.parent / "reproducibility_journal.csv"
            )
        except InvalidGitRepositoryError:
            logging.warning(
                "Could not find a git repository. Saving the journal to the results directory."
                "To add to the journal, git clone agentlab and use `pip install -e .`"
            )
            journal_path = RESULTS_DIR / "reproducibility_journal.csv"

    logging.info(f"Appending to journal {journal_path}")

    if len(report_df) != len(info["agent_names"]):
        raise ValueError(
            "Mismatch between the number of agents in reproducibility info and the summary report."
        )

    report_df = _verify_report(
        report_df, info["agent_names"], strict_reproducibility=strict_reproducibility
    )

    rows = []
    headers = None
    if journal_path.exists():
        headers = _get_csv_headers(journal_path)

    if headers is None:  # first creation
        headers = list(info.keys())
        headers[headers.index("agent_names")] = "agent_name"
        rows.append(headers)

    for agent_name in info["agent_names"]:
        info_copy = info.copy()
        del info_copy["agent_names"]
        info_copy["agent_name"] = agent_name

        _add_result_to_info(info_copy, report_df)

        rows.append([str(info_copy[key]) for key in headers])

    with open(journal_path, "a", newline="") as file:
        writer = csv.writer(file)
        for row in rows:
            writer.writerow(row)


import csv
import logging
import os
import platform
from datetime import datetime
from importlib import metadata
from pathlib import Path

import bgym
import pandas as pd
from bgym import Benchmark
from git import InvalidGitRepositoryError, Repo
from git.config import GitConfigParser

import agentlab
from agentlab.experiments.exp_utils import RESULTS_DIR

def _get_repo(module):
    return Repo(Path(module.__file__).resolve().parent, search_parent_directories=True)

def _get_benchmark_version(
    benchmark: Benchmark, allow_bypass_benchmark_version: bool = False
) -> str:
    benchmark_name = benchmark.name

    if hasattr(benchmark, "get_version"):
        return benchmark.get_version()

    # in between 2 pull requests
    if benchmark_name.startswith("miniwob"):
        return metadata.distribution("browsergym.miniwob").version
    elif benchmark_name.startswith("workarena"):
        return metadata.distribution("browsergym.workarena").version
    elif benchmark_name.startswith("webarena"):
        return metadata.distribution("browsergym.webarena").version
    elif benchmark_name.startswith("visualwebarena"):
        return metadata.distribution("browsergym.visualwebarena").version
    elif benchmark_name.startswith("weblinx"):


import os
import platform
from datetime import datetime
from importlib import metadata
from pathlib import Path

import bgym
import pandas as pd
from bgym import Benchmark
from git import InvalidGitRepositoryError, Repo
from git.config import GitConfigParser

import agentlab
from agentlab.experiments.exp_utils import RESULTS_DIR

def _get_repo(module):
    return Repo(Path(module.__file__).resolve().parent, search_parent_directories=True)

def _get_benchmark_version(
    benchmark: Benchmark, allow_bypass_benchmark_version: bool = False
) -> str:
    benchmark_name = benchmark.name

    if hasattr(benchmark, "get_version"):
        return benchmark.get_version()

    # in between 2 pull requests
    if benchmark_name.startswith("miniwob"):
        return metadata.distribution("browsergym.miniwob").version
    elif benchmark_name.startswith("workarena"):
        return metadata.distribution("browsergym.workarena").version
    elif benchmark_name.startswith("webarena"):
        return metadata.distribution("browsergym.webarena").version
    elif benchmark_name.startswith("visualwebarena"):
        return metadata.distribution("browsergym.visualwebarena").version
    elif benchmark_name.startswith("weblinx"):
        try:
            return metadata.distribution("weblinx_browsergym").version
        except metadata.PackageNotFoundError:
            return "0.0.1rc1"
    elif benchmark_name.startswith("assistantbench"):
        return metadata.distribution("browsergym.assistantbench").version
    else:
        if allow_bypass_benchmark_version:
            return "bypassed"
        else:
            raise ValueError(f"Unknown benchmark {benchmark_name}")

def _get_git_username(repo: Repo) -> str:
    """
    Retrieves the first available Git username from various sources.

    Note: overlycomplex designed by Claude and not fully tested.

    This function checks multiple locations for the Git username in the following order:
    1. Repository-specific configuration
    2. GitHub API (if the remote is a GitHub repository)
    3. Global Git configuration
    4. System Git configuration
    5. Environment variables (GIT_AUTHOR_NAME and GIT_COMMITTER_NAME)

    Args:
        repo (Repo): A GitPython Repo object representing the Git repository.

    Returns:
        str: The first non-None username found, or None if no username is found.


    elif benchmark_name.startswith("workarena"):
        return metadata.distribution("browsergym.workarena").version
    elif benchmark_name.startswith("webarena"):
        return metadata.distribution("browsergym.webarena").version
    elif benchmark_name.startswith("visualwebarena"):
        return metadata.distribution("browsergym.visualwebarena").version
    elif benchmark_name.startswith("weblinx"):
        try:
            return metadata.distribution("weblinx_browsergym").version
        except metadata.PackageNotFoundError:
            return "0.0.1rc1"
    elif benchmark_name.startswith("assistantbench"):
        return metadata.distribution("browsergym.assistantbench").version
    else:
        if allow_bypass_benchmark_version:
            return "bypassed"
        else:
            raise ValueError(f"Unknown benchmark {benchmark_name}")

def _get_git_username(repo: Repo) -> str:
    """
    Retrieves the first available Git username from various sources.

    Note: overlycomplex designed by Claude and not fully tested.

    This function checks multiple locations for the Git username in the following order:
    1. Repository-specific configuration
    2. GitHub API (if the remote is a GitHub repository)
    3. Global Git configuration
    4. System Git configuration
    5. Environment variables (GIT_AUTHOR_NAME and GIT_COMMITTER_NAME)

    Args:
        repo (Repo): A GitPython Repo object representing the Git repository.

    Returns:
        str: The first non-None username found, or None if no username is found.
    """
    # Repository-specific configuration
    try:
        username = repo.config_reader().get_value("user", "name", None)
        if username:
            return username
    except Exception:
        pass

    try:
        # GitHub username
        remote_url = repo.remotes.origin.url
        if "github.com" in remote_url:
            import json
            import re
            import urllib.request

            match = re.search(r"github\.com[:/](.+)/(.+)\.git", remote_url)
            if match:
                owner, repo_name = match.groups()
                api_url = f"https://api.github.com/repos/{owner}/{repo_name}"
                with urllib.request.urlopen(api_url) as response:
                    data = json.loads(response.read().decode())
                    username = data["owner"]["login"]
                    if username:
                        return username
    except Exception:
        pass

    try:
        # Global configuration
        username = GitConfigParser(repo.git.config("--global", "--list"), read_only=True).get_value(
            "user", "name", None
        )
        if username:
            return username
    except Exception:
        pass

    try:
        # System configuration
        username = GitConfigParser(repo.git.config("--system", "--list"), read_only=True).get_value(
            "user", "name", None
        )
        if username:
            return username
    except Exception:
        pass

    # Environment variables
    return os.environ.get("GIT_AUTHOR_NAME") or os.environ.get("GIT_COMMITTER_NAME")

def _get_git_info(module, changes_white_list=()) -> tuple[str, list[tuple[str, Path]]]:
    """
    Retrieve comprehensive git information for the given module.

    This function attempts to find the git repository containing the specified
    module and returns the current commit hash and a comprehensive list of all
    files that contribute to the repository's state.

    Args:
        module: The Python module object to check for git information.
        changes_white_list: A list of file paths to ignore when checking for changes.

    Returns:
        tuple: A tuple containing two elements:
            - str or None: The current git commit hash, or None if not a git repo.
            - list of tuple: A list of (status, Path) tuples for all modified files.
              Empty list if not a git repo. Status can be 'M' (modified), 'A' (added),
              'D' (deleted), 'R' (renamed), 'C' (copied), 'U' (updated but unmerged),


        )
        if username:
            return username
    except Exception:
        pass

    try:
        # System configuration
        username = GitConfigParser(repo.git.config("--system", "--list"), read_only=True).get_value(
            "user", "name", None
        )
        if username:
            return username
    except Exception:
        pass

    # Environment variables
    return os.environ.get("GIT_AUTHOR_NAME") or os.environ.get("GIT_COMMITTER_NAME")

def _get_git_info(module, changes_white_list=()) -> tuple[str, list[tuple[str, Path]]]:
    """
    Retrieve comprehensive git information for the given module.

    This function attempts to find the git repository containing the specified
    module and returns the current commit hash and a comprehensive list of all
    files that contribute to the repository's state.

    Args:
        module: The Python module object to check for git information.
        changes_white_list: A list of file paths to ignore when checking for changes.

    Returns:
        tuple: A tuple containing two elements:
            - str or None: The current git commit hash, or None if not a git repo.
            - list of tuple: A list of (status, Path) tuples for all modified files.
              Empty list if not a git repo. Status can be 'M' (modified), 'A' (added),
              'D' (deleted), 'R' (renamed), 'C' (copied), 'U' (updated but unmerged),
              or '??' (untracked).
    """

    try:
        repo = _get_repo(module)

        git_hash = repo.head.object.hexsha

        modified_files = []

        # Staged changes
        staged_changes = repo.index.diff(repo.head.commit)
        for change in staged_changes:
            modified_files.append((change.change_type, Path(change.a_path)))

        # Unstaged changes
        unstaged_changes = repo.index.diff(None)
        for change in unstaged_changes:
            modified_files.append((change.change_type, Path(change.a_path)))

        # Untracked files
        untracked_files = repo.untracked_files
        for file in untracked_files:
            modified_files.append(("??", Path(file)))

        # wildcard matching from white list
        modified_files_filtered = []
        for status, file in modified_files:
            if any(file.match(pattern) for pattern in changes_white_list):
                continue
            modified_files_filtered.append((status, file))

        return git_hash, modified_files_filtered
    except InvalidGitRepositoryError:
        return None, []

def get_reproducibility_info(
    agent_names: str | list[str],
    benchmark: Benchmark,
    study_id: str = "",
    comment=None,
    changes_white_list=(  # Files that are often modified during experiments but do not affect reproducibility
        "*/reproducibility_script.py",
        "*reproducibility_journal.csv",
        "*main.py",
        "*inspect_results.ipynb",
    ),
    ignore_changes=False,
    allow_bypass_benchmark_version=False,
):
    """
    Retrieve a dict of information that could influence the reproducibility of an experiment.
    """
    from browsergym import core


        for change in unstaged_changes:
            modified_files.append((change.change_type, Path(change.a_path)))

        # Untracked files
        untracked_files = repo.untracked_files
        for file in untracked_files:
            modified_files.append(("??", Path(file)))

        # wildcard matching from white list
        modified_files_filtered = []
        for status, file in modified_files:
            if any(file.match(pattern) for pattern in changes_white_list):
                continue
            modified_files_filtered.append((status, file))

        return git_hash, modified_files_filtered
    except InvalidGitRepositoryError:
        return None, []

def get_reproducibility_info(
    agent_names: str | list[str],
    benchmark: Benchmark,
    study_id: str = "",
    comment=None,
    changes_white_list=(  # Files that are often modified during experiments but do not affect reproducibility
        "*/reproducibility_script.py",
        "*reproducibility_journal.csv",
        "*main.py",
        "*inspect_results.ipynb",
    ),
    ignore_changes=False,
    allow_bypass_benchmark_version=False,
):
    """
    Retrieve a dict of information that could influence the reproducibility of an experiment.
    """
    from browsergym import core

    import agentlab

    if isinstance(agent_names, str):
        agent_names = [agent_names]

    try:
        repo = _get_repo(agentlab)
    except InvalidGitRepositoryError:
        repo = None

    info = {
        "git_user": _get_git_username(repo),
        "agent_names": agent_names,
        "benchmark": benchmark.name,
        "study_id": study_id,
        "comment": comment,
        "benchmark_version": _get_benchmark_version(benchmark, allow_bypass_benchmark_version),
        "date": datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
        "os": f"{platform.system()} ({platform.version()})",
        "python_version": platform.python_version(),
        "playwright_version": metadata.distribution("playwright").version,
    }

    def add_git_info(module_name, module):
        git_hash, modified_files = _get_git_info(module, changes_white_list)

        modified_files_str = "\n".join([f"  {status}: {file}" for status, file in modified_files])

        if len(modified_files) > 0:
            msg = (
                f"Module {module_name} has uncommitted changes. "
                f"Modified files:  \n{modified_files_str}\n"
            )
            if ignore_changes:
                logging.warning(
                    msg + "Ignoring changes as requested and proceeding to experiments."
                )
            else:
                raise ValueError(
                    msg + "Please commit or stash your changes before running the experiment."
                )

        info[f"{module_name}_version"] = module.__version__
        info[f"{module_name}_git_hash"] = git_hash
        info[f"{module_name}__local_modifications"] = modified_files_str

    add_git_info("agentlab", agentlab)
    add_git_info("browsergym", core)
    return info

def assert_compatible(info: dict, old_info: dict, raise_if_incompatible=True):
    """Make sure that the two info dicts are compatible."""
    # TODO may need to adapt if there are multiple agents, and the re-run on
    # error only has a subset of agents. Hence old_info.agent_name != info.agent_name
    for key in info.keys():
        if key in ("date", "avg_reward", "std_err", "n_completed", "n_err"):
            continue
        if info[key] != old_info[key]:
            _raise_or_warn(
                f"Reproducibility info already exist and is not compatible."
                f"Key {key} has changed from {old_info[key]} to {info[key]}."
                f"Set strict_reproducibility=False to bypass this error.",
                raise_error=raise_if_incompatible,
            )

def _raise_or_warn(msg, raise_error=True):
    if raise_error:


                f"Modified files:  \n{modified_files_str}\n"
            )
            if ignore_changes:
                logging.warning(
                    msg + "Ignoring changes as requested and proceeding to experiments."
                )
            else:
                raise ValueError(
                    msg + "Please commit or stash your changes before running the experiment."
                )

        info[f"{module_name}_version"] = module.__version__
        info[f"{module_name}_git_hash"] = git_hash
        info[f"{module_name}__local_modifications"] = modified_files_str

    add_git_info("agentlab", agentlab)
    add_git_info("browsergym", core)
    return info

def assert_compatible(info: dict, old_info: dict, raise_if_incompatible=True):
    """Make sure that the two info dicts are compatible."""
    # TODO may need to adapt if there are multiple agents, and the re-run on
    # error only has a subset of agents. Hence old_info.agent_name != info.agent_name
    for key in info.keys():
        if key in ("date", "avg_reward", "std_err", "n_completed", "n_err"):
            continue
        if info[key] != old_info[key]:
            _raise_or_warn(
                f"Reproducibility info already exist and is not compatible."
                f"Key {key} has changed from {old_info[key]} to {info[key]}."
                f"Set strict_reproducibility=False to bypass this error.",
                raise_error=raise_if_incompatible,
            )

def _raise_or_warn(msg, raise_error=True):
    if raise_error:
        raise ValueError(msg)
    else:
        logging.warning(msg)

def _verify_report(report_df: pd.DataFrame, agent_names=list[str], strict_reproducibility=True):

    report_df = report_df.reset_index()

    unique_agent_names = report_df["agent.agent_name"].unique()
    if set(agent_names) != set(unique_agent_names):
        raise ValueError(
            f"Agent names in the report {unique_agent_names} do not match the agent names {agent_names}."
        )
    if len(set(agent_names)) != len(agent_names):
        raise ValueError(f"Duplicate agent names {agent_names}.")


    add_git_info("browsergym", core)
    return info

def assert_compatible(info: dict, old_info: dict, raise_if_incompatible=True):
    """Make sure that the two info dicts are compatible."""
    # TODO may need to adapt if there are multiple agents, and the re-run on
    # error only has a subset of agents. Hence old_info.agent_name != info.agent_name
    for key in info.keys():
        if key in ("date", "avg_reward", "std_err", "n_completed", "n_err"):
            continue
        if info[key] != old_info[key]:
            _raise_or_warn(
                f"Reproducibility info already exist and is not compatible."
                f"Key {key} has changed from {old_info[key]} to {info[key]}."
                f"Set strict_reproducibility=False to bypass this error.",
                raise_error=raise_if_incompatible,
            )

def _raise_or_warn(msg, raise_error=True):
    if raise_error:
        raise ValueError(msg)
    else:
        logging.warning(msg)

def _verify_report(report_df: pd.DataFrame, agent_names=list[str], strict_reproducibility=True):

    report_df = report_df.reset_index()

    unique_agent_names = report_df["agent.agent_name"].unique()
    if set(agent_names) != set(unique_agent_names):
        raise ValueError(
            f"Agent names in the report {unique_agent_names} do not match the agent names {agent_names}."
        )
    if len(set(agent_names)) != len(agent_names):
        raise ValueError(f"Duplicate agent names {agent_names}.")

    report_df = report_df.set_index("agent.agent_name", inplace=False)

    for idx in report_df.index:
        n_err = report_df.loc[idx, "n_err"].item()
        n_completed, n_total = report_df.loc[idx, "n_completed"].split("/")
        if n_err > 0:


    # error only has a subset of agents. Hence old_info.agent_name != info.agent_name
    for key in info.keys():
        if key in ("date", "avg_reward", "std_err", "n_completed", "n_err"):
            continue
        if info[key] != old_info[key]:
            _raise_or_warn(
                f"Reproducibility info already exist and is not compatible."
                f"Key {key} has changed from {old_info[key]} to {info[key]}."
                f"Set strict_reproducibility=False to bypass this error.",
                raise_error=raise_if_incompatible,
            )

def _raise_or_warn(msg, raise_error=True):
    if raise_error:
        raise ValueError(msg)
    else:
        logging.warning(msg)

def _verify_report(report_df: pd.DataFrame, agent_names=list[str], strict_reproducibility=True):

    report_df = report_df.reset_index()

    unique_agent_names = report_df["agent.agent_name"].unique()
    if set(agent_names) != set(unique_agent_names):
        raise ValueError(
            f"Agent names in the report {unique_agent_names} do not match the agent names {agent_names}."
        )
    if len(set(agent_names)) != len(agent_names):
        raise ValueError(f"Duplicate agent names {agent_names}.")

    report_df = report_df.set_index("agent.agent_name", inplace=False)

    for idx in report_df.index:
        n_err = report_df.loc[idx, "n_err"].item()
        n_completed, n_total = report_df.loc[idx, "n_completed"].split("/")
        if n_err > 0:
            _raise_or_warn(
                f"Experiment {idx} has {n_err} errors. Please rerun the study and make sure all tasks are completed.",
                raise_error=strict_reproducibility,
            )
        if n_completed != n_total:
            _raise_or_warn(
                f"Experiment {idx} has {n_completed} completed tasks out of {n_total}. "
                f"Please rerun the study and make sure all tasks are completed.",
                raise_error=strict_reproducibility,
            )
    return report_df

def _get_csv_headers(file_path: str) -> list[str]:
    with open(file_path, "r", newline="") as file:
        reader = csv.reader(file)
        try:
            headers = next(reader)
        except StopIteration:
            headers = None
    return headers

def _add_result_to_info(info: dict, report_df: pd.DataFrame):
    """Extracts the results from the report and adds them to the info dict inplace"""

    for key in ("avg_reward", "std_err", "n_err", "n_completed"):
        value = report_df.loc[info["agent_name"], key]
        if hasattr(value, "item"):
            value = value.item()
        info[key] = value