        start_index += len(lines[i])
    start_index += start_col

    # End index always comes from the value node
    if hasattr(node, "end_lineno") and hasattr(node, "end_col_offset"):
        end_index = 0
        for i in range(node.end_lineno - 1):
            end_index += len(lines[i])
        end_index += node.end_col_offset
    else:
        # Fallback estimation
        if hasattr(ast, "get_source_segment"):
            segment = ast.get_source_segment(source, node)
            end_index = start_index + len(segment) if segment else start_index + 1
        else:
            end_index = start_index + 1

    return start_index, end_index

def find_bids_and_xy_pairs(args: list[ArgInfo]) -> list[ArgInfo]:
    """
    Find bid arguments and x,y coordinate pairs from a list of ArgInfo objects.

    Args:
        args: List of ArgInfo objects from parse_function_calls

    Returns:
        List of ArgInfo objects containing:
        - Original bid arguments (unchanged)
        - Merged x,y pairs with joint names, tuple values, and combined indices

    Rules for x,y pairs:
    - Must be consecutive arguments
    - Must end with 'x' and 'y' respectively
    - Must have the same prefix (everything before 'x'/'y')
    - Merged name: prefix + "_xy"
    - Merged value: (x_value, y_value) as tuple of floats
    - Merged indices: start of x to stop of y
    """
    result = []
    i = 0

    while i < len(args):
        current_arg = args[i]

        # Check if current arg name ends with 'bid'
        if current_arg.name.endswith("bid"):
            result.append(current_arg)
            i += 1
            continue

        # Check for x,y pair
        if i + 1 < len(args) and current_arg.name.endswith("x") and args[i + 1].name.endswith("y"):

            next_arg = args[i + 1]

            # Extract prefixes (everything before 'x' and 'y')
            current_prefix = current_arg.name[:-1]  # Remove 'x'
            next_prefix = next_arg.name[:-1]  # Remove 'y'

            # Check if they have the same prefix and are from the same function
            if (
                current_prefix == next_prefix
                and current_arg.function_name == next_arg.function_name
            ):

                # Create merged ArgInfo for x,y pair
                merged_name = f"{current_prefix}xy"

                # Convert values to floats and create tuple
                try:
                    x_val = float(current_arg.value)
                    y_val = float(next_arg.value)
                    merged_value = (x_val, y_val)
                except (ValueError, TypeError):
                    # If conversion fails, keep original values
                    merged_value = (current_arg.value, next_arg.value)

                merged_arg = ArgInfo(
                    function_name=current_arg.function_name,
                    name=merged_name,
                    value=merged_value,
                    type="tuple",
                    start_index=current_arg.start_index,
                    stop_index=next_arg.stop_index,
                )

                result.append(merged_arg)
                i += 2  # Skip both x and y args
                continue

        # If no special handling, skip this argument
        i += 1

    return result

def overlay_cross(
    img: Image.Image,
    coord: tuple[float, float],
    color: Union[str, tuple[int, int, int]] = "red",
    length: int = 7,
    width: int = 1,
) -> Image.Image:
    draw = ImageDraw.Draw(img)

    x, y = coord
    half_len = length // 2

    # Draw horizontal line
    draw.line([x - half_len, y, x + half_len, y], fill=color, width=width)
    # Draw vertical line
    draw.line([x, y - half_len, x, y + half_len], fill=color, width=width)

    return img


                merged_arg = ArgInfo(
                    function_name=current_arg.function_name,
                    name=merged_name,
                    value=merged_value,
                    type="tuple",
                    start_index=current_arg.start_index,
                    stop_index=next_arg.stop_index,
                )

                result.append(merged_arg)
                i += 2  # Skip both x and y args
                continue

        # If no special handling, skip this argument
        i += 1

    return result

def overlay_cross(
    img: Image.Image,
    coord: tuple[float, float],
    color: Union[str, tuple[int, int, int]] = "red",
    length: int = 7,
    width: int = 1,
) -> Image.Image:
    draw = ImageDraw.Draw(img)

    x, y = coord
    half_len = length // 2

    # Draw horizontal line
    draw.line([x - half_len, y, x + half_len, y], fill=color, width=width)
    # Draw vertical line
    draw.line([x, y - half_len, x, y + half_len], fill=color, width=width)

    return img

def overlay_rectangle(
    img: Image.Image,
    bbox: tuple[float, float, float, float],
    color: Union[str, tuple[int, int, int]] = "red",
    width: int = 1,
    dashed: bool = True,
) -> Image.Image:
    draw = ImageDraw.Draw(img)

    x, y, w, h = bbox

    if dashed:
        # Draw dashed rectangle
        linedashed(draw, x, y, x + w, y, color, width)
        linedashed(draw, x + w, y, x + w, y + h, color, width)
        linedashed(draw, x + w, y + h, x, y + h, color, width)
        linedashed(draw, x, y + h, x, y, color, width)
    else:


def overlay_cross(
    img: Image.Image,
    coord: tuple[float, float],
    color: Union[str, tuple[int, int, int]] = "red",
    length: int = 7,
    width: int = 1,
) -> Image.Image:
    draw = ImageDraw.Draw(img)

    x, y = coord
    half_len = length // 2

    # Draw horizontal line
    draw.line([x - half_len, y, x + half_len, y], fill=color, width=width)
    # Draw vertical line
    draw.line([x, y - half_len, x, y + half_len], fill=color, width=width)

    return img

def overlay_rectangle(
    img: Image.Image,
    bbox: tuple[float, float, float, float],
    color: Union[str, tuple[int, int, int]] = "red",
    width: int = 1,
    dashed: bool = True,
) -> Image.Image:
    draw = ImageDraw.Draw(img)

    x, y, w, h = bbox

    if dashed:
        # Draw dashed rectangle
        linedashed(draw, x, y, x + w, y, color, width)
        linedashed(draw, x + w, y, x + w, y + h, color, width)
        linedashed(draw, x + w, y + h, x, y + h, color, width)
        linedashed(draw, x, y + h, x, y, color, width)
    else:
        draw.rectangle([x, y, x + w, y + h], outline=color, width=width)

    return img

