tree.py

import asyncio
import heapq
import json
import os
import random
import re
import string
from collections import deque
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple

from general_prompt import (
    BLOCKSWORLD_EXECUTABILITY_EXAMPLE,
    EXECUTABILITY_CHECK_PROMPT,
    GENERATE_AVAILABLE_ACTIONS_PROMPT,
    GOAL_EXTRACTOR_PROMPT_ACP,
    GOAL_EXTRACTOR_PROMPT_PLANBENCH,
    INITIAL_STATE_EXTRACTOR_PROMPT_PLANBENCH,
    PROGRESSION_PROMPT,
    STATE_CHECK_PROMPT,
    STATE_SCORE_PROMPT,
)
from utils import AsyncRequestSender, ContentExtraction, RequestSender, readPDDL

# Import the new heuristic computation system
try:
    from heuristic_computer import clear_lora_model_cache as new_clear_lora_model_cache
    from heuristic_computer import (
        compute_heuristic_batch_with_lora as new_compute_heuristic_batch_with_lora,
    )
    from heuristic_computer import (
        compute_heuristic_with_lora as new_compute_heuristic_with_lora,
    )

    HEURISTIC_SYSTEM_AVAILABLE = True
except ImportError as exc:
    HEURISTIC_SYSTEM_AVAILABLE = False
    new_compute_heuristic_batch_with_lora = None
    HEURISTIC_IMPORT_ERROR = exc


def _load_pddl(filename: str) -> str:
    """尝试多个相对路径加载 PDDL 文件，全部失败则返回空字符串（不阻断启动）。"""
    search_paths = [
        f"../domain_pddl/{filename}",
        f"./domain_pddl/{filename}",
        f"domain_pddl/{filename}",
        os.path.join(
            os.path.dirname(os.path.abspath(__file__)), "domain_pddl", filename
        ),
    ]
    for path in search_paths:
        try:
            content = readPDDL(path)
            return content
        except FileNotFoundError:
            continue
    print(f"[警告] 无法加载 PDDL 文件: {filename}，已跳过。")
    return ""


DOMAIN_DESCRIPTION_BLOCKSWORLD = _load_pddl("blocksworld.pddl")
DOMAIN_DESCRIPTION_LOGISTICS = _load_pddl("logistics.pddl")
DOMAIN_DESCRIPTION_DEPOT = _load_pddl("depots.pddl")

DOMAIN_MAP = {
    "blocksworld": DOMAIN_DESCRIPTION_BLOCKSWORLD,
    "logistics": DOMAIN_DESCRIPTION_LOGISTICS,
    "depots": DOMAIN_DESCRIPTION_DEPOT,
    "depot": DOMAIN_DESCRIPTION_DEPOT,
}


def _get_domain_description(domain: str) -> str:
    """根据领域名称返回对应的 PDDL 领域描述字符串。"""
    return DOMAIN_MAP.get(domain.lower(), DOMAIN_DESCRIPTION_BLOCKSWORLD)


def normalize_state_for_search(state: str) -> str:
    """Normalize a controlled-NL state into an order-insensitive search key."""
    predicates = re.split(r"[.,]|\band\b", state.strip().lower())
    return "|".join(sorted(predicate.strip() for predicate in predicates if predicate.strip()))


def record_better_expansion(
    closed_best_g: Dict[str, float], normalized_state: str, g_cost: float
) -> bool:
    """Record an expansion and allow reopening only when g strictly improves."""
    previous = closed_best_g.get(normalized_state)
    if previous is not None and previous <= g_cost:
        return False
    closed_best_g[normalized_state] = g_cost
    return True


class Node:
    """
    Description
    -----------
    Node of the tree, each node represents a state.

    Parameters
    ----------
    node_id : the ID of the node, which is a unique identifier for the node.
    parent_id : the parent id of this node.
    children_id : children of this node, corresponding to available_actions.
    available_actions : list, each element is a list of [available_action, ...].
    state : the state of the tree node.
    action_taken : the last action taken from parent to this node.
    history_actions : the actions taken from root to this node.
    g_cost : the cost from root to this node.
    h_cost : the cost from this node to goal, estimated by LLM.
    status : whether this node is active or pruned.
    flag : whether this state is the goal state.
    pruned_reason : if pruned, the reason for pruning.
    """

    def __init__(
        self,
        node_id: int,
        parent_id: Optional[int] = None,
        state: str = "",
        query: str = "",
        g_cost: float = 0.0,
        h_cost: float = 0.0,
        **kwargs,
    ):

        self.node_id: int = node_id
        self.parent_id: Optional[int] = parent_id
        self.children_id: List[int] = []
        self.generated_actions: List[str] = kwargs.get("generated_actions", [])
        self.available_actions: Optional[List[List[Any]]] = kwargs.get(
            "available_actions", []
        )
        self.state: str = state
        self.state_score: float = kwargs.get("state_score", 5.0)
        self.action_taken: Optional[List[str]] = kwargs.get("action_taken", [])
        self.history_actions: List[str] = kwargs.get("history_actions", [])
        self.query: str = query
        self.g_cost: float = g_cost
        self.h_cost: float = h_cost
        self.flag: bool = False  # True means this state is the goal state
        self.status: bool = True  # True (1) = active, False (0) = pruned
        self.pruned_reason: Optional[str] = (
            None  # INFINITE_HEURISTIC_VALUE or DOMAIN_KNOWLEDGE_VIOLATION
        )

    @property
    def f_cost(self) -> float:
        """count total cost f, f = g+h = actual cost + heuristic cost."""
        h_val = self.h_cost if self.h_cost is not None else 0.0
        return self.g_cost + h_val

    def __lt__(self, other: "Node") -> bool:
        """compare nodes for priority queue (A* algorithm)."""
        return self.f_cost < other.f_cost

    def to_dict(self) -> Dict[str, Any]:
        """Convert node to dictionary for JSON serialization."""
        return {
            "node_id": self.node_id,
            "parent_id": self.parent_id,
            "children_id": self.children_id,
            "generated_actions": self.generated_actions,
            "available_actions": self.available_actions,
            "state": self.state,
            "state_score": self.state_score,
            "action_taken": self.action_taken,
            "history_actions": self.history_actions,
            "query": self.query,
            "g_cost": self.g_cost,
            "h_cost": self.h_cost,
            "f_cost": self.f_cost,
            "flag": self.flag,
            "status": self.status,
            "pruned_reason": self.pruned_reason,
        }


