import asyncio
import copy
import json
import os
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable, Optional
from uuid import uuid4

from agents.apply_diff import apply_diff
from tqdm.auto import tqdm

from relai import AsyncRELAI
from relai._client import is_context_length_exceeded_error
from relai.critico.critico import Critico, CriticoLog
from relai.simulator import AgentLog, AsyncAgent, AsyncSimulator
from relai.utils import no_trace

from ..schema.visual import ConfigOptVizSchema, ConfigSchema, GraphOptVizSchema, ParamSchema, RunSchema
from .graph import get_current_param_graph
from .params import get_current_params
from .utils import ProportionalSampler, extract_code, get_full_func_name

AgentLogCompressor = Callable[[AgentLog], AgentLog]


class Maestro:
    """
    Maestro automatically optimizes an AI agent to maximize its Critico score, navigating the space of
    configurations to intelligently improve performance on the chosen criteria.
    """

    def __init__(
        self,
        client: AsyncRELAI,
        agent_fn: AsyncAgent,
        goal: Optional[str] = None,
        max_memory: int = 20,
        name: str = "No Name",
        log_to_platform: bool = True,
    ):
        """
        Args:
            client (AsyncRELAI): An instance of the AsyncRELAI client to interact with the RELAI platform.
            agent_fn (AsyncAgent): The agent function to be optimized.
            goal (str, optional): Optional description of the goal of optimization. If None, increasing evaluation score
                will be considered as the only goal. Defaults to None.
            max_memory (int, optional): Control the maximum number of previous optimization history visible at each
                optimization step. Defaults to 20.
            name (str, optional): Name of the configuration optimization visualization on RELAI platform.
                Defaults to "No Name".
            log_to_platform (bool): Whether to log optimization progress and results on RELAI platform.
                Defaults to True.
        """
        self.setups: list[dict] = []
        self.agent_fn: AsyncAgent = agent_fn
        self.log: list[dict] = []
        self.max_memory: int = max_memory
        self._client: AsyncRELAI = client
        self.goal: str = goal if goal is not None else "Higher scores"
        self.log_to_platform: bool = log_to_platform
        self.config_opt_viz_id: str | None = None
        self.name: str = name

        self.versions: list[dict] = [
            {
                "idx": 0,
                "params": copy.deepcopy(get_current_params()),
                "log": copy.deepcopy(self.log),
                "visits": 0,
                "average_score": 0.0,
                "utc_timestamp": datetime.now(timezone.utc).isoformat(),
                "validation_runs": [],
                "validation_count": 0,
            }
        ]
        self.current_version: int = 0
        self.total_visits: int = 0
        self.inf: float = 1e208
        self.score_min: float = self.inf
        self.score_max: float = -self.inf

    def add_setup(self, simulator: AsyncSimulator, critico: Critico, weight: float = 1):
        """
        Add a new setup consisting of a simulator and a critico to Maestro.

        Args:
            simulator (AsyncSimulator): An AsyncSimulator to run the agent in the new setup.
            critico (Critico): A Critico with evaluators for the new setup.
            weight (float): A positive float representing the weight of this setup in comparson to others.
                Defaults to 1.
        """

        if not isinstance(simulator, AsyncSimulator):
            raise ValueError("Currently only AsyncSimulator is supported in Maestro.")
        if weight <= 0:
            raise ValueError(f"`weight` must be a positive float, got {weight}.")

        self.setups.append({"simulator": simulator, "critico": critico, "weight": weight})

    def _serialize_past_proposals(self) -> str:
        """
        Serialize the previous optimization history of the current version.

        Returns:
            str: A string representing the previous optimization history of the current version.
        """
        return json.dumps(self.log[-self.max_memory :], indent=2)

    async def _select(self, explore: bool = False):
        """
        Select the best performing or the most promising version (of the agent to optimize)
        according to past evidences and set it as the current version.

        Args:
            explore (bool): If False, setting the best performing version as current;
                If True, setting the most promising version as current. Defaults to False.
        """
        selected_version = await self._client.select_version(
            {
                "versions": [{"average_score": v["average_score"], "visits": v["visits"]} for v in self.versions],
                "explore": explore,
                "min_score": self.score_min,
                "max_score": self.score_max,
                "total_visits": self.total_visits,
            }
        )
        self.current_version = selected_version
        get_current_params().sync(self.versions[selected_version]["params"])
        self.log = copy.deepcopy(self.versions[selected_version]["log"])

    def _update_score(self, score: float):
        """
        Adding an evaluation score of the current version to evidence for version selection.

        Args:
            score (float): The evaluation score of the current version to be added as evidence
            for version selection.
        """
        self.total_visits += 1
        self.versions[self.current_version]["average_score"] = (
            self.versions[self.current_version]["average_score"] * self.versions[self.current_version]["visits"]
            + score
        ) / (self.versions[self.current_version]["visits"] + 1.0)
        self.versions[self.current_version]["visits"] += 1

        self.score_min = min(self.score_min, score)
        self.score_max = max(self.score_max, score)

    def _serialize_agent_outputs(self, agent_outputs: Any) -> str:
        """
        Serialize agent outputs.

        Args:
            agent_outputs (Any): The agent outputs to serialize.

        Returns:
            str: A string representing the agent outputs.
        """
        if isinstance(agent_outputs, dict):
            ret = ""
            for key, value in agent_outputs.items():
                ret += f"{key}:\n {value}\n"
            return ret
        else:
            return str(agent_outputs)

    def _compress_test_case(
        self, test_case: dict[str, Any], agent_log: AgentLog, agent_log_compressor: AgentLogCompressor
    ) -> dict[str, Any]:
        compressed_log = agent_log_compressor(agent_log)
        return {
            "input": str(compressed_log.simulation_tape.agent_inputs),
            "log": compressed_log.simulation_tape.extras["relai_log"],
            "trace_id": compressed_log.trace_id,
            "output": self._serialize_agent_outputs(compressed_log.agent_outputs),
            "eval_score": test_case["eval_score"],
            "eval_feedback": test_case["eval_feedback"],
        }

    def _compress_test_cases(
        self,
        test_cases: list[dict[str, Any]],
        agent_logs: list[AgentLog],
        agent_log_compressor: AgentLogCompressor,
    ) -> list[dict[str, Any]]:
        return [
            self._compress_test_case(test_case, agent_log, agent_log_compressor)
            for test_case, agent_log in zip(test_cases, agent_logs)
        ]

    def _compress_review_test_cases(
        self,
        test_cases_updated: list[dict[str, Any]],
        agent_logs: list[AgentLog],
        agent_logs_updated: list[AgentLog],
        agent_log_compressor: AgentLogCompressor,
    ) -> list[dict[str, Any]]:
        compressed_test_cases_updated = []
        for test_case_updated, agent_log, agent_log_updated in zip(test_cases_updated, agent_logs, agent_logs_updated):
            compressed_previous = self._compress_test_case(test_case_updated, agent_log, agent_log_compressor)
            compressed_current = self._compress_test_case(test_case_updated, agent_log_updated, agent_log_compressor)
            compressed_current["previous_log"] = compressed_previous["log"]
            compressed_current["previous_output"] = compressed_previous["output"]
            compressed_current["previous_eval_score"] = test_case_updated["previous_eval_score"]
            compressed_current["previous_eval_feedback"] = test_case_updated["previous_eval_feedback"]
            compressed_test_cases_updated.append(compressed_current)
        return compressed_test_cases_updated

    async def _run_with_compression_fallback(
        self,
        compress_mode: bool,
        agent_log_compressor: AgentLogCompressor | None,
        regular_call: Callable[[], Awaitable[Any]],
        compressed_call: Callable[[], Awaitable[Any]],
    ) -> tuple[Any, bool]:
        if compress_mode:
            return await compressed_call(), True

        try:
            return await regular_call(), False
        except Exception as e:
            if agent_log_compressor is None or not is_context_length_exceeded_error(e):
                raise e
            return await compressed_call(), True

    def _require_agent_log_compressor(self, agent_log_compressor: AgentLogCompressor | None) -> AgentLogCompressor:
        if agent_log_compressor is None:
            raise ValueError("`agent_log_compressor` is required for compressed execution.")
        return agent_log_compressor

    async def _evaluate(
        self, awaitables: list[Awaitable], criticos: list[Critico], verbose: bool = True, print_flag: str = ""
    ) -> tuple[list[dict[str, Any]], list[AgentLog]]:
        """
        Run and evaluate the current version of the agent through a set of awaitables.

        Args:
            awaitables (list[Awaitable]): A list of awaitables, each representing a run of the agent
            criticos (list[Critico]): A list of Critico objects, each corresponding to an awaitable
            verbose (bool): If True, additional information will be printed during evaluation.
                Defaults to True.
            print_flag (str): A string to be put next to the printed info when `verbose` is True.
                Used to distinguish printed info from different types of evaluations.

        Returns:
            list[dict]: A list of dictionary. Each dictionary contains evaluation info corresponding to
                one `Sample` object and has the following keys: "input", "log", "output", "eval_score",
                "eval_feedback".
        """
        RELAI_EVAL_BATCH_SIZE = int(os.getenv("RELAI_EVAL_BATCH_SIZE", 10))

        get_current_param_graph().clear()
        test_cases = []

        async def _evaluate_awaitable(awaitable: Awaitable, critico: Critico) -> tuple[dict[str, Any], AgentLog]:
            agent_log: AgentLog = (await awaitable)[0]

            # Handle errors gracefully
            if agent_log.agent_outputs.get("[Exception]") is not None:
                return {
                    "input": str(agent_log.simulation_tape.agent_inputs),
                    "log": agent_log.simulation_tape.extras["relai_log"],
                    "trace_id": agent_log.trace_id,
                    "output": self._serialize_agent_outputs(agent_log.agent_outputs),
                    "eval_score": 0.0,
                    "eval_feedback": (
                        "The following error is encountered during agent execution and should be fixed:\n"
                        f"{agent_log.agent_outputs.get('[Exception]')}"
                    ),
                }, agent_log

            with no_trace():
                eval_result: CriticoLog = (await critico.evaluate([agent_log]))[0]

            return {
                "input": str(agent_log.simulation_tape.agent_inputs),
                "log": agent_log.simulation_tape.extras["relai_log"],
                "trace_id": agent_log.trace_id,
                "output": self._serialize_agent_outputs(agent_log.agent_outputs),
                "eval_score": eval_result.aggregate_score,
                "eval_feedback": eval_result.aggregate_feedback,
            }, agent_log

        batches_awaitable = [
            awaitables[i : i + RELAI_EVAL_BATCH_SIZE] for i in range(0, len(awaitables), RELAI_EVAL_BATCH_SIZE)
        ]
        batches_critico = [
            criticos[i : i + RELAI_EVAL_BATCH_SIZE] for i in range(0, len(criticos), RELAI_EVAL_BATCH_SIZE)
        ]
        test_cases = []
        agent_logs = []
        for batch_awaitable, batch_critico in zip(batches_awaitable, batches_critico):
            tasks = [_evaluate_awaitable(sample, critico) for sample, critico in zip(batch_awaitable, batch_critico)]
            batch_test_cases = await asyncio.gather(*tasks)
            test_cases.extend([test_case for test_case, _ in batch_test_cases])
            agent_logs.extend([agent_log for _, agent_log in batch_test_cases])

        if verbose:
            for test_case in test_cases:
                print("=================agent excution result===================")
                print(f"- input:\n{test_case['input']}\n")
                print(f"- log{print_flag}:\n{test_case['log']}\n")
                print(f"- output{print_flag}:\n{test_case['output']}\n")
                print(f"- eval score{print_flag}:\n{test_case['eval_score']}\n")
                print(f"- eval feedback{print_flag}:\n{test_case['eval_feedback']}\n")
                print("=========================================================\n\n")

        return test_cases, agent_logs

    async def _iterate(
        self,
        batch_size: int,
        sampler: ProportionalSampler,
        config_context: str | None = None,
        agent_log_compressor: AgentLogCompressor | None = None,
        verbose: bool = True,
        group_id: str | None = None,
        pbar: tqdm | None = None,
        agent_version_uuid: str | None = None,
        environment_uuid: str | None = None,
    ) -> bool:
        """
        An iterate step will propose changes to the current version of the agent and
        conduct a preliminary examination of the proposed changes.

        It returns True if the proposed changes pass the preliminary examination and
        False otherwise.

        Args:
            batch_size (int): 2 * `batch_size` samples will be obtained from the main Critico agent,
                i.e. `critico`, where `batch_size` of them will be used to propose changes and the other
                `batch_size` of them will be used for preliminary examinations.
            sampler (ProportionalSampler): Sampler to use for selecting setups.
            config_context (str, optional): Additional context or information about the configs/parameters to assist
                in the optimization process, such as the role of different parameters or the specific formats/features
                supported. This can be especially helpful for optimizing parameters with customized formats.
            agent_log_compressor (Callable[[AgentLog], AgentLog], optional): A function that compresses the AgentLog
                when it is too long to be congested. It takes an AgentLog as input and returns a compressed
                AgentLog. This can be especially helpful for optimizing agents with long execution traces.
            verbose (bool): If True, additional information will be printed during the iterate step.
                Defaults to True.
            group_id (str, optional): An optional group ID to associate all runs together. If not provided,
                a new UUID will be generated.
            pbar (tqdm, optional): A progress bar to display the progress of the iteration. Defaults to None.
            agent_version_uuid (str, optional): An optional agent version UUID to associate with the runs.
            environment_uuid (str, optional): An optional environment UUID to associate with the runs.

        Returns:
            bool: True if the proposed changes pass the preliminary examination and False otherwise.

        Raises:
            ValueError: If no setup (simulator, critico) has been added to Maestro.
        """
        if len(self.setups) == 0:
            raise ValueError(
                "No setup (simulator, critico) has been added to Maestro. Please add at least one setup before optimization."
            )

        group_id = uuid4().hex if group_id is None else group_id

        setups = sampler.sample(batch_size * 2)
        awaitables = []
        criticos = []
        for setup in setups:
            simulator = setup["simulator"]
            critico = setup["critico"]
            awaitables.append(
                simulator.run(
                    num_runs=1,
                    group_id=group_id,
                    agent_version_uuid=agent_version_uuid,
                    environment_uuid=environment_uuid,
                    catch_exception=True,
                )
            )
            criticos.append(critico)

        test_cases, agent_logs = await self._evaluate(awaitables=awaitables, criticos=criticos, verbose=verbose)

        if pbar is not None:
            pbar.update(len(test_cases))

        compress_mode = False

        async def _propose_values_uncompressed():
            return await self._client.propose_values(
                {
                    "params": get_current_params().export(),
                    "serialized_past_proposals": self._serialize_past_proposals(),
                    "test_cases": test_cases[:batch_size],
                    "goal": self.goal
                    if config_context is None
                    else f"{self.goal}\nAdditional context:\n<context>\n{config_context}\n</context>",
                    "param_graph": get_current_param_graph().export(),
                }
            )

        async def _propose_values_compressed():
            compressor = self._require_agent_log_compressor(agent_log_compressor)
            compressed_test_cases = self._compress_test_cases(test_cases, agent_logs, compressor)
            return await self._client.propose_values(
                {
                    "params": get_current_params().export(),
                    "serialized_past_proposals": self._serialize_past_proposals(),
                    "test_cases": compressed_test_cases[:batch_size],
                    "goal": self.goal
                    if config_context is None
                    else f"{self.goal}\nAdditional context:\n<context>\n{config_context}\n</context>",
                    "param_graph": get_current_param_graph().export(),
                }
            )

        (analysis, proposed_values), compress_mode = await self._run_with_compression_fallback(
            compress_mode=compress_mode,
            agent_log_compressor=agent_log_compressor,
            regular_call=_propose_values_uncompressed,
            compressed_call=_propose_values_compressed,
        )

        changes = []
        for param, value in proposed_values.items():
            changes.append(
                {"param": param, "previous value": get_current_params().__getattr__(param), "new value": value}
            )
            if verbose:
                print("=" * 60)
                print("- proposed param change:", param)
                print("")
                print("- previous value:\n\n", get_current_params().__getattr__(param))
                print("")
                print("- new value:\n\n", value)
                print("=" * 60)

        self.log.append({"proposal id": len(self.log), "proposed changes": changes})

        # Decide if the change should be kept

        for name, proposed_value in proposed_values.items():
            get_current_params().update(name, proposed_value)

        new_awaitables = []
        new_criticos = []
        for test_case, agent_log, setup in zip(test_cases, agent_logs, setups):
            simulator = setup["simulator"]
            critico = setup["critico"]
            new_awaitables.append(
                simulator.rerun(
                    [agent_log.simulation_tape],
                    group_id=group_id,
                    agent_version_uuid=agent_version_uuid,
                    environment_uuid=environment_uuid,
                    catch_exception=True,
                )
            )
            new_criticos.append(critico)

        test_cases_updated, agent_logs_updated = await self._evaluate(
            awaitables=new_awaitables, criticos=new_criticos, verbose=verbose, print_flag=" (changed)"
        )

        if pbar is not None:
            pbar.update(len(test_cases_updated))

        for sample_id in range(0, batch_size * 2):
            test_cases_updated[sample_id]["previous_log"] = test_cases[sample_id]["log"]
            test_cases_updated[sample_id]["previous_output"] = test_cases[sample_id]["output"]
            test_cases_updated[sample_id]["previous_eval_score"] = test_cases[sample_id]["eval_score"]
            test_cases_updated[sample_id]["previous_eval_feedback"] = test_cases[sample_id]["eval_feedback"]

        for change in changes:
            get_current_params().update(change["param"], change["previous value"])

        previous_score = 0
        new_score = 0
        for test_case in test_cases_updated[batch_size:]:
            previous_score += test_case["previous_eval_score"]
            new_score += test_case["eval_score"]
        previous_score = float(previous_score) / batch_size
        new_score = float(new_score) / batch_size

        async def _review_values_uncompressed():
            return await self._client.review_values(
                {
                    "params": get_current_params().export(),
                    "serialized_past_proposals": self._serialize_past_proposals(),
                    "proposal": changes,
                    "test_cases": test_cases_updated[:batch_size],
                    "holdout_test_cases": test_cases_updated[batch_size:],
                    "previous_score": previous_score,
                    "new_score": new_score,
                    "goal": self.goal
                    if config_context is None
                    else f"{self.goal}\nAdditional context:\n<context>\n{config_context}\n</context>",
                    "analysis": analysis,
                }
            )

        async def _review_values_compressed():
            compressor = self._require_agent_log_compressor(agent_log_compressor)
            compressed_test_cases_updated = self._compress_review_test_cases(
                test_cases_updated, agent_logs, agent_logs_updated, compressor
            )
            return await self._client.review_values(
                {
                    "params": get_current_params().export(),
                    "serialized_past_proposals": self._serialize_past_proposals(),
                    "proposal": changes,
                    "test_cases": compressed_test_cases_updated[:batch_size],
                    "holdout_test_cases": compressed_test_cases_updated[batch_size:],
                    "previous_score": previous_score,
                    "new_score": new_score,
                    "goal": self.goal
                    if config_context is None
                    else f"{self.goal}\nAdditional context:\n<context>\n{config_context}\n</context>",
                    "analysis": analysis,
                }
            )

        review_decision, compress_mode = await self._run_with_compression_fallback(
            compress_mode=compress_mode,
            agent_log_compressor=agent_log_compressor,
            regular_call=_review_values_uncompressed,
            compressed_call=_review_values_compressed,
        )

        if review_decision["accepted"]:
            self.log[-1]["status"] = "ACCEPTED"
            self.log[-1]["review comment"] = review_decision["full comment"]
            for change in changes:
                get_current_params().update(change["param"], change["new value"])

        else:
            self.log[-1]["status"] = "REJECTED"
            self.log[-1]["review comment"] = review_decision["full comment"]

        if verbose:
            print("previous avg score: ", previous_score)
            print("new avg score: ", new_score)
            print("accepted: ", review_decision["accepted"])
            print("review comment:\n", review_decision["full comment"])
            print("-" * 60 + "\n\n")

        return review_decision["accepted"]

    async def optimize_config(
        self,
        total_rollouts: int,
        batch_size: int = 8,
        explore_radius: int = 5,
        explore_factor: float = 0.5,
        config_context: str | None = None,
        agent_log_compressor: AgentLogCompressor | None = None,
        group_id: str | None = None,
        verbose: bool = True,
        agent_version_uuid: str | None = None,
        environment_uuid: str | None = None,
    ) -> dict[str, Any]:
        """
        Optimize the configs (parameters) of the agent.

        Args:
            total_rollouts (int): Total number of rollouts to use for optimization.
            batch_size (int): Base batch size to use for individual optimization steps. Defaults to 8.
            explore_radius (int): A positive integer controlling the aggressiveness of exploration during optimization.
                A larger `explore_radius` encourages the optimizer to make more substantial changes between successive configurations.
                Defaults to 5.
            explore_factor (float): A float between 0 to 1 controlling the exploration-exploitation trade-off.
                A higher `explore_factor` allocates more rollouts to discover new configs,
                while a lower value allocates more rollouts to ensure the discovered configs are thoroughly evaluated.
                Defaults to 0.5.
            config_context (str, optional): Additional context or information about the configs/parameters to assist
                in the optimization process, such as the role of different parameters or the specific formats/features
                supported. This can be especially helpful for optimizing parameters with customized formats.
            agent_log_compressor (Callable[[AgentLog], AgentLog], optional): A function that compresses the AgentLog
                when it is too long to be congested. It takes an AgentLog as input and returns a compressed
                AgentLog. This can be especially helpful for optimizing agents with long execution traces.
            group_id (str, optional): An optional group ID to associate all runs together. If not provided,
                a new UUID will be generated.
            verbose (bool, optional): If True, related information will be printed during the optimization step.
                Defaults to True.
            agent_version_uuid (str, optional): An optional agent version UUID to associate with the runs.
            environment_uuid (str, optional): An optional environment UUID to associate with the runs.

        Raises:
            ValueError: If the input parameters are not valid.

        Returns:
            dict[str, Any]: Dictionary representation of all optimized parameters.
        """
        if not isinstance(batch_size, int) or batch_size <= 0:
            raise ValueError(f"`batch_size` must be a positive integer, got {batch_size}.")

        if not isinstance(explore_radius, int) or explore_radius <= 0:
            raise ValueError(f"`explore_radius` must be a positive integer, got {explore_radius}.")

        if explore_factor <= 0 or explore_factor >= 1:
            raise ValueError(f"`explore_factor` must be a float between 0 and 1, got {explore_factor}.")

        group_size = (batch_size + 1) // 2
        # total_rollouts = (iterate_steps * group_size * 4 + select_steps * group_size) * num_rounds
        # explore_factor = (iterate_steps * group_size * 4) / (iterate_steps * group_size * 4 + select_steps * group_size)
        iterate_steps: int = explore_radius
        select_steps: int = int(explore_radius * 4 * (1 - explore_factor) / explore_factor)
        num_rounds: int = int(total_rollouts / (iterate_steps * group_size * 4 + select_steps * group_size))
        total_rollouts = num_rounds * (iterate_steps * group_size * 4 + select_steps * group_size)

        print("\n" + "=" * 30 + "optimize_config settings" + "=" * 30)
        print("  total_rollouts: ", total_rollouts)
        print("  (adjusted) batch_size: ", group_size * 2)
        print("  explore_radius: ", explore_radius)
        print("  explore_factor: ", explore_factor)
        print("-" * 60)
        print("  num_rounds: ", num_rounds)
        print("  steps per round for config discovery: ", iterate_steps)
        print("  steps per round for config evaluation: ", select_steps)

        print("=" * 80 + "\n\n")

        if num_rounds == 0:
            raise ValueError(
                f"`total_rollouts` is too small for the given `batch_size` {batch_size}, `explore_radius` {explore_radius}, and `explore_factor` {explore_factor}. "
                f"Please increase `total_rollouts` to at least {iterate_steps * group_size * 4 + select_steps * group_size}."
            )

        sampler = ProportionalSampler(
            elements=self.setups,
            weights=[setup["weight"] for setup in self.setups],
        )
        if group_id is None:
            group_id = "Maestro-Config-" + uuid4().hex
        pbar = tqdm(total=total_rollouts, desc="Total rollouts consumed for config optimization")

        for round in range(num_rounds):
            print("\n\n" + "=" * 30 + f" Round {round + 1}/{num_rounds} begins" + "=" * 30)
            print("Total versions accepted: ", len(self.versions))
            print("Rebase to version: ", self.current_version)
            print(
                "Score for the current base version: %s based on %s rollouts"
                % (
                    self.versions[self.current_version]["average_score"],
                    self.versions[self.current_version]["visits"] * group_size,
                )
            )
            print("\n\n")

            new_version = False
            for _ in range(iterate_steps):
                changes_accepted = await self._iterate(
                    batch_size=group_size,
                    config_context=config_context,
                    agent_log_compressor=agent_log_compressor,
                    verbose=verbose,
                    sampler=sampler,
                    group_id=group_id,
                    pbar=pbar,
                    agent_version_uuid=agent_version_uuid,
                    environment_uuid=environment_uuid,
                )
                if changes_accepted:
                    new_version = True

            if new_version:
                self.current_version = len(self.versions)
                self.versions.append(
                    {
                        "idx": self.current_version,
                        "params": copy.deepcopy(get_current_params()),
                        "log": copy.deepcopy(self.log[-self.max_memory :]),
                        "visits": 0,
                        "average_score": 0.0,
                        "utc_timestamp": datetime.now(timezone.utc).isoformat(),
                        "validation_runs": [],
                        "validation_count": 0,
                    }
                )

            # Evaluate existing versions
            for _ in range(select_steps):
                await self._select(explore=True)

                setups = sampler.sample(group_size)
                awaitables = []
                criticos = []
                for setup in setups:
                    simulator = setup["simulator"]
                    critico = setup["critico"]
                    awaitables.append(
                        simulator.run(
                            num_runs=1,
                            group_id=group_id,
                            agent_version_uuid=agent_version_uuid,
                            environment_uuid=environment_uuid,
                            catch_exception=True,
                        )
                    )
                    criticos.append(critico)

                test_cases_validation, _ = await self._evaluate(
                    awaitables=awaitables, criticos=criticos, verbose=verbose, print_flag="(validation)"
                )

                if pbar is not None:
                    pbar.update(len(test_cases_validation))

                validation_score = 0.0
                for test_case in test_cases_validation:
                    validation_score += test_case["eval_score"]

                    self.versions[self.current_version]["validation_runs"].append(
                        RunSchema(
                            log=test_case["log"],
                            trace_id=test_case["trace_id"],
                            input=test_case["input"],
                            output=test_case["output"],
                            eval_score=test_case["eval_score"],
                            eval_feedback=test_case["eval_feedback"],
                        )
                    )

                    # Kept up to 100 runs per config for browsing
                    self.versions[self.current_version]["validation_runs"] = self.versions[self.current_version][
                        "validation_runs"
                    ][-100:]

                    self.versions[self.current_version]["validation_count"] += 1

                validation_score /= len(test_cases_validation)
                self._update_score(validation_score)

            # Switch to the current version with highest score
            await self._select(explore=False)

            print("\n\n" + "=" * 30 + f" Round {round + 1}/{num_rounds} finishes" + "=" * 30)
            print("Total versions accepted: ", len(self.versions))
            print("Best version index: ", self.current_version)
            print(
                "Score for the best version: %s based on %s rollouts"
                % (
                    self.versions[self.current_version]["average_score"],
                    self.versions[self.current_version]["visits"] * group_size,
                )
            )

            print(
                "All versions: ",
                {
                    i: {
                        "score": self.versions[i]["average_score"],
                        "rollouts evaluated": self.versions[i]["visits"] * group_size,
                    }
                    for i in range(len(self.versions))
                },
            )
            print("--------------------")

            async def sync_to_platform():
                payload = ConfigOptVizSchema(
                    name=self.name,
                    configs=[
                        ConfigSchema(
                            params={
                                key: ParamSchema(**config["params"]._params[key]) for key in config["params"].all()
                            },
                            validation_score=config["average_score"],
                            validation_count=config["validation_count"],
                            validation_runs=[],  # currently not in use; upload a placeholder to reduce payload size
                            utc_timestamp=config["utc_timestamp"],
                        )
                        for config in self.versions
                    ],
                    current_config=self.current_version,
                    validation_score_over_version=[
                        (idx, config["average_score"], config["validation_count"])
                        for idx, config in enumerate(self.versions)
                    ],
                    validation_score_over_time=[
                        (config["utc_timestamp"], config["average_score"], config["validation_count"])
                        for config in self.versions
                    ],
                )

                self.config_opt_viz_id = await self._client.update_config_opt_visual(
                    config_viz=payload, uuid=self.config_opt_viz_id
                )

                return self.config_opt_viz_id

            if self.log_to_platform:
                await sync_to_platform()
                print(
                    (
                        f"Results of round {round + 1}/{num_rounds} uploaded to RELAI platform, "
                        f"visualization id: {self.config_opt_viz_id}\n\n\n"
                    )
                )

        return get_current_params().export()

    async def optimize_structure(
        self,
        total_rollouts: int,
        batch_size: int = 10,
        description: Optional[str] = None,
        code_paths: Optional[list[str]] = None,
        code: Optional[str] = None,
        code_verifier: Optional[Callable[[str], tuple[bool, str]]] = None,
        code_context: Optional[str] = None,
        agent_log_compressor: AgentLogCompressor | None = None,
        group_id: str | None = None,
        verbose: bool = True,
        agent_version_uuid: str | None = None,
        environment_uuid: str | None = None,
    ) -> str:
        """
        Propose structural changes (i.e. changes that cannot be achieved by setting parameters alone) to
        improve the agent.

        Args:
            total_rollouts (int): Total number of rollouts to use for optimization.
            batch_size (int): Number of rollouts to use in each batch.
                Generally, a moderate batch_size (e.g. 10-20) is required and recommended.
                For agents with longer execution traces: Try reducing batch_size if an error is raised.
            description (str, optional): Text description of the current structure/workflow/... of the agent.
            code_paths (list[str], optional): A list of paths corresponding to code files containing
                the implementation of the agent.
            code (str, optional): A string representing the code (python, json, etc.) of the agent.
                If both `code_paths` and `code` are provided, `code_paths` will be ignored.
            code_verifier (Callable[[str], tuple[bool, str]], optional):
                A function that verifies the provided code and raises errors if any issues are found.
                It takes a string (code) as input and returns a tuple of a boolean and a string.
                If the boolean is True, the code is verified successfully; otherwise, the string contains
                the issues identified during verification. The verifier will only be applied if `code` or
                `code_paths` is provided, and the optimizer will try to generate code that passes the verification.
            code_context (str, optional): Additional context or information about the code to assist generating
                code that passes the verification. This is especially useful when `code_verifier` is provided and
                when the code is in json or certain domain-specific languages/formats.
            agent_log_compressor (Callable[[AgentLog], AgentLog], optional): A function that compresses the AgentLog
                when it is too long to be congested. It takes an AgentLog as input and returns a compressed
                AgentLog. This can be especially helpful for optimizing agents with long execution traces.
            group_id (str, optional): An optional group ID to associate all runs together. If not provided,
                a new UUID will be generated.
            verbose (bool, optional): If True, additional information will be printed during the optimization.
                Defaults to True.
            agent_version_uuid (str, optional): An optional agent version UUID to associate with the runs.
            environment_uuid (str, optional): An optional environment UUID to associate with the runs.

        Returns:
            str: Suggestion for structural changes to the agent.
        """

        print("\n" + "=" * 30 + "optimize_structure settings" + "=" * 30)
        print("  total_rollouts: ", total_rollouts)
        print("=" * 80 + "\n\n")

        if code is None and code_paths is not None:
            code = extract_code(code_paths=code_paths)

        sampler = ProportionalSampler(
            elements=self.setups,
            weights=[setup["weight"] for setup in self.setups],
        )
        if group_id is None:
            group_id = "Maestro-Struct-" + uuid4().hex

        print("=" * 80)
        print("Running the agent to collect traces...\n\n")

        setups = sampler.sample(total_rollouts)
        awaitables = []
        criticos = []
        for setup in setups:
            simulator = setup["simulator"]
            critico = setup["critico"]
            awaitables.append(
                simulator.run(
                    num_runs=1,
                    group_id=group_id,
                    agent_version_uuid=agent_version_uuid,
                    environment_uuid=environment_uuid,
                    catch_exception=True,
                )
            )
            criticos.append(critico)

        test_cases, agent_logs = await self._evaluate(awaitables=awaitables, criticos=criticos, verbose=verbose)
        compress_mode = False

        if batch_size < len(test_cases):
            print("Selecting and clustering samples for optimization...\n\n")
            active_tags = {}

            for test_case, agent_log in zip(test_cases, agent_logs):

                async def _process_uncompressed():
                    return await self._client.process_test_case(
                        {
                            "test_case": test_case,
                            "active_tags": active_tags,
                            "threshold": batch_size * 3,
                            "agent_name": get_full_func_name(self.agent_fn),
                            "agent_code": code,
                            "structure_description": description,
                            "params": get_current_params().export(),
                            "goal": self.goal,
                            "constrained": False,
                        }
                    )

                async def _process_compressed():
                    compressor = self._require_agent_log_compressor(agent_log_compressor)
                    compressed_test_case = self._compress_test_case(test_case, agent_log, compressor)
                    return await self._client.process_test_case(
                        {
                            "test_case": compressed_test_case,
                            "active_tags": active_tags,
                            "threshold": batch_size * 3,
                            "agent_name": get_full_func_name(self.agent_fn),
                            "agent_code": code,
                            "structure_description": description,
                            "params": get_current_params().export(),
                            "goal": self.goal,
                            "constrained": False,
                        }
                    )

                (_, active_tags), compress_mode = await self._run_with_compression_fallback(
                    compress_mode=compress_mode,
                    agent_log_compressor=agent_log_compressor,
                    regular_call=_process_uncompressed,
                    compressed_call=_process_compressed,
                )

            for tag in active_tags:
                active_tags[tag] = 0

            tags_per_test_case = []
            for test_case, agent_log in zip(test_cases, agent_logs):

                async def _tag_uncompressed():
                    return await self._client.process_test_case(
                        {
                            "test_case": test_case,
                            "active_tags": active_tags,
                            "threshold": batch_size * 3,
                            "agent_name": get_full_func_name(self.agent_fn),
                            "agent_code": code,
                            "structure_description": description,
                            "params": get_current_params().export(),
                            "goal": self.goal,
                            "constrained": True,
                        }
                    )

                async def _tag_compressed():
                    compressor = self._require_agent_log_compressor(agent_log_compressor)
                    compressed_test_case = self._compress_test_case(test_case, agent_log, compressor)
                    return await self._client.process_test_case(
                        {
                            "test_case": compressed_test_case,
                            "active_tags": active_tags,
                            "threshold": batch_size * 3,
                            "agent_name": get_full_func_name(self.agent_fn),
                            "agent_code": code,
                            "structure_description": description,
                            "params": get_current_params().export(),
                            "goal": self.goal,
                            "constrained": True,
                        }
                    )

                (tags, _), compress_mode = await self._run_with_compression_fallback(
                    compress_mode=compress_mode,
                    agent_log_compressor=agent_log_compressor,
                    regular_call=_tag_uncompressed,
                    compressed_call=_tag_compressed,
                )

                tags_per_test_case.append(tags)
                for tag in tags:
                    if tag in active_tags:
                        active_tags[tag] += 1

            to_sort = []
            for idx in range(len(test_cases)):
                score = 0
                for tag in tags_per_test_case[idx]:
                    score = max(score, active_tags.get(tag, 0))
                to_sort.append((score, idx))
            to_sort.sort(reverse=True)
            selected_test_cases = [test_cases[idx] for score, idx in to_sort[:batch_size]]
            selected_agent_logs = [agent_logs[idx] for score, idx in to_sort[:batch_size]]
        else:
            selected_test_cases = test_cases
            selected_agent_logs = agent_logs

        print("=" * 80)
        print("Optimizing structure...\n\n")

        async def _optimize_structure_uncompressed():
            return await self._client.optimize_structure(
                {
                    "agent_name": get_full_func_name(self.agent_fn),
                    "agent_code": code
                    if code_context is None
                    else (f"{code}\nAdditional context about the code:\n<context>\n{code_context}\n</context>"),
                    "structure_description": description,
                    "params": get_current_params().export(),
                    "serialized_past_proposals": self._serialize_past_proposals(),
                    "test_cases": selected_test_cases,
                    "goal": self.goal,
                    "max_correction_rounds": 1,
                }
            )

        async def _optimize_structure_compressed():
            compressor = self._require_agent_log_compressor(agent_log_compressor)
            compressed_test_cases = self._compress_test_cases(selected_test_cases, selected_agent_logs, compressor)
            return await self._client.optimize_structure(
                {
                    "agent_name": get_full_func_name(self.agent_fn),
                    "agent_code": code
                    if code_context is None
                    else (f"{code}\nAdditional context about the code:\n<context>\n{code_context}\n</context>"),
                    "structure_description": description,
                    "params": get_current_params().export(),
                    "serialized_past_proposals": self._serialize_past_proposals(),
                    "test_cases": compressed_test_cases,
                    "goal": self.goal,
                    "max_correction_rounds": 1,
                }
            )

        suggestion, compress_mode = await self._run_with_compression_fallback(
            compress_mode=compress_mode,
            agent_log_compressor=agent_log_compressor,
            regular_call=_optimize_structure_uncompressed,
            compressed_call=_optimize_structure_compressed,
        )

        if (code is not None) and (code_verifier is not None):
            print("Code verifier provided. Generating code...\n\n")
            suggested_code = ""
            suggested_patch = ""
            verified = False
            issues = None
            correction_round = 0
            max_correction_rounds = 3

            while not verified and correction_round < max_correction_rounds:
                suggested_patch = await self._client.apply_structure_to_code(
                    {
                        "agent_name": get_full_func_name(self.agent_fn),
                        "agent_code": code,
                        "suggested_structure": suggestion,
                        "previous_suggested_code": suggested_patch,
                        "issues": issues,
                        "code_context": code_context,
                    }
                )
                try:
                    suggested_code = apply_diff(code, suggested_patch)
                    verified, issues = code_verifier(suggested_code)
                except Exception as e:
                    verified = False
                    issues = f"Error applying suggested patch: {e}"

                if not verified:
                    print(
                        (
                            f"Code verification failed in round {correction_round + 1}/{max_correction_rounds}\n"
                            f"issues: {issues}\n\n"
                        )
                    )
                correction_round += 1

            if verified:
                print("Code verified successfully!\n\n")
                suggestion_to_upload = (
                    f"[Code verified successfully]\n"
                    f"<proposal>\n{suggestion}\n</proposal>\n\n"
                    f"<suggested_patch>\n{suggested_patch}\n</suggested_patch>\n\n"
                    f"<issues>\n{issues}\n</issues>\n\n"
                )
                suggestion = suggested_code
            else:
                print("Code verification failed.")
                suggestion_to_upload = (
                    f"[Code verification failed]\n"
                    f"<proposal>\n{suggestion}\n</proposal>\n\n"
                    f"<suggested_patch>\n{suggested_patch}\n</suggested_patch>\n\n"
                    f"<issues>\n{issues}\n</issues>\n\n"
                )
                suggestion = suggestion_to_upload
        else:
            suggestion_to_upload = suggestion

        async def sync_to_platform():
            payload = GraphOptVizSchema(
                name=self.name,
                proposal=suggestion_to_upload,
                runs=[],  # currently not in use; upload a placeholder to reduce payload size
            )

            return await self._client.update_graph_opt_visual(payload)

        print("=" * 40 + "suggestion" + "=" * 40)
        print(suggestion)
        print("=" * 90 + "\n\n")

        if self.log_to_platform:
            uid = await sync_to_platform()
            print(f"Results uploaded to RELAI platform, visualization id: {uid}\n\n\n")

        return suggestion