# Adapted from https://stackoverflow.com/questions/51908563/dotted-or-dashed-line-with-python-pillow/58885306#58885306
def linedashed(
    draw: PIL.ImageDraw.Draw, x0, y0, x1, y1, fill, width, dash_length=4, nodash_length=8
):
    line_dx = x1 - x0  # delta x (can be negative)
    line_dy = y1 - y0  # delta y (can be negative)
    line_length = math.hypot(line_dx, line_dy)  # line length (positive)
    if line_length == 0:
        return  # Avoid division by zero in case the line length is 0
    pixel_dx = line_dx / line_length  # x add for 1px line length
    pixel_dy = line_dy / line_length  # y add for 1px line length
    dash_start = 0
    while dash_start < line_length:
        dash_end = dash_start + dash_length
        if dash_end > line_length:
            dash_end = line_length
        draw.line(
            (


linedashed

    width: int = 1,
    dashed: bool = True,
) -> Image.Image:
    draw = ImageDraw.Draw(img)

    x, y, w, h = bbox

    if dashed:
        # Draw dashed rectangle
        linedashed(draw, x, y, x + w, y, color, width)
        linedashed(draw, x + w, y, x + w, y + h, color, width)
        linedashed(draw, x + w, y + h, x, y + h, color, width)
        linedashed(draw, x, y + h, x, y, color, width)
    else:
        draw.rectangle([x, y, x + w, y + h], outline=color, width=width)

    return img

# Adapted from https://stackoverflow.com/questions/51908563/dotted-or-dashed-line-with-python-pillow/58885306#58885306
def linedashed(
    draw: PIL.ImageDraw.Draw, x0, y0, x1, y1, fill, width, dash_length=4, nodash_length=8
):
    line_dx = x1 - x0  # delta x (can be negative)
    line_dy = y1 - y0  # delta y (can be negative)
    line_length = math.hypot(line_dx, line_dy)  # line length (positive)
    if line_length == 0:
        return  # Avoid division by zero in case the line length is 0
    pixel_dx = line_dx / line_length  # x add for 1px line length
    pixel_dy = line_dy / line_length  # y add for 1px line length
    dash_start = 0
    while dash_start < line_length:
        dash_end = dash_start + dash_length
        if dash_end > line_length:
            dash_end = line_length
        draw.line(
            (
                round(x0 + pixel_dx * dash_start),
                round(y0 + pixel_dy * dash_start),
                round(x0 + pixel_dx * dash_end),
                round(y0 + pixel_dy * dash_end),
            ),
            fill=fill,
            width=width,
        )
        dash_start += dash_length + nodash_length

def annotate_action(
    img: Image.Image, action_string: str, properties: dict[str, tuple], colormap: str = "tab10"
) -> str:
    """
    Annotate an image with overlays for action arguments and return colored HTML.

    Args:
        img: PIL Image to modify in place
        action_string: String containing function calls
        properties: Dict mapping bid strings to bounding boxes (x1, y1, x2, y2)
        colormap: Matplotlib colormap name for auto-color selection

    Returns:
        HTML string with arguments colored to match overlays
    """
    # Parse function calls to get all arguments
    all_args = parse_function_calls(action_string)


    pixel_dx = line_dx / line_length  # x add for 1px line length
    pixel_dy = line_dy / line_length  # y add for 1px line length
    dash_start = 0
    while dash_start < line_length:
        dash_end = dash_start + dash_length
        if dash_end > line_length:
            dash_end = line_length
        draw.line(
            (
                round(x0 + pixel_dx * dash_start),
                round(y0 + pixel_dy * dash_start),
                round(x0 + pixel_dx * dash_end),
                round(y0 + pixel_dy * dash_end),
            ),
            fill=fill,
            width=width,
        )
        dash_start += dash_length + nodash_length

def annotate_action(
    img: Image.Image, action_string: str, properties: dict[str, tuple], colormap: str = "tab10"
) -> str:
    """
    Annotate an image with overlays for action arguments and return colored HTML.

    Args:
        img: PIL Image to modify in place
        action_string: String containing function calls
        properties: Dict mapping bid strings to bounding boxes (x1, y1, x2, y2)
        colormap: Matplotlib colormap name for auto-color selection

    Returns:
        HTML string with arguments colored to match overlays
    """
    # Parse function calls to get all arguments
    all_args = parse_function_calls(action_string)

    # Filter to get bids and xy pairs
    filtered_args = find_bids_and_xy_pairs(all_args)

    # Get colormap
    cmap = plt.get_cmap(colormap)

    # Track colors for each filtered argument
    colors = []

    # Add overlays to image
    for i, arg_info in enumerate(filtered_args):
        # Get color from colormap
        color_rgb = cmap(i % cmap.N)
        color_255 = tuple(int(c * 255) for c in color_rgb[:3])  # Convert to 0-255 range

        colors.append(color_rgb[:3])  # Store normalized RGB for HTML

        if arg_info.name.endswith("xy"):
            # Handle x,y coordinate pairs
            x, y = arg_info.value
            overlay_cross(img, (x, y), color_255, length=9, width=3)

        elif arg_info.name.endswith("bid"):
            # Handle bid arguments with bounding boxes
            bid_value = arg_info.value
            if bid_value in properties:

                bbox = properties[bid_value]["bbox"]
                if bbox:
                    overlay_rectangle(img, bbox, color_255, width=3)

    # Generate colored HTML
    html = create_colored_html(action_string, filtered_args, colors)

    return html

def create_colored_html(action_string: str, filtered_args: list, colors: list) -> str:
    """
    Create HTML with colored arguments using start/stop indices.

    Args:
        action_string: Original action string
        filtered_args: List of ArgInfo objects with start_index/stop_index
        colors: List of RGB tuples, same length as filtered_args

    Returns:
        HTML string with colored spans
    """
    # Sort args by start position for sequential processing
    sorted_pairs = sorted(zip(filtered_args, colors), key=lambda x: x[0].start_index)

    # Build HTML with colored spans
    html_parts = []
    last_end = 0


if arg_info.name.endswith("xy"):
            # Handle x,y coordinate pairs
            x, y = arg_info.value
            overlay_cross(img, (x, y), color_255, length=9, width=3)

        elif arg_info.name.endswith("bid"):
            # Handle bid arguments with bounding boxes
            bid_value = arg_info.value
            if bid_value in properties:

                bbox = properties[bid_value]["bbox"]
                if bbox:
                    overlay_rectangle(img, bbox, color_255, width=3)

    # Generate colored HTML
    html = create_colored_html(action_string, filtered_args, colors)

    return html

def create_colored_html(action_string: str, filtered_args: list, colors: list) -> str:
    """
    Create HTML with colored arguments using start/stop indices.

    Args:
        action_string: Original action string
        filtered_args: List of ArgInfo objects with start_index/stop_index
        colors: List of RGB tuples, same length as filtered_args

    Returns:
        HTML string with colored spans
    """
    # Sort args by start position for sequential processing
    sorted_pairs = sorted(zip(filtered_args, colors), key=lambda x: x[0].start_index)

    # Build HTML with colored spans
    html_parts = []
    last_end = 0

    for arg_info, color_rgb in sorted_pairs:
        # Add uncolored text before this argument
        html_parts.append(action_string[last_end : arg_info.start_index])

        # Get the argument text
        arg_text = action_string[arg_info.start_index : arg_info.stop_index]

        # Convert color to hex
        color_hex = "#{:02x}{:02x}{:02x}".format(
            int(color_rgb[0] * 255), int(color_rgb[1] * 255), int(color_rgb[2] * 255)
        )

        # Add colored span
        html_parts.append(f'{arg_text}')

        last_end = arg_info.stop_index

    # Add remaining text
    html_parts.append(action_string[last_end:])

    return "".join(html_parts)


import base64
from io import BytesIO
from pathlib import Path

from agentlab.experiments.loop import ExpResult
from agentlab.experiments.study import get_most_recent_study
from agentlab.llm.llm_utils import BaseMessage as AgentLabBaseMessage

def exp_result_to_html(
    exp_result: ExpResult,
    steps_open: bool = True,
    som_open: bool = False,
    axtree_open: bool = False,
    html_open: bool = False,
    prompt_open: bool = False,
    embed_images: bool = True,
) -> str:
    """
    Convert an ExpResult to HTML with collapsible sections.

    Args:
        exp_result: ExpResult object containing experiment data
        steps_open: Whether step sections start expanded (default: True)
        som_open: Whether SOM screenshot sections start expanded (default: False)
        axtree_open: Whether AXTree sections start expanded (default: False)
        html_open: Whether HTML sections start expanded (default: False)
        prompt_open: Whether Prompt sections start expanded (default: False)
        embed_images: Whether to embed images as base64 or use file paths (default: True)

    Returns:
        str: HTML string with collapsible episode visualization
    """
    # Get basic episode info
    env_args = exp_result.exp_args.env_args
    steps_info = exp_result.steps_info

    # Build HTML structure
    html_parts = []

    # Add CSS for styling (unchanged)
    html_parts.append(
        """

    """
    )

    # Start main container
    html_parts.append('
')

    # Episode title and metadata
    html_parts.append(f'
{env_args.task_name}
')
    html_parts.append(f'
Seed: {env_args.task_seed}
')

    # Goal section - format like xray
    if len(steps_info) > 0:
        try:
            goal = steps_info[0].obs.get("goal_object", "No goal specified")
            goal_html = _format_goal(goal)
            html_parts.append(
                f'
Goal:
{goal_html}
'
            )
        except (IndexError, AttributeError):
            pass

    # Process each step
    for i, step_info in enumerate(steps_info):
        if step_info.action is None and i == 0:
            continue  # Skip initial reset step if no action

        # Step container with collapsible wrapper
        step_open_attr = "open" if steps_open else ""
        html_parts.append(f"")
        html_parts.append(f"Step {i}")
        html_parts.append('
')

        # Screenshot (flat in step)
        screenshot_html = _get_screenshot_html(exp_result, i, embed_images)
        html_parts.append(screenshot_html)

        # Action (flat in step)
        if step_info.action is not None:
            html_parts.append(
                f'
Action:
{_escape_html(step_info.action)}
'
            )

        # Think (flat in step)
        think_content = step_info.agent_info.get("think", "")
        if think_content:
            html_parts.append(
                f'
Think:
{_escape_html(think_content)}
'
            )

        # SOM Screenshot (nested collapsible)
        som_screenshot_html = _get_som_screenshot_html(exp_result, i, som_open, embed_images)
        if som_screenshot_html:
            html_parts.append(som_screenshot_html)

        # AXTree (nested collapsible)
        axtree_content = step_info.obs.get("axtree_txt", "")
        if axtree_content:
            axtree_open_attr = "open" if axtree_open else ""
            html_parts.append(
                f"""

                AXTree

{_escape_html(axtree_content)}

            """
            )

        # HTML (nested collapsible)
        html_content = step_info.obs.get("dom_txt", "")
        if html_content:
            html_open_attr = "open" if html_open else ""
            html_parts.append(
                f"""

                HTML

{_escape_html(html_content)}

            """
            )

        # Prompt/Chat messages (nested collapsible) - format like xray
        chat_messages = step_info.agent_info.get("chat_messages", [])
        if chat_messages:
            prompt_open_attr = "open" if prompt_open else ""
            chat_html = _format_chat_messages_like_xray(chat_messages)
            html_parts.append(
                f"""

                Prompt

                    {chat_html}

            """
            )

        # Close step container
        html_parts.append("
")  # step-content
        html_parts.append("")  # step

    # Close main container
    html_parts.append("
")

    return "".join(html_parts)

def _get_screenshot_html(exp_result, step: int, embed_images: bool) -> str:
    """Get HTML for main screenshot at given step."""
    try:
        if embed_images:
            screenshot = exp_result.get_screenshot(step, som=False)
            return _image_to_html(screenshot, f"Screenshot {step}")
        else:
            screenshot_path = exp_result.get_screenshot_path(step, som=False)
            return _path_to_html(screenshot_path, f"Screenshot {step}")
    except (FileNotFoundError, IndexError):
        return "
Screenshot not available
"

def _get_som_screenshot_html(exp_result, step: int, som_open: bool, embed_images: bool) -> str:
    """Get HTML for SOM screenshot if available."""
    try:
        if embed_images:
            screenshot_som = exp_result.get_screenshot(step, som=True)
            if screenshot_som:
                som_open_attr = "open" if som_open else ""
                som_html = _image_to_html(screenshot_som, f"SOM Screenshot {step}")
                return f"""

                    Screenshot_som[{step}]

                        {som_html}

                """
        else:
            screenshot_path = exp_result.get_screenshot_path(step, som=True)
            if screenshot_path and screenshot_path.exists():
                som_open_attr = "open" if som_open else ""
                som_html = _path_to_html(screenshot_path, f"SOM Screenshot {step}")
                return f"""

                    Screenshot_som[{step}]

                        {som_html}

                """
    except (FileNotFoundError, IndexError):
        pass
    return ""

def _image_to_html(image, alt_text: str) -> str:
    """Convert PIL Image to HTML img tag with base64 encoding."""
    if image is None:
        return f"
{alt_text} not available
"

    # Convert PIL Image to base64
    buffer = BytesIO()
    image.save(buffer, format="PNG")
    img_str = base64.b64encode(buffer.getvalue()).decode()

    return f''

def _path_to_html(image_path, alt_text: str) -> str:
    """Convert image path to HTML img tag."""
    if image_path is None or not image_path.exists():
        return f"
{alt_text} not available
"

    # Convert to absolute path and use file:// protocol
    abs_path = image_path.resolve()
    return f''

# Rest of the helper functions remain unchanged...
def _format_goal(goal) -> str:
    """Format goal object like xray does - using code blocks."""
    if goal is None:
        return "
No goal specified
"

    # Format like xray's AgentLabBaseMessage approach
    goal_str = str(AgentLabBaseMessage("", goal))

    return f"
{_escape_html(goal_str)}
"

def _format_chat_messages_like_xray(messages) -> str:
    """Format chat messages like xray does - with proper role separation."""
    if not messages:
        return "
No chat messages
"

    formatted_parts = []

    for i, msg in enumerate(messages):
        message_html = []

        if hasattr(msg, "role") and hasattr(msg, "content"):
            # Handle BaseMessage objects
            role = getattr(msg, "role", "unknown")
            content = getattr(msg, "content", str(msg))

            message_html.append(f'
{role.upper()}
')
            message_html.append(f'
{_escape_html(str(content))}
')

        elif isinstance(msg, dict):
            # Handle dict messages
            role = msg.get("role", "unknown")
            content = msg.get("content", str(msg))

            message_html.append(f'
{role.upper()}
')

            if isinstance(content, list):
                # Handle multi-part content like xray
                for part in content:
                    if isinstance(part, dict):
                        if part.get("type") == "text":
                            message_html.append(
                                f'
{_escape_html(part.get("text", ""))}
'
                            )
                        elif part.get("type") == "image":
                            message_html.append('
[IMAGE]
')
                        elif part.get("type") == "tool_use":
                            tool_str = _format_tool_call_like_xray(part)
                            message_html.append(
                                f'
{_escape_html(tool_str)}
'
                            )
                        else:
                            message_html.append(
                                f'
{_escape_html(str(part))}
'
                            )
                    else:
                        message_html.append(
                            f'
{_escape_html(str(part))}
'
                        )
            else:
                message_html.append(f'
{_escape_html(str(content))}
')
        else:
            # Handle other message types
            message_html.append(f'
{_escape_html(str(msg))}
')

        formatted_parts.append(f'
{"".join(message_html)}
')

    return "".join(formatted_parts)

def _format_tool_call_like_xray(tool_item: dict) -> str:
    """Format tool calls like xray does."""
    name = tool_item.get("name", "unknown")
    input_data = tool_item.get("input", {})
    call_id = tool_item.get("call_id", "unknown")

    return f"Tool Call: {name} `{input_data}` (call_id: {call_id})"

def _escape_html(text: str) -> str:
    """Escape HTML special characters."""
    if not isinstance(text, str):
        text = str(text)

    return (
        text.replace("&", "&")
        .replace("<", "<")
        .replace(">", ">")
        .replace('"', """)
        .replace("'", "'")
    )

if __name__ == "__main__":

    from agentlab.experiments.exp_utils import RESULTS_DIR

    result_dir = get_most_recent_study(RESULTS_DIR, contains=None)
    for exp_dir in result_dir.iterdir():
        if exp_dir.is_dir():
            break

    print(f"Using first exp_dir in most recent study:\n{exp_dir}")
    exp_result = ExpResult(exp_dir=exp_dir)

    page = exp_result_to_html(exp_result, embed_images=False)

    output_file = exp_dir / "episode.html"
    print(f"Writing HTML to\n{output_file}")
    output_file.write_text(page)
    # cmd open output_file using subprocess
    import subprocess

    subprocess.run(["open", str(output_file)])  # macOS command to open HTML file


import base64
from io import BytesIO
from pathlib import Path

from agentlab.experiments.loop import ExpResult
from agentlab.experiments.study import get_most_recent_study
from agentlab.llm.llm_utils import BaseMessage as AgentLabBaseMessage

def exp_result_to_html(
    exp_result: ExpResult,
    steps_open: bool = True,
    som_open: bool = False,
    axtree_open: bool = False,
    html_open: bool = False,
    prompt_open: bool = False,
    embed_images: bool = True,
) -> str:
    """
    Convert an ExpResult to HTML with collapsible sections.

    Args:
        exp_result: ExpResult object containing experiment data
        steps_open: Whether step sections start expanded (default: True)
        som_open: Whether SOM screenshot sections start expanded (default: False)
        axtree_open: Whether AXTree sections start expanded (default: False)
        html_open: Whether HTML sections start expanded (default: False)
        prompt_open: Whether Prompt sections start expanded (default: False)
        embed_images: Whether to embed images as base64 or use file paths (default: True)

    Returns:
        str: HTML string with collapsible episode visualization
    """
    # Get basic episode info
    env_args = exp_result.exp_args.env_args
    steps_info = exp_result.steps_info

    # Build HTML structure
    html_parts = []

    # Add CSS for styling (unchanged)
    html_parts.append(
        """

    """
    )

    # Start main container
    html_parts.append('
')

    # Episode title and metadata
    html_parts.append(f'
{env_args.task_name}
')
    html_parts.append(f'
Seed: {env_args.task_seed}
')

    # Goal section - format like xray
    if len(steps_info) > 0:
        try:
            goal = steps_info[0].obs.get("goal_object", "No goal specified")
            goal_html = _format_goal(goal)
            html_parts.append(
                f'
Goal:
{goal_html}
'
            )
        except (IndexError, AttributeError):
            pass

    # Process each step
    for i, step_info in enumerate(steps_info):
        if step_info.action is None and i == 0:
            continue  # Skip initial reset step if no action

        # Step container with collapsible wrapper
        step_open_attr = "open" if steps_open else ""
        html_parts.append(f"")
        html_parts.append(f"Step {i}")
        html_parts.append('
')

        # Screenshot (flat in step)
        screenshot_html = _get_screenshot_html(exp_result, i, embed_images)
        html_parts.append(screenshot_html)

        # Action (flat in step)
        if step_info.action is not None:
            html_parts.append(
                f'
Action:
{_escape_html(step_info.action)}
'
            )

        # Think (flat in step)
        think_content = step_info.agent_info.get("think", "")
        if think_content:
            html_parts.append(
                f'
Think:
{_escape_html(think_content)}
'
            )

        # SOM Screenshot (nested collapsible)
        som_screenshot_html = _get_som_screenshot_html(exp_result, i, som_open, embed_images)
        if som_screenshot_html:
            html_parts.append(som_screenshot_html)

        # AXTree (nested collapsible)
        axtree_content = step_info.obs.get("axtree_txt", "")
        if axtree_content:
            axtree_open_attr = "open" if axtree_open else ""
            html_parts.append(
                f"""

                AXTree

{_escape_html(axtree_content)}

            """
            )

        # HTML (nested collapsible)
        html_content = step_info.obs.get("dom_txt", "")
        if html_content:
            html_open_attr = "open" if html_open else ""
            html_parts.append(
                f"""

                HTML

{_escape_html(html_content)}

            """
            )

        # Prompt/Chat messages (nested collapsible) - format like xray
        chat_messages = step_info.agent_info.get("chat_messages", [])
        if chat_messages:
            prompt_open_attr = "open" if prompt_open else ""
            chat_html = _format_chat_messages_like_xray(chat_messages)
            html_parts.append(
                f"""

                Prompt

                    {chat_html}

            """
            )

        # Close step container
        html_parts.append("
")  # step-content
        html_parts.append("")  # step

    # Close main container
    html_parts.append("
")

    return "".join(html_parts)

def _get_screenshot_html(exp_result, step: int, embed_images: bool) -> str:
    """Get HTML for main screenshot at given step."""
    try:
        if embed_images:
            screenshot = exp_result.get_screenshot(step, som=False)
            return _image_to_html(screenshot, f"Screenshot {step}")
        else:
            screenshot_path = exp_result.get_screenshot_path(step, som=False)
            return _path_to_html(screenshot_path, f"Screenshot {step}")
    except (FileNotFoundError, IndexError):
        return "
Screenshot not available
"

def _get_som_screenshot_html(exp_result, step: int, som_open: bool, embed_images: bool) -> str:
    """Get HTML for SOM screenshot if available."""
    try:
        if embed_images:
            screenshot_som = exp_result.get_screenshot(step, som=True)


f"""

                Prompt

                    {chat_html}

            """
            )

        # Close step container
        html_parts.append("
")  # step-content
        html_parts.append("")  # step

    # Close main container
    html_parts.append("
")

    return "".join(html_parts)

def _get_screenshot_html(exp_result, step: int, embed_images: bool) -> str:
    """Get HTML for main screenshot at given step."""
    try:
        if embed_images:
            screenshot = exp_result.get_screenshot(step, som=False)
            return _image_to_html(screenshot, f"Screenshot {step}")
        else:
            screenshot_path = exp_result.get_screenshot_path(step, som=False)
            return _path_to_html(screenshot_path, f"Screenshot {step}")
    except (FileNotFoundError, IndexError):
        return "
Screenshot not available
"

def _get_som_screenshot_html(exp_result, step: int, som_open: bool, embed_images: bool) -> str:
    """Get HTML for SOM screenshot if available."""
    try:
        if embed_images:
            screenshot_som = exp_result.get_screenshot(step, som=True)
            if screenshot_som:
                som_open_attr = "open" if som_open else ""
                som_html = _image_to_html(screenshot_som, f"SOM Screenshot {step}")
                return f"""

                    Screenshot_som[{step}]

                        {som_html}

                """
        else:
            screenshot_path = exp_result.get_screenshot_path(step, som=True)


# Close main container
    html_parts.append("
")

    return "".join(html_parts)

def _get_screenshot_html(exp_result, step: int, embed_images: bool) -> str:
    """Get HTML for main screenshot at given step."""
    try:
        if embed_images:
            screenshot = exp_result.get_screenshot(step, som=False)
            return _image_to_html(screenshot, f"Screenshot {step}")
        else:
            screenshot_path = exp_result.get_screenshot_path(step, som=False)
            return _path_to_html(screenshot_path, f"Screenshot {step}")
    except (FileNotFoundError, IndexError):
        return "
Screenshot not available
"

def _get_som_screenshot_html(exp_result, step: int, som_open: bool, embed_images: bool) -> str:
    """Get HTML for SOM screenshot if available."""
    try:
        if embed_images:
            screenshot_som = exp_result.get_screenshot(step, som=True)
            if screenshot_som:
                som_open_attr = "open" if som_open else ""
                som_html = _image_to_html(screenshot_som, f"SOM Screenshot {step}")
                return f"""

                    Screenshot_som[{step}]

                        {som_html}

                """
        else:
            screenshot_path = exp_result.get_screenshot_path(step, som=True)
            if screenshot_path and screenshot_path.exists():
                som_open_attr = "open" if som_open else ""
                som_html = _path_to_html(screenshot_path, f"SOM Screenshot {step}")
                return f"""

                    Screenshot_som[{step}]

                        {som_html}

                """
    except (FileNotFoundError, IndexError):
        pass
    return ""

def _image_to_html(image, alt_text: str) -> str:
    """Convert PIL Image to HTML img tag with base64 encoding."""
    if image is None:
        return f"
{alt_text} not available
"

    # Convert PIL Image to base64
    buffer = BytesIO()
    image.save(buffer, format="PNG")
    img_str = base64.b64encode(buffer.getvalue()).decode()

    return f''

def _path_to_html(image_path, alt_text: str) -> str:
    """Convert image path to HTML img tag."""
    if image_path is None or not image_path.exists():
        return f"
{alt_text} not available
"


                """
        else:
            screenshot_path = exp_result.get_screenshot_path(step, som=True)
            if screenshot_path and screenshot_path.exists():
                som_open_attr = "open" if som_open else ""
                som_html = _path_to_html(screenshot_path, f"SOM Screenshot {step}")
                return f"""

                    Screenshot_som[{step}]

                        {som_html}

                """
    except (FileNotFoundError, IndexError):
        pass
    return ""

def _image_to_html(image, alt_text: str) -> str:
    """Convert PIL Image to HTML img tag with base64 encoding."""
    if image is None:
        return f"
{alt_text} not available
"

    # Convert PIL Image to base64
    buffer = BytesIO()
    image.save(buffer, format="PNG")
    img_str = base64.b64encode(buffer.getvalue()).decode()

    return f''

def _path_to_html(image_path, alt_text: str) -> str:
    """Convert image path to HTML img tag."""
    if image_path is None or not image_path.exists():
        return f"
{alt_text} not available
"

    # Convert to absolute path and use file:// protocol
    abs_path = image_path.resolve()
    return f''

# Rest of the helper functions remain unchanged...
def _format_goal(goal) -> str:
    """Format goal object like xray does - using code blocks."""
    if goal is None:
        return "
No goal specified
"

    # Format like xray's AgentLabBaseMessage approach
    goal_str = str(AgentLabBaseMessage("", goal))


                """
    except (FileNotFoundError, IndexError):
        pass
    return ""

def _image_to_html(image, alt_text: str) -> str:
    """Convert PIL Image to HTML img tag with base64 encoding."""
    if image is None:
        return f"
{alt_text} not available
"

    # Convert PIL Image to base64
    buffer = BytesIO()
    image.save(buffer, format="PNG")
    img_str = base64.b64encode(buffer.getvalue()).decode()

    return f''

def _path_to_html(image_path, alt_text: str) -> str:
    """Convert image path to HTML img tag."""
    if image_path is None or not image_path.exists():
        return f"
{alt_text} not available
"

    # Convert to absolute path and use file:// protocol
    abs_path = image_path.resolve()
    return f''

# Rest of the helper functions remain unchanged...
def _format_goal(goal) -> str:
    """Format goal object like xray does - using code blocks."""
    if goal is None:
        return "
No goal specified
"

    # Format like xray's AgentLabBaseMessage approach
    goal_str = str(AgentLabBaseMessage("", goal))

    return f"
{_escape_html(goal_str)}
"

def _format_chat_messages_like_xray(messages) -> str:
    """Format chat messages like xray does - with proper role separation."""
    if not messages:
        return "
No chat messages
"

    formatted_parts = []


# Convert PIL Image to base64
    buffer = BytesIO()
    image.save(buffer, format="PNG")
    img_str = base64.b64encode(buffer.getvalue()).decode()

    return f''

def _path_to_html(image_path, alt_text: str) -> str:
    """Convert image path to HTML img tag."""
    if image_path is None or not image_path.exists():
        return f"
{alt_text} not available
"

    # Convert to absolute path and use file:// protocol
    abs_path = image_path.resolve()
    return f''

# Rest of the helper functions remain unchanged...
def _format_goal(goal) -> str:
    """Format goal object like xray does - using code blocks."""
    if goal is None:
        return "
No goal specified
"

    # Format like xray's AgentLabBaseMessage approach
    goal_str = str(AgentLabBaseMessage("", goal))

    return f"
{_escape_html(goal_str)}
"

def _format_chat_messages_like_xray(messages) -> str:
    """Format chat messages like xray does - with proper role separation."""
    if not messages:
        return "
No chat messages
"

    formatted_parts = []

    for i, msg in enumerate(messages):
        message_html = []

        if hasattr(msg, "role") and hasattr(msg, "content"):
            # Handle BaseMessage objects
            role = getattr(msg, "role", "unknown")
            content = getattr(msg, "content", str(msg))

            message_html.append(f'
{role.upper()}
')
            message_html.append(f'
{_escape_html(str(content))}
')


if image_path is None or not image_path.exists():
        return f"
{alt_text} not available
"

    # Convert to absolute path and use file:// protocol
    abs_path = image_path.resolve()
    return f''

# Rest of the helper functions remain unchanged...
def _format_goal(goal) -> str:
    """Format goal object like xray does - using code blocks."""
    if goal is None:
        return "
No goal specified
"

    # Format like xray's AgentLabBaseMessage approach
    goal_str = str(AgentLabBaseMessage("", goal))

    return f"
{_escape_html(goal_str)}
"

def _format_chat_messages_like_xray(messages) -> str:
    """Format chat messages like xray does - with proper role separation."""
    if not messages:
        return "
No chat messages
"

    formatted_parts = []

    for i, msg in enumerate(messages):
        message_html = []

        if hasattr(msg, "role") and hasattr(msg, "content"):
            # Handle BaseMessage objects
            role = getattr(msg, "role", "unknown")
            content = getattr(msg, "content", str(msg))

            message_html.append(f'
{role.upper()}
')
            message_html.append(f'
{_escape_html(str(content))}
')

        elif isinstance(msg, dict):
            # Handle dict messages
            role = msg.get("role", "unknown")
            content = msg.get("content", str(msg))

            message_html.append(f'
{role.upper()}
')

            if isinstance(content, list):
                # Handle multi-part content like xray
                for part in content:
                    if isinstance(part, dict):
                        if part.get("type") == "text":
                            message_html.append(
                                f'
{_escape_html(part.get("text", ""))}
'
                            )
                        elif part.get("type") == "image":
                            message_html.append('
[IMAGE]
')
                        elif part.get("type") == "tool_use":
                            tool_str = _format_tool_call_like_xray(part)
                            message_html.append(
                                f'
{_escape_html(tool_str)}
'
                            )
                        else:
                            message_html.append(
                                f'
{_escape_html(str(part))}
'
                            )
                    else:
                        message_html.append(
                            f'
{_escape_html(str(part))}
'
                        )
            else:
                message_html.append(f'
{_escape_html(str(content))}
')
        else:
            # Handle other message types
            message_html.append(f'
{_escape_html(str(msg))}
')

        formatted_parts.append(f'
{"".join(message_html)}
')

    return "".join(formatted_parts)

def _format_tool_call_like_xray(tool_item: dict) -> str:
    """Format tool calls like xray does."""
    name = tool_item.get("name", "unknown")
    input_data = tool_item.get("input", {})
    call_id = tool_item.get("call_id", "unknown")

    return f"Tool Call: {name} `{input_data}` (call_id: {call_id})"

def _escape_html(text: str) -> str:
    """Escape HTML special characters."""
    if not isinstance(text, str):
        text = str(text)

    return (
        text.replace("&", "&")
        .replace("<", "<")
        .replace(">", ">")


)
                        else:
                            message_html.append(
                                f'
{_escape_html(str(part))}
'
                            )
                    else:
                        message_html.append(
                            f'
{_escape_html(str(part))}
'
                        )
            else:
                message_html.append(f'
{_escape_html(str(content))}
')
        else:
            # Handle other message types
            message_html.append(f'
{_escape_html(str(msg))}
')

        formatted_parts.append(f'
{"".join(message_html)}
')

    return "".join(formatted_parts)

def _format_tool_call_like_xray(tool_item: dict) -> str:
    """Format tool calls like xray does."""
    name = tool_item.get("name", "unknown")
    input_data = tool_item.get("input", {})
    call_id = tool_item.get("call_id", "unknown")

    return f"Tool Call: {name} `{input_data}` (call_id: {call_id})"

def _escape_html(text: str) -> str:
    """Escape HTML special characters."""
    if not isinstance(text, str):
        text = str(text)

    return (
        text.replace("&", "&")
        .replace("<", "<")
        .replace(">", ">")
        .replace('"', """)
        .replace("'", "'")
    )

if __name__ == "__main__":

    from agentlab.experiments.exp_utils import RESULTS_DIR


else:
                message_html.append(f'
{_escape_html(str(content))}
')
        else:
            # Handle other message types
            message_html.append(f'
{_escape_html(str(msg))}
')

        formatted_parts.append(f'
{"".join(message_html)}
')

    return "".join(formatted_parts)

def _format_tool_call_like_xray(tool_item: dict) -> str:
    """Format tool calls like xray does."""
    name = tool_item.get("name", "unknown")
    input_data = tool_item.get("input", {})
    call_id = tool_item.get("call_id", "unknown")

    return f"Tool Call: {name} `{input_data}` (call_id: {call_id})"

def _escape_html(text: str) -> str:
    """Escape HTML special characters."""
    if not isinstance(text, str):
        text = str(text)

    return (
        text.replace("&", "&")
        .replace("<", "<")
        .replace(">", ">")
        .replace('"', """)
        .replace("'", "'")
    )

if __name__ == "__main__":

    from agentlab.experiments.exp_utils import RESULTS_DIR

    result_dir = get_most_recent_study(RESULTS_DIR, contains=None)
    for exp_dir in result_dir.iterdir():
        if exp_dir.is_dir():
            break

    print(f"Using first exp_dir in most recent study:\n{exp_dir}")
    exp_result = ExpResult(exp_dir=exp_dir)

    page = exp_result_to_html(exp_result, embed_images=False)

    output_file = exp_dir / "episode.html"
    print(f"Writing HTML to\n{output_file}")
    output_file.write_text(page)
    # cmd open output_file using subprocess


import fnmatch
import json
import random
import re
import traceback
import warnings
from collections import defaultdict
from logging import warn
from pathlib import Path

import numpy as np
import pandas as pd
from IPython.display import display
from tqdm import tqdm

from agentlab.experiments.loop import ExpResult, get_exp_result, yield_all_exp_results

# TODO find a more portable way to code set_task_category_as_index at least
# handle dynamic imports. We don't want to always import workarena
# from browsergym.workarena import TASK_CATEGORY_MAP

warnings.filterwarnings("ignore", category=pd.errors.PerformanceWarning)

pd.set_option("display.multi_sparse", False)

AGENT_NAME_KEY = "agent.agent_name"
TASK_KEY = "env.task_name"

def get_constants_and_variables(df: pd.DataFrame, drop_constants: bool = False):
    """Filter out constants from the dataframe."""

    constants = {}
    variable_keys = []
    for col in df.columns:
        try:
            nuniq = df[col].nunique(dropna=False)
        except TypeError:
            nuniq = 0  # non hashable types are considered variables
        if nuniq == 1:
            if isinstance(df[col].iloc[0], np.generic):
                val = df[col].iloc[0].item()
            else:
                val = df[col].iloc[0]
            constants[col] = val
            if drop_constants:
                df = df.drop(col, axis=1)
        else:
            variable_keys.append(col)

    return constants, variable_keys, df

def set_index_from_variables(
    df: pd.DataFrame,
    index_white_list=("agent.*",),
    index_black_list=("*model_url*", "*extra*", "*._*"),
    task_key=TASK_KEY,
    add_agent_and_benchmark=True,
):
    """Set the index, inplace, to env.task_name and all variables.

    Introspects `df` to find all fields that are variable and set the index to
    those fields. This will allow to easily groupby and compare results. To
    filter undersired variables from the index, use index_white_list and
    index_black_list.

    Args:
        df: The dataframe to modify
        index_white_list: List of wildard patterns to match variables that
            should be included in the index.
        index_black_list: List of wildard patterns to match variables that
            should be excluded from the index.
        task_key: The key to use as the first level of the index.
        add_agent_and_benchmark: If True, add agent.agent_name and env.benchmark
    """
    df.reset_index(inplace=True)
    constants, variables, _ = get_constants_and_variables(df)

    index_variables = []
    if add_agent_and_benchmark:
        index_variables.append("agent.agent_name")
        if "env.benchmark" not in df.columns:
            df["env.benchmark"] = df[TASK_KEY].map(_benchmark_from_task_name)
        index_variables.append("env.benchmark")

    for var in variables:
        white = any([fnmatch.fnmatch(var, pattern) for pattern in index_white_list])
        black = any([fnmatch.fnmatch(var, pattern) for pattern in index_black_list])

        if white and (not black) and (var not in index_variables):
            index_variables.append(var)

    for var in index_variables:
        if df[var].isnull().any():
            warn(
                f"Variable {var} contains NaN or None values. This will be replaced by the string 'None' to avoid some pandas bug."
            )
            df[var] = df[var].fillna("None")

    # agent_variables = [var for var in variables if var.startswith("agent.")]
    df.set_index([task_key] + index_variables, inplace=True)
    df.sort_index(inplace=True)

def load_result_df(
    exp_dir,
    progress_fn=tqdm,
    set_index=True,
    result_df=None,
    index_white_list=("agent.*",),
    index_black_list=("*model_url*", "*extra*", "*._*"),
    remove_args_suffix=True,
):
    """Load the result dataframe.

    Will set the index to env.task_name and all columens that are not constant and
    starts with agent. This will allow to easily groupby and compare
    results. This index can be changed later using df.set_index.

    Args:
        exp_dir: Path to the experiment directory
        progress_fn: Progress function to use when loading the results
        set_index: If True, set the index to env.task_name and variable agent
        result_df: If not None, speed up the loading process by reusing
            alreading loaded objects.
        index_white_list: List of wildard patterns to match variables that
            should be included in the index.
        index_black_list: List of wildard patterns to match variables that
            should be excluded from the index.
        remove_args_suffix: If True, remove the _args suffix from the columns

    Returns:
        pd.DataFrame: The result dataframe
    """

    if result_df is not None:
        result_list = list(result_df["exp_result"])
    else:
        result_list = list(yield_all_exp_results(exp_dir, progress_fn=progress_fn))

    if len(result_list) == 0:
        return None

    if progress_fn is not None:
        result_list = progress_fn(result_list, desc="Loading results")

    df = pd.DataFrame([exp_result.get_exp_record() for exp_result in result_list])

    if remove_args_suffix:
        df.columns = [col.replace("_args", "") for col in df.columns]

    if set_index:
        set_index_from_variables(df, index_white_list, index_black_list)
    return df

def reduce_episodes(result_df: pd.DataFrame) -> pd.DataFrame:
    """Reduce the dataframe to a single row per episode and summarize some of the columns."""

    levels = list(range(result_df.index.nlevels))
    return result_df.groupby(level=levels).apply(summarize)

def report_2d(df: pd.DataFrame, reduce_fn: callable = reduce_episodes, n_row_keys=1):
    """Generic function to create a 2d report based on the dataframe.

    The code is simple but can be a bit cryptic. This is best explained in the
    following  3 steps:
    1) Groupby: Will use the existing multi-index to groupby. Make sure to set the
       an index to the desired keys before calling this function.
    2) Reduce: Uses the reduce_fn to reduce the content of each group to a single
       variable, creating a 1D series indexed by its original index.
    3) Unstack: Produce a 2D table such that the first n_row_keys are used to
       specify how many dimensions are used for the rows. The remaining
       dimensions are used for the columns.

    Args:
        df: The dataframe to reduce
        reduce_fn: The function to use to reduce the sub dataframe. By default
            this is reduce_episodes.
        n_row_keys: The number of keys to use for the rows.

    Returns:
        pd.DataFrame: The 2D report
    """

    levels = list(range(df.index.nlevels))
    reduced_df = df.groupby(level=levels).apply(reduce_fn)  # type: pd.Series
    return reduced_df.unstack(level=levels[n_row_keys:])

def report_constant_and_variables(df, show_stack_traces=True):
    constants, variables, _ = get_constants_and_variables(df)
    print("Constants:")
    for k, v in constants.items():
        print(f"    {k}: {v}")

    print("\nVariables:")
    for var in variables:
        if not show_stack_traces and var == "stack_trace":
            continue

        # get unique with count and sort by count descending
        unique_counts = df[var].value_counts().sort_values(ascending=False)

        print(f"    {var}: n_unique={len(unique_counts)}")
        for i, (val, count) in enumerate(unique_counts.items()):
            print(f"        {count}x : {val}")
            if i >= 2:
                break
        if len(unique_counts) > 3:
            print("        ...\n")

def get_std_err(df, metric):
    """Get the standard error for a binary metric."""
    # extract non missing values
    data = df[metric].dropna().values

    # asser either 0 or 1
    if np.all(np.isin(data, [0, 1])):
        mean = np.mean(data)
        std_err = np.sqrt(mean * (1 - mean) / len(data))
        return mean, std_err
    else:
        return get_sample_std_err(df, metric)

def get_sample_std_err(df, metric):
    """Get the standard error for a binary metric."""
    # extract non missing values
    data = df[metric].dropna().values

    mean = np.mean(data)
    std_err = np.std(data, ddof=1) / np.sqrt(len(data))
    if np.isnan(std_err):
        std_err = np.float64(0)
    return mean, std_err

def summarize(sub_df):
    if "cum_reward" not in sub_df:
        record = dict(
            avg_reward=np.nan,
            std_err=np.nan,
            # avg_raw_reward=np.nan,
            avg_steps=np.nan,
            n_completed=f"0/{len(sub_df)}",
            n_err=0,
        )
    else:
        err = sub_df["err_msg"].notnull()
        n_completed = err.copy()
        for col in ["truncated", "terminated"]:
            if col in sub_df:
                n_completed = n_completed | sub_df[col]
        n_completed = n_completed.sum()

        if n_completed == 0:
            return None

        _mean_reward, std_reward = get_std_err(sub_df, "cum_reward")

        # sanity check, if there is an error the reward should be zero
        assert sub_df[sub_df["err_msg"].notnull()]["cum_reward"].sum() == 0

        record = dict(
            avg_reward=sub_df["cum_reward"].mean(skipna=True).round(3),
            std_err=std_reward.round(3),
            # avg_raw_reward=sub_df["cum_raw_reward"].mean(skipna=True).round(3),
            avg_steps=sub_df["n_steps"].mean(skipna=True).round(3),
            n_completed=f"{n_completed}/{len(sub_df)}",
            n_err=err.sum(skipna=True),
        )
        if "stats.cum_cost" in sub_df:
            record["cum_cost"] = sub_df["stats.cum_cost"].sum(skipna=True).round(4)
        if "stats.cum_effective_cost" in sub_df:
            record["cum_effective_cost"] = (
                sub_df["stats.cum_effective_cost"].sum(skipna=True).round(4)
            )
            record.pop("cum_cost", None)

    return pd.Series(record)

def summarize_stats(sub_df):
    """Summarize the stats columns."""

    # make sure there are completed runs
    err = sub_df["err_msg"].notnull()
    n_completed = err.copy()
    for col in ["truncated", "terminated"]:
        if col in sub_df:
            n_completed = n_completed | sub_df[col]
    n_completed = n_completed.sum()

    if n_completed == 0:
        return None

    record = dict(
        avg_reward=sub_df["cum_reward"].mean(skipna=True).round(3),
    )
    for key in sub_df.keys():
        if key.startswith("stats."):
            key_ = key.split(".")[1]
            op = key_.split("_")[0]
            if op == "cum":
                record[key_] = sub_df[key].sum(skipna=True)
            elif op == "max":
                record[key_] = sub_df[key].max(skipna=True)
            else:
                raise ValueError(f"Unknown stats operation: {op}")
    return pd.Series(record)

def _find_diff(tuple1, tuple2):
    """return the list of index wher tuple1 != tuple2"""
    return [i for i, (a, b) in enumerate(zip(tuple1, tuple2)) if a != b]

def _extract_ablation_study(report: pd.DataFrame, progression=False):
    """Reduce the multi-index to a change description compared to the previous row."""
    names = report.index.names
    report = report.copy()
    # report.sort_index(inplace=True)

    reference_index = None
    for index in report.index:
        if reference_index is not None:
            diffs = _find_diff(reference_index, index)

            if progression:
                change = "↳ " + ", ".join([f"{names[i]}={index[i]}" for i in diffs])
            else:
                changes = []
                for i in diffs:
                    val = index[i]
                    if isinstance(val, bool):
                        changes.append(("+" if val else "-") + names[i])
                    else:
                        changes.append(f"{names[i]}←{val}")
                change = ", ".join(changes)
        else:
            change = "Initial Configuration"
        report.loc[index, "change"] = change
        if progression:
            reference_index = index
        else:
            reference_index = report.index[0]

    report = report.reset_index()
    report = report.set_index(["change"])

    # delete columns related to old index
    return report.drop(names, axis=1)

def ablation_report(result_df: pd.DataFrame, reduce_fn=summarize, progression=False):
    """Reduce the multi-index to a change description compared to the previous row.

    *NOTE*: This assumes that this experiments was launched with make_ablation_study.

    Rows will be sorted according to the average ExpArgs.order for all
    experiments associated with the multi-index.

    Args:
        result_df: The result dataframe as returned by load_result_df.
        reduce_fn: The function to use to reduce the sub dataframe. By default
            this is summarize.
        progression: If True, the change description will be the progression

    Returns:
        A dataframe with the change description as index.
    """
    report = global_report(result_df, reduce_fn=reduce_fn)
    report = _sort_order(result_df, report)
    report = _extract_ablation_study(report, progression=progression)
    return report

def _get_avg_order(df: pd.DataFrame, row: pd.Series):
    """Return the average order for the given row."""
    df = df.reset_index(level=0, drop=True, inplace=False)
    # df.sort_index(inplace=True)

    sub_df = df.loc[row.name]
    orders = [get_exp_result(exp_dir).exp_args.order for exp_dir in sub_df.exp_dir]
    orders = [order for order in orders if order is not None]
    if len(orders) == 0:
        return None
    return np.mean(orders)

def _sort_order(result_df, report):
    """Add a column to the report with the average order for each agent and sort."""

    def add_order(row):
        return _get_avg_order(result_df, row)

    report["avg_order"] = report.apply(add_order, axis=1)
    return report.sort_values("avg_order", ascending=True)

def global_report(
    result_df: pd.DataFrame,
    reduce_fn=summarize,
    rename_index=lambda name: name.replace("agent.flags.", ""),
):
    """Produce a report that summarize all tasks and all episodes for each
    agent.

    Args:
        result_df: The result dataframe as returned by load_result_df.
        reduce_fn: The function to use to reduce the sub dataframe. By default
            this is summarize.
        rename_index: Function to rename the index. By default we remove the prefix
            "agent.flags."

    Returns:
        pd.DataFrame: The report
    """

    levels = list(range(result_df.index.nlevels))

    if len(levels) == 1:
        print("Only one configuration is found, returning a per-task report.")
        report = report_2d(result_df, reduce_fn=reduce_fn)
        report.loc["[ALL TASKS]"] = reduce_fn(result_df)
    else:
        print(
            "Found multiple configuration, averaging across tasks and returning a per-agent report."
        )
        report = result_df.groupby(level=levels[1:]).apply(reduce_fn)

        if rename_index is not None:
            index_names = [rename_index(name) for name in report.index.names]
            report = report.rename_axis(index=index_names)

        # if has key avg_reward
        if "avg_reward" in report.columns:
            report = report.sort_values("avg_reward", ascending=False)

    return report

def _rename_bool_flags(report: pd.DataFrame, true_str="✓", false_str="-"):
    """Rename the boolean flags to be more compact and readable."""
    map_bool = lambda x: true_str if x is True else false_str if x is False else x
    if isinstance(report.index, pd.MultiIndex):
        report.index = report.index.set_levels(
            [[map_bool(i) for i in level] for level in report.index.levels]
        )
    return report

def flag_report(report: pd.DataFrame, metric: str = "avg_reward", round_digits: int = 2):
    # for all index in the multi-index with boolean value, get the average for
    # True and the average for False separately. Produce a new dataframe with
    # the average for True and False for each index and the ratio between the
    # two as a new column.

    # check the number of levels
    if report.index.nlevels <= 1:
        print(f"Only {report.index.nlevels} levels in the index, cannot produce flag report.")
        return

    report = report.copy()
    report = report.reset_index()

    records = []
    for col in report.columns:
        if report[col].dtype == bool:
            avg_true = report[report[col]][metric].mean()
            avg_false = report[~report[col]][metric].mean()
            ratio = avg_true / avg_false
            records.append(dict(hparam=col, avg_true=avg_true, avg_false=avg_false, ratio=ratio))

    flag_report = pd.DataFrame(records).set_index("hparam")
    flag_report.sort_values("ratio", ascending=False, inplace=True)
    if round_digits is not None:
        flag_report = flag_report.round(round_digits)

    return flag_report

def display_report(
    report: pd.DataFrame,
    apply_shrink_columns: bool = True,
    copy_to_clipboard: bool = True,
    rename_bool_flags: bool = True,
    print_only: str = None,
):
    """Display the report in a nicer-ish format.

    To be able to wrap col names we need to use set_wrap_stype, which returns a
    styled df, and doesn't behave like a normal df. For encapsulate the displaying in
    this function.

    Args:
        report: The report to display
        apply_shrink_columns: Make the column more compat by replacing
            underscores with newlines
        copy_to_clipboard: Copy the report to the clipboard
        rename_bool_flags: Rename the boolean flags to be more compact and readable
        print_only: Print only the given column
    """
    report = report.copy()

    if apply_shrink_columns:
        report = shrink_columns(report)

    if rename_bool_flags:
        report = _rename_bool_flags(report)

    # if copy_to_clipboard:
    #     to_clipboard(report)

    columns = list(report.columns)

    report.reset_index(inplace=True)

    if print_only:
        columns = [print_only] + columns
        report = report[columns]

    styled_report = set_wrap_style(report)

    display(styled_report)

def shrink_columns(df, also_wrap_index=True):
    """Make the column names more compact by replacing underscores with newlines"""
    df = df.copy()

    df.columns = [col.replace("_", "\n") for col in df.columns]
    if also_wrap_index:
        df.index.names = [name.replace("_", "\n") for name in df.index.names]

    # Define a formatter function that formats float numbers without trailing zeros
    def formatter(x):
        if isinstance(x, float):
            return "{:.10f}".format(x).rstrip("0").rstrip(".")
        return x

    return df.map(formatter)

def set_wrap_style(df):
    return df.style.set_table_styles([{"selector": "th", "props": [("white-space", "pre-wrap")]}])

# ------------
# Error Utils
# ------------

def map_err_key(err_msg: str):
    if err_msg is None:
        return err_msg

    # remove logs from the message if any
    err_msg = err_msg[: err_msg.find("=== logs ===")].rstrip()
    regex_replacements = [
        (
            r"your messages resulted in \d+ tokens",
            "your messages resulted in x tokens",
        ),
        (
            r"(?<=Exception uncaught by agent or environment in task\s)([^\s]+)",
            "<task_name>.",
        ),
    ]

    for pattern, replacement in regex_replacements:
        err_msg = re.sub(pattern, replacement, err_msg)
    return err_msg

def error_report(df: pd.DataFrame, max_stack_trace=10, use_log=False):
    """Report the error message for each agent."""

    if "err_key" not in df:
        df["err_key"] = df["err_msg"].map(map_err_key)

    unique_counts = df["err_key"].value_counts().sort_values(ascending=False)
    report = []
    for err_key, count in unique_counts.items():
        report.append("-------------------")
        report.append(f"## {count
# ... truncated ...


import numpy as np
import pandas as pd
from IPython.display import display
from tqdm import tqdm

from agentlab.experiments.loop import ExpResult, get_exp_result, yield_all_exp_results

# TODO find a more portable way to code set_task_category_as_index at least
# handle dynamic imports. We don't want to always import workarena
# from browsergym.workarena import TASK_CATEGORY_MAP

warnings.filterwarnings("ignore", category=pd.errors.PerformanceWarning)

pd.set_option("display.multi_sparse", False)

AGENT_NAME_KEY = "agent.agent_name"
TASK_KEY = "env.task_name"

def get_constants_and_variables(df: pd.DataFrame, drop_constants: bool = False):
    """Filter out constants from the dataframe."""

    constants = {}
    variable_keys = []
    for col in df.columns:
        try:
            nuniq = df[col].nunique(dropna=False)
        except TypeError:
            nuniq = 0  # non hashable types are considered variables
        if nuniq == 1:
            if isinstance(df[col].iloc[0], np.generic):
                val = df[col].iloc[0].item()
            else:
                val = df[col].iloc[0]
            constants[col] = val
            if drop_constants:
                df = df.drop(col, axis=1)
        else:
            variable_keys.append(col)

    return constants, variable_keys, df

def set_index_from_variables(
    df: pd.DataFrame,
    index_white_list=("agent.*",),
    index_black_list=("*model_url*", "*extra*", "*._*"),
    task_key=TASK_KEY,
    add_agent_and_benchmark=True,
):
    """Set the index, inplace, to env.task_name and all variables.

    Introspects `df` to find all fields that are variable and set the index to
    those fields. This will allow to easily groupby and compare results. To
    filter undersired variables from the index, use index_white_list and
    index_black_list.

    Args:
        df: The dataframe to modify
        index_white_list: List of wildard patterns to match variables that
            should be included in the index.


    variable_keys = []
    for col in df.columns:
        try:
            nuniq = df[col].nunique(dropna=False)
        except TypeError:
            nuniq = 0  # non hashable types are considered variables
        if nuniq == 1:
            if isinstance(df[col].iloc[0], np.generic):
                val = df[col].iloc[0].item()
            else:
                val = df[col].iloc[0]
            constants[col] = val
            if drop_constants:
                df = df.drop(col, axis=1)
        else:
            variable_keys.append(col)

    return constants, variable_keys, df

def set_index_from_variables(
    df: pd.DataFrame,
    index_white_list=("agent.*",),
    index_black_list=("*model_url*", "*extra*", "*._*"),
    task_key=TASK_KEY,
    add_agent_and_benchmark=True,
):
    """Set the index, inplace, to env.task_name and all variables.

    Introspects `df` to find all fields that are variable and set the index to
    those fields. This will allow to easily groupby and compare results. To
    filter undersired variables from the index, use index_white_list and
    index_black_list.

    Args:
        df: The dataframe to modify
        index_white_list: List of wildard patterns to match variables that
            should be included in the index.
        index_black_list: List of wildard patterns to match variables that
            should be excluded from the index.
        task_key: The key to use as the first level of the index.
        add_agent_and_benchmark: If True, add agent.agent_name and env.benchmark
    """
    df.reset_index(inplace=True)
    constants, variables, _ = get_constants_and_variables(df)

    index_variables = []
    if add_agent_and_benchmark:
        index_variables.append("agent.agent_name")
        if "env.benchmark" not in df.columns:
            df["env.benchmark"] = df[TASK_KEY].map(_benchmark_from_task_name)
        index_variables.append("env.benchmark")

    for var in variables:
        white = any([fnmatch.fnmatch(var, pattern) for pattern in index_white_list])
        black = any([fnmatch.fnmatch(var, pattern) for pattern in index_black_list])

        if white and (not black) and (var not in index_variables):
            index_variables.append(var)

    for var in index_variables:
        if df[var].isnull().any():
            warn(
                f"Variable {var} contains NaN or None values. This will be replaced by the string 'None' to avoid some pandas bug."
            )
            df[var] = df[var].fillna("None")

    # agent_variables = [var for var in variables if var.startswith("agent.")]
    df.set_index([task_key] + index_variables, inplace=True)
    df.sort_index(inplace=True)

def load_result_df(
    exp_dir,
    progress_fn=tqdm,
    set_index=True,
    result_df=None,
    index_white_list=("agent.*",),
    index_black_list=("*model_url*", "*extra*", "*._*"),
    remove_args_suffix=True,
):
    """Load the result dataframe.

    Will set the index to env.task_name and all columens that are not constant and
    starts with agent. This will allow to easily groupby and compare
    results. This index can be changed later using df.set_index.

    Args:
        exp_dir: Path to the experiment directory
        progress_fn: Progress function to use when loading the results


    for var in variables:
        white = any([fnmatch.fnmatch(var, pattern) for pattern in index_white_list])
        black = any([fnmatch.fnmatch(var, pattern) for pattern in index_black_list])

        if white and (not black) and (var not in index_variables):
            index_variables.append(var)

    for var in index_variables:
        if df[var].isnull().any():
            warn(
                f"Variable {var} contains NaN or None values. This will be replaced by the string 'None' to avoid some pandas bug."
            )
            df[var] = df[var].fillna("None")

    # agent_variables = [var for var in variables if var.startswith("agent.")]
    df.set_index([task_key] + index_variables, inplace=True)
    df.sort_index(inplace=True)

def load_result_df(
    exp_dir,
    progress_fn=tqdm,
    set_index=True,
    result_df=None,
    index_white_list=("agent.*",),
    index_black_list=("*model_url*", "*extra*", "*._*"),
    remove_args_suffix=True,
):
    """Load the result dataframe.

    Will set the index to env.task_name and all columens that are not constant and
    starts with agent. This will allow to easily groupby and compare
    results. This index can be changed later using df.set_index.

    Args:
        exp_dir: Path to the experiment directory
        progress_fn: Progress function to use when loading the results
        set_index: If True, set the index to env.task_name and variable agent
        result_df: If not None, speed up the loading process by reusing
            alreading loaded objects.
        index_white_list: List of wildard patterns to match variables that
            should be included in the index.
        index_black_list: List of wildard patterns to match variables that
            should be excluded from the index.
        remove_args_suffix: If True, remove the _args suffix from the columns

    Returns:
        pd.DataFrame: The result dataframe
    """

    if result_df is not None:
        result_list = list(result_df["exp_result"])
    else:
        result_list = list(yield_all_exp_results(exp_dir, progress_fn=progress_fn))

    if len(result_list) == 0:
        return None

    if progress_fn is not None:
        result_list = progress_fn(result_list, desc="Loading results")

    df = pd.DataFrame([exp_result.get_exp_record() for exp_result in result_list])

    if remove_args_suffix:
        df.columns = [col.replace("_args", "") for col in df.columns]

    if set_index:
        set_index_from_variables(df, index_white_list, index_black_list)
    return df

def reduce_episodes(result_df: pd.DataFrame) -> pd.DataFrame:
    """Reduce the dataframe to a single row per episode and summarize some of the columns."""

    levels = list(range(result_df.index.nlevels))
    return result_df.groupby(level=levels).apply(summarize)

def report_2d(df: pd.DataFrame, reduce_fn: callable = reduce_episodes, n_row_keys=1):
    """Generic function to create a 2d report based on the dataframe.

    The code is simple but can be a bit cryptic. This is best explained in the
    following  3 steps:
    1) Groupby: Will use the existing multi-index to groupby. Make sure to set the
       an index to the desired keys before calling this function.
    2) Reduce: Uses the reduce_fn to reduce the content of each group to a single
       variable, creating a 1D series indexed by its original index.
    3) Unstack: Produce a 2D table such that the first n_row_keys are used to
       specify how many dimensions are used for the rows. The remaining


        result_list = list(result_df["exp_result"])
    else:
        result_list = list(yield_all_exp_results(exp_dir, progress_fn=progress_fn))

    if len(result_list) == 0:
        return None

    if progress_fn is not None:
        result_list = progress_fn(result_list, desc="Loading results")

    df = pd.DataFrame([exp_result.get_exp_record() for exp_result in result_list])

    if remove_args_suffix:
        df.columns = [col.replace("_args", "") for col in df.columns]

    if set_index:
        set_index_from_variables(df, index_white_list, index_black_list)
    return df

def reduce_episodes(result_df: pd.DataFrame) -> pd.DataFrame:
    """Reduce the dataframe to a single row per episode and summarize some of the columns."""

    levels = list(range(result_df.index.nlevels))
    return result_df.groupby(level=levels).apply(summarize)

def report_2d(df: pd.DataFrame, reduce_fn: callable = reduce_episodes, n_row_keys=1):
    """Generic function to create a 2d report based on the dataframe.

    The code is simple but can be a bit cryptic. This is best explained in the
    following  3 steps:
    1) Groupby: Will use the existing multi-index to groupby. Make sure to set the
       an index to the desired keys before calling this function.
    2) Reduce: Uses the reduce_fn to reduce the content of each group to a single
       variable, creating a 1D series indexed by its original index.
    3) Unstack: Produce a 2D table such that the first n_row_keys are used to
       specify how many dimensions are used for the rows. The remaining
       dimensions are used for the columns.

    Args:
        df: The dataframe to reduce
        reduce_fn: The function to use to reduce the sub dataframe. By default
            this is reduce_episodes.
        n_row_keys: The number of keys to use for the rows.


    if progress_fn is not None:
        result_list = progress_fn(result_list, desc="Loading results")

    df = pd.DataFrame([exp_result.get_exp_record() for exp_result in result_list])

    if remove_args_suffix:
        df.columns = [col.replace("_args", "") for col in df.columns]

    if set_index:
        set_index_from_variables(df, index_white_list, index_black_list)
    return df

def reduce_episodes(result_df: pd.DataFrame) -> pd.DataFrame:
    """Reduce the dataframe to a single row per episode and summarize some of the columns."""

    levels = list(range(result_df.index.nlevels))
    return result_df.groupby(level=levels).apply(summarize)

def report_2d(df: pd.DataFrame, reduce_fn: callable = reduce_episodes, n_row_keys=1):
    """Generic function to create a 2d report based on the dataframe.

    The code is simple but can be a bit cryptic. This is best explained in the
    following  3 steps:
    1) Groupby: Will use the existing multi-index to groupby. Make sure to set the
       an index to the desired keys before calling this function.
    2) Reduce: Uses the reduce_fn to reduce the content of each group to a single
       variable, creating a 1D series indexed by its original index.
    3) Unstack: Produce a 2D table such that the first n_row_keys are used to
       specify how many dimensions are used for the rows. The remaining
       dimensions are used for the columns.

    Args:
        df: The dataframe to reduce
        reduce_fn: The function to use to reduce the sub dataframe. By default
            this is reduce_episodes.
        n_row_keys: The number of keys to use for the rows.

    Returns:
        pd.DataFrame: The 2D report
    """

    levels = list(range(df.index.nlevels))
    reduced_df = df.groupby(level=levels).apply(reduce_fn)  # type: pd.Series
    return reduced_df.unstack(level=levels[n_row_keys:])

def report_constant_and_variables(df, show_stack_traces=True):
    constants, variables, _ = get_constants_and_variables(df)
    print("Constants:")
    for k, v in constants.items():
        print(f"    {k}: {v}")

    print("\nVariables:")
    for var in variables:
        if not show_stack_traces and var == "stack_trace":
            continue

        # get unique with count and sort by count descending
        unique_counts = df[var].value_counts().sort_values(ascending=False)

        print(f"    {var}: n_unique={len(unique_counts)}")
        for i, (val, count) in enumerate(unique_counts.items()):
            print(f"        {count}x : {val}")
            if i >= 2:


       variable, creating a 1D series indexed by its original index.
    3) Unstack: Produce a 2D table such that the first n_row_keys are used to
       specify how many dimensions are used for the rows. The remaining
       dimensions are used for the columns.

    Args:
        df: The dataframe to reduce
        reduce_fn: The function to use to reduce the sub dataframe. By default
            this is reduce_episodes.
        n_row_keys: The number of keys to use for the rows.

    Returns:
        pd.DataFrame: The 2D report
    """

    levels = list(range(df.index.nlevels))
    reduced_df = df.groupby(level=levels).apply(reduce_fn)  # type: pd.Series
    return reduced_df.unstack(level=levels[n_row_keys:])

def report_constant_and_variables(df, show_stack_traces=True):
    constants, variables, _ = get_constants_and_variables(df)
    print("Constants:")
    for k, v in constants.items():
        print(f"    {k}: {v}")

    print("\nVariables:")
    for var in variables:
        if not show_stack_traces and var == "stack_trace":
            continue

        # get unique with count and sort by count descending
        unique_counts = df[var].value_counts().sort_values(ascending=False)

        print(f"    {var}: n_unique={len(unique_counts)}")
        for i, (val, count) in enumerate(unique_counts.items()):
            print(f"        {count}x : {val}")
            if i >= 2:
                break
        if len(unique_counts) > 3:
            print("        ...\n")

def get_std_err(df, metric):
    """Get the standard error for a binary metric."""
    # extract non missing values
    data = df[metric].dropna().values

    # asser either 0 or 1
    if np.all(np.isin(data, [0, 1])):
        mean = np.mean(data)
        std_err = np.sqrt(mean * (1 - mean) / len(data))
        return mean, std_err
    else:
        return get_sample_std_err(df, metric)

def get_sample_std_err(df, metric):
    """Get the standard error for a binary metric."""
    # extract non missing values
    data = df[metric].dropna().values


    for k, v in constants.items():
        print(f"    {k}: {v}")

    print("\nVariables:")
    for var in variables:
        if not show_stack_traces and var == "stack_trace":
            continue

        # get unique with count and sort by count descending
        unique_counts = df[var].value_counts().sort_values(ascending=False)

        print(f"    {var}: n_unique={len(unique_counts)}")
        for i, (val, count) in enumerate(unique_counts.items()):
            print(f"        {count}x : {val}")
            if i >= 2:
                break
        if len(unique_counts) > 3:
            print("        ...\n")

def get_std_err(df, metric):
    """Get the standard error for a binary metric."""
    # extract non missing values
    data = df[metric].dropna().values

    # asser either 0 or 1
    if np.all(np.isin(data, [0, 1])):
        mean = np.mean(data)
        std_err = np.sqrt(mean * (1 - mean) / len(data))
        return mean, std_err
    else:
        return get_sample_std_err(df, metric)

def get_sample_std_err(df, metric):
    """Get the standard error for a binary metric."""
    # extract non missing values
    data = df[metric].dropna().values

    mean = np.mean(data)
    std_err = np.std(data, ddof=1) / np.sqrt(len(data))
    if np.isnan(std_err):
        std_err = np.float64(0)
    return mean, std_err

def summarize(sub_df):
    if "cum_reward" not in sub_df:
        record = dict(
            avg_reward=np.nan,
            std_err=np.nan,
            # avg_raw_reward=np.nan,


            if i >= 2:
                break
        if len(unique_counts) > 3:
            print("        ...\n")

def get_std_err(df, metric):
    """Get the standard error for a binary metric."""
    # extract non missing values
    data = df[metric].dropna().values

    # asser either 0 or 1
    if np.all(np.isin(data, [0, 1])):
        mean = np.mean(data)
        std_err = np.sqrt(mean * (1 - mean) / len(data))
        return mean, std_err
    else:
        return get_sample_std_err(df, metric)

def get_sample_std_err(df, metric):
    """Get the standard error for a binary metric."""
    # extract non missing values
    data = df[metric].dropna().values

    mean = np.mean(data)
    std_err = np.std(data, ddof=1) / np.sqrt(len(data))
    if np.isnan(std_err):
        std_err = np.float64(0)
    return mean, std_err

def summarize(sub_df):
    if "cum_reward" not in sub_df:
        record = dict(
            avg_reward=np.nan,
            std_err=np.nan,
            # avg_raw_reward=np.nan,
            avg_steps=np.nan,
            n_completed=f"0/{len(sub_df)}",
            n_err=0,
        )
    else:
        err = sub_df["err_msg"].notnull()
        n_completed = err.copy()
        for col in ["truncated", "terminated"]:
            if col in sub_df:
                n_completed = n_completed | sub_df[col]
        n_completed = n_completed.sum()


summarize

    if np.all(np.isin(data, [0, 1])):
        mean = np.mean(data)
        std_err = np.sqrt(mean * (1 - mean) / len(data))
        return mean, std_err
    else:
        return get_sample_std_err(df, metric)

def get_sample_std_err(df, metric):
    """Get the standard error for a binary metric."""
    # extract non missing values
    data = df[metric].dropna().values

    mean = np.mean(data)
    std_err = np.std(data, ddof=1) / np.sqrt(len(data))
    if np.isnan(std_err):
        std_err = np.float64(0)
    return mean, std_err

def summarize(sub_df):
    if "cum_reward" not in sub_df:
        record = dict(
            avg_reward=np.nan,
            std_err=np.nan,
            # avg_raw_reward=np.nan,
            avg_steps=np.nan,
            n_completed=f"0/{len(sub_df)}",
            n_err=0,
        )
    else:
        err = sub_df["err_msg"].notnull()
        n_completed = err.copy()
        for col in ["truncated", "terminated"]:
            if col in sub_df:
                n_completed = n_completed | sub_df[col]
        n_completed = n_completed.sum()

        if n_completed == 0:
            return None

        _mean_reward, std_reward = get_std_err(sub_df, "cum_reward")

        # sanity check, if there is an error the reward should be zero
        assert sub_df[sub_df["err_msg"].notnull()]["cum_reward"].sum() == 0

        record = dict(
            avg_reward=sub_df["cum_reward"].mean(skipna=True).round(3),
            std_err=std_reward.round(3),
            # avg_raw_reward=sub_df["cum_raw_reward"].mean(skipna=True).round(3),
            avg_steps=sub_df["n_steps"].mean(skipna=True).round(3),
            n_completed=f"{n_completed}/{len(sub_df)}",
            n_err=err.sum(skipna=True),
        )
        if "stats.cum_cost" in sub_df:
            record["cum_cost"] = sub_df["stats.cum_cost"].sum(skipna=True).round(4)
        if "stats.cum_effective_cost" in sub_df:
            record["cum_effective_cost"] = (
                sub_df["stats.cum_effective_cost"].sum(skipna=True).round(4)
            )
            record.pop("cum_cost", None)

    return pd.Series(record)

def summarize_stats(sub_df):
    """Summarize the stats columns."""

    # make sure there are completed runs
    err = sub_df["err_msg"].notnull()
    n_completed = err.copy()
    for col in ["truncated", "terminated"]:
        if col in sub_df:
            n_completed = n_completed | sub_df[col]
    n_completed = n_completed.sum()

    if n_completed == 0:
        return None

    record = dict(
        avg_reward=sub_df["cum_reward"].mean(skipna=True).round(3),
    )
    for key in sub_df.keys():


        record = dict(
            avg_reward=sub_df["cum_reward"].mean(skipna=True).round(3),
            std_err=std_reward.round(3),
            # avg_raw_reward=sub_df["cum_raw_reward"].mean(skipna=True).round(3),
            avg_steps=sub_df["n_steps"].mean(skipna=True).round(3),
            n_completed=f"{n_completed}/{len(sub_df)}",
            n_err=err.sum(skipna=True),
        )
        if "stats.cum_cost" in sub_df:
            record["cum_cost"] = sub_df["stats.cum_cost"].sum(skipna=True).round(4)
        if "stats.cum_effective_cost" in sub_df:
            record["cum_effective_cost"] = (
                sub_df["stats.cum_effective_cost"].sum(skipna=True).round(4)
            )
            record.pop("cum_cost", None)

    return pd.Series(record)

def summarize_stats(sub_df):
    """Summarize the stats columns."""

    # make sure there are completed runs
    err = sub_df["err_msg"].notnull()
    n_completed = err.copy()
    for col in ["truncated", "terminated"]:
        if col in sub_df:
            n_completed = n_completed | sub_df[col]
    n_completed = n_completed.sum()

    if n_completed == 0:
        return None

    record = dict(
        avg_reward=sub_df["cum_reward"].mean(skipna=True).round(3),
    )
    for key in sub_df.keys():
        if key.startswith("stats."):
            key_ = key.split(".")[1]
            op = key_.split("_")[0]
            if op == "cum":
                record[key_] = sub_df[key].sum(skipna=True)
            elif op == "max":
                record[key_] = sub_df[key].max(skipna=True)
            else:
                raise ValueError(f"Unknown stats operation: {op}")
    return pd.Series(record)

def _find_diff(tuple1, tuple2):
    """return the list of index wher tuple1 != tuple2"""
    return [i for i, (a, b) in enumerate(zip(tuple1, tuple2)) if a != b]

def _extract_ablation_study(report: pd.DataFrame, progression=False):
    """Reduce the multi-index to a change description compared to the previous row."""
    names = report.index.names
    report = report.copy()
    # report.sort_index(inplace=True)

    reference_index = None
    for index in report.index:
        if reference_index is not None:
            diffs = _find_diff(reference_index, index)

            if progression:
                change = "↳ " + ", ".join([f"{names[i]}={index[i]}" for i in diffs])


    if n_completed == 0:
        return None

    record = dict(
        avg_reward=sub_df["cum_reward"].mean(skipna=True).round(3),
    )
    for key in sub_df.keys():
        if key.startswith("stats."):
            key_ = key.split(".")[1]
            op = key_.split("_")[0]
            if op == "cum":
                record[key_] = sub_df[key].sum(skipna=True)
            elif op == "max":
                record[key_] = sub_df[key].max(skipna=True)
            else:
                raise ValueError(f"Unknown stats operation: {op}")
    return pd.Series(record)

def _find_diff(tuple1, tuple2):
    """return the list of index wher tuple1 != tuple2"""
    return [i for i, (a, b) in enumerate(zip(tuple1, tuple2)) if a != b]

def _extract_ablation_study(report: pd.DataFrame, progression=False):
    """Reduce the multi-index to a change description compared to the previous row."""
    names = report.index.names
    report = report.copy()
    # report.sort_index(inplace=True)

    reference_index = None
    for index in report.index:
        if reference_index is not None:
            diffs = _find_diff(reference_index, index)

            if progression:
                change = "↳ " + ", ".join([f"{names[i]}={index[i]}" for i in diffs])
            else:
                changes = []
                for i in diffs:
                    val = index[i]
                    if isinstance(val, bool):


        avg_reward=sub_df["cum_reward"].mean(skipna=True).round(3),
    )
    for key in sub_df.keys():
        if key.startswith("stats."):
            key_ = key.split(".")[1]
            op = key_.split("_")[0]
            if op == "cum":
                record[key_] = sub_df[key].sum(skipna=True)
            elif op == "max":
                record[key_] = sub_df[key].max(skipna=True)
            else:
                raise ValueError(f"Unknown stats operation: {op}")
    return pd.Series(record)

def _find_diff(tuple1, tuple2):
    """return the list of index wher tuple1 != tuple2"""
    return [i for i, (a, b) in enumerate(zip(tuple1, tuple2)) if a != b]

def _extract_ablation_study(report: pd.DataFrame, progression=False):
    """Reduce the multi-index to a change description compared to the previous row."""
    names = report.index.names
    report = report.copy()
    # report.sort_index(inplace=True)

    reference_index = None
    for index in report.index:
        if reference_index is not None:
            diffs = _find_diff(reference_index, index)

            if progression:
                change = "↳ " + ", ".join([f"{names[i]}={index[i]}" for i in diffs])
            else:
                changes = []
                for i in diffs:
                    val = index[i]
                    if isinstance(val, bool):
                        changes.append(("+" if val else "-") + names[i])
                    else:
                        changes.append(f"{names[i]}←{val}")
                change = ", ".join(changes)
        else:
            change = "Initial Configuration"
        report.loc[index, "change"] = change
        if progression:
            reference_index = index
        else:
            reference_index = report.index[0]

    report = report.reset_index()
    report = report.set_index(["change"])

    # delete columns related to old index
    return report.drop(names, axis=1)

def ablation_report(result_df: pd.DataFrame, reduce_fn=summarize, progression=False):
    """Reduce the multi-index to a change description compared to the previous row.

    *NOTE*: This assumes that this experiments was launched with make_ablation_study.

    Rows will be sorted according to the average ExpArgs.order for all
    experiments associated with the multi-index.

    Args:
        result_df: The result dataframe as returned by load_result_df.
        reduce_fn: The function to use to reduce the sub dataframe. By default
            this is summarize.
        progression: If True, the change description will be the progression

    Returns:
        A dataframe with the change description as index.
    """
    report = global_report(result_df, reduce_fn=reduce_fn)


                    if isinstance(val, bool):
                        changes.append(("+" if val else "-") + names[i])
                    else:
                        changes.append(f"{names[i]}←{val}")
                change = ", ".join(changes)
        else:
            change = "Initial Configuration"
        report.loc[index, "change"] = change
        if progression:
            reference_index = index
        else:
            reference_index = report.index[0]

    report = report.reset_index()
    report = report.set_index(["change"])

    # delete columns related to old index
    return report.drop(names, axis=1)

def ablation_report(result_df: pd.DataFrame, reduce_fn=summarize, progression=False):
    """Reduce the multi-index to a change description compared to the previous row.

    *NOTE*: This assumes that this experiments was launched with make_ablation_study.

    Rows will be sorted according to the average ExpArgs.order for all
    experiments associated with the multi-index.

    Args:
        result_df: The result dataframe as returned by load_result_df.
        reduce_fn: The function to use to reduce the sub dataframe. By default
            this is summarize.
        progression: If True, the change description will be the progression

    Returns:
        A dataframe with the change description as index.
    """
    report = global_report(result_df, reduce_fn=reduce_fn)
    report = _sort_order(result_df, report)
    report = _extract_ablation_study(report, progression=progression)
    return report

def _get_avg_order(df: pd.DataFrame, row: pd.Series):
    """Return the average order for the given row."""
    df = df.reset_index(level=0, drop=True, inplace=False)
    # df.sort_index(inplace=True)

    sub_df = df.loc[row.name]
    orders = [get_exp_result(exp_dir).exp_args.order for exp_dir in sub_df.exp_dir]
    orders = [order for order in orders if order is not None]
    if len(orders) == 0:
        return None
    return np.mean(orders)

def _sort_order(result_df, report):
    """Add a column to the report with the average order for each agent and sort."""

    def add_order(row):
        return _get_avg_order(result_df, row)


    *NOTE*: This assumes that this experiments was launched with make_ablation_study.

    Rows will be sorted according to the average ExpArgs.order for all
    experiments associated with the multi-index.

    Args:
        result_df: The result dataframe as returned by load_result_df.
        reduce_fn: The function to use to reduce the sub dataframe. By default
            this is summarize.
        progression: If True, the change description will be the progression

    Returns:
        A dataframe with the change description as index.
    """
    report = global_report(result_df, reduce_fn=reduce_fn)
    report = _sort_order(result_df, report)
    report = _extract_ablation_study(report, progression=progression)
    return report

def _get_avg_order(df: pd.DataFrame, row: pd.Series):
    """Return the average order for the given row."""
    df = df.reset_index(level=0, drop=True, inplace=False)
    # df.sort_index(inplace=True)

    sub_df = df.loc[row.name]
    orders = [get_exp_result(exp_dir).exp_args.order for exp_dir in sub_df.exp_dir]
    orders = [order for order in orders if order is not None]
    if len(orders) == 0:
        return None
    return np.mean(orders)

def _sort_order(result_df, report):
    """Add a column to the report with the average order for each agent and sort."""

    def add_order(row):
        return _get_avg_order(result_df, row)

    report["avg_order"] = report.apply(add_order, axis=1)
    return report.sort_values("avg_order", ascending=True)

def global_report(
    result_df: pd.DataFrame,
    reduce_fn=summarize,
    rename_index=lambda name: name.replace("agent.flags.", ""),
):
    """Produce a report that summarize all tasks and all episodes for each
    agent.


    """
    report = global_report(result_df, reduce_fn=reduce_fn)
    report = _sort_order(result_df, report)
    report = _extract_ablation_study(report, progression=progression)
    return report

def _get_avg_order(df: pd.DataFrame, row: pd.Series):
    """Return the average order for the given row."""
    df = df.reset_index(level=0, drop=True, inplace=False)
    # df.sort_index(inplace=True)

    sub_df = df.loc[row.name]
    orders = [get_exp_result(exp_dir).exp_args.order for exp_dir in sub_df.exp_dir]
    orders = [order for order in orders if order is not None]
    if len(orders) == 0:
        return None
    return np.mean(orders)

def _sort_order(result_df, report):
    """Add a column to the report with the average order for each agent and sort."""

    def add_order(row):
        return _get_avg_order(result_df, row)

    report["avg_order"] = report.apply(add_order, axis=1)
    return report.sort_values("avg_order", ascending=True)

def global_report(
    result_df: pd.DataFrame,
    reduce_fn=summarize,
    rename_index=lambda name: name.replace("agent.flags.", ""),
):
    """Produce a report that summarize all tasks and all episodes for each
    agent.

    Args:
        result_df: The result dataframe as returned by load_result_df.
        reduce_fn: The function to use to reduce the sub dataframe. By default
            this is summarize.
        rename_index: Function to rename the index. By default we remove the prefix
            "agent.flags."

    Returns:
        pd.DataFrame: The report
    """