def extract_initial_state(raw_text: str, max_retries: int = 3, **kwargs) -> str:
    """
    Extract initial state from raw text.

    Parameters
    ----------
    raw_text : str
        Raw text containing initial state description
    max_retries : int
        Maximum number of retry attempts
    **kwargs : dict
        LLM parameters (model, api_key, base_url)

    Returns
    -------
    str
        Extracted initial state
    """

    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }

    prompt = (
        f"{DOMAIN_DESCRIPTION_BLOCKSWORLD}\n"
        f"[Initial State needed to be extracted]{raw_text}\n"
        f"{INITIAL_STATE_EXTRACTOR_PROMPT_PLANBENCH}"
    )

    for attempt in range(max_retries):
        response = RequestSender(prompt, **LLM_params)
        extracted = ContentExtraction(response)

        # Check if extraction was successful
        if not extracted.startswith("Extraction failed"):
            return extracted

        print(
            f"[Warning] Initial state extraction failed (attempt {attempt + 1}/{max_retries}): {extracted}"
        )

    # All retries failed
    print(f"[Error] Initial state extraction failed after {max_retries} attempts")
    return extracted


def extract_goal(raw_text: str, max_retries: int = 3, **kwargs) -> str:
    """
    Extract goal from raw text.

    Parameters
    ----------
    raw_text : str
        Raw text containing goal description
    max_retries : int
        Maximum number of retry attempts
    **kwargs : dict
        LLM parameters (model, api_key, base_url)

    Returns
    -------
    str
        Extracted goal
    """

    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }

    prompt = (
        f"[Raw text that needed to be extracted as goal]{raw_text}\n"
        f"{GOAL_EXTRACTOR_PROMPT_PLANBENCH}\n"
    )

    for attempt in range(max_retries):
        response = RequestSender(prompt, **LLM_params)
        goal = ContentExtraction(response)

        # Check if extraction was successful
        if not goal.startswith("Extraction failed"):
            return goal

        print(
            f"[Warning] Goal extraction failed (attempt {attempt + 1}/{max_retries}): {goal}"
        )

    # All retries failed
    print(f"[Error] Goal extraction failed after {max_retries} attempts")
    return goal


def check_goal_state(
    state: str,
    goal: str,
    max_retries: int = 3,
    domain_description: str = None,
    **kwargs,
) -> bool:
    """
    Check if the current state satisfies the goal.

    Parameters
    ----------
    state : str
        Current state description
    goal : str
        Goal state description
    max_retries : int
        Maximum number of retry attempts
    **kwargs : dict
        LLM parameters (model, api_key, base_url)

    Returns
    -------
    bool
        True if state satisfies goal, False otherwise
    """
    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }

    domain_desc = domain_description or DOMAIN_DESCRIPTION_BLOCKSWORLD
    prompt = (
        f"{domain_desc}\n"
        f"[Current State]{state}\n"
        f"[Goal State]{goal}\n"
        f"{STATE_CHECK_PROMPT}\n"
    )

    for attempt in range(max_retries):
        response = RequestSender(prompt, **LLM_params)
        extracted = ContentExtraction(response)

        # Check if extraction was successful
        if not extracted.startswith("Extraction failed"):
            result = extracted.strip().lower()
            return result == "true"

        print(
            f"[Warning] Goal check extraction failed (attempt {attempt + 1}/{max_retries}): {extracted}"
        )

    # All retries failed
    print(f"[Error] Goal check extraction failed after {max_retries} attempts")
    return False


def generate_available_actions(
    state: str,
    max_retries: int = 3,
    last_action: str = None,
    domain_description: str = None,
    **kwargs,
) -> List[str]:
    """
    Generate all available actions from the current state.

    Parameters
    ----------
    state : str
        Current state description
    max_retries : int
        Maximum number of retry attempts
    last_action : str, optional
        The last action taken (to avoid generating inverse actions)
    **kwargs : dict
        LLM parameters (model, api_key, base_url)

    Returns
    -------
    List[str]
        List of available actions
    """
    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }

    # Build prompt with optional last action hint
    last_action_str = last_action if last_action else "None"
    domain_desc = domain_description or DOMAIN_DESCRIPTION_BLOCKSWORLD

    prompt = (
        f"{domain_desc}\n"
        f"[Current State]{state}\n"
        f"[Last executed action]{last_action_str}\n"
        f"{GENERATE_AVAILABLE_ACTIONS_PROMPT}\n"
    )

    for attempt in range(max_retries):
        response = RequestSender(prompt, **LLM_params)

        # Check if request failed
        if not response.startswith("Request Failed"):
            # Extract content from \box{...}
            extracted = ContentExtraction(response)

            # Debug: print raw response and extracted content
            # print(f"[DEBUG] Raw LLM response (first 500 chars): {response[:500]}...")
            # print(f"[DEBUG] Extracted content: {extracted}")

            if extracted.startswith("Extraction failed"):
                print(
                    f"[Warning] Action extraction failed (attempt {attempt + 1}/{max_retries}): {extracted}"
                )
                continue

            # Split actions by period and clean them
            actions = []
            for action in extracted.split("."):
                action = action.strip()
                if action:  # Only add non-empty actions
                    actions.append(action)

            print(f"[DEBUG] Parsed actions: {actions}")

            if not actions:
                print(
                    f"[Warning] No actions parsed from extracted content (attempt {attempt + 1}/{max_retries})"
                )
                continue

            return actions

        print(
            f"[Warning] Generate actions request failed (attempt {attempt + 1}/{max_retries}): {response}"
        )

    # All retries failed
    print(f"[Error] Generate actions failed after {max_retries} attempts")
    return []


def check_action_executability(
    state: str,
    action: str,
    max_retries: int = 3,
    domain_description: str = None,
    **kwargs,
) -> bool:
    """
    Check if an action is executable in the current state.

    Parameters
    ----------
    state : str
        Current state description
    action : str
        Action to check
    max_retries : int
        Maximum number of retry attempts
    **kwargs : dict
        LLM parameters (model, api_key, base_url)

    Returns
    -------
    bool
        True if action is executable, False otherwise
    """
    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }

    domain_desc = domain_description or DOMAIN_DESCRIPTION_BLOCKSWORLD
    prompt = (
        f"{domain_desc}\n"
        f"{BLOCKSWORLD_EXECUTABILITY_EXAMPLE}\n"
        f"[Current State]{state}\n"
        f"[Action]{action}\n"
        f"{EXECUTABILITY_CHECK_PROMPT}\n"
    )

    for attempt in range(max_retries):
        response = RequestSender(prompt, **LLM_params)
        extracted = ContentExtraction(response)

        # Check if extraction was successful
        if not extracted.startswith("Extraction failed"):
            is_executable = _parse_boxed_bool(extracted)
            # print(
            #     f"[DEBUG] Executability check for '{action}': extracted='{extracted}', result={is_executable}"
            # )
            return is_executable

        print(
            f"[Warning] Executability check extraction failed (attempt {attempt + 1}/{max_retries}): {extracted}"
        )
        # print(f"[DEBUG] Raw response: {response[:300]}...")

    # All retries failed
    print(f"[Error] Executability check failed after {max_retries} attempts")
    return False


async def async_check_action_executability(
    state: str,
    action: str,
    max_retries: int = 3,
    domain_description: str = None,
    **kwargs,
) -> bool:
    """Async version of check_action_executability for concurrent execution."""
    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }
    domain_desc = domain_description or DOMAIN_DESCRIPTION_BLOCKSWORLD
    prompt = (
        f"{domain_desc}\n"
        f"{BLOCKSWORLD_EXECUTABILITY_EXAMPLE}\n"
        f"[Current State]{state}\n"
        f"[Action]{action}\n"
        f"{EXECUTABILITY_CHECK_PROMPT}\n"
    )
    for attempt in range(max_retries):
        response = await AsyncRequestSender(prompt, **LLM_params)
        extracted = ContentExtraction(response)
        if not extracted.startswith("Extraction failed"):
            return _parse_boxed_bool(extracted)
        print(
            f"[Warning] Async executability check failed (attempt {attempt + 1}/{max_retries})"
        )
    return False


async def async_take_action(
    state: str,
    action: str,
    max_retries: int = 3,
    domain_description: str = None,
    **kwargs,
) -> str:
    """Async version of take_action for concurrent execution."""
    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }
    domain_desc = domain_description or DOMAIN_DESCRIPTION_BLOCKSWORLD
    from general_prompt import PROGRESSION_EXAMPLE_PROMPT
    prompt = (
        f"{domain_desc}\n"
        f"{PROGRESSION_EXAMPLE_PROMPT}\n"
        f"[Current State]{state}\n"
        f"[Action]{action}\n"
        f"{PROGRESSION_PROMPT}\n"
    )
    for attempt in range(max_retries):
        response = await AsyncRequestSender(prompt, **LLM_params)
        extracted = ContentExtraction(response)
        if not extracted.startswith("Extraction failed"):
            return extracted
        print(
            f"[Warning] Async state progression failed (attempt {attempt + 1}/{max_retries})"
        )
    return state  # Return original state if all retries fail


async def async_check_goal_state(
    state: str,
    goal: str,
    max_retries: int = 3,
    domain_description: str = None,
    **kwargs,
) -> bool:
    """Async version of check_goal_state for concurrent execution."""
    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }
    domain_desc = domain_description or DOMAIN_DESCRIPTION_BLOCKSWORLD
    prompt = (
        f"{domain_desc}\n"
        f"[Current State]{state}\n"
        f"[Goal State]{goal}\n"
        f"{STATE_CHECK_PROMPT}\n"
    )
    for attempt in range(max_retries):
        response = await AsyncRequestSender(prompt, **LLM_params)
        extracted = ContentExtraction(response)
        if not extracted.startswith("Extraction failed"):
            return extracted.strip().lower() == "true"
        print(f"[Warning] Async goal check failed (attempt {attempt + 1}/{max_retries})")
    return False


def _parse_boxed_bool(extracted: str) -> bool:
    """Parse boxed boolean answers while tolerating light punctuation noise."""
    normalized = extracted.strip().lower().strip(" \t\r\n.,;:!?")
    return normalized == "true"


async def async_evaluate_state(
    state: str,
    goal: str,
    max_retries: int = 3,
    domain_description: str = None,
    **kwargs,
) -> float:
    """Async version of evaluate_state for concurrent execution (used in ToT)."""
    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }
    domain_desc = domain_description or DOMAIN_DESCRIPTION_BLOCKSWORLD
    prompt = (
        f"{domain_desc}\n"
        f"[Current State]{state}\n"
        f"[Goal State]{goal}\n"
        f"{STATE_SCORE_PROMPT}\n"
    )
    for attempt in range(max_retries):
        response = await AsyncRequestSender(prompt, **LLM_params)
        extracted = ContentExtraction(response)
        if not extracted.startswith("Extraction failed"):
            try:
                score = float(extracted.strip())
                return max(0.0, min(10.0, score))
            except ValueError:
                print(
                    f"[Warning] Could not parse score '{extracted}' (attempt {attempt + 1}/{max_retries})"
                )
                continue
        print(f"[Warning] Async state evaluation failed (attempt {attempt + 1}/{max_retries})")
    raise RuntimeError(
        f"asynchronous state evaluation failed after {max_retries} attempts"
    )


async def async_generate_available_actions(
    state: str,
    max_retries: int = 3,
    last_action: str = None,
    domain_description: str = None,
    **kwargs,
) -> List[str]:
    """Async version of generate_available_actions for concurrent execution."""
    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }
    last_action_str = last_action if last_action else "None"
    domain_desc = domain_description or DOMAIN_DESCRIPTION_BLOCKSWORLD
    prompt = (
        f"{domain_desc}\n"
        f"[Current State]{state}\n"
        f"[Last executed action]{last_action_str}\n"
        f"{GENERATE_AVAILABLE_ACTIONS_PROMPT}\n"
    )
    for attempt in range(max_retries):
        response = await AsyncRequestSender(prompt, **LLM_params)
        if not response.startswith("Request Failed"):
            extracted = ContentExtraction(response)
            if extracted.startswith("Extraction failed"):
                print(
                    f"[Warning] Async action extraction failed (attempt {attempt + 1}/{max_retries})"
                )
                continue
            actions = [a.strip() for a in extracted.split(".") if a.strip()]
            if not actions:
                print(
                    f"[Warning] No actions from async generation (attempt {attempt + 1}/{max_retries})"
                )
                continue
            return actions
        print(f"[Warning] Async generate actions failed (attempt {attempt + 1}/{max_retries})")
    return []


def take_action(
    state: str,
    action: str,
    max_retries: int = 3,
    domain_description: str = None,
    **kwargs,
) -> str:
    """
    Execute an action and return the new state.

    Parameters
    ----------
    state : str
        Current state description
    action : str
        Action to execute
    max_retries : int
        Maximum number of retry attempts
    **kwargs : dict
        LLM parameters (model, api_key, base_url)

    Returns
    -------
    str
        New state after executing the action
    """
    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }

    domain_desc = domain_description or DOMAIN_DESCRIPTION_BLOCKSWORLD
    prompt = (
        f"{domain_desc}\n"
        f"[Current State]{state}\n"
        f"[Action]{action}\n"
        f"{PROGRESSION_PROMPT}\n"
    )

    for attempt in range(max_retries):
        response = RequestSender(prompt, **LLM_params)
        extracted = ContentExtraction(response)

        # Check if extraction was successful
        if not extracted.startswith("Extraction failed"):
            return extracted

        print(
            f"[Warning] State progression extraction failed (attempt {attempt + 1}/{max_retries}): {extracted}"
        )

    # All retries failed
    print(f"[Error] State progression failed after {max_retries} attempts")
    return state  # Return original state if all retries fail


def evaluate_state(
    state: str,
    goal: str,
    max_retries: int = 3,
    domain_description: str = None,
    **kwargs,
) -> float:
    """
    USED IN TOT.
    Evaluate how promising a state is for reaching the goal.
    Returns a score between 0 and 10, where higher is better.

    Parameters
    ----------
    state : str
        Current state description
    goal : str
        Goal state description
    max_retries : int
        Maximum number of retry attempts
    **kwargs : dict
        LLM parameters (model, api_key, base_url)

    Returns
    -------
    float
        Score between 0 and 10 indicating how close the state is to the goal
    """
    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }

    domain_desc = domain_description or DOMAIN_DESCRIPTION_BLOCKSWORLD
    prompt = (
        f"{domain_desc}\n"
        f"[Current State]{state}\n"
        f"[Goal State]{goal}\n"
        f"{STATE_SCORE_PROMPT}\n"
    )

    for attempt in range(max_retries):
        response = RequestSender(prompt, **LLM_params)
        extracted = ContentExtraction(response)

        # Check if extraction was successful
        if not extracted.startswith("Extraction failed"):
            try:
                score = float(extracted.strip())
                # Ensure score is in valid range
                return max(0.0, min(10.0, score))
            except ValueError:
                print(
                    f"[Warning] Could not parse score '{extracted}' (attempt {attempt + 1}/{max_retries})"
                )
                continue

        print(
            f"[Warning] State evaluation extraction failed (attempt {attempt + 1}/{max_retries}): {extracted}"
        )

    raise RuntimeError(f"state evaluation failed after {max_retries} attempts")


def tot_search(
    initial_state: str,
    goal: str,
    max_depth: int = 10,
    beam_width: int = 3,
    thoughts_per_state: int = 5,
    max_nodes: int = 300,
    save_tree: bool = True,
    output_path: str = None,
    save_incremental: bool = True,
    domain: str = "blocksworld",
    **kwargs,
) -> Tuple[Optional[List[str]], Dict[int, Node], int]:
    """
    Perform Tree of Thought (ToT) search using beam search strategy.

    ToT generates multiple possible next actions (thoughts) for each state,
    evaluates them, and keeps only the top-k most promising ones.

    Parameters
    ----------
    initial_state : str
        Initial state description
    goal : str
        Goal state description
    max_depth : int
        Maximum search depth
    beam_width : int
        Number of best nodes to keep at each depth level (beam search width)
    thoughts_per_state : int
        Number of candidate actions to generate and evaluate per state
    max_nodes : int
        Maximum total number of nodes to create
    save_tree : bool
        Whether to save the search tree to a JSON file
    output_path : str, optional
        Output file path for the search tree
    domain : str
        Planning domain name (e.g., "blocksworld", "logistics", "depots")
    **kwargs : dict
        LLM parameters (model, api_key, base_url)

    Returns
    -------
    Tuple[Optional[List[str]], Dict[int, Node], int]
        Tuple of (list of actions to reach goal or None, dictionary of all nodes,
        number of nodes actually explored/expanded)
    """
    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }

    domain_description = _get_domain_description(domain)
    print(f"[INFO] Using domain: {domain}")

    # Initialize the root node
    root = Node(
        node_id=0,
        parent_id=None,
        state=initial_state,
        query=goal,
        g_cost=0.0,
        h_cost=0.0,
        history_actions=[],
        state_score=5.0,
    )

    # Tracking structures
    nodes_dict = {0: root}
    next_node_id = 1
    current_beam = [root]  # Nodes at current depth level
    nodes_explored_count = 0  # Count of nodes actually expanded/checked

    print(f"Starting Tree of Thought (ToT) search...")
    print(f"Initial state: {initial_state}")
    print(f"Goal: {goal}")
    print(
        f"Parameters: max_depth={max_depth}, beam_width={beam_width}, thoughts_per_state={thoughts_per_state}, max_nodes={max_nodes}\n"
    )

    # Helper function to save tree before returning
    def _save_and_return(
        result: Optional[List[str]], explored_count: int, found_goal: bool = False
    ):
        if save_tree:
            metadata = {
                "initial_state": initial_state,
                "goal": goal,
                "max_depth": max_depth,
                "beam_width": beam_width,
                "thoughts_per_state": thoughts_per_state,
                "max_nodes": max_nodes,
                "domain": domain,
                "solution_found": found_goal,
                "solution": result,
                "nodes_explored": explored_count,
                "nodes_generated": len(nodes_dict),
            }
            save_search_tree(
                nodes_dict, output_path, search_type="tot", metadata=metadata
            )
        return result, nodes_dict, explored_count

    def _save_incremental():
        if save_incremental and save_tree:
            metadata = {
                "initial_state": initial_state,
                "goal": goal,
                "max_depth": max_depth,
                "beam_width": beam_width,
                "thoughts_per_state": thoughts_per_state,
                "max_nodes": max_nodes,
                "domain": domain,
                "solution_found": False,
                "solution": None,
                "nodes_explored": nodes_explored_count,
                "nodes_generated": len(nodes_dict),
            }
            save_search_tree(
                nodes_dict, output_path, search_type="tot", metadata=metadata
            )

    # Beam search through depth levels
    for depth in range(max_depth):
        print(f"\n{'=' * 60}")
        print(f"[DEPTH {depth}] Processing {len(current_beam)} nodes in beam")
        print(f"Total nodes so far: {len(nodes_dict)}/{max_nodes}")

        if not current_beam:
            print("Beam is empty, search terminated")
            break

        # --- PHASE 1: Concurrent goal checks for all beam nodes ---
        async def _check_goals(nodes):
            tasks = [
                async_check_goal_state(
                    node.state, goal, domain_description=domain_description, **LLM_params
                )
                for node in nodes
            ]
            return await asyncio.gather(*tasks)

        goal_results = asyncio.run(_check_goals(current_beam))
        nodes_explored_count += len(current_beam)
        for node, is_goal in zip(current_beam, goal_results):
            if is_goal:
                print(f"\n{'=' * 60}")
                print("GOAL FOUND!")
                print(f"Solution path: {node.history_actions}")
                print(f"Number of actions: {len(node.history_actions)}")
                print(f"Total nodes created: {len(nodes_dict)}")
                print(f"Total nodes explored: {nodes_explored_count}")
                node.flag = True
                return _save_and_return(
                    node.history_actions, nodes_explored_count, found_goal=True
                )

        if len(nodes_dict) >= max_nodes:
            print(f"\n[LIMIT REACHED] Maximum total nodes ({max_nodes}) reached")
            break

        # --- PHASE 2: Concurrent action generation for all beam nodes ---
        async def _gen_all(nodes):
            tasks = [
                async_generate_available_actions(
                    node.state,
                    domain_description=domain_description,
                    **LLM_params,
                )
                for node in nodes
            ]
            return await asyncio.gather(*tasks)

        all_actions_per_node = asyncio.run(_gen_all(current_beam))
        for node, actions in zip(current_beam, all_actions_per_node):
            print(f"\nNode {node.node_id}: {len(actions)} potential actions generated")

        # --- PHASE 3: Concurrent executability checks ---
        # Collect (beam_idx, action) candidate pairs, take up to thoughts_per_state*2 per node
        candidate_pairs: List[tuple] = []
        for beam_idx, (node, actions) in enumerate(zip(current_beam, all_actions_per_node)):
            for action in actions[: thoughts_per_state * 2]:
                candidate_pairs.append((beam_idx, action))

        async def _check_exec_all(pairs):
            tasks = [
                async_check_action_executability(
                    current_beam[bi].state,
                    act,
                    domain_description=domain_description,
                    **LLM_params,
                )
                for bi, act in pairs
            ]
            return await asyncio.gather(*tasks)

        exec_results = asyncio.run(_check_exec_all(candidate_pairs))

        # Collect executable actions per beam node, limited to thoughts_per_state
        executable_per_node: dict = {bi: [] for bi in range(len(current_beam))}
        for (bi, action), ok in zip(candidate_pairs, exec_results):
            if ok and len(executable_per_node[bi]) < thoughts_per_state:
                executable_per_node[bi].append(action)
                print(f"  ✓ [Node {current_beam[bi].node_id}] {action}")
            else:
                print(f"  ✗ [Node {current_beam[bi].node_id}] {action}")

        # --- PHASE 4: Concurrent take_action for all expansion pairs ---
        expansion_pairs: List[tuple] = []
        for bi in range(len(current_beam)):
            for action in executable_per_node[bi]:
                if next_node_id + len(expansion_pairs) >= max_nodes:
                    break
                expansion_pairs.append((bi, action))
            if next_node_id + len(expansion_pairs) >= max_nodes:
                break

        async def _take_all(pairs):
            tasks = [
                async_take_action(
                    current_beam[bi].state,
                    act,
                    domain_description=domain_description,
                    **LLM_params,
                )
                for bi, act in pairs
            ]
            return await asyncio.gather(*tasks)

        new_states = asyncio.run(_take_all(expansion_pairs))
        print(f"  Completed {len(new_states)} concurrent state transitions")

        # --- PHASE 5: Concurrent evaluate_state for all new states ---
        async def _eval_all(states):
            tasks = [
                async_evaluate_state(
                    state, goal, domain_description=domain_description, **LLM_params
                )
                for state in states
            ]
            return await asyncio.gather(*tasks)

        state_scores = asyncio.run(_eval_all(new_states))
        print(f"  Completed {len(state_scores)} concurrent state evaluations")

        # --- Create nodes ---
        next_level_candidates = []
        for (bi, action), new_state, state_score in zip(
            expansion_pairs, new_states, state_scores
        ):
            if next_node_id >= max_nodes:
                print(f"\n[LIMIT REACHED] Would exceed max_nodes ({max_nodes})")
                break
            current_node = current_beam[bi]
            print(f"  Node {current_node.node_id} → '{action}' score={state_score:.2f}/10.0")
            new_node = Node(
                node_id=next_node_id,
                parent_id=current_node.node_id,
                state=new_state,
                query=goal,
                g_cost=current_node.g_cost + 1,
                h_cost=10.0 - state_score,
                action_taken=[action],
                history_actions=current_node.history_actions + [action],
                state_score=state_score,
            )
            nodes_dict[next_node_id] = new_node
            current_node.children_id.append(next_node_id)
            next_level_candidates.append(new_node)
            print(f"  Created Node {next_node_id}")
            next_node_id += 1

        # Save incremental snapshot after each depth layer
        _save_incremental()

        # Select top beam_width candidates for next level based on state_score
        if next_level_candidates:
            next_level_candidates.sort(key=lambda n: n.state_score, reverse=True)
            current_beam = next_level_candidates[:beam_width]
            print(f"\n{'=' * 60}")
            print(
                f"[BEAM SELECTION] Selected top {len(current_beam)} nodes for next level:"
            )
            for i, node in enumerate(current_beam):
                print(
                    f"  {i + 1}. Node {node.node_id} - Score: {node.state_score:.2f} - Action: {node.action_taken}"
                )
        else:
            current_beam = []

    print(f"\n{'=' * 60}")
    print("ToT search completed - No solution found")
    print(f"Total nodes generated: {len(nodes_dict)}")
    print(f"Total nodes explored: {nodes_explored_count}")
    print(f"Maximum depth reached: {depth}")
    return _save_and_return(None, nodes_explored_count)


def bfs_search(
    initial_state: str,
    goal: str,
    max_depth: int = 10,
    max_width: int = 8,
    max_nodes: int = 300,
    save_tree: bool = True,
    output_path: str = None,
    save_incremental: bool = True,
    domain: str = "blocksworld",
    **kwargs,
) -> Tuple[Optional[List[str]], Dict[int, Node], int]:
    """
    Perform BFS search to find a plan from initial state to goal state.

    Parameters
    ----------
    initial_state : str
        Initial state description
    goal : str
        Goal state description
    max_depth : int
        Maximum search depth (maximum number of layers)
    max_width : int
        Maximum number of nodes to expand per layer (breadth limit)
    max_nodes : int
        Maximum total number of nodes to create
    save_tree : bool
        Whether to save the search tree to a JSON file (default: True)
    output_path : str, optional
        Output file path for the search tree. If None, generates a timestamped filename
    domain : str
        Planning domain name (e.g., "blocksworld", "logistics", "depots")
    **kwargs : dict
        LLM parameters (model, api_key, base_url)

    Returns
    -------
    Tuple[Optional[List[str]], Dict[int, Node], int]
        Tuple of (list of actions to reach goal or None, dictionary of all nodes,
        number of nodes actually explored/expanded)
    """
    LLM_params = {
        "model": kwargs.get("model", "gpt-4o"),
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }

    domain_description = _get_domain_description(domain)
    print(f"[INFO] Using domain: {domain}")

    # Initialize the root node
    root = Node(
        node_id=0,
        parent_id=None,
        state=initial_state,
        query=goal,
        g_cost=0.0,
        h_cost=0.0,
        history_actions=[],
    )

    # BFS queue and tracking structures
    queue = deque([root])
    nodes_dict = {0: root}  # Store all nodes by ID
    next_node_id = 1
    nodes_explored_count = 0  # Count of nodes actually popped and processed

    print(f"Starting BFS search...")
    print(f"Initial state: {initial_state}")
    print(f"Goal: {goal}")
    print(
        f"Limits: max_depth={max_depth}, max_width={max_width}, max_nodes={max_nodes}\n"
    )

    current_depth = 0
    nodes_at_current_depth = 0

    # Helper function to save tree before returning
    def _save_and_return(
        result: Optional[List[str]], explored_count: int, found_goal: bool = False
    ):
        if save_tree:
            metadata = {
                "initial_state": initial_state,
                "goal": goal,
                "max_depth": max_depth,
                "max_width": max_width,
                "max_nodes": max_nodes,
                "domain": domain,
                "solution_found": found_goal,
                "solution": result,
                "nodes_explored": explored_count,
                "nodes_generated": len(nodes_dict),
            }
            save_search_tree(
                nodes_dict, output_path, search_type="bfs", metadata=metadata
            )
        return result, nodes_dict, explored_count

    def _save_incremental():
        if save_incremental and save_tree:
            metadata = {
                "initial_state": initial_state,
                "goal": goal,
                "max_depth": max_depth,
                "max_width": max_width,
                "max_nodes": max_nodes,
                "domain": domain,
                "solution_found": False,
                "solution": None,
                "nodes_explored": nodes_explored_count,
                "nodes_generated": len(nodes_dict),
            }
            save_search_tree(
                nodes_dict, output_path, search_type="bfs", metadata=metadata
            )

    while queue:
        # Check total nodes limit
        if len(nodes_dict) >= max_nodes:
            print(f"\n{'=' * 60}")
            print(f"[LIMIT REACHED] Maximum total nodes ({max_nodes}) reached")
            print(f"Total nodes created: {len(nodes_dict)}")
            return _save_and_return(None, nodes_explored_count)

        # Get the next node from queue
        current_node = queue.popleft()

        # Track depth changes
        if current_node.g_cost > current_depth:
            current_depth = int(current_node.g_cost)
            nodes_at_current_depth = 0
            print(f"\n{'=' * 60}")
            print(f"[DEPTH {current_depth}] Starting new layer")

        # Check width limit for current depth
        if nodes_at_current_depth >= max_width:
            print(
                f"[LIMIT REACHED] Maximum width ({max_width}) reached at depth {current_depth}, skipping remaining nodes at this depth"
            )
            # Skip all remaining nodes at this depth
            while queue and queue[0].g_cost == current_depth:
                queue.popleft()
            continue

        nodes_at_current_depth += 1
        nodes_explored_count += 1  # This node is actually being processed

        print(f"\n{'=' * 60}")
        print(
            f"Exploring Node {current_node.node_id} (depth={current_node.g_cost}, {nodes_at_current_depth}/{max_width} at this depth)"
        )
        print(f"State: {current_node.state[:100]}...")  # Print first 100 chars
        print(f"History: {current_node.history_actions}")
        print(f"Total nodes so far: {len(nodes_dict)}/{max_nodes}")

        # Step 1: Check if this is the goal state
        if check_goal_state(
            current_node.state, goal, domain_description=domain_description, **LLM_params
        ):
            print(f"\n{'=' * 60}")
            print("GOAL FOUND!")
            print(f"Solution path: {current_node.history_actions}")
            print(f"Number of actions: {len(current_node.history_actions)}")
            print(f"Total nodes created: {len(nodes_dict)}")
            print(f"Total nodes explored: {nodes_explored_count}")
            current_node.flag = True
            return _save_and_return(
                current_node.history_actions, nodes_explored_count, found_goal=True
            )

        # Check depth limit
        if current_node.g_cost >= max_depth:
            print(f"[LIMIT REACHED] Max depth {max_depth} reached, skipping expansion")
            continue

        # Step 2: Generate all available actions
        print(f"\nGenerating available actions...")
        available_actions = generate_available_actions(
            current_node.state, domain_description=domain_description, **LLM_params
        )
        current_node.generated_actions = available_actions
        print(f"Found {len(available_actions)} potential actions: {available_actions}")

        # Step 3: Concurrent executability checks
        async def _check_exec_bfs(actions):
            tasks = [
                async_check_action_executability(
                    current_node.state,
                    a,
                    domain_description=domain_description,
                    **LLM_params,
                )
                for a in actions
            ]
            return await asyncio.gather(*tasks)

        exec_results = asyncio.run(_check_exec_bfs(available_actions))

        executable_actions = []
        for action, ok in zip(available_actions, exec_results):
            if ok:
                executable_actions.append(action)
                print(f"  ✓ {action}")
            else:
                print(f"  ✗ {action} (not executable)")

        print(f"\n{len(executable_actions)} executable actions: {executable_actions}")

        # Step 4: Concurrent take_action (limited to remaining node budget)
        remaining = max_nodes - next_node_id
        candidate_actions = executable_actions[:remaining]

        async def _take_bfs(actions):
            tasks = [
                async_take_action(
                    current_node.state,
                    a,
                    domain_description=domain_description,
                    **LLM_params,
                )
                for a in actions
            ]
            return await asyncio.gather(*tasks)

        new_states = asyncio.run(_take_bfs(candidate_actions))
        print(f"  Completed {len(new_states)} concurrent state transitions")

        # Step 5: Create nodes
        for action, new_state in zip(candidate_actions, new_states):
            new_node = Node(
                node_id=next_node_id,
                parent_id=current_node.node_id,
                state=new_state,
                query=goal,
                g_cost=current_node.g_cost + 1,
                h_cost=0.0,
                action_taken=[action],
                history_actions=current_node.history_actions + [action],
            )
            nodes_dict[next_node_id] = new_node
            current_node.children_id.append(next_node_id)
            queue.append(new_node)
            print(f"  Created Node {next_node_id}: {new_state[:80]}...")
            next_node_id += 1

        # Save incremental snapshot after expanding each node
        _save_incremental()

    print(f"\n{'=' * 60}")
    print("BFS search completed - No solution found")
    print(f"Total nodes generated: {len(nodes_dict)}")
    print(f"Total nodes explored: {nodes_explored_count}")
    print(f"Maximum depth reached: {current_depth}")
    return _save_and_return(None, nodes_explored_count)


def clear_lora_model_cache():
    """
    Clear the cached LoRA model to free up memory.
    Call this function if you want to reload the model or free up GPU memory.
    """
    if HEURISTIC_SYSTEM_AVAILABLE:
        new_clear_lora_model_cache()


def compute_heuristic_with_lora(
    state: str, goal: str, max_retries: int = 3, **kwargs
) -> float:
    """
    Use local qwen3-blocksworld-lora model to compute heuristic value (h) for A* search.
    The model receives natural language state and goal, returns estimated distance to goal.

    Delegates to the strict shared heuristic system. Import or model failures are
    surfaced to the caller; fallback is available only through its explicit flag.

    Parameters
    ----------
    state : str
        Current state description in natural language
    goal : str
        Goal state description in natural language
    max_retries : int
        Maximum number of retry attempts
    **kwargs : dict
        Additional parameters (model_path, base_model_path, etc.)

    Returns
    -------
    float
        Heuristic value (estimated cost to reach goal from current state)
    """
    if not HEURISTIC_SYSTEM_AVAILABLE:
        raise RuntimeError(
            f"strict heuristic system could not be imported: {HEURISTIC_IMPORT_ERROR}"
        )
    return new_compute_heuristic_with_lora(state, goal, max_retries, **kwargs)


def astar_search(
    initial_state: str,
    goal: str,
    max_depth: int = 20,
    max_nodes: int = 500,
    save_tree: bool = True,
    output_path: str = None,
    save_incremental: bool = True,
    domain: str = "blocksworld",
    **kwargs,
) -> Tuple[Optional[List[str]], Dict[int, Node], int]:
    """
    Perform learned-heuristic best-first search from initial state to goal state.

    The learned h_add approximation is not guaranteed to be admissible or consistent,
    so this routine should be evaluated as satisficing search. Closed states are
    reopened when a strictly better g-cost is discovered.

    Parameters
    ----------
    initial_state : str
        Initial state description
    goal : str
        Goal state description
    max_depth : int
        Maximum search depth
    max_nodes : int
        Maximum total number of nodes to create
    save_tree : bool
        Whether to save the search tree to a JSON file
    output_path : str, optional
        Output file path for the search tree
    save_incremental : bool
        Whether to save the tree incrementally after each node creation
    domain : str
        Domain name ("blocksworld", "logistics", "depots") used to select the
        correct PDDL description for action/state prompts.
    **kwargs : dict
        LLM parameters (model, api_key, base_url)
        - heuristic_model: model for computing heuristic (default: qwen3-blocksworld-lora)
        - action_model: model for action generation and state progression
        - check_executability: whether to verify generated actions with LLM before progression

    Returns
    -------
    Tuple[Optional[List[str]], Dict[int, Node], int]
        Tuple of (list of actions to reach goal or None, dictionary of all nodes, number of explored nodes)
    """
    # Resolve domain description once for all async helper calls
    domain_desc = _get_domain_description(domain)
    print(f"DEBUG: astar_search using domain='{domain}', length of domain_desc={len(domain_desc)}")

    # Separate models for different tasks
    heuristic_model = kwargs.get("heuristic_model", "qwen3-lora")
    action_model = kwargs.get("action_model", kwargs.get("model", "gpt-4o"))
    check_executability = kwargs.get("check_executability", True)

    LLM_params = {
        "model": action_model,
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
    }

    heuristic_params = {
        "model": heuristic_model,
        "api_key": kwargs.get("api_key", "null"),
        "base_url": kwargs.get("base_url", "null"),
        "lora_model_path": kwargs.get("lora_model_path", "./fine-tuning/qwen3-lora"),
        "base_model_path": kwargs.get(
            "base_model_path", "./fine-tuning/qwen3-1.7b-local"
        ),
        "allow_heuristic_fallback": kwargs.get("allow_heuristic_fallback", False),
        "heuristic_precision": kwargs.get("heuristic_precision", "auto"),
        "heuristic_max_seq_length": kwargs.get("heuristic_max_seq_length", 2048),
        "h_type": kwargs.get("h_type", "hadd"),
    }

    # Compute initial heuristic
    initial_h = compute_heuristic_with_lora(initial_state, goal, **heuristic_params)

    # Initialize the root node
    root = Node(
        node_id=0,
        parent_id=None,
        state=initial_state,
        query=goal,
        g_cost=0.0,
        h_cost=initial_h,
        history_actions=[],
    )

    # A* uses priority queue (min-heap) based on f_cost
    open_set = []  # Priority queue: (f_cost, node_id)
    heapq.heappush(open_set, (root.f_cost, 0))

    nodes_dict = {0: root}
    closed_best_g: Dict[str, float] = {}
    state_to_best_g = {}  # Map state -> best g_cost seen so far
    next_node_id = 1
    expansion_count = 0

    # Initialize state_to_best_g with normalized initial state
    state_to_best_g[normalize_state_for_search(initial_state)] = 0.0

    print(f"Starting A* search...")
    print(f"Initial state: {initial_state}")
    print(f"Goal: {goal}")
    print(f"Initial heuristic: {initial_h}")
    print(f"Limits: max_depth={max_depth}, max_nodes={max_nodes}")
    print(f"Heuristic model: {heuristic_model}")
    print(f"Action model: {action_model}\n")

    # Helper function to save tree incrementally
    def _save_incremental():
        if save_incremental and save_tree:
            metadata = {
                "initial_state": initial_state,
                "goal": goal,
                "max_depth": max_depth,
                "max_nodes": max_nodes,
                "domain": domain,
                "heuristic_model": heuristic_model,
                "action_model": action_model,
                "check_executability": check_executability,
                "solution_found": False,
                "solution": None,
                "nodes_explored": expansion_count,
                "unique_states_expanded": len(closed_best_g),
                "nodes_generated": len(nodes_dict),
                "search_in_progress": True,
            }
            save_search_tree(
                nodes_dict, output_path, search_type="astar", metadata=metadata
            )

    # Helper function to save tree before returning
    def _save_and_return(result: Optional[List[str]], found_goal: bool = False):
        if save_tree:
            metadata = {
                "initial_state": initial_state,
                "goal": goal,
                "max_depth": max_depth,
                "max_nodes": max_nodes,
                "domain": domain,
                "heuristic_model": heuristic_model,
                "action_model": action_model,
                "check_executability": check_executability,
                "solution_found": found_goal,
                "solution": result,
                "nodes_explored": expansion_count,
                "unique_states_expanded": len(closed_best_g),
                "nodes_generated": len(nodes_dict),
                "search_in_progress": False,
            }
            save_search_tree(
                nodes_dict, output_path, search_type="astar", metadata=metadata
            )
        return result, nodes_dict, expansion_count
    event_loop = asyncio.new_event_loop()
    try:
        asyncio.set_event_loop(event_loop)

        while open_set:
            # Get node with lowest f_cost
            _, current_id = heapq.heappop(open_set)
            current_node = nodes_dict[current_id]

            normalized_state = normalize_state_for_search(current_node.state)
            if current_node.g_cost > state_to_best_g.get(normalized_state, float("inf")):
                continue
            if not record_better_expansion(
                closed_best_g, normalized_state, current_node.g_cost
            ):
                continue
            expansion_count += 1

            print(f"\n{'=' * 60}")
            print(f"Expanding Node {current_node.node_id}")
            print(
                f"g={current_node.g_cost:.1f}, h={current_node.h_cost:.1f}, f={current_node.f_cost:.1f}"
            )
            print(f"State: {current_node.state[:100]}...")
            print(f"History: {current_node.history_actions}")
            print(
                f"Open set size: {len(open_set)}, "
                f"Unique expanded states: {len(closed_best_g)}"
            )

            # Check if this is the goal state
            if check_goal_state(
                current_node.state, goal, domain_description=domain_desc, **LLM_params
            ):
                print(f"\n{'=' * 60}")
                print("GOAL FOUND!")
                print(f"Solution path: {current_node.history_actions}")
                print(f"Number of actions: {len(current_node.history_actions)}")
                print(f"Total nodes created: {len(nodes_dict)}")
                print(f"Node expansions: {expansion_count}")
                current_node.flag = True
                return _save_and_return(
                    current_node.history_actions, found_goal=True
                )

            # Check depth limit
            if current_node.g_cost >= max_depth:
                print(
                    f"[LIMIT REACHED] Max depth {max_depth} reached, skipping expansion"
                )
                current_node.status = False
                current_node.pruned_reason = "MAX_DEPTH_REACHED"
                continue

            # Still pop and goal-check queued nodes after reaching the generation cap.
            if len(nodes_dict) >= max_nodes:
                print(f"[LIMIT REACHED] Maximum total nodes ({max_nodes}) reached")
                current_node.status = False
                current_node.pruned_reason = "MAX_NODES_REACHED"
                continue

            # Generate available actions
            print(f"\nGenerating available actions...")
            available_actions = generate_available_actions(
                current_node.state, domain_description=domain_desc, **LLM_params
            )
            current_node.generated_actions = available_actions
            print(f"Found {len(available_actions)} potential actions")

            if check_executability:
                # --- 3.1  Concurrent executability checks ---
                async def _check_all(actions):
                    tasks = [
                        async_check_action_executability(
                            current_node.state,
                            a,
                            domain_description=domain_desc,
                            **LLM_params,
                        )
                        for a in actions
                    ]
                    return await asyncio.gather(*tasks)

                executability_results = event_loop.run_until_complete(
                    _check_all(available_actions)
                )
                executable_actions = []
                for action, ok in zip(available_actions, executability_results):
                    if ok:
                        executable_actions.append(action)
                        print(f"  ✓ {action}")
                    else:
                        print(f"  ✗ {action} (not executable)")

                print(f"\n{len(executable_actions)} executable actions")
            else:
                executable_actions = available_actions
                print("Skipping executability checks; trusting generated actions")

            # Enforce node limit before expansion
            remaining_slots = max_nodes - len(nodes_dict)
            if remaining_slots <= 0:
                print(f"\n[LIMIT REACHED] Would exceed max_nodes ({max_nodes})")
                current_node.available_actions = executable_actions
                continue
            candidate_actions = executable_actions[:remaining_slots]

            # Record available_actions in current node
            current_node.available_actions = executable_actions

            # --- 3.1  Concurrent state transitions ---
            async def _take_all(actions):
                tasks = [
                    async_take_action(
                        current_node.state,
                        a,
                        domain_description=domain_desc,
                        **LLM_params,
                    )
                    for a in actions
                ]
                return await asyncio.gather(*tasks)

            new_states_raw = event_loop.run_until_complete(_take_all(candidate_actions))
            print(f"  Completed {len(new_states_raw)} concurrent state transitions")

            # Filter states already visited with a better g_cost
            new_g_cost = current_node.g_cost + 1
            valid_actions: List[str] = []
            valid_states: List[str] = []
            for action, new_state in zip(candidate_actions, new_states_raw):
                normalized_new_state = normalize_state_for_search(new_state)
                if normalized_new_state in state_to_best_g:
                    if state_to_best_g[normalized_new_state] <= new_g_cost:
                        print(
                            f"  ✗ State already visited with g={state_to_best_g[normalized_new_state]:.1f}, skipping"
                        )
                        continue
                state_to_best_g[normalized_new_state] = new_g_cost
                valid_actions.append(action)
                valid_states.append(new_state)

            # --- 3.2  Batch heuristic computation ---
            if valid_states:
                print(
                    f"  Computing heuristics for {len(valid_states)} states in batch..."
                )
                if HEURISTIC_SYSTEM_AVAILABLE and new_compute_heuristic_batch_with_lora:
                    h_costs = new_compute_heuristic_batch_with_lora(
                        valid_states, goal, **heuristic_params
                    )
                else:
                    h_costs = [
                        compute_heuristic_with_lora(s, goal, **heuristic_params)
                        for s in valid_states
                    ]
                print(f"  Batch heuristics: {h_costs}")
            else:
                h_costs = []

            # Create nodes for all valid new states
            for action, new_state, h_cost in zip(valid_actions, valid_states, h_costs):
                print(f"\nExpanding action: {action}")
                print(f"  New state: {new_state[:150]}...")
                print(f"  ✓ Heuristic: {h_cost}")

                new_node = Node(
                    node_id=next_node_id,
                    parent_id=current_node.node_id,
                    state=new_state,
                    query=goal,
                    g_cost=new_g_cost,
                    h_cost=h_cost,
                    action_taken=[action],
                    history_actions=current_node.history_actions + [action],
                )

                nodes_dict[next_node_id] = new_node
                current_node.children_id.append(next_node_id)
                heapq.heappush(open_set, (new_node.f_cost, next_node_id))

                print(f"  Created Node {next_node_id} (f={new_node.f_cost:.1f})")
                _save_incremental()
                next_node_id += 1

        print(f"\n{'=' * 60}")
        print("A* search completed - No solution found")
        print(f"Total nodes created: {len(nodes_dict)}")
        print(f"Node expansions: {expansion_count}")
        return _save_and_return(None)
    finally:
        try:
            event_loop.close()
        finally:
            asyncio.set_event_loop(None)


def save_search_tree(
    nodes_dict: Dict[int, Node],
    output_path: str = None,
    search_type: str = "bfs",
    metadata: Dict[str, Any] = None,
) -> str:
    """
    Save all nodes from the search tree to a JSON file.

    Parameters
    ----------
    nodes_dict : Dict[int, Node]
        Dictionary mapping node_id to Node objects
    output_path : str, optional
        Output file path. If None, generates a timestamped filename
    search_type : str
        Type of search algorithm used (e.g., "bfs", "astar")
    metadata : Dict[str, Any], optional
        Additional metadata to save (e.g., initial_state, goal, search params)

    Returns
    -------
    str
        Path to the saved JSON file
    """
    # Generate default output path if not provided
    if output_path is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_dir = "search_results"
        os.makedirs(output_dir, exist_ok=True)
        output_path = os.path.join(output_dir, f"{search_type}_tree_{timestamp}.json")

    # Convert all nodes to dictionaries
    nodes_list = [node.to_dict() for node in nodes_dict.values()]

    # Build the output structure
    output_data = {
        "search_type": search_type,
        "total_nodes": len(nodes_dict),
        "timestamp": datetime.now().isoformat(),
        "metadata": metadata or {},
        "nodes": nodes_list,
    }

    # Save to JSON file
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(output_data, f, indent=2, ensure_ascii=False)

    print(f"[INFO] Search tree saved to: {output_path}")
    return output_path