#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : %FFILE%
# Author : %USER%
# Email  : %MAIL%
# Date   : %MONTH%/%DAY%/%YEAR%
#
# This file is part of Project Concepts.
# Distributed under terms of the %LICENSE% license.


concepts/.vim-template:.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 06/23/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 11/01/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Generic implementation for a few basic algorithms, such as A* search."""

concepts/algorithm/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : configuration_space.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/01/2019
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
from math import ceil

__all__ = [
    'ConfigurationSpace',
    'ProxyConfigurationSpace',
    'BoxConfigurationSpace',
    'ConstrainedConfigurationSpace',
    'BoxConstrainedConfigurationSpace',
    'ProblemSpace',
    'CollisionFreeProblemSpace'
]


class ConfigurationSpace(object):
    def sample(self):
        raise NotImplementedError()

    def difference(self, config1, config2):
        raise NotImplementedError()

    def distance(self, config1, config2):
        raise NotImplementedError()

    def validate_config(self, config):
        raise NotImplementedError()

    def gen_path(self, config1, config2):
        raise NotImplementedError()


class ProxyConfigurationSpace(ConfigurationSpace):
    def __init__(self, proxy):
        self.proxy = proxy

    def __getattr__(self, item):
        return getattr(self.proxy, item)

    def sample(self):
        return self.proxy.sample()

    def difference(self, config1, config2):
        return self.proxy.difference(config1, config2)

    def distance(self, config1, config2):
        return self.proxy.distance(config1, config2)

    def validate_config(self, config):
        return self.proxy.validate_config(config)

    def gen_path(self, config1, config2):
        raise NotImplementedError()


class BoxConfigurationSpace(ConfigurationSpace):
    def __init__(self, cspace_ranges, cspace_max_stepdiff):
        super().__init__()
        self.cspace_ranges = cspace_ranges

        if isinstance(cspace_max_stepdiff, (int, float)):
            cspace_max_stepdiff = (cspace_max_stepdiff,) * len(self.cspace_ranges)
        self.cspace_max_stepdiff = np.array(cspace_max_stepdiff, dtype='float32')

    def sample(self):
        return tuple(r.sample() for r in self.cspace_ranges)

    def difference(self, one: tuple, two: tuple):
        a = [self.cspace_ranges[i].difference(one[i], two[i]) for i in range(len(self.cspace_ranges))]
        return a

    def distance(self, one: tuple, two: tuple):
        a = [abs(self.cspace_ranges[i].difference(one[i], two[i])) for i in range(len(self.cspace_ranges))]
        return sum(a)

    def validate_config(self, config):
        return len(config) == len(self.cspace_ranges) and all(
            [self.cspace_ranges[i].contains(config[i]) for i in range(len(self.cspace_ranges))]
        )

    def gen_path_l1(self, config1: tuple, config2: tuple):
        # TODO(Jiayuan Mao @ 2023/03/09): bring this back. In panda planning, sometimes the input configurations are invalid...
        # assert self.validate_config(config1) and self.validate_config(config2), 'Invalid config pair: {} and {}.'.format(config1, config2)
        diffs = [self.cspace_ranges[i].difference(config1[i], config2[i]) for i in range(len(self.cspace_ranges))]
        samples = max([int(ceil(abs(diff) / max_diff)) + 1 for diff, max_diff in zip(diffs, self.cspace_max_stepdiff)])
        samples = max(2, samples)
        linear_interpolation = [diff / (samples - 1.0) for diff in diffs]
        path = [config1]
        for s in range(1, samples - 1):
            sample = tuple(
                [
                    self.cspace_ranges[i].in_range(config1[i] + s * linear_interpolation[i])
                    for i in range(len(self.cspace_ranges))
                ]
            )
            path.append(sample)
        return True, path + [config2]

    def gen_path_l2(self, config1: tuple, config2: tuple):
        diffs = [self.cspace_ranges[i].difference(config1[i], config2[i]) for i in range(len(self.cspace_ranges))]
        diffs_norm = np.linalg.norm(diffs, ord=2)
        nr_steps = int(ceil(diffs_norm / self.cspace_max_stepdiff[0])) + 1
        nr_steps = max(2, nr_steps)

        linear_interpolation = [diff / (nr_steps - 1.0) for diff in diffs]
        path = [config1]
        for s in range(1, nr_steps - 1):
            sample = tuple(
                [
                    self.cspace_ranges[i].in_range(config1[i] + s * linear_interpolation[i])
                    for i in range(len(self.cspace_ranges))
                ]
            )
            path.append(sample)

        # max_step_size = 0
        # for c1, c2 in zip(path, path[1:]):
        #     max_step_size = max(max_step_size, self.distance(c1, c2))
        # print(f'gen_path_l2: max_step_size = {max_step_size}; nr_configs = {len(path)}')

        return True, path + [config2]

    def gen_path(self, config1: tuple, config2: tuple):
        return self.gen_path_l1(config1, config2)


class ConstrainedConfigurationSpace(ProxyConfigurationSpace):
    def __init__(self, cspace, constrain_config_func=None, extra_validate_config_func=None):
        super().__init__(cspace)
        self.constrain_config_func = constrain_config_func
        self.extra_validate_config_func = extra_validate_config_func

    def sample(self):
        return self.constrain_config(self.proxy.sample())

    def validate_config(self, config):
        if self.extra_validate_config_func is not None:
            return self.proxy.validate_config(config) and self.extra_validate_config_func(self, config)
        else:
            return self.proxy.validate_config(config)

    def constrain_config(self, config):
        if self.constrain_config_func is not None:
            return self.constrain_config_func(self, config)

        raise NotImplementedError()

    def gen_path(self, config1: tuple, config2: tuple):
        success, path = self.proxy.gen_path(config1, config2)
        for i, config in enumerate(path):
            path[i] = self.constrain_config(config)
        return success, path


class BoxConstrainedConfigurationSpace(ConstrainedConfigurationSpace):
    def __init__(
        self,
        cspace,
        constrain_config_func=None,
        extra_validate_config_func=None,
        cspace_max_stepdiff_proj=3,
        gen_path_atol=1e-5,
    ):

        assert isinstance(cspace, BoxConfigurationSpace)

        super().__init__(cspace, constrain_config_func, extra_validate_config_func)
        self._cspace_max_stepdiff_proj = cspace_max_stepdiff_proj
        self._gen_path_atol = gen_path_atol

    @property
    def cspace_ranges(self):
        return self.proxy.cspace_ranges

    @property
    def cspace_max_stepdiff(self):
        return self.proxy.cspace_max_stepdiff

    @property
    def cspace_max_stepdiff_proj(self):
        if isinstance(self._cspace_max_stepdiff_proj, (int, float)):
            return tuple(x * self._cspace_max_stepdiff_proj for x in self.cspace_max_stepdiff)
        return self._cspace_max_stepdiff_proj

    def gen_path(self, config1, config2):
        assert self.validate_config(config1) and self.validate_config(config2)

        config = config1
        path = [config]
        success = False

        for i in range(1000):
            config = config + np.fmin(self.cspace_max_stepdiff, np.array(config2) - np.array(config))
            config = self.constrain_config(config)

            diffs = [self.cspace_ranges[i].difference(path[-1][i], config[i]) for i in range(len(self.cspace_ranges))]
            valid = all(abs(diff) < constrain for diff, constrain in zip(diffs, self.cspace_max_stepdiff_proj))
            if valid:
                path.append(config)
                if np.allclose(config, config2):
                    success = True
                    break
                if self.distance(config, config2) >= self.distance(path[-1], config2) - self._gen_path_atol:
                    break
            else:
                break

        return success, path


class ProblemSpace(object):
    def __init__(self, cspace):
        self.cspace = cspace

    def sample(self, nr_trials=None):
        if nr_trials is None:
            nr_trials = self.sample.default_nr_trials
        for i in range(nr_trials):
            sample = self.cspace.sample()
            if self.validate_config(sample):
                return sample
        raise ValueError('Unable to produce a sample.')

    sample.default_nr_trials = 100

    def difference(self, config1, config2):
        return self.cspace.difference(config1, config2)

    def distance(self, config1, config2):
        return self.cspace.distance(config1, config2)

    def validate_config(self, config):
        return self.cspace.validate_config(config)

    def distance_path(self, path):
        distance = 0
        for i in range(len(path) - 1):
            distance += self.cspace.distance(path[i], path[i + 1])
        return distance

    def validate_path(self, config1, config2):
        success, _, _ = self.try_extend_path(config1, config2)
        return success

    def try_extend_path(self, start, end):
        success, path = self.cspace.gen_path(start, end)
        if path is None:
            return False, start, path

        safe_path = [path[0]]
        for i in range(1, len(path)):
            configuration = path[i]
            # jacinle.log_function.print(f'Validate {i}/{len(path)}: {configuration} -> {self.validate_config(configuration)}')
            if not self.validate_config(configuration):
                return False, safe_path[-1], safe_path
            safe_path.append(configuration)
        return success, safe_path[-1], safe_path


class CollisionFreeProblemSpace(ProblemSpace):
    def __init__(self, cspace, collide_fn=None):
        super().__init__(cspace)
        self.collide_fn = collide_fn

    def validate_config(self, configuration):
        return self.cspace.validate_config(configuration) and not self.collide(configuration)

    def collide(self, configuration):
        if self.collide_fn is not None:
            return self.collide_fn(configuration)
        raise NotImplementedError()


concepts/algorithm/configuration_space.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 06/14/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/algorithm/rrt/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : rrt.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/01/2019
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
import numpy.random as npr

from typing import Optional, Tuple, List
from concepts.algorithm.configuration_space import ProblemSpace

__all__ = ['RRTNode', 'RRT', 'smooth_path', 'get_smooth_path', 'optimize_path', 'rrt', 'birrt']

"""
Basic RRT algorithm.

The following algorithms and data structures have a generic interface. The are not designed specifically for a
specific robot.
"""


class RRTNode(object):
    def __init__(self, config, parent=None):
        self.config = config
        self.parent = parent
        self.children = list()

    def add_to_children(self, other):
        self.children.append(other)
        return self

    def attach_to(self, other):
        self.parent = other
        self.parent.add_to_children(self)
        return self

    def backtrace(self, config=True):
        path = list()

        def dfs(x):
            if x.parent is not None:
                dfs(x.parent)
            path.append(x.config if config else x)

        try:
            dfs(self)
            return path
        finally:
            del dfs

    @classmethod
    def from_states(cls, states):
        if isinstance(states, list):
            return [cls(s) for s in states]
        else:
            return cls(states)

    def __repr__(self):
        return f'RRTNode(config={self.config}, parent={self.parent})'


def traverse_rrt_bfs(nodes):
    queue = nodes.copy()
    results = list()

    while len(queue) > 0:
        x = queue[0]
        queue = queue[1:]
        results.append(x)
        for y in x.children:
            queue.append(y)

    return results


class RRT(object):
    """The RRT tree."""

    def __init__(self, pspace, roots):
        if isinstance(roots, RRTNode):
            roots = [roots]
        else:
            assert isinstance(roots, list)

        self.pspace = pspace
        self.roots = roots
        self.size = len(roots)

    def extend(self, parent, child_config):
        child = RRTNode(child_config).attach_to(parent)
        self.size += 1
        return child

    def nearest(self, config, pspace=None):
        if pspace is None:
            pspace = self.pspace

        best_node, best_value = None, np.inf
        for node in traverse_rrt_bfs(self.roots):
            distance = pspace.distance(node.config, config)
            if distance < best_value:
                best_value = distance
                best_node = node

        return best_node


def smooth_path(pspace, path, nr_attemps=100, use_fine_path: bool = False):
    if nr_attemps is None:
        nr_attemps = 100

    if use_fine_path:
        path = get_smooth_path(pspace, path)

    for i in range(nr_attemps):
        if len(path) <= 2:
            break

        # Use the uniform pair sampling method.
        a, b = npr.randint(0, len(path) - 1), npr.randint(0, len(path) - 1)
        if a > b:
            a, b = b, a + 1
        else:
            b = b + 1
        # The original version:
        # a = npr.randint(0, len(path) - 1)
        # b = npr.randint(a + 1, len(path))
        if use_fine_path:
            success, _, subpath = pspace.try_extend_path(path[a], path[b])
            if success:
                path = path[:a + 1] + subpath[1:] + path[b:]
        else:
            if pspace.validate_path(path[a], path[b]):
                path = path[:a + 1] + path[b:]

    return path


def get_smooth_path(pspace, path):
    cpath = [path[0]]
    for config in path[1:]:
        cpath.extend(pspace.cspace.gen_path(cpath[-1], config)[1][1:])
    path = cpath
    return path


def optimize_path_forward(pspace, path):
    spath = [path[0]]
    ptr = 1
    while ptr + 1 < len(path):
        if pspace.validate_path(spath[-1], path[ptr + 1]):
            ptr += 1
        else:
            _, sub_path = pspace.cspace.gen_path(spath[-1], path[ptr])
            spath.append(sub_path[1])
            if len(sub_path) == 2:
                ptr += 1
    spath.extend(pspace.cspace.gen_path(spath[-1], path[ptr])[1][1:])
    return spath


def optimize_path(pspace, path, *args, **kwargs):
    path = get_smooth_path(path)
    while True:
        length = len(path)
        path = optimize_path_forward(pspace, path)
        path.reverse()
        path = optimize_path_forward(pspace, path)
        path.reverse()
        if len(path) >= length:
            break
    return path


def rrt(pspace, start_state, goal_state, nr_iterations=1000, p_sample_goal=0.05, nr_smooth_iterations=None):
    rrt = RRT(pspace, RRTNode.from_states(start_state))

    for i in range(nr_iterations):
        sample_goal = False
        if npr.uniform() < p_sample_goal:
            sample_goal = True
            next_config = goal_state
        else:
            next_config = pspace.sample()

        node = rrt.nearest(next_config)
        success, next_config, _ = pspace.try_extend_path(node.config, next_config)
        if next_config is not None:
            new_node = rrt.extend(node, next_config)
            if sample_goal and success:
                return smooth_path(pspace, new_node.backtrace(), nr_smooth_iterations), rrt

    return None, rrt


def birrt(pspace: ProblemSpace, start_state, goal_state, nr_iterations=1000, nr_smooth_iterations=None, smooth_fine_path=False) -> Tuple[Optional[List[np.ndarray]], Tuple[RRT, RRT]]:
    rrt_start = RRT(pspace, RRTNode.from_states(start_state))
    rrt_goal = RRT(pspace, RRTNode.from_states(goal_state))
    swapped = False

    if True:
        # Try to directly connecting the starting state and the goal state.
        success, next_config, _ = pspace.try_extend_path(start_state, goal_state)
        if success:
            rrt_start.extend(rrt_start.nearest(goal_state), goal_state)
            return [start_state, goal_state], (rrt_start, rrt_goal)

    for i in range(nr_iterations):
        next_config = pspace.sample()

        node_start = rrt_start.nearest(next_config)
        success, next_config, _ = pspace.try_extend_path(node_start.config, next_config)
        if next_config is not None:
            new_node_start = rrt_start.extend(node_start, next_config)

            node_goal = rrt_goal.nearest(next_config)
            success, next_config, _ = pspace.try_extend_path(node_goal.config, next_config)
            if success:
                path1 = new_node_start.backtrace()
                path2 = node_goal.backtrace()
                path = list(path1) + list(reversed(path2))

                # still add the new node to the rrt_goal for better visualization.
                _ = rrt_goal.extend(node_goal, next_config)

                if swapped:
                    path.reverse()
                    rrt_start, rrt_goal = rrt_goal, rrt_start
                return smooth_path(pspace, path, nr_smooth_iterations, use_fine_path=smooth_fine_path), (rrt_start, rrt_goal)
            elif next_config is not None:
                _ = rrt_goal.extend(node_goal, next_config)

        rrt_start, rrt_goal, swapped = rrt_goal, rrt_start, not swapped

    return None, (rrt_start, rrt_goal)


concepts/algorithm/rrt/rrt.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/03/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/algorithm/search/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : heuristic_search.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 11/17/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Generic implementation of A* search."""

from dataclasses import dataclass
from typing import TypeVar, List, Callable, Iterator, Tuple, Set
import heapq as hq

__all__ = ['SearchNode', 'QueueNode', 'run_heuristic_search', 'backtrace_plan']

State = TypeVar('State')
Action = TypeVar('Action')


@dataclass
class SearchNode(object):
    """A node object corresponding to the current search state, containing the state, the parent node, the last action, the cost, and the depth."""

    state: State
    """The current state."""

    parent: 'SearchNode'
    """The parent node."""

    action: Action
    """The action that leads to the current state."""

    cost: float
    """The cost of the path from the root to the current state."""

    g: float
    """The estimated cost-to-go."""


@dataclass
class QueueNode(object):
    """A node object in the queue, containing the priority and the search node."""

    priority: float
    """The priority of the node."""

    node: SearchNode
    """The search node."""

    @property
    def state(self):
        return self.node.state

    def __iter__(self):
        yield self.priority
        yield self.node

    def __lt__(self, other):  # so we don't need the tie-breaker.
        return self.priority < other.priority


def run_heuristic_search(
    initial_state: State,
    check_goal: Callable[[State], bool],
    get_priority: Callable[[State, int], float],
    get_successors: Callable[[State], Iterator[Tuple[Action, State, float]]],
    check_visited: bool = True,
    max_expansions: int = 10000,
    max_depth: int = 1000
):
    """A generic implementation for heuristic search.

    Args:
        initial_state: the initial state.
        check_goal: a function mapping from state to bool, returning True if the state has satisfied the goal.
        get_priority: a function mapping from state to float, returning the priority of the state.
            Smaller priority means higher priority (will be visited first).
        get_successors: a function mapping from state to an iterator of (action, state, cost) tuples.
        check_visited: whether to check if the state has been visited. Set to false if the State representation
            is not hashable.
        max_expansions: the maximum number of expansions.
        max_depth: the maximum depth of the search tree.

    Returns:
        A tuple of (state_sequence, action_sequence, cost_sequence, nr_expansions).

    Raises:
        ValueError: if the search fails.
    """
    queue: List[QueueNode] = list()
    visited: Set[State] = set()

    def push_node(node: SearchNode):
        hq.heappush(queue, QueueNode(get_priority(node.state, node.g), node))
        if check_visited:
            visited.add(node.state)

    root_node = SearchNode(state=initial_state, parent=None, action=None, cost=None, g=0)
    push_node(root_node)
    nr_expansions = 0

    while len(queue) > 0 and nr_expansions < max_expansions:
        priority, node = hq.heappop(queue)

        if node.g > max_depth:
            raise RuntimeError('Failed to find a plan (maximum depth reached).')

        if check_goal(node.state):
            return backtrace_plan(node, nr_expansions)
        nr_expansions += 1
        for action, child_state, cost in get_successors(node.state):
            if check_visited and child_state in visited:
                continue
            child_node = SearchNode(state=child_state, parent=node, action=action, cost=cost, g=node.g + cost)
            push_node(child_node)

    raise RuntimeError('Failed to find a plan (maximum expansion reached).')


def backtrace_plan(node: SearchNode, nr_expansions: int) -> Tuple[List[State], List[Action], List[float], int]:
    """Backtrace the plan from the goal node.

    Args:
        node: a search node where the goal is satisfied.
        nr_expansions: the number of expansions. This value will be returned by this function.

    Returns:
        a tuple of (state_sequence, action_sequence, cost_sequence, nr_expansions).
    """
    state_sequence = []
    action_sequence = []
    cost_sequence = []

    while node.parent is not None:
        action_sequence.insert(0, node.action)
        state_sequence.insert(0, node.state)
        cost_sequence.insert(0, node.cost)
        node = node.parent
    state_sequence.insert(0, node.state)

    return state_sequence, action_sequence, cost_sequence, nr_expansions


concepts/algorithm/search/heuristic_search.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/13/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data loaders for various datasets."""


concepts/benchmark/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 06/13/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from .quickaccess import *

concepts/benchmark/algorithm_env/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : graph.py
# Author : Honghua Dong
# Email  : dhh19951@gmail.com
# Date   : 05/07/2018
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import copy
import numpy as np

import jacinle.random as random

__all__ = ['Graph', 'random_generate_graph', 'random_generate_graph_dnc', 'random_generate_special_graph']


class Graph(object):
    """Store a graph using adjacency matrix."""

    def __init__(self, nr_nodes, edges, coordinates=None):
        """Initialize a graph.

        Args:
            nr_nodes: The Number of nodes in the graph.
            edges: The adjacency matrix of the graph.
        """
        edges = edges.astype('int32')
        assert edges.min() >= 0 and edges.max() <= 1
        self._nr_nodes = nr_nodes
        self._edges = edges
        self._coordinates = coordinates
        self._shortest = None
        self.extra_info = {}

    def get_edges(self):
        return copy.copy(self._edges)

    def get_coordinates(self):
        return self._coordinates

    def get_relations(self):
        """Return edges and identity matrix."""
        return np.stack([self.get_edges(), np.eye(self._nr_nodes)], axis=-1)

    def has_edge(self, x, y):
        return self._edges[x, y] == 1

    def get_out_degree(self):
        """Return the out degree of each node."""
        return np.sum(self._edges, axis=1)

    def get_shortest(self):
        """Return the length of shortest path between nodes."""
        if self._shortest is not None:
            return self._shortest

        n = self._nr_nodes
        edges = self.get_edges()

        # n + 1 indicates unreachable.
        shortest = np.ones((n, n)) * (n + 1)
        shortest[np.where(edges == 1)] = 1
        # Make sure that shortest[x, x] = 0
        shortest -= shortest * np.eye(n)
        shortest = shortest.astype(np.int32)

        # Floyd Algorithm
        for k in range(n):
            for i in range(n):
                for j in range(n):
                    if i != j:
                        shortest[i, j] = min(shortest[i, j], shortest[i, k] + shortest[k, j])
        self._shortest = shortest
        return self._shortest

    def get_connectivity(self, k=None, exclude_self=True):
        """Return the k-connectivity for each pair of nodes. It will return the full connectivity matrix if k is None or k < 0.
        When exclude_self is True, the diagonal elements will be 0.

        Args:
            k: the k-connectivity. Default: ``None``(full connectivity).
            exclude_self: exclude the diagonal elements. Default: ``True``.

        Returns:
            conn: The connectivity matrix.
        """
        shortest = self.get_shortest()
        if k is None or k < 0:
            k = self._nr_nodes
        k = min(k, self._nr_nodes)
        conn = (shortest <= k).astype(np.int32)
        if exclude_self:
            n = self._nr_nodes
            inds = np.where(~np.eye(n, dtype=np.bool_))
            conn = conn[inds]
            conn.resize(n, n - 1)
        return conn


def random_generate_special_graph(n: int, graph_type: str, directed: bool = False) -> Graph:
    """Randomly generate a special type graph.

    For list graph, the nodes are randomly permuted and connected in order. If the graph is directed, the edges are
    directed from the first node to the last node.

    Args:
        n: The number of nodes in the graph.
        graph_type: The type of the graph, e.g. list, tree. Currently only support list.
        directed: Directed or Undirected graph. Default: ``False``(undirected).

    Returns:
        graph: Generated graph.
    """
    if graph_type == 'list':
        nodes = random.permutation(n)
        edges = np.zeros((n, n))
        for i in range(n - 1):
            x, y = nodes[i], nodes[i + 1]
            if directed:
                edges[x, y] = 1
            else:
                edges[x, y] = edges[y, x] = 1
        graph = Graph(n, edges)
        graph.extra_info['nodes_list'] = nodes
        return graph
    else:
        assert False, "not supported graph type: {}".format(graph_type)


def random_generate_graph(n, p, directed=False) -> Graph:
    """Randomly generate a graph by sampling the existence of each edge.
    Each edge between nodes has the probability :math:`p`(directed) or :math:`p^2`(undirected) to not exist.

    This paradigm is also called the Erdős–Rényi model.

    Args:
        n: The number of nodes in the graph.
        p: the probability that a edge doesn't exist in directed graph.
        directed: Directed or Undirected graph. Default: ``False``(undirected)

    Returns:
        graph: Generated graph.
    """
    edges = (random.rand(n, n) < p).astype(np.float32)
    edges -= edges * np.eye(n)
    if not directed:
        edges = np.maximum(edges, edges.T)
    return Graph(n, edges)


def random_generate_graph_dnc(n, p=None, directed=False) -> Graph:
    """Random graph generation method as in DNC, the Differentiable Neural Computer paper.
    Sample :math:`n` nodes in a unit square. sample out-degree (:math:`m`) of each nodes,
    connect to :math:`m` nearest neighbors (Euclidean distance) in the unit square.

    Args:
        n: The number of nodes in the graph.
        p: Control the range of the sample of out-degree. Default: :math:`(1, n // 3)`
            (float): :math:`(1, int(n * p))`
            (int): :math:`(1, p)`
            (tuple): :math:`(p[0], p[1])`
        directed: Directed or Undirected graph. Default: False (undirected)

    Returns:
        graph: A randomly generated graph.
    """
    edges = np.zeros((n, n), dtype=np.float32)
    pos = random.rand(n, 2)

    def dist(x, y):
        return ((x - y) ** 2).mean()

    if type(p) is tuple:
        lower, upper = p
    else:
        lower = 1
        if p is None:
            upper = n // 3
        elif type(p) is int:
            upper = p
        elif type(p) is float:
            upper = int(n * p)
        else:
            assert False
        upper = max(upper, 1)
    lower = max(lower, 1)
    upper = min(upper, n - 1)

    for i in range(n):
        d = []
        k = random.randint(upper - lower + 1) + lower
        for j in range(n):
            if i != j:
                d.append((dist(pos[i], pos[j]), j))
        d.sort()
        for j in range(k):
            edges[i, d[j][1]] = 1
    if not directed:
        edges = np.maximum(edges, edges.T)
    return Graph(n, edges, pos)

concepts/benchmark/algorithm_env/graph.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : graph_env.py
# Author : Honghua Dong
# Email  : dhh19951@gmail.com
# Date   : 04/27/2018
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np

import jacinle.random as random
from jacinle.utils.meta import notnone_property
from jaclearn.rl.env import SimpleRLEnvBase

from concepts.benchmark.algorithm_env.graph import random_generate_graph, random_generate_graph_dnc, random_generate_special_graph

__all__ = ['GraphEnvBase', 'PathGraphEnv']


class GraphEnvBase(SimpleRLEnvBase):
    """Graph Env Base."""

    def __init__(self, nr_nodes, p=0.5, directed=False, gen_method='edge'):
        """Initialize the environment.

        Args:
            n: The number of nodes in the graph.
            p: Parameter for random generation. (Default 0.5)
                (edge method): The probability that a edge doesn't exist in directed graph.
                (dnc method): Control the range of the sample of out-degree.
                other methods: Unused.
            directed: Directed or Undirected graph. Default: `False`(undirected)
            gen_method: Use which method to randomly generate a graph.
                'edge': By sampling the existance of each edge.
                'dnc': Sample out-degree (:math:`m`) of each nodes, and link to nearest neighbors in the unit square.
                'list': generate a chain-like graph.
        """

        super().__init__()
        self._nr_nodes = nr_nodes
        self._p = p
        self._directed = directed
        self._gen_method = gen_method
        self._graph = None

    @notnone_property
    def graph(self):
        """The generated graph."""
        return self._graph

    def _restart(self):
        """ Restart the environment. """
        self._gen_graph()

    def _gen_graph(self):
        """ generate the graph by specified method. """
        n = self._nr_nodes
        p = self._p
        if self._gen_method == 'edge':
            self.graph = random_generate_graph(n, p, self._directed)
        elif self._gen_method == 'dnc':
            self.graph = random_generate_graph_dnc(n, p, self._directed)
        else:
            self._graph = random_generate_special_graph(n, self._gen_method, self._directed)


class PathGraphEnv(GraphEnvBase):
    """Env for Finding a path from starting node to the destination."""
    def __init__(self, nr_nodes, dist_range, p=0.5, directed=False, gen_method='edge'):
        """Initialize the environment.

        Args:
            nr_nodes: The number of nodes in the graph.
            p: Parameter for random generation. (Default: 0.5)
                (edge method): The probability that an edge doesn't exist in directed graph.
                (dnc method): Control the range of the sample of out-degree.
                other methods: Unused.
            directed: Directed or Undirected graph. Default: `False` (undirected)
            gen_method: Use which method to randomly generate a graph.
                'edge': By sampling the existance of each edge.
                'dnc': Sample out-degree (:math:`m`) of each nodes, and link to nearest neighbors in the unit square.
                'list': generate a chain-like graph.
            dist_range: The sampling range of distance between starting node and the destination.
        """
        super().__init__(nr_nodes, p, directed, gen_method)
        self._dist_range = dist_range

    @property
    def dist(self):
        """The distance between starting node and the destination."""
        return self._dist

    def _restart(self):
        """ Restart the environment. """
        super()._restart()
        self._dist = self._sample_dist()
        self._task = None
        while True:
            self._task = self._gen()
            if self._task is not None:
                break
            # If fail to find two nodes with sampled distance, generate another graph.
            self._gen_graph()
        self._current = self._task[0]
        self._set_current_state(self._task)
        self._steps = 0

    def _sample_dist(self):
        lower, upper = self._dist_range
        upper = min(upper, self._nr_nodes - 1)
        return random.randint(upper - lower + 1) + lower

    def _gen(self):
        """ Sample the starting node and the destination according to the distance. """
        dist_matrix = self._graph.get_shortest()
        st, ed = np.where(dist_matrix == self.dist)
        if len(st) == 0:
            return None
        ind = random.randint(len(st))
        return st[ind], ed[ind]

    def _action(self, target):
        """
            Move to the target node from current node if has_edge(current -> target).
            Returns: reward, is_over
        """
        if self._current == self._task[1]:
            return 1, True
        if self._graph.has_edge(self._current, target):
            self._current = target
        self._set_current_state((self._current, self._task[1]))
        if self._current == self._task[1]:
            return 1, True
        self._steps += 1
        if self._steps >= self.dist:
            return 0, True
        return 0, False


concepts/benchmark/algorithm_env/graph_env.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : quickaccess.py
# Author : Honghua Dong
# Email  : dhh19951@gmail.com
# Date   : 05/11/2018
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from jaclearn.rl.env import RLEnvBase, ProxyRLEnvBase
from jaclearn.rl.space import DiscreteActionSpace
from jaclearn.rl.proxy import LimitLengthProxy

from concepts.benchmark.algorithm_env.sort_envs import ListSortingEnv
from concepts.benchmark.algorithm_env.graph_env import PathGraphEnv

__all__ = ['get_sort_env', 'get_path_env', 'make']


class _MapActionProxy(ProxyRLEnvBase):
    def __init__(self, other, mapping):
        super().__init__(other)
        self._mapping = mapping

    def map_action(self, action):
        assert action < len(self._mapping)
        return self._mapping[action]

    def _get_action_space(self):
        return DiscreteActionSpace(len(self._mapping))

    def _action(self, action):
        return self.proxy.action(self.map_action(action))


def _map_graph_action(p, n, exclude_self=True):
    mapping = [(i, j) for i in range(n) for j in range(n) if (i != j or not exclude_self)]
    p = _MapActionProxy(p, mapping)
    return p


def get_sort_env(n: int, exclude_self: bool = True) -> ListSortingEnv:
    """Get a sorting environment with n elements.

    Args:
        n: number of elements.
        exclude_self: whether to exclude swap(i, i) actions.

    Returns:
        A sorting environment.
    """
    env_cls = ListSortingEnv
    p = env_cls(n)
    p = LimitLengthProxy(p, n * 2)
    p = _map_graph_action(p, n, exclude_self=exclude_self)
    return p


def get_path_env(n, dist_range, prob_edge=0.5, directed=False, gen_method='edge'):
    """Get a path-finding environment with n nodes."""
    env_cls = PathGraphEnv
    p = env_cls(n, dist_range, prob_edge, directed=directed, gen_method=gen_method)
    # p = LimitLengthProxy(p, max_len)
    return p


def make(task: str, *args, **kwargs) -> RLEnvBase:
    if task == 'sort':
        return get_sort_env(*args, **kwargs)
    elif task == 'path':
        return get_path_env(*args, **kwargs)
    else:
        raise ValueError('Unknown task: {}.'.format(task))

concepts/benchmark/algorithm_env/quickaccess.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : sort_envs.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 06/09/2018
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np

import jacinle.random as random
from jacinle.utils.meta import notnone_property
from jaclearn.rl.env import SimpleRLEnvBase


class ListSortingEnv(SimpleRLEnvBase):
    """Env for sorting a random permutation."""

    def __init__(self, nr_numbers: int):
        """Initialize the environment.

        Args:
            nr_numbers: The number of numbers in the array.
        """
        super().__init__()
        self._nr_numbers = nr_numbers
        self._array = None

    @notnone_property
    def array(self):
        """The underlying array to be sorted."""
        return self._array

    def get_state(self):
        """Compute the state given the array."""
        x, y = np.meshgrid(self.array, self.array)
        number_relations = np.stack([x < y, x == y, x > y], axis=-1).astype('float')
        index = np.array(list(range(self._nr_numbers)))
        x, y = np.meshgrid(index, index)
        position_relations = np.stack([x < y, x == y, x > y], axis=-1).astype('float')
        return np.concatenate([number_relations, position_relations], axis=-1)

    def _calculate_optimal(self):
        """Calculate the optimal number of steps for sorting the array."""
        a = self._array
        b = [0 for i in range(len(a))]
        cnt = 0
        for i, x in enumerate(a):
            if b[i] == 0:
                j = x
                b[i] = 1
                while b[j] == 0:
                    b[j] = 1
                    j = a[j]
                assert i == j
                cnt += 1
        return len(a) - cnt

    def _restart(self):
        """Restart: Generate a random permutation."""
        self._array = random.permutation(self._nr_numbers)
        self._set_current_state(self.get_state())
        self.optimal = self._calculate_optimal()

    def _action(self, action):
        """Action: Swap the numbers at the index :math:`i` and :math:`j`. Returns: reward, is_over.

        Args:
            action: A tuple of two numbers, indicating the indices to be swapped.

        Returns:
            reward: 1 if the array is sorted, 0 otherwise.
            is_over: True if the array is sorted, False otherwise.
        """
        a = self._array
        i, j = action
        x, y = a[i], a[j]
        a[i], a[j] = y, x
        self._set_current_state(self.get_state())
        for i in range(self._nr_numbers):
            if a[i] != i:
                return 0, False
        return 1, True

concepts/benchmark/algorithm_env/sort_envs.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 06/13/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/benchmark/blocksworld/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/26/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/benchmark/clevr/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : clevr_constants.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/29/2018
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""
Basic concepts in the CLEVR dataset.
"""

__all__ = ['g_attribute_concepts', 'g_relational_concepts', 'g_synonyms']


g_attribute_concepts = {
    'color': ['gray', 'red', 'blue', 'green', 'brown', 'purple', 'cyan', 'yellow'],
    'material': ['rubber', 'metal'],
    'shape': ['cube', 'sphere', 'cylinder'],
    'size': ['small', 'large']
}

g_relational_concepts = {
    'spatial_relation': ['left', 'right', 'front', 'behind']
}

g_synonyms = {
    "thing": ["thing", "object"],
    "sphere": ["sphere", "ball", "spheres", "balls"],
    "cube": ["cube", "block", "cubes", "blocks"],
    "cylinder": ["cylinder", "cylinders"],
    "large": ["large", "big"],
    "small": ["small", "tiny"],
    "metal": ["metallic", "metal", "shiny"],
    "rubber": ["rubber", "matte"],
}


def read_concepts():
    attribute_concepts = []
    for k in g_attribute_concepts.keys():
        attribute_concepts.extend(g_attribute_concepts[k])
    for k in g_synonyms.keys():
        attribute_concepts.extend(g_synonyms[k])
    attribute_concepts = list(set(attribute_concepts))

    relational_concepts = g_relational_concepts['spatial_relation']
    multi_relational_concepts = g_relational_concepts['spatial_relation']

    return attribute_concepts, relational_concepts, multi_relational_concepts


concepts/benchmark/clevr/clevr_constants.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : dataset.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/26/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import os.path as osp

import nltk
import numpy as np
from PIL import Image

import jacinle.io as io
import jaclearn.vision.coco.mask_utils as mask_utils

from jacinle.logging import get_logger
from jacinle.utils.container import GView
from jactorch.data.dataset import FilterableDatasetUnwrapped, FilterableDatasetView

from concepts.benchmark.common.vocab import Vocab
from concepts.benchmark.clevr.clevr_constants import g_attribute_concepts, g_relational_concepts

logger = get_logger(__file__)

__all__ = ['CLEVRDatasetUnwrapped', 'CLEVRDatasetFilterableView', 'make_dataset', 'CLEVRCustomTransferDataset', 'make_custom_transfer_dataset']


class CLEVRDatasetUnwrapped(FilterableDatasetUnwrapped):
    def __init__(self, scenes_json, questions_json, image_root, image_transform, vocab_json, output_vocab_json, question_transform=None, incl_scene=True, incl_raw_scene=False):
        super().__init__()

        self.scenes_json = scenes_json
        self.questions_json = questions_json
        self.image_root = image_root
        self.image_transform = image_transform
        self.vocab_json = vocab_json
        self.output_vocab_json = output_vocab_json
        self.question_transform = question_transform

        self.incl_scene = incl_scene
        self.incl_raw_scene = incl_raw_scene

        logger.info('Loading scenes from: "{}".'.format(self.scenes_json))
        self.scenes = io.load_json(self.scenes_json)['scenes']
        if isinstance(self.questions_json, (tuple, list)):
            self.questions = list()
            for filename in self.questions_json:
                logger.info('Loading questions from: "{}".'.format(filename))
                self.questions.extend(io.load_json(filename)['questions'])
        else:
            logger.info('Loading questions from: "{}".'.format(self.questions_json))
            self.questions = io.load_json(self.questions_json)['questions']

        if self.vocab_json is not None:
            logger.info('Loading vocab from: "{}".'.format(self.vocab_json))
            self.vocab = Vocab.from_json(self.vocab_json)
        else:
            logger.info('Building the vocab.')
            self.vocab = Vocab.from_dataset(self, keys=['question_tokenized'])

        if output_vocab_json is not None:
            logger.info('Loading output vocab from: "{}".'.format(self.output_vocab_json))
            self.output_vocab = Vocab.from_json(self.output_vocab_json)
        else:
            logger.info('Building the output vocab.')
            self.output_vocab = Vocab.from_dataset(self, keys=['answer'], single_word=True)

    def _get_metainfo(self, index):
        question = self.questions[index]
        scene = self.scenes[question['image_index']]
        question['scene'] = scene
        question['program'] = question.pop('program', None)  # In CLEVR-Humans, there is no program.

        question['image_index'] = question['image_index']
        question['image_filename'] = self._get_image_filename(scene)
        question['question_index'] = index
        question['question_tokenized'] = nltk.word_tokenize(question['question'])
        question['question_type'] = get_question_type(question['program'])

        # question['scene_complexity'] = len(scene['objects'])
        # question['program_complexity'] = len(question['program'])

        return question

    def _get_image_filename(self, scene: dict) -> str:
        return scene['image_filename']

    def __getitem__(self, index):
        metainfo = GView(self.get_metainfo(index))
        feed_dict = GView()

        # metainfo annotations
        if self.incl_scene:
            feed_dict.update(annotate_objects(metainfo.scene))
            if 'objects' in feed_dict:
                # NB(Jiayuan Mao): in some datasets_v1, object information might be completely unavailable.
                feed_dict.objects_raw = feed_dict.objects.copy()
            feed_dict.update(annotate_scene(metainfo.scene))

        if self.incl_raw_scene:
            feed_dict.scene = metainfo.scene

        # image
        feed_dict.image_index = metainfo.image_index
        feed_dict.image_filename = metainfo.image_filename
        if self.image_root is not None and feed_dict.image_filename is not None:
            feed_dict.image = Image.open(osp.join(self.image_root, feed_dict.image_filename)).convert('RGB')
            feed_dict.image, feed_dict.objects = self.image_transform(feed_dict.image, feed_dict.objects)

        # question
        feed_dict.question_index = metainfo.question_index
        feed_dict.question_raw = metainfo.question
        feed_dict.question_raw_tokenized = metainfo.question_tokenized
        feed_dict.question = metainfo.question_tokenized
        feed_dict.question_type = metainfo.question_type
        feed_dict.answer = canonize_answer(metainfo.answer, None)

        if self.question_transform is not None:
            self.question_transform(feed_dict)
        feed_dict.question = np.array(self.vocab.map_sequence(feed_dict.question), dtype=np.int64)

        return feed_dict.raw()

    def __len__(self):
        return len(self.questions)


class CLEVRDatasetFilterableView(FilterableDatasetView):
    def filter_program_size_raw(self, max_length):
        def filt(question):
            return question['program'] is None or len(question['program']) <= max_length

        return self.filter(filt, 'filter-program-size-clevr[{}]'.format(max_length))

    def filter_scene_size(self, max_scene_size):
        def filt(question):
            return len(question['scene']['objects']) <= max_scene_size

        return self.filter(filt, 'filter-scene-size[{}]'.format(max_scene_size))

    def filter_question_type(self, *, allowed=None, disallowed=None):
        def filt(question):
            if allowed is not None:
                return question['question_type'] in allowed
            elif disallowed is not None:
                return question['question_type'] not in disallowed

        if allowed is not None:
            return self.filter(filt, 'filter-question-type[allowed={' + (','.join(list(allowed))) + '}]')
        elif disallowed is not None:
            return self.filter(filt, 'filter-question-type[disallowed={' + (','.join(list(disallowed))) + '}]')
        else:
            raise ValueError('Must provide either allowed={...} or disallowed={...}.')

    def make_dataloader(self, batch_size, shuffle, drop_last, nr_workers):
        from jactorch.data.dataloader import JacDataLoader
        from jactorch.data.collate import VarLengthCollateV2

        collate_guide = {
            'scene': 'skip',
            'objects_raw': 'skip',
            'objects': 'concat',

            'image_index': 'skip',
            'image_filename': 'skip',
            'question_index': 'skip',
            'question_metainfo': 'skip',

            'question_raw': 'skip',
            'question_raw_tokenized': 'skip',
            'question_type': 'skip',
            'question': 'pad',

            'program_raw': 'skip',
            'program_seq': 'skip',
            'program_tree': 'skip',
            'program_qsseq': 'skip',
            'program_qstree': 'skip',

            'question_type': 'skip',
            'answer': 'skip',
        }

        # Scene annotations.
        for attr_name in g_attribute_concepts:
            collate_guide['attribute_' + attr_name] = 'concat'
            collate_guide['attribute_relation_' + attr_name] = 'concat'
        for attr_name in g_relational_concepts:
            collate_guide['relation_' + attr_name] = 'concat'

        return JacDataLoader(
            self, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last,
            num_workers=nr_workers, pin_memory=True,
            collate_fn=VarLengthCollateV2(collate_guide)
        )


class CLEVRCustomTransferDataset(FilterableDatasetUnwrapped):
    def __init__(self, scenes_json, questions_json, image_root, image_transform, query_list_key, custom_fields, output_vocab_json=None, incl_scene=True, incl_raw_scene=False):
        super().__init__()

        self.scenes_json = scenes_json
        self.questions_json = questions_json
        self.image_root = image_root
        self.image_transform = image_transform
        self.query_list_key = query_list_key
        self.custom_fields = custom_fields
        self.output_vocab_json = output_vocab_json

        self.incl_scene = incl_scene
        self.incl_raw_scene = incl_raw_scene

        logger.info('Loading scenes from: "{}".'.format(self.scenes_json))
        self.scenes = io.load_json(self.scenes_json)['scenes']
        if isinstance(self.questions_json, (tuple, list)):
            self.questions = list()
            for filename in self.questions_json:
                logger.info('Loading questions from: "{}".'.format(filename))
                self.questions.extend(io.load_json(filename)[query_list_key])
        else:
            logger.info('Loading questions from: "{}".'.format(self.questions_json))
            self.questions = io.load_json(self.questions_json)[query_list_key]

        if output_vocab_json is not None:
            logger.info('Loading output vocab from: "{}".'.format(self.output_vocab_json))
            self.output_vocab = Vocab.from_json(self.output_vocab_json)
        else:
            logger.info('Building the output vocab.')
            self.output_vocab = Vocab.from_dataset(self, keys=['answer'], single_word=True)

    def _get_scene_index_from_question(self, question):
        if 'image_index' in question:
            return question['image_index']
        if 'scene_index' in question:
            return question['scene_index']
        raise KeyError('Cannot find scene index from question.')

    def _get_metainfo(self, index):
        question = self.questions[index]
        scene = self.scenes[self._get_scene_index_from_question(question)]
        question['scene'] = scene
        question['program'] = question.pop('program', None)  # In CLEVR-Humans, there is no program.

        question['image_index'] = self._get_scene_index_from_question(question)
        question['image_filename'] = self._get_image_filename(scene)
        question['question_index'] = index
        question['question'] = question['question']
        question['answer'] = question['answer']
        question['question_type'] = question.get('question_type', self.query_list_key)

        for field_name in self.custom_fields:
            question[field_name] = scene[field_name]

        return question

    def _get_image_filename(self, scene: dict) -> str:
        return scene['image_filename']

    def __getitem__(self, index):
        metainfo = GView(self.get_metainfo(index))
        feed_dict = GView()

        # metainfo annotations
        feed_dict.update(annotate_objects(metainfo.scene))
        if 'objects' in feed_dict:
            # NB(Jiayuan Mao): in some datasets_v1, object information might be completely unavailable.
            feed_dict.objects_raw = feed_dict.objects.copy()

        if self.incl_scene:
            feed_dict.update(annotate_scene(metainfo.scene))
        if self.incl_raw_scene:
            feed_dict.scene = metainfo.scene

        # image
        feed_dict.image_index = metainfo.image_index
        feed_dict.image_filename = metainfo.image_filename
        if self.image_root is not None and feed_dict.image_filename is not None:
            feed_dict.image = Image.open(osp.join(self.image_root, feed_dict.image_filename)).convert('RGB')
            feed_dict.image, feed_dict.objects = self.image_transform(feed_dict.image, feed_dict.objects)

        # question
        feed_dict.question_raw = metainfo.question
        feed_dict.question_type = metainfo.question_type
        feed_dict.answer = metainfo.answer
        return feed_dict.raw()

    def __len__(self):
        return len(self.questions)

    def make_dataloader(self, batch_size, shuffle, drop_last, nr_workers):
        from jactorch.data.dataloader import JacDataLoader
        from jactorch.data.collate import VarLengthCollateV2

        collate_guide = {
            'scene': 'skip',
            'objects_raw': 'skip',
            'objects': 'concat',

            'image_index': 'skip',
            'image_filename': 'skip',
            'question_index': 'skip',
            'question_metainfo': 'skip',

            'question_raw': 'skip',
            'question_type': 'skip',
            'answer': 'skip',
        }

        for field in self.custom_fields:
            collate_guide[field] = 'skip'

        # Scene annotations.
        for attr_name in g_attribute_concepts:
            collate_guide['attribute_' + attr_name] = 'concat'
            collate_guide['attribute_relation_' + attr_name] = 'concat'
        for attr_name in g_relational_concepts:
            collate_guide['relation_' + attr_name] = 'concat'

        return JacDataLoader(
            self, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last,
            num_workers=nr_workers, pin_memory=True,
            collate_fn=VarLengthCollateV2(collate_guide)
        )


def make_dataset(scenes_json, questions_json, image_root, *, image_transform=None, vocab_json=None, output_vocab_json=None, filterable_view_cls=None, **kwargs):
    if filterable_view_cls is None:
        filterable_view_cls = CLEVRDatasetFilterableView

    if image_transform is None:
        import jactorch.transforms.bbox as T
        image_transform = T.Compose([
            T.NormalizeBbox(),
            T.Resize(256),
            T.DenormalizeBbox(),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    return filterable_view_cls(CLEVRDatasetUnwrapped(
        scenes_json, questions_json, image_root, image_transform, vocab_json, output_vocab_json, **kwargs
    ))


def make_custom_transfer_dataset(scenes_json, questions_json, image_root, query_list_key, custom_fields, *, image_transform=None, output_vocab_json=None, **kwargs):
    if image_transform is None:
        import jactorch.transforms.bbox as T
        image_transform = T.Compose([
            T.NormalizeBbox(),
            T.Resize(256),
            T.DenormalizeBbox(),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    return CLEVRCustomTransferDataset(
        scenes_json, questions_json, image_root,
        query_list_key=query_list_key, custom_fields=custom_fields,
        image_transform=image_transform, output_vocab_json=output_vocab_json, **kwargs
    )


def annotate_scene(scene):
    feed_dict = dict()

    if not _is_object_annotation_available(scene):
        return feed_dict

    for attr_name, concepts in g_attribute_concepts.items():
        concepts2id = {v: i for i, v in enumerate(concepts)}
        values = list()
        for obj in scene['objects']:
            assert attr_name in obj
            values.append(concepts2id[obj[attr_name]])
        values = np.array(values, dtype='int64')
        feed_dict['attribute_' + attr_name] = values
        lhs, rhs = np.meshgrid(values, values)
        compare_label = (lhs == rhs).astype('float32')
        compare_label[np.diag_indices_from(compare_label)] = 0
        feed_dict['attribute_relation_' + attr_name] = compare_label.reshape(-1)

    nr_objects = len(scene['objects'])
    for attr_name, concepts in g_relational_concepts.items():
        concept_values = []
        for concept in concepts:
            values = np.zeros((nr_objects, nr_objects), dtype='float32')
            assert concept in scene['relationships']
            this_relation = scene['relationships'][concept]
            assert len(this_relation) == nr_objects
            for i, this_row in enumerate(this_relation):
                for j in this_row:
                    values[j, i] = 1
            concept_values.append(values)
        concept_values = np.stack(concept_values, -1)
        feed_dict['relation_' + attr_name] = concept_values.reshape(-1, concept_values.shape[-1])

    return feed_dict


def canonize_answer(answer, question_type):
    if answer in ('yes', 'no'):
        answer = (answer == 'yes')
    elif isinstance(answer, str) and answer.isdigit():
        answer = int(answer)
        assert 0 <= answer <= 10
    return answer


def annotate_objects(scene):
    if 'objects' not in scene and 'objects_detection' not in scene:
        return dict()

    boxes = [mask_utils.toBbox(i['mask']) for i in _get_object_masks(scene)]
    if len(boxes) == 0:
        return {'objects': np.zeros((0, 4), dtype='float32')}
    boxes = np.array(boxes)
    boxes[:, 2] += boxes[:, 0]
    boxes[:, 3] += boxes[:, 1]
    return {'objects': boxes.astype('float32')}


def _is_object_annotation_available(scene):
    if len(scene['objects']) > 0 and 'mask' in scene['objects'][0]:
        return True
    return False


def _get_object_masks(scene):
    """Backward compatibility: in self-generated clevr scenes, the groundtruth masks are provided;
    while in the clevr test data, we use Mask R-CNN to detect all the masks, and stored in `objects_detection`."""
    if 'objects_detection' not in scene:
        return scene['objects']
    if _is_object_annotation_available(scene):
        return scene['objects']
    return scene['objects_detection']


g_last_op_to_question_type = {
    'query_color': 'query_attr',
    'query_shape': 'query_attr',
    'query_material': 'query_attr',
    'query_size': 'query_attr',
    'exist': 'exist',
    'count': 'count',
    'equal_integer': 'cmp_number',
    'greater_than': 'cmp_number',
    'less_than': 'cmp_number',
    'equal_color': 'cmp_attr',
    'equal_shape': 'cmp_attr',
    'equal_material': 'cmp_attr',
    'equal_size': 'cmp_attr',
}


def get_op_type(op):
    if 'type' in op:
        return op['type']
    return op['function']


def get_question_type(program):
    if program is None:
        return 'unk'
    return g_last_op_to_question_type[get_op_type(program[-1])]


concepts/benchmark/clevr/dataset.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/26/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/benchmark/common/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : vocab.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/26/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import torch

import jacinle.io as io
import jaclearn.embedding.constant as const
from jacinle.utils.tqdm import tqdm

__all__ = ['Vocab', 'gen_vocab', 'gen_vocab_from_words']


class Vocab(object):
    def __init__(self, word2idx=None):
        self.word2idx = word2idx if word2idx is not None else dict()
        self._idx2word = None

    @classmethod
    def from_json(cls, json_file):
        return cls(io.load_json(json_file))

    @classmethod
    def from_dataset(cls, dataset, keys, extra_words=None, single_word=False):
        return gen_vocab(dataset, keys, extra_words=extra_words, cls=cls, single_word=single_word)

    @classmethod
    def from_list(cls, dataset, extra_words=None, single_word=False):
        return gen_vocab(dataset, extra_words=extra_words, cls=cls, single_word=single_word)

    def dump_json(self, json_file):
        io.dump_json(json_file, self.word2idx)

    def check_json_consistency(self, json_file):
        rhs = io.load_json(json_file)
        for k, v in self.word2idx.items():
            if not (k in rhs and rhs[k] == v):
                return False
        return True

    def words(self):
        return self.word2idx.keys()

    @property
    def idx2word(self):
        if self._idx2word is None or len(self.word2idx) != len(self._idx2word):
            self._idx2word = {v: k for k, v in self.word2idx.items()}
        return self._idx2word

    def __len__(self):
        return len(self.word2idx)

    def __iter__(self):
        return iter(self.word2idx.keys())

    def add(self, word):
        self.add_word(word)

    def add_word(self, word):
        self.word2idx[word] = len(self.word2idx)

    def map(self, word):
        return self.word2idx.get(
            word,
            self.word2idx.get(const.EBD_UNKNOWN, -1)
        )

    def map_sequence(self, sequence, add_be=False):
        if isinstance(sequence, str):
            sequence = sequence.split()
        sequence = [self.map(w) for w in sequence]
        if add_be:
            sequence.insert(0, self.word2idx[const.EBD_BOS])
            sequence.append(self.word2idx[const.EBD_EOS])
        return sequence

    def map_fields(self, feed_dict, fields):
        feed_dict = feed_dict.copy()
        for k in fields:
            if k in feed_dict:
                feed_dict[k] = self.map(feed_dict[k])
        return feed_dict

    def invmap_sequence(self, sequence, proc_be=False):
        if proc_be:
            raise NotImplementedError()
        if torch.is_tensor(sequence):
            sequence = sequence.detach().cpu().tolist()
        return [self.idx2word[int(x)] for x in sequence]


def gen_vocab(dataset, keys=None, extra_words=None, cls=None, single_word=False):
    """Generate a Vocabulary instance from a dataset.

    By default this function will retrieve the data using the `get_metainfo` function,
    or it will fallback to `dataset[i]` if the function does not exist.

    The function should return a dictionary. Users can specify a list of keys that will
    be returned by the `get_metainfo` function. This function will split the string indexed
    by these keys and add tokens to the vocabulary.
    If the argument `keys` is not specified, this function assumes the return of `get_metainfo`
    to be a string.

    By default, this function will add four additional tokens:
    EBD_PAD, EBD_BOS, EBD_EOS, and EBD_UNK. Users can specify additional extra tokens using the
    extra_words argument.
    """
    if cls is None:
        cls = Vocab

    all_words = set()
    for i in tqdm(len(dataset), desc='Building the vocab'):
        if hasattr(dataset, 'get_metainfo'):
            metainfo = dataset.get_metainfo(i)
        else:
            metainfo = dataset[i]

        if keys is None:
            for w in metainfo.split():
                all_words.add(w)
        else:
            for k in keys:
                if single_word:
                    all_words.add(str(metainfo[k]))
                elif isinstance(metainfo[k], str):
                    for w in metainfo[k].split():
                        all_words.add(w)
                else:
                    for w in metainfo[k]:
                        all_words.add(w)

    vocab = cls()
    vocab.add(const.EBD_ALL_ZEROS)
    for w in sorted(all_words):
        vocab.add(w)
    for w in [const.EBD_UNKNOWN, const.EBD_BOS, const.EBD_EOS]:
        vocab.add(w)

    if extra_words is not None:
        for w in extra_words:
            vocab.add(w)

    return vocab


def gen_vocab_from_words(words, extra_words=None, cls=None):
    """Generate a Vocabulary instance from a list of words."""
    if cls is None:
        cls = Vocab
    vocab = cls()
    vocab.add(const.EBD_ALL_ZEROS)
    for w in sorted(words):
        vocab.add(w)
    for w in [const.EBD_UNKNOWN, const.EBD_BOS, const.EBD_EOS]:
        vocab.add(w)
    if extra_words is not None:
        for w in extra_words:
            vocab.add(w)
    return vocab


concepts/benchmark/common/vocab.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/06/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/benchmark/gridworld/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/06/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/benchmark/gridworld/crafting_world/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : crafting_world_env.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/23/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
import os.path as osp
from PIL import Image
from typing import Sequence

from concepts.benchmark.gridworld.crafting_world.crafting_world_rules import MINING_RULES, CRAFTING_RULES, get_all_locations
from concepts.benchmark.gridworld.crafting_world.crafting_world_gen.utils import underline_to_pascal
from concepts.pdsketch.domain import Domain
from concepts.pdsketch.operator import OperatorApplier

SKIP_CRAFTING_LOCATION_CHECK = True


class CraftingWorldSimulator(object):
    def __init__(self):
        self.nr_grids = 15
        self.nr_inventory = 3
        self.agent_pos = 1
        self.objects = dict()  # str: (str, int), name: (type, pos)
        self.inventory = dict()  # int: Optional[Tuple[str, str]]  # (type, name)
        self.hypothetical = set()  # str

    def reset_from_state(self, objects, state):
        agent_at = list(state['agent-at'])[0][0]
        self.agent_pos = int(agent_at[1:])

        self.nr_grids = len(objects['tile'])
        self.nr_inventory = nr_inventory = len(objects['inventory'])

        self.objects = dict()
        self.inventory = {i: None for i in range(1, 1 + nr_inventory)}

        for obj_name, obj_loc in state.get('object-at', []):
            obj_type = None
            for obj_name2, obj_type2 in state['object-of-type']:
                if obj_name2 == obj_name:
                    obj_type = obj_type2
                    break
            assert obj_type is not None
            self.objects[obj_name] = (obj_type, int(obj_loc[1:]))

        for inv_id, obj_name in state.get('inventory-holding', []):
            obj_type = None
            for obj_name2, obj_type2 in state['object-of-type']:
                if obj_name2 == obj_name:
                    obj_type = obj_type2
            assert obj_type is not None
            self.inventory[int(inv_id[1:])] = (obj_type, obj_name)

        for obj_name, obj_type in state['object-of-type']:
            if obj_type == 'Hypothetical':
                self.hypothetical.add(obj_name)

    def move_to(self, pos):
        self.agent_pos = max(1, min(self.nr_grids, pos))
        return True

    def move_left(self):
        self.agent_pos = max(1, self.agent_pos - 1)
        return True

    def move_right(self):
        self.agent_pos = min(self.nr_grids, self.agent_pos + 1)
        return True

    def pick_up(self, inventory, obj_name):
        if self.inventory[inventory] is not None:
            return False
        if self.objects[obj_name][1] != self.agent_pos:
            return False

        self.inventory[inventory] = self.objects[obj_name][0], obj_name
        del self.objects[obj_name]
        return True

    def place_down(self, inventory):
        if self.inventory[inventory] is None:
            return False

        obj_type, obj_name = self.inventory[inventory]
        self.objects[obj_name] = obj_type, self.agent_pos

    def mine(self, obj_name, inventory, hypothetical_object_name, tool_inventory=None):
        if self.objects[obj_name][1] != self.agent_pos:
            return False
        if self.inventory[inventory] is not None:
            return False
        if hypothetical_object_name not in self.hypothetical:
            return False
        if tool_inventory is not None and self.inventory[tool_inventory] is None:
            return False

        obj_type, _ = self.objects[obj_name]

        for rule in MINING_RULES:
            if underline_to_pascal(rule['location']) == obj_type:
                if tool_inventory is None:
                    if len(rule['holding']) == 0:
                        new_obj_type = underline_to_pascal(rule['create'])
                        self.inventory[inventory] = (new_obj_type, hypothetical_object_name)
                        self.hypothetical.remove(hypothetical_object_name)
                        return True
                else:
                    tool_type, _ = self.inventory[tool_inventory]
                    if len(rule['holding']) == 0 or (len(rule['holding']) == 1 and underline_to_pascal(rule['holding'][0]) == tool_type):
                        new_obj_type = underline_to_pascal(rule['create'])
                        self.inventory[inventory] = (new_obj_type, hypothetical_object_name)
                        self.hypothetical.remove(hypothetical_object_name)
                        return True

        return False

    def craft(self, obj_name, inventory, hypothetical_object_name, ingredients_inventory, target_type=None):
        if SKIP_CRAFTING_LOCATION_CHECK:
            if self.objects[obj_name][1] != self.agent_pos:
                return False
        if self.inventory[inventory] is not None:
            return False
        if hypothetical_object_name not in self.hypothetical:
            return False
        for ingredient_inventory in ingredients_inventory:
            if self.inventory[ingredient_inventory] is None:
                return False

        obj_type, _ = self.objects[obj_name]
        # print('Checking crafting', inventory, hypothetical_object_name, target_type, ingredients_inventory)

        for rule in CRAFTING_RULES:
            if target_type is not None and underline_to_pascal(rule['create']) != target_type:
                continue
            # print('  checking crafting rule', rule['location'], rule['recipe'], rule['create'])
            # if not SKIP_CRAFTING_LOCATION_CHECK:
            #     print(f'    matching crafting location', underline_to_pascal(rule['location']), obj_type)
            if underline_to_pascal(rule['location']) == obj_type or SKIP_CRAFTING_LOCATION_CHECK:
                if len(rule['recipe']) == len(ingredients_inventory):
                    current_holding_types = set()
                    for ingredient_inventory in ingredients_inventory:
                        ingredient_type, _ = self.inventory[ingredient_inventory]
                        current_holding_types.add(ingredient_type)
                    target_holding_types = set()
                    for ingredient_type in rule['recipe']:
                        target_holding_types.add(underline_to_pascal(ingredient_type))
                    # print(f'    matching crafting recipe current={current_holding_types}, target={target_holding_types}')
                    if current_holding_types == target_holding_types:
                        new_obj_type = underline_to_pascal(rule['create'])
                        self.inventory[inventory] = (new_obj_type, hypothetical_object_name)
                        self.hypothetical.remove(hypothetical_object_name)
                        for ingredient_inventory in ingredients_inventory:
                            self.hypothetical.add(self.inventory[ingredient_inventory][1])
                            self.inventory[ingredient_inventory] = None
                        return True
        return False


class CraftingWorldRenderer(object):
    def __init__(self, map_w: int, map_h: int, max_inventory: int):
        self._map_w = map_w
        self._map_h = map_h
        self.basic_canvas = np.zeros((map_h * 17 + 1, map_w * 17 + 1, 3), dtype=np.uint8) + 255
        self._inventory_h = int(max_inventory / map_w) + (1 if max_inventory % map_w != 0 else 0)
        self.inventory_canvas = np.zeros((self._inventory_h * 17 + 1, map_w * 17 + 1, 3), dtype=np.uint8) + 255
        self._init_basic_canvas()
        self._block_images = dict()
        self._init_resource_png()

    def _init_basic_canvas(self):
        # Draw grid
        for i in range(self._map_w + 1):
            self.basic_canvas[:, i * 17, :] = 128
        for i in range(self._map_h + 1):
            self.basic_canvas[i * 17, :, :] = 128

        # Draw grid for the inventory
        for i in range(self._inventory_h + 1):
            self.inventory_canvas[i * 17, :, :] = 128
        for i in range(self._map_w + 1):
            self.inventory_canvas[:, i * 17, :] = 128

    def _init_resource_png(self):
        filename = osp.join(osp.dirname(__file__), 'assets', 'BlockCSS.png')
        block_image = np.array(Image.open(filename))
        filename = osp.join(osp.dirname(__file__), 'assets', 'ItemCSS.png')
        item_image = np.array(Image.open(filename))
        filename = osp.join(osp.dirname(__file__), 'assets', 'EntityCSS.png')
        entity_image = np.array(Image.open(filename))

        def extract(image, x, y):
            return image[y:y+16, x:x+16, :]

        def make_small_item_image(image):
            item_image = np.zeros_like(image)
            # item_image[8:8+8, :8, :] = image[::2, ::2, :]
            item_image[8:8+8, :8, :3] = image[::2, ::2, :3] * (image[::2, ::2, 3:4] / 255) + 255 * (1 - image[::2, ::2, 3:4] / 255)
            item_image[8:8+8, :8, 3] = 255
            # add a small border
            item_image[8:8+8, 0, :] = 128
            item_image[8:8+8, 7, :] = 128
            item_image[8, :8, :] = 128
            item_image[15, :8, :] = 128
            return item_image

        def make_small_villager_image(image):
            villager_image = np.zeros((image.shape[0], image.shape[1], 4))
            villager_image[8:8+8, 8:8+8, :3] = image[::2, ::2, :]
            villager_image[8:8+8, 8:8+8, 3] = (image[::2, ::2].mean(axis=-1) < 255-32) * 255
            return villager_image

        self._block_images['GoldOreVein'] = extract(block_image, 752, 448)
        self._block_images['CoalOreVein'] = extract(block_image, 704, 448)
        self._block_images['IronOreVein'] = extract(block_image, 0, 464)
        self._block_images['CobblestoneStash'] = extract(block_image, 32, 416)
        self._block_images['BeetrootCrop'] = extract(item_image, 96, 480)
        self._block_images['Chicken'] = extract(entity_image, 32, 432)
        self._block_images['Sheep'] = extract(entity_image, 144, 432)
        self._block_images['Tree'] = np.array(Image.open(osp.join(osp.dirname(__file__), 'assets', 'Jungle_Tree.png')).resize((16, 16)))
        self._block_images['SugarCanePlant'] = extract(block_image, 432, 576)
        self._block_images['PotatoPlant'] = extract(block_image, 304, 576)
        self._block_images['WorkStation'] = extract(block_image, 0, 512)
        self._block_images['Furnace'] = extract(block_image, 464, 464)

        self._block_images['Coal'] = extract(item_image, 176, 544)
        self._block_images['Coal/Small'] = make_small_item_image(self._block_images['Coal'])
        self._block_images['Feather'] = extract(item_image, 240, 544)
        self._block_images['Feather/Small'] = make_small_item_image(self._block_images['Feather'])
        self._block_images['Axe'] = extract(item_image, 160, 464)  # Stone Axe
        self._block_images['Axe/Small'] = make_small_item_image(self._block_images['Axe'])
        self._block_images['Pickaxe'] = extract(item_image, 192, 464)  # Stone Pickaxe
        self._block_images['Pickaxe/Small'] = make_small_item_image(self._block_images['Pickaxe'])
        self._block_images['SugarCane'] = extract(item_image, 176, 576)
        self._block_images['SugarCane/Small'] = make_small_item_image(self._block_images['SugarCane'])
        self._block_images['Potato'] = extract(item_image, 224, 496)
        self._block_images['Potato/Small'] = make_small_item_image(self._block_images['Potato'])
        self._block_images['CookedPotato'] = extract(item_image, 64, 480)
        self._block_images['CookedPotato/Small'] = make_small_item_image(self._block_images['CookedPotato'])
        self._block_images['Beetroot'] = extract(item_image, 96, 480)
        self._block_images['Beetroot/Small'] = make_small_item_image(self._block_images['Beetroot'])
        self._block_images['BeetrootSoup'] = extract(item_image, 112, 480)
        self._block_images['BeetrootSoup/Small'] = make_small_item_image(self._block_images['BeetrootSoup'])
        self._block_images['Bed'] = extract(block_image, 208, 528)
        self._block_images['Bed/Small'] = make_small_item_image(self._block_images['Bed'])
        self._block_images['IronOre'] = extract(item_image, 0, 880)
        self._block_images['IronOre/Small'] = make_small_item_image(self._block_images['IronOre'])
        self._block_images['IronIngot'] = extract(item_image, 112, 560)
        self._block_images['IronIngot/Small'] = make_small_item_image(self._block_images['IronIngot'])
        self._block_images['GoldOre'] = extract(item_image, 240, 864)
        self._block_images['GoldOre/Small'] = make_small_item_image(self._block_images['GoldOre'])
        self._block_images['GoldIngot'] = extract(item_image, 48, 560)
        self._block_images['GoldIngot/Small'] = make_small_item_image(self._block_images['GoldIngot'])
        self._block_images['Cobblestone'] = extract(item_image, 240, 800)  # actually Netherite Ingot
        self._block_images['Cobblestone/Small'] = make_small_item_image(self._block_images['Cobblestone'])
        self._block_images['Sword'] = extract(item_image, 0, 448)  # Iron Sword
        self._block_images['Sword/Small'] = make_small_item_image(self._block_images['Sword'])
        self._block_images['Shears'] = extract(item_image, 144, 464)
        self._block_images['Shears/Small'] = make_small_item_image(self._block_images['Shears'])
        self._block_images['Stick'] = extract(item_image, 80, 48)
        self._block_images['Stick/Small'] = make_small_item_image(self._block_images['Stick'])
        self._block_images['Boat'] = extract(item_image, 112, 592)  # Oak Boat
        self._block_images['Boat/Small'] = make_small_item_image(self._block_images['Boat'])
        self._block_images['Bowl'] = extract(item_image, 208, 528)
        self._block_images['Bowl/Small'] = make_small_item_image(self._block_images['Bowl'])
        self._block_images['Wood'] = extract(block_image, 352, 416)  # Oak Log
        self._block_images['Wood/Small'] = make_small_item_image(self._block_images['Wood'])
        self._block_images['WoodPlank'] = extract(block_image, 224, 464)  # Oak Planks
        self._block_images['WoodPlank/Small'] = make_small_item_image(self._block_images['WoodPlank'])
        self._block_images['Wool'] = extract(block_image, 0, 32)  # White Wool
        self._block_images['Wool/Small'] = make_small_item_image(self._block_images['Wool'])
        self._block_images['Arrow'] = extract(item_image, 80, 32)
        self._block_images['Arrow/Small'] = make_small_item_image(self._block_images['Arrow'])
        self._block_images['Paper'] = extract(item_image, 96, 544)
        self._block_images['Paper/Small'] = make_small_item_image(self._block_images['Paper'])

        self._block_images['agent'] = np.array(Image.open(osp.join(osp.dirname(__file__), 'assets', 'Plains_Villager_Base.png')).resize((16, 16)))
        self._block_images['agent/Small'] = make_small_villager_image(self._block_images['agent'])

    def render(self, simulator: CraftingWorldSimulator):
        map_canvas = self.basic_canvas.copy()
        inventory_canvas = self.inventory_canvas.copy()

        def draw_primitive(canvas, obj_pos, image, mask=None):
            obj_x = obj_pos % self._map_w
            obj_y = int(obj_pos / self._map_w)
            if isinstance(image, np.ndarray) and image.shape[2] == 4:
                image, mask = image[:, :, :3], image[:, :, 3]
            if mask is not None:
                canvas[obj_y * 17 + 1: (obj_y + 1) * 17, obj_x * 17 + 1: (obj_x + 1) * 17, :] = np.clip((
                    image * (mask[:, :, None] / 255) +
                    canvas[obj_y * 17 + 1: (obj_y + 1) * 17, obj_x * 17 + 1: (obj_x + 1) * 17, :] * (1 - mask[:, :, None] / 255)
                ), 0, 255).astype('uint8')
            else:
                canvas[obj_y * 17 + 1: (obj_y + 1) * 17, obj_x * 17 + 1: (obj_x + 1) * 17, :] = image

        def draw(canvas, obj_name, obj_type, obj_pos, use_small=True):
            if obj_pos > self._map_w * self._map_h:
                raise ValueError(f'Object {obj_name} is out of bounds')
            obj_pos = obj_pos - 1

            if obj_type.endswith('Station'):
                obj_type = 'WorkStation'

            if use_small and obj_type + '/Small' in self._block_images:
                image = self._block_images[obj_type + '/Small']
            elif obj_type in self._block_images:
                image = self._block_images[obj_type]
            else:
                print(f'Warning: unknown object type {obj_type}.')
                image = 0

            draw_primitive(canvas, obj_pos, image)

        # Draw objects
        location_types = get_all_locations()
        for obj_name, (obj_type, obj_pos) in simulator.objects.items():
            if obj_type in location_types:
                draw(map_canvas, obj_name, obj_type, obj_pos)
        for obj_name, (obj_type, obj_pos) in simulator.objects.items():
            if obj_type not in location_types:
                draw(map_canvas, obj_name, obj_type, obj_pos)
        for i, content in simulator.inventory.items():
            if content is not None:
                obj_type, obj_name = content
                draw(inventory_canvas, obj_name, obj_type, i, use_small=False)

        # Highlight the border of the agent position
        agent_x = (simulator.agent_pos - 1) % self._map_w
        agent_y = int((simulator.agent_pos - 1) / self._map_w)
        map_canvas[agent_y * 17: (agent_y + 1) * 17 + 1, agent_x * 17, :] = (29, 209, 77)
        map_canvas[agent_y * 17: (agent_y + 1) * 17 + 1, (agent_x + 1) * 17, :] = (29, 209, 77)
        map_canvas[agent_y * 17, agent_x * 17: (agent_x + 1) * 17 + 1, :] = (29, 209, 77)
        map_canvas[(agent_y + 1) * 17, agent_x * 17: (agent_x + 1) * 17 + 1, :] = (29, 209, 77)
        draw_primitive(map_canvas, simulator.agent_pos - 1, self._block_images['agent/Small'])

        # Concatenate the two canvases
        canvas = np.concatenate((map_canvas, np.zeros((8, map_canvas.shape[1], 3), dtype='uint8') + 255, inventory_canvas), axis=0)
        return canvas


class CraftingWorldPDDLExecutor(object):
    def __init__(self, domain: Domain, simulator: CraftingWorldSimulator):
        self.domain = domain
        self.simulator = simulator

    def execute(self, plan: Sequence[OperatorApplier]):
        last_failed_operator = None
        for i, action in enumerate(plan):
            rv = self.step(action)
            if not rv:
                last_failed_operator = i
                break

    def step(self, action: OperatorApplier) -> bool:
        action_name = action.operator.name
        action_args = action.arguments
        if action_name == "move-right":
            self.simulator.move_right()
        elif action_name == "move-left":
            self.simulator.move_left()
        elif action_name == "move-to":
            self.simulator.move_to(int(action_args[1][1:]))
        elif action_name == "pick-up":
            try:
                self.simulator.pick_up(int(_find_string_start_with(action_args, "i", first=True)[1:]), _find_string_start_with(action_args, "o", first=True))
            except KeyError as e:
                print(f'  pick-up {action_args} failed. Reason: {e}')
                return False
        elif action_name == "place-down":
            self.simulator.place_down(int(_find_string_start_with(action_args, "i", first=True)[1:]))
        elif action_name.startswith('mine'):
            # Trying mining.
            inventory_indices = [int(x[1:]) for x in _find_string_start_with(action_args, "i")]
            object_indices = _find_string_start_with(action_args, "o")

            hypothetical_object = [x for x in object_indices if x in self.simulator.hypothetical]
            if len(hypothetical_object) != 1:
                return False
            hypothetical_object = hypothetical_object[0]

            empty_inventory = [x for x in inventory_indices if self.simulator.inventory[x] is None]
            if len(empty_inventory) != 1:
                return False
            empty_inventory = empty_inventory[0]

            target_object = [
                x for x in object_indices if x in self.simulator.objects and self.simulator.objects[x][1] == self.simulator.agent_pos
            ]
            if len(target_object) != 1:
                return False
            target_object = target_object[0]

            tool_inventory = list(set(inventory_indices) - set([empty_inventory]))

            rv = self.simulator.mine(target_object, empty_inventory, hypothetical_object, tool_inventory=tool_inventory[0] if len(tool_inventory) > 0 else None)
            if not rv:
                return False
        elif action_name.startswith('craft'):
            inventory_indices = [int(x[1:]) for x in _find_string_start_with(action_args, "i")]
            object_indices = _find_string_start_with(action_args, "o")

            hypothetical_object = [x for x in object_indices if x in self.simulator.hypothetical]
            if len(hypothetical_object) != 1:
                return False
            hypothetical_object = hypothetical_object[0]

            empty_inventory = [x for x in inventory_indices if self.simulator.inventory[x] is None]
            if len(empty_inventory) != 1:
                return False
            empty_inventory = empty_inventory[0]

            target_object = [ x for x in object_indices if x in self.simulator.objects and self.simulator.objects[x][1] == self.simulator.agent_pos ]
            if len(target_object) != 1:
                return False
            target_object = target_object[0]

            ingredients = list(set(inventory_indices) - {empty_inventory})

            target_type = None
            hypothetical_object_index = action_args.index(hypothetical_object)
            hypothetical_object_varname = self.domain.operators[action_name].arguments[hypothetical_object_index].name
            for effect in self.domain.operators[action_name].effects:
                if (
                    effect.assign_expr.predicate.function.name == 'object-of-type' and
                    effect.assign_expr.predicate.arguments[0].name == hypothetical_object_varname and
                    effect.assign_expr.predicate.arguments[1].__class__.__name__ == 'ObjectConstantExpression' and
                    effect.assign_expr.value.__class__.__name__ == 'ConstantExpression' and
                    effect.assign_expr.value.constant.item() == 1
                ):
                    # print('  Found target type', effect.assign_expr)
                    target_type = effect.assign_expr.predicate.arguments[1].name
                    break

            rv = self.simulator.craft(target_object, empty_inventory, hypothetical_object, ingredients_inventory=ingredients, target_type=target_type)

            if not rv:
                return False
        else:
            return False

        return True


def _find_string_start_with(list_of_string, start, first=False):
    rv = list()
    for s in list_of_string:
        if s.startswith(start):
            if first:
                return s
            rv.append(s)
    return rv

concepts/benchmark/gridworld/crafting_world/crafting_world_env.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import numpy.random as npr
from typing import Any, Tuple, Dict

from concepts.benchmark.gridworld.crafting_world.crafting_world_rules import get_all_mining_tools, get_all_mining_locations, get_all_mining_outcomes
from concepts.benchmark.gridworld.crafting_world.crafting_world_gen.utils import underline_to_pascal, underline_to_space


PROBLEM_PDDL_TEMPLATE = """
(define
 (problem crafting-world-v20230106-p{problem_id})
 (:domain crafting-world-v20230106)
 (:objects
   {objects_str}
 )
 (:init
   {init_str}
 )
 (:goal {goal_str} )
)
"""


def gen_linear_tile(n: int) -> Tuple[str, str]:
    """Generate the object string and init string for a linear 1D map.

    .. code::

        t1 t2 t3 t4 ... tn

    Args:
        n: the length of the map.

    Returns:
        a tuple of the object string and init string.
    """

    object_str = ""
    init_str = ""
    for i in range(1, n + 1):
        object_str += f"t{i} - tile\n"
        # if i > 1:
        #     init_str += f"(tile-right t{i - 1} t{i})\n"
        #     init_str += f"(tile-left t{i} t{i - 1})\n"
    return object_str, init_str


def gen_locations_and_objects(n: int, inventory_size: int) -> Tuple[str, str, int, list, list]:
    """Sample a linear map by randomly placing mining locations and tools on the map."""

    mining_locations = get_all_mining_locations()
    mining_tools = get_all_mining_tools()

    nr_objects = len(mining_locations) + len(mining_tools) + inventory_size

    map_location = [None for _ in range(n)]
    map_objects = [[] for _ in range(n)]

    assert len(mining_locations) < n, 'Map is too small to fit all mining locations: {} < {}'.format(n, len(mining_locations))

    object_str = ''
    init_str = ''

    mining_location_assignments = npr.choice(np.arange(n), len(mining_locations), replace=False)
    for i, x in enumerate(mining_location_assignments):
        map_location[x] = (i + 1, mining_locations[i])
        object_str += f'o{i + 1} - object\n'
        init_str += '(object-of-type o{} {})\n'.format(i + 1, underline_to_pascal(mining_locations[i]))
        init_str += '(object-at o{} t{})\n'.format(i + 1, x + 1)

    tool_location_assignments = npr.choice(np.arange(n), len(mining_tools), replace=True)
    for i, x in enumerate(tool_location_assignments):
        map_objects[x].append((i + 1 + len(mining_locations), mining_tools[i]))
        object_str += f'o{i + 1 + len(mining_locations)} - object\n'
        init_str += '(object-of-type o{} {})\n'.format(i + 1 + len(mining_locations), underline_to_pascal(mining_tools[i]))
        init_str += '(object-at o{} t{})\n'.format(i + 1 + len(mining_locations), x + 1)

    for i in range(inventory_size):
        object_str += f'o{i + 1 + len(mining_locations) + len(mining_tools)} - object\n'
        object_str += f'i{i + 1} - inventory\n'
        init_str += '(object-of-type o{} Hypothetical)\n'.format(i + 1 + len(mining_locations) + len(mining_tools))
        init_str += '(inventory-empty i{})\n'.format(i + 1)

    return object_str, init_str, nr_objects, map_location, map_objects


def gen_v20230204_instance_record(problem_id: str, split: str, n: int = 15, inventory_size: int = 3) -> Dict[str, Any]:
    import jacinle

    mining_outcomes = get_all_mining_outcomes()
    goal = npr.choice(mining_outcomes)

    obj_object_str, obj_init_str, nr_objects, map_location, map_objects = gen_locations_and_objects(n, inventory_size)
    map_object_str, map_init_str = gen_linear_tile(n)
    object_str = obj_object_str + map_object_str
    init_str = obj_init_str + map_init_str
    init_str += '(agent-at t1)\n'

    goal_nl = 'Mine {} from the map.'.format(underline_to_space(goal))
    # goal_str = f'(and (inventory-holding i{inventory_size} o{nr_objects})\n(object-of-type o{nr_objects} {underline_to_pascal(goal)}))\n'
    goal_str = f'(exists (?i - inventory) (exists (?o - object) (and (inventory-holding ?i ?o) (object-of-type ?o {underline_to_pascal(goal)})) ))'

    problem_pddl = PROBLEM_PDDL_TEMPLATE.format(
        problem_id=problem_id,
        objects_str=jacinle.indent_text(object_str, indent_format='   ').strip(),
        init_str=jacinle.indent_text(init_str, indent_format='   ').strip(),
        goal_str=jacinle.indent_text(goal_str, indent_format='   ').strip(),
    )

    return dict(
        problem_id=problem_id,
        split=split,
        problem_pddl=problem_pddl,
        goal_nl=goal_nl,
        goal=goal,
        map_location=map_location,
        map_objects=map_objects,
        nr_objects=nr_objects,
        inventory_size=inventory_size,
    )


concepts/benchmark/gridworld/crafting_world/crafting_world_gen/cw_20230204_minining_only.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import numpy.random as npr
from typing import Any, Tuple, Dict

from concepts.benchmark.gridworld.crafting_world.crafting_world_rules import get_all_crafting_ingradients, get_all_crafting_locations, get_all_crafting_outcomes
from concepts.benchmark.gridworld.crafting_world.crafting_world_gen.utils import underline_to_pascal, underline_to_space


PROBLEM_PDDL_TEMPLATE = """
(define
 (problem crafting-world-v20230106-p{problem_id})
 (:domain crafting-world-v20230106)
 (:objects
   {objects_str}
 )
 (:init
   {init_str}
 )
 (:goal {goal_str} )
)
"""


def gen_linear_tile(n: int) -> Tuple[str, str]:
    """Generate the object string and init string for a linear 1D map.

    .. code::

        t1 t2 t3 t4 ... tn

    Args:
        n: the length of the map.

    Returns:
        a tuple of the object string and init string.
    """

    object_str = ""
    init_str = ""
    for i in range(1, n + 1):
        object_str += f"t{i} - tile\n"
        # if i > 1:
        #     init_str += f"(tile-right t{i - 1} t{i})\n"
        #     init_str += f"(tile-left t{i} t{i - 1})\n"
    return object_str, init_str


def gen_locations_and_objects(n: int, inventory_size: int) -> Tuple[str, str, int, list, list]:
    """Sample a linear map by randomly placing mining locations and tools on the map."""

    crafting_locations = get_all_crafting_locations()
    crafting_ingredients = get_all_crafting_ingradients()

    nr_objects = len(crafting_locations) + len(crafting_ingredients) + inventory_size

    map_location = [None for _ in range(n)]
    map_objects = [[] for _ in range(n)]

    assert len(crafting_locations) < n, 'Map is too small to fit all crafting locations: {} < {}'.format(n, len(crafting_locations))

    object_str = ''
    init_str = ''

    crafting_location_assignments = npr.choice(np.arange(n), len(crafting_locations), replace=False)
    for i, x in enumerate(crafting_location_assignments):
        map_location[x] = (i + 1, crafting_locations[i])
        object_str += f'o{i + 1} - object\n'
        init_str += '(object-of-type o{} {})\n'.format(i + 1, underline_to_pascal(crafting_locations[i]))
        init_str += '(object-at o{} t{})\n'.format(i + 1, x + 1)

    crafting_ingradients_assignments = npr.choice(np.arange(n), len(crafting_ingredients), replace=True)
    for i, x in enumerate(crafting_ingradients_assignments):
        map_objects[x].append((i + 1 + len(crafting_locations), crafting_ingredients[i]))
        object_str += f'o{i + 1 + len(crafting_locations)} - object\n'
        init_str += '(object-of-type o{} {})\n'.format(i + 1 + len(crafting_locations), underline_to_pascal(crafting_ingredients[i]))
        init_str += '(object-at o{} t{})\n'.format(i + 1 + len(crafting_locations), x + 1)

    for i in range(inventory_size):
        object_str += f'o{i + 1 + len(crafting_locations) + len(crafting_ingredients)} - object\n'
        object_str += f'i{i + 1} - inventory\n'
        init_str += '(object-of-type o{} Hypothetical)\n'.format(i + 1 + len(crafting_locations) + len(crafting_ingredients))
        init_str += '(inventory-empty i{})\n'.format(i + 1)

    return object_str, init_str, nr_objects, map_location, map_objects


def gen_v20230829_instance_record(problem_id: str, split: str, n: int = 15, inventory_size: int = 3) -> Dict[str, Any]:
    import jacinle

    crafting_outcomes = get_all_crafting_outcomes()
    goal = npr.choice(crafting_outcomes)

    obj_object_str, obj_init_str, nr_objects, map_location, map_objects = gen_locations_and_objects(n, inventory_size)
    map_object_str, map_init_str = gen_linear_tile(n)
    object_str = obj_object_str + map_object_str
    init_str = obj_init_str + map_init_str
    init_str += '(agent-at t1)\n'

    goal_nl = 'Craft a {}.'.format(underline_to_space(goal))
    # goal_str = f'(and (inventory-holding i{inventory_size} o{nr_objects})\n(object-of-type o{nr_objects} {underline_to_pascal(goal)}))\n'
    goal_str = f'(exists (?i - inventory) (exists (?o - object) (and (inventory-holding ?i ?o) (object-of-type ?o {underline_to_pascal(goal)})) ))'

    problem_pddl = PROBLEM_PDDL_TEMPLATE.format(
        problem_id=problem_id,
        objects_str=jacinle.indent_text(object_str, indent_format='   ').strip(),
        init_str=jacinle.indent_text(init_str, indent_format='   ').strip(),
        goal_str=jacinle.indent_text(goal_str, indent_format='   ').strip(),
    )

    return dict(
        problem_id=problem_id,
        split=split,
        problem_pddl=problem_pddl,
        goal_nl=goal_nl,
        goal=goal,
        map_location=map_location,
        map_objects=map_objects,
        nr_objects=nr_objects,
        inventory_size=inventory_size,
    )


concepts/benchmark/gridworld/crafting_world/crafting_world_gen/cw_20230829_crafting_only.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import numpy.random as npr
from typing import Any, Optional, Tuple, Dict

from concepts.benchmark.gridworld.crafting_world.crafting_world_rules import get_all_mining_locations, get_all_mining_outcomes
from concepts.benchmark.gridworld.crafting_world.crafting_world_rules import get_all_crafting_ingradients, get_all_crafting_locations, get_all_crafting_outcomes, CRAFTING_RULES
from concepts.benchmark.gridworld.crafting_world.crafting_world_gen.utils import underline_to_pascal, underline_to_space


PROBLEM_PDDL_TEMPLATE = """
(define
 (problem crafting-world-v20230106-p{problem_id})
 (:domain crafting-world-v20230106)
 (:objects
   {objects_str}
 )
 (:init
   {init_str}
 )
 (:goal {goal_str} )
)
"""


def gen_linear_tile(n: int) -> Tuple[str, str]:
    """Generate the object string and init string for a linear 1D map.

    .. code::

        t1 t2 t3 t4 ... tn

    Args:
        n: the length of the map.

    Returns:
        a tuple of the object string and init string.
    """

    object_str = ""
    init_str = ""
    for i in range(1, n + 1):
        object_str += f"t{i} - tile\n"
        # if i > 1:
        #     init_str += f"(tile-right t{i - 1} t{i})\n"
        #     init_str += f"(tile-left t{i} t{i - 1})\n"
    return object_str, init_str


def gen_locations_and_objects(n: int, inventory_size: int, station_agnostic: bool = True) -> Tuple[str, str, int, list, list]:
    """Sample a linear map by randomly placing mining locations and tools on the map."""

    all_locations = get_all_mining_locations() + get_all_crafting_locations(use_only_workstation=station_agnostic)
    mining_tools = ['axe', 'pickaxe']  #  + ['coal'] * 5

    nr_objects = len(all_locations) + len(mining_tools) + inventory_size

    map_location = [None for _ in range(n)]
    map_objects = [[] for _ in range(n)]

    assert len(all_locations) < n, 'Map is too small to fit all locations: {} < {}'.format(n, len(all_locations))

    object_str = ''
    init_str = ''

    location_assignments = npr.choice(np.arange(n), len(all_locations), replace=False)
    for i, x in enumerate(location_assignments):
        map_location[x] = (i + 1, all_locations[i])
        object_str += f'o{i + 1} - object\n'
        init_str += '(object-of-type o{} {})\n'.format(i + 1, underline_to_pascal(all_locations[i]))
        init_str += '(object-at o{} t{})\n'.format(i + 1, x + 1)

    mining_tool_assignments = npr.choice(np.arange(n), len(mining_tools), replace=False)
    for i, x in enumerate(mining_tool_assignments):
        map_objects[x].append((i + 1 + len(all_locations), mining_tools[i]))
        object_str += f'o{i + 1 + len(all_locations)} - object\n'
        init_str += '(object-of-type o{} {})\n'.format(i + 1 + len(all_locations), underline_to_pascal(mining_tools[i]))
        init_str += '(object-at o{} t{})\n'.format(i + 1 + len(all_locations), x + 1)

    for i in range(inventory_size):
        object_str += f'o{i + 1 + len(all_locations) + len(mining_tools)} - object\n'
        object_str += f'i{i + 1} - inventory\n'
        init_str += '(object-of-type o{} Hypothetical)\n'.format(i + 1 + len(all_locations) + len(mining_tools))
        init_str += '(inventory-empty i{})\n'.format(i + 1)

    return object_str, init_str, nr_objects, map_location, map_objects


def gen_v20230913_instance_record(problem_id: str, split: str, n: int = 25, inventory_size: int = 8, goal_index: Optional[int] = None, station_agnostic: bool = True) -> Dict[str, Any]:
    import jacinle

    all_outcomes = get_all_crafting_outcomes() + get_all_mining_outcomes()
    if goal_index is not None:
        assert goal_index < len(all_outcomes), 'Invalid goal index: {} >= {}'.format(goal_index, len(all_outcomes))
        goal = all_outcomes[goal_index]
    else:
        goal = npr.choice(all_outcomes)

    obj_object_str, obj_init_str, nr_objects, map_location, map_objects = gen_locations_and_objects(n, inventory_size)
    map_object_str, map_init_str = gen_linear_tile(n)
    object_str = obj_object_str + map_object_str
    init_str = obj_init_str + map_init_str
    init_str += '(agent-at t1)\n'

    goal_nl = 'Mine or craft a {}.'.format(underline_to_space(goal))
    # goal_str = f'(and (inventory-holding i{inventory_size} o{nr_objects})\n(object-of-type o{nr_objects} {underline_to_pascal(goal)}))\n'
    goal_str = f'(exists (?i - inventory) (exists (?o - object) (and (inventory-holding ?i ?o) (object-of-type ?o {underline_to_pascal(goal)})) ))'

    problem_pddl = PROBLEM_PDDL_TEMPLATE.format(
        problem_id=problem_id,
        objects_str=jacinle.indent_text(object_str, indent_format='   ').strip(),
        init_str=jacinle.indent_text(init_str, indent_format='   ').strip(),
        goal_str=jacinle.indent_text(goal_str, indent_format='   ').strip(),
    )

    return dict(
        problem_id=problem_id,
        split=split,
        problem_pddl=problem_pddl,
        goal_nl=goal_nl,
        goal=goal,
        map_location=map_location,
        map_objects=map_objects,
        nr_objects=nr_objects,
        inventory_size=inventory_size,
    )

concepts/benchmark/gridworld/crafting_world/crafting_world_gen/cw_20230913_mixed.py
concepts-ai/Concepts
def underline_to_pascal(s: str) -> str:
    """Converts a string from underlined to pascal case.

    Example:
        >>> underline_to_pascal("hello_world")
        "HelloWorld"

    Args:
        s: the string to convert.

    Returns:
        the converted string.
    """
    return ''.join([w.capitalize() for w in s.split('_')])


def underline_to_space(s: str) -> str:
    """Converts a string from underlined to a space-separated and lower case string.

    Example:
        >>> underline_to_space("Hello_world")
        "hello world"
    """
    return ' '.join([w.lower() for w in s.split('_')])

concepts/benchmark/gridworld/crafting_world/crafting_world_gen/utils.py
concepts-ai/Concepts
from functools import lru_cache
from typing import List

MINING_RULES = [
    dict(
        rule_name='mine_iron_ore',
        create='iron_ore',
        action='mine',
        location='iron_ore_vein',
        recipe=[],
        holding=['pickaxe'],
    ),
    dict(
        rule_name='mine_coal',
        create='coal',
        action='mine',
        location='coal_ore_vein',
        recipe=[],
        holding=['pickaxe'],
    ),
    dict(
        rule_name='mine_cobblestone',
        create='cobblestone',
        action='mine',
        location='cobblestone_stash',
        recipe=[],
        holding=['pickaxe'],
    ),
    dict(
        rule_name='mine_wood',
        create='wood',
        action='mine',
        location='tree',
        recipe=[],
        holding=['axe'],
    ),
    dict(
        rule_name='mine_feather',
        create='feather',
        action='mine',
        location='chicken',
        recipe=[],
        holding=['sword'],
    ),
    dict(
        rule_name='mine_wool1',
        create='wool',
        action='mine',
        location='sheep',
        recipe=[],
        holding=['shears'],
    ),
    dict(
        rule_name='mine_wool2',
        create='wool',
        action='mine',
        location='sheep',
        recipe=[],
        holding=['sword'],
    ),
    dict(
        rule_name='mine_potato',
        create='potato',
        action='mine',
        location='potato_plant',
        recipe=[],
        holding=[],
    ),
    dict(
        rule_name='mine_beetroot',
        create='beetroot',
        action='mine',
        location='beetroot_crop',
        recipe=[],
        holding=[],
    ),
    dict(
        rule_name='mine_gold_ore',
        create='gold_ore',
        action='mine',
        location='gold_ore_vein',
        recipe=[],
        holding=['pickaxe'],
    ),
    dict(
        rule_name='mine_sugar_cane',
        create='sugar_cane',
        action='mine',
        location='sugar_cane_plant',
        recipe=[],
        holding=[],
    ),
]

CRAFTING_RULES = [
    dict(
        rule_name='craft_wood_plank',
        create='wood_plank',
        action='craft',
        location='work_station',
        recipe=['wood'],
        holding=[],
    ),
    dict(
        rule_name='craft_stick',
        create='stick',
        action='craft',
        location='work_station',
        recipe=['wood_plank'],
        holding=[],
    ),
    dict(
        rule_name='craft_arrow',
        create='arrow',
        action='craft',
        location='weapon_station',
        recipe=['feather', 'stick'],
        holding=[],
    ),
    dict(
        rule_name='craft_sword',
        create='sword',
        action='craft',
        location='weapon_station',
        recipe=['iron_ingot', 'stick'],
        holding=[],
    ),
    dict(
        rule_name='craft_shears1',
        create='shears',
        action='craft',
        location='tool_station',
        recipe=['iron_ingot'],
        holding=[],
    ),
    dict(
        rule_name='craft_shears2',
        create='shears',
        action='craft',
        location='tool_station',
        recipe=['gold_ingot'],
        holding=[],
    ),
    dict(
        rule_name='craft_iron_ingot',
        create='iron_ingot',
        action='craft',
        location='furnace',
        recipe=['iron_ore', 'coal'],
        holding=[],
    ),
    dict(
        rule_name='craft_gold_ingot',
        create='gold_ingot',
        action='craft',
        location='furnace',
        recipe=['gold_ore', 'coal'],
        holding=[],
    ),
    dict(
        rule_name='craft_bed',
        create='bed',
        action='craft',
        location='bed_station',
        recipe=['wood_plank', 'wool'],
        holding=[],
    ),
    dict(
        rule_name='craft_boat',
        create='boat',
        action='craft',
        location='boat_station',
        recipe=['wood_plank'],
        holding=[],
    ),
    dict(
        rule_name='craft_bowl1',
        create='bowl',
        action='craft',
        location='food_station',
        recipe=['wood_plank'],
        holding=[],
    ),
    dict(
        rule_name='craft_bowl2',
        create='bowl',
        action='craft',
        location='food_station',
        recipe=['iron_ingot'],
        holding=[],
    ),
    dict(
        rule_name='craft_cooked_potato',
        create='cooked_potato',
        action='craft',
        location='furnace',
        recipe=['potato', 'coal'],
        holding=[],
    ),
    dict(
        rule_name='craft_beetroot_soup',
        create='beetroot_soup',
        action='craft',
        location='food_station',
        recipe=['beetroot', 'bowl'],
        holding=[],
    ),
    dict(
        rule_name='craft_paper',
        create='paper',
        action='craft',
        location='work_station',
        recipe=['sugar_cane'],
        holding=[],
    )
]


@lru_cache()
def get_all_mining_tools() -> List[str]:
    """Return a list of tools that can be used in mining actions."""

    tools = list()
    for rule in MINING_RULES:
        for holding in rule['holding']:
            if holding not in tools:
                tools.append(holding)
    return tools


@lru_cache()
def get_all_mining_locations() -> List[str]:
    """Return a list of locations that can be mined."""

    locations = list()
    for rule in MINING_RULES:
        if rule['location'] not in locations:
            locations.append(rule['location'])
    return locations


@lru_cache()
def get_all_mining_outcomes() -> List[str]:
    """Return a list of outcomes that can be mined."""

    outcomes = list()
    for rule in MINING_RULES:
        if rule['create'] not in outcomes:
            outcomes.append(rule['create'])
    return outcomes


@lru_cache()
def get_all_crafting_ingradients() -> List[str]:
    """Return a list of tools that can be used in crafting actions."""

    tools = list()
    for rule in CRAFTING_RULES:
        for holding in rule['recipe']:
            if holding not in tools:
                tools.append(holding)
    return tools

@lru_cache()
def get_all_crafting_locations(use_only_workstation: bool = False) -> List[str]:
    """Return a list of locations that can be crafted."""

    if use_only_workstation:
        return ['work_station']

    locations = list()
    for rule in CRAFTING_RULES:
        if rule['location'] not in locations:
            locations.append(rule['location'])
    return locations


@lru_cache()
def get_all_crafting_outcomes() -> List[str]:
    """Return a list of outcomes that can be crafted."""

    outcomes = list()
    for rule in CRAFTING_RULES:
        if rule['create'] not in outcomes:
            outcomes.append(rule['create'])
    return outcomes


@lru_cache()
def get_all_locations() -> List[str]:
    """Return a list of all locations."""
    return get_all_mining_locations() + get_all_crafting_locations()
concepts/benchmark/gridworld/crafting_world/crafting_world_rules.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/06/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import os.path as osp

__all__ = ['get_domain_filename']


def get_domain_filename():
    return osp.join(osp.dirname(__file__), 'domain.pddl')


def get_station_agnostic_domain_filename():
    return osp.join(osp.dirname(__file__), 'domain-station-agnostic.pddl')

concepts/benchmark/gridworld/crafting_world/crafting_world_teleport/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : domain-gen.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/06/2023
#
# This file is part of HACL-PyTorch.
# Distributed under terms of the MIT license.

import concepts.benchmark.gridworld.crafting_world.crafting_world_rules as rules


def underline_to_pascal(s):
    return ''.join([w.capitalize() for w in s.split('_')])


mining_template_0 = """
 (:action {action_name}
  :parameters (?targetinv - inventory ?x - object ?target - object ?t - tile)
  :precondition (and
    (agent-at ?t)
    (object-at ?x ?t)
    (object-of-type ?x {target_type})
    (inventory-empty ?targetinv)
    (object-of-type ?target Hypothetical)
  )
  :effect (and
    (not (inventory-empty ?targetinv))
    (inventory-holding ?targetinv ?target)
    (not (object-of-type ?target Hypothetical))
    (object-of-type ?target {create_type})
  )
 )
 (:regression {action_name}-1 [always]
  :parameters ((forall ?target-inventory - inventory) (forall ?target - object) (forall ?target-resource - object) (forall ?t - tile))
  :goal (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {create_type}))))
  :precondition (and (object-at ?target-resource ?t) (object-of-type ?target-resource {target_type}) (inventory-empty ?target-inventory) (object-of-type ?target Hypothetical))
  :rule (then
    (achieve (agent-at ?t))
    ({action_name} ?target-inventory ?target-resource ?target ?t)
  )
 )"""


mining_template_1 = """
 (:action {action_name}
  :parameters (?toolinv - inventory ?targetinv - inventory ?x - object ?tool - object ?target - object ?t - tile)
  :precondition (and
    (agent-at ?t)
    (object-at ?x ?t)
    (object-of-type ?x {target_type})
    (inventory-holding ?toolinv ?tool)
    (object-of-type ?tool {holding})
    (inventory-empty ?targetinv)
    (object-of-type ?target Hypothetical)
  )
  :effect (and
    (not (inventory-empty ?targetinv))
    (inventory-holding ?targetinv ?target)
    (not (object-of-type ?target Hypothetical))
    (object-of-type ?target {create_type})
  )
 )
 (:regression {action_name}-1 [always]
  :parameters ((forall ?target-inventory - inventory) (forall ?target - object) (forall ?target-resource - object) (forall ?t - tile) (forall ?holding - object) (forall ?holding-inventory - inventory))
  :goal (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {create_type}))))
  :precondition (and
    (object-at ?target-resource ?t) (object-of-type ?target-resource {target_type}) (inventory-empty ?target-inventory) (object-of-type ?target Hypothetical)
    (inventory-holding ?holding-inventory ?holding) (object-of-type ?holding {holding})
  )
  :rule (then
    (achieve (agent-at ?t))
    ({action_name} ?holding-inventory ?target-inventory ?target-resource ?holding ?target ?t)
  )
 )
 (:regression {action_name}-2 [always]
  :goal (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {create_type}))))
  :rule (then
    (achieve (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {holding})))))
    (achieve (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {create_type})))))
  )
 )
 """

crafting_template_1 = """
 (:action {action_name}
  :parameters (?ingredientinv1 - inventory ?targetinv - inventory ?station - object ?ingredient1 - object ?target - object ?t - tile)
  :precondition (and
    (agent-at ?t)
    (object-at ?station ?t)
    (object-of-type ?station {station_type})
    (inventory-holding ?ingredientinv1 ?ingredient1)
    (object-of-type ?ingredient1 {ingredient1_type})
    (inventory-empty ?targetinv)
    (object-of-type ?target Hypothetical)
  )
  :effect (and
    (not (inventory-empty ?targetinv))
    (inventory-holding ?targetinv ?target)
    (not (object-of-type ?target Hypothetical))
    (object-of-type ?target {create_type})
    (not (inventory-holding ?ingredientinv1 ?ingredient1))
    (inventory-empty ?ingredientinv1)
    (not (object-of-type ?ingredient1 {ingredient1_type}))
    (object-of-type ?ingredient1 Hypothetical)
  )
 )
 (:regression {action_name}-1 [always]
  :parameters ((forall ?target-inventory - inventory) (forall ?target - object) (forall ?target-resource - object) (forall ?t - tile) (forall ?ingredient1 - object) (forall ?ingredient1-inventory - inventory))
  :goal (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {create_type}))))
  :precondition (and
    (object-at ?target-resource ?t) (object-of-type ?target-resource {station_type}) (inventory-empty ?target-inventory) (object-of-type ?target Hypothetical)
    (inventory-holding ?ingredient1-inventory ?ingredient1) (object-of-type ?ingredient1 {ingredient1_type})
  )
  :rule (then
    (achieve (agent-at ?t))
    ({action_name} ?ingredient1-inventory ?target-inventory ?target-resource ?ingredient1 ?target ?t)
  )
 )
 (:regression {action_name}-2 [always]
  :goal (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {create_type}))))
  :rule (then
    (achieve (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {ingredient1_type})))))
    (achieve (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {create_type})))))
  )
 )
 """

crafting_template_2 = """
 (:action {action_name}
  :parameters (?ingredientinv1 - inventory ?ingredientinv2 - inventory ?targetinv - inventory ?station - object ?ingredient1 - object ?ingredient2 - object ?target - object ?t - tile)
  :precondition (and
    (agent-at ?t)
    (object-at ?station ?t)
    (object-of-type ?station {station_type})
    (inventory-holding ?ingredientinv1 ?ingredient1)
    (object-of-type ?ingredient1 {ingredient1_type})
    (inventory-holding ?ingredientinv2 ?ingredient2)
    (object-of-type ?ingredient2 {ingredient2_type})
    (inventory-empty ?targetinv)
    (object-of-type ?target Hypothetical)
  )
  :effect (and
    (not (inventory-empty ?targetinv))
    (inventory-holding ?targetinv ?target)
    (not (object-of-type ?target Hypothetical))
    (object-of-type ?target {create_type})
    (not (inventory-holding ?ingredientinv1 ?ingredient1))
    (inventory-empty ?ingredientinv1)
    (not (object-of-type ?ingredient1 {ingredient1_type}))
    (object-of-type ?ingredient1 Hypothetical)
    (not (inventory-holding ?ingredientinv2 ?ingredient2))
    (inventory-empty ?ingredientinv2)
    (not (object-of-type ?ingredient2 {ingredient2_type}))
    (object-of-type ?ingredient2 Hypothetical)
  )
 )
 (:regression {action_name}-1 [always]
  :parameters ((forall ?target-inventory - inventory) (forall ?target - object) (forall ?target-resource - object) (forall ?t - tile) (forall ?ingredient1 - object) (forall ?ingredient1-inventory - inventory) (forall ?ingredient2 - object) (forall ?ingredient2-inventory - inventory))
  :goal (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {create_type}))))
  :precondition (and
    (object-at ?target-resource ?t) (object-of-type ?target-resource {station_type}) (inventory-empty ?target-inventory) (object-of-type ?target Hypothetical)
    (inventory-holding ?ingredient1-inventory ?ingredient1) (object-of-type ?ingredient1 {ingredient1_type})
    (inventory-holding ?ingredient2-inventory ?ingredient2) (object-of-type ?ingredient2 {ingredient2_type})
  )
  :rule (then
    (achieve (agent-at ?t))
    ({action_name} ?ingredient1-inventory ?ingredient2-inventory ?target-inventory ?target-resource ?ingredient1 ?ingredient2 ?target ?t)
  )
 )
 (:regression {action_name}-2 [always]
  :goal (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {create_type}))))
  :rule (then
    (achieve (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {ingredient1_type})))))
    (achieve (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {ingredient2_type})))))
    (achieve (exists (?i - inventory) (exists (?x - object) (and (inventory-holding ?i ?x) (object-of-type ?x {create_type})))))
  )
 )
 """


def main():
    mining_rules = ''
    for r in rules.MINING_RULES:
        action_name = r['rule_name'].replace('_', '-')
        create_type = underline_to_pascal(r['create'])
        target_type = underline_to_pascal(r['location'])
        holding = r['holding'][0] if len(r['holding']) == 1 else None
        if holding is not None:
            holding = underline_to_pascal(holding)

        if holding is None:
            mining_rules += mining_template_0.format(action_name=action_name, create_type=create_type, target_type=target_type)
        else:
            mining_rules += mining_template_1.format(action_name=action_name, create_type=create_type, target_type=target_type, holding=holding)

    crafting_rules = ''
    for r in rules.CRAFTING_RULES:
        action_name = r['rule_name'].replace('_', '-')
        create_type = underline_to_pascal(r['create'])
        station_type = underline_to_pascal(r['location'])
        recipe = list(map(underline_to_pascal, r['recipe']))

        if len(recipe) == 1:
            ingredient1_type = recipe[0]
            crafting_rules += crafting_template_1.format(action_name=action_name, create_type=create_type, station_type=station_type, ingredient1_type=ingredient1_type)
        elif len(recipe) == 2:
            ingredient1_type, ingredient2_type = recipe
            crafting_rules += crafting_template_2.format(action_name=action_name, create_type=create_type, station_type=station_type, ingredient1_type=ingredient1_type, ingredient2_type=ingredient2_type)
        else:
            raise ValueError('Invalid recipe length: {}'.format(len(recipe)))

    with open('./domain.pddl-template') as f:
        template = f.read()
    with open('./domain.pddl', 'w') as f:
        f.write(template.format(mining_rules=mining_rules, crafting_rules=crafting_rules))
    print('Generated: domain.pddl')


def main_station_agnostic():
    mining_rules = ''
    for r in rules.MINING_RULES:
        action_name = r['rule_name'].replace('_', '-')
        create_type = underline_to_pascal(r['create'])
        target_type = underline_to_pascal(r['location'])
        holding = r['holding'][0] if len(r['holding']) == 1 else None
        if holding is not None:
            holding = underline_to_pascal(holding)

        if holding is None:
            mining_rules += mining_template_0.format(action_name=action_name, create_type=create_type, target_type=target_type)
        else:
            mining_rules += mining_template_1.format(action_name=action_name, create_type=create_type, target_type=target_type, holding=holding)

    crafting_rules = ''
    for r in rules.CRAFTING_RULES:
        action_name = r['rule_name'].replace('_', '-')
        create_type = underline_to_pascal(r['create'])
        station_type = 'WorkStation'
        recipe = list(map(underline_to_pascal, r['recipe']))

        if len(recipe) == 1:
            ingredient1_type = recipe[0]
            crafting_rules += crafting_template_1.format(action_name=action_name, create_type=create_type, station_type=station_type, ingredient1_type=ingredient1_type)
        elif len(recipe) == 2:
            ingredient1_type, ingredient2_type = recipe
            crafting_rules += crafting_template_2.format(action_name=action_name, create_type=create_type, station_type=station_type, ingredient1_type=ingredient1_type, ingredient2_type=ingredient2_type)
        else:
            raise ValueError('Invalid recipe length: {}'.format(len(recipe)))

    with open('./domain.pddl-template') as f:
        template = f.read()
    with open('./domain-station-agnostic.pddl', 'w') as f:
        f.write(template.format(mining_rules=mining_rules, crafting_rules=crafting_rules))
    print('Generated: domain.pddl')



if __name__ == '__main__':
    main()
    main_station_agnostic()


concepts/benchmark/gridworld/crafting_world/crafting_world_teleport/domain-gen.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/26/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/benchmark/humanmotion/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : dataset.py
# Author : Joy Hsu
# Email  : joycj@stanford.edu
# Date   : 03/23/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import os.path as osp

import jacinle.io as io

from jacinle.logging import get_logger
from jacinle.utils.container import GView
from jactorch.data.dataset import FilterableDatasetUnwrapped, FilterableDatasetView

from concepts.benchmark.humanmotion.utils import nsclseq_to_nscltree, nsclseq_to_nsclqsseq, nscltree_to_nsclqstree, program_to_nsclseq

import numpy as np
import math



logger = get_logger(__file__)

__all__ = ['NSTrajDataset', 'MotionClassificationDataset']


class NSTrajDatasetUnwrapped(FilterableDatasetUnwrapped):
    def __init__(self, data_dir, data_split_file, split, data_source, no_gt_segments, filter_supervision, max_frames=150):
        super().__init__()

        self.labels_json = osp.join(data_dir, 'motion_concepts.json')
        self.questions_json = osp.join(data_dir, 'questions.json')
        self.joints_root = osp.join(data_dir, 'motion_sequences')

        self.labels = io.load_json(self.labels_json)
        self.questions = io.load_json(self.questions_json)
        self.split_question_ids = io.load_json(data_split_file)[split]

        self.data_source = data_source

        self.max_frames = max_frames
        self.no_gt_segments = no_gt_segments
        self.filter_supervision = filter_supervision

    def _get_questions(self):
        return self.questions

    def _get_metainfo(self, index):
        question = self.questions[self.split_question_ids[index]]

        # program section
        has_program = False
        if 'program_nsclseq' in question:
            question['program_raw'] = question['program_nsclseq']
            question['program_seq'] = question['program_nsclseq']
            has_program = True
        elif 'program' in question:
            question['program_raw'] = question['program']
            question['program_seq'] = program_to_nsclseq(question['program'])
            has_program = True

        if has_program:
            question['program_tree'] = nsclseq_to_nscltree(question['program_seq'])
            question['program_qsseq'] = nsclseq_to_nsclqsseq(question['program_seq'])
            question['program_qstree'] = nscltree_to_nsclqstree(question['program_tree'])

        return question

    def __getitem__(self, index):
        metainfo = GView(self.get_metainfo(index))
        feed_dict = GView()

        question = self.questions[self.split_question_ids[index]]

        if 'program_raw' in metainfo:
            feed_dict.program_raw = metainfo.program_raw
            feed_dict.program_seq = metainfo.program_seq
            feed_dict.program_tree = metainfo.program_tree
            feed_dict.program_qsseq = metainfo.program_qsseq
            feed_dict.program_qstree = metainfo.program_qstree

        if '/' in question['answer']:
            question['answer'] = question['answer'].split('/')[0]
        feed_dict.answer = question['answer']
        feed_dict.question_type = question['query_type']
        feed_dict.segment_boundaries = []
        feed_dict.question_text = question['question']

        # process joints
        if self.data_source == 'teach_synth':
            id_name = 'teach_synth_id'
            motion_id = question[id_name]
            num_segments = len(self.labels[motion_id]['labels'])
        else:
            id_name = 'babel_id'
            motion_id = question[id_name]
            num_segments = len(self.labels[motion_id])
        feed_dict.babel_id = motion_id

        joints = np.load(osp.join(self.joints_root, motion_id, 'joints.npy')) # T, V, C
        # change shape of joints to match model
        joints = joints[:, :, :, np.newaxis] # T, V, C, M
        joints = joints.transpose(2, 0, 1, 3) # C, T, V, M

        # label info
        if self.data_source == 'teach_synth':
            labels_frame_info = self.labels[motion_id]['labels']
        else:
            labels_frame_info = self.labels[motion_id]

        if 'filter_answer_0' in question:
            filter_segment = labels_frame_info[question['filter_answer_0']]
            if filter_segment['end_f'] > np.shape(joints)[1]: # right now end frame can be slightly off (dataset issue)
                filter_segment['end_f'] = np.shape(joints)[1]
            feed_dict.filter_boundaries = [(filter_segment['start_f'], filter_segment['end_f'])]
            if 'filter_answer_1' in question:
                filter_segment = labels_frame_info[question['filter_answer_1']]
                if filter_segment['end_f'] > np.shape(joints)[1]: # right now end frame can be slightly off (dataset issue)
                    filter_segment['end_f'] = np.shape(joints)[1]
                feed_dict.filter_boundaries.append((filter_segment['start_f'], filter_segment['end_f']))

        if not self.no_gt_segments:
            joints_combined = np.zeros((num_segments, 3, self.max_frames, 22, 1), dtype=np.float32) # num_segs, C, T, V, M

            for seg_i, seg in enumerate(labels_frame_info):
                if seg['end_f'] > np.shape(joints)[1]: # right now end frame can be slightly off (dataset issue)
                    seg['end_f'] = np.shape(joints)[1]
                num_frames = seg['end_f'] - seg['start_f']

                if num_frames > self.max_frames: # clip segments to max_frames
                    num_frames = self.max_frames

                joints_combined[seg_i, :, :num_frames, :, :] = joints[:, seg['start_f']: seg['start_f'] + num_frames, :, :]

                feed_dict.segment_boundaries.append((seg['start_f'], (seg['start_f'] + num_frames)))

            feed_dict.joints = joints_combined
            feed_dict.num_segs = num_segments
        else:
            total_num_frames = np.shape(joints)[1]
            num_frames_per_seg = 45
            overlap_frames = 15
            num_segments = math.ceil(total_num_frames / num_frames_per_seg)
            feed_dict['info'] = []
            joints_combined = np.zeros((num_segments, 3, num_frames_per_seg + overlap_frames*2, 22, 1), dtype=np.float32) # num_segs, C, T, V, M
            for i in range(num_segments):
                start_f = i * num_frames_per_seg
                end_f = (i + 1) * num_frames_per_seg
                if end_f > total_num_frames: end_f = total_num_frames

                missing_before_context = overlap_frames - start_f if start_f < overlap_frames else 0
                existing_after_context = total_num_frames - end_f
                if existing_after_context > overlap_frames: existing_after_context = overlap_frames

                joints_combined[i, :, missing_before_context:overlap_frames+(end_f - start_f)+existing_after_context, :, :] = joints[:, start_f - (overlap_frames - missing_before_context):end_f + existing_after_context, :, :]

                feed_dict.segment_boundaries.append((start_f - (overlap_frames - missing_before_context), end_f + existing_after_context))

            feed_dict.joints = joints_combined
            feed_dict.num_segs = num_segments

        return feed_dict.raw()

    def __len__(self):
        return len(self.split_question_ids)

class NSTrajDatasetFilterableView(FilterableDatasetView):
    def filter_questions(self, allowed):
        def filt(question):
            return question['query_type'] in allowed

        return self.filter(filt, 'filter-question-type[allowed={{{}}}]'.format(','.join(list(allowed))))


    def make_dataloader(self, batch_size, shuffle, drop_last, nr_workers):
        from jactorch.data.dataloader import JacDataLoader
        from jactorch.data.collate import VarLengthCollateV2

        collate_guide = {
            'joints': 'concat',
            'answer': 'skip',
            'segment_boundaries': 'skip',
            'filter_boundaries': 'skip',

            'program_raw': 'skip',
            'program_seq': 'skip',
            'program_tree': 'skip',
            'program_qsseq': 'skip',
            'program_qstree': 'skip',
        }

        return JacDataLoader(
            self, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last,
            num_workers=nr_workers, pin_memory=True,
            collate_fn=VarLengthCollateV2(collate_guide))


def NSTrajDataset(*args, **kwargs):
    return NSTrajDatasetFilterableView(NSTrajDatasetUnwrapped(*args, **kwargs))

class MotionClassificationDatasetUnwrapped(FilterableDatasetUnwrapped):
    def __init__(self, labels_json, joints_root, symbolic, max_frames=150):
        super().__init__()
        self.labels_json = labels_json
        self.joints_root = joints_root
        self.act2idx_json = 'nscl/datasets/babel/action_label_2_idx.json'

        self.labels = io.load_json(self.labels_json)
        self.babel_ids = list(self.labels.keys())
        self.act2idx = io.load_json(self.act2idx_json)

        self.symbolic = symbolic
        self.max_frames = max_frames

    def __getitem__(self, index):
        feed_dict = GView()

        babel_id = self.babel_ids[index]
        actions = []

        num_segments = len(self.labels[babel_id])

        joints = np.load(osp.join(self.joints_root, f'{babel_id}.npy')) # T, V, C

        # change shape of joints to match model
        joints = joints[:, :, :, np.newaxis] # T, V, C, M
        joints = joints.transpose(2, 0, 1, 3) # C, T, V, M

        joints_combined = np.zeros((num_segments, 3, self.max_frames, 22, 1), dtype=np.float32) # num_segs, C, T, V, M

        # TODO: take out segments that only have one frame (to be consistent with unbatched version)
        for seg_i, seg in enumerate(self.labels[babel_id]):
            if seg['end_f'] > np.shape(joints)[1]: # right now end frame can be slightly off (dataset issue)
                seg['end_f'] = np.shape(joints)[1]
            num_frames = seg['end_f'] - seg['start_f']

            if num_frames > self.max_frames: # clip segments to max_frames
                num_frames = self.max_frames

            joints_combined[seg_i, :, :num_frames, :, :] = joints[:, seg['start_f']: seg['start_f'] + num_frames, :, :]

            action = seg['action']
            if self.symbolic: # option 1: add all actions
                for i in range(len(action)):
                    action[i] = self.act2idx[action[i]]
                actions.append(action)
            else: # option 2: pick randomly among actions
                random_action_idx = np.random.randint(len(action))
                action_idx = self.act2idx[action[random_action_idx].replace('.', '')] # remove periods from action names
                actions.append([action_idx])

        feed_dict.actions = actions
        feed_dict.joints = joints_combined
        feed_dict.id = babel_id
        feed_dict.num_segs = num_segments

        return feed_dict.raw()

    def __len__(self):
        return len(self.babel_ids)

class MotionClassificationDatsetFilterableView(FilterableDatasetView):
    def make_dataloader(self, batch_size, shuffle, drop_last, nr_workers):
        from jactorch.data.dataloader import JacDataLoader
        from jactorch.data.collate import VarLengthCollateV2

        collate_guide = {
            'joints': 'concat',
            'actions': 'skip',
        }

        gdef.update_collate_guide(collate_guide)

        return JacDataLoader(
            self, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last,
            num_workers=nr_workers, pin_memory=True,
            collate_fn=VarLengthCollateV2(collate_guide))

def MotionClassificationDataset(*args, **kwargs):
    return MotionClassificationDatsetFilterableView(MotionClassificationDatasetUnwrapped(*args, **kwargs))




concepts/benchmark/humanmotion/dataset.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : humanmotion_constants.py
# Author : Joy Hsu
# Email  : joycj@stanford.edu
# Date   : 04/04/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

attribute_concepts_mapping = {
    'Motion': ['cartwheel_Action', 'catch_Action', 'crawl_Action', 'crouch_Action', 'flip_Action', 'jog_Action', 'jump_Action', 'jumping_jacks_Action', 'kick_Action', 'kneel_Action', 'knock_Action', 'place_something_Action', 'run_Action', 'sit_Action', 'squat_Action', 'stand_up_Action', 'take_Action', 'throw_Action', 'walk_Action', 'wave_Action'],
    'Part': ['left_arm_Action', 'left_foot_Action', 'left_hand_Action', 'left_leg_Action', 'right_arm_Action', 'right_foot_Action', 'right_hand_Action', 'right_leg_Action'],
    'Direction': ['backwards_Action', 'forward_Action', 'left_Action', 'right_Action']
}

concepts/benchmark/humanmotion/humanmotion_constants.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : utils.py
# Author : Joy Hsu
# Email  : joycj@stanford.edu
# Date   : 03/24/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from collections import defaultdict
from copy import deepcopy

__all__ = ['nsclseq_to_nscltree', 'nsclseq_to_nsclqsseq', 'nscltree_to_nsclqstree', 'program_to_nsclseq']

def nsclseq_to_nscltree(seq_program):
    def dfs(sblock):
        tblock = deepcopy(sblock)
        input_ids = tblock.pop('inputs')
        tblock['inputs'] = [dfs(seq_program[i]) for i in input_ids]
        return tblock

    try:
        return dfs(seq_program[-1])
    finally:
        del dfs
        
def nsclseq_to_nsclqsseq(seq_program):
    qs_seq = deepcopy(seq_program)
    cached = defaultdict(list)

    '''for sblock in qs_seq:
        for param_type in gdef.parameter_types:
            if param_type in sblock:
                sblock[param_type + '_idx'] = len(cached[param_type])
                sblock[param_type + '_values'] = cached[param_type]
                cached[param_type].append(sblock[param_type])'''

    return qs_seq

def nscltree_to_nsclqstree(tree_program):
    qs_tree = deepcopy(tree_program)
    cached = defaultdict(list)

    '''for tblock in iter_nscltree(qs_tree):
        for param_type in gdef.parameter_types:
            if param_type in tblock:
                tblock[param_type + '_idx'] = len(cached[param_type])
                tblock[param_type + '_values'] = cached[param_type]
                cached[param_type].append(tblock[param_type])'''

    return qs_tree


def iter_nscltree(tree_program):
    yield tree_program
    for i in tree_program['inputs']:
        yield from iter_nscltree(i)


def get_clevr_pblock_op(block):
    if 'type' in block:
        return block['type']
    assert 'function' in block
    return block['function']


def get_clevr_op_attribute(op):
    return "_".join(op.split('_')[1:])


def program_to_nsclseq(program):
    nscl_program = list()
    mapping = dict()

    for block_id, block in enumerate(program):
        op = get_clevr_pblock_op(block)
        current = None
        if op == 'scene':
            current = dict(op='scene')
        elif op.startswith('filter'):
            concept = block['value_inputs'][0]
            last = nscl_program[mapping[block['inputs'][0]]]
            if last['op'] == 'filter':
                last['concept'].append(concept)
            else:
                current = dict(op='filter', concept=[concept])
        elif op.startswith('relate'):
            concept = block['value_inputs'][0]
            current = dict(op='relate', relational_concept=[concept])
        elif op.startswith('same'):
            attribute = get_clevr_op_attribute(op)
            current = dict(op='relate_attribute_equal', attribute=attribute)
        elif op in ('intersect', 'union'):
            current = dict(op=op)
        elif op == 'unique':
            pass  # We will ignore the unique operations.
        else:
            if op.startswith('query'):
                if block_id == len(program) - 1:
                    attribute = get_clevr_op_attribute(op)
                    current = dict(op='query', attribute=attribute)
            elif op.startswith('equal') and op != 'equal_integer':
                attribute = get_clevr_op_attribute(op)
                current = dict(op='query_attribute_equal', attribute=attribute)
            elif op == 'exist':
                current = dict(op='exist')
            elif op == 'count':
                if block_id == len(program) - 1:
                    current = dict(op='count')
            elif op == 'equal_integer':
                current = dict(op='count_equal')
            elif op == 'less_than':
                current = dict(op='count_less')
            elif op == 'greater_than':
                current = dict(op='count_greater')
            else:
                raise ValueError('Unknown CLEVR operation: {}.'.format(op))

        if current is None:
            assert len(block['inputs']) == 1
            mapping[block_id] = mapping[block['inputs'][0]]
        else:
            current['inputs'] = list(map(mapping.get, block['inputs']))

            if '_output' in block:
                current['output'] = deepcopy(block['_output'])

            nscl_program.append(current)
            mapping[block_id] = len(nscl_program) - 1

    return nscl_program

concepts/benchmark/humanmotion/utils.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 06/13/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/benchmark/inductive_reasoning/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : boolean_normal_form.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 03/18/2018
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import collections
import jacinle.random as random

__all__ = ['CNF', 'DNF', 'random_generate_cnf']


class NormalForm(object):
    merge1 = None
    merge2 = None
    init1 = None
    init2 = None
    split1 = None
    split2 = None

    def __init__(self, nr_variables, exprs, varnames=None):
        self.nr_variables = nr_variables
        self.exprs = exprs
        self.idx2name = varnames

    def __call__(self, assigns=None, **kwargs):
        if assigns is None:
            return self.eval(kwargs)
        return self.eval(assigns)

    def eval(self, assigns):
        assert self.nr_variables == len(assigns)

        if isinstance(assigns, collections.Mapping):
            assert self.idx2name is not None
            assigns = [assigns[self.idx2name[k]] for k in range(self.nr_variables)]
        else:
            assert isinstance(assigns, collections.Sequence)
        assigns = list(map(bool, assigns))

        v1 = type(self).init1
        for e1, neg1 in self.exprs:
            v2 = type(self).init2
            for e2, neg2 in e1:
                if neg2:
                    v2 = type(self).merge2(v2, not assigns[e2])
                else:
                    v2 = type(self).merge2(v2, assigns[e2])
            if neg1:
                v1 = type(self).merge1(v1, not v2)
            else:
                v1 = type(self).merge1(v1, v2)
        return v1

    @classmethod
    def from_string(cls, expr):
        expr = expr.strip().replace(' ', '')  # canonize
        closures = []
        variables = set()

        ands = expr.split(cls.split1)
        for a in ands:
            this_closure = []
            if a[0] == '!':
                a = a[1:]
                nega = True
            else:
                nega = False

            assert a[0] == '(' and a[-1] == ')'
            a = a[1:-1]
            ors = a.split(cls.split2)
            for o in ors:
                if o.startswith('!'):
                    o = o[1:]
                    nego = True
                else:
                    nego = False
                this_closure.append([o, nego])
                variables.add(o)
            closures.append([this_closure, nega])

        variables = sorted(variables)
        var2idx = {v: i for i, v in enumerate(variables)}
        nr_variables = len(variables)
        for a, _ in closures:
            for o in a:
                o[0] = var2idx[o[0]]
        return cls(nr_variables, closures, varnames=variables)

    def to_string(self):
        if self.idx2name is not None:
            idx2name = self.idx2name
        else:
            assert self.nr_variables <= 26
            idx2name = [chr(97 + i) for i in range(26)]

        s1 = ''
        for e1, neg1 in self.exprs:
            if len(s1) != 0:
                s1 += type(self).split1
            s2 = ''
            for e2, neg2 in e1:
                if len(s2) != 0:
                    s2 += type(self).split2
                if neg2:
                    s2 += '!' + idx2name[e2]
                else:
                    s2 += idx2name[e2]
            if neg1:
                s2 = '!(' + s2 + ')'
            else:
                s2 = '(' + s2 + ')'
            s1 += s2
        return s1


class CNF(NormalForm):
    init1 = True
    init2 = False
    split1 = '&'
    split2 = '|'
    merge1 = lambda x, y: x and y
    merge2 = lambda x, y: x or y


class DNF(NormalForm):
    init1 = False
    init2 = True
    split1 = '|'
    split2 = '&'
    merge1 = lambda x, y: x or y
    merge2 = lambda x, y: x and y


def _canonize_max(m):
    if type(m) is int:
        m = [max(m // 2, 1), max(m, 1)]
    else:
        assert m[0] > 0 and m[1] > 0
    return m


def random_generate_cnf(nr_variables, max_ands, max_ors):
    max_ands = _canonize_max(max_ands)
    max_ors = _canonize_max(max_ors)

    closures = []
    nr_ands = random.randint(max_ands[0], max_ands[1] + 1)
    for i in range(nr_ands):
        nr_ors = random.randint(max_ors[0], max_ors[1] + 1)
        vars = random.randint(0, nr_variables, size=nr_ors)
        nots = random.randint(0, 2, size=nr_ors)
        this_closures = [(v, n) for v, n in zip(vars, nots)]
        closures.append((this_closures, random.randint(0, 2)))
    return CNF(nr_variables, closures)

concepts/benchmark/inductive_reasoning/boolean_normal_form.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : boolean_normal_form_dataset.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 03/18/2018
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
from torch.utils.data.dataset import Dataset


class NormalFormDataset(Dataset):
    """Learning a general normal form."""

    def __init__(self, normal_form, include=None, exclude=None):
        assert normal_form.nr_variables <= 16

        self.normal_form = normal_form
        self.include = include
        self.exclude = exclude

        if self.include is not None:
            self.include = list(self.include)
        elif self.exclude is not None:
            self.exclude = set(self.exclude)
            self.include = [i for i in range(1 << self.normal_form.nr_variables) if i not in self.exclude]
        else:
            self.include = list(range(1 << self.normal_form.nr_variables))

    def __getitem__(self, item):
        assigns = _binary_decomposition(self.include[item], self.normal_form.nr_variables)
        result = self.normal_form(assigns)
        return dict(input=np.array(assigns, dtype=np.float32), label=float(result))

    def __len__(self):
        return len(self.include)


class TruthTableDataset(Dataset):
    """Learning a truth table."""
    def __init__(self, nr_variables, table):
        assert nr_variables <= 16
        assert 1 << nr_variables == len(table)
        self.nr_variables = nr_variables
        self.table = table

    def __getitem__(self, item):
        assigns = _binary_decomposition(item, self.nr_variables)
        result = self.table[item]
        return dict(input=np.array(assigns, dtype=np.float32), label=float(result))

    def __len__(self):
        return len(self.table)


class ParityDataset(Dataset):
    """Learning the parity function."""
    def __init__(self, nr_variables):
        assert nr_variables <= 16
        self.nr_variables = nr_variables

    def __getitem__(self, item):
        assigns = _binary_decomposition(item, self.nr_variables)
        result = sum(assigns) % 2
        return dict(input=np.array(assigns, dtype=np.float32), label=float(result))

    def __len__(self):
        return 1 << self.nr_variables


def _binary_decomposition(v, n):
    return [bool(v & (1 << i)) for i in range(n)]

concepts/benchmark/inductive_reasoning/boolean_normal_form_dataset.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : family.py
# Author : Honghua Dong
# Email  : dhh19951@gmail.com
# Date   : 05/07/2018
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
import jacinle.random as random

__all__ = ['Family', 'random_generate_family']


class Family(object):
    """A data structure that stores the relationship between N people in a family."""

    def __init__(self, nr_people: int, relations: np.ndarray):
        """Initialize a family with relations.

        Args:
            nr_people: number of people in the family.
            relations: a 3D array of shape (nr_people, nr_people, 6), where
                relations[i, j, 0] = 1 if j is the husband of i, 0 otherwise.
                relations[i, j, 1] = 1 if j is the wife of i, 0 otherwise.
                relations[i, j, 2] = 1 if j is the father of i, 0 otherwise.
                relations[i, j, 3] = 1 if j is the mother of i, 0 otherwise.
                relations[i, j, 4] = 1 if j is the son of i, 0 otherwise.
                relations[i, j, 5] = 1 if j is the daughter of i, 0 otherwise.
        """
        self._n = nr_people
        self._relations = relations

    @property
    def father(self) -> np.ndarray:
        return self._relations[:, :, 2]

    @property
    def mother(self) -> np.ndarray:
        return self._relations[:, :, 3]

    @property
    def son(self) -> np.ndarray:
        return self._relations[:, :, 4]

    @property
    def daughter(self) -> np.ndarray:
        return self._relations[:, :, 5]

    def has_father(self) -> np.ndarray:
        return self.father.max(axis=1)

    def has_daughter(self) -> np.ndarray:
        return self.daughter.max(axis=1)

    def has_sister(self) -> np.ndarray:
        return _clip_mul(self.father, self.daughter).max(axis=1)

    def get_parents(self) -> np.ndarray:
        return np.clip(self.father + self.mother, 0, 1)

    def get_grandfather(self) -> np.ndarray:
        return _clip_mul(self.get_parents(), self.father)

    def get_grandmother(self) -> np.ndarray:
        return _clip_mul(self.get_parents(), self.mother)

    def get_grandparents(self) -> np.ndarray:
        parents = self.get_parents()
        return _clip_mul(parents, parents)

    def get_uncle(self) -> np.ndarray:
        return _clip_mul(self.get_grandparents(), self.son)

    def get_maternal_great_uncle(self) -> np.ndarray:
        return _clip_mul(_clip_mul(self.get_grandmother(), self.mother), self.son)


def random_generate_family(n, p_marriage=0.8, verbose=False) -> Family:
    assert n > 0
    ids = list(random.permutation(n))

    single_m = []
    single_w = []
    couples = [None]
    rel = np.zeros((n, n, 6))  # husband, wife, father, mother, son, daughter
    fathers = [None for i in range(n)]
    mothers = [None for i in range(n)]

    def add_couple(man, woman):
        couples.append((man, woman))
        rel[woman, man, 0] = 1  # husband
        rel[man, woman, 1] = 1  # wife
        if verbose:
            print('couple', man, woman)

    def add_child(parents, child, gender):
        father, mother = parents
        fathers[child] = father
        mothers[child] = mother
        rel[child, father, 2] = 1  # father
        rel[child, mother, 3] = 1  # mother
        if gender == 0:  # son
            rel[father, child, 4] = 1
            rel[mother, child, 4] = 1
        else:  # daughter
            rel[father, child, 5] = 1
            rel[mother, child, 5] = 1
        if verbose:
            print('child', father, mother, child, gender)

    def check_relations(man, woman):
        if fathers[man] is None or fathers[woman] is None:
            return True
        if fathers[man] == fathers[woman]:
            return False

        def same_parent(x, y):
            return fathers[x] is not None and fathers[y] is not None and fathers[x] == fathers[y]

        for x in [fathers[man], mothers[man]]:
            for y in [fathers[woman], mothers[woman]]:
                if same_parent(man, y) or same_parent(woman, x) or same_parent(x, y):
                    return False
        return True

    while len(ids) > 0:
        x = ids.pop()
        gender = random.randint(2)
        parents = random.choice_list(couples)
        if gender == 0:
            single_m.append(x)
        else:
            single_w.append(x)
        if parents is not None:
            add_child(parents, x, gender)

        if random.rand() < p_marriage and len(single_m) > 0 and len(single_w) > 0:
            mi = random.randint(len(single_m))
            wi = random.randint(len(single_w))
            man = single_m[mi]
            woman = single_w[wi]
            if check_relations(man, woman):
                add_couple(man, woman)
                del single_m[mi]
                del single_w[wi]

    return Family(n, rel)


def _clip_mul(x, y):
    return np.clip(np.matmul(x, y), 0, 1)

concepts/benchmark/inductive_reasoning/family.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : graph_dataset.py
# Author : Honghua Dong
# Email  : dhh19951@gmail.com
# Date   : 05/07/2018
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np

from torch.utils.data.dataset import Dataset
from torchvision import datasets

import jacinle.random as random

from concepts.benchmark.algorithm_env.graph import random_generate_graph, random_generate_graph_dnc
from concepts.benchmark.inductive_reasoning.family import random_generate_family

__all__ = ['GraphOutDegreeDataset', 'GraphConnectivityDataset', 'GraphAdjacentDataset', 'FamilyTreeDataset']


class GraphDatasetBase(Dataset):
    def __init__(self, nr_nodes, p, epoch_size, directed=False, gen_method='dnc'):
        if type(nr_nodes) is int:
            self.nr_nodes = (max(nr_nodes // 2, 1), nr_nodes)
        else:
            self.nr_nodes = tuple(nr_nodes)
        self.p = p
        self.epoch_size = epoch_size
        self.directed = directed
        self.gen_method = gen_method

    def _gen_graph(self, item):
        nr_nodes = item % (self.nr_nodes[1] - self.nr_nodes[0] + 1) + self.nr_nodes[0]
        if type(self.p) is float:
            p = self.p
        else:
            p = self.p[0] + random.rand() * (self.p[1] - self.p[0])
        gen_graph = random_generate_graph_dnc if self.gen_method == 'dnc' else random_generate_graph
        return gen_graph(nr_nodes, p, directed=self.directed)

    def __len__(self):
        return self.epoch_size


class GraphOutDegreeDataset(GraphDatasetBase):
    def __init__(self, nr_nodes, p, epoch_size, degree=2, directed=False, gen_method='dnc'):
        super().__init__(nr_nodes, p, epoch_size, directed, gen_method)
        self.degree = degree

    def __getitem__(self, item):
        graph = self._gen_graph(item)
        return dict(
            n=graph._nr_nodes,
            relations=np.expand_dims(graph.get_edges(), axis=-1),
            target=(graph.get_out_degree() == self.degree).astype('float'),
        )


class GraphConnectivityDataset(GraphDatasetBase):
    def __init__(self, nr_nodes, p, epoch_size, dist_limit=None, directed=False, gen_method='dnc'):
        super().__init__(nr_nodes, p, epoch_size, directed, gen_method)
        self.dist_limit = dist_limit

    def __getitem__(self, item):
        graph = self._gen_graph(item)
        return dict(
            n=graph._nr_nodes,
            relations=np.expand_dims(graph.get_edges(), axis=-1),
            # relations=graph.get_relations(),
            target=graph.get_connectivity(self.dist_limit, exclude_self=True),
        )


class GraphAdjacentDataset(GraphDatasetBase):
    def __init__(self, nr_nodes, p, epoch_size, nr_colors, directed=False, gen_method='dnc',
                 is_mnist_colors=False, is_train=True):

        super().__init__(nr_nodes, p, epoch_size, directed, gen_method)
        self._nr_colors = nr_colors
        self._mnist_colors = is_mnist_colors
        if is_mnist_colors:
            assert nr_colors == 10
            transform = None
            self.mnist = datasets.MNIST('../data', train=is_train, download=True, transform=transform)

    def __getitem__(self, item):
        graph = self._gen_graph(item)
        n = graph._nr_nodes
        if self._mnist_colors:
            m = self.mnist.__len__()
            digits = []
            colors = []
            for i in range(n):
                x = random.randint(m)
                digit, color = self.mnist.__getitem__(x)
                digits.append(np.array(digit)[np.newaxis])
                colors.append(color)
            digits, colors = np.array(digits), np.array(colors)
        else:
            colors = random.randint(self._nr_colors, size=n)
        states = np.zeros((n, self._nr_colors))
        adjacent = np.zeros((n, self._nr_colors))
        for i in range(n):
            states[i, colors[i]] = 1
            adjacent[i, colors[i]] = 1
            for j in range(n):
                if graph.has_edge(i, j):
                    adjacent[i, colors[j]] = 1
        if self._mnist_colors:
            states = digits
        return dict(
            n=n,
            relations=np.expand_dims(graph.get_edges(), axis=-1),
            states=states,
            colors=colors,
            target=adjacent,
            # connectivity=graph.get_connectivity(self.dist_limit, exclude_self=True),
        )


class FamilyTreeDataset(Dataset):
    def __init__(self, nr_people, epoch_size, task, p_marriage=0.8, balance_sample=False):
        super().__init__()
        if type(nr_people) is int:
            self.nr_people = (max(nr_people // 2, 1), nr_people)
        else:
            self.nr_people = tuple(nr_people)
        self.epoch_size = epoch_size
        self.task = task
        self.p_marriage = p_marriage
        self.balance_sample = balance_sample
        self.data = []

        assert task in ['has-father', 'has-daughter', 'has-sister', 'parents', 'grandparents', 'uncle', 'maternal-great-uncle']

    def _gen_family(self, item):
        nr_people = item % (self.nr_people[1] - self.nr_people[0] + 1) + self.nr_people[0]
        return random_generate_family(nr_people, self.p_marriage)

    def __getitem__(self, item):
        while len(self.data) == 0:
            family = self._gen_family(item)
            relations = family._relations[:, :, 2:]
            if self.task == 'has-father':
                target = family.has_father()
            elif self.task == 'has-daughter':
                target = family.has_daughter()
            elif self.task == 'has-sister':
                target = family.has_sister()
            elif self.task == 'parents':
                target = family.get_parents()
            elif self.task == 'grandparents':
                target = family.get_grandparents()
            elif self.task == 'uncle':
                target = family.get_uncle()
            elif self.task == 'maternal-great-uncle':
                target = family.get_maternal_great_uncle()
            else:
                assert False, "{} is not supported.".format(self.task)

            if not self.balance_sample:
                return dict(n=family._n, relations=relations, target=target)

            def get_position(x):
                return list(np.vstack(np.where(x)).T)

            def append_data(pos, target):
                states = np.zeros((family._n, 2))
                states[pos[0], 0] = states[pos[1], 1] = 1
                self.data.append(dict(n=family._n, relations=relations, states=states, target=target))

            positive = get_position(target == 1)
            if len(positive) == 0:
                continue
            negative = get_position(target == 0)
            np.random.shuffle(negative)
            negative = negative[:len(positive)]
            for i in positive:
                append_data(i, 1)
            for i in negative:
                append_data(i, 0)

        return self.data.pop()

    def __len__(self):
        return self.epoch_size

concepts/benchmark/inductive_reasoning/graph_dataset.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/01/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/benchmark/manip_tabletop/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : pybullet_tabletop.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/01/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import os.path as osp
from typing import Optional, List

import jacinle.io as io

from concepts.simulator.pybullet.client import BulletClient
from concepts.simulator.pybullet.world import WorldSaver
from concepts.simulator.pybullet.components.ur5.ur5_robot import UR5Robot
from concepts.simulator.pybullet.components.panda.panda_robot import PandaRobot
from concepts.simulator.pybullet.rotation_utils import rpy
from concepts.utils.typing_utils import Vec2f, Vec3f, Vec4f, BroadcastableVec3f

__all__ = ['TableTopEnv']


class TableTopEnv(object):
    """TableTopEnv is a environment for manipulating tools in a 3D but table-top-only environment.
    So there is only minimal amount of 3D information involved.

    The environment will be simulated by pybullet.
    """

    def __init__(self, client: Optional[BulletClient] = None):
        if client is None:
            client = BulletClient()

        self.client = client
        self.robot = None
        self.saver = WorldSaver(client.w, save=False)

    @property
    def world(self):
        return self.client.w

    @property
    def w(self):
        return self.client.w

    @property
    def p(self):
        return self.client.p

    def reset(self) -> None:
        self.saver.restore()

    PLANE_FILE = 'assets://basic/plane/plane.urdf'
    WORKSPACE_FILE = 'assets://basic/plane/workspace.urdf'

    PLANE_LARGE_FILE = 'assets://basic/plane/plane_large.urdf'
    WORKSPACE_LARGE_FILE = 'assets://basic/plane/workspace_large.urdf'

    def add_workspace(self, large: bool = False) -> int:
        """Add a table-top workspace to the environment.

        Returns:
            the body id of the collision shape of the body.
        """

        if not large:
            # The actual table with collision shape.
            plane_id = self.client.load_urdf(type(self).PLANE_FILE, (0, 0, -0.001), static=True, body_name='table')
            # Just a visual shape for the table.
            self.client.load_urdf(type(self).WORKSPACE_FILE, (1, 0, 0), static=True, body_name='workspace')
        else:
            plane_id = self.client.load_urdf(type(self).PLANE_LARGE_FILE, (0, 0, -0.001), static=True, body_name='table')
            self.client.load_urdf(type(self).WORKSPACE_LARGE_FILE, (1, 0, 0), static=True, body_name='workspace')
        return plane_id

    def add_robot(self, robot: str = 'panda', pos: Optional[Vec3f] = None) -> int:
        """Add a robot to the environment.

        Args:
            robot: the type of the robot. Currently only ``['ur5', 'panda']`` are supported.
            pos: the initial position of the robot. If not given, the robot will be placed at the origin.

        Returns:
            the body id of the robot.
        """
        if robot == 'ur5':
            self.robot = UR5Robot(self.client, pos=pos)
        elif robot == 'panda':
            self.robot = PandaRobot(self.client, pos=pos)
        else:
            raise ValueError(f'Unknown robot type: {robot}.')

        return self.robot.get_robot_body_id()

    def add_region(
        self, size_2d: Vec2f, location_2d: Vec2f, name: str = 'region', *,
        color_rgba: Vec4f = (0.5, 0.5, 0.5, 1)
    ) -> int:
        """Add a visual-only region indicator to the environment.

        Args:
            size_2d: the size of the region.
            location_2d: the location of the region, asssumed to be on the table (z = 0).
            name: the name of the region.
            color_rgba: the color of the region.

        Returns:
            the body id of the region.
        """
        visual_shape = self.p.createVisualShape(self.p.GEOM_BOX, halfExtents=[size_2d[0] / 2, size_2d[1] / 2, 0.0001], rgbaColor=color_rgba)
        shape = self.p.createMultiBody(
            baseMass=0,
            baseVisualShapeIndex=visual_shape,
            basePosition=[location_2d[0], location_2d[1], 0.001],
            baseOrientation=(0, 0, 0, 1)
        )
        self.world.notify_update(shape, name, 'region')
        return shape

    def add_box(
        self, size: BroadcastableVec3f, location_2d: Vec2f, name: str = 'box', *,
        static: bool = False, z_height: float = 0,
        mass: float = 0.2, lateral_friction: float = 1.0,
        color_rgba: Vec4f = (1, 0.34, 0.34, 1.),
        quat: Optional[Vec4f] = None
    ) -> int:
        if isinstance(size, float):
            size = (size, size, size)
        return self.client.load_urdf_template(
            'assets://basic/box/box-template.urdf',
            {'DIM': size, 'MASS': mass, 'COLOR': color_rgba, 'LATERAL_FRICTION': lateral_friction},
            pos=(location_2d[0], location_2d[1], size[2] / 2 + z_height),
            quat=quat,
            body_name=name,
            static=static,
        )

    # TODO(Jiayuan Mao @ 08/12): use better typing like Vector2f.
    def add_container(
        self, size_2d: Vec2f, depth: float, location_2d: Vec2f, name: str = 'container', *,
        static=True,
        color_rgba: Vec4f = (0.5, 0.5, 0.5, 1)
    ) -> int:
        return self.client.load_urdf_template(
            'assets://basic/container/container-template.urdf', {
                'DIM': (size_2d[0], size_2d[1], depth),
                'HALF': (size_2d[0] / 2, size_2d[1] / 2, depth / 2)
            },
            (location_2d[0], location_2d[1], depth / 2),
            rgba=(0.5, 1.0, 0.5, 1.0),
            body_name=name,
            static=static,
        )

    def add_bar(
        self, size_2d: Vec2f, thickness: float, location_2d: Vec2f, name: str = 'bar-shape', *,
        static: bool = False,
        quat: Vec4f = (0, 0, 0, 1)
    ) -> int:
        return self.client.load_urdf_template(
            'assets://basic/simple-tools/bar-shape-template.urdf',
            {'DIM': (size_2d[0], size_2d[1], thickness)},
            (location_2d[0], location_2d[1], thickness / 2),
            quat,
            body_name=name,
            static=static,
        )

    def add_l_shape(
        self, size1_2d: Vec2f, size2_2d: Vec2f, thickness: float, location_2d: Vec2f, name: str = 'l-shape', *,
        static: bool = False,
        quat: Vec4f = (0, 0, 0, 1)
    ) -> int:
        return self.client.load_urdf_template(
            'assets://basic/simple-tools/l-shape-template.urdf', {
                'DIMX': (size1_2d[0], size1_2d[1], thickness),
                'HALFX': (size1_2d[0] / 2, size1_2d[1] / 2, thickness / 2),
                'DIMY': (size2_2d[0], size2_2d[1], thickness),
                'HALFY': (size2_2d[0] / 2, size2_2d[1] / 2, thickness / 2),
                'DISP': ((size1_2d[0] + size2_2d[0]) / 2, (size2_2d[1] - size1_2d[1]) / 2, 0.0)
            },
            (location_2d[0], location_2d[1], thickness / 2),
            quat,
            body_name=name,
            static=static,
        )

    def add_l_shape_with_tip(
        self, size1_2d: Vec2f, size2_2d: Vec2f, size3_2d: Vec2f, thickness: float, location_2d: Vec2f, name: str = 'l-shape-with-tip', *,
        static=False,
        quat: Vec4f = (0, 0, 0, 1)
    ) -> int:
        return self.client.load_urdf_template(
            'assets://basic/simple-tools/l-shape-with-tip-template.urdf', {
                'DIMX': (size1_2d[0], size1_2d[1], thickness),
                'HALFX': (size1_2d[0] / 2, size1_2d[1] / 2, thickness / 2),
                'DIMY': (size2_2d[0], size2_2d[1], thickness),
                'HALFY': (size2_2d[0] / 2, size2_2d[1] / 2, thickness / 2),
                'DIMZ': (size3_2d[0], size3_2d[1], thickness),
                'HALFZ': (size3_2d[0] / 2, size3_2d[1] / 2, thickness / 2),
                'DISPY': ((size1_2d[0] + size2_2d[0]) / 2, (size2_2d[1] - size1_2d[1]) / 2, 0.0),
                'DISPZ': ((-size2_2d[0] - size3_2d[0]) / 2, (size2_2d[1] - size3_2d[1]) / 2, 0.0)
            },
            (location_2d[0], location_2d[1], thickness / 2),
            quat,
            body_name=name,
            static=static,
        )

    def add_t_shape(
        self, size1_2d: Vec2f, size2_2d: Vec2f, thickness: float, location_2d: Vec2f, name: str = 't-shape', *,
        static: bool = False,
        quat: Vec4f = (0, 0, 0, 1)
    ) -> int:
        return self.client.load_urdf_template(
            'assets://basic/simple-tools/t-shape-template.urdf', {
                'DIMX': (size1_2d[0], size1_2d[1], thickness),
                'HALFX': (size1_2d[0] / 2, size1_2d[1] / 2, thickness / 2),
                'DIMY': (size2_2d[0], size2_2d[1], thickness),
                'HALFY': (size2_2d[0] / 2, size2_2d[1] / 2, thickness / 2),
                'DISP': ((size1_2d[0] + size2_2d[0]) / 2, 0, 0.0)
            },
            (location_2d[0], location_2d[1], thickness / 2),
            quat,
            body_name=name,
            static=static,
        )

    def add_t_shape_with_tip(
        self, size1_2d: Vec2f, size2_2d: Vec2f, size3_2d: Vec2f, thickness: float, location_2d: Vec2f, name: str = 't-shape-with-tip', *,
        static: bool = False,
        quat: Vec4f = (0, 0, 0, 1)
    ) -> int:
        return self.client.load_urdf_template(
            'assets://basic/simple-tools/t-shape-with-tip-template.urdf', {
                'DIMX': (size1_2d[0], size1_2d[1], thickness),
                'HALFX': (size1_2d[0] / 2, size1_2d[1] / 2, thickness / 2),
                'DIMY': (size2_2d[0], size2_2d[1], thickness),
                'HALFY': (size2_2d[0] / 2, size2_2d[1] / 2, thickness / 2),
                'DIMZ': (size3_2d[0], size3_2d[1], thickness),
                'HALFZ': (size3_2d[0] / 2, size3_2d[1] / 2, thickness / 2),
                'DISPY': ((size1_2d[0] + size2_2d[0]) / 2, 0.0, 0.0),
                'DISPZ1': ((-size2_2d[0] - size3_2d[0]) / 2, (size2_2d[1] - size3_2d[1]) / 2, 0.0),
                'DISPZ2': ((-size2_2d[0] - size3_2d[0]) / 2, (-size2_2d[1] + size3_2d[1]) / 2, 0.0)
            },
            (location_2d[0], location_2d[1], thickness / 2),
            quat,
            body_name=name,
            static=static,
        )

    def add_plate(
        self, scale: float, location_2d: Vec2f, name: str = 'plate', *,
        static: bool = False, z_height: float = 0.0
    ) -> int:
        return self.client.load_urdf(
            'assets://objects/kitchenware/plate1/model_normalized.obj.urdf',
            (location_2d[0], location_2d[1], 0.063 * scale + z_height),
            rpy(90, 0, 0),
            scale=scale,
            body_name=name,
            static=static,
        )

    def add_thin_plate(
        self, scale: float, location_2d: Vec2f, name: str = 'plate', *,
        static: bool = False, z_height: float = 0.0
    ) -> int:
        return self.client.load_urdf(
            'assets://objects/ycb/029_plate/textured.obj.urdf',
            (location_2d[0], location_2d[1], 0.063 * scale + z_height),
            scale=scale,
            body_name=name,
            static=static,
        )

    @property
    def igibson_root(self) -> str:
        return osp.join(self.client.assets_root, 'igibson')

    def find_igibson_object_by_category(self, category: str) -> List[str]:
        path = osp.join(self.client.assets_root, 'igibson', 'objects', category)
        all_files = io.lsdir(path, '*/*.urdf')
        assert len(all_files) > 0, f'No urdf files found in {path}.'
        return all_files

    def add_igibson_object_by_category(
        self, category: str, location_3d: Vec3f, scale: float = 1.0, name: Optional[str] = None, *,
        urdf_file: Optional[str] = None,
        static: bool = False, quat: Vec4f = (0, 0, 0, 1),
    ):
        if name is None:
            name = category

        if urdf_file is None:
            all_files = self.find_igibson_object_by_category(category)
            urdf_file = all_files[0]

        return self.client.load_urdf(
            urdf_file,
            location_3d,
            quat,
            body_name=name,
            scale=scale,
            static=static,
        )


concepts/benchmark/manip_tabletop/pybullet_tabletop.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/13/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Basic data structures and utilities (e.g., parser, executor) for domain-specific languages."""

concepts/dsl/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : constraint.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 03/10/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data structures for Constraint Satisification Problems (CSPs).

In general, a CSP can be represented as a tuple of variables, domains (for each variable), and constraints.
Here, we only consider constraints of the following form: ``function(*arguments) == rv``. Internally:

- Each variable can either be a determined :class:`~concepts.dsl.tensor_value.TensorValue` or an :class:`OptimisticValue` (i.e., undetermined value).
    For :class:`TensorValue`, the representation contains its data type and the actual value.
    For :class:`OptimisticValue`, the representation contains its data type and a single integer, representing the index of the variable.
- Each constraint is represented as a :class:`Constraint` object, which contains the function, arguments, and the return value.
    The function is represented as a string (only for equality constraints), or BoolOpType (and, or, not), or QuantificationOpType (forall, exists),
        or :class:`~concepts.dsl.dsl_functions.Function` objects. Both arguments and return values can be :class:`DeterminedValue` or :class:`OptimisticValue`.
- The index of optimistic values starts from -2147483647 (i.e., ``OPTIM_MAGIC_NUMBER``). Therefore, they can be directly stored in tensors without any conversion,
    which is helpful for implementing CSP-based expression executors. There are a few helper functions to convert between optimistic values and their indices.
"""

from enum import IntEnum
from dataclasses import dataclass
from typing import Any, Optional, Union, Sequence, Tuple, List, Set, Dict

import torch
import jacinle
from jacinle.utils.enum import JacEnum

from concepts.dsl.dsl_types import BOOL, INT64, TensorValueTypeBase, NamedTensorValueType, PyObjValueType, Variable
from concepts.dsl.dsl_functions import Function
from concepts.dsl.tensor_value import TensorValue
from concepts.dsl.expression import BoolOpType, QuantificationOpType

__all__ = [
    'OPTIM_MAGIC_NUMBER', 'OPTIM_MAGIC_NUMBER_UPPER', 'OPTIM_MAGIC_NUMBER_MAGIC',
    'is_optimistic_value', 'optimistic_value_id', 'is_optimistic_value',
    'OptimisticValue', 'cvt_opt_value', 'SimulationFluentConstraintFunction',
    'Constraint', 'EqualityConstraint', 'ConstraintSatisfactionProblem',
    'AssignmentType', 'Assignment', 'AssignmentDict', 'print_assignment_dict', 'ground_assignment_value'
]


OPTIM_MAGIC_NUMBER = -2147483647
OPTIM_MAGIC_NUMBER_UPPER = (-2147483648) / 2
OPTIM_MAGIC_NUMBER_MAGIC = -2147483648


def is_optimistic_value(v: Union[int, torch.Tensor]):
    """Check if a value is an optimistic identifier."""
    return v < OPTIM_MAGIC_NUMBER_UPPER


def optimistic_value_id(v: Union[int, 'OptimisticValue']):
    """Get the optimistic identifier from an optimistic value."""
    if isinstance(v, OptimisticValue):
        return v.identifier - OPTIM_MAGIC_NUMBER
    return v - OPTIM_MAGIC_NUMBER


def maybe_optimistic_string(v: int):
    """If `v` is an optimistic value, return a string representation of the optimistic identifier, otherwise return `str(v)`."""
    if is_optimistic_value(v):
        return '@' + str(optimistic_value_id(v))
    return str(v)


class OptimisticValue(object):
    """An optimistic value object holds a pair of data type and an optimistic identifier."""

    def __init__(self, dtype: TensorValueTypeBase, identifier: int):
        """Initialize the OptimisticValue object.

        Args:
            dtype: the type of the value.
            identifier: the optimistic identifier.
        """

        if not isinstance(dtype, (NamedTensorValueType, PyObjValueType)) and dtype != BOOL and dtype != INT64:
            raise TypeError('OptimisticValue only supports NamedTensorValueType, PyObjValueType, BOOL, and INT64.')

        self.dtype = dtype
        self.identifier = identifier
        self.__post_init__()

    dtype: TensorValueTypeBase
    """The dtype of the optimistic value."""

    identifier: int
    """The optimistic identifier."""

    # TODO(Jiayuan Mao @ 2023/08/16): remove this check.
    def __post_init__(self):
        assert isinstance(self.identifier, int)

    def __str__(self) -> str:
        return f'O[{self.dtype}]{{@{optimistic_value_id(self.identifier)}}}'

    def __repr__(self) -> str:
        return self.__str__()


def cvt_opt_value(value: Union[OptimisticValue, TensorValue, bool, int], dtype: Optional[TensorValueTypeBase] = None) -> Union[TensorValue, OptimisticValue]:
    """Convert a value to OptimisticValue or a (determined) TensorValue. Acceptable types are:

    - OptimisticValue
    - TensorValue, which will be converted to OptimisticValue if it corresponds to an optimistic value.
    - bool, which will be converted to TensorValue.
    - int, which will be converted to OptimisticValue if it is an optimistic identifier, otherwise TensorValue.

    Args:
        value: the value to be converted.
        dtype: the data type of the value. If not specified (and the input is a TensorValue), the data type of the value will be used.

    Returns:
        The converted value.
    """

    if isinstance(value, OptimisticValue):
        return value
    elif isinstance(value, TensorValue):
        assert dtype is None or value.dtype == dtype
        v = value.single_elem()
        if isinstance(v, OptimisticValue):
            return v
        return value
    elif isinstance(value, bool):
        assert dtype is None or dtype == BOOL
        return TensorValue.from_scalar(value, BOOL)
    elif isinstance(value, int):
        if is_optimistic_value(value):
            assert dtype is not None, 'dtype must be specified for optimistic value'
            return OptimisticValue(dtype, value)
        if dtype == BOOL:
            raise RuntimeError('Should not use int to represent bool. This could be a bug in the code; report to the developers.')
            # return TensorValue.from_scalar(bool(value), BOOL)
        assert dtype == INT64
        return TensorValue.from_scalar(value, INT64)
    else:
        raise TypeError('Unknown value type: {} (type={}).'.format(value, type(value)))


@dataclass
class SimulationFluentConstraintFunction(object):
    """SimulationConstraint is a special kind of constraint that asserts the return value of the function is the grounding
    of a predicate after executing action ``action_index`` in the simulation. This is a special kind of constraint that
    has to be evaluated in a simulation environment. Therefore it is listed separately from the other domain-general constraints.
    """

    action_index: int
    """The index of the action that is executed in the simulation."""

    predicate: Function
    """The predicate to be grounded."""

    arguments: Tuple[int, ...]
    """The arguments to the predicate."""

    @property
    def name(self) -> str:
        return f'SimulationConstraint({self.action_index}, {self.predicate}, {self.arguments})'


class Constraint(object):
    """A constraint is basically a equality expression: ``function(*arguments) == rv``."""

    EQUAL = '__EQ__'
    """Magic name for equality constraints."""

    def __init__(
        self, function: Union[str, BoolOpType, QuantificationOpType, Function, SimulationFluentConstraintFunction],
        arguments: Sequence[Union[TensorValue, OptimisticValue]],
        rv: Union[TensorValue, OptimisticValue],
        note: Any = None
    ):
        """Initialize a constraint. Each constraint take the form of:

        .. code-block:: python

            function(*arguments) == rv

        Args:
            function: the function name, or a BoolOpType or QuantifierType, or a Function object.
            arguments: the arguments of the function.
            rv: the expected return value of the function.
            note: an optional note for the constraint.
        """

        self.function = function
        self.arguments = tuple(map(cvt_opt_value, arguments))
        self.rv = cvt_opt_value(rv)
        self.note = note

    function: Union[str, BoolOpType, QuantificationOpType, Function, SimulationFluentConstraintFunction]
    """The function identifier: either a string (currently only for equality constraints), or a BoolOpType or QuantificationOpType (for Boolean expressions), or a Function object."""

    arguments: Tuple[Union[TensorValue, OptimisticValue], ...]
    """The arguments to the function."""

    rv: Union[TensorValue, OptimisticValue]
    """The expected return value of the function."""

    note: Any
    """An optional note for the constraint."""

    def constraint_str(self):
        """Return the string representation of the constraint."""
        argument_strings = [x.format(short=True) if isinstance(x, TensorValue) else str(x) for x in self.arguments]
        if self.is_equal_constraint and isinstance(self.rv, TensorValue):
            if self.rv.item():
                return '__EQ__(' + ', '.join(argument_strings) + ')'
            else:
                return '__NEQ__(' + ', '.join(argument_strings) + ')'
        else:
            if isinstance(self.function, (str, JacEnum)):
                name = str(self.function)
            else:
                name = self.function.name
            return name + '(' + ', '.join(argument_strings) + ') == ' + str(self.rv)

    def __str__(self):
        ret = self.constraint_str()
        if self.note is not None:
            ret += '  # ' + str(self.note)
        return ret

    __repr__ = jacinle.repr_from_str

    @property
    def is_equal_constraint(self) -> bool:
        """Check if the constraint is an equality constraint."""
        return isinstance(self.function, str) and self.function == Constraint.EQUAL

    @classmethod
    def from_function(cls, function: Function, args: Sequence[Union[bool, int, torch.Tensor, TensorValue, Any]], rv: Union[bool, int, torch.Tensor, TensorValue, Any]) -> 'Constraint':
        """Create a constraint given a function, arguments, and return value.

        Args:
            function: the function object.
            args: the arguments of the function. The arguments can be bool, int, or torch.Tensor.
            rv: the return value of the function. The return value can be bool, int, or torch.Tensor.

        Returns:
            The created constraint.
        """
        _cvt = cvt_opt_value
        args = [_cvt(x, var.dtype if isinstance(var, Variable) else var) for x, var in zip(args, function.arguments)]
        rv = _cvt(rv, function.return_type)
        return cls(function, args, rv)


class EqualityConstraint(Constraint):
    """A special constraint for equality constraints. It is equivalent to:

    .. code-block:: python

        Constraint(Constraint.EQUAL, [left, right], rv)

    Basically, it states

    .. code-block:: python

        (left == right) == rv

    Therefore, when rv is True, it states that left and right are equal, and when rv is False, it states that left and right are not equal.
    """

    def __init__(self, left: Union[TensorValue, OptimisticValue], right: Union[TensorValue, OptimisticValue], rv: Optional[Union[TensorValue, OptimisticValue]] = None):
        """Initialize an equality constraint.

        Args:
            left: the left hand side of the equality constraint.
            right: the right hand side of the equality constraint.
            rv: the expected return value of the equality constraint. If None, it will be set to True (i.e., left == right).
        """
        super().__init__(Constraint.EQUAL, [left, right], rv if rv is not None else cvt_opt_value(True, BOOL))

    function: str
    """The function identifier, which is always ``Constraint.EQUAL``."""

    arguments: Tuple[Union[TensorValue, OptimisticValue], Union[TensorValue, OptimisticValue]]
    rv: Union[TensorValue, OptimisticValue]
    note: Any

    @classmethod
    def from_bool(cls, left: Union[bool, int, OptimisticValue], right: Union[bool, int, OptimisticValue], rv: Optional[Union[bool, int, OptimisticValue]] = None) -> 'EqualityConstraint':
        """Create an equality constraint from bool or optimistic values (represented as integers).

        Args:
            left: the left hand side of the equality constraint.
            right: the right hand side of the equality constraint.
            rv: the expected return value of the equality constraint. If None, it will be set to True (i.e., left == right).

        Returns:
            The created equality constraint.
        """

        def _cvt(x):
            if x is None:
                return x
            if isinstance(x, OptimisticValue):
                return x
            elif isinstance(x, bool):
                return cvt_opt_value(x, BOOL)
            else:
                assert isinstance(x, int) and is_optimistic_value(x)
                return OptimisticValue(BOOL, x)
        return cls(_cvt(left), _cvt(right), _cvt(rv))


class ConstraintSatisfactionProblem(object):
    """A constraint satisfaction problem.

    A constraint satisfaction problem is a set of constraints, and a set of variables.
    The solution to a constraint satisfaction problem is to find a set of values for the variables that satisfy all the constraints.
    """

    def __init__(
        self,
        index2actionable: Optional[Dict[int, bool]] = None,
        index2type: Optional[Dict[int, TensorValueTypeBase]] = None,
        index2domain: Optional[Dict[int, Set[Any]]] = None,
        constraints: Optional[List[Constraint]] = None,
        counter: int = 0,
    ):
        """Initialize a constraint satisfaction problem.

        Args:
            counter: the counter for generating new variable indices.
            index2actionable: a mapping from variable indices to whether the variable is actionable.
            index2type: a mapping from variable indices to the type of the variable.
            index2domain: a mapping from variable indices to the domain of the variable.
            constraints: a list of constraints.
        """

        self.index2actionable = index2actionable if index2actionable is not None else dict()
        self.index2type = index2type if index2type is not None else dict()
        self.index2domain = index2domain if index2domain is not None else dict()
        self.constraints = constraints if constraints is not None else list()
        self._optim_var_counter = counter

    index2actionable: Dict[int, bool]
    """A mapping from variable indices to whether the variable is actionable."""

    index2type: Dict[int, Union[TensorValueTypeBase, PyObjValueType]]
    """A mapping from variable indices to the type of the variable."""

    index2domain: Dict[int, Set[Any]]
    """A mapping from variable indices to the domain of the variable."""

    constraints: List[Constraint]
    """A list of constraints."""

    def clone(self, constraints: Optional[List[Constraint]] = None) -> 'ConstraintSatisfactionProblem':
        """Clone the constraint satisfaction problem.

        Args:
            constraints: the constraints to be replaced into the cloned constraint satisfaction problem. If None, the constraints of the original constraint satisfaction problem will be used.
        """

        if constraints is None:
            constraints = self.constraints.copy()
        return type(self)(self.index2actionable.copy(), self.index2type.copy(), self.index2domain.copy(), constraints, self._optim_var_counter)

    def new_actionable_var(self, dtype: Union[TensorValueTypeBase, PyObjValueType], domain: Optional[Set[Any]] = None, wrap: bool = False) -> Union[int, OptimisticValue]:
        """Create a new actionable variable.

        Args:
            dtype: the type of the variable.
            domain: the domain of the variable. If None, it will be assumed to the full domain of the type.
            wrap: whether to wrap the variable index into an OptimisticValue.
        """
        identifier = self.new_var(dtype, domain)
        self.index2actionable[identifier] = True
        if wrap:
            return OptimisticValue(dtype, identifier)
        return identifier

    def new_var(self, dtype: Union[TensorValueTypeBase, PyObjValueType], domain: Optional[Set[Any]] = None, wrap: bool = False) -> Union[int, OptimisticValue]:
        """Create a new variable.

        Args:
            dtype: the type of the variable.
            domain: the domain of the variable. If None, it will be assumed to the full domain of the type.
            wrap: whether to wrap the variable index into an OptimisticValue.

        Returns:
            The index of the new variable (int) if wrap is False, or the wrapped OptimisticValue if wrap is True.
        """

        identifier = OPTIM_MAGIC_NUMBER + self._optim_var_counter
        self.index2type[identifier] = dtype
        self._optim_var_counter += 1
        if domain is not None:
            self.index2domain[identifier] = domain

        if wrap:
            return OptimisticValue(dtype, identifier)
        return identifier

    def get_type(self, identifier: int) -> Union[TensorValueTypeBase, PyObjValueType]:
        """Get the type of a variable."""
        return self.index2type[identifier]

    def get_domain(self, identifier: int) -> Optional[Set[Any]]:
        """Get the domain of a variable."""
        return self.index2domain.get(identifier, None)

    def add_domain_value(self, identifier: int, value: Any):
        """Add a value to the domain of a variable."""
        if identifier not in self.index2domain:
            self.index2domain[identifier] = set()
        self.index2domain[identifier].add(value)

    def add_constraint(self, c: Constraint, note: Optional[Any] = None):
        """Add a constraint to the constraint satisfaction problem.

        Args:
            c: the constraint to be added.
            note: the note of the constraint.
        """
        if note is not None:
            c.note = note
        self.constraints.append(c)

    def __str__(self) -> str:
        fmt = 'ConstraintSatisfactionProblem{\n'
        fmt += '  Actionable Variables:\n    ' + '\n    '.join([f'@{x - OPTIM_MAGIC_NUMBER} - {self.index2type[x]}' for x in self.index2actionable]) + '\n'
        fmt += '  Constraints:\n'
        for c in self.constraints:
            fmt += f'    {c}\n'
        fmt += '}'
        return fmt

    def __repr__(self):
        return self.__str__()


class AssignmentType(IntEnum):
    """See class `Assignment`."""
    EQUAL = 0
    VALUE = 1
    IGNORE = 2


@dataclass
class Assignment(object):
    """An assignment of a variable."""

    t: AssignmentType
    """the type of the assignment. There are three types:

        EQUAL: The variable is equal to another variable.
        VALUE: The variable is equal to a value.
        IGNORE: The variable is ignored. This happens when the variable does not appear in any constraint.
    """

    d: Union[bool, int, None, TensorValue]
    """the value of the variable.

        - EQUAL: The variable is equal to another variable.
        - VALUE: The variable is equal to a value.
        - IGNORE: The variable is ignored. This happens when the variable does not appear in any constraint.
    """

    generator_index: Optional[Tuple[str, int]] = None

    # TODO(Jiayuan Mao @ 2023/08/16): remove this check.
    def __post_init__(self):
        assert self.d is None or isinstance(self.d, (TensorValue, bool)), f'Invalid type of d: {type(self.d)}'

    @property
    def assignment_type(self):
        """Alias of `assignment.t`."""
        return self.t

    @property
    def data(self):
        """Alias of `assignment.d`."""
        return self.d


AssignmentDict = Dict[int, Assignment]
"""A mapping from variable indices to assignment values."""


def print_assignment_dict(assignments: AssignmentDict):
    """Print an assignment dictionary.

    Args:
        assignments: the assignment dictionary.
    """
    print('AssignmentDict{')
    for k, v in assignments.items():
        while v.t is AssignmentType.EQUAL and v.d in assignments:
            v = assignments[v.d]
        if v.t is AssignmentType.VALUE:
            print('  @{}: {}'.format(optimistic_value_id(k), v.d))
        elif v.t is AssignmentType.EQUAL:
            print('  @{}: {}'.format(optimistic_value_id(k), f'@{optimistic_value_id(v.d)}'))
        elif v.t is AssignmentType.IGNORE:
            print('  @{}: {}'.format(optimistic_value_id(k), 'IGNORE'))
        else:
            raise ValueError(f'Unknown assignment type: {v.t}')
    print('}')


def ground_assignment_value(assignments: AssignmentDict, identifier: int) -> Any:
    """Get the value of a variable based on the assignment dictionary. It will follow the EQUAL assignment types.

    Args:
        assignments: the assignment dictionary.
        identifier: the identifier of the variable.

    Returns:
        the value of the variable in the assignment dict.
    """
    while identifier in assignments and assignments[identifier].t is AssignmentType.EQUAL:
        identifier = assignments[identifier].d
    assert identifier in assignments
    return assignments[identifier].d


concepts/dsl/constraint.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : dsl_domain.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""The base calss for domain definitions: containing custom types, constants, and functions."""

import inspect
from typing import Optional, Union, Dict, Callable
from concepts.dsl.dsl_types import ObjectType, ValueType, ObjectConstant
from concepts.dsl.dsl_functions import Function
from concepts.dsl.value import Value

__all__ = ['DSLDomainBase']


class DSLDomainBase(object):
    """The baseclass for the domain definition of a domain-specific language."""

    def __init__(self, name: Optional[str] = None):
        """Initialize the domain.

        Args:
            name: the name of the domain. If not specified, the name of the class will be used.
        """
        if name is None:
            name = type(self).__name__

        self.name = name
        self.types = dict()
        self.functions = dict()
        self.constants = dict()

    name: str
    """The name of the domain."""

    types: Dict[str, Union[ObjectType, ValueType]]
    """The types defined in the domain, as a dictionary from type names to types."""

    functions: Dict[str, Function]
    """The functions defined in the domain, as a dictionary from function names to functions."""

    constants: Dict[str, Union[ObjectConstant, Value]]
    """The constants defined in the domain, as a dictionary from the name to the :class:`~concepts.dsl.value.Value` objects."""

    def define_type(self, t: Union[ObjectType, ValueType]) -> Union[ObjectType, ValueType]:
        """Define a type in the domain.

        Args:
            t: the type to be defined.

        Returns:
            the type that is defined.
        """
        if not isinstance(t, (ObjectType, ValueType)):
            raise TypeError('Types can only be object types and value types.')
        if t.typename in self.types:
            raise NameError(f'Type {t.typename} already defined.')
        self.types[t.typename] = t

        return t

    def define_function(self, function: Union[Function, Callable], implementation: bool = True) -> Function:
        """Define a function in the domain.

        Args:
            function: the function to be defined.
            implementation: whether to store the function body of `function` as the implementation of the function.

        Returns:
            the function that is defined.
        """
        if not isinstance(function, Function):
            assert inspect.isfunction(function)
            function = Function.from_function(function, implementation=implementation)
        if function.name in self.functions:
            raise ValueError(f'Function {function.name} already defined.')
        self.functions[function.name] = function

        return function

    def define_const(self, dtype: Union[ObjectType, ValueType], value: str):
        """Define a constant with the given type and value.

        Args:
            dtype: the type of the constant.
            value: the value of the constant. The value should be a string that is the name of the constant.
        """
        if value in self.constants:
            raise ValueError(f'Constant {value} already defined.')

        if isinstance(dtype, ObjectType):
            self.constants[value] = ObjectConstant(dtype, value)
        else:
            self.constants[value] = Value(dtype, value)

    def get_function(self, name: str) -> Function:
        """Get the function with the given name.

        Args:
            name: the name of the function.

        Returns:
            the function with the given name.
        """
        if name not in self.functions:
            raise KeyError(f'Function {name} not found.')
        return self.functions[name]

    def __str__(self) -> str:
        return f'{type(self).__name__}({self.name})'

    def __repr__(self) -> str:
        return str(self)


concepts/dsl/dsl_domain.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : dsl_functions.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/19/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data structures representing functions in a DSL.

Most importantly, this file contains the following classes:

- :class:`FunctionType`: the underlying type of a function, including argument types and return types.
- :class:`Function`: the function object, which is a callable object that can be used in expressions. They have names and types.

This file also implements a data structure for overloaded functions: :class:`OverloadedFunctionType`.
Internally, it contains a list of :class:`FunctionType` objects, and it is used to represent overloaded functions.
There are a few argument resolution methods implemented for both :class:`FunctionType` and :class:`OverloadedFunctionType`.
"""

import itertools
import contextlib
import inspect
import re
from typing import TYPE_CHECKING, Any, Union, Sequence, Tuple, List, Dict, Optional, Callable
from dataclasses import dataclass

import jacinle
from jacinle.utils.cache import cached_property
from jacinle.utils.defaults import option_context
from jacinle.utils.printing import indent_text

from concepts.dsl.dsl_types import TypeBase, ObjectType, ValueType, ConstantType, UnionType, Variable
from concepts.dsl.dsl_types import FormatContext, get_format_context

if TYPE_CHECKING:
    from concepts.dsl.expression import Expression, FunctionApplicationExpression

__all__ = [
    'FunctionArgumentResolutionError', 'FunctionArgumentResolutionContext', 'get_function_argument_resolution_context',
    'FunctionArgumentUnset', 'AnonymousFunctionArgumentGenerator',
    'FunctionArgumentType', 'FunctionArgumentListType', 'FunctionReturnType', 'FunctionType',
    'OverloadedFunctionResolution', 'OverloadedFunctionAmbiguousResolutions', 'OverloadedFunctionType',
    'FunctionTyping',
    'FunctionOverriddenCallList', 'FunctionDerivedExpressionList', 'FunctionResolvedFromRecord', 'Function',
]


class FunctionArgumentResolutionError(Exception):
    """Exception raised when the function argument resolution fails."""
    pass


class FunctionArgumentResolutionContext(option_context(
    '_FunctionArgumentResolutionContext',
    check_missing=True, check_type=True, check_overloaded_ambiguity=True, exc_verbose=True
)):
    """A context manager for controlling the function argument resolution.

    Attributes:
        check_missing (bool): whether to check if the function argument is missing.
        check_type (bool): whether to check if the function argument type is correct.
        check_overloaded_ambiguity (bool): whether to check if the function argument resolution is ambiguous.
        exc_verbose (bool):wWhether to print verbose error message.
    """

    @contextlib.contextmanager
    def exc(self, exc_type=None, from_=None):
        if self.exc_verbose:
            yield
        else:
            if exc_type is None:
                exc_type = FunctionArgumentResolutionError
            if from_ is not None:
                raise exc_type() from from_
            raise exc_type()


get_function_argument_resolution_context: Callable[[], FunctionArgumentResolutionContext] = FunctionArgumentResolutionContext.get_default
"""Get the current function argument resolution context."""


FunctionArgumentUnset = object()
"""A placeholder indicating that the argument is not specified."""


class AnonymousFunctionArgumentGenerator(object):
    """A generator for anonymous function arguments."""

    def __init__(self, template='_t{i:d}'):
        self.template = template
        self.counter = 0

    @property
    def nr_generated(self) -> int:
        return self.counter

    def gen(self, n: Optional[int] = None) -> Union[str, List[str]]:
        if n is None:
            self.counter += 1
            return self.template.format(i=self.counter)
        return [self.gen() for _ in range(n)]


FunctionArgumentType = Union[ObjectType, ValueType, 'FunctionType']
"""Acceptable types for function arguments. See the documentation of `FunctionType` for more details."""

FunctionArgumentListType = Union[Sequence[FunctionArgumentType], Sequence[Variable], Dict[str, FunctionArgumentType]]
"""Acceptable types for function argument lists. See the documentation of `FunctionType` for more details."""

FunctionReturnType = Union[ValueType, 'FunctionType', Sequence[Union[ValueType, 'FunctionType']]]
"""Acceptable types for function return types. See the documentation of `FunctionType` for more details."""


class FunctionType(TypeBase):
    """FunctionType defines the signature of a function."""

    argument_types: Tuple[FunctionArgumentType, ...]
    """The types of the arguments."""

    argument_names: Tuple[str, ...]
    """The names of the arguments."""

    arguments: Tuple[Variable, ...]
    """The argument list composed of `Variable` instances."""

    arguments_dict: Dict[str, FunctionArgumentType]
    """The arguments as a dict, as mappings from argument names to argument types."""

    arguments_name2index: Dict[str, int]
    """The mapping from argument names to argument indices."""

    return_type: Union[FunctionReturnType, Tuple[FunctionReturnType, ...]]
    """The return type of the function type."""

    return_name: Optional[Union[str, Tuple[str, ...]]]
    """The name of the return value."""

    is_singular_return: bool
    """Whether there is only one return value."""

    is_cacheable: bool
    """Whether the function is cacheable."""

    def __init__(
        self,
        arguments: FunctionArgumentListType,
        return_type: FunctionReturnType,
        argument_names: Optional[Sequence[str]] = None,
        return_name: Optional[Union[str, Sequence[str]]] = None,
        alias: Optional[str] = None
    ):
        """Initialize the function type.

        There are four ways to specify the arguments of the function type:

            1. A list of types, in which case the name of each argument is the index of the argument, using the format `#{index}`.
            2. A list of types as the `arguments`, and a list of names as the `argument_names`.
            3. A list of variables, in which case the name of each argument is the name of the variable.
            4. A dictionary of {name: type}, in which case the order of the arguments is the order of the keys.

        The return type can be either a single type or a tuple of types (multi-return types).

        Args:
            arguments: The arguments of the function type.
                When it is a list, the name of the arguments will be automatically generated.
                When it is a dict, the name of the arguments will be the keys of the dict.
            return_type: The return type of the function type.
            argument_names: The names of the arguments.
            return_name: The name of the return value.
            alias: The alias name of the function type.
        """

        self.arguments = None  # noqa
        self.arguments_dict = None  # noqa

        if isinstance(arguments, (list, tuple)):
            if len(arguments) == 0:
                self.arguments = tuple()
                self.arguments_dict = dict()
                self.argument_names = tuple()
                self.argument_types = tuple()
            elif isinstance(arguments[0], Variable):
                assert argument_names is None, 'Cannot specify both `arguments` and `argument_names`.'
                self.argument_types = tuple(arg.dtype for arg in arguments)
                self.argument_names = tuple(arg.name for arg in arguments)
                self.arguments = arguments
            else:
                if argument_names is None:
                    argument_names = tuple('#' + str(i) for i in range(len(arguments)))
                else:
                    assert len(arguments) == len(argument_names), 'The length of `arguments` and `argument_names` must be the same.'
                self.argument_types = tuple(arguments)
                self.argument_names = tuple(argument_names)
        elif isinstance(arguments, dict):
            assert argument_names is None, 'Cannot specify both `arguments` and `argument_names`.'
            self.argument_names = tuple(arguments.keys())
            self.argument_types = tuple(arguments.values())
            self.arguments_dict = arguments.copy()
        else:
            raise TypeError(f'Invalid argument types: {arguments}. Must be a list or a dict.')

        if self.arguments is None:
            self.arguments = tuple(Variable(name, dtype) for name, dtype in zip(self.argument_names, self.argument_types))
        if self.arguments_dict is None:
            self.arguments_dict = {name: dtype for name, dtype in zip(self.argument_names, self.argument_types)}

        if isinstance(return_type, ValueType):
            self.return_type = return_type
            self.return_name = return_name
        else:
            self.return_type = tuple(return_type)
            if return_name is None:
                self.return_name = None
            else:
                self.return_name = tuple(return_name)
                assert len(self.return_type) == len(self.return_name), 'The length of `return_type` and `return_name` must be the same.'

            if len(self.return_type) == 1:
                self.return_type = self.return_type[0]
                if self.return_name is not None:
                    self.return_name = self.return_name[0]

        self.is_singular_return = isinstance(self.return_type, ValueType)
        self.is_cacheable = self._gen_is_cacheale()

        super().__init__(self._gen_typename(), alias=alias)

    def _gen_typename(self) -> str:
        return '(' + ', '.join([str(arg) for arg in self.arguments]) + ') -> ' + str(self.return_type)

    def _gen_is_cacheale(self):
        for arg_def in self.arguments:
            if isinstance(arg_def, ValueType):
                return False
        return True

    @property
    def nr_arguments(self) -> int:
        """Return the number of arguments."""
        return len(self.argument_types)

    @cached_property
    def nr_object_arguments(self) -> int:
        """Return the number of arguments that are ObjectType-ed."""
        return len(list(filter(lambda x: isinstance(x, ValueType), self.argument_types)))

    @cached_property
    def nr_value_arguments(self) -> int:
        """Return the number of arguments that are ValueType-ed."""
        return len(list(filter(lambda x: isinstance(x, ValueType), self.argument_types)))

    @cached_property
    def nr_variable_arguments(self) -> int:
        """Return the number of arguments that are VariableType-ed (i.e. is ValueType-ed but not ConstantType-ed."""
        return len(list(filter(lambda x: isinstance(x, ValueType) and not isinstance(x, ConstantType), self.argument_types)))

    @cached_property
    def nr_constant_arguments(self) -> int:
        """Return the number of arguments that are ConstantType-ed."""
        return len(list(filter(lambda x: isinstance(x, ConstantType), self.argument_types)))

    @cached_property
    def arguments_name2index(self):
        assert self.argument_names is not None
        return {v: k for k, v in enumerate(self.argument_names)}

    @classmethod
    def from_annotation(cls, function: Callable, sig: Optional[inspect.Signature] = None) -> Union['FunctionType', 'OverloadedFunctionType']:
        """Create a FunctionType from a function annotation.

        Args:
            function: The function.
            sig: The signature of the function.

        Returns:
            Union[FunctionType, OverloadedFunctionType]: The function type.
        """

        if sig is None:
            sig = inspect.signature(function)

        argument_types = list()
        argument_names = list()
        for i, (name, param) in enumerate(sig.parameters.items()):
            if i == 0 and name == 'self':
                continue  # is an instancemethod.
            if i == 0 and name == 'cls':
                continue  # is a classmethod.
            argument_names.append(name)
            argument_types.append(param.annotation)

        return_type = sig.return_annotation

        if inspect._empty in argument_types or return_type is inspect._empty:
            raise FunctionArgumentResolutionError(f'Incomplete argument and return type annotation for {function}.')

        function_type = cls(argument_types, return_type, argument_names=argument_names)

        for arg_type in function_type.argument_types:
            if isinstance(arg_type, UnionType):
                return OverloadedFunctionType.from_function_type_with_union_arguments(function_type)

        return function_type

    def resolve_args(self, *args: Any, **kwargs: Any) -> List[Any]:
        """Resolve the arguments to the function type.

        If you want to specify a specific "positional" argument by its index, use `_{index}` as the name of the argument.

        Args:
            *args: The positional arguments.
            **kwargs: The keyword arguments.

        Returns:
            A list of argument values.
        """

        resolution_context = get_function_argument_resolution_context()

        # Construct a mapping from the name of arguments to their indices.
        name2index = {f'#{i}': i for i in range(self.nr_arguments)}
        name2index.update(self.arguments_name2index)

        arguments = [FunctionArgumentUnset for _ in range(self.nr_arguments)]
        if len(args) + len(kwargs) > self.nr_arguments:
            with resolution_context.exc():
                raise FunctionArgumentResolutionError(f'Function {self} takes {len(self.argument_types)} arguments, got {len(args) + len(kwargs)}.')

        for i in range(len(args)):
            arguments[i] = args[i]
        for k, v in kwargs.items():
            if k not in name2index:
                with resolution_context.exc():
                    raise FunctionArgumentResolutionError(f'Got unknown keyword argument: {k} when invoking function {self}.')
            i = name2index[k]
            if arguments[i] is not FunctionArgumentUnset:
                with resolution_context.exc():
                    raise FunctionArgumentResolutionError(f'Got duplicated argument for keyword argument: {k} when invoking function {self}.')
            arguments[i] = v

        if resolution_context.check_missing:
            for i in range(self.nr_arguments):
                if arguments[i] is FunctionArgumentUnset:
                    with resolution_context.exc():
                        raise FunctionArgumentResolutionError(f'Missing argument {self.argument_names[i]} when invoking function {self}.')

        if resolution_context.check_type:
            from concepts.dsl.expression import get_types
            argument_types = get_types(arguments)
            for i in range(self.nr_arguments):
                if argument_types[i] is not FunctionArgumentUnset and not argument_types[i].downcast_compatible(self.argument_types[i]):
                    with resolution_context.exc():
                        raise FunctionArgumentResolutionError(f'Typecheck failed for argument {self.argument_names[i]} while invoking the function {self}.\nInvoked with types: {argument_types}.')

        return arguments


@dataclass
class OverloadedFunctionResolution(object):
    """The data structure for storing the result of resolving an overloaded function."""

    type_index: int
    """The index of the function type that matches the expected signature."""

    ftype: FunctionType
    """The function type that matches the expected signature."""

    arguments: List[Any]
    """The resolved arguments."""


class OverloadedFunctionAmbiguousResolutions(list, List[OverloadedFunctionResolution]):
    pass


class OverloadedFunctionType(TypeBase):
    types: Tuple[FunctionType]

    def __init__(
        self,
        types: Sequence[Union['OverloadedFunctionType', FunctionType]],
        alias: Optional[str] = None
    ):
        types_flatten: List[FunctionType] = list()
        for ftype in types:
            if isinstance(ftype, OverloadedFunctionType):
                types_flatten.extend(ftype.types)
            else:
                assert isinstance(type, FunctionType)
                types_flatten.append(ftype)

        self.types = tuple(types_flatten)
        super().__init__(self._gen_typename(), alias=alias)

    def _gen_typename(self) -> str:
        return 'Overloaded{' + ','.join([x.typename for x in self.types]) + '}'

    @property
    def nr_types(self) -> int:
        """Return the number of sub-types."""
        return len(self.types)

    @classmethod
    def from_function_type_with_union_arguments(cls, function_type: FunctionType):
        """Create an OverloadedFunctionType from a FunctionType with Union-Typed arguments."""
        product_bases = list()
        for arg_type in function_type.argument_types:
            if isinstance(arg_type, UnionType):
                product_bases.append(arg_type.types)
            else:
                product_bases.append([arg_type])

        product_types = tuple(
            FunctionType(
                arg_type, function_type.return_type,
                argument_names=function_type.argument_names
            ) for arg_type in itertools.product(*product_bases)
        )
        return cls(product_types, alias=function_type.typename)

    def resolve_type_and_args(self, *args: Any, **kwargs: Any) -> Union[OverloadedFunctionResolution, OverloadedFunctionAmbiguousResolutions]:
        """Resolve the exact sub-function type being called and the argument list.

        Args:
            *args: The positional arguments.
            **kwargs: The keyword arguments.

        Returns:
            A :class:`OverloadedFunctionResolution` object if the resolution is unambiguous, or a :class:`OverloadedFunctionAmbiguousResolutions` object if the resolution is ambiguous.
            The ambiguity resolution object will only be returned if the ``check_overloaded_ambiguity`` flag is set to False in the :class:`FunctionArgumentResolutionContext`.

            - The :class:`OverloadedFunctionResolution` object contains the index of the sub-function type being called, the sub-function type, and the resolved argument list.
            - The :class:`OverloadedFunctionAmbiguousResolutions` object contains a list of :class:`OverloadedFunctionResolution` objects.
        """
        resolution_context = get_function_argument_resolution_context()

        success_results = list()
        exceptions = list()
        for i, ftype in enumerate(self.types):
            try:
                arguments = ftype.resolve_args(*args, **kwargs)
                success_results.append(OverloadedFunctionResolution(i, ftype, arguments))
            except FunctionArgumentResolutionError as e:
                exceptions.append(e)

        if len(success_results) == 1:
            return success_results[0]
        elif len(success_results) == 0:
            with resolution_context.exc():
                fmt = 'Failed to resolve overloaded function{}.\n'.format('' if self.typename is None else ' ' + self.typename)
                fmt += 'Detailed messages are:\n'
                for ftype, r in zip(self.types, exceptions):
                    this_fmt = 'Trying ' + str(ftype) + ':\n'
                    this_fmt += indent_text(str(r))
                    fmt += indent_text(this_fmt) + '\n'
                raise FunctionArgumentResolutionError(fmt.rstrip())
        else:
            if resolution_context.check_overloaded_ambiguity:
                with resolution_context.exc():
                    fmt = 'Got ambiguous application of overloaded function{}.\n'.format('' if self.typename is None else ' ' + self.typename)
                    fmt += 'Candidates are:\n'
                    for r in success_results:
                        fmt += indent_text(str(r[1])) + '\n'
                    fmt += 'Invoked with arguments: {}.'.format(str(success_results[0][2]))
                    raise FunctionArgumentResolutionError(fmt)
            else:
                return OverloadedFunctionAmbiguousResolutions(success_results)


class _FunctionTypingSugarInner(object):
    def __init__(self, return_type):
        self.return_type = return_type

    def __call__(self, *args, **kwargs):
        if len(args) == 0 and len(kwargs) == 0:
            return FunctionType(tuple(), self.return_type)
        elif len(args) != 0:
            assert len(kwargs) == 0, 'Only support all positional arguments or all positional keyword arguments.'
            return FunctionType(args, self.return_type)
        elif len(kwargs) != 0:
            assert len(args) == 0, 'Only support all positional arguments or all positional keyword arguments.'
            return FunctionType(tuple(kwargs.values()), self.return_type, tuple(kwargs.keys()))
        raise ValueError('Unreachable.')


class _FunctionTypingSugar(object):
    def __getitem__(self, return_type):
        return _FunctionTypingSugarInner(return_type)


"""FunctionTyping is a language-sugar constructor for function types.
For example: `FunctionTypingp[BOOL](INT64, INT64)` creates a function type with two INT64 arguments and a BOOL return type."""
FunctionTyping = _FunctionTypingSugar()


class FunctionOverriddenCallList(list, List[Callable]):
    """A data structure that holds multiple overridden __call__ implementations for a function.

    This is only useful when we are partial evaluating a function (and when the actual function type can not be resolved.)
    """
    pass


class FunctionDerivedExpressionList(list, List['Expression']):
    """A data structure that holds multiple derived expressions for a function."""
    pass


@dataclass
class FunctionResolvedFromRecord(object):
    function: Callable
    ftype_index: Union[int, Tuple[int, ...]]


class Function(object):
    """A function object holds a function type and an optional overridden __call__.

    The function object holds an additional field called `overridden_call`, which isa callable function.
    This field is used to override the __call__ method of the function object.

    By default, the __call__ function returns a FunctionApplication object, which contains the name of the function
    and a list of arguments. However, when `overridden_call` is set, the __call__ method will return the result of
    calling `overridden_call` with the arguments.
    """

    def __init__(
        self,
        name: str,
        ftype: Union[FunctionType, OverloadedFunctionType],
        derived_expression: Optional[Union['Expression', FunctionDerivedExpressionList]] = None,
        overridden_call: Optional[Union[Callable, FunctionOverriddenCallList]] = None,
        resolved_from: Optional[FunctionResolvedFromRecord] = None,
        function_body: Optional[Union[Callable, Sequence[Callable]]] = None
    ):
        """
        Args:
            name: the name of the function.
            ftype: the function type.
            derived_expression: the expression that this function is derived from.
            overridden_call: the overridden call function.
            resolved_from: the record of the function that this function is resolved from. This is used for handling
                partial evaluation and function specialization (for overloadded functions).
            function_body: the function body.
        """

        self.ftype = ftype
        self.derived_expression = derived_expression
        self.overridden_call = overridden_call
        self.resolved_from = resolved_from

        if isinstance(self.ftype, OverloadedFunctionType) and isinstance(self.overridden_call, FunctionOverriddenCallList):
            assert self.ftype.nr_types == len(self.overridden_call)

        self.name = name
        self.function_body = function_body  # the function body defined during the declaration.

        if self.derived_expression is None and self.overridden_call is not None:
            if self.is_overloaded:
                self.derived_expression = FunctionDerivedExpressionList()
                for i in range(self.ftype.nr_types):
                    self.derived_expression.append(_gen_expression_from_overridden_call(self.ftype.types[i], self.overridden_call[i]))
            else:
                self.derived_expression = _gen_expression_from_overridden_call(ftype, self.overridden_call)

        self.is_derived = self.derived_expression is not None

    def set_function_name(self, function_name: str):
        """Set the function name."""
        self.name = function_name

    def set_function_body(self, function_body: Callable):
        """Set the function body."""
        self.function_body = function_body

    """Argument and return type of the function (when the function is not an overloaded one)."""

    @property
    def arguments(self) -> Tuple[Variable]:
        assert not self.is_overloaded
        return self.ftype.arguments

    @property
    def nr_arguments(self) -> int:
        assert not self.is_overloaded
        return self.ftype.nr_arguments

    @property
    def return_type(self) -> Union[ValueType, Tuple[ValueType, ...]]:
        assert not self.is_overloaded
        return self.ftype.return_type

    """When the function is overloaded, the following functions are used for get the "overridden calls" for each function type."""

    @property
    def is_overloaded(self) -> bool:
        """Return True if the function is overloaded."""
        return isinstance(self.ftype, OverloadedFunctionType)

    def get_overridden_call(self, ftype_index: Optional[int] = None) -> Optional[Callable]:
        """Get the overridden call function."""
        if isinstance(self.overridden_call, FunctionOverriddenCallList):
            assert ftype_index is not None
            return self.overridden_call[ftype_index]
        return self.overridden_call

    def get_sub_function(self, ftype_index: int) -> 'Function':
        assert self.is_overloaded
        assert 0 <= ftype_index < self.ftype.nr_types
        return type(self)(
            self.name,
            self.ftype.types[ftype_index],
            self.get_overridden_call(ftype_index),
            resolved_from=FunctionResolvedFromRecord(self, ftype_index),
            function_body=self.function_body[ftype_index] if self.function_body is not None else None
        )

    @cached_property
    def all_sub_functions(self) -> List['Function']:
        assert self.is_overloaded
        return [self.get_sub_function(i) for i in range(self.ftype.nr_types)]

    @classmethod
    def from_function(cls, function: Callable, implementation: bool = True, sig: Optional[inspect.Signature] = None):
        """Create a function object from an actual Python function.

        Args:
            function: The function.
            implementation: Whether the function is an implementation. Defaults to True.
            sig: The signature of the function. Defaults to None.
        """
        ftype = FunctionType.from_annotation(function, sig=sig)
        return cls(function.__name__, ftype, function_body=function if implementation else None)

    def __call__(self, *args, **kwargs):
        if self.overridden_call is not None:
            if isinstance(self.ftype, OverloadedFunctionType):
                ftype_index, function_type, resolved_args = self.ftype.resolve_type_and_args(*args, **kwargs)
                return self.get_overridden_call(ftype_index)(*resolved_args)
            else:
                resolved_args = self.ftype.resolve_args(*args, **kwargs)
                return self.overridden_call(*resolved_args)

        if isinstance(self.ftype, OverloadedFunctionType):
            ftype_index, function_type, resolved_args = self.ftype.resolve_type_and_args(*args, **kwargs)
            function = Function(
                self.name + f'_{ftype_index}',
                function_type,
                overridden_call=None,  # Must be none.
                resolved_from=FunctionResolvedFromRecord(self, ftype_index),
            )
        else:
            resolved_args = self.ftype.resolve_args(*args, **kwargs)
            function = self

        from concepts.dsl.expression import FunctionApplicationExpression, cvt_expression_list
        return FunctionApplicationExpression(function, cvt_expression_list(resolved_args, function.ftype.argument_types))

    def __str__(self):
        if self.is_derived and not self.is_overloaded:
            with FormatContext(type_format_cls=False).as_default():
                if get_format_context().function_format_lambda:
                    fmt = ''.join(['lam ' + str(n.name) + '.' for n in self.arguments])
                    with FormatContext(object_format_type=False).as_default():
                        fmt += str(self.derived_expression)
                else:
                    fmt = 'def ' + self.name + '(' + ', '.join([str(x) for x in self.arguments]) + '): '
                    with FormatContext(object_format_type=False).as_default():
                        fmt += 'return ' + indent_text(str(self.derived_expression)).lstrip()
        else:
            if isinstance(self.ftype, OverloadedFunctionType):
                fmt = '\n'.join([f'{func_type}' for func_type in self.ftype.types])
                if self.name is not None:
                    fmt = re.sub(r'^' + re.escape(self.name) + ' (overloaded): ', '', fmt, flags=re.MULTILINE)
                    fmt = self.name + ': ' + '\n' + indent_text(fmt)
            else:
                fmt = f'{self.name}{self.ftype}'

        return fmt

    __repr__ = jacinle.repr_from_str

    def remap_arguments(self, remapping: List[int]) -> 'Function':
        """
        Generate a new Function object with a different argument order.
        Specifically, remapping is a permutation. The i-th argument to the new function will be the remapping[i]-th
        argument in the old function.

        Args:
            remapping: The remapping.

        Returns:
            The new function.
        """
        if isinstance(self.ftype, OverloadedFunctionType):
            raise NotImplementedError('Argument remapping for overloaded functions are not implemented.')

        new_argument_types = [self.ftype.argument_types[i] for i in remapping]

        def new_overridden_call(*args):
            remapped_args = [None for _ in range(len(args))]
            for i, arg in enumerate(args):
                remapped_args[remapping[i]] = arg
            return self(*remapped_args)

        return Function(
            self.name,
            FunctionType(
                new_argument_types, self.ftype.return_type,
            ), overridden_call=new_overridden_call, resolved_from=self.resolved_from
        )

    def partial(self, *args, execute_fully_bound_functions=False, **kwargs) -> Union['Function', 'FunctionApplicationExpression']:
        if self.name == '__lambda__':
            new_name = '__lambda__'
        else:
            new_name = f'{self.name}_partial'
        new_overridden_call = None
        new_resolved_from = None

        if self.is_overloaded:
            with FunctionArgumentResolutionContext(
                check_missing=False,
                check_overloaded_ambiguity=False
            ).as_default():
                types_and_arguments = self.ftype.resolve_type_and_args(*args, **kwargs)

            if not isinstance(types_and_arguments, OverloadedFunctionAmbiguousResolutions):
                ftype_index, function_type, resolved_args = types_and_arguments
                unmapped_arguments = [i for i, arg in enumerate(resolved_args) if arg is FunctionArgumentUnset]
                if len(unmapped_arguments) == 0:
                    return self._apply_with_resolved_args(resolved_args, ftype_index, function_type)
                new_type = _gen_partial_function_type(function_type, unmapped_arguments)
                new_resolved_from = FunctionResolvedFromRecord(self, ftype_index)
            else:
                # Block BEGIN {{{
                # If there is one specific resolution s.t. all variables are grounded, use it.

                all_grounded_resolutions = list()
                for ftype_index, function_type, resolved_args in types_and_arguments:
                    unmapped_arguments = [i for i, arg in enumerate(resolved_args) if arg is FunctionArgumentUnset]
                    if len(unmapped_arguments) == 0:
                        all_grounded_resolutions.append((ftype_index, function_type, resolved_args))

                if len(all_grounded_resolutions) == 1:
                    ftype_index, function_type, resolved_args = all_grounded_resolutions[0]
                    return self._apply_with_resolved_args(resolved_args, ftype_index, function_type)
                elif len(all_grounded_resolutions) > 1:
                    with get_function_argument_resolution_context().exc():
                        fmt = 'Got ambiguous application of overloaded function{}.\n'.format(
                            '' if self.name is None else ' ' + self.name)
                        fmt += 'Candidates are:\n'
                        for r in all_grounded_resolutions:
                            fmt += indent_text(str(r[1])) + '\n'
                        fmt += 'Invoked with arguments: {}.'.format(str(all_grounded_resolutions[0][2]))
                        raise FunctionArgumentResolutionError(fmt)

                # }}} Block END.

                possible_resolution_ids = list()
                possible_resolutions = list()
                possible_overridden_calls = list()
                for ftype_index, function_type, resolved_args in types_and_arguments:
                    unmapped_arguments = [i for i, arg in enumerate(resolved_args) if arg is FunctionArgumentUnset]
                    new_subtype = _gen_partial_function_type(
                        function_type, unmapped_arguments
                    )
                    possible_resolution_ids.append(ftype_index)
                    possible_resolutions.append(new_subtype)
                    possible_overridden_calls.append(_gen_partial_overriden_call(
                        new_subtype, resolved_args, self
                    ))
                new_type = OverloadedFunctionType(possible_resolutions)
                new_resolved_from = FunctionResolvedFromRecord(self, tuple(possible_resolution_ids))
                new_overridden_call = FunctionOverriddenCallList(possible_overridden_calls)
        else:
            with FunctionArgumentResolutionContext(check_missing=False).as_default():
                resolved_args = self.ftype.resolve_args(*args, **kwargs)
            unmapped_arguments = [i for i, arg in enumerate(resolved_args) if arg is FunctionArgumentUnset]

            if execute_fully_bound_functions:
                if len(unmapped_arguments) == 0:
                    return self._apply_with_resolved_args(resolved_args)

            new_type = _gen_partial_function_type(self.ftype, unmapped_arguments)
            new_overridden_call = _gen_partial_overriden_call(
                new_type, resolved_args, self
            )

        return Function(
            new_name, new_type,
            overridden_call=new_overridden_call, resolved_from=new_resolved_from
        )

    def _apply_with_resolved_args(
        self, resolved_args,
        resolved_ftype_id=None, resolved_function_type=None
    ):

        if self.overridden_call is not None:
            return self.get_overridden_call(resolved_ftype_id)(*resolved_args)

        if isinstance(self.ftype, OverloadedFunctionType):
            function = Function(
                self.name + f'_{resolved_ftype_id}',
                resolved_function_type,
                overridden_call=None,  # Must be none.
                resolved_from=FunctionResolvedFromRecord(self, resolved_ftype_id)
            )
        else:
            function = self

        from concepts.dsl.expression import FunctionApplicationExpression, cvt_expression_list
        return FunctionApplicationExpression(function, cvt_expression_list(resolved_args, function.ftype.argument_types))


def _gen_expression_from_overridden_call(ftype: FunctionType, overridden_call: Callable):
    from concepts.dsl.expression import ExpressionDefinitionContext
    ctx = ExpressionDefinitionContext()
    with ctx.as_default():
        arguments = [ctx[arg] for arg in ftype.arguments]
        return overridden_call(*arguments)


def _gen_partial_function_type(old_type: FunctionType, unmapped_arguments):
    new_argument_types = [old_type.argument_types[i] for i in unmapped_arguments]
    new_return_type = old_type.return_type
    new_argument_names = None
    if old_type.argument_names is not None:
        new_argument_names = [old_type.argument_names[i] for i in unmapped_arguments]
    new_type = FunctionType(new_argument_types, new_return_type, argument_names=new_argument_names)
    return new_type


def _gen_partial_overriden_call(new_type, resolved_args, call):
    assert isinstance(new_type, FunctionType)

    def partial_overriden_call(*new_args, **new_kwargs):
        new_resolved_args = new_type.resolve_args(*new_args, **new_kwargs)
        new_full_args = resolved_args.copy()
        j = 0
        for i in range(len(resolved_args)):
            if new_full_args[i] is FunctionArgumentUnset:
                new_full_args[i] = new_resolved_args[j]
                j += 1
        return call(*new_full_args)

    return partial_overriden_call

concepts/dsl/dsl_functions.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : dsl_types.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/09/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data structures for custom types in a DSL.

All types extend the base class :class:`TypeBase`.

:class:`TypeBase` has four derived types::

    - :class:`ObjectType`: corresponds to an object in the physical world.
    - :class:`ValueType`: corresponds to an encoding of a feature of the object.
    - :class:`~concepts.dsl.dsl_functions.FunctionType`: corresponds to a function over the encodings.
    - :class:`~concepts.dsl.dsl_functions.OverloadedFunctionType`: corresponds to a collection of functions that share the same name.

Furthermore, :class:`ValueType` has four derived types::

    - :class:`ConstantType`: corresponds to a constant.
    - :class:`SequenceType`: corresponds to a sequence of values.
    - :class:`TensorValueTypeBase`: corresponds to a tensor-based encoding.
    - :class:`PyObjValueType`: corresponds to an arbitrary Python object.
"""

import torch
from typing import TYPE_CHECKING, Optional, Union, Iterable, Sequence, Tuple

import jacinle
from jacinle.logging import get_logger
from jacinle.utils.defaults import option_context

if TYPE_CHECKING:
    from concepts.dsl.dsl_functions import FunctionType

logger = get_logger(__file__)

__all__ = [
    'FormatContext', 'get_format_context',
    'TypeBase', 'UnionType', 'ObjectType', 'ValueType',
    'AutoType', 'AnyType', 'UnionTyping', 'is_any_type', 'is_auto_type',  # Special types and type constructors.
    'TensorValueTypeBase', 'PyObjValueType', 'ConstantType', 'SequenceValueType',  # Derived from ValueType.
    'ScalarValueType', 'BOOL', 'INT64', 'FLOAT32', 'VectorValueType', 'NamedTensorValueType',  # Derived from TensorValueTypeBase.
    'QINDEX', 'Variable', 'ObjectConstant', 'UnnamedPlaceholder'  # Placeholder-like objects.
]


class FormatContext(option_context(
    '_FormatContext',
    type_format_cls=False, object_format_type=True, function_format_lambda=False, expr_max_length=120, type_varname_sep=': '
)):
    """FormatContext is a context manager that controls the format of types and objects."""

    type_format_cls: bool
    """If True, the type will be formatted as ``Type[typename]``."""

    object_format_type: bool
    """If True, the object format will include the type."""

    function_format_lambda: bool
    """If True, the function format will be lambda-function styled. Otherwise, it will be python-function styled."""

    expr_max_length: int
    """The maximum length of the expression per line."""

    type_varname_sep: str
    """The separator between the type and the variable name."""


get_format_context = FormatContext.get_default


class TypeBase(object):
    """Base class for all types."""

    def __init__(self, typename: str, alias: Optional[str] = None, parent_type: Optional['TypeBase'] = None):
        """Initialize the type.

        Args:
            typename: The name of the type.
            alias: The alias of the type.
        """

        self._typename = typename
        self._alias = alias
        self._parent_type = parent_type

    @property
    def typename(self) -> str:
        """The (full) typename of the type."""
        return self._typename

    @property
    def alias(self) -> Optional[str]:
        """An optional alias of the type."""
        return self._alias

    @property
    def parent_type(self) -> Optional['TypeBase']:
        """The parent type of the type."""
        return self._parent_type

    @property
    def parent_typename(self):
        """Return the typename of the parent type."""
        return self._parent_type.typename

    def __str__(self) -> str:
        if get_format_context().type_format_cls:
            return self.long_str()
        else:
            return self.short_str()

    __repr__ = jacinle.repr_from_str

    def short_str(self) -> str:
        """Return the short string representation of the type."""
        if self.alias is not None:
            return self.alias
        return self.typename

    def long_str(self) -> str:
        """Return the long string representation of the type."""
        return f'Type[{self.short_str()}]'

    def downcast_compatible(self, other: 'TypeBase') -> bool:
        """Check if the type is downcast-compatible with another type.
        That is, if the type is a subtype of the other type.
        """
        return self.typename == other.typename or self == AnyType or self == AutoType

    def __eq__(self, other: 'TypeBase') -> bool:
        return self.typename == other.typename

    def __ne__(self, other: 'TypeBase') -> bool:
        return not (self == other)

    def __hash__(self) -> int:
        return hash(self.typename)


"""AnyType corresponds to the union of all types."""
AnyType = TypeBase('AnyType')


"""AutoType means the type will be automatically inferred later."""
AutoType = TypeBase('AutoType')


def is_any_type(t: TypeBase) -> bool:
    """Check if the type is AnyType."""
    return t == AnyType


def is_auto_type(t: TypeBase) -> bool:
    """Check if the type is AutoType."""
    return t == AutoType


class UnionType(TypeBase):
    """The UnionType is a type that is the union of multiple types."""

    def __init__(self, *types: TypeBase, alias: Optional[str] = None):
        """Initialize the union type.

        Args:
            types: The types in the union.
            alias: The alias of the union type.
        """
        self.types = tuple(types)
        super().__init__(self.long_str(), alias=alias)

    types: Tuple[TypeBase, ...]
    """The underlying types of the union type."""

    def short_str(self) -> str:
        return ' | '.join(t.short_str() for t in self.types)

    def long_str(self) -> str:
        return 'Union[' + ', '.join(t.long_str() for t in self.types) + ']'

    def downcast_compatible(self, other: TypeBase) -> bool:
        raise NotImplementedError('Cannot downcast to a union type.')


class _UnionTypingSugar(object):
    def __getitem__(self, item: Tuple[TypeBase]):
        assert isinstance(item, tuple)
        return UnionType(*item)


UnionTyping = _UnionTypingSugar()


class ObjectType(TypeBase):
    """The ObjectType corresponds to the type of "real-world" objects."""

    def __init__(self, typename: str, parent_types: Optional[Sequence['ObjectType']] = None, alias: Optional[str] = None):
        """Initialize the object type.

        Args:
            typename: The name of the object type.
            alias: The alias of the object type.
        """
        super().__init__(typename, alias=alias)

        self.parent_types = tuple(parent_types) if parent_types is not None else tuple()

    parent_types: Tuple['ObjectType', ...]
    """The parent types of the object type."""

    def iter_parent_types(self) -> Iterable['ObjectType']:
        """Iterate over all parent types.

        Yields:
            the parent types following the inheritance order.
        """
        for parent_type in self.parent_types:
            yield parent_type
            yield from parent_type.iter_parent_types()

    def long_str(self) -> str:
        if len(self.parent_types) == 0:
            return f'OT[{self.typename}]'

        return f'OT[{self.typename}, parent={", ".join(t.typename for t in self.parent_types)}]'


class ValueType(TypeBase):
    """The ValueType corresponds to a value of a certain type."""

    def long_str(self) -> str:
        return f'VT[{self.typename}]'

    def assignment_type(self) -> 'ValueType':
        """Return the value type for assignment."""
        return self


class TensorValueTypeBase(ValueType):
    """A value type refers to a type of some features associated with physical objects."""

    def __init__(self, typename: str, quantized: bool, alias: Optional[str] = None):
        super().__init__(typename, alias)
        self._quantized = quantized

    @property
    def quantized(self) -> bool:
        """Whether the tensor is quantized."""
        return self._quantized

    def ndim(self) -> int:
        """Return the number of dimensions of the value type."""
        raise NotImplementedError()

    def size(self) -> int:
        """Return the total size of the value type."""
        raise NotImplementedError()

    def size_tuple(self) -> Tuple[int]:
        """Return the size of the value type as a tuple."""
        raise NotImplementedError()

    def tensor_dtype(self) -> torch.dtype:
        """Return the corresponding PyTorch tensor dtype."""
        raise NotImplementedError()

    def is_intrinsically_quantized(self) -> bool:
        """Return whether the value type can be quantized.
        In this project, quantizable means we can transform the value into a single integer scalar.
        Therefore, only BOOL-valued and INTxx-valued tensors are quantizable.

        Returns:
            True if the value type can be quantized.
        """
        raise NotImplementedError()

    @staticmethod
    def from_tensor(tensor: torch.Tensor) -> 'TensorValueTypeBase':
        """Create a value type from a pytorch tensor.

        Args:
            tensor: the tensor.

        Returns:
            the tensor value type.
        """
        if tensor.dtype in (torch.float32, torch.float64):
            return FLOAT32
        elif tensor.dtype in (torch.int32, torch.int64):
            return INT64
        elif tensor.dtype == torch.bool:
            return BOOL
        else:
            raise TypeError(f'Unsupported tensor type: {tensor.dtype}.')


class PyObjValueType(ValueType):
    def __init__(self, pyobj_type: Union[type, str], alias: Optional[str] = None):
        self.pyobj_type = pyobj_type
        super().__init__(str(pyobj_type), alias=alias)

    pyobj_type: Union[type, str]
    """The underlying Python object type."""


class SequenceValueType(ValueType):
    """The SequenceValueType corresponds to a sequence of values of a certain type."""

    def __init__(self, value_type: ValueType, alias: Optional[str] = None):
        """Initialize the sequence value type.

        Args:
            value_type: The value type of the sequence.
        """
        super().__init__(f'Sequence[{value_type.typename}]', alias=alias)
        self.value_type = value_type

    value_type: ValueType
    """The value type of the sequence."""

    def long_str(self) -> str:
        return f'SequenceT[{self.value_type.typename}]'

    def assignment_type(self) -> 'ValueType':
        """Return the value type for assignment."""
        return self


class ConstantType(ValueType):
    """The ConstantType corresponds to a constant value."""

    def long_str(self) -> str:
        return f'CT[{self.typename}]'


class ScalarValueType(TensorValueTypeBase):
    TENSOR_DTYPE_MAPPING = {
        'int32': torch.int32,
        'int64': torch.int64,
        'float32': torch.float32,
        'float64': torch.float64,
        'bool': torch.int64,
    }

    TENSOR_DTYPE_QUANTIZED = {
        'int32': True,
        'int64': True,
        'float32': False,
        'float64': False,
        'bool': True
    }

    def __init__(self, typename):
        if typename not in type(self).TENSOR_DTYPE_MAPPING:
            raise TypeError(f'Unknown scalar value type: {typename}.')
        super().__init__(typename, quantized=type(self).TENSOR_DTYPE_QUANTIZED[typename])

    def ndim(self) -> int:
        return 0

    def size(self) -> int:
        return 1

    def size_tuple(self) -> Tuple[int]:
        return tuple()

    def tensor_dtype(self) -> torch.dtype:
        return type(self).TENSOR_DTYPE_MAPPING[self.typename]

    INTRINSICALLY_QUANTIZED_DTYPES = ['bool', 'int32', 'int64']

    def is_intrinsically_quantized(self) -> bool:
        return self.typename in self.INTRINSICALLY_QUANTIZED_DTYPES

    def __eq__(self, o: object) -> bool:
        if isinstance(o, ScalarValueType):
            return self.typename == o.typename
        return False

    def __hash__(self):
        return hash(self.typename)


BOOL = ScalarValueType('bool')
INT64 = ScalarValueType('int64')
FLOAT32 = ScalarValueType('float32')


class VectorValueType(TensorValueTypeBase):
    def __init__(self, dtype: ScalarValueType, dim: int = 0, choices: int = 0, factors: int = 1, alias: Optional[str] = None):
        assert isinstance(dtype, ScalarValueType), 'Currently only support 1D vectors.'
        self.dtype = dtype
        self.dim = dim
        self.choices = choices
        self.factors = factors

        typename = f'vector[{self.dtype}, dim={self.dim}, choices={self.choices}, factors={self.factors}]'
        quantized = choices > 0
        super().__init__(typename, quantized=quantized, alias=alias)

    dtype: ScalarValueType
    """The scalar value type of the vector."""

    dim: int
    """The dimension of the vector."""

    choices: int
    """The number of choices for vector values per factor."""

    factors: int
    """The number of factors of the vector."""

    def ndim(self) -> int:
        return 1

    def size(self) -> int:
        return self.dim * self.factors

    def size_tuple(self) -> Tuple[int]:
        return (self.size(), )

    def assignment_type(self) -> 'VectorValueType':
        if self._quantized:
            return VectorValueType(FLOAT32, self.choices, 0, self.factors)
        return self

    def tensor_dtype(self) -> torch.dtype:
        return self.dtype.tensor_dtype()

    def is_intrinsically_quantized(self) -> bool:
        return False

    def long_str(self):
        if self.alias is not None:
            return self.alias + f'({self.typename})'
        return self.typename


class NamedTensorValueType(TensorValueTypeBase):
    def __init__(self, typename: str, base_type: TensorValueTypeBase):
        super().__init__(typename, base_type._quantized)
        self.base_type = base_type

    base_type: TensorValueTypeBase
    """The base tensor value type."""

    def ndim(self) -> int:
        return self.base_type.ndim()

    def size(self) -> int:
        return self.base_type.size()

    def size_tuple(self) -> Tuple[int]:
        return self.base_type.size_tuple()

    def assignment_type(self) -> TensorValueTypeBase:
        return self

    def tensor_dtype(self) -> torch.dtype:
        return self.base_type.tensor_dtype()

    def is_intrinsically_quantized(self) -> bool:
        return self.base_type.is_intrinsically_quantized()

    def __eq__(self, o: object) -> bool:
        if isinstance(o, NamedTensorValueType):
            return self.typename == o.typename
        return False

    def __hash__(self):
        return hash(self.typename)


QINDEX = slice(None)
"""The "quantified index" object, used to indicate a quantified dimension in a tensor value type. Internally it is `slice(None)`."""


class _Placeholder(object):
    def __init__(self, name: str, dtype: Union[ObjectType, ValueType, 'FunctionType']):
        self.name = name
        self.dtype = dtype

    name: str
    """The name of the placeholder."""

    dtype: Union[ObjectType, ValueType, 'FunctionType']
    """The data type of the placeholder."""

    @property
    def typename(self):
        return self.dtype.typename

    def __str__(self):
        sep = get_format_context().type_varname_sep
        return f'{self.name}{sep}{self.dtype}'

    __repr__ = jacinle.repr_from_str


class Variable(_Placeholder):
    """The class representing a variable in a function."""

    name: str
    """The name of the variable."""

    dtype: Union[ObjectType, ValueType, 'FunctionType']
    """The data type of the variable."""

    quantifier_flag: Optional[str] = None
    """Additional quantifier flag for the variable. This flag will be set by the parser indicating the quantifier scope of this variable.
    Currently this is only used in PDSketch for "quantified variable" in regression rules. """

    def set_quantifier_flag(self, flag: str):
        self.quantifier_flag = flag

    def __str__(self):
        if self.quantifier_flag is None:
            return super().__str__()
        else:
            return f'({self.quantifier_flag} {super().__str__()})'


class ObjectConstant(_Placeholder):
    """The class representing a constant object in a DSL."""

    name: str
    """The name of the constant."""

    dtype: ObjectType
    """The data type of the constant."""


class UnnamedPlaceholder(_Placeholder):
    def __init__(self, dtype: Union[ObjectType, ValueType]):
        super().__init__('_', dtype)

    name: str
    """The name of the unnamed placeholder. Always equal to '_'."""

    dtype: Union[ObjectType, ValueType]
    """The data type of the unnamed placeholder."""

    def __str__(self):
        return '??' + str(self.dtype)


concepts/dsl/dsl_types.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""The baseclass for all executors of domain-specific languages and a few implementations for simple function domains and tensor-based value representations."""

concepts/dsl/executors/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : executor_base.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""The baseclass for all executors of domain-specific languages."""

from abc import ABC, abstractmethod
from typing import Any, Dict, Callable

from concepts.dsl.dsl_domain import DSLDomainBase
from concepts.dsl.expression import Expression

__all__ = ['DSLExecutorBase', 'DSLExecutionError']


class DSLExecutorBase(ABC):
    """The baseclass for all executors of domain-specific languages."""

    def __init__(self, domain: DSLDomainBase):
        """Initialize the executor.

        Args:
            domain: the domain of the executor.
        """

        self._domain = domain
        self._function_implementations = dict()

    @property
    def domain(self) -> DSLDomainBase:
        """The domain of the executor."""
        return self._domain

    @property
    def function_implementations(self) -> Dict[str, Callable]:
        """The implementations of functions, which is a mapping from function names to implementations."""
        return self._function_implementations

    def register_function(self, name: str, func: Callable):
        """Register an implementation for a function to the executor. Alias for :meth:`register_function_implementation`.

        Args:
            name: the name of the function.
            func: the implementation of the function.
        """
        self.register_function_implementation(name, func)

    def register_function_implementation(self, name: str, func: Callable):
        """Register an implementation for a function.

        Args:
            name: the name of the function.
            func: the implementation of the function.
        """
        self._function_implementations[name] = func

    def get_function_implementation(self, name: str) -> Callable:
        """Get the implementation of a function. When the executor does not have an implementation for the function,
        the implementation of the function in the domain will be returned. If that is also None, a `KeyError` will be
        raised.

        Args:
            name: the name of the function.

        Returns:
            the implementation of the function.
        """

        if name in self._function_implementations:
            return self._function_implementations[name]
        elif name in self._domain.functions:
            function_body = self._domain.functions[name].function_body
            if function_body is None:
                raise KeyError(f'No implementation for function {name}.')
            return function_body
        raise KeyError(f'No implementation for function {name}.')

    def has_function_implementation(self, name: str) -> bool:
        """Check whether the executor has an implementation for a function.

        Args:
            name: the name of the function.

        Returns:
            whether the executor has an implementation for the function.
        """
        return name in self._function_implementations or (name in self.domain.functions and self.domain.functions[name].function_body is not None)

    @abstractmethod
    def execute(self, expression: Expression) -> Any:
        """Execute an expression.

        Args:
            expression: the expression to execute.

        Returns:
            the result of the execution.
        """
        raise NotImplementedError()


class DSLExecutionError(Exception):
    """The exception raised when an execution error occurs."""


concepts/dsl/executors/executor_base.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : function_domain_executor.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""A simple executor for a domain that only contains function applications. See the tutorial for more details."""

import functools
import contextlib
from typing import Any, Optional, Union, Callable
from jacinle.logging import get_logger

from concepts.dsl.dsl_functions import FunctionType, Function
from concepts.dsl.value import ValueBase, Value
from concepts.dsl.expression import Expression, FunctionApplicationExpression, ConstantExpression
from concepts.dsl.function_domain import FunctionDomain
from concepts.dsl.parsers.parser_base import ParserBase
from concepts.dsl.parsers.function_expression_parser import FunctionExpressionParser
from concepts.dsl.executors.executor_base import DSLExecutorBase

logger = get_logger(__file__)

__all__ = ['FunctionDomainExecutor']


class FunctionDomainExecutor(DSLExecutorBase):
    """The executor for simple function DSLs."""

    def __init__(self, domain: FunctionDomain, parser: Optional[ParserBase] = None):
        """Initialize the executor.

        Args:
            domain: the function domain.
            parser: the parser of the DSL. If not specified, the default
                :class:`~concepts.dsl.parsers.function_expression_parser.FunctionExpressionParser` will be used.
        """
        super().__init__(domain)
        self.parser = parser

        for function_name, function in domain.functions.items():
            assert not function.is_overloaded
            if not function.is_derived:
                if hasattr(self, function_name):
                    self.register_function(function_name, self.unwrap_values(getattr(self, function_name)))
                    logger.info('Function {} automatically registered.'.format(function_name))

        self._grounding = None

    _domain: FunctionDomain

    @property
    def domain(self) -> FunctionDomain:
        """The function domain of the executor."""
        return self._domain

    @property
    def grounding(self) -> Any:
        """The grounding of the current execution."""
        return self._grounding

    parser: ParserBase
    """The parser of the DSL."""

    @contextlib.contextmanager
    def with_grounding(self, grounding: Any):
        """A context manager for setting the grounding of the executor.

        Args:
            grounding: the grounding to set.
        """
        old_grounding = self.grounding
        self._grounding = grounding
        try:
            yield
        finally:
            self._grounding = old_grounding

    def parse_expression(self, string: str) -> Expression:
        """Parse an expression from a string.

        Args:
            string: the string to parse.

        Returns:
            The parsed expression.
        """
        if self.parser is None:
            self.parser = FunctionExpressionParser(self.domain)
        return self.parser.parse_expression(string)

    def execute(self, expr: Union[str, Expression], grounding: Optional[Any] = None) -> Value:
        """Execute an expression.

        Args:
            expr: the expression to execute.
            grounding: the grounding of the expression.

        Returns:
            The result of the execution.
        """

        if isinstance(expr, str):
            expr = self.parse_expression(expr)

        grounding = grounding if grounding is not None else self._grounding
        with self.with_grounding(grounding):
            return self._execute(expr)

    def execute_function(self, function: Function, *args: Value, grounding: Optional[Any] = None) -> Value:
        """Execute a function with a list of arguments.

        Args:
            function: the function to execute.
            *args: the arguments of the function.
            grounding: the grounding of the function.

        Returns:
            The result of the execution.
        """

        expression = function(*args)
        return self.execute(expression, grounding)

    def _execute(self, expr: Expression) -> Value:
        """Internal implementation of the executor. This method will be called by the public method :meth:`execute`.
        This function basically implements a depth-first search on the expression tree.

        Args:
            expr: the expression to execute.

        Returns:
            The result of the execution.
        """
        if isinstance(expr, FunctionApplicationExpression):
            func = self.function_implementations[expr.function.name]
            args = [self._execute(arg) for arg in expr.arguments]
            return func(*args)
        elif isinstance(expr, ConstantExpression):
            assert isinstance(expr.constant, Value)
            return expr.constant
        else:
            raise NotImplementedError(f'Unknown expression type: {type(expr)}')

    def unwrap_values(self, func_or_ftype: Union[Callable, FunctionType]) -> Callable:
        """A function decorator that automatically unwraps the values of the arguments of the function.
        Basically, this decorator will unwrap the values of the arguments of the function, and then wrap the result with the
        :class:`concepts.dsl.value.Value` class.

        There are two ways to use this decorator. The first way is to use it as a function decorator:
        In this case, the wrapped function should have the same name as the DSL function it implements.

            >>> domain = FunctionDomain()
            >>> # Assume domain has a function named "add" with two arguments.
            >>> executor = FunctionDomainExecutor(domain)
            >>> @executor.unwrap_values
            >>> def add(a, b):
            >>>     return a + b
            >>> executor.register_function('add', add)

        The second way is to use it as function that generates a function decorator:

            >>> domain = FunctionDomain()
            >>> # Assume domain has a function named "add" with two arguments.
            >>> executor = FunctionDomainExecutor(domain)
            >>> @executor.unwrap_values(domain.functions['add'].ftype)
            >>> def add(a, b):
            >>>     return a + b
            >>> executor.register_function('add', executor.unwrap_values(add))

        Args:
            func_or_ftype: the function to wrap, or the function type of the function to wrap.

        Returns:
            The decorated function or a function decorator.
        """

        if isinstance(func_or_ftype, FunctionType):
            ftype = func_or_ftype
        else:
            if func_or_ftype.__name__ not in self.domain.functions:
                raise NameError(f'Function {func_or_ftype.__name__} is not registered in the domain.')
            ftype = self.domain.functions[func_or_ftype.__name__].ftype

        def wrapper(func):
            @functools.wraps(func)
            def wrapped(*args, **kwargs):
                args = [arg.value if isinstance(arg, Value) else arg for arg in args]
                kwargs = {k: v.value if isinstance(v, Value) else v for k, v in kwargs.items()}
                rv = func(*args, **kwargs)

                if isinstance(ftype.return_type, tuple):
                    return tuple(
                        Value(ftype.return_type[i], rv[i]) if not isinstance(rv[i], ValueBase) else rv[i] for i in range(len(rv))
                    )
                elif isinstance(rv, ValueBase):
                    return rv
                else:
                    return Value(ftype.return_type, rv)
            return wrapped

        if isinstance(func_or_ftype, FunctionType):
            return wrapper
        else:
            return wrapper(func_or_ftype)


concepts/dsl/executors/function_domain_executor.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : tensor_value_executor.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 11/03/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Tensor-based expression executor.

The high-level interface for tensor-based expression is that we can execute an expression with a given state and a set of
bounded variables. The executor will return a tensor value.

The state is represented using :class:`concepts.dsl.tensor_state.TensorState` or :class:`concepts.dsl.tensor_state.NamedObjectTensorState`, which internally stores a dictionary
mapping from string (the state variable name, e.g., ``is_hot``) to a :class:`concepts.dsl.tensor_value.TensorValue` class.

The bounded variables are essentially a dictionary mapping from strings (the variable name, e.g., ``x``) to its value. There are
two types of values: (1) a :class:`concepts.dsl.tensor_value.TensorValue` class, which represents an actual value (e.g., a vector representation);
(2) a :class:`StateObjectReference` instance or a QINDEX (a.k.a., ``slice(None)``), which represents a reference to an object in the state.

With the bounded variables, the expressions can have variables, which are essentially placeholders for the actual values. For example,

.. code-block:: python

    domain = FunctionDomain()
    # Define an object type `person`.
    domain.define_type(ObjectType('person'))
    # Define a state variable `is_friend` with type `person x person -> bool`.
    domain.define_function(Function('is_friend', FunctionType([ObjectType('person'), ObjectType('person')], BOOL)))

    x = VariableExpression(Variable('x', ObjectType('person')))
    y = VariableExpression(Variable('y', ObjectType('person')))
    relation = FunctionApplication(domain.functions['is_friend'], [x, y])

Then we can execute the expression with a given state and bounded variables:

.. code-block:: python

    # See the documentation for namedObjectTensorState for more details.
    state = NamedObjectTensorState({
        'is_friend': TensorValue(BOOL, ['x', 'y'], torch.tensor([[1, 1, 1], [1, 1, 0], [1, 0, 1]], dtype=torch.bool))
    }, object_names={
        'Alice': ObjectType('person'),
        'Bob': ObjectType('person'),
        'Charlie': ObjectType('person'),
    })
    executor = SimpleFunctionTensorValueExecutor(domain)

    # For both of the following lines, the result is a tensor value with value `True`.
    # Use the constructed expression:
    executor.execute(relation, state, {'x': 'Alice', 'y': 'Bob'})
    # To use the default parser:
    executor.execute('is_friend(x, y)', state, {'x': 'Alice', 'y': 'Bob'})
"""

import contextlib
from typing import Optional, Union, Tuple, Sequence, Dict

import torch

from concepts.dsl.dsl_types import ObjectType, NamedTensorValueType, PyObjValueType, Variable, UnnamedPlaceholder, QINDEX
from concepts.dsl.dsl_domain import DSLDomainBase
from concepts.dsl.function_domain import FunctionDomain
from concepts.dsl.tensor_value import TensorValue, scalar
from concepts.dsl.tensor_state import StateObjectReference, TensorState, NamedObjectTensorState
from concepts.dsl.expression import Expression, VariableExpression, ObjectConstantExpression, ConstantExpression, FunctionApplicationExpression
from concepts.dsl.constraint import OptimisticValue, ConstraintSatisfactionProblem, OPTIM_MAGIC_NUMBER
from concepts.dsl.parsers.parser_base import ParserBase
from concepts.dsl.parsers.function_expression_parser import FunctionExpressionParser
from concepts.dsl.executors.executor_base import DSLExecutorBase
from concepts.dsl.executors.value_quantizers import ValueQuantizer, PyObjectStore

__all__ = [
    'BoundedVariablesDict', 'BoundedVariablesDictCompatible',
    'compose_bvdict', 'compose_bvdict_args', 'get_bvdict',
    'TensorValueExecutorReturnType', 'TensorValueExecutorBase', 'FunctionDomainTensorValueExecutor'
]


BoundedVariablesDict = Dict[str, Dict[str, Union[StateObjectReference, slice, TensorValue]]]
"""Internal representation of a bounded variable dictionary. It stores a nested two-layer dictionary, where the first layer
stores the type of the object, and the second layer stores the name of the object. The value can be either a :class:`concepts.dsl.tensor_value.TensorValue`
or a :class:`StateObjectReference` instance (representing the reference to a single object)."""

BoundedVariablesDictCompatibleKeyType = Union[str, Variable]
BoundedVariablesDictCompatibleValueType = Union[str, int, slice, bool, float, torch.Tensor, TensorValue, StateObjectReference]

BoundedVariablesDictCompatible = Union[
    None, Sequence[Variable],
    Dict[BoundedVariablesDictCompatibleKeyType, BoundedVariablesDictCompatibleValueType],
    BoundedVariablesDict
]
"""Compatible types with :class:`BoundedVariablesDict`. They can be converted to :class:`BoundedVariablesDict` using :func:`compose_bvdict`."""


def compose_bvdict(input_dict: BoundedVariablesDictCompatible, state: Optional[TensorState] = None) -> BoundedVariablesDict:
    """Compose a bounded variable dict from raw inputs.

    Args:
        input_dict: the input dict. There are three types of inputs:

            1. A sequence of :class:`concepts.dsl.dsl_types.Variable` instances, which represents a set of variables with no values.
            2. A dictionary mapping from :class:`concepts.dsl.dsl_types.Variable` instances to the actual value.
            3. A dictionary mapping from strings (the name of the variables) to values.

            Acceptable values are:

            1. A :class:`str`, which represents a reference to an object in the state (so the state must be object-named).
            2. An integer, which represents a reference to an object in the state (so the state must be object-named).
            3. A QINDEX (a.k.a., ``slice(None)``), which represents all objects in the state of a given type (so the state must be object-named).
            4. A :class:`concepts.dsl.tensor_value.TensorValue` instance, which represents an actual value.
            5. A :class:`StateObjectReference` instance, which represents a reference to an object in the state (so the state must be object-named).
            6. A :class:`bool`, :class:`int`, :class:`float`, or :class:`torch.Tensor` instance, which represents an actual value. They will be converted to a :class:`concepts.dsl.tensor_value.TensorValue` instance.

        state: the state.

    Returns:
        a dictionary mapping from strings (the typename) to a dictionary mapping from strings (the name of the variables) to values.
    """

    if input_dict is None:
        return dict()

    if isinstance(input_dict, dict):
        if len(input_dict) == 0:
            return input_dict

        sample_value = next(iter(input_dict.values()))
        if isinstance(sample_value, dict):
            return input_dict

        output_dict = dict()
        for var, value in input_dict.items():
            if isinstance(var, Variable):
                # Part 1: the variable corresponds to an object.
                if isinstance(var.dtype, ObjectType):
                    if isinstance(value, int):
                        assert isinstance(state, NamedObjectTensorState)
                        value = StateObjectReference(state.object_type2name[var.dtype.typename][value], value)
                        output_dict.setdefault(var.typename, dict()).setdefault(var.name, value)
                    elif isinstance(value, str):
                        assert isinstance(state, NamedObjectTensorState)
                        value = StateObjectReference(value, state.get_typed_index(value))
                        output_dict.setdefault(var.typename, dict()).setdefault(var.name, value)
                    elif isinstance(value, slice):
                        output_dict.setdefault(var.typename, dict()).setdefault(var.name, value)
                    elif isinstance(value, StateObjectReference):
                        output_dict.setdefault(var.typename, dict()).setdefault(var.name, value)
                    else:
                        raise TypeError(f'Invalid value type for variable {var}: {type(value)}.')
                # Part 2: the variable corresponds to a Python object.
                elif isinstance(var.dtype, PyObjValueType):
                    if isinstance(value, TensorValue):
                        pass
                    else:
                        value = TensorValue.from_scalar(value, var.dtype)
                    typename = var.dtype.typename
                    output_dict.setdefault(typename, {})[var.name] = value
                # Part 3: the variable corresponds to a PyTorch tensor.
                elif isinstance(var.dtype, NamedTensorValueType):
                    if isinstance(value, TensorValue):
                        pass
                    elif isinstance(value, (bool, int, float, torch.Tensor)):
                        value = TensorValue.from_scalar(value, var.dtype)
                    elif isinstance(value, UnnamedPlaceholder):
                        value = TensorValue.from_optimistic_value_int(OPTIM_MAGIC_NUMBER, var.dtype)  # Just a placeholder.
                    else:
                        raise TypeError(f'Invalid value type for variable {var}: {type(value)}.')
                    output_dict.setdefault(var.dtype.typename, {})[var.name] = value
                else:
                    raise TypeError(f'Invalid variable type: {var.dtype}.')
            elif isinstance(var, OptimisticValue):
                raise RuntimeError('Invalid branch; OptimisticValue should be handled in the previous branch. Report a bug to the developers.')
            elif isinstance(var, str) and isinstance(value, str):
                assert state is not None
                typename, value_index = state.get_typename(value), state.get_typed_index(value)
                value = StateObjectReference(value, value_index)
                output_dict.setdefault(typename, dict()).setdefault(var, value)
            else:
                raise TypeError(f'Invalid KV pair: {var} -> {value}.')
        return output_dict
    elif isinstance(input_dict, (list, tuple)):
        # The input dict is a list of variables.
        assert isinstance(input_dict, (list, tuple))
        output_dict = dict()
        for var in input_dict:
            assert isinstance(var, Variable)
            output_dict.setdefault(var.typename, dict()).setdefault(var.name, QINDEX)
        return output_dict
    else:
        raise TypeError(f'Invalid input type: {type(input_dict)}.')


def compose_bvdict_args(arguments_def: Sequence[Variable], arguments: Sequence[BoundedVariablesDictCompatibleValueType], state: Optional[TensorState] = None) -> BoundedVariablesDict:
    """Compose a bounded variable dict, but from a list of arguments. This function is useful when we want to compose a bounded variable dict from a list of arguments to a function.

    Args:
        arguments_def: the definition of the arguments, including their name and dtypes.
        arguments: the actual arguments.
        state: the state.

    Returns:
        a bounded variable dictionary.
    """
    return compose_bvdict(dict(zip(arguments_def, arguments)), state=state)


def get_bvdict(bvdict: BoundedVariablesDict, variable: Variable) -> Union[StateObjectReference, slice, TensorValue]:
    """Get the value of a variable from a bounded variable dict.

    Args:
        bvdict: the bounded variable dict.
        variable: the variable.

    Returns:
        the value of the variable.
    """
    return bvdict[variable.typename][variable.name]


TensorValueExecutorReturnTypeElem = Union[TensorValue, slice, StateObjectReference, None]
TensorValueExecutorReturnType = Union[TensorValueExecutorReturnTypeElem, Tuple[TensorValueExecutorReturnTypeElem, ...]]


class TensorValueExecutorBase(DSLExecutorBase):
    """The base class for tensor value executors."""

    def __init__(self, domain: DSLDomainBase, parser: Optional[ParserBase] = None):
        """Initialize the base class for tensor value executors.

        Args:
            domain: the domain of the executor.
            parser: the parser to use. If None, no parser will be used.
        """
        super().__init__(domain)

        self._parser = parser

        self._state = None
        self._csp = None
        self._bounded_variables = dict()
        self._optimistic_execution = False
        self._value_quantizer = ValueQuantizer(self)
        self._pyobj_store = PyObjectStore(self)

    @property
    def parser(self) -> Optional[ParserBase]:
        """The parser for the domain."""
        return self._parser

    @property
    def state(self) -> Optional[TensorState]:
        """The current state of the environment."""
        return self._state

    @property
    def csp(self) -> Optional[ConstraintSatisfactionProblem]:
        """The CSP that describes the constraints in past executions."""
        return self._csp

    @property
    def optimistic_execution(self) -> bool:
        """Whether to execute the expression optimistically (i.e., treat all CSP constraints True)."""
        return self._optimistic_execution

    @property
    def bounded_variables(self) -> BoundedVariablesDict:
        """The bounded variables for the execution. Note that most of the time you should use the :meth:`get_bounded_variable` method to get values for the bounded variable."""
        return self._bounded_variables

    @property
    def value_quantizer(self) -> ValueQuantizer:
        """The value quantizer."""
        return self._value_quantizer

    @property
    def pyobj_store(self) -> PyObjectStore:
        """The Python object store."""
        return self._pyobj_store

    @contextlib.contextmanager
    def with_state(self, state: Optional[TensorState] = None):
        """A context manager to temporarily set the state of the executor."""
        old_state = self._state
        self._state = state
        yield
        self._state = old_state

    @contextlib.contextmanager
    def with_csp(self, csp: ConstraintSatisfactionProblem):
        """A context manager to temporarily set the CSP of the executor."""
        old_csp = self._csp
        self._csp = csp
        yield
        self._csp = old_csp

    @contextlib.contextmanager
    def with_bounded_variables(self, bvdict: BoundedVariablesDictCompatible):
        """A context manager to set the bounded variables for the executor.

        Args:
            bvdict: the bounded variables.
        """
        old_bvdict = self._bounded_variables
        self._bounded_variables = compose_bvdict(bvdict, state=self._state)
        yield
        self._bounded_variables = old_bvdict

    @contextlib.contextmanager
    def new_bounded_variables(self, bvdict: BoundedVariablesDictCompatible):
        """A context manager to add additional bounded variables to the executor.

        Args:
            bvdict: the new bounded variables.
        """
        bvdict = compose_bvdict(bvdict, state=self._state)
        for typename, variables in bvdict.items():
            for name, value in variables.items():
                if typename not in self._bounded_variables:
                    self._bounded_variables[typename] = dict()
                assert name not in self._bounded_variables[typename], f'Variable {name} already exists in bounded variables.'
                self._bounded_variables[typename][name] = value
        yield
        for typename, variables in bvdict.items():
            for name in variables:
                del self._bounded_variables[typename][name]

    def get_bounded_variable(self, variable: Variable) -> Union[TensorValue, slice, StateObjectReference]:
        """Get the value of a bounded variable.

        Args:
            variable: the variable.

        Returns:
            the value of the variable.
        """
        return get_bvdict(self._bounded_variables, variable)

    def set_value_quantizer(self, value_quantizer: ValueQuantizer):
        """Set the value quantizer for the executor.

        Args:
            value_quantizer: the value quantizer.
        """
        self._value_quantizer = value_quantizer

    def reset_value_quantizer(self):
        """Reset the value quantizer to the default value quantizer."""
        self._value_quantizer = ValueQuantizer(self)

    def set_pyobj_store(self, pyobj_store: PyObjectStore):
        """Set the Python object store for the executor.

        Args:
            pyobj_store: the Python object store.
        """
        self._pyobj_store = pyobj_store

    def reset_pyobj_store(self):
        """Reset the Python object store."""
        self._pyobj_store = PyObjectStore(self)

    def set_parser(self, parser: ParserBase):
        """Set the parser for the executor.

        Args:
            parser: the parser.
        """
        self._parser = parser

    def parse(self, expression: Union[Expression, str]):
        """Parse an expression.

        Args:
            expression: the expression to parse. When the input is already an expression, it will be returned directly.

        Returns:
            the parsed expression.
        """
        if isinstance(expression, Expression):
            return expression
        if self._parser is None:
            raise ValueError('No parser is set for the executor.')
        return self._parser.parse_expression(expression)

    def execute(
        self, expression: Union[Expression, str],
        state: Optional[TensorState] = None,
        bounded_variables: Optional[BoundedVariablesDictCompatible] = None,
        csp: Optional[ConstraintSatisfactionProblem] = None,
        optimistic_execution: bool = False
    ) -> TensorValueExecutorReturnType:
        """Execute an expression.

        Args:
            expression: the expression to execute.
            state: the state to use. If None, the current state of the executor will be used.
            bounded_variables: the bounded variables to use. If None, the current bounded variables of the executor will be used.
            csp: the constraint satisfaction problem to use. If None, the current CSP of the executor will be used.
            optimistic_execution: whether to execute the expression optimistically (i.e., treat all CSP constraints True).

        Returns:
            the TensorValue object.
        """
        if isinstance(expression, str):
            expression = self.parse(expression)

        state = state if state is not None else self._state
        csp = csp if csp is not None else self._csp
        bounded_variables = bounded_variables if bounded_variables is not None else self._bounded_variables
        with self.with_state(state), self.with_csp(csp), self.with_bounded_variables(bounded_variables):
            try:
                self._optimistic_execution, backup = optimistic_execution, self._optimistic_execution
                return self._execute(expression)
            finally:
                self._optimistic_execution = backup

    def _execute(self, expression: Expression) -> TensorValueExecutorReturnType:
        raise NotImplementedError()

    @contextlib.contextmanager
    def checkpoint_storage(self):
        """Checkpoint the storages of the executor, including the value quantizer and the Python object store."""
        with self.value_quantizer.checkpoint(), self.pyobj_store.checkpoint():
            yield


class FunctionDomainTensorValueExecutor(TensorValueExecutorBase):
    """Similar to :class:`~concepts.dsl.executors.function_domain_executor.FunctionDomainExecutor`, but works for :class:`~concepts.dsl.tensor_value.TensorValue`.
    The two of the main differences are:

    1. The :meth:`execute` method returns a :class:`~concepts.dsl.tensor_value.TensorValue` object instead of a :class:`~concepts.dsl.value.Value` object.
    2. The class supports binding variables to values during execution. See the documentation for this file and tutorials for details.
    """

    def __init__(self, domain: FunctionDomain, parser: Optional[ParserBase] = None):
        """Initialize a tensor value executor for a function domain.

        Args:
            domain: the domain of the executor.
            parser: the parser to use. If not specified, no parser will be used.
        """

        if parser is None:
            parser = FunctionExpressionParser(domain, allow_variable=True, escape_string=True)

        super().__init__(domain, parser)

    _domain: FunctionDomain

    @property
    def domain(self) -> FunctionDomain:
        """The function domain of the executor."""
        return self._domain

    def _execute(self, expr: Expression) -> TensorValueExecutorReturnType:
        if isinstance(expr, VariableExpression):
            variable = expr.variable
            return self._bounded_variables[variable.dtype.typename][variable.name]
        elif isinstance(expr, ObjectConstantExpression):
            constant = expr.constant
            assert isinstance(self._state, NamedObjectTensorState)
            return StateObjectReference(
                constant.name,
                self._state.get_typed_index(constant.name, constant.dtype.typename)
            )
        elif isinstance(expr, ConstantExpression):
            assert isinstance(expr.constant, TensorValue)
            return expr.constant
        elif isinstance(expr, FunctionApplicationExpression):
            func = expr.function
            args = [self._execute(arg) for arg in expr.arguments]
            if func.name in self._state.features:
                args = [arg.index if isinstance(arg, StateObjectReference) else arg for arg in args]
                return self._state.features[func.name][tuple(args)]
            else:
                assert self.has_function_implementation(func.name)
                return self.get_function_implementation(func.name)(*args)


concepts/dsl/executors/tensor_value_executor.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : value_quantizers.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/15/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import collections
import contextlib
from typing import TYPE_CHECKING, Any, Optional, Union, Tuple, Iterable, Sequence, List, Mapping, Dict

import torch

from concepts.dsl.dsl_types import TensorValueTypeBase, ScalarValueType, NamedTensorValueType, PyObjValueType
from concepts.dsl.dsl_domain import DSLDomainBase
from concepts.dsl.tensor_value import TensorValue, scalar
from concepts.dsl.tensor_state import TensorState

if TYPE_CHECKING:
    from concepts.dsl.executors.tensor_value_executor import TensorValueExecutorBase


__all__ = ['ValueQuantizer', 'PyObjectStore']


class ValueQuantizer(object):
    values: Dict[str, Union[List[Any], Dict[Any, int]]]
    """the value dictionary is a mapping from typename to a list of values. When we quantize a value `v`, we basically return the index of `v` in the list."""

    def __init__(self, executor: 'TensorValueExecutorBase'):
        """Initialize a value quantizer.

        Args:
            executor: the executor for the domain to use. We need the executor to access the underlying domain and
                the corresponding hash / equality functions.
        """

        self.executor = executor
        self.domain = executor.domain
        self.values = dict()

    def quantize(self, typename: str, value: Union[torch.Tensor, TensorValue]) -> int:
        """Quantize a single value. This API is used to quantize a single value, and it is the lowest level API in this class.
        Most of the time, you should use other higher-level APIs such as :meth:`quantize_tensor` and :meth:`quantize_value` instead.

        Args:
            typename: the typename of the value.
            value: the value to quantize.

        Returns:
            the quantized value, as a single integer.
        """

        if not isinstance(value, TensorValue):
            value = TensorValue(self.domain.types[typename], [], value)
        use_hash = self.executor.has_function_implementation(f'type::{typename}::hash')
        if typename not in self.values:
            self.values[typename] = dict() if use_hash else list()

        if use_hash:
            hash_value = self.executor.get_function_implementation(f'type::{typename}::hash')(value, wrap_rv=False)
            if hash_value not in self.values[typename]:
                self.values[typename][hash_value] = len(self.values[typename])
            return self.values[typename][hash_value]
        else:
            for i, v in enumerate(self.values[typename]):
                if bool(self.executor.get_function_implementation(f'type::{typename}::equal')(v, value, wrap_rv=False)):
                    return i
            self.values[typename].append(value)
            return len(self.values[typename]) - 1

    def quantize_tensor(self, dtype: NamedTensorValueType, tensor: torch.Tensor) -> torch.Tensor:
        """Quantize a PyTorch tensor. The main difference between this function and the :meth:`quantize` function is that
        this function will quantize all the elements in the tensor.

        Args:
            dtype: the dtype of the tensor.
            tensor: the tensor to quantize.

        Returns:
            the quantized tensor.
        """

        if dtype.base_type.typename in ScalarValueType.INTRINSICALLY_QUANTIZED_DTYPES:
            raise RuntimeError('This branch should not be reached. Please report this bug to the developers.')
            # if torch.dtype in (torch.float16, torch.float32, torch.float64):
            #     return torch.round(tensor).to(torch.int64)

        tensor_flatten = tensor.reshape((-1,) + dtype.size_tuple())
        quantized_values = [self.quantize(dtype.typename, v) for v in tensor_flatten]
        quantized_tensor = torch.tensor(quantized_values, dtype=torch.int64, device=tensor_flatten.device)
        quantized_tensor = quantized_tensor.reshape(tensor.shape[:-dtype.ndim()])
        return quantized_tensor

    def quantize_value(self, value: TensorValue) -> TensorValue:
        """Quantize a single TensorValue object.

        Args:
            value: the value to quantize.

        Returns:
            the quantized value.
        """
        if value.dtype.is_intrinsically_quantized():
            return value.simple_quantize()
        assert isinstance(value.dtype, NamedTensorValueType)
        return TensorValue(value.dtype, value.batch_variables, self.quantize_tensor(value.dtype, value.tensor), value.batch_dims, quantized=True)

    def quantize_dict_list(self, continuous_values: Mapping[str, Sequence[Union[torch.Tensor, TensorValue]]]) -> Mapping[str, Sequence[TensorValue]]:
        """Quantize a dictionary of lists of values. The return is a dictionary that maps from the same keys to a list of quantized values.
        Note that the return type is a dictionary of lists of :class:`concepts.dsl.tensor_value.TensorValue`, not a dictionary of lists of integers.

        Args:
            continuous_values: the dictionary of lists of values to quantize.

        Returns:
            the quantized dictionary of lists of values.
        """
        output_dict = dict()
        for typename, values in continuous_values.items():
            output_dict[typename] = set()
            for v in values:
                output_dict[typename].add(self.quantize(typename, v))
            output_dict[typename] = [TensorValue(self.domain.types[typename], [], x, quantized=True) for x in output_dict[typename]]
        return output_dict

    def quantize_state(self, state: TensorState, includes=None, excludes=None) -> TensorState:
        """Quantize a TensorState object. Users can specify which variables to quantize by using the `includes` and `excludes` arguments.
        Meanwhile, this function will read the function definition in the domain to determine whether the feature is a state variable.
        If the feature is not a state variable, it will not be quantized.

        Args:
            state: the state to quantize.
            includes: the variables to include in the quantization. If this argument is not None, only the variables in this list will be quantized.
            excludes: the variables to exclude in the quantization. If this argument is not None, the variables in this list will not be quantized.

        Returns:
            the quantized state.
        """
        state = state.clone()
        for feature_name in state.features.all_feature_names:
            if (includes is not None and feature_name not in includes) or (excludes is not None and feature_name in excludes):
                rv = state.features[feature_name]
            else:
                function = self.domain.functions[feature_name]
                if hasattr(function, 'is_state_variable') and not function.is_state_variable:
                    rv = state.features[feature_name]
                else:
                    rv = self.quantize_value(state.features[feature_name])
            state.features[feature_name] = rv
        return state

    def unquantize(self, typename: str, value: int) -> TensorValue:
        """The lowest-level API to unquantize a single value. Most of the time, you should use other higher-level APIs.

        Args:
            typename: the typename of the value.
            value: the value to unquantize.

        Returns:
            the unquantized value.
        """
        return self.values[typename][value]

    def unquantize_tensor(self, dtype: NamedTensorValueType, tensor: torch.Tensor) -> torch.Tensor:
        """Unquantize a PyTorch tensor. The main difference between this function and the :meth:`unquantize` function is that
        this function will unquantize all the elements in the tensor.

        Args:
            dtype: the dtype of the tensor.
            tensor: the tensor to unquantize.

        Returns:
            the unquantized tensor.
        """

        if dtype.is_intrinsically_quantized():
            return tensor

        assert dtype.typename in self.values, f'Unknown typename: {dtype.typename}.'
        lookup_table = self.values[dtype.typename]
        output = [lookup_table[x].tensor for x in tensor.flatten().tolist()]
        output = torch.stack(output, dim=0).reshape(tensor.shape + dtype.size_tuple())
        return output

    def unquantize_value(self, value: TensorValue) -> TensorValue:
        """Unquantize a single TensorValue object.

        Args:
            value: the value to unquantize.

        Returns:
            the unquantized value.
        """
        dtype = value.dtype
        if isinstance(dtype, PyObjValueType):
            return value

        assert isinstance(dtype, TensorValueTypeBase)
        if dtype.is_intrinsically_quantized():
            return value
        else:
            assert isinstance(dtype, NamedTensorValueType)
            return TensorValue(dtype, value.batch_variables, self.unquantize_tensor(dtype, value.tensor), value.batch_dims, quantized=False)

    # def unquantize_csp(self, csp: ConstraintSatisfactionProblem) -> ConstraintSatisfactionProblem:
    #     """Unquantize a ConstraintSatisfactionProblem object. Essnetially this function will unquantize all the determined values in the CSP.

    #     Args:
    #         csp: the CSP to unquantize.

    #     Returns:
    #         the unquantized CSP.
    #     """
    #     def _cvt(arg):
    #         if isinstance(arg, DeterminedValue):
    #             if not arg.quantized:
    #                 return arg
    #             elif arg.dtype.quantized:
    #                 if arg.dtype == BOOL:
    #                     return DeterminedValue(BOOL, bool(arg.value), True)
    #                 return DeterminedValue(arg.dtype, int(arg.value), True)
    #             else:
    #                 assert isinstance(arg.dtype, NamedTensorValueType) and isinstance(arg.value, int)
    #                 return DeterminedValue(arg.dtype, self.unquantize(arg.dtype.typename, arg.value), False)
    #         else:
    #             return arg

    #     csp = csp.clone()
    #     for i, c in enumerate(csp.constraints):
    #         new_args = tuple(map(_cvt, c.arguments))
    #         new_rv = _cvt(c.rv)
    #         csp.constraints[i] = Constraint(c.function, new_args, new_rv, note=c.note)
    #     return csp

    @contextlib.contextmanager
    def checkpoint(self):
        """A context manager that can be used to checkpoint all the quantized values. This is useful when you performs a series of executions and want to restore the quantized values to save memory."""
        old_values = self.values.copy()
        yield
        self.values = old_values


class PyObjectStore(object):
    """A store for Python objects. This class is used to map Python objects to integers (so that they can be stored in a TensorValue object) and vice versa."""

    executor: 'TensorValueExecutorBase'
    """The executor that uses this object store."""

    domain: DSLDomainBase
    """The underlying domain for the PyObjectStore."""

    def __init__(self, executor: 'TensorValueExecutorBase'):
        """Initialize the PyObjectStore object.

        Args:
            executor: the executor that uses this object store.
        """
        self.executor = executor
        self.domain = executor.domain
        self.objects = collections.defaultdict(list)

    def add(self, typename: str, python_object: Any) -> int:
        """The lowest-API to add a Python object to the store. Most of the time, you should use other higher-level APIs such as :meth:`make_value` and :meth:`make_batched_value`.

        Args:
            typename: the typename of the Python object.
            python_object: the Python object to add.

        Returns:
            the index of the Python object in the store.
        """
        self.objects[typename].append(python_object)
        return len(self.objects[typename]) - 1

    def retrieve(self, typename: str, index: int) -> Any:
        """The lowest-API to retrieve a Python object from the store. Most of the time, you should use other higher-level APIs such as :meth:`retrieve_value`."""
        return self.objects[typename][index]

    def retrieve_value(self, value: TensorValue) -> Union[Any, List[Any]]:
        """Retrieve one or multiple Python object from the store. When the value is a batched value, this function will return a list of (or nested lists of) Python objects."""
        return _tensor2pyobj_list(self, value.dtype.typename, value.tensor)

    def make_value(self, typename: str, python_object: Any) -> TensorValue:
        """Make a TensorValue object from a single Python object.

        Args:
            typename: the typename of the Python object.
            python_object: the Python object to add.

        Returns:
            the TensorValue object.
        """
        index = self.add(typename, python_object)
        return scalar(index, self.domain.types[typename])

    def make_batched_value(self, typename: str, pyobj_list: List[Any], batch_variables: Optional[Sequence[str]] = None) -> TensorValue:
        """Make a TensorValue object from a list of (or nested lists of) Python objects.

        Args:
            typename: the typename of the Python object.
            pyobj_list: the list of (or nested lists of) Python objects.
            batch_variables: the batch variables of the TensorValue object.
        """

        sizes = _get_pyobj_list_size(pyobj_list)
        tensor = torch.zeros(sizes, dtype=torch.int64)

        for indices, pyobj in _iterate_pyobj_list(pyobj_list):
            tensor[indices] = self.add(typename, pyobj)

        return TensorValue(
            self.domain.types[typename],
            batch_variables if batch_variables is not None else len(sizes),
            tensor, batch_dims=0, quantized=True
        )

    @contextlib.contextmanager
    def checkpoint(self):
        """A context manager that can be used to checkpoint all the stored Python objects. This is useful when you performs a series of executions and want to restore the objects to save memory."""
        old_objects = self.objects.copy()
        yield
        self.objects = old_objects


def _tensor2pyobj_list(pyobj_store: PyObjectStore, typename: str, indices_tensor: torch.Tensor) -> Union[Any, List[Any]]:
    if indices_tensor.dim() == 0:
        return pyobj_store.retrieve(typename, indices_tensor.item())
    else:
        return [_tensor2pyobj_list(pyobj_store, typename, indices_tensor[i]) for i in range(indices_tensor.size(0))]


def _get_pyobj_list_size(pyobj_list: Union[Any, List[Any]]) -> Tuple[int, ...]:
    if type(pyobj_list) is list:
        assert len(pyobj_list) > 0
        return (len(pyobj_list), ) + _get_pyobj_list_size(pyobj_list[0])
    else:
        return tuple()


def _iterate_pyobj_list(pyobj_list: Union[Any, List[Any]], indices: Tuple[int, ...] = tuple()) -> Iterable[Tuple[Tuple[int, ...], Any]]:
    """Given a nested list of pyobjs, yield a tuple of (indices, pyobj) for each pyobj."""
    if type(pyobj_list) is list:
        for i, pyobj in enumerate(pyobj_list):
            yield from _iterate_pyobj_list(pyobj, indices + (i,))
    else:
        yield indices, pyobj_list



concepts/dsl/executors/value_quantizers.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : expression.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/19/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data structures for expressions in a DSL.

All classes extend the basic :class:`Expression` class. They can be categorized into the following groups:

- :class:`ObjectOutputExpression` and :class:`ValueOutputExpression` are the expressions that output objects or values.
- :class:`VariableExpression` which is the expression that refers to a variable.
- :class:`VariableAssignmentExpression` which assigns a value to a variable.

Under the :class:`ValueOutputExpression` category, there are a few sub-categories:

- :class:`ConstantExpression` which is the expression that outputs a constant value.
- :class:`FunctionApplicationExpression` which represents the application of a function.
- :class:`BoolExpression` which represents Boolean operations (and, or, not).
- :class:`QuantificationExpression` which represents quantification (forall, exists).
- :class:`GeneralizedQuantificationExpression` which represents generalized quantification (iota, all, counting quantifiers).
- :class:`PredicateEqualExpression` which represents the equality test between a state variable and a value.
- :class:`ConditionalSelectExpression` which represents the conditional selection for some computed value.
- :class:`DeicticSelectExpression` which represents the deictic selection for some computed value (i.e., forall quantifiers).

Under the :class:`ObjectOutputExpression` category, there are a few sub-categories:

- :class:`ObjectConstantExpression` which is the expression that outputs a constant object.

Under the :class:`VariableAssignmentExpression` category, there are a few sub-categories:

- :class:`AssignExpression` which is the expression that assigns a value to a state variable.
- :class:`ConditionalAssignExpression` which is the expression that assigns a value to a state variable conditionally.
- :class:`DeicticAssignExpression` which is the expression that assigns values to state variables with deictic expressions (i.e., forall quantifiers).

The most important classes are: :class:`VariableExpression`, :class:`ObjectConstantExpression`, :class:`ConstantExpression`, and :class:`FunctionApplicationExpression`.
"""

import contextlib
from abc import ABC, abstractmethod
from typing import Any, Optional, Union, Iterable, Tuple, Sequence, List, Dict, Callable
from functools import lru_cache

import torch
import jacinle

from jacinle.utils.enum import JacEnum
from jacinle.utils.printing import indent_text
from jacinle.utils.defaults import wrap_custom_as_default, gen_get_default

from concepts.dsl.dsl_types import FormatContext, get_format_context
from concepts.dsl.dsl_types import TypeBase, ObjectType, ValueType, AutoType, TensorValueTypeBase, PyObjValueType, BOOL, FLOAT32, INT64, ObjectConstant, Variable
from concepts.dsl.dsl_functions import FunctionType, Function, FunctionArgumentUnset, AnonymousFunctionArgumentGenerator
from concepts.dsl.dsl_domain import DSLDomainBase
from concepts.dsl.value import ValueBase, Value
from concepts.dsl.tensor_value import TensorValue

try:
    from typing import TypeGuard
except ImportError:
    class _DummyTypeGuard:
        def __getitem__(self, item):
            return bool
    TypeGuard = _DummyTypeGuard()


__all__ = [
    'DSLExpressionError', 'Expression', 'ExpressionDefinitionContext', 'get_expression_definition_context',
    'ObjectOutputExpression', 'ValueOutputExpression', 'VariableExpression', 'VariableAssignmentExpression',
    'ObjectConstantExpression', 'ConstantExpression',
    'FunctionApplicationError', 'FunctionApplicationExpression',
    'ConditionalSelectExpression', 'DeicticSelectExpression',
    'BoolOpType', 'BoolExpression', 'AndExpression', 'OrExpression', 'NotExpression',
    'QuantificationOpType', 'QuantificationExpression', 'ForallExpression', 'ExistsExpression',
    'PredicateEqualExpression', 'AssignExpression', 'ConditionalAssignExpression', 'DeicticAssignExpression',
    'cvt_expression', 'cvt_expression_list', 'get_type', 'get_types',
    'is_object_output_expression', 'is_variable_assignment_expression', 'is_variable_assignment_expression',
    'is_and_expr', 'is_or_expr', 'is_not_expr', 'is_forall_expr', 'is_exists_expr',
    'iter_exprs', 'find_free_variables'
]


class DSLExpressionError(Exception):
    pass


class Expression(ABC):
    """Expression is an abstract class for all expressions in the DSL.
    An important note about Expression is that the class itself does not contain any "implementation."

    For example, the expression `and(x, y, z)` does not contain any information about how to compute
    the conjunction (e.g., taking the product of the three values).
    The actual implementation of the expression will be provided by the `Executor` classes.
    """

    @property
    @abstractmethod
    def return_type(self) -> Optional[Union[ObjectType, ValueType, FunctionType, Tuple[Union[ValueType, FunctionType], ...]]]:
        raise NotImplementedError()

    def check_arguments(self):
        ctx = get_expression_definition_context()
        if ctx is None or ctx.check_arguments:
            self._check_arguments()

    def _check_arguments(self):
        pass

    @abstractmethod
    def __str__(self) -> str:
        raise NotImplementedError()

    __repr__ = jacinle.repr_from_str

    @lru_cache(maxsize=10)
    def cached_string(self, max_length: Optional[int] = None):
        if max_length is None:
            return str(self)
        else:
            with FormatContext(expr_max_length=max_length).as_default():
                return str(self)


class ExpressionDefinitionContext(object):
    def __init__(
        self, *variables: Variable,
        domain: Optional['DSLDomainBase'] = None,
        scope: Optional[str] = None,
        is_effect_definition: bool = False,
        allow_auto_predicate_def: bool = True,
        check_arguments: bool = True,
    ):
        """Initialize the context.

        Args:
            variables: The variables that are available in the expression.
            domain: the domain of the expression.
            scope: the current definition scope (e.g., in a function). This variable will be used to generate unique names for the functions.
            is_effect_definition: whether the expression is defined in an effect of an operator.
            allow_auto_predicate_def: whether to enable automatic predicate definition.
            check_arguments: whether to check the arguments of the functions.
        """

        self.variables = list(variables)
        self.variable_name2obj = {v.name: v for v in self.variables}
        self.domain = domain
        self.scope = scope
        self.anonymous_argument_generator = AnonymousFunctionArgumentGenerator()
        self.is_effect_definition_stack = [is_effect_definition]
        self.allow_auto_predicate_def = allow_auto_predicate_def
        self.check_arguments = check_arguments

    OPTION_NAMES = ['allow_auto_predicate_def', 'check_arguments']

    variables: List[Variable]
    """The list of variables."""

    variable_name2obj: Dict[str, Variable]
    """The mapping from variable names to variables."""

    domain: Optional['DSLDomainBase']
    """The domain of the expression."""

    scope: Optional[str]
    """The current definition scope (e.g., in a function). This variable will be used to generate unique names for the functions."""

    anonymous_argument_generator: AnonymousFunctionArgumentGenerator
    """The anonymous argument generator."""

    is_effect_definition_stack: List[bool]
    """Whether the expression is defined in an effect of an operator."""

    allow_auto_predicate_def: bool
    """Whether to enable automatic predicate definition."""

    check_arguments: bool
    """Whether to check the arguments of the functions."""

    @wrap_custom_as_default
    def as_default(self):
        yield self

    def __getitem__(self, variable: Union[str, Variable]) -> 'VariableExpression':
        return self.wrap_variable(variable)

    def wrap_variable(self, variable: Union[str, Variable]) -> 'VariableExpression':
        if isinstance(variable, Variable):
            return VariableExpression(variable)
        variable_name = variable
        if variable_name == '??':
            return VariableExpression(Variable('??', AutoType))
        if variable_name not in self.variable_name2obj:
            raise ValueError('Unknown variable: {}; available variables: {}.'.format(variable_name, self.variables))
        return VariableExpression(self.variable_name2obj[variable_name])

    def gen_random_named_variable(self, dtype) -> Variable:
        """Generate a variable expression with a random name. This utility is useful in "flatten_expression". See the doc for that function for details."""
        name = self.anonymous_argument_generator.gen()
        return Variable(name, dtype)

    @contextlib.contextmanager
    def with_variables(self, *args: Variable):
        """Reset the list of variables."""
        old_variables = self.variables.copy()
        self.variables = list(args)
        self.variable_name2obj = {v.name: v for v in self.variables}
        yield
        self.variables = old_variables
        self.variable_name2obj = {v.name: v for v in self.variables}

    @contextlib.contextmanager
    def new_variables(self, *args: Variable):
        """Adding a list of new variables."""
        for arg in args:
            if arg.name in self.variable_name2obj:
                raise ValueError(f'Variable {arg.name} already exists.')
            self.variables.append(arg)
            self.variable_name2obj[arg.name] = arg
        yield self
        for arg in reversed(args):
            self.variables.pop()
            del self.variable_name2obj[arg.name]

    @contextlib.contextmanager
    def mark_is_effect_definition(self, is_effect_definition: bool):
        self.is_effect_definition_stack.append(is_effect_definition)
        yield self
        self.is_effect_definition_stack.pop()

    @property
    def is_effect_definition(self) -> bool:
        return self.is_effect_definition_stack[-1]

    @contextlib.contextmanager
    def options(self, **kwargs):
        for k, v in kwargs.items():
            if k not in self.OPTION_NAMES:
                raise ValueError(f'Unknown option {k}.')
        old_options = {k: getattr(self, k) for k in kwargs}
        for k, v in kwargs.items():
            setattr(self, k, v)
        yield self
        for k, v in old_options.items():
            setattr(self, k, v)


get_expression_definition_context: Callable[[], Optional[ExpressionDefinitionContext]] = gen_get_default(ExpressionDefinitionContext)


class ObjectOutputExpression(Expression):
    @property
    def return_type(self) -> ObjectType:
        raise NotImplementedError()

    def __str__(self) -> str:
        raise NotImplementedError()


class ValueOutputExpression(Expression):
    @property
    def return_type(self) -> ValueType:
        raise NotImplementedError()

    def __str__(self) -> str:
        raise NotImplementedError()


class VariableExpression(Expression):
    def __init__(self, variable: Variable):
        self.variable = variable

    variable: Variable
    """The variable."""

    @property
    def name(self) -> str:
        return self.variable.name

    @property
    def dtype(self) -> Union[ObjectType, ValueType, FunctionType]:
        return self.variable.dtype

    @property
    def return_type(self) -> Union[ObjectType, ValueType, FunctionType]:
        return self.variable.dtype

    def __str__(self) -> str:
        return f'V::{self.name}'


class VariableAssignmentExpression(Expression, ABC):
    @property
    def return_type(self):
        return None


class ObjectConstantExpression(ObjectOutputExpression):
    def __init__(self, constant: ObjectConstant):
        self.constant = constant

    constant: ObjectConstant
    """The object constant."""

    @property
    def name(self) -> str:
        """The name of the object."""
        return self.constant.name

    @property
    def dtype(self) -> ObjectType:
        """The type of the object."""
        return self.constant.dtype

    @property
    def return_type(self) -> ObjectType:
        return self.constant.dtype

    def __str__(self) -> str:
        return f'OBJ::{self.name}'


class ConstantExpression(ValueOutputExpression):
    """Constant expression always returns a constant value."""

    constant: ValueBase
    """The constant."""

    def __init__(self, value: Union[bool, int, float, str, torch.Tensor, Any, ValueBase], dtype: Optional[ValueType] = None):
        if isinstance(value, ValueBase):
            self.constant = value
        else:
            assert dtype is not None
            if isinstance(dtype, (TensorValueTypeBase, PyObjValueType)):
                self.constant = TensorValue.from_scalar(value, dtype)
            else:
                self.constant = Value(dtype, value)

    @property
    def return_type(self) -> ValueType:
        assert isinstance(self.constant.dtype, ValueType)
        return self.constant.dtype

    @classmethod
    def true(cls):
        return cls(torch.tensor(1, dtype=torch.int64), BOOL)

    @classmethod
    def false(cls):
        return cls(torch.tensor(0, dtype=torch.int64), BOOL)

    @classmethod
    def int64(cls, value):
        return cls(torch.tensor(value, dtype=torch.int64), INT64)

    @classmethod
    def float32(cls, value):
        return cls(torch.tensor(value, dtype=torch.float32), FLOAT32)

    def __str__(self):
        if isinstance(self.constant, TensorValue) and self.constant.is_single_elem:
            return f'Const::{self.constant.single_elem()}'
        return str(self.constant)


ConstantExpression.TRUE = ConstantExpression.true()
ConstantExpression.FALSE = ConstantExpression.false()


class FunctionApplicationError(Exception):
    def __init__(self, index: int, expect, got):
        msg = f'Argument #{index} type does not match: expect {expect}, got {got}.'
        super().__init__(msg)

        self.index = index
        self.expect = expect
        self.got = got


class FunctionApplicationExpression(ValueOutputExpression):
    """Function application expression represents the application of a function over a list of arguments."""

    def __init__(self, function: Function, arguments: Iterable[Expression]):
        self.function = function
        self.arguments = tuple(arguments)
        self.check_arguments()

    function: Function
    """The function to be applied."""

    arguments: Tuple[Expression, ...]
    """The list of arguments to the function."""

    def _check_arguments(self):
        try:
            if len(self.function.arguments) != len(self.arguments):
                raise TypeError('Argument number mismatch: expect {}, got {}.'.format(len(self.function.arguments), len(self.arguments)))
            for i, (arg_def, arg) in enumerate(zip(self.function.arguments, self.arguments)):
                if isinstance(arg_def, Variable):
                    if isinstance(arg, VariableExpression):
                        if not arg.dtype.downcast_compatible(arg_def.dtype):
                            raise FunctionApplicationError(i, arg_def.dtype, arg.dtype)
                    elif isinstance(arg, ObjectConstantExpression):
                        if not arg.dtype.downcast_compatible(arg_def.dtype):
                            raise FunctionApplicationError(i, arg_def.dtype, arg.dtype)
                    elif isinstance(arg, ConstantExpression):
                        if not arg.return_type.downcast_compatible(arg_def.dtype):
                            raise FunctionApplicationError(i, arg_def.dtype, arg.return_type)
                    elif isinstance(arg, (FunctionApplicationExpression, GeneralizedQuantificationExpression)):
                        if not arg.return_type.downcast_compatible(arg_def.dtype):
                            raise FunctionApplicationError(i, arg_def.dtype, arg.return_type)
                    else:
                        raise FunctionApplicationError(i, 'VariableExpression or ObjectConstantExpression or ConstantExpression or FunctionApplication', type(arg))
                elif isinstance(arg_def, ValueType):
                    if isinstance(arg, ValueOutputExpression):
                        pass
                    elif isinstance(arg, VariableExpression) and isinstance(arg.return_type, ValueType):
                        pass
                    elif isinstance(arg, ConstantExpression):
                        pass
                    else:
                        raise FunctionApplicationError(i, 'ValueOutputExpression', type(arg))
                    if arg_def != arg.return_type:
                        raise FunctionApplicationError(i, arg_def, arg.return_type)
                else:
                    raise TypeError('Unknown argument definition type: {}.'.format(type(arg_def)))
        except (TypeError, FunctionApplicationError) as e:
            error_header = 'Error during applying {}.\n'.format(str(self.function))
            try:
                arguments_str = ', '.join(str(arg) for arg in self.arguments)
                error_header += ' Arguments: {}\n'.format(arguments_str)
            except Exception:  # noqa
                pass
            raise TypeError(error_header + str(e)) from e

    @property
    def return_type(self) -> ValueType:
        return self.function.return_type

    def __str__(self) -> str:
        fmt = self.function.name + '('
        arg_fmt = [str(x) for x in self.arguments]
        arg_fmt_len = [len(x) for x in arg_fmt]

        ctx = get_format_context()

        # The following criterion is just an approximation. A more principled way is to pass the current indent level
        # to the recursive calls to str(x).
        if ctx.expr_max_length > 0 and (sum(arg_fmt_len) + len(fmt) + 1 > ctx.expr_max_length):
            if sum(arg_fmt_len) > ctx.expr_max_length:
                fmt += '\n' + ',\n'.join([indent_text(x) for x in arg_fmt]) + '\n'
            else:
                fmt += '\n' + ', '.join(arg_fmt) + '\n'
        else:
            fmt += ', '.join(arg_fmt)
        fmt += ')'
        return fmt


class ConditionalSelectExpression(ValueOutputExpression):
    """Conditional select expression represents the selection of a value based on a condition."""

    def __init__(self, predicate: ValueOutputExpression, condition: ValueOutputExpression):
        self.predicate = predicate
        self.condition = condition
        self.check_arguments()

    predicate: ValueOutputExpression
    """The predicate expression."""

    condition: ValueOutputExpression
    """The condition expression."""

    def _check_arguments(self):
        if isinstance(self.condition, ValueOutputExpression) and self.condition.return_type == BOOL:
            pass
        elif isinstance(self.condition, VariableExpression) and self.condition.return_type.downcast_compatible(BOOL):
            pass
        else:
            raise TypeError('Condition must be a boolean expression.')

    @property
    def return_type(self) -> ValueType:
        return self.predicate.return_type

    def __str__(self):
        return f'cond-select({self.predicate} if {self.condition})'


class DeicticSelectExpression(ValueOutputExpression):
    def __init__(self, variable: Variable, expr: ValueOutputExpression):
        self.variable = variable
        self.expression = expr
        self.check_arguments()

    variable: Variable
    """The new quantified variable."""

    expression: ValueOutputExpression
    """The internal expression."""

    def _check_arguments(self):
        assert isinstance(self.variable.dtype, ObjectType)

    @property
    def return_type(self) -> ValueType:
        return self.expression.return_type

    def __str__(self):
        return f'deictic-select({self.variable}: {self.expression})'


class BoolOpType(JacEnum):
    AND = 'and'
    OR = 'or'
    NOT = 'not'


class BoolExpression(ValueOutputExpression):
    def __init__(self, bool_op_type: BoolOpType, arguments: Sequence[ValueOutputExpression]):
        self.bool_op = bool_op_type
        self.arguments = tuple(arguments)
        self.check_arguments()

    bool_op: BoolOpType
    """The boolean operation. Can be AND, OR, NOT."""

    arguments: Tuple[ValueOutputExpression, ...]
    """The list of arguments."""

    def _check_arguments(self):
        if self.bool_op is BoolOpType.NOT:
            assert len(self.arguments) == 1, f'Number of arguments for NotOp should be 1, got: {len(self.arguments)}.'
        for i, arg in enumerate(self.arguments):
            assert isinstance(arg, ValueOutputExpression), f'BoolOp only accepts ValueOutputExpressions, got argument #{i} of type {type(arg)}.'

    @property
    def return_type(self) -> ValueType:
        return self.arguments[0].return_type

    def __str__(self):
        arguments = ', '.join([str(arg) for arg in self.arguments])
        return f'{self.bool_op.value}({arguments})'


class AndExpression(BoolExpression):
    bool_op: BoolOpType
    """The boolean operation. Must be :py:attr:`BoolOpType.AND`."""

    arguments: Tuple[ValueOutputExpression, ...]

    def __init__(self, *arguments: ValueOutputExpression):
        super().__init__(BoolOpType.AND, arguments)


class OrExpression(BoolExpression):
    bool_op: BoolOpType
    """The boolean operation. Must be :py:attr:`BoolOpType.OR`."""

    arguments: Tuple[ValueOutputExpression, ...]

    def __init__(self, *arguments: ValueOutputExpression):
        super().__init__(BoolOpType.OR, arguments)


class NotExpression(BoolExpression):
    bool_op: BoolOpType
    """The boolean operation. Must be :py:attr:`BoolOpType.NOT`."""

    arguments: Tuple[ValueOutputExpression]
    """The list of arguments. Must contain exactly one argument."""

    def __init__(self, arg: ValueOutputExpression):
        super().__init__(BoolOpType.NOT, [arg])


class QuantificationOpType(JacEnum):
    FORALL = 'forall'
    EXISTS = 'exists'


class QuantificationExpression(ValueOutputExpression):
    def __init__(self, quantification_op: QuantificationOpType, variable: Variable, expr: ValueOutputExpression):
        self.quantification_op = quantification_op
        self.variable = variable
        self.expression = expr

        self.check_arguments()

    quantification_op: QuantificationOpType
    """The quantification operation. Can be FORALL or EXISTS."""

    variable: Variable
    """The quantified variable."""

    expression: ValueOutputExpression
    """The internal expression."""

    def _check_arguments(self):
        assert isinstance(self.expression, ValueOutputExpression), f'QuantificationOp only accepts ValueOutputExpressions, got type {type(self.expression)}.'
        assert isinstance(self.variable.dtype, ObjectType)

    @property
    def return_type(self) -> ValueType:
        return self.expression.return_type

    def __str__(self):
        return f'{self.quantification_op.value}({self.variable}: {self.expression})'


class GeneralizedQuantificationExpression(ValueOutputExpression):
    def __init__(self, quantification_op: Any, variable: Variable, expr: ValueOutputExpression, return_type: Optional[ValueType] = None):
        self.quantification_op = quantification_op
        self.variable = variable
        self.expression = expr

        self._return_type = return_type if return_type is not None else self.expression.return_type

        self.check_arguments()

    quantification_op: Any
    """The quantification operation. It can be any data type."""

    variable: Variable
    """The quantified variable."""

    expression: ValueOutputExpression
    """The internal expression."""

    def _check_arguments(self):
        assert isinstance(self.expression, ValueOutputExpression), f'QuantificationOp only accepts ValueOutputExpressions, got type {type(self.expression)}.'
        assert isinstance(self.variable.dtype, ObjectType)

    @property
    def return_type(self) -> ValueType:
        return self._return_type

    def __str__(self):
        return f'{self.quantification_op}({self.variable}: {self.expression})'


class ForallExpression(QuantificationExpression):
    def __init__(self, variable: Variable, expr: ValueOutputExpression):
        super().__init__(QuantificationOpType.FORALL, variable, expr)

    quantification_op: QuantificationOpType
    """The quantification operation. Must be :py:attr:`QuantificationOpType.FORALL`."""

    variable: Variable
    expression: ValueOutputExpression


class ExistsExpression(QuantificationExpression):
    def __init__(self, variable: Variable, expr: ValueOutputExpression):
        super().__init__(QuantificationOpType.EXISTS, variable, expr)

    quantification_op: QuantificationOpType
    """The quantification operation. Must be :py:attr:`QuantificationOpType.EXISTS`."""

    variable: Variable
    expression: ValueOutputExpression


class _PredicateValueExpression(Expression, ABC):
    def __init__(self, predicate: Union[VariableExpression, FunctionApplicationExpression], value: ValueOutputExpression):
        self.predicate = predicate
        self.value = value
        self.check_arguments()

    def _check_arguments(self):
        try:
            rtype = self.predicate.return_type
            if rtype.assignment_type() != self.value.return_type:
                raise FunctionApplicationError(0, f'{self.predicate.return_type}(assignment type is {rtype.assignment_type()})', self.value.return_type)
        except TypeError as e:
            raise e
        except FunctionApplicationError as e:
            error_header = 'Error during _PredicateValueExpression checking: feature = {} value = {}.\n'.format(str(self.predicate), str(self.value))
            raise TypeError(
                error_header +
                f'Value type does not match: expect: {e.expect}, got {e.got}.'
            ) from e


class PredicateEqualExpression(ValueOutputExpression, _PredicateValueExpression):
    predicate: Union[VariableExpression, FunctionApplicationExpression]
    """The predicate expression."""

    value: ValueOutputExpression
    """The value expression."""

    def _check_arguments(self):
        super()._check_arguments()
        if not isinstance(self.predicate, (VariableExpression, FunctionApplicationExpression)):
            raise TypeError(f'PredicateEqualOp only support dest type VariableExpression or FunctionApplication, got {type(self.predicate)}.')

    @property
    def return_type(self):
        return BOOL

    def __str__(self):
        return f'equal({self.predicate}, {self.value})'


class AssignExpression(_PredicateValueExpression, VariableAssignmentExpression):
    def __init__(self, predicate: FunctionApplicationExpression, value: ValueOutputExpression):
        _PredicateValueExpression.__init__(self, predicate, value)

    predicate: Union[VariableExpression, FunctionApplicationExpression]
    """The predicate expression, must be a :class:`FunctionApplicationExpression` which refers to a state variable."""

    value: ValueOutputExpression
    """The expression for the value to assign to the state variable."""

    def _check_arguments(self):
        super()._check_arguments()
        assert isinstance(self.predicate, FunctionApplicationExpression), 'AssignOp only support dest type FunctionApplication, got {}.'.format(type(self.predicate))

    def __str__(self):
        return f'assign({self.predicate}: {self.value})'


class ConditionalAssignExpression(_PredicateValueExpression, VariableAssignmentExpression):
    def __init__(self, feature: FunctionApplicationExpression, value: ValueOutputExpression, condition: ValueOutputExpression):
        self.condition = condition
        _PredicateValueExpression.__init__(self, feature, value)

    predicate: Union[VariableExpression, FunctionApplicationExpression]
    """The predicate expression, must be a :class:`FunctionApplicationExpression` which refers to a state variable."""

    value: ValueOutputExpression
    """The expression for the value to assign to the state variable."""

    condition: ValueOutputExpression
    """The condition expression."""

    def _check_arguments(self):
        super()._check_arguments()
        assert isinstance(self.condition, ValueOutputExpression) and self.condition.return_type == BOOL

    def __str__(self):
        return f'cond-assign({self.predicate}: {self.value} if {self.condition})'


class DeicticAssignExpression(VariableAssignmentExpression):
    def __init__(self, variable: Variable, expr: Union[VariableAssignmentExpression]):
        self.variable = variable
        self.expression = expr
        self.check_arguments()

    variable: Variable
    """The quantified variable."""

    expression: VariableAssignmentExpression
    """The internal expression."""

    def _check_arguments(self):
        assert isinstance(self.variable.dtype, ObjectType)

    def __str__(self):
        return f'deictic-assign({self.variable}: {self.expression})'


ExpressionCompatible = Union[Expression, Variable, str, ObjectConstant, bool, int, float, torch.Tensor, ValueBase]


def cvt_expression(expr: ExpressionCompatible, dtype: Optional[Union[ObjectType, ValueType]] = None) -> Expression:
    """Convert an expression compatible object to an expression. Acceptable types are:

    * :class:`Expression`.
    * :class:`Variable`: return a :class:`VariableExpression`.
    * :class:`str`: return a :class:`ConstantExpression` with the given constant string name, or a :class:`ObjectConstantExpression` if the dtype is a :class:`ObjectType`.
    * :class:`ObjectConstant`: return a :class:`ObjectConstantExpression`.
    * :class:`bool`, :class:`int`, :class:`float`, :class:`torch.Tensor`: return a :class:`ConstantExpression`.
    * :class:`ValueBase`: return a :class:`ConstantExpression`.

    Args:
        expr: the expression compatible object.
        dtype: the expected data type of the expression. If not given, the dtype will be inferred from the given object.

    Returns:
        the converted expression.

    Raises:
        TypeError: if the given object is not an expression compatible object.
    """

    if isinstance(expr, Expression):
        return expr
    elif isinstance(expr, Variable):
        return VariableExpression(expr)
    elif isinstance(expr, str):
        if isinstance(dtype, ObjectType):
            return ObjectConstantExpression(ObjectConstant(expr, dtype or AutoType))
        elif isinstance(dtype, ValueType):
            return ConstantExpression(Value(dtype or AutoType, expr))
    elif isinstance(expr, ObjectConstant):
        return ObjectConstantExpression(expr)
    elif isinstance(expr, bool):
        return ConstantExpression(torch.tensor(int(expr), dtype=torch.int64), dtype or BOOL)
    elif isinstance(expr, int):
        return ConstantExpression(torch.tensor(expr, dtype=torch.int64), dtype or INT64)
    elif isinstance(expr, float):
        return ConstantExpression(torch.tensor(expr, dtype=torch.float32), dtype or FLOAT32)
    elif isinstance(expr, torch.Tensor):
        if expr.dtype == torch.int64:
            return ConstantExpression(expr, dtype or INT64)
        elif expr.dtype == torch.float32:
            return ConstantExpression(expr, dtype or FLOAT32)
        else:
            raise TypeError(f'Unsupported tensor type: {expr.dtype}.')
    elif isinstance(expr, ValueBase):
        if isinstance(expr.dtype, ValueType):
            return ConstantExpression(expr, expr.dtype)
        else:
            raise TypeError(f'Unsupported value type: {expr.dtype}.')
    else:
        raise TypeError(f'Non-compatible expression type {type(expr)} for expression "{expr}".')


def cvt_expression_list(arguments: Sequence[ExpressionCompatible], dtypes: Optional[Sequence[Union[ObjectType, ValueType]]] = None) -> List[Expression]:
    """Convert a list of expression compatible objects to a list of expressions.

    Args:
        arguments: the list of expression compatible objects.
        dtypes: the list of expected data types of the expressions. If not given, the dtypes will be inferred from the given objects.
            It can be a single data type, in which case all the expressions will be converted to this data type.

    Returns:
        the list of converted expressions.
    """
    if dtypes is None:
        arguments = [cvt_expression(arg) for arg in arguments]
    else:
        arguments = [cvt_expression(arg, dtype) for arg, dtype in zip(arguments, dtypes)]
    return arguments


def get_type(value: Any) -> Union[TypeBase, Tuple[TypeBase, ...]]:
    """Get the type of the given value."""
    if value is FunctionArgumentUnset:
        return FunctionArgumentUnset
    elif isinstance(value, Function):
        return value.ftype
    elif isinstance(value, Expression):
        return value.return_type
    elif isinstance(value, ValueBase):
        return value.dtype
    elif isinstance(value, (bool, int, float, str)):
        return AutoType
    else:
        raise ValueError(f'Unknown value type: {type(value)}.')


def get_types(args=None, kwargs=None):
    """Get the types of the given arguments and keyword arguments."""
    ret = list()
    if args is not None:
        ret.append(tuple(get_type(v) for v in args))
    if kwargs is not None:
        ret.append({k: get_type(v) for k, v in kwargs.items()})
    if len(ret) == 1:
        return ret[0]
    return tuple(ret)


def is_object_output_expression(expr: Expression) -> TypeGuard[ObjectOutputExpression]:
    return isinstance(expr, ObjectOutputExpression) or (isinstance(expr, VariableExpression) and isinstance(expr.variable.dtype, ObjectType))


def is_value_output_expression(expr: Expression) -> TypeGuard[ValueOutputExpression]:
    return isinstance(expr, ValueOutputExpression)


def is_variable_assignment_expression(expr: Expression) -> bool:
    return isinstance(expr, VariableAssignmentExpression)


def is_and_expr(expr: Expression) -> TypeGuard[AndExpression]:
    return isinstance(expr, BoolExpression) and expr.bool_op is BoolOpType.AND


def is_or_expr(expr: Expression) -> TypeGuard[OrExpression]:
    return isinstance(expr, BoolExpression) and expr.bool_op is BoolOpType.OR


def is_not_expr(expr: Expression) -> TypeGuard[NotExpression]:
    return isinstance(expr, BoolExpression) and expr.bool_op is BoolOpType.NOT


def is_constant_bool_expr(expr: Expression) -> TypeGuard[ConstantExpression]:
    if isinstance(expr, ConstantExpression) and expr.return_type == BOOL:
        return True
    return False


def is_forall_expr(expr: Expression) -> TypeGuard[ForallExpression]:
    return isinstance(expr, QuantificationExpression) and expr.quantification_op is QuantificationOpType.FORALL


def is_exists_expr(expr: Expression) -> TypeGuard[ExistsExpression]:
    return isinstance(expr, QuantificationExpression) and expr.quantification_op is QuantificationOpType.EXISTS


def iter_exprs(expr: Expression) -> Iterable[Expression]:
    """Iterate over all sub-expressions of the input."""
    yield expr
    if isinstance(expr, FunctionApplicationExpression):
        for arg in expr.arguments:
            yield from iter_exprs(arg)
    elif isinstance(expr, BoolExpression):
        for arg in expr.arguments:
            yield from iter_exprs(arg)
    elif isinstance(expr, QuantificationExpression):
        yield from iter_exprs(expr.expression)
    elif isinstance(expr, GeneralizedQuantificationExpression):
        yield from iter_exprs(expr.expression)
    elif isinstance(expr, PredicateEqualExpression):
        yield from iter_exprs(expr.predicate)
        yield from iter_exprs(expr.value)
    elif isinstance(expr, AssignExpression):
        yield from iter_exprs(expr.value)
    elif isinstance(expr, ConditionalSelectExpression):
        yield from iter_exprs(expr.predicate)
        yield from iter_exprs(expr.condition)
    elif isinstance(expr, ConditionalAssignExpression):
        yield from iter_exprs(expr.value)
        yield from iter_exprs(expr.condition)
    elif isinstance(expr, (DeicticSelectExpression, DeicticAssignExpression)):
        yield from iter_exprs(expr.expression)
    elif isinstance(expr, (VariableExpression, ConstantExpression, ObjectConstantExpression)):
        pass
    else:
        raise TypeError('Unknown expression type: {}.'.format(type(expr)))


def find_free_variables(expr: Expression) -> Tuple[Variable, ...]:
    free_variables = dict()
    bounded_variables = dict()

    def dfs(e: Expression):
        if isinstance(e, VariableExpression):
            if e.variable.name not in bounded_variables:
                free_variables[e.variable.name] = e.variable
        elif isinstance(e, (QuantificationExpression, GeneralizedQuantificationExpression)):
            bounded_variables[e.variable.name] = e.variable
            dfs(e.expression)
            del bounded_variables[e.variable.name]
        elif isinstance(e, FunctionApplicationExpression):
            [dfs(arg) for arg in e.arguments]
        elif isinstance(e, BoolExpression):
            [dfs(arg) for arg in e.arguments]
        elif isinstance(e, PredicateEqualExpression):
            dfs(e.predicate)
            dfs(e.value)
        elif isinstance(e, AssignExpression):
            dfs(e.value)
        elif isinstance(e, ConditionalSelectExpression):
            dfs(e.predicate)
            dfs(e.condition)
        elif isinstance(e, ConditionalAssignExpression):
            dfs(e.value)
            dfs(e.condition)
        elif isinstance(e, (DeicticSelectExpression, DeicticAssignExpression)):
            bounded_variables[e.variable.name] = e.variable
            dfs(e.expression)
            del bounded_variables[e.variable.name]
        elif isinstance(e, (ConstantExpression, ObjectConstantExpression)):
            pass
        else:
            raise TypeError('Unknown expression type: {}.'.format(type(e)))

    dfs(expr)
    return tuple(free_variables.values())


concepts/dsl/expression.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : expression_visitor.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/30/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""A visitor for iterating over expressions."""

from typing import Any

import concepts.dsl.expression as E
from concepts.dsl.expression import Expression


class ExpressionVisitor(object):
    """A visitor for iterating over expressions."""

    def visit(self, expr: Expression) -> Any:
        """The main entry point of the visitor. It will call the corresponding method for the given expression type.

        Args:
            expr: the expression to visit.

        Returns:
            the result of the visit.
        """
        if isinstance(expr, E.VariableExpression):
            return self.visit_variable_expression(expr)
        elif isinstance(expr, E.ObjectConstantExpression):
            return self.visit_object_constant_expression(expr)
        elif isinstance(expr, E.ConstantExpression):
            return self.visit_constant_expression(expr)
        elif isinstance(expr, E.FunctionApplicationExpression):
            return self.visit_function_application_expression(expr)
        elif isinstance(expr, E.ConditionalSelectExpression):
            return self.visit_conditional_select_expression(expr)
        elif isinstance(expr, E.DeicticSelectExpression):
            return self.visit_deictic_select_expression(expr)
        elif isinstance(expr, E.BoolExpression):
            return self.visit_bool_expression(expr)
        elif isinstance(expr, E.QuantificationExpression):
            return self.visit_quantification_expression(expr)
        elif isinstance(expr, E.GeneralizedQuantificationExpression):
            return self.visit_generalized_quantification_expression(expr)
        elif isinstance(expr, E.PredicateEqualExpression):
            return self.visit_predicate_equal_expression(expr)
        elif isinstance(expr, E.AssignExpression):
            return self.visit_assign_expression(expr)
        elif isinstance(expr, E.ConditionalAssignExpression):
            return self.visit_conditional_assign_expression(expr)
        elif isinstance(expr, E.DeicticAssignExpression):
            return self.visit_deictic_assign_expression(expr)
        else:
            raise TypeError(f'Unknown expression type: {type(expr)}.')

    def visit_variable_expression(self, expr: E.VariableExpression) -> Any:
        raise NotImplementedError()

    def visit_object_constant_expression(self, expr: E.ObjectConstantExpression) -> Any:
        raise NotImplementedError()

    def visit_constant_expression(self, expr: E.ConstantExpression) -> Any:
        raise NotImplementedError()

    def visit_function_application_expression(self, expr: E.FunctionApplicationExpression) -> Any:
        raise NotImplementedError()

    def visit_conditional_select_expression(self, expr: E.ConditionalSelectExpression) -> Any:
        raise NotImplementedError()

    def visit_deictic_select_expression(self, expr: E.DeicticSelectExpression) -> Any:
        raise NotImplementedError()

    def visit_bool_expression(self, expr: E.BoolExpression) -> Any:
        raise NotImplementedError()

    def visit_quantification_expression(self, expr: E.QuantificationExpression) -> Any:
        raise NotImplementedError()

    def visit_generalized_quantification_expression(self, expr: E.GeneralizedQuantificationExpression) -> Any:
        raise NotImplementedError()

    def visit_predicate_equal_expression(self, expr: E.PredicateEqualExpression) -> Any:
        raise NotImplementedError()

    def visit_assign_expression(self, expr: E.AssignExpression) -> Any:
        raise NotImplementedError()

    def visit_conditional_assign_expression(self, expr: E.ConditionalAssignExpression) -> Any:
        raise NotImplementedError()

    def visit_deictic_assign_expression(self, expr: E.DeicticAssignExpression) -> Any:
        raise NotImplementedError()


concepts/dsl/expression_visitor.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : function_domain.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/18/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Function domain contains the definition for a collection of types and function types."""

import inspect
import collections
import contextlib
from typing import Any, Optional, Union, List, Dict, Callable
from jacinle.logging import get_logger
from jacinle.utils.printing import indent_text

from concepts.dsl.dsl_types import TypeBase, AutoType, ObjectType, ValueType, UnionType, Variable
from concepts.dsl.dsl_functions import FunctionType, OverloadedFunctionType, Function, FunctionArgumentResolutionError, get_function_argument_resolution_context
from concepts.dsl.dsl_domain import DSLDomainBase
from concepts.dsl.expression import ExpressionDefinitionContext, ConstantExpression, FunctionApplicationExpression, get_type
from concepts.dsl.value import Value

logger = get_logger(__file__)

__all__ = ['FunctionDomain', 'resolve_lambda_function_type']


class FunctionDomain(DSLDomainBase):
    def __init__(self, name: Optional[str] = None):
        super().__init__(name)

        self.__define_overloaded_functions = collections.defaultdict(list)

    AllowedTypes = Union[ObjectType, ValueType, FunctionType]

    name: str
    types: Dict[str, Union[ObjectType, ValueType]]
    functions: Dict[str, Function]
    constants: Dict[str, Value]

    def define_overloaded_function(self, name: str, overloads: List[Union[Function, Callable]], implementation: bool = True):
        """Define a function with multiple overloads.

        Args:
            name: the name of the function.
            overloads: the overloads of the function. Each overload should be a python function or a :class:`Function` instance.
            implementation: whether to store the function body of `function` as the implementation of the function.
        """
        overloads_qs = list()
        overloads_function_body = list()
        for f in overloads:
            if isinstance(f, Function):
                overloads_qs.append(f.ftype)
                overloads_function_body.append(f.function_body)
            else:
                overloads_qs.append(FunctionType.from_annotation(f))
                overloads_function_body.append(f)

        self.functions[name] = Function(
            name, OverloadedFunctionType(overloads_qs),
            overridden_call=None, resolved_from=None,
        )

    def lam(self, lambda_expression: Callable, name: str = '__lambda__', typing_cues: Optional[Dict[str, Union[str, TypeBase]]] = None) -> Function:
        """Parse a lambda expression into a Function instance.

        Args:
            lambda_expression: the lambda expression.
            name: the name of the function.
            typing_cues: the typing cues for the function. It should be a dictionary that maps the argument names to the types.
            If you want to specify the return type, use the key ``'return'``.

        Returns:
            The parsed :class:`Function` instance.
        """

        # TODO (Jiayuan Mao): fix this.
        # sig = inspect.signature(lambda_expression)
        # argument_names = list(sig.parameters.keys())
        # if len(argument_names) == 0:
        #     return lambda_expression()

        try:
            function_type = FunctionType.from_annotation(lambda_expression)
        except FunctionArgumentResolutionError:
            function_type = resolve_lambda_function_type(self, lambda_expression, typing_cues)

        return Function(name, function_type, overridden_call=lambda_expression)

    def __getattr__(self, name: str) -> Union[TypeBase, Value, Function]:
        if name.startswith('t_'):
            return self.types[name[2:]]
        elif name.startswith('c_'):
            return self.constants[name[2:]]
        elif name.startswith('f_'):
            return self.functions[name[2:]]
        raise AttributeError(name)

    @contextlib.contextmanager
    def define(self, implementation=True):
        """A context manager that defines the types, constants, and functions in the domain.

        Usage:

        .. code-block:: python

            with domain.define():
                A = domain.define_type('A', ...)
                B = domain.define_type('B', ...)

                # Definition of constants. Equivalent to `domain.define_const(domain.types['A'], 'a')`.
                a: A
                b: B

                # Definition of functions.
                def f(a: A, b: B) -> A:
                    pass

                # Definition of overloaded functions.
                @domain.overload
                def g(a: A, b: B) -> A:
                    pass

                @domain.overload
                def g(b: B, a: A) -> B:
                    pass

        Note that the definition for constants is only allowed when the code is executed at the global scope.
        """
        locals_before = inspect.stack()[2][0].f_locals.copy()
        annotations_before = locals_before.get('__annotations__', dict()).copy()
        yield self
        locals_after = inspect.stack()[2][0].f_locals.copy()
        annotations_after = locals_after.get('__annotations__', dict()).copy()

        new_functions = {
            k for k in locals_after
            if k not in locals_before and not isinstance(locals_after[k], FunctionDomain)
        }
        new_annotations = {
            k for k in annotations_after
            if k not in annotations_before or annotations_after[k] != annotations_before[k]
        }

        if len(new_annotations) == 0:
            logger.warning('ts.define() for constants is only allowed at the global scope.')

        functions = list()
        for func_name in new_functions:
            var = locals_after[func_name]
            if not inspect.isfunction(var):
                raise ValueError('Support only function definitions in the DEFINE body, got {}.'.format(func_name))
            if func_name in self.__define_overloaded_functions:
                continue
            functions.append((var.__code__.co_firstlineno, var))

        for func_name, overloads in self.__define_overloaded_functions.items():
            lineno = overloads[0].__code__.co_firstlineno
            if len(overloads) > 1:
                functions.append((lineno, (func_name, overloads)))
            else:
                functions.append((lineno, overloads[0]))
        functions.sort()

        for _, f in functions:
            if isinstance(f, tuple):
                self.define_overloaded_function(*f, implementation=implementation)
            else:
                self.define_function(f, implementation=implementation)

        for const_name in new_annotations:
            variable_type = annotations_after[const_name]
            self.define_const(variable_type, const_name)

        self.__define_overloaded_functions = collections.defaultdict(list)

    def overload(self, function):
        """Overload a function. See the docstring for :meth:`define` for usage."""
        self.__define_overloaded_functions[function.__name__].append(function)
        return function

    def format_summary(self) -> str:
        """Format the summary of the domain."""
        fmt = 'Types:\n'
        for type in self.types.values():
            fmt += '  ' + str(type) + '\n'
        fmt += 'Constants:\n'
        for const in self.constants.values():
            fmt += '  ' + str(const) + '\n'
        fmt += 'Functions:\n'
        for function in self.functions.values():
            fmt += '  ' + str(function).replace('\n', '\n  ') + '\n'

        fmt = 'TypeSystem: {}\n'.format(self.name) + indent_text(fmt.rstrip())
        return fmt

    def print_summary(self):
        """Print the summary of the domain."""
        print(self.format_summary())

    def serialize(self, program: Union[FunctionApplicationExpression, Value]) -> Dict[str, Any]:
        """Serialize a program into a dictionary.

        Args:
            program: the program to be serialized.

        Returns:
            the serialized program.
        """

        def dfs(p):
            if isinstance(p, FunctionApplicationExpression):
                record = dict()
                record['__serialize_type__'] = 'function'
                record['function'] = p.function.name
                record['args'] = list()
                for arg in p.arguments:
                    record['args'].append(dfs(arg))
            elif isinstance(p, ConstantExpression):
                record = dict()
                record['__serialize_type__'] = 'object'
                record['type'] = p.constant.dtype.typename
                record['value'] = p.constant.value
            else:
                raise TypeError('Unserializable object: {}.'.format(type(p)))
            return record
        return dfs(program)

    def deserialize(self, dictionary: Dict[str, Any]) -> Union[FunctionApplicationExpression, Value]:
        """Deserialize a program from a dictionary.

        Args:
            dictionary: the dictionary to be deserialized.

        Returns:
            the deserialized program.
        """
        def dfs(d):
            assert '__serialize_type__' in d
            stype = d['__serialize_type__']
            if stype == 'function':
                func = self.functions[d['function']]
                args = list()
                for arg in d['args']:
                    args.append(dfs(arg))
                return func(*args)
            elif stype == 'object':
                return ConstantExpression(Value(self.types[d['type']], d['value']))
            else:
                raise TypeError('Unknown serialized type: {}.'.format(stype))
        return dfs(dictionary)


def _canonize_type(fn_domain: FunctionDomain, t: Optional[Union[str, TypeBase]]) -> Optional[TypeBase]:
    if t is None:
        return None
    if isinstance(t, UnionType):
        return UnionType(*tuple(_canonize_type(fn_domain, t) for t in t.types))
    if isinstance(t, TypeBase):
        return t

    assert isinstance(t, str)
    return fn_domain.types[t]


def _canonize_signature(fn_domain: FunctionDomain, signature: inspect.Signature) -> inspect.Signature:
    """Canonize the signature of a function. Specifically, it converts the type names to the corresponding type objects.
    For example, given a domain with types ``A`` and ``B``, the signature ``def f(a: 'A', b: 'B') -> A`` will be converted to
    the signature of ``def f(a: A, b: B) -> A``.

    Args:
        fn_domain: the function domain.
        signature: the signature to be canonized.

    Returns:
        the canonized signature.
    """
    params = [
        inspect.Parameter(v.name, v.kind, default=v.default, annotation=_canonize_type(fn_domain, v.annotation))
        for k, v in signature.parameters.items()
    ]
    return_annotation = _canonize_type(fn_domain, signature.return_annotation)
    return inspect.Signature(params, return_annotation=return_annotation)


def resolve_lambda_function_type(fn_domain: FunctionDomain, lambda_expression: Callable, typing_cues: Dict[str, Union[str, TypeBase]]) -> FunctionType:
    """Resolve the function type of a lambda expression. It enumerates all possible types for the arguments and the return value,
    and then checks whether the lambda expression satisfies the type constraints. See the docstring for :meth:`FunctionDomain.lam` for usage.

    Args:
        fn_domain: The function domain.
        lambda_expression: The lambda expression.
        typing_cues: The typing cues. The key is the variable name, and the value is the type.
            If you want to specify the typing cue for the return type, use the key 'return'.

    Returns:
        FunctionType: The resolved function type.
    """
    sig = inspect.signature(lambda_expression)
    parameters = list(sig.parameters.keys())
    parameter_types = {k: AutoType for k in parameters}

    for i, (name, param) in enumerate(sig.parameters.items()):
        if param.annotation is not sig.empty:
            parameter_types[param.name] = param.annotation
    return_type = sig.return_annotation if sig.return_annotation is not sig.empty else AutoType

    if typing_cues is not None:
        for k, v in typing_cues.items():
            assert k in parameter_types
            parameter_types[k] = v

        if 'return' in typing_cues:
            return_type = typing_cues['return']

    ctx = ExpressionDefinitionContext()
    with ctx.as_default():
        for k in sig.parameters:
            if parameter_types[k] is AutoType:
                exceptions = list()
                success_types = list()

                for k_type in fn_domain.types.values():
                    parameter_types[k] = k_type
                    parameter_grounding = list()

                    for param in parameters:
                        t = parameter_types[param]
                        parameter_grounding.append(ctx.wrap_variable(Variable(param, t)))

                    try:
                        output = lambda_expression(*parameter_grounding)
                        success_types.append(k_type)

                        if return_type is AutoType:
                            return_type = get_type(output)

                    except FunctionArgumentResolutionError as e:
                        exceptions.append((k_type, e))

                if len(success_types) == 1:
                    parameter_types[k] = success_types[0]
                elif len(success_types) == 0:
                    with get_function_argument_resolution_context().exc():
                        fmt = 'Failed to infer argument type for {}.\n'.format(k)
                        fmt += 'Detailed messages are:\n'
                        for t, e in exceptions:
                            fmt += indent_text('Trying {}:\n'.format(t) + indent_text(str(e)) + '\n')
                        raise FunctionArgumentResolutionError(fmt.rstrip())
                else:
                    with get_function_argument_resolution_context().exc():
                        fmt = 'Got ambiguous type for {}.\n'.format(k)
                        fmt += 'Candidates are:\n'
                        for r in success_types:
                            fmt += indent_text(str(r)) + '\n'
                        raise FunctionArgumentResolutionError(fmt.strip())

    if return_type is AutoType and len(parameter_types) == 0:
        return_type = get_type(lambda_expression())

    return FunctionType([parameter_types[i] for i in parameters], return_type, parameters)


concepts/dsl/function_domain.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/10/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Simple search-based algorithms for learning expressions in a DSL."""


concepts/dsl/learning/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : function_domain_search.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/10/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""An enumerative search algorithm to generate candidate functions and expressions in a simple function domain."""

import itertools
from collections import defaultdict
from dataclasses import dataclass
from typing import Any, Optional, Union, Iterable, Sequence, Tuple, List, Callable

from concepts.dsl.dsl_types import ConstantType, ValueType
from concepts.dsl.dsl_functions import FunctionType, Function
from concepts.dsl.dsl_domain import DSLDomainBase
from concepts.dsl.value import Value
from concepts.dsl.expression import ConstantExpression, FunctionApplicationExpression, VariableExpression
from concepts.dsl.function_domain import FunctionDomain
from concepts.dsl.executors.function_domain_executor import FunctionDomainExecutor

__all__ = [
    'FunctionDomainExpressionSearchResult', 'FunctionDomainExpressionEnumerativeSearcher', 'gen_merge_functions',
    'gen_expression_search_result_from_expressions',
    'FunctionArgumentStat', 'stat_function', 'canonize_function_parameters',
    'learn_expression_from_examples'
]


_Types = FunctionDomain.AllowedTypes


@dataclass
class FunctionDomainExpressionSearchResult(object):
    expression: Union[ConstantExpression, Function, FunctionApplicationExpression]
    """The expression that is enumerated."""

    depth: int
    """The depth of the expression."""

    nr_constant_arguments: int
    """The number of constant arguments in the expression."""

    nr_variable_arguments: int
    """The number of variable arguments in the expression."""

    nr_function_arguments: int
    """The number of function arguments in the expression."""


class FunctionDomainExpressionEnumerativeSearcher(object):
    """An enumerator of expressions and functions for a function domain."""

    def __init__(self, domain: DSLDomainBase):
        """Initialize the searcher.

        Args:
            domain: the domain of the semantics.
        """
        self.domain = domain
        self._constant_type_cache = self._gen_constant_type_cache()
        self._function_type_cache = self._gen_function_type_cache()

    def _gen_constant_type_cache(self):
        """Generate a dictionary that maps types to constants."""
        cache = defaultdict(list)
        for const in self.domain.constants.values():
            cache[const.dtype].append(const)
        return cache

    def _gen_function_type_cache(self):
        """Generate a dictionary that maps types to functions (by return types)."""
        cache = defaultdict(list)
        for func in self.domain.functions.values():
            if func.is_overloaded:
                for f in func.all_sub_functions:
                    cache[f.return_type].append(f)
            else:
                cache[func.ftype.return_type].append(func)
        return cache

    def gen(
        self,
        return_type: Optional[Union[_Types, Tuple[_Types, ...], List[_Types]]] = None,
        *,
        max_depth: int = 3,
        max_variable_arguments: int = 2,
        max_constant_arguments: int = 1,
        max_function_arguments: int = 0,
        search_constants: bool = False,
        hash_function: Optional[Callable[[Union[Function, FunctionApplicationExpression]], Any]] = None,
        verbose: bool = False,
    ):
        return self.gen_constant_expressions(return_type) + self.gen_function_application_expressions(
            return_type, max_depth=max_depth, max_variable_arguments=max_variable_arguments,
            max_constant_arguments=max_constant_arguments, max_function_arguments=max_function_arguments,
            search_constants=search_constants, hash_function=hash_function, verbose=verbose
        )

    def gen_constant_expressions(
        self,
        return_type: Optional[Union[_Types, Tuple[_Types, ...], List[_Types]]] = None
    ) -> List[FunctionDomainExpressionSearchResult]:
        """Generate constant expressions of a set of given types.

        Args:
            return_type: the return type of the expressions. If None, all types are allowed.
                It can be a single type, a tuple of types, or a list of types.

        Returns:
            A list of constant expressions.
        """
        constants = list()
        for c in self.domain.constants.values():
            if (
                return_type is None or
                (isinstance(return_type, ConstantType) and c.dtype == return_type) or
                (isinstance(return_type, (tuple, list)) and c.dtype in return_type)
            ):
                constants.append(FunctionDomainExpressionSearchResult(
                    ConstantExpression(c), 1, 0, 0, 0
                ))
        return constants

    def gen_function_application_expressions(
        self,
        return_type: Optional[Union[_Types, Tuple[_Types, ...], List[_Types]]] = None,
        *,
        max_depth: int = 3,
        max_variable_arguments: int = 2,
        max_constant_arguments: int = 1,
        max_function_arguments: int = 0,
        search_constants: bool = False,
        hash_function: Optional[Callable[[Union[Function, FunctionApplicationExpression]], Any]] = None,
        verbose: bool = False,
    ) -> List[FunctionDomainExpressionSearchResult]:
        """Generate functions and function application expressions of a set of given types.

        Args:
            max_depth: the maximum depth of the expressions.
            return_type: the return type of the expressions. If None, all types are allowed.
                It can be a single type, a tuple of types, or a list of types.
            max_variable_arguments: the maximum number of variable arguments of the functions.
            max_constant_arguments: the maximum number of constant arguments of the functions.
                Note that when ``search_constants`` is True, this parameter corresponds to the maximum number of
                constant arguments bound to the function / function application expression.
            max_function_arguments: the maximum number of arguments of the functions.
            search_constants: whether to search for constants.
            verbose: whether to print the search progress.

        Returns:
            A list of function application expressions.
        """

        ftcache = self._function_type_cache
        current = {i: defaultdict(list) for i in range(1, max_depth + 1)}

        def iter_ddl_values(dd):
            for v in dd.values():
                yield from v

        for rtype, functions in ftcache.items():
            for f in functions:
                current[1][f.ftype.return_type].append(gen_merge_functions(f))

        if max_function_arguments > 0:
            for ret_type in self.domain.types.values():
                current[1][ret_type].extend(self._gen_function_primitives(ret_type, max_function_arguments))

        if verbose:
            print('-' * 20, 'Depth', 1, '-' * 100)
            for rtype, functions in current[1].items():
                for f in functions:
                    print(rtype, '\t', f)

        for depth in range(2, max_depth + 1):
            for depth1 in range(1, depth):
                for f1 in list(iter_ddl_values(current[depth1])):
                    for i in range(f1.nr_arguments):
                        for depth2 in range(1, depth - depth1 + 1):
                            for f2 in current[depth2][f1.ftype.argument_types[i]]:
                                current[depth][f1.ftype.return_type].append(
                                    gen_merge_functions(f1, i, f2)
                                )

            if verbose:
                print('-' * 20, 'Depth', depth, '-' * 100)
                for rtype, functions in current[depth].items():
                    for f in functions:
                        print(rtype, '\t', f)

        expressions = list()
        for depth, vs in current.items():
            for f in iter_ddl_values(vs):
                if isinstance(f, Function):
                    stat = stat_function(f)
                    if stat.nr_constant_arguments <= max_constant_arguments and \
                            stat.nr_variable_arguments <= max_variable_arguments and \
                            stat.nr_function_arguments <= max_function_arguments:
                        expressions.append(FunctionDomainExpressionSearchResult(
                            canonize_function_parameters(f, ignore_permutation=True),
                            depth, stat.nr_constant_arguments, stat.nr_variable_arguments, stat.nr_function_arguments
                        ))
                else:
                    expressions.append(FunctionDomainExpressionSearchResult(
                        f, depth, 0, 0, 0
                    ))

        expressions = self._unique_function_expressions(expressions, hash_function=hash_function)

        if search_constants:
            expressions = self._bind_constants_to_expressions(expressions)

        if return_type is not None:
            output_expressions = list()
            for result in expressions:
                if _match_return_type(result.expression, return_type):
                    output_expressions.append(result)
            expressions = output_expressions

        return expressions

    def _gen_function_primitives(self, ret_type, nr_function_arguments):
        def function_call(func, *args):
            return func(*args)

        def gen():
            types = tuple(self.domain.types.values())
            for repeat in range(1, nr_function_arguments + 1):
                for arg_types in itertools.product(types, repeat=repeat):
                    yield Function(
                        '__lambda__',
                        FunctionType(
                            [FunctionType(arg_types, ret_type), ] + list(arg_types),
                            ret_type
                        ),
                        overridden_call=function_call,
                    )

        return tuple(gen())

    def _unique_function_expressions(
        self,
        functions: List[FunctionDomainExpressionSearchResult],
        hash_function: Optional[Callable[[Union[Function, FunctionApplicationExpression]], Any]] = None
    ) -> List[FunctionDomainExpressionSearchResult]:
        """Return a list of unique functions and function application expressions.

        Args:
            functions: a list of input expressions.
            hash_function: a custom function that generates a hash for a function or function application expression.
                If None, the default ``str`` function is used.

        Returns:
            A list of expressions without duplicates.
        """

        if hash_function is None:
            hash_function = str

        unique_functions = list()
        unique_function_hashes = set()
        for r in functions:
            h = hash_function(r.expression)
            if h not in unique_function_hashes:
                unique_functions.append(r)
                unique_function_hashes.add(h)

        return unique_functions

    def _bind_constants_to_expressions(
        self,
        functions: List[FunctionDomainExpressionSearchResult],
    ) -> List[FunctionDomainExpressionSearchResult]:
        output_functions = list()
        for result in functions:
            if not isinstance(result.expression, Function):
                output_functions.append(result)
                continue

            f = result.expression
            constant_denotations, constant_types = list(), list()
            for index, argument_type in enumerate(f.ftype.argument_types):
                if isinstance(argument_type, ConstantType):
                    constant_denotations.append(f'#{index}')
                    constant_types.append(argument_type)

            for const in itertools.product(*[
                self._constant_type_cache[t] for t in constant_types
            ]):
                const_mapping = {k: ConstantExpression(v) for k, v in zip(constant_denotations, const)}
                partial_func = f.partial(**const_mapping, execute_fully_bound_functions=True)
                new_r = partial_func
                output_functions.append(FunctionDomainExpressionSearchResult(
                    new_r,
                    result.depth,
                    result.nr_constant_arguments, result.nr_variable_arguments, result.nr_function_arguments
                ))

        return output_functions


def _match_return_type(f: Union[Function, FunctionApplicationExpression], return_types: Union[_Types, Tuple[_Types, ...], List[_Types]]) -> bool:
    """Check if the type of a function or a function application expression matches a set of given return type."""
    if not isinstance(return_types, (tuple, list)):
        return_types = (return_types, )

    if isinstance(f, Function):
        for return_type in return_types:
            if f.ftype.typename == return_type.typename:
                return True
    elif isinstance(f, FunctionApplicationExpression):
        return f.return_type in return_types
    else:
        raise TypeError(f'Expected Function or FunctionApplicationExpression, got {type(f)}.')


def gen_merge_functions(f1: Function, arg_index=None, f2: Optional[Function] = None) -> Function:
    """Generate merge functions. Specifically, given two functions f1 and f2,
    this function generates a new function. For example, given f1(x, y) and f2(z),
    with arg_index = 1, this function generates:

    .. code-block:: python

        def merged(x, z):
            return f1(x, f2(z))

    That is, we first apply f2 with the input arguments, and then apply f1 with the
    rest of the arguments and the output of f2.

    A special case is when f2 is None. In this case, we generate a function that
    is simply a wrapper of f1.

    Args:
        f1: the first function.
        arg_index: the index of the argument of f1 that we want to merge with f2.
        f2: the second function.

    Returns:
        A new function.
    """
    if arg_index is None:
        return Function(
            '__lambda__',
            FunctionType(f1.ftype.argument_types, f1.ftype.return_type), overridden_call=f1
        )
    else:
        f1_arg_types = f1.ftype.argument_types
        f2_arg_types = f2.ftype.argument_types
        arg_types = f2_arg_types + f1_arg_types[:arg_index] + f1_arg_types[arg_index + 1:]

        def new_function_call(*args):
            f2_args = args[:f2.nr_arguments]
            f1_args = list(args[f2.nr_arguments:])

            f2_ret = f2(*f2_args)
            f1_args.insert(arg_index, f2_ret)
            return f1(*f1_args)

        return Function(
            '__lambda__',
            FunctionType(arg_types, f1.ftype.return_type),
            overridden_call=new_function_call
        )


def gen_expression_search_result_from_expressions(expressions: Iterable[Union[ConstantExpression, Function, FunctionApplicationExpression]]) -> List[FunctionDomainExpressionSearchResult]:
    """Generate a list of FunctionDomainExpressionSearchResult from a list of expressions."""
    results = list()
    for expression in expressions:
        if isinstance(expression, (ConstantExpression, FunctionApplicationExpression)):
            results.append(FunctionDomainExpressionSearchResult(expression, 0, 0, 0, 0))
        else:
            stat = stat_function(expression)
            results.append(FunctionDomainExpressionSearchResult(expression, 0, stat.nr_constant_arguments, stat.nr_variable_arguments, stat.nr_function_arguments))
    return results


@dataclass
class FunctionArgumentStat(object):
    """Statistics for the argument list of a function."""

    nr_constant_arguments: int
    nr_variable_arguments: int
    nr_function_arguments: int


def stat_function(f: Function) -> FunctionArgumentStat:
    """Return the number of constants, variables, and functions in a function."""
    nr_variable_arguments = 0
    nr_constant_arguments = 0
    nr_function_arguments = 0
    for arg_t in f.ftype.argument_types:
        if isinstance(arg_t, ConstantType):
            nr_constant_arguments += 1
        elif isinstance(arg_t, FunctionType):
            nr_function_arguments += 1
            nr_variable_arguments += 1  # function arg is also variable arg
        else:
            nr_variable_arguments += 1
    return FunctionArgumentStat(
        nr_constant_arguments=nr_constant_arguments,
        nr_variable_arguments=nr_variable_arguments,
        nr_function_arguments=nr_function_arguments
    )


def canonize_function_parameters(
    f: Union[ConstantExpression, FunctionApplicationExpression, Function],
    ignore_permutation: bool = False
) -> Union[ConstantExpression, FunctionApplicationExpression, Function]:
    """Return a new function object with argument reordered: functions, variables, constants.

    Args:
        f: the function to be canonized. If the function is a function application expression or a constant expression,
            this function returns the same object.

    Returns:
        A new function object with argument reordered: functions, variables, constants.
    """
    if isinstance(f, (ConstantExpression, FunctionApplicationExpression)):
        return f
    assert isinstance(f, Function)

    # NB(Jiayuan Mao @ 2022/12/10): if the function is a primitive function, just return it.
    if f.overridden_call is None:
        return f

    assert not f.is_overloaded
    assert isinstance(f.derived_expression, FunctionApplicationExpression)
    function_args, variable_args, constant_args = list(), list(), list()

    if ignore_permutation:
        def walk(node: FunctionApplicationExpression):
            if isinstance(node.function, VariableExpression):
                new_index = f.ftype.argument_names.index(node.function.name)
                if new_index not in function_args:
                    function_args.append(new_index)
            for arg in node.arguments:
                if isinstance(arg, VariableExpression):
                    new_index = f.ftype.argument_names.index(arg.name)
                    if isinstance(arg.dtype, ConstantType):
                        if new_index not in constant_args:
                            constant_args.append(new_index)
                    elif isinstance(arg.dtype, ValueType):
                        if new_index not in variable_args:
                            variable_args.append(new_index)
                    else:
                        raise TypeError('Unknown type for anonymous argument #{}, type = {}.'.format(arg.name, arg.dtype))
                elif isinstance(arg, FunctionApplicationExpression):
                    walk(arg)
                else:
                    raise TypeError('Unknown type for anonymous argument type {}.'.format(type(arg)))
        walk(f.derived_expression)
    else:
        for arg_index, arg_t in enumerate(f.ftype.argument_types):
            if isinstance(arg_t, FunctionType):
                function_args.append(arg_index)
            elif isinstance(arg_t, ConstantType):
                constant_args.append(arg_index)
            elif isinstance(arg_t, ValueType):
                variable_args.append(arg_index)
            else:
                raise TypeError('Unknown type for argument {}, type = {}.'.format(arg_index, arg_t))
    new_argument_mapping = list(function_args) + list(variable_args) + list(constant_args)

    f = f.remap_arguments(new_argument_mapping)
    return f


def learn_expression_from_examples(
    domain: FunctionDomain,
    executor: FunctionDomainExecutor,
    input_output: Sequence[Tuple[Sequence[Value], Value, Any]],
    criterion: Callable[[Value, Value], bool],
    candidate_expressions: Optional[Iterable[FunctionDomainExpressionSearchResult]] = None,
    max_depth: int = 3,
    max_function_arguments: int = 0,
    search_constants: bool = True,
) -> Function:
    """Learn a function from examples.

    Args:
        domain: the function domain.
        executor: the executor of the function domain.
        input_output: a sequence of (input, output, grounding) pairs.
        criterion: a function that takes two values and returns True if they are equal.
        candidate_expressions: a sequence of candidate functions. If None, we will generate
            candidate functions automatically.
        max_depth: the maximum depth of the function. Only used when ``candidate_expressions`` is None.
        max_function_arguments: the maximum number of arguments of a function. Only used when
            ``candidate_expressions`` is None.
        search_constants: whether to search for constants. Only used when ``candidate_expressions`` is None.

    Returns:
        A function.
    """

    assert len(input_output) > 0, 'No input-output pairs are given.'
    sample_input, sample_output, _ = input_output[0]
    if len(sample_input) == 0:
        target_type = sample_output.dtype
    else:
        target_type = FunctionType([v.dtype for v in sample_input], sample_output.dtype)

    if candidate_expressions is None:
        domain = FunctionDomainExpressionEnumerativeSearcher(domain)
        candidate_expressions = domain.gen_function_application_expressions(
            target_type,
            max_depth=max_depth,
            max_function_arguments=max_function_arguments,
            search_constants=search_constants
        )

    if isinstance(target_type, FunctionType):
        def score_function(result: FunctionDomainExpressionSearchResult) -> float:
            f = result.expression
            if isinstance(f, FunctionApplicationExpression):
                return -1
            return sum(
                criterion(executor.execute_function(f, *input, grounding=grounding), output)
                for input, output, grounding in input_output
            )
    else:
        def score_function(result: FunctionDomainExpressionSearchResult) -> float:
            f = result.expression
            if isinstance(f, Function):
                return -1
            return sum(
                criterion(executor.execute(f, grounding), output)
                for input, output, grounding in input_output
            )
    return max(candidate_expressions, key=score_function).expression


concepts/dsl/learning/function_domain_search.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Baseclass for all parsers of domain-specific languages and the implementation of a simple function domain parser."""

concepts/dsl/parsers/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : fol_python_parser.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/05/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import ast
from typing import Optional, Tuple, Sequence
from concepts.dsl.dsl_types import TypeBase, ObjectType, BOOL, INT64, Variable
from concepts.dsl.dsl_functions import Function, FunctionType
from concepts.dsl.dsl_domain import DSLDomainBase
from concepts.dsl.expression import ExpressionDefinitionContext, ValueOutputExpression, get_expression_definition_context, get_types
from concepts.dsl.expression import NotExpression, AndExpression, OrExpression, ForallExpression, ExistsExpression, FunctionApplicationExpression
from concepts.dsl.parsers.parser_base import ParserBase


class FOLPythonParser(ParserBase):
    """A parser to parse first-order logic (FOL) expressions in Python syntax. Currrently supported features:

    - logic operations: and, or, not
    - quantifiers: forall, exists (see below for the syntax)
    - function application: ``f(x, y, z)``
    - function definition: ``def f(x, y, z): return x + y + z``

    The syntax for quantifiers is as follows:
        .. code-block:: python

            from typing import Any, Type, Callable

            def forall(dtype: Type, func: Callable[[Any], bool]) -> bool: ...
            def exists(dtype: Type, func: Callable[[Any], bool]) -> bool: ...


    Examples:
        .. code-block:: python

            from concepts.dsl.dsl_types import ObjectType
            from concepts.dsl.function_domain import FunctionDomain

            domain = FunctionDomain()
            domain.define_type(ObjectType('Person'))

            parser = FOLPythonParser(domain, inplace_definition=True)
            parser.parse_expression('exists(Person, lambda x: is_phd(x))')

            function_string = '''
            def is_grandfather(x: Person, y: Person) -> bool:
                # x is the grandfather of y
                return exists(Person, lambda z: is_father(x, z) and is_parent(y, z))
            '''
            parser.parse_function(function_string)
    """

    def __init__(self, domain: DSLDomainBase, inplace_definition: bool = False, inplace_polymorphic_function: bool = False, inplace_definition_type: bool = False):
        """Initialize the parser.

        Args:
            domain: the domain to use.
            inplace_definition: whether to allow expressions to contain functions that are not defined in the domain.
                If set to True, the parser will automatically define these functions.
            inplace_polymorphic_function: whether inplace functions are assumed to be polymorphic.
            inplace_definition_type: whether to allow expressions to contain types that are not defined in the domain.
                If set to True, the parser will automatically define these types.
        """
        self.domain = domain
        self.inplace_definition = inplace_definition
        self.inplace_polymorphic_function = inplace_polymorphic_function
        self.inplace_definition_type = inplace_definition_type

    domain: DSLDomainBase
    """The domain for types and functions."""

    inplace_definition: bool
    """Whether to allow functions to be defined in-place."""

    def parse_domain_string(self, string: str) -> DSLDomainBase:
        raise NotImplementedError('FOLPythonParser does not support parsing domain definition.')

    def parse_expression(self, string: str, arguments: Sequence[Variable] = tuple()) -> ValueOutputExpression:
        module = ast.parse(string)
        expression = ast_get_expression(module)
        return self.parse_expression_ast(expression, arguments)

    def parse_expression_ast(self, expression: ast.AST, arguments: Sequence[Variable] = tuple()) -> ValueOutputExpression:
        with ExpressionDefinitionContext(*arguments, domain=self.domain).as_default():
            return self._parse_expression_inner(expression)

    def parse_function(self, string: str) -> Function:
        module = ast.parse(string)
        function = ast_get_function_definition(module)
        arguments, return_type, body = ast_get_simple_function(function)

        arguments = [Variable(arg.name, self._get_type_from_domain(self.domain, arg.dtype)) for arg in arguments]
        return_type = self._get_type_from_domain(return_type, return_type)

        return Function(
            function.name,
            FunctionType(arguments, return_type),
            self.parse_expression_ast(body, arguments)
        )

    def _parse_expression_inner(self, expression: ast.AST) -> ValueOutputExpression:
        ctx = get_expression_definition_context()
        if isinstance(expression, ast.Call):
            function_name = ast_get_literal_or_class_name(expression.func)
            if self._is_quantification_expression_name(function_name):
                return self._parse_quantification_expression(function_name, expression)
            else:  # function is a regular function.
                return self._parse_function_application(function_name, expression)
        elif isinstance(expression, ast.UnaryOp) and isinstance(expression.op, ast.Not):
            return NotExpression(self._parse_expression_inner(expression.operand))
        elif isinstance(expression, ast.BoolOp):
            if isinstance(expression.op, ast.And):
                return AndExpression(*[self._parse_expression_inner(value) for value in expression.values])
            elif isinstance(expression.op, ast.Or):
                return OrExpression(*[self._parse_expression_inner(value) for value in expression.values])
            else:
                raise NotImplementedError(f'BoolOp {expression.op} is not supported.')
        elif isinstance(expression, ast.Name):
            return ctx[expression.id]
        else:
            raise NotImplementedError(f'Expression {expression} is not supported. Full expression: {ast.dump(expression)}.')

    def _is_quantification_expression_name(self, function_name: str) -> bool:
        return function_name in ('forall', 'exists')

    def _parse_quantification_expression(self, function_name: str, expression: ast.Call) -> ValueOutputExpression:
        ctx = get_expression_definition_context()

        assert len(expression.args) in (2, 3), f'Expect two or three arguments for quantification expressions, got {len(expression.args)}: {ast.dump(expression)}'
        assert isinstance(expression.args[0], (ast.Name, ast.Constant)), f'Expect a variable type name for quantification expressions, got {expression.args[0]}: {ast.dump(expression)}'
        if len(expression.args) == 3:
            assert isinstance(expression.args[1], ast.Constant), f'Expect an integer for counting quantifiers, got {expression.args[1]}: {ast.dump(expression)}'
            assert isinstance(expression.args[1].value, int), f'Expect an integer for counting quantifiers, got {expression.args[1].value}: {ast.dump(expression)}'

        assert isinstance(expression.args[-1], ast.Lambda), f'Expect a lambda expression for quantification expressions, got {expression.args[1]}: {ast.dump(expression)}'
        var_type = ast_get_literal_or_class_name(expression.args[0])
        var_type = self._get_type_from_domain(ctx.domain, var_type)
        assert isinstance(var_type, ObjectType), f'Expect an object type for quantification expressions, got {var_type}.'

        lambda_expression = expression.args[-1]
        assert len(lambda_expression.args.args) == 1, f'Expect one argument for quantification expressions, got {len(lambda_expression.args.args)}.'
        assert lambda_expression.args.args[0].annotation is None, f'Expect no type annotation for lambda functions in quantification expressions, got {lambda_expression.args.args[0].annotation}.'

        var_name = lambda_expression.args.args[0].arg
        var = Variable(var_name, var_type)

        return self._parse_quantification_expression_inner(function_name, var, lambda_expression.body, counting_quantifier=expression.args[1].value if len(expression.args) == 3 else None)

    def _parse_quantification_expression_inner(self, function_name: str, var: Variable, lambda_body: ast.AST, counting_quantifier: Optional[int] = None) -> ValueOutputExpression:
        assert counting_quantifier is None, 'Counting quantifiers are not supported yet.'
        ctx = get_expression_definition_context()
        with ctx.new_variables(var):
            if function_name == 'forall':
                return ForallExpression(var, self._parse_expression_inner(lambda_body))
            else:
                return ExistsExpression(var, self._parse_expression_inner(lambda_body))

    def _parse_function_application(self, function_name: str, expression: ast.Call) -> ValueOutputExpression:
        ctx = get_expression_definition_context()

        parsed_args = [self._parse_expression_inner(arg) for arg in expression.args]
        function = None
        if function_name not in ctx.domain.functions:
            if self.inplace_definition:
                if self.inplace_polymorphic_function:
                    function_name = function_name + '_' + '_'.join([arg.return_type.typename for arg in parsed_args])

                if function_name in ctx.domain.functions:
                    function = ctx.domain.functions[function_name]
                else:
                    function = Function(function_name, FunctionType(get_types(parsed_args), BOOL))
                    ctx.domain.define_function(function)
            else:
                raise KeyError(f'Function {function_name} is not defined in the domain.')
        else:
            function = ctx.domain.functions[function_name]
        return FunctionApplicationExpression(function, parsed_args)

    def _get_type_from_domain(self, domain: DSLDomainBase, name: str) -> ObjectType:
        if name == 'bool':
            return BOOL
        elif name == 'int':
            return INT64
        else:
            if name not in domain.types:
                if self.inplace_definition_type:
                    domain.define_type(ObjectType(name))
                else:
                    raise ValueError(f'Undefined type {name}.')

            return domain.types[name]


def ast_get_literal_or_class_name(const: ast.AST) -> str:
    """Get the literal value or identifier name of a constant.

    Args:
        const: the constant, should be either :class:`ast.Constant` or :class:`ast.Name`.

    Returns:
        the literal value or identifier name.
    """
    if isinstance(const, ast.Constant):
        return const.value
    elif isinstance(const, ast.Name):
        return const.id
    else:
        raise TypeError(f'Expect an ast.Constant or a ast.Name, got {type(const)}.')


def ast_get_function_definition(module: ast.Module) -> ast.FunctionDef:
    """Get the single function definition in the module.

    Args:
        module: the module. It should contains exactly one function definition.

    Returns:
        the function definition.
    """
    assert len(module.body) == 1, f'Expect one single function definition, got {len(module.body)}.'
    assert isinstance(module.body[0], ast.FunctionDef), f'Expect a function definition, got {module.body[0]}.'
    return module.body[0]


def ast_get_expression(module: ast.Module) -> ast.AST:
    """Get the single expression in the module.

    Args:
        module: the module. It should contains exactly one expression.

    Returns:
        the expression.
    """
    assert len(module.body) == 1, f'Expect one single expression, got {len(module.body)}.'
    assert isinstance(module.body[0], ast.Expr), f'Expect an expression, got {module.body[0]}.'
    return module.body[0].value


def ast_get_simple_function(function: ast.FunctionDef) -> Tuple[Tuple[Variable, ...], TypeBase, ast.AST]:
    """Get the arguments, return type, and body of a simple function. This function only works for "simple functions".
    That is, the function body contains only a single return statement. This function imposes strong restrictions:

    - The function should have a single return statement.
    - All arguments and return type should be annotated with either a single class name or a string. (It does not support
      type hints like ``List[int]``.)

    Args:
        function: the function definition.

    Returns:
        - arguments: the arguments of the function as a tuple of :class:`~concepts.dsl.dsl_types.Variable` (types are strings).
        - return_type: the return type of the function (string).
        - body: the body of the function.
    """
    arguments = list()
    for arg in function.args.args:
        assert arg.annotation is not None, f'Expect type annotation for argument {arg.arg}.'
        arguments.append(Variable(arg.arg, ast_get_literal_or_class_name(arg.annotation)))

    assert function.returns is not None, f'Expect return type annotation for function {function.name}.'
    return_type = ast_get_literal_or_class_name(function.returns)

    assert len(function.body) == 1, f'Expect one single return statement, got {len(function.body)}.'
    return_statement = function.body[0]
    assert isinstance(return_statement, ast.Return), f'Expect a return statement, got {return_statement}.'

    return tuple(arguments), return_type, return_statement.value


concepts/dsl/parsers/fol_python_parser.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : function_expression_parser.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""The parser for simple function expressions."""

import lark
import json

from dataclasses import dataclass
from concepts.dsl.dsl_types import ObjectType, Variable, ObjectConstant
from concepts.dsl.function_domain import FunctionDomain
from concepts.dsl.value import Value
from concepts.dsl.expression import Expression, ConstantExpression, VariableExpression, ObjectConstantExpression, FunctionApplicationExpression
from concepts.dsl.parsers.parser_base import ParserBase

# lark.v_args
inline_args = lark.v_args(inline=True)

__all__ = ['FunctionExpressionTransformer', 'FunctionExpressionParser']


@dataclass
class _Placeholder(object):
    name: str
    placeholder_type: str


class FunctionExpressionTransformer(lark.Transformer):
    """The lark transformer for the simple function expression parser."""

    def __init__(self, domain: FunctionDomain, escape_string: bool = True):
        """Initialize the transformer.

        Args:
            domain: The domain to use.
            escape_string: Whether to escape the string.
        """
        super().__init__()
        self.domain = domain
        self.escape_string = escape_string

    def start(self, args):
        return args[0]

    @inline_args
    def function_application(self, function_name, *args):
        function = self.domain.get_function(function_name)

        if function.is_overloaded:
            raise NotImplementedError()
        if function.nr_arguments != len(args):
            raise ValueError(f'Function {function_name} expects {function.nr_arguments} arguments, but {len(args)} are given.')

        args = list(args)
        for i, arg in enumerate(args):
            if isinstance(arg, _Placeholder):
                arg_type = function.ftype.argument_types[i]
                if arg.placeholder_type == 'c':
                    if isinstance(arg_type, ObjectType):
                        args[i] = ObjectConstantExpression(ObjectConstant(arg.name, arg_type))
                    else:
                        args[i] = ConstantExpression(Value(arg_type, arg.name))
                elif arg.placeholder_type == 'v':
                    args[i] = VariableExpression(Variable(arg.name, arg_type))
                else:
                    raise ValueError(f'Unknown placeholder type {arg.placeholder_type}.')

        return FunctionApplicationExpression(function, args)

    def function_name(self, function_name):
        return function_name[0].value

    def argument(self, argument):
        return argument[0]

    def variable(self, variable):
        return _Placeholder(variable[0].value, 'v')

    def constant(self, constant):
        if self.escape_string:
            return _Placeholder(json.loads(constant[0].value), 'c')
        return _Placeholder(constant[0].value, 'c')


class FunctionExpressionParser(ParserBase):
    """The simple function expression parser.

    This parser works for simple function expressions, for example: ``f(x, y)``.

    Each function name is a string composed of letters, digits, and _, and each argument is an expression or a string.
    When ``escape_string`` is set to ``True``, the string should be escaped with double quotes.

        >>> from concepts.dsl.function_domain import FunctionDomain
        >>> from concepts.dsl.parsers.function_expression_parser import FunctionExpressionParser
        >>> domain = FunctionDomain()
        >>> parser = FunctionExpressionParser(domain)
        >>> parser.parse_expression('f(g(), "string")')

    """

    GRAMMAR = r"""
start: function_application
function_application: function_name "(" (argument ("," argument)*)? ")"
function_name: STRING

%import common.WS
%ignore WS

%import common.LETTER
%import common.DIGIT
STRING: LETTER ("_"|"-"|LETTER|DIGIT)*

%import common.ESCAPED_STRING
"""

    def __init__(self, domain: FunctionDomain, allow_variable: bool = False, escape_string: bool = True):
        """Initialize the parser.

        Args:
            domain: the domain to use.
            allow_variable: whether to allow variable.
            escape_string: whether to escape the string.
        """

        self.allow_variable = allow_variable
        self.escape_string = escape_string if escape_string is not None else allow_variable

        if self.allow_variable:
            assert self.escape_string, 'If allow_variable is True, escape_string must be True.'

        self.parser = lark.Lark(
            self.GRAMMAR +
            ('constant: ESCAPED_STRING\n' if escape_string else 'constant: STRING\n') +
            ('variable: STRING\nargument: function_application | constant | variable' if allow_variable else 'argument: function_application | constant'),
            parser='lalr', start='start'
        )
        self.transformer = FunctionExpressionTransformer(domain)

    def parse_expression(self, string: str) -> Expression:
        """Parse an expression from a string.

        Args:
            string: The string to parse.

        Returns:
            The parsed expression.
        """
        tree = self.parser.parse(string)
        return self.transformer.transform(tree)


concepts/dsl/parsers/function_expression_parser.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : parser_base.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""The baseclass for all parsers of domain-specific languages."""

from concepts.dsl.dsl_domain import DSLDomainBase
from concepts.dsl.expression import Expression

__all__ = ['ParserBase']


class ParserBase(object):
    """The baseclass for all parsers of domain-specific languages."""

    def parse_domain_file(self, path: str) -> DSLDomainBase:
        """Parse a domain from a file.

        Args:
            path: the path to the file.

        Returns:
            the parsed domain.
        """
        with open(path) as f:
            return self.parse_domain_string(f.read())

    def parse_domain_string(self, string: str) -> DSLDomainBase:
        """Parse a domain from a string.

        Args:
            string: the string to parse.

        Returns:
            the parsed domain.
        """
        raise NotImplementedError()

    def parse_expression(self, string: str) -> Expression:
        """Parse an expression from a string.

        Args:
            string: the string to parse.

        Returns:
            the parsed expression.
        """
        raise NotImplementedError()


concepts/dsl/parsers/parser_base.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : tensor_state.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 11/05/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data structures for tensor states. A tensor state is composed of tensor values.
Intuitively, each state can be represented as a dictionary, mapping from feature names to tensors.
Semantically, this dictionary represents the state of the system (e.g., a scene, a collection of robot and object states in an environment, etc.).
"""

from jacinle.utils.printing import indent_text
from dataclasses import dataclass
from typing import Any, Optional, Union, Iterable, Tuple, List, Mapping, Dict

from concepts.dsl.dsl_types import ObjectType, AutoType
from concepts.dsl.tensor_value import TensorValue, concat_tvalues


__all__ = [
    'StateObjectReference',
    'MultidimensionalArrayInterface', 'TensorValueDict',
    'TensorStateBase',
    'NamedObjectStateMixin', 'TensorState', 'NamedObjectTensorState',
    'ObjectNameArgument', 'ObjectTypeArgument',
    'concat_states',
]


@dataclass
class StateObjectReference(object):
    """The StateObjectReference class represents a reference to an object in the state. It contains the name of the object and
    an index of the object in the state. Note that this index is depending on the type of the object. See the documentation for
    :class:`~concepts.dsl.tensor_state.NamedObjectTensorState` for more details.
    """

    name: str
    index: int


class MultidimensionalArrayInterface(object):
    """
    A multi-dimensional array inferface. At a high-level, this can be interpreted as a dictionary that maps
    feature names (keys) to multi-diemsntional tensors (value).
    """

    def __init__(self, all_feature_names: Iterable[str] = tuple()):
        self.all_feature_names = set(all_feature_names)

    def clone(self) -> 'MultidimensionalArrayInterface':
        """Clone the multidimensional array interface."""
        raise NotImplementedError()

    def get_feature(self, name: str) -> TensorValue:
        """Get the feature tensor with the given name."""
        raise NotImplementedError()

    def _set_feature_impl(self, name: str, feature: TensorValue):
        """Set the feature tensor with the given name. It is guaranteed that the name is in the all_feature_names."""
        raise NotImplementedError()

    def set_feature(self, name: str, feature: TensorValue):
        """Set the feature tensor with the given name."""
        if name not in self.all_feature_names:
            self.all_feature_names.add(name)
        self._set_feature_impl(name, feature)

    def update_feature(self, other_tensor_dict: Mapping[str, TensorValue]):
        """Update the feature tensors with the given tensor dict."""
        for key, value in other_tensor_dict.items():
            self.set_feature(key, value)

    def __contains__(self, item: str) -> bool:
        """Check if the given feature name is in the interface."""
        return item in self.all_feature_names

    def __getitem__(self, name: str) -> TensorValue:
        """Get the feature tensor with the given name."""
        return self.get_feature(name)

    def __setitem__(self, key, value):
        """Set the feature tensor with the given name."""
        self.set_feature(key, value)

    def keys(self) -> Iterable[str]:
        """Get the feature names."""
        return self.all_feature_names

    def values(self) -> Iterable[TensorValue]:
        """Get the feature tensors."""
        for key in self.all_feature_names:
            yield self.get_feature(key)

    def items(self) -> Iterable[Tuple[str, TensorValue]]:
        """Get the feature name-tensor pairs."""
        for key in self.all_feature_names:
            yield key, self.get_feature(key)


class TensorValueDict(MultidimensionalArrayInterface):
    """Basic tensor dict implementation."""

    def __init__(self, tensor_dict: Optional[Dict[str, TensorValue]] = None):
        if tensor_dict is None:
            tensor_dict = dict()

        super().__init__(tensor_dict.keys())
        self.tensor_dict = tensor_dict

    def clone(self) -> 'TensorValueDict':
        return type(self)({k: v.clone() for k, v in self.tensor_dict.items()})

    def get_feature(self, name: str) -> TensorValue:
        return self.tensor_dict[name]

    def _set_feature_impl(self, name, feature: TensorValue):
        self.tensor_dict[name] = feature


class TensorStateBase(object):
    """A state representation is essentially a mapping from feature names to tensors."""

    @property
    def batch_dims(self) -> int:
        raise NotImplementedError()

    @property
    def features(self) -> MultidimensionalArrayInterface:
        raise NotImplementedError()

    def clone(self) -> 'TensorStateBase':
        raise NotImplementedError()

    def __getitem__(self, name: str):
        return self.features[name]

    def __str__(self) -> str:
        raise NotImplementedError()

    def __repr__(self):
        return self.__str__()


ObjectNameArgument = Union[Iterable[str], Mapping[str, ObjectType]]
ObjectTypeArgument = Optional[Iterable[ObjectType]]


class NamedObjectStateMixin(object):
    """A state type mixin with named objects."""

    def __init__(self, object_names: ObjectNameArgument, object_types: ObjectTypeArgument = None):
        """A state type mixin with named objects.
        The object names can be either a list of names, or a mapping from names to :class:`ObjectType`'s.

            - If the `object_names` is a list of names, then the user should also specify a list of object types.
            - If the `object_names` is a mapping from names to :class:`ObjectType`'s, then the `object_types` argument should be None.

        Args:
            object_names: the object names.
            object_types: the object types.
        """
        if isinstance(object_names, Mapping):
            assert object_types is None, 'object_types should be None if object_names is a mapping.'
            self.object_names = tuple(object_names.keys())
            self.object_types = tuple(object_names.values())
        else:
            assert object_types is not None, 'object_types should not be None if object_names is not a mapping.'
            self.object_names = tuple(object_names)
            self.object_types = tuple(object_types)

        self.object_type2name: Dict[str, List[str]] = dict()
        self.object_name2index: Dict[Tuple[str, str], int] = dict()
        self.object_name2defaultindex: Dict[str, Tuple[str, int]] = dict()

        for name, obj_type in zip(self.object_names, self.object_types):
            self.object_type2name.setdefault(obj_type.typename, list()).append(name)
            self.object_name2index[name, obj_type.typename] = len(self.object_type2name[obj_type.typename]) - 1
            self.object_name2defaultindex[name] = obj_type.typename, len(self.object_type2name[obj_type.typename]) - 1
            for t in obj_type.iter_parent_types():
                self.object_type2name.setdefault(t.typename, list()).append(name)
                self.object_name2index[name, t.typename] = len(self.object_type2name[t.typename]) - 1

    @property
    def nr_objects(self) -> int:
        """The number of objects in the current state."""
        return len(self.object_types)

    def get_typename(self, name: str) -> str:
        """Get the typename of the object with the given name."""
        return self.object_name2defaultindex[name][0]

    def get_typed_index(self, name, typename: Optional[str] = None) -> int:
        """Get the typed index of the object with the given name.
        There is an additional typename argument to specify the type of the object.
        Because the same object can have multiple types (due to inheritence), the object can have multiple typed indices, one for each type.
        When the typename is not specified, the default type of the object is used (i.e., the most specific type).

        Args:
            name: the name of the object.
            typename: the typename of the object. If not specified, the default type of the object is used (i.e. the most specific type).

        Returns:
            the typed index of the object.
        """
        if typename is None or typename == AutoType.typename:
            return self.object_name2defaultindex[name][1]
        return self.object_name2index[name, typename]

    def get_nr_objects_by_type(self, typename: str) -> int:
        """Get the number of objects with the given type."""
        return len(self.object_type2name[typename])


class TensorState(TensorStateBase):
    """A state representation is essentially a mapping from feature names to tensors."""

    def __init__(self, features: Optional[Union[Mapping[str, Any], TensorValueDict]] = None, batch_dims: int = 0, internals: Optional[Dict[str, Any]] = None):
        """Initialize a state.

        Args:
            features: the features of the state.
            batch_dims: the number of batch dimensions.
            internals: the internal state of the state.
        """

        if features is None:
            features = dict()
        if internals is None:
            internals = dict()

        if isinstance(features, TensorValueDict):
            self._features = features
        else:
            self._features = TensorValueDict(features)
        self._batch_dims = batch_dims

        self._internals = dict(internals)

    @property
    def batch_dims(self) -> int:
        """The number of batchified dimensions. For the basic State, it should be 0."""
        return self._batch_dims

    @property
    def features(self) -> TensorValueDict:
        return self._features

    @property
    def internals(self) -> Dict[str, Any]:
        """Additional internal information about the state."""
        return self._internals

    def clone(self) -> 'TensorState':
        return type(self)(features=self._features.clone(), batch_dims=self._batch_dims, internals=self.clone_internals())

    def clone_internals(self):
        """Clone the internals."""
        return self.internals.copy()

    def summary_string(self) -> str:
        """Get a summary string of the state. The main difference between this and __str__ is that this function only formats the shape of intermediate tensors."""
        fmt = f'''{type(self).__name__}{{
  states:
'''
        for p in self.features.all_feature_names:
            feature = self.features[p]
            fmt += f'    {p}: {feature.format(content=False)}\n'
        fmt += self.extra_state_str()
        fmt += '}'
        return fmt

    def __str__(self):
        fmt = f'''{type(self).__name__}{{
  states:
'''
        for p in self.features.all_feature_names:
            tensor = self.features[p]
            fmt += f'    - {p}'
            fmt += ': ' + indent_text(str(tensor), level=2).strip() + '\n'
        fmt += self.extra_state_str()
        fmt += '}'
        return fmt

    def extra_state_str(self) -> str:
        """Extra state string."""
        return ''


class NamedObjectTensorState(TensorState, NamedObjectStateMixin):
    """A state type with named objects."""

    def __init__(self, features: Optional[Union[Mapping[str, Any], MultidimensionalArrayInterface]], object_names: ObjectNameArgument, object_types: ObjectTypeArgument = None, batch_dims: int = 0, internals: Optional[Mapping[str, Any]] = None):
        """Initialize the state.

        Args:
            features: the features of the state.
            object_types: the types of the objects.
            object_names: the names of the objects. If the object_names is a mapping, the object_types should be None.
            batch_dims: the number of batchified dimensions.
            internals: the internals of the state.
        """

        TensorState.__init__(self, features, batch_dims, internals)
        NamedObjectStateMixin.__init__(self, object_names, object_types)

    def clone(self) -> 'NamedObjectTensorState':
        return type(self)(features=self._features.clone(), object_types=self.object_types, object_names=self.object_names, batch_dims=self._batch_dims, internals=self.clone_internals())

    def extra_state_str(self) -> str:
        """Extra state string: add the objects."""
        if self.object_names is not None:
            objects_str = [f'{name} - {dtype.typename}' for name, dtype in zip(self.object_names, self.object_types)]
        else:
            objects_str = self.object_names
        return '  objects: ' + ', '.join(objects_str) + '\n'


def concat_states(*args: TensorState) -> TensorState:
    """Concatenate a list of states into a batch state.

    Args:
        *args: a list of states.

    Returns:
        a new state, which is the concatenation of the input states.
        This new state will have a new batch dimension.
    """

    if len(args) == 0:
        raise ValueError('No states to concatenate.')

    all_features = list(args[0].features.all_feature_names)

    # 1. Sanity checks.
    for state in args[1:]:
        assert len(all_features) == len(state.features.all_feature_names)
        for feature in all_features:
            assert feature in state.features.all_feature_names

    # 2. Put the same feature into a list.
    features = {feature_name: list() for feature_name in all_features}
    for state in args:
        for key, value in state.features.tensor_dict.items():
            features[key].append(value)

    # 3. Actually, compute the features.
    feature_names = list(features.keys())
    for feature_name in feature_names:
        features[feature_name] = concat_tvalues(*features[feature_name])

    # 4. Create the new state.
    state = args[0]
    kwargs: Dict[str, Any] = dict()
    if isinstance(state, NamedObjectTensorState):
        kwargs = dict(object_types=state.object_types, object_names=state.object_names)

    kwargs['features'] = features
    kwargs['batch_dims'] = args[0].batch_dims + 1
    return type(state)(**kwargs)


concepts/dsl/tensor_state.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : tensor_value.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 11/02/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data structures and simple shape-related operations for tensor values. Internally, we use torch.Tensor to represent tensor values."""

from dataclasses import dataclass
from copy import deepcopy
from typing import TYPE_CHECKING, Any, Optional, Union, Tuple, Iterable, Sequence, List

import warnings
import numpy as np
import torch
import torch.nn.functional as F
import jactorch
from jacinle.utils.printing import indent_text

from concepts.dsl.dsl_types import TensorValueTypeBase, BOOL, INT64, FLOAT32, NamedTensorValueType, ScalarValueType, VectorValueType, PyObjValueType, QINDEX
from concepts.dsl.value import ValueBase

if TYPE_CHECKING:
    from concepts.dsl.constraint import OptimisticValue


__all__ = [
    'TensorizedPyObjValues', 'MaskedTensorStorage', 'TensorValue',  # Basic tensorized representations.
    'from_tensor', 'vector_values', 'scalar',  # Tensor value creation.
    'concat_tvalues', 'index_tvalue', 'set_index_tvalue', 'expand_as_tvalue', 'expand_tvalue',
    'simple_quantize_tvalue'
]


class TensorizedPyObjValues(object):
    """A value object with nested lists of Python objects. Note that this implies that we can not use list as the Python object type when shape is not specified."""

    def __init__(self, dtype: PyObjValueType, values: Union[Any, List[Any], List[List[Any]], List[List[List[Any]]]], shape: Optional[Tuple[int, ...]] = None):
        self.dtype = dtype
        if isinstance(values, np.ndarray):
            self.values = values
            self.shape = shape if shape is not None else self._infer_shape()
        elif shape is None:
            self.values = np.array(values, dtype=object)
            self.shape = shape if shape is not None else self._infer_shape()
        else:
            self.values = np.empty(shape, dtype=object)
            _recursive_assign(self.values, values)
            self.shape = shape

    @property
    def ndim(self):
        return len(self.shape)

    def numel(self) -> int:
        return int(np.prod(self.shape))

    def item(self):
        assert self.ndim == 0
        return self.values.item()

    def fast_index(self, arguments: Tuple[int, ...]):
        # assert len(arguments) == self.ndim
        # value = self.values
        # for i in arguments:
        #     value = value[i]
        # return value
        return self.values[arguments]

    def _infer_shape(self):
        """Infer the shape of the values by recursively checking the first element."""

        # if self.dtype.pyobj_type is list:
        #     raise TypeError('BatchedPyObjValue can not have list as the Python object type when shape is not specified.')

        # def _infer_shape_recursive(values):
        #     if type(values) is list:
        #         return (len(values), ) + _infer_shape_recursive(values[0])
        #     else:
        #         return tuple()

        # return _infer_shape_recursive(self.values)

        return self.values.shape

    @classmethod
    def make_empty(cls, dtype: PyObjValueType, shape: Iterable[int]):
        shape = tuple(shape)
        values = np.empty(shape, dtype=object)
        return cls(dtype, values, shape)

        # def _make_empty_recursive(s):
        #     if len(s) == 0:
        #         return None
        #     return [_make_empty_recursive(s[1:]) for _ in range(s[0])]
        # return cls(dtype, _make_empty_recursive(shape), shape)

    def clone(self, deep_copy: bool = False):
        """Recursively clone the values."""

        if deep_copy:
            return deepcopy(self)
        else:
            return type(self)(self.dtype, self.values.copy(), self.shape)

        # if deep_copy:
        #     new_values = deepcopy(self.values)
        # else:
        #     def _clone_recursive(values):
        #         if type(values) is list:
        #             return [_clone_recursive(v) for v in values]
        #         else:
        #             return values
        #     new_values = _clone_recursive(self.values)

        # return TensorizedPyObjValues(self.dtype, new_values, self.shape)

    def permute(self, indices: Tuple[int, ...]):
        """Permute the axes in the nested list of values."""
        return type(self)(self.dtype, np.transpose(self.values, indices), tuple(self.shape[i] for i in indices))

    def expand(self, target_size: Tuple[int, ...]):
        return type(self)(self.dtype, np.broadcast_to(self.values, target_size), target_size)

    def __str__(self):
        if self.values.ndim == 0:
            return str(self.values.item())
        return f'TensorizedPyObjValues[{self.dtype}, shape={self.shape}]'

    def __repr__(self):
        return self.__str__()


def _recursive_assign(np_array: np.ndarray, values: Union[Any, List[Any], List[List[Any]], List[List[List[Any]]]]):
    if len(np_array.shape) == 0:
        np_array[()] = values
    elif len(np_array.shape) == 1:
        for i, v in enumerate(values):
            np_array[i] = v
    else:
        for i, v in enumerate(values):
            _recursive_assign(np_array[i], v)


@dataclass
class MaskedTensorStorage(object):
    """A storage for quantized tensors."""

    value: Union[torch.Tensor, TensorizedPyObjValues]
    """The unquantized tensor."""

    mask: Optional[torch.Tensor] = None
    """The mask of the value. If not None, entry = 1 if the value is valid, and 0 otherwise."""

    optimistic_values: Optional[torch.Tensor] = None
    """The optimistic values for the tensor. 0 for non-optimistic values."""

    quantized_values: Optional[torch.Tensor] = None
    """The quantized values for the tensor. -1 for non-quantized values."""


class TensorValue(ValueBase):
    """A value object with an internal :class:`torch.Tensor` storage."""

    def __init__(
        self,
        dtype: Union[TensorValueTypeBase, PyObjValueType],
        batch_variables: Union[Iterable[str], int],
        tensor: Union[torch.Tensor, TensorizedPyObjValues, MaskedTensorStorage],
        batch_dims: int = 0,
        quantized: Optional[bool] = None,
        *,
        _check_tensor: bool = False,
        _mask_certified_flag: bool = True
    ):
        """Instantiate a Value object for storing intermediate computation results.

        The tensor is assumed to have the following layout: ``tensor[B1, B2, B3, ..., V1, V2, V3, ..., D1, D2, D3, ...]``.

        - The first `batch_dims` dimensions are "batch".
        - The next `len(batch_variables)` dimensions are "variables".
        - The next `dtype.ndim()` dimensions are data dimensions (e.g., images, vectors).
        - A special case is that `dtype.ndim()` can be zero (scalar).

        Args:
            dtype: The data type of the Value object.
            batch_variables: A sequence of variables that are processed in "batch." This typically corresponds to "quantified variables."
                It can also be a single integer, indicating the number of batched variables.
            tensor: The actual tensor
            batch_dims: The additional batch dimensions at the beginning. Defaults to 0.
            _check_tensor: internal flag, whether to run the tensor shape/type sanity check.
            _mask_certified_flag: internal flag, indicating whether self.tensor_mask is guaranteed to be the correct mask.
                This flag will be marked false when we do expand_as.
        """
        super().__init__(dtype)

        if isinstance(batch_variables, int):
            batch_variables = tuple([f'?#{i}' for i in range(batch_variables)])
        self.batch_variables = tuple(batch_variables)
        self.batch_dims = batch_dims

        self.tensor_mask = None
        self.tensor_optimistic_values = None
        self.tensor_quantized_values = None
        self.quantized = quantized if quantized is not None else self._guess_quantized()

        if quantized is not None:
            if not hasattr(TensorValue, '_deprecated_quantized_warning'):
                TensorValue._deprecated_quantized_warning = True
                warnings.warn('TensorValue.quantized is deprecated. Please remove any usage of this flag. This flag will be removed by 2023/12/31', DeprecationWarning)

        if isinstance(tensor, MaskedTensorStorage):
            self.tensor = tensor.value
            self.tensor_mask = tensor.mask
            self.tensor_optimistic_values = tensor.optimistic_values
            self.tensor_quantized_values = tensor.quantized_values
        elif isinstance(tensor, TensorizedPyObjValues):
            assert isinstance(self.dtype, PyObjValueType)
            self.tensor = tensor
        elif isinstance(tensor, torch.Tensor):
            # TODO (Joy Hsu @ 2023/10/04): assert TensorValueTypeBase, PyObjValueType.
            self.tensor = tensor

        assert isinstance(self.tensor, (torch.Tensor, TensorizedPyObjValues))
        if _check_tensor:
            self._check_tensor('tensor')
            self._check_tensor('tensor_optimistic_values', is_index=True)
            self._check_tensor('tensor_quantized_values', is_index=True)
            self._check_tensor('tensor_mask', is_index=True)

        self._mask_certified_flag = _mask_certified_flag

        self._backward_function = None
        self._backward_args = tuple()
        self.tensor_grad = None

    dtype: Union[TensorValueTypeBase, PyObjValueType]
    """The data type of the Value object."""

    batch_variables: Tuple[str, ...]
    """The list of batch variable names."""

    batch_dims: int
    """Additional batch dimensions at the beginning."""

    quantized: bool
    """Whether the values in self.tensor is quantized."""

    tensor: Union[torch.Tensor, TensorizedPyObjValues]
    """The internal tensor storage."""

    tensor_mask: Optional[torch.Tensor]
    """A mask of the tensor, indicating which elements are valid."""

    tensor_optimistic_values: Optional[torch.Tensor]
    """The optimistic values for the tensor. 0 for non-optimistic values."""

    tensor_quantized_values: Optional[torch.Tensor]
    """The quantized values for the tensor. -1 for non-quantized values."""

    def _guess_quantized(self) -> bool:
        return isinstance(self.dtype, TensorValueTypeBase) and self.dtype.is_intrinsically_quantized()

    @property
    def total_batch_dims(self):
        return self.batch_dims + len(self.batch_variables)

    @property
    def nr_variables(self):
        return len(self.batch_variables)

    @property
    def is_torch_tensor(self):
        return isinstance(self.tensor, torch.Tensor)

    @property
    def is_tensorized_pyobj(self):
        return isinstance(self.tensor, TensorizedPyObjValues)

    def get_variable_size(self, variable_name_or_index: Union[str, int]) -> int:
        if isinstance(variable_name_or_index, str):
            variable_name_or_index = self.batch_variables.index(variable_name_or_index)
        if self.is_torch_tensor:
            return self.tensor.size(variable_name_or_index + self.batch_dims)
        return self.tensor.shape[variable_name_or_index + self.batch_dims]

    def to_masked_tensor_storage(self, clone_tensor: bool = False) -> MaskedTensorStorage:
        if clone_tensor:
            return MaskedTensorStorage(_maybe_clone_tensor(self.tensor), _maybe_clone_tensor(self.tensor_mask), _maybe_clone_tensor(self.tensor_optimistic_values), _maybe_clone_tensor(self.tensor_quantized_values))
        return MaskedTensorStorage(self.tensor, self.tensor_mask, self.tensor_optimistic_values, self.tensor_quantized_values)

    def clone(self, clone_tensor=True):
        storage = self.to_masked_tensor_storage(clone_tensor=clone_tensor)
        return type(self)(
            self.dtype, self.batch_variables, storage, self.batch_dims,
            _check_tensor=False, _mask_certified_flag=self._mask_certified_flag
        )

    def rename_batch_variables(self, new_variables: Sequence[str], clone: bool = False):
        assert len(self.batch_variables) == len(new_variables)
        rv = self.clone() if clone else self
        rv.batch_variables = tuple(new_variables)
        return rv

    def init_tensor_optimistic_values(self):
        if self.tensor_optimistic_values is None:
            if self.is_torch_tensor:
                self.tensor_optimistic_values = torch.zeros(self.tensor.shape[:self.total_batch_dims], dtype=torch.int64, device=self.tensor.device)
            else:
                self.tensor_optimistic_values = torch.zeros(self.tensor.shape, dtype=torch.int64)

    # Simple creation.

    @classmethod
    def from_tensor(cls, value: torch.Tensor, dtype: Optional[TensorValueTypeBase] = None, batch_variables: Optional[Iterable[str]] = None) -> 'TensorValue':
        return from_tensor(value, dtype, batch_variables)

    @classmethod
    def from_tensorized_pyobj(cls, value: TensorizedPyObjValues, dtype: Optional[PyObjValueType] = None, batch_variables: Optional[Iterable[str]] = None) -> 'TensorValue':
        return from_tensorized_pyobj(value, dtype, batch_variables)

    @classmethod
    def from_values(cls, *args: Any, dtype: Optional[TensorValueTypeBase] = None) -> 'TensorValue':
        return vector_values(*args, dtype=dtype)

    @classmethod
    def from_scalar(cls, value: Any, dtype: Optional[Union[TensorValueTypeBase, PyObjValueType]] = None) -> 'TensorValue':
        return scalar(value, dtype=dtype)

    @classmethod
    def make_empty(cls, dtype: Union[TensorValueTypeBase, PyObjValueType], batch_variables: Iterable[str] = tuple(), batch_sizes: Iterable[int] = tuple(), batch_dims: int = 0):
        if isinstance(dtype, TensorValueTypeBase):
            tensor = torch.zeros(tuple(batch_sizes) + dtype.size_tuple(), dtype=dtype.tensor_dtype())
        else:
            tensor = TensorizedPyObjValues.make_empty(dtype, batch_sizes)
        return cls(dtype, batch_variables, tensor, batch_dims=batch_dims)

    @classmethod
    def from_optimistic_value(cls, value: 'OptimisticValue') -> 'TensorValue':
        rv = cls.make_empty(value.dtype)
        rv.init_tensor_optimistic_values()
        rv.tensor_optimistic_values[...] = value.identifier
        return rv

    @classmethod
    def from_optimistic_value_int(cls, identifier: int, dtype: Union[TensorValueTypeBase, PyObjValueType]) -> 'TensorValue':
        rv = cls.make_empty(dtype)
        rv.init_tensor_optimistic_values()
        rv.tensor_optimistic_values[...] = identifier
        return rv

    # Simple tensor access.

    @property
    def is_scalar(self):
        return self.total_batch_dims == 0 and (isinstance(self.dtype, PyObjValueType) or self.dtype.ndim() == 0)

    def item(self) -> Union[torch.Tensor, Any, 'OptimisticValue']:
        from .constraint import is_optimistic_value, OptimisticValue

        assert self.is_scalar
        if self.tensor_optimistic_values is not None and is_optimistic_value(self.tensor_optimistic_values.item()):
            return OptimisticValue(self.dtype, int(self.tensor_optimistic_values.item()))
        if self.is_tensorized_pyobj:
            return self.tensor.values.item()
        return self.tensor.item()

    @property
    def is_single_elem(self):
        return self.total_batch_dims == 0

    def single_elem(self) -> Union[torch.Tensor, Any, 'OptimisticValue']:
        from .constraint import is_optimistic_value, OptimisticValue

        assert self.is_single_elem
        if self.tensor_optimistic_values is not None and is_optimistic_value(self.tensor_optimistic_values.item()):
            return OptimisticValue(self.dtype, int(self.tensor_optimistic_values.item()))
        if self.is_tensorized_pyobj:
            return self.tensor.values.item()
        return self.tensor

    def fast_index(self, arguments: Tuple[int, ...], wrap: bool = True) -> Union[torch.Tensor, Any, 'OptimisticValue']:
        from .constraint import is_optimistic_value, OptimisticValue

        arguments = tuple(arguments)
        assert len(arguments) == self.total_batch_dims
        if self.tensor_optimistic_values is not None and is_optimistic_value(self.tensor_optimistic_values[arguments]):
            return OptimisticValue(self.dtype, int(self.tensor_optimistic_values[arguments]))
        if self.is_tensorized_pyobj:
            if wrap:
                return TensorValue.from_scalar(self.tensor.fast_index(arguments), self.dtype)
            return self.tensor.fast_index(arguments)
        if wrap:
            return TensorValue.from_scalar(self.tensor[arguments], self.dtype)
        return self.tensor[arguments]

    def __bool__(self):
        if self.dtype == BOOL:
            return bool(self.item())
        else:
            raise TypeError('Cannot convert Value object {} into bool.'.format(self))

    # Value indexing.

    def index(self, key: 'ValueIndexType') -> 'TensorValue':
        if self.is_torch_tensor:
            return index_tvalue(self, key)
        else:
            return index_pyobj_value(self, key)

    def set_index(self, key: 'ValueIndexType', value: 'ValueSetIndexType') -> 'TensorValue':
        if self.is_torch_tensor:
            return set_index_tvalue(self, key, value)
        else:
            return set_index_pyobj_value(self, key, value)

    def __getitem__(self, item) -> 'TensorValue':
        return self.index(item)

    def __setitem__(self, key: 'ValueIndexType', value: 'ValueSetIndexType'):
        self.set_index(key, value)

    # Value shape manipulation.

    def expand(self, batch_variables: Iterable[str], batch_sizes: Iterable[int]) -> 'TensorValue':
        # NB(Jiayuan Mao @ 2023/08/16): expand_tvalue work for both pyobj and torch tensor.
        return expand_tvalue(self, batch_variables, batch_sizes)

    def expand_as(self, other: 'TensorValue') -> 'TensorValue':
        return expand_as_tvalue(self, other)

    # Simple value quantization.

    def is_simple_quantizable(self) -> bool:
        return is_tvalue_simple_quantizable(self)

    def simple_quantize(self) -> 'TensorValue':
        return simple_quantize_tvalue(self)

    # Stringify functions.

    STR_MAX_TENSOR_SIZE = 100

    @classmethod
    def set_print_options(cls, max_tensor_size=STR_MAX_TENSOR_SIZE):
        cls.STR_MAX_TENSOR_SIZE = max_tensor_size

    def format(self, content: bool = True, short: bool = False) -> str:
        from .constraint import is_optimistic_value, optimistic_value_id

        if short:
            if self.total_batch_dims == 0 and self.tensor_optimistic_values is not None and is_optimistic_value(self.tensor_optimistic_values.item()):
                return f'@{optimistic_value_id(self.tensor_optimistic_values.item())}'

            if self.is_scalar:
                if self.is_torch_tensor:
                    return str(self.item())
                else:
                    value = self.tensor.values.item()
                    return str(f'{type(value).__name__}@{hex(id(value))}')
            if content:
                if self.is_torch_tensor:
                    return str(self.tensor.tolist())
                else:
                    return str(self.tensor.values.tolist())
            else:
                axes = ', '.join(('B', ) * self.batch_dims + self.batch_variables)
                return f'{self.dtype}<{axes}>'

        if self.total_batch_dims == 0 and self.tensor_optimistic_values is not None and is_optimistic_value(self.tensor_optimistic_values.item()):
            return f'OptValue[{self.dtype}]{{@{optimistic_value_id(self.single_elem().identifier)}}}'

        axes = ', '.join(('B', ) * self.batch_dims + self.batch_variables)
        quantized_flag = ', quantized' if self.quantized else ''
        backward_flag = ', backward=' + str(self._backward_function) if self._backward_function is not None else ''
        tensor_content = ''

        if content:
            if self.tensor_grad is not None:
                tensor_content = str(self.tensor) + '\ngrad:\n' + str(self.tensor_grad)
            else:
                tensor_content = str(self.tensor)

            if self.tensor.numel() < type(self).STR_MAX_TENSOR_SIZE:
                if self.tensor.numel() < 10 and self.tensor.ndim <= 1:
                    tensor_content = '{' + tensor_content.replace('\n', ' ') + '}'
                else:
                    tensor_content = '{\n' + indent_text(tensor_content) + '\n}'

            return f'Value[{self.dtype}, axes=[{axes}], tdtype={self.tensor.dtype}, tdshape={tuple(self.tensor.shape)}{quantized_flag}{backward_flag}]{tensor_content}'

        return f'Value[{self.dtype}, axes=[{axes}], tdtype={self.tensor.dtype}, tdshape={tuple(self.tensor.shape)}{quantized_flag}{backward_flag}]'

    def short_str(self):
        return self.format(content=True, short=True)

    def __str__(self):
        return self.format(content=True)

    def __repr__(self):
        return self.__str__()

    # Internal dtype and shape sanity check.

    def _check_tensor(self, tensor_name: str, is_index: Optional[bool] = False):
        tensor = getattr(self, tensor_name)
        if tensor is None:
            return
        try:
            if isinstance(self.dtype, NamedTensorValueType):
                dtype = self.dtype.base_type
            else:
                dtype = self.dtype

            if isinstance(dtype, ScalarValueType):
                if self.total_batch_dims == 0:
                    if dtype in (BOOL, INT64):
                        assert isinstance(tensor, (bool, int, torch.Tensor))
                        if isinstance(tensor, (bool, int)):
                            setattr(self, tensor_name, torch.tensor(tensor, dtype=torch.int64))
                        else:
                            if dtype == BOOL and tensor.dtype == torch.bool:
                                pass
                            elif dtype == BOOL and tensor.dtype == torch.int64:
                                setattr(self, tensor_name, tensor.to(torch.bool))
                            elif dtype == INT64 and tensor.dtype == torch.int64:
                                pass
                            else:
                                assert tensor.dtype == torch.float32
                    else:
                        raise TypeError('Unsupported dtype: {}.'.format(self.dtype))
                else:
                    assert torch.is_tensor(tensor)
                    assert tensor.ndim == len(self.batch_variables) + self.batch_dims
                    if dtype in (BOOL, INT64):
                        if dtype == BOOL and tensor.dtype == torch.bool:
                            pass
                        elif dtype == BOOL and tensor.dtype == torch.int64:
                            setattr(self, tensor_name, tensor.to(torch.bool))
                        elif dtype == INT64 and tensor.dtype == torch.int64:
                            pass
                        else:
                            assert tensor.dtype == torch.float32
                    else:
                        raise TypeError('Unsupported dtype: {}.'.format(self.dtype))
            elif isinstance(dtype, VectorValueType):
                assert torch.is_tensor(tensor)
                if is_index:
                    assert tensor.ndim == len(self.batch_variables) + self.batch_dims
                else:
                    assert tensor.ndim == len(self.batch_variables) + dtype.ndim() + self.batch_dims
            elif isinstance(dtype, PyObjValueType):
                assert isinstance(tensor, TensorizedPyObjValues)
                assert tensor.dtype == dtype
                assert tensor.ndim == len(self.batch_variables) + self.batch_dims
            else:
                raise NotImplementedError('Unsupported dtype: {}.'.format(self.dtype))
        except AssertionError as e:
            axes = ', '.join(self.batch_variables)
            raise ValueError(f'Tensor {tensor_name} shape/dtype mismatch for Value[{self.dtype}, axes=[{axes}], batch_dims={self.batch_dims}, tdtype={tensor.dtype}, tdshape={tuple(tensor.shape)}]') from e


    TRUE: 'TensorValue'
    FALSE: 'TensorValue'


def _maybe_clone_tensor(tensor: Optional[Union[torch.Tensor, TensorizedPyObjValues]]) -> Optional[torch.Tensor]:
    return tensor if tensor is None else tensor.clone()


def from_tensor(value: torch.Tensor, dtype: Optional[TensorValueTypeBase] = None, batch_variables: Optional[List[str]] = None) -> TensorValue:
    if dtype is None:
        dtype = TensorValueTypeBase.from_tensor(value)
    if batch_variables is None:
        batch_variables = list()
    return TensorValue(dtype, batch_variables, value)


def from_tensorized_pyobj(value: TensorizedPyObjValues, dtype: Optional[PyObjValueType] = None, batch_variables: Optional[List[str]] = None) -> TensorValue:
    if dtype is None:
        dtype = PyObjValueType(value.dtype.pyobj_type, value.dtype.shape)
    if batch_variables is None:
        batch_variables = len(value.shape)
    return TensorValue(dtype, batch_variables, value)


def vector_values(*args: Any, dtype: Optional[TensorValueTypeBase] = None) -> TensorValue:
    if dtype is None:
        dtype = VectorValueType(FLOAT32, len(args), 0)
    assert isinstance(dtype, TensorValueTypeBase)
    tensor_dtype = dtype.tensor_dtype()
    return TensorValue(dtype, [], torch.tensor(args, dtype=tensor_dtype))


def scalar(value: Union[TensorValue, torch.Tensor, bool, float, int, Any], dtype: Optional[Union[TensorValueTypeBase, PyObjValueType]] = None) -> TensorValue:
    if isinstance(value, TensorValue):
        return value

    if dtype is None:
        if not isinstance(value, torch.Tensor):
            value = torch.tensor(value)
        dtype = TensorValueTypeBase.from_tensor(value)

    if torch.is_tensor(value):
        return TensorValue(dtype, [], value)

    if isinstance(dtype, PyObjValueType):
        if isinstance(value, int):
            raise RuntimeError('Invalid branch: report to developers.')
        else:
            return TensorValue(dtype, [], TensorizedPyObjValues(dtype, value, shape=()))

    assert isinstance(dtype, TensorValueTypeBase)
    tensor_dtype = dtype.tensor_dtype()
    return TensorValue(dtype, [], torch.tensor(value, dtype=tensor_dtype))


def t(value: torch.Tensor, dtype: Optional[TensorValueTypeBase] = None, batch_variables: Optional[List[str]] = None) -> TensorValue:
    """Alias for :func:`from_tensor`."""
    return from_tensor(value, dtype, batch_variables)


def v(*args: Any, dtype: Optional[TensorValueTypeBase] = None) -> TensorValue:
    """Alias for :func:`values`."""
    return vector_values(*args, dtype=dtype)


def s(value: Any, dtype: Optional[Union[TensorValueTypeBase, PyObjValueType]] = None) -> TensorValue:
    """Alias for :func:`scalar`."""
    return scalar(value, dtype)


TensorValue.TRUE = TensorValue.from_scalar(True, BOOL)
TensorValue.FALSE = TensorValue.from_scalar(False, BOOL)



def concat_tvalues(*args: TensorValue):  # will produce a Value with batch_dims == 1, but the input can be either 0-batch or 1-batch.
    assert len(args) > 0
    include_tensor_optimistic_values = any([v.tensor_optimistic_values is not None for v in args])
    include_tensor_quantized_values = any([v.tensor_quantized_values is not None for v in args])

    # Sanity check.
    for value in args[1:]:
        assert value.is_torch_tensor
        assert value.dtype == args[0].dtype
        assert value.batch_variables == args[0].batch_variables
        if include_tensor_optimistic_values:
            pass  # we have default behavior for None.
        if include_tensor_quantized_values:
            assert value.tensor_quantized_values is not None

    device = args[0].tensor.device

    # Collect all tensors.
    all_tensor = [v.tensor for v in args]
    all_tensor_mask = [v.tensor_mask for v in args]
    all_tensor_optimistic_values = [v.tensor_optimistic_values for v in args]
    all_tensor_quantized_values = [v.tensor_quantized_values for v in args]

    target_shape = tuple([
        max([v.get_variable_size(i) for v in args])
        for i in range(args[0].nr_variables)
    ])
    for i in range(len(args)):
        tensor, tensor_mask, tensor_optimistic_values, tensor_quantized_values = all_tensor[i], all_tensor_mask[i], all_tensor_optimistic_values[i], all_tensor_quantized_values[i]
        all_tensor[i] = _pad_tensor(tensor, target_shape, args[i].dtype, args[i].batch_dims)

        if tensor_mask is None:
            tensor_mask = torch.ones(target_shape, dtype=torch.bool, device=device)
        else:
            tensor_mask = _pad_tensor(tensor_mask, target_shape, args[i].dtype, args[i].batch_dims)
        all_tensor_mask[i] = tensor_mask

        if include_tensor_optimistic_values:
            if tensor_optimistic_values is None:
                tensor_optimistic_values = torch.zeros(target_shape, dtype=torch.int64, device=device)
            else:
                tensor_optimistic_values = _pad_tensor(tensor_optimistic_values, target_shape, args[i].dtype, args[i].batch_dims)
            all_tensor_optimistic_values[i] = tensor_optimistic_values

        if include_tensor_quantized_values:
            tensor_quantized_values = _pad_tensor(tensor_quantized_values, target_shape, args[i].dtype, args[i].batch_dims)
            all_tensor_quantized_values[i] = tensor_quantized_values

        if args[0].batch_dims == 0:
            all_tensor[i] = all_tensor[i].unsqueeze(0)
            all_tensor_mask[i] = all_tensor_mask[i].unsqueeze(0)
            all_tensor_optimistic_values[i] = all_tensor_optimistic_values[i].unsqueeze(0) if all_tensor_optimistic_values[i] is not None else None
            all_tensor_quantized_values[i] = all_tensor_quantized_values[i].unsqueeze(0) if all_tensor_quantized_values[i] is not None else None
        else:
            assert args[0].batch_dims == 1

    masked_tensor_storage = MaskedTensorStorage(
        torch.cat(all_tensor, dim=0),
        torch.cat(all_tensor_mask, dim=0),
        torch.cat(all_tensor_optimistic_values, dim=0) if include_tensor_optimistic_values else None,
        torch.cat(all_tensor_quantized_values, dim=0) if include_tensor_quantized_values else None
    )
    return TensorValue(args[0].dtype, args[0].batch_variables, masked_tensor_storage, batch_dims=1)


def _pad_tensor(tensor: torch.Tensor, target_shape: Iterable[int], dtype: TensorValueTypeBase, batch_dims: int, constant_value: float = 0.0):
    target_shape = tuple(target_shape)
    paddings = list()
    for size, max_size in zip(tensor.size()[batch_dims:], target_shape):
        paddings.extend((max_size - size, 0))
    if tensor.dim() - batch_dims == len(target_shape):
        pass
    elif tensor.dim() - batch_dims == len(target_shape) + dtype.ndim():
        paddings.extend([0 for _ in range(dtype.ndim() * 2)])
    else:
        raise ValueError('Shape error during tensor padding.')
    paddings.reverse()  # no need to add batch_dims.
    return F.pad(tensor, paddings, "constant", constant_value)


ValueIndexElementType = Union[int, slice, List[int], torch.Tensor]
ValueIndexType = Optional[Union[ValueIndexElementType, Tuple[ValueIndexElementType, ...]]]
ValueSetIndexType = Union[bool, int, float, np.ndarray, torch.Tensor, TensorValue]


def index_tvalue(value: TensorValue, indices: ValueIndexType) -> TensorValue:
    assert value.is_torch_tensor, 'Use index_pyobj_value for non-tensor values.'
    if isinstance(indices, (int, list, torch.Tensor)) or indices == QINDEX:
        indices = (indices, )
    assert indices is None or isinstance(indices, tuple)

    if indices is None:
        return value.clone()

    assert len(indices) == value.nr_variables
    batch_variables = list()
    for i, idx in enumerate(indices):
        if idx == QINDEX:
            batch_variables.append(value.batch_variables[i])

    if value.batch_dims == 0:
        return type(value)(
            value.dtype, batch_variables, MaskedTensorStorage(
                value.tensor[indices] if value.tensor is not None else None,
                value.tensor_mask[indices] if value.tensor_mask is not None else None,
                value.tensor_optimistic_values[indices] if value.tensor_optimistic_values is not None else None,
                value.tensor_quantized_values[indices] if value.tensor_quantized_values is not None else None
            ), batch_dims=value.batch_dims, _mask_certified_flag=value._mask_certified_flag
        )
    elif value.batch_dims == 1:
        indices = (jactorch.batch,) + indices
        return type(value)(
            value.dtype, batch_variables, MaskedTensorStorage(
                jactorch.bvindex(value.tensor)[indices][0] if value.tensor is not None else None,
                jactorch.bvindex(value.tensor_mask)[indices][0] if value.tensor_mask is not None else None,
                jactorch.bvindex(value.tensor_optimistic_values)[indices][0] if value.tensor_optimistic_values is not None else None,
                jactorch.bvindex(value.tensor_quantized_values)[indices][0] if value.tensor_quantized_values is not None else None
            ), batch_dims=value.batch_dims, _mask_certified_flag=value._mask_certified_flag
        )
    else:
        raise NotImplementedError('Unsupported batched dims: {}.'.format(value.batch_dims))


def set_index_tvalue(lhs: TensorValue, indices: ValueIndexType, rhs: ValueSetIndexType) -> TensorValue:
    if isinstance(indices, (int, torch.Tensor)) or indices == QINDEX:
        indices = (indices, )
    if isinstance(indices, list):
        raise RuntimeError('Undefined behavior for list indices. Use either tuple([...], ) for single-dimension indexing or tuple(..., ...) for multi-dimension indexing.')
    assert indices is None or isinstance(indices, tuple)

    optimistic_values, quantized_values = None, None
    if isinstance(rhs, (bool, int, float)):
        rhs = torch.tensor(rhs, dtype=lhs.tensor.dtype, device=lhs.tensor.device)
    elif isinstance(rhs, np.ndarray):
        rhs = torch.from_numpy(rhs).to(lhs.tensor.dtype).to(lhs.tensor.device)
    if isinstance(rhs, TensorValue):
        optimistic_values, quantized_values = rhs.tensor_optimistic_values, rhs.tensor_quantized_values
        rhs = rhs.tensor

    if lhs.batch_dims == 0:
        if indices is None or len(indices) == 0:
            lhs.tensor = rhs
            lhs.tensor_optimistic_values = optimistic_values
            lhs.tensor_quantized_values = quantized_values
        else:
            # We have to this cloning. Consider the following case:
            # v[0] = v[1]
            indices = indices[0] if len(indices) == 1 else indices

            lhs.tensor = lhs.tensor.clone()
            lhs.tensor[indices] = rhs
            if optimistic_values is not None:
                lhs.init_tensor_optimistic_values()
                lhs.tensor_optimistic_values = lhs.tensor_optimistic_values.clone()
                lhs.tensor_optimistic_values[indices] = optimistic_values
            if lhs.tensor_quantized_values is not None:
                assert quantized_values is not None
                lhs.tensor_quantized_values = lhs.tensor_quantized_values.clone()
                lhs.tensor_quantized_values[indices] = quantized_values
    else:
        raise NotImplementedError('Unsupported batched dims: {}.'.format(lhs.batch_dims))

    return lhs


def index_pyobj_value(value: TensorValue, indices: ValueIndexType) -> TensorValue:
    if isinstance(indices, (int, slice)):
        indices = (indices, )

    assert indices is None or isinstance(indices, tuple)

    if indices is None:
        return value.clone()

    assert len(indices) == value.nr_variables
    batch_variables = list()
    sizes = list()
    for i, idx in enumerate(indices):
        if idx == QINDEX:
            batch_variables.append(value.batch_variables[i])
            sizes.append(value.get_variable_size(i))

    if value.batch_dims == 0:
        values = _recursive_index_pyobj_value(value.tensor.values, indices)
        values = TensorizedPyObjValues(value.dtype, values, sizes)
        return type(value)(value.dtype, batch_variables, values, batch_dims=value.batch_dims)
    elif value.batch_dims == 1:
        values = [_recursive_index_pyobj_value(value.tensor.values[i], [indexer[i] for indexer in indices]) for i in range(len(value.tensor.values))]
        assert isinstance(values, np.ndarray)
        values = TensorizedPyObjValues(value.dtype, values, values.shape)
        return type(value)(value.dtype, batch_variables, values, batch_dims=value.batch_dims)
    else:
        raise NotImplementedError('Unsupported batched dims: {}.'.format(value.batch_dims))


def _recursive_index_pyobj_value(np_array, indices):
    return np_array[indices]
    # if len(indices) == 1:
    #     if isinstance(indices[0], int):
    #         return nested_list[indices[0]]
    #     elif indices[0] == QINDEX:
    #         return nested_list
    #     else:
    #         raise ValueError(f'Invalid index: {indices[0]}.')

    # if isinstance(indices[0], int):
    #     return _recursive_index_pyobj_value(nested_list[indices[0]], indices[1:])
    # elif indices[0] == QINDEX:
    #     return [_recursive_index_pyobj_value(nested_list[i], indices[1:]) for i in range(len(nested_list))]
    # else:
    #     raise ValueError(f'Invalid index: {indices[0]}.')


def set_index_pyobj_value(lhs: TensorValue, indices: ValueIndexType, rhs: ValueSetIndexType) -> TensorValue:
    if isinstance(indices, (int, slice)):
        indices = (indices, )
    assert indices is None or isinstance(indices, tuple)

    is_scalar_rhs = not isinstance(rhs, TensorValue)

    if lhs.batch_dims == 0:
        _recursive_set_index_pyobj_value(lhs.tensor.values, indices, rhs, is_scalar_rhs)
    else:
        raise NotImplementedError('Unsupported batched dims: {}.'.format(lhs.batch_dims))
    return lhs


def _recursive_set_index_pyobj_value(np_array, indices: Tuple[Union[int, slice], ...], rhs, is_scalar_rhs):
    np_array[indices] = rhs

    # if len(indices) == 1:
    #     if isinstance(indices[0], int):
    #         nested_list[indices[0]] = rhs
    #     elif indices[0] == QINDEX:
    #         if is_scalar_rhs:
    #             for i in range(len(nested_list)):
    #                 nested_list[i] = rhs
    #         else:
    #             assert len(nested_list) == len(rhs.tensor.values)
    #             for i in range(len(nested_list)):
    #                 nested_list[i] = rhs[i]
    #     else:
    #         raise ValueError(f'Invalid index: {indices[0]}.')
    # else:
    #     if isinstance(indices[0], int):
    #         _recursive_set_index_pyobj_value(nested_list[indices[0]], indices[1:], rhs, is_scalar_rhs)
    #     elif indices[0] == QINDEX:
    #         for i in range(len(nested_list)):
    #             _recursive_set_index_pyobj_value(nested_list[i], indices[1:], rhs[i] if not is_scalar_rhs else rhs, is_scalar_rhs)
    #     else:
    #         raise ValueError(f'Invalid index: {indices[0]}.')


def expand_as_tvalue(value: TensorValue, other: TensorValue) -> TensorValue:
    return expand_tvalue(value, other.batch_variables, other.tensor.size()[other.batch_dims: other.batch_dims + len(other.batch_variables)])


def expand_tvalue(value: TensorValue, batch_variables: Iterable[str], batch_sizes: Iterable[int]) -> TensorValue:
    batch_variables = tuple(batch_variables)
    batch_sizes = tuple(batch_sizes)

    data = value.tensor
    masked_tensor_storage_kwargs = {
        'mask': value.tensor_mask,
        'optimistic_values': value.tensor_optimistic_values,
        'quantized_values': value.tensor_quantized_values,
    }

    current_batch_variables = list(value.batch_variables)
    for var in batch_variables:
        if var not in current_batch_variables:
            data = data.unsqueeze(value.batch_dims)
            for k, v in masked_tensor_storage_kwargs.items():
                if v is not None:
                    masked_tensor_storage_kwargs[k] = v.unsqueeze(value.batch_dims)
            current_batch_variables.insert(0, var)

    indices = list()
    sizes = list()

    # process the first "batch" dims.
    for i in range(value.batch_dims):
        indices.append(i)
        sizes.append(data.size(i))

    # process the next "variables" dims.
    for var, size in zip(batch_variables, batch_sizes):
        indices.append(current_batch_variables.index(var) + value.batch_dims)
        sizes.append(size)

    for k, v in masked_tensor_storage_kwargs.items():
        if v is not None:
            # corner case for "scalar" storage.
            if v.ndim > 0:
                v = v.permute(indices)
            masked_tensor_storage_kwargs[k] = v.expand(sizes)

    # process the last "data" dims.
    if isinstance(value.dtype, TensorValueTypeBase):
        for i in range(value.dtype.ndim()):
            indices.append(value.batch_dims + len(batch_variables))
            sizes.append(data.size(value.batch_dims + len(batch_variables) + i))

    # corner case for "scalar" storage.
    if len(indices) > 0:
        data = data.permute(indices)
    data = data.expand(sizes)

    rv = TensorValue(
        value.dtype, batch_variables,
        MaskedTensorStorage(data, **masked_tensor_storage_kwargs),
        batch_dims=value.batch_dims, _mask_certified_flag=False
    )

    return rv


def expand_pyobj_value(value: TensorValue, batch_variables: Iterable[str], batch_sizes: Iterable[int]) -> TensorValue:
    pass


def is_tvalue_simple_quantizable(value: TensorValue) -> bool:
    """A value is simple quantizable if it is either quantized, or it has tensor_quantized_values, or it has intrinsically quantized dtypes."""
    return (
        value.tensor_quantized_values is not None or
        value.dtype.is_intrinsically_quantized()
    )


def simple_quantize_tvalue(value: TensorValue) -> TensorValue:
    """Quantize a value into a quantized value using the simple rules.

    The simple rules are:
        1. If the value is already quantized, return the value itself.
        2. If the value has tensor_quantized_values, return the tensor_indices.
        3. If the value has intrinsically quantized dtypes, quantize the value.
            - For BOOL, quantize the value with (value > 0.5).to(torch.int64).
            - For INT64, quantize the value with value.to(torch.int64).

    Args:
        value: the value to be quantized.

    Returns:
        the quantized value.
    """
    assert value.dtype.is_intrinsically_quantized() or value.tensor_quantized_values is not None

    if value.tensor_quantized_values is not None:
        rv = value.tensor_quantized_values
    elif value.dtype == BOOL:
        rv = (value.tensor > 0.5).to(torch.int64)
    elif value.dtype == INT64:
        rv = value.tensor.to(torch.int64)
    elif isinstance(value.dtype, NamedTensorValueType) and value.dtype.base_type == BOOL:
        rv = (value.tensor > 0.5).to(torch.int64)
    elif isinstance(value.dtype, NamedTensorValueType) and value.dtype.base_type == INT64:
        rv = value.tensor.to(torch.int64)
    else:
        raise TypeError('Unable to quantize value. Need either tensor_indices, or intrinsically quantized dtypes: {}.'.format(str(value)))

    rv = MaskedTensorStorage(rv, value.tensor_mask)
    return type(value)(value.dtype, value.batch_variables, rv, value.batch_dims, _mask_certified_flag=value._mask_certified_flag)



concepts/dsl/tensor_value.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : value.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/18/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""The baseclass for all value representations of domain-specific languages."""

from typing import Any, Union

from concepts.dsl.dsl_types import TypeBase, ObjectType, ValueType
from concepts.dsl.dsl_types import get_format_context

__all__ = ['ValueBase', 'Value']


class ValueBase(object):
    """The baseclass for all value representations of domain-specific languages."""

    dtype: TypeBase
    """The type of the value."""

    def __init__(self, dtype: TypeBase):
        """Initialize the Value object.

        Args:
            dtype: the type of the value.
        """
        self.dtype = dtype


class Value(ValueBase):
    """A simple value object holds a pair of dtype and a value."""

    dtype: Union[ObjectType, ValueType]
    """The type of the value."""

    value: Any
    """The value."""

    def __init__(self, dtype: Union[ObjectType, ValueType], value: Any):
        """Initialize the Value object.

        Args:
            dtype: the type of the value.
            value: the value.
        """
        super().__init__(dtype)
        self.value = value
        self._check_type()

    def _check_type(self):
        """Check if the value is of the correct type. Default to no-op."""
        pass

    def __str__(self):
        if get_format_context().object_format_type:
            return f'V({self.value}, dtype={self.dtype})'
        else:
            return self.value

    def __repr__(self):
        return f'V({self.value}, dtype={self.dtype})'

concepts/dsl/value.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/12/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Interfaces for hardware devices such as robots and cameras."""

concepts/hw_interface/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/11/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/hw_interface/franka/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : server.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/11/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
from typing import Tuple
from jacinle.logging import get_logger
from jacinle.comm.service import Service, SocketClient

logger = get_logger(__file__)

__all__ = ['FrankaService', 'FrankaClient']


class FrankaService(Service):
    DEFAULT_PORTS = [12345, 12346]

    arm: 'franka_interface.ArmInterface'
    camera: 'RealSenseDevice'
    camera_extrinsics: np.ndarray

    def initialize(self):
        self.arm = None
        self.camera = None
        self.camera_extrinsics = None

        if self.configs.get('debug', False):
            logger.critical('Running in debug mode.')
            return

        try:
            import rospy
            import franka_interface
            rospy.init_node('franka_interface')
            self.arm = franka_interface.ArmInterface()
        except ImportError:
            logger.warning('Failed to import franka_interface. Please install it first.')

        if self.configs.get('camera', None) == 'D435':
            from concepts.hw_interface.realsense.device import RealSenseDevice
            self.camera = RealSenseDevice.find_devices('D435')[0]
            self.camera.start_pipeline()

        if self.configs.get('camera_extrinsics', None) is not None:
            with open(self.configs.get('camera_extrinsics')) as f:
                self.camera_extrinsics = np.asarray(list(map(float, f.read().split())))

    def finalize(self):
        if self.camera is not None:
            self.camera.stop_pipeline()

    @property
    def gripper(self):
        if self.arm is None:
            raise RuntimeError('FrankaInterface is not initialized.')
        return self.arm.hand

    def call(self, command, *args, **kwargs):
        if self.arm is None:
            logger.critical(f'[DEBUG] Received command {command} with args {args}.')
            return

        function = getattr(self, command)
        logger.critical(f'Received command {command}.')
        return function(*args, **kwargs)

    def move_qpos(self, qpos: np.ndarray, timeout: float = 10) -> None:
        """Move the robot to a given joint position."""
        qpos_dict = self.arm.convertToDict(qpos)
        self.arm.move_to_joint_positions(qpos_dict, timeout)

    def move_home(self, timeout: float = 10) -> None:
        """Move the robot to the home position."""
        self.arm.move_to_neutral(timeout)

    def move_qpos_to_touch(self, qpos: np.ndarray, timeout: float = 10) -> None:
        """Move the robot to a touch position."""
        qpos_dict = self.arm.convertToDict(qpos)
        self.arm.move_to_touch(qpos_dict, timeout)

    def move_qpos_from_touch(self, qpos: np.ndarray, timeout: float = 10) -> None:
        """Move the robot from a touch position to a given joint position."""
        qpos_dict = self.arm.convertToDict(qpos)
        self.arm.move_from_touch(qpos_dict, timeout)

    def move_qpos_trajectory(self, qpos_list, timeout: float = 10) -> None:
        qpos_dict = [self.arm.convertToDict(qpos) for qpos in qpos_list]
        self.arm.execute_position_path(qpos_dict, timeout)

    def reset_errors(self) -> None:
        """Reset the error state of the robot."""
        self.arm.resetErrors()

    def open_gripper(self) -> None:
        """Open the gripper."""
        self.gripper.open()

    def close_gripper(self) -> None:
        """Close the gripper."""
        self.gripper.close()

    def grasp_gripper(self, width: float, force: float = 40):
        """Grasp an object with the gripper."""
        self.gripper.grasp(width, force)

    def get_qpos(self) -> np.ndarray:
        """Return the current joint position."""
        return np.array(self.arm.convertToList(self.arm.joint_angles()), dtype=np.float32)

    def get_ee_pose(self) -> Tuple[np.ndarray, np.ndarray]:
        """Return the end-effector pose."""
        rv = self.arm.endpoint_pose()
        o = rv['orientation']
        return np.array(rv['position'], dtype=np.float32), np.array([o.x, o.y, o.z, o.w], dtype=np.float32)

    def set_cart_impedance_freemotion(self):
        """Set the robot to Cartesian impedance mode with free motion."""
        self.arm.set_cart_impedance_pose(
            self.arm.endpoint_pose(),
            stiffness=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        )

    def capture_image(self) -> Tuple[np.ndarray, np.ndarray]:
        """Return the color and depth image."""
        if self.camera is None:
            raise RuntimeError('Camera is not initialized.')
        return self.camera.capture_images()

    def get_camera_intrinsics(self) -> Tuple[np.ndarray, np.ndarray]:
        """Return the color and depth intrinsics matrices."""
        if self.camera is None:
            raise RuntimeError('Camera is not initialized.')
        return self.camera.color_intrinsics, self.camera.depth_intrinsics

    def get_camera_extrinsics(self) -> np.ndarray:
        """Return the color and depth intrinsics matrices."""
        if self.camera is None:
            raise RuntimeError('Camera is not initialized.')
        return self.camera_extrinsics


class FrankaClient(SocketClient):
    def __init__(self, server: str, name='franka-client', port_pair=None, automatic_reset_errors: bool = True):
        if port_pair is None:
            port_pair = FrankaService.DEFAULT_PORTS

        connection = (f'tcp://{server}:{port_pair[0]}', f'tcp://{server}:{port_pair[1]}')
        super().__init__(name, connection)

        self.automatic_reset_errors = automatic_reset_errors

    def move_qpos(self, qpos: np.ndarray, timeout: float = 10) -> None:
        self.call('move_qpos', qpos, timeout)
        if self.automatic_reset_errors:
            self.call('reset_errors')

    def move_home(self, timeout: float = 10) -> None:
        self.call('move_home', timeout)
        if self.automatic_reset_errors:
            self.call('reset_errors')

    def move_qpos_to_touch(self, qpos: np.ndarray, timeout: float = 10) -> None:
        self.call('move_qpos_to_touch', qpos, timeout)
        if self.automatic_reset_errors:
            self.call('reset_errors')

    def move_qpos_from_touch(self, qpos: np.ndarray, timeout: float = 10) -> None:
        self.call('move_qpos_from_touch', qpos, timeout)
        if self.automatic_reset_errors:
            self.call('reset_errors')

    def move_qpos_trajectory(self, qpos_list, timeout: float = 10) -> None:
        current_qpos = self.get_qpos()
        qpos_list = [current_qpos] + qpos_list
        self.call('move_qpos_trajectory', qpos_list, timeout)
        if self.automatic_reset_errors:
            self.call('reset_errors')

    def reset_errors(self) -> None:
        return self.call('reset_errors')

    def open_gripper(self) -> None:
        self.call('open_gripper')
        self.call('reset_errors')

    def close_gripper(self) -> None:
        self.call('close_gripper')
        self.call('reset_errors')

    def grasp_gripper(self, width: float = 0.05, force: float = 40):
        self.call('grasp_gripper', width, force)
        self.call('reset_errors')

    def get_qpos(self) -> np.ndarray:
        return self.call('get_qpos')

    def get_ee_pose(self) -> Tuple[np.ndarray, np.ndarray]:
        return self.call('get_ee_pose')

    def set_cart_impedance_freemotion(self):
        return self.call('set_cart_impedance_freemotion')

    def capture_image(self) -> Tuple[np.ndarray, np.ndarray]:
        return self.call('capture_image')

    def get_camera_intrinsics(self) -> Tuple[np.ndarray, np.ndarray]:
        return self.call('get_camera_intrinsics')

    def get_camera_extrinsics(self) -> Tuple[np.ndarray, np.ndarray]:
        return self.call('get_camera_extrinsics')


concepts/hw_interface/franka/server.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/05/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/hw_interface/george_vision/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : object_centric_vision.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/05/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import os.path as osp
from typing import Optional, Union, Tuple, List, Dict
from functools import cached_property
from dataclasses import dataclass

import jacinle
import jacinle.io as io
import open3d as o3d
import numpy as np
import cv2
import matplotlib.pyplot as plt
from concepts.utils.rotationlib import axisangle2quat, quat2mat

from concepts.hw_interface.george_vision.segmentation_models import random_colored_mask, visualize_instance_segmentation, ImageBasedPCDSegmentationModel, PointGuidedImageSegmentationModel, InstanceSegmentationResult
from concepts.simulator.pybullet.client import BulletClient
from concepts.simulator.pybullet.urdf_utils.obj2urdf import ObjectUrdfBuilder


__all__ = [
    'CameraTransformation', 'ObjectDetectionResult', 'ObjectCentricVisionPipeline',
    'get_pointcloud', 'filter_pointcloud_range', 'ransac_table_detection', 'threshold_table_detection',
    'project_pointcloud_to_plane', 'object_reconstruction_by_extrusion', 'remove_outliers', 'mesh_reconstruction_alpha_shape',
    'canonize_mesh_center_', 'compute_transformation_from_plane_equation', 'visualize_with_camera_matrix',
    'load_scene_in_pybullet',
]


class CameraTransformation(object):
    """Camera transformation class."""

    def __init__(self, *, intrinsics=None, extrinsics=None):
        if intrinsics is None:
            intrinsics = type(self).default_intrinsics()

        self.intrinsics = intrinsics
        self.extrinsics = extrinsics

    intrinsics: np.ndarray
    """Intrinsic matrix of the camera."""

    extrinsics: Optional[np.ndarray]
    """Extrinsic matrix of the camera."""

    @classmethod
    def from_extrinsics_file(cls, filename):
        """Load the camera transformation from an extrinsics file."""
        return cls(intrinsics=None, extrinsics=io.load(filename))

    @classmethod
    def from_intrinsics_and_extrinsics_file(cls, intrinsics_file, extrinsics_file):
        """Load the camera transformation from an intrinsics file and an extrinsics file."""
        return cls(intrinsics=io.load(intrinsics_file), extrinsics=io.load(extrinsics_file))

    @classmethod
    def default_intrinsics(cls):
        """Returns the default intrinsic matrix."""
        intrinsics = np.array([
            [1.36275940e+03, 0.00000000e+00, 9.39736938e+02],
            [0.00000000e+00, 1.36269678e+03, 5.67980225e+02],
            [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]
        ])
        return intrinsics

    def w2c(self, x: Union[float, np.ndarray], y: Union[float, np.ndarray], z: Union[float, np.ndarray]):
        """Converts a point in the world coordinate system to the camera coordinate system. The inputs can be either scalars or arrays.

        Args:
            x: x coordinate(s) in the world coordinate system.
            y: y coordinate(s) in the world coordinate system.
            z: z coordinate(s) in the world coordinate system.

        Returns:
            a tuple of (u, v, z) in the camera coordinate system.
        """
        if isinstance(x, float) and isinstance(y, float) and isinstance(z, float):
            x, y, z, _ = np.dot(self.extrinsics, np.array([x, y, z, 1]))
            u, v, z = np.dot(self.intrinsics, np.array([x, y, z]))
            u, v = int(u / z), int(v / z)
            return u, v
        else:
            x, y, z = map(lambda x: np.asarray(x).reshape(-1), [x, y, z])
            xyz = np.stack([x, y, z, np.ones_like(x)], axis=-1)
            xyz = np.dot(self.extrinsics, xyz.transpose(1, 0)).transpose(1, 0)[:, :3]
            uvz = np.dot(self.intrinsics, xyz.transpose(1, 0)).transpose(1, 0)
            u, v, z = uvz[:, 0] / uvz[:, 2], uvz[:, 1] / uvz[:, 2], uvz[:, 2]
            return u.astype(np.int32), v.astype(np.int32), z

    def c2w(self, u: Union[float, np.ndarray], v: Union[float, np.ndarray], d: Union[float, np.ndarray]):
        """Converts a point in the camera coordinate system to the world coordinate system. The inputs can be either scalars or arrays.

        Args:
            u: u coordinate(s) in the camera coordinate system.
            v: v coordinate(s) in the camera coordinate system.
            d: depth(s) in the camera coordinate system.

        Returns:
            a tuple of (x, y, z) in the world coordinate system.
        """
        if isinstance(u, float) and isinstance(v, float) and isinstance(d, float):
            x, y, z = np.dot(np.linalg.inv(self.intrinsics), np.array([u * d, v * d, d]))
            x, y, z, _ = np.dot(np.linalg.inv(self.extrinsics), np.array([x, y, z, 1]))
            return x, y, z
        else:
            u, v, d = map(lambda x: np.asarray(x).reshape(-1), [u, v, d])
            xyz = np.stack([u * d, v * d, d], axis=-1)
            xyz = np.dot(np.linalg.inv(self.intrinsics), xyz.transpose(1, 0)).transpose(1, 0)
            xyz = np.stack([xyz[:, 0], xyz[:, 1], xyz[:, 2], np.ones_like(xyz[:, 0])], axis=-1)
            xyz = np.dot(np.linalg.inv(self.extrinsics), xyz.transpose(1, 0)).transpose(1, 0)
            x, y, z = xyz[:, 0], xyz[:, 1], xyz[:, 2]
            return x, y, z

    @property
    def camera_pose(self) -> Tuple[np.ndarray, np.ndarray]:
        """Returns the camera pose in the world coordinate system (i.e. the extrinsics matrix in the form of (translation, rotation)).

        Returns:
            A tuple of (translation, rotation). Rotation is represented by a matrix.
        """
        rotation = self.extrinsics[:3, :3]
        translation = self.extrinsics[:3, 3]
        return translation, rotation


@dataclass
class ObjectDetectionResult(object):
    """The result of object detection."""

    label: str
    """The label of the object."""

    pcd: o3d.geometry.PointCloud
    """The point cloud of the object."""

    reconstructed_pcd: Optional[o3d.geometry.PointCloud] = None
    """The reconstructed point cloud of the object."""

    reconstructed_mesh: Optional[o3d.geometry.TriangleMesh] = None
    """The mesh of the object."""

    pcd_mask: Optional[np.ndarray] = None

    @property
    def mesh(self) -> o3d.geometry.TriangleMesh:
        """Returns the mesh of the object. If the mesh is not reconstructed, an error will be raised."""

        if self.reconstructed_mesh is None:
            raise ValueError('The mesh is not reconstructed yet. Please call construct_object_meshes() first.')
        return self.reconstructed_mesh

    @property
    def mesh_bbox(self) -> o3d.geometry.AxisAlignedBoundingBox:
        """Returns the bounding box of the mesh of the object. If the mesh is not reconstructed, an error will be raised."""

        if self.reconstructed_mesh is None:
            raise ValueError('The mesh is not reconstructed yet. Please call construct_object_meshes() first.')
        return self.reconstructed_mesh.get_axis_aligned_bounding_box()


class ObjectCentricVisionPipeline(object):
    """The pipeline for object-centric vision."""

    def __init__(self, color_image: np.ndarray, depth_image: np.ndarray):
        """Initializes the pipeline.

        Args:
            color_image: The color image.
            depth_image: The depth image.
        """
        self.color_image = color_image
        self.depth_image = depth_image
        self.camera_transform = None
        self.foreground_ranges = None
        self._pcd = None
        self._foreground_segmentation = None
        self._table_model = None
        self._table_segmentation = None
        self._image_detected_objects = None
        self._detected_objects = None

    @property
    def pcd(self) -> o3d.geometry.PointCloud:
        """The point cloud of the scene."""
        if self._pcd is None:
            raise ValueError('Point cloud is not constructed yet. Please call construct_pointcloud() first.')
        return self._pcd

    @property
    def foreground_segmentation(self) -> np.ndarray:
        """The segmentation of the foreground, in the form of a boolean array of the same length as the point cloud."""
        if self._foreground_segmentation is None:
            raise ValueError('Background segmentation is not constructed yet. Please call construct_foreground_by_cropping() first.')
        return self._foreground_segmentation

    @cached_property
    def foreground_pcd(self) -> o3d.geometry.PointCloud:
        """The point cloud of the foreground."""
        return self.pcd.select_by_index(np.where(self.foreground_segmentation)[0])

    @property
    def table_model(self) -> np.ndarray:
        """The model of the table, in the form of a length-4 array representing the plane equation ax + by + cz + d = 0."""
        if self._table_model is None:
            raise ValueError('Table model is not constructed yet. Please call construct_table_segmentation() first.')
        return self._table_model

    @property
    def table_segmentation(self) -> np.ndarray:
        """The segmentation of the table, in the form of a boolean array of the same length as the point cloud."""
        if self._table_segmentation is None:
            raise ValueError('Table segmentation is not constructed yet. Please call construct_table_segmentation() first.')
        return self._table_segmentation

    @cached_property
    def tabletop_segmentation(self) -> np.ndarray:
        """The segmentation of the tabletop, in the form of a boolean array of the same length as the point cloud. Tabletop is defined as the foreground points that are not table points."""
        return ~self.table_segmentation & self.foreground_segmentation

    @cached_property
    def table_pcd(self) -> o3d.geometry.PointCloud:
        """The point cloud of the table."""
        return self.pcd.select_by_index(np.where(self.table_segmentation)[0])

    @cached_property
    def tabletop_pcd(self) -> o3d.geometry.PointCloud:
        """The point cloud of the tabletop. Tabletop is defined as the foreground points that are not table points."""
        return self.pcd.select_by_index(np.where(self.tabletop_segmentation)[0])

    @property
    def image_detected_objects(self) -> InstanceSegmentationResult:
        """The result of object detection in the image."""
        if self._image_detected_objects is None:
            raise ValueError('Object detection is not performed yet. Please call construct_object_segmentation() first.')
        return self._image_detected_objects

    @property
    def detected_objects(self) -> List[ObjectDetectionResult]:
        """The result of object detection in the scene."""
        if self._detected_objects is None:
            raise ValueError('Object detection is not performed yet. Please call construct_object_segmentation() first.')
        return self._detected_objects

    def construct_pointcloud(self, camera_transform: Optional[CameraTransformation] = None) -> o3d.geometry.PointCloud:
        """Constructs the point cloud from the color and depth images.

        Args:
            camera_transform: the camera transformation. If None, reconstruct the point cloud in the camera coordinate system.

        Returns:
            the point cloud.
        """
        self.camera_transform = camera_transform
        self._pcd = get_pointcloud(self.color_image, self.depth_image, camera_transform)
        return self._pcd

    def construct_foreground_by_cropping(
        self,
        x_range: Optional[Tuple[float, float]] = None,
        y_range: Optional[Tuple[float, float]] = None,
        z_range: Optional[Tuple[float, float]] = None,
        depth_range: Optional[Tuple[float, float]] = None
    ) -> np.ndarray:
        """Constructs the foreground segmentation by cropping the point cloud.

        Args:
            x_range: the range of x-axis in the world coordinate system.
            y_range: the range of y-axis in the world coordinate system.
            z_range: the range of z-axis in the world coordinate system.
            depth_range: the range of depth in the camera coordinate system.

        Returns:
            the foreground segmentation as a boolean array of shape (N,). N is the number of points in the point cloud.
        """
        depth_mask = np.ones(len(self.pcd.points), dtype=np.bool_)
        if depth_range is not None:
            depth_mask = depth_mask & (self.depth_image.flatten() > depth_range[0] * 1000) & (self.depth_image.flatten() < depth_range[1] * 1000)

        self.foreground_ranges = (x_range, y_range, z_range)
        self._foreground_segmentation = filter_pointcloud_range(self.pcd, x_range, y_range, z_range, return_condition=True) & depth_mask

        return self._foreground_segmentation

    def construct_table_segmentation(self, method: str = 'ransac', **kwargs):
        """Constructs the table segmentation.

        Args:
            method: the method to use for table segmentation. Available methods are {'ransac', 'threshold'}.

        Returns:
            the table segmentation as a boolean array of shape (N,). N is the number of points in the point cloud.
        """
        pcd = self.foreground_pcd
        if method == 'ransac':
            plane_model, raw_table_segmentation = ransac_table_detection(pcd, **kwargs, return_condition=True)
        elif method == 'threshold':
            plane_model, raw_table_segmentation = threshold_table_detection(pcd, **kwargs, return_condition=True)
        else:
            raise ValueError(f'Unknown method {method}. Available methods are {{ransac, threshold}}.')

        table_segmentation = np.zeros(len(self.pcd.points), dtype=np.bool_)
        table_segmentation[self.foreground_segmentation] = raw_table_segmentation
        self._table_model = plane_model
        self._table_segmentation = table_segmentation

        return self._table_segmentation

    def construct_object_segmentation(
        self, image_segmentation_model: Optional[ImageBasedPCDSegmentationModel] = None, detection_min_points: int = 50, enforce_detected_points: bool = True,
        dbscan: bool = True, dbscan_eps: float = 0.01, dbscan_min_samples: int = 500, dbscan_min_points: int = 1500, dbscan_with_sam: bool = True, sam_model: Optional[PointGuidedImageSegmentationModel] = None, sam_max_points: int = 100000,
        outlier_filter: bool = True, outlier_filter_radius: float = 0.005, outlier_filter_min_neighbors: int = 50,
        min_height: float = 0.015, verbose: bool = False
    ):
        """Constructs the object segmentation. The segmentation is performed in the following steps:

        1. If image_segmentation_model is not None, use the image segmentation model to segment the image.
        2. If dbscan is True, use DBSCAN to cluster the remaining points from the previous step.
        3. If outlier_filter is True, use radius outlier filter to remove outliers for each object obtained from the previous two steps.
        4. Use height filter to remove objects that are too thin.

        Args:
            image_segmentation_model: the image segmentation model. If None, skip the image segmentation step.
            detection_min_points: the minimum number of points for an object to be considered as a valid object.
            enforce_detected_points: if True, keep all points that are detected by the image segmentation model.
                Otherwise, only keep the points that are detected by the image segmentation model and are also in the "tabletop" segmentation, which equals foreground setminus table.
            dbscan: if True, use DBSCAN to cluster the remaining points from the previous step.
            dbscan_eps: the maximum distance between two samples for one to be considered as in the neighborhood of the other.
            dbscan_min_samples: the number of samples in a neighborhood for a point to be considered as a core point.
            dbscan_min_points: the minimum number of points for an object to be considered as a valid object.
            dbscan_with_sam: if True, use SAM to filter the points before DBSCAN.
            sam_model: the SAM model. If None, skip the SAM step.
            sam_max_points: the maximum number of points to use for SAM. This option is useful to avoid the table being detected as an object.
            outlier_filter: if True, use radius outlier filter to remove outliers for each object obtained from the previous two steps.
            outlier_filter_radius: the radius of the sphere that will determine which points are neighbors.
            outlier_filter_min_neighbors: the minimum number of neighbors that a point must have to be considered as an inlier.
            min_height: the minimum height of an object.
            verbose: if True, print out debug information.
        """
        remain_condition = self.tabletop_segmentation.copy()
        detected_objects = list()

        if image_segmentation_model is not None:
            results = image_segmentation_model.segment_image(self.color_image)
            self._image_detected_objects = results

            for i in range(results.nr_objects):
                label = results.pred_cls[i]

                if 'table' in label:
                    continue

                mask = results.masks[i].flatten()
                if enforce_detected_points:
                    mask = mask & self.foreground_segmentation
                else:
                    mask = mask & self.tabletop_segmentation

                if verbose:
                    print('Detected object with label', results.pred_cls[i], 'and mask size', mask.sum(), 'points.')

                if mask.sum() < detection_min_points or mask.sum() > sam_max_points:
                    continue
                pcd = self.pcd.select_by_index(np.where(mask)[0])
                detected_objects.append(ObjectDetectionResult(label, pcd, pcd_mask=mask))
                remain_condition = remain_condition & (~mask)

        # Run DBSCAN on the remaining points
        if dbscan:
            pcd = self.pcd.select_by_index(np.where(remain_condition)[0])
            if verbose:
                print('Running DBSCAN on remaining points with size', len(pcd.points))
            labels = np.asarray(pcd.cluster_dbscan(eps=dbscan_eps, min_points=dbscan_min_samples, print_progress=verbose))
            unique_labels = np.unique(labels)
            for label in unique_labels:
                if label == -1:
                    continue

                mask = labels == label
                label_name = f'unknown_{label}'
                global_mask = np.zeros(len(self.pcd.points), dtype=np.bool_)
                global_mask[remain_condition] = mask
                if global_mask.sum() < dbscan_min_points:
                    continue

                if dbscan_with_sam and sam_model is not None:
                    y, x = np.where(global_mask.reshape(self.color_image.shape[:2]))
                    y_mean, x_mean = int(y.mean()), int(x.mean())
                    if verbose:
                        print('  Running SAM on DBSCAN cluster with size', len(y), 'points.')
                    mask = sam_model.segment_from_point(self.color_image, (x_mean, y_mean))
                    if verbose:
                        print('  SAM mask size:', mask.sum())

                    if mask.sum() > sam_max_points:
                        if verbose:
                            print('  SAM mask too large, skipping.')
                        continue

                    mask = mask.flatten()
                    if enforce_detected_points:
                        mask = mask & self.foreground_segmentation
                    else:
                        mask = mask & self.tabletop_segmentation
                    global_mask = mask

                pcd = self.pcd.select_by_index(np.where(global_mask)[0])
                # pcd.paint_uniform_color(np.random.rand(3))
                detected_objects.append(ObjectDetectionResult(label_name, pcd, pcd_mask=global_mask))
                if verbose:
                    print('Detected object with label', label_name, 'and mask size', mask.sum(), 'points.')

        if verbose:
            print('Detected', len(detected_objects), 'objects before filtering.')

        # Remove outliers and thin objects
        filtered_objects = list()
        for obj in detected_objects:
            if outlier_filter:
                obj.pcd = remove_outliers(obj.pcd, radius=outlier_filter_radius, min_neighbors=outlier_filter_min_neighbors)
            if not dbscan_with_sam:
                min_bounds = obj.pcd.get_min_bound()
                max_bounds = obj.pcd.get_max_bound()

                if max_bounds[2] - min_bounds[2] < min_height:
                    if verbose:
                        print('Object', obj.label, 'is too thin, skipping.')
                    continue

            filtered_objects.append(obj)

        if verbose:
            print('Detected', len(filtered_objects), 'objects.')
        self._detected_objects = filtered_objects
        return self._detected_objects

    def construct_object_completion_by_extrusion(self, min_height: Optional[float] = 0.05, verbose: bool = False):
        """Construct object completion by extrusion the top surface of the table.

        Args:
            min_height: the minimum height of an object. If specified as None, skip the thin object filtering step.
                Otherwise, if an object is too thin, it will be extruded to the minimum height.
            verbose: if True, print out debug information.
        """
        for obj in self.detected_objects:
            if verbose:
                print('Shape reconstructing object', obj.label, 'with', len(obj.pcd.points), 'points.')
            obj.reconstructed_pcd = object_reconstruction_by_extrusion(obj.pcd, self.table_model, min_height=min_height)

    def construct_object_meshes(self, alpha: float = 0.1, verbose: bool = False):
        """Construct meshes for all detected objects using alpha shape.

        Args:
            alpha: the alpha value for alpha shape.
            verbose: if True, print out debug information.
        """
        for obj in self.detected_objects:
            if verbose:
                print('Mesh reconstruction for object', obj.label, 'with', len(obj.pcd.points), 'points.')
            obj.reconstructed_mesh = mesh_reconstruction_alpha_shape(obj.reconstructed_pcd, alpha=alpha)

    def export_scene(self, output_dir: str, verbose: bool = False):
        """Export the scene to a directory. All the metadata will be saved as "metainfo.json".

        Args:
            output_dir: the output directory.
        """
        io.mkdir(output_dir)
        metainfo = {
            'table_model': self.table_model.tolist(),
            'table_transformation': compute_transformation_from_plane_equation(*self.table_model).tolist(),
            'objects': list()
        }
        cv2.imwrite(osp.join(output_dir, 'color.png'), self.color_image[:, :, ::-1])
        cv2.imwrite(osp.join(output_dir, 'depth.png'), self.depth_image)
        for idx, detection in jacinle.tqdm_gofor(self.detected_objects, desc='Exporting objects'):
            det_metainfo = {
                'index': idx,
                'label': detection.label,
                'unique_label': f'o{idx:02d}_{detection.label}',
                'pcd': f'{idx}_pcd.ply',
                'mesh': f'{idx}_mesh.ply',
                'mesh_obj': f'{idx}_mesh.obj',
                'urdf': f'{idx}_mesh.obj.urdf',
            }
            metainfo['objects'].append(det_metainfo)
            o3d.io.write_point_cloud(osp.join(output_dir, det_metainfo['pcd']), detection.pcd)
            mesh = detection.reconstructed_mesh
            mesh, center = canonize_mesh_center_(mesh)
            det_metainfo['pos'] = center.tolist()
            o3d.io.write_triangle_mesh(osp.join(output_dir, det_metainfo['mesh']), mesh)
            o3d.io.write_triangle_mesh(osp.join(output_dir, det_metainfo['mesh_obj']), mesh, write_triangle_uvs=True)

            with jacinle.cond_with(jacinle.suppress_output(), not verbose):  # if verbose, print out the output of urdf builder
                builder = ObjectUrdfBuilder(output_dir)
                builder.build_urdf(osp.join(output_dir, det_metainfo['mesh_obj']), force_overwrite=True, decompose_concave=True, force_decompose=False, center=None)

        io.dump_json(osp.join(output_dir, 'metainfo.json'), metainfo)

    def visualize_pcd(self, name: str):
        """Visualize a point cloud. It will also visualize two coordinate frames: the camera frame and the world frame.

        Args:
            name: the name of the point cloud to visualize. Can be one of the following:
                - 'full': the full point cloud.
                - 'foreground': the foreground point cloud.
                - 'table': the table point cloud.
                - 'tabletop': the tabletop point cloud.
        """
        geometries = list()
        if name == 'full':
            geometries.append(self.pcd)
        elif name == 'foreground':
            geometries.append(self.foreground_pcd)
            # visualize the bounding box
            bbox = o3d.geometry.AxisAlignedBoundingBox(
                min_bound=(self.foreground_ranges[0][0], self.foreground_ranges[1][0], self.foreground_ranges[2][0]),
                max_bound=(self.foreground_ranges[0][1], self.foreground_ranges[1][1], self.foreground_ranges[2][1])
            )
            bbox.color = (0, 1, 0)
            geometries.append(bbox)
        elif name == 'table':
            pcd = self.table_pcd
            # pcd.paint_uniform_color([1, 0, 0])
            geometries.append(pcd)
        elif name == 'tabletop':
            pcd = self.tabletop_pcd
            geometries.append(pcd)
        else:
            raise ValueError(f'Unknown name {name}. Available names are {{full, foreground, table, tabletop}}.')

        world_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.05, origin=[0, 0, 0])
        geometries.append(world_frame)
        if self.camera_transform is not None:
            camera = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.05).transform(np.linalg.inv(self.camera_transform.extrinsics))
            geometries.append(camera)

        # o3d.visualization.draw_geometries(geometries)
        visualize_with_camera_matrix(geometries, np.linalg.inv(self.camera_transform.extrinsics))

    def visualize_objects_raw_detection(self):
        """Visualize the raw detection results in 2D."""
        visualize_instance_segmentation(self.color_image, self.image_detected_objects)

    def visualize_objects_2d(self, rect_th: int = 3, text_size: int = 1, text_th: int = 3):
        """Visualize the detected objects back-projected to the color image.

        Args:
            rect_th: the thickness of the rectangle.
            text_size: the size of the text.
            text_th: the thickness of the text.
        """
        image = self.color_image.copy()
        for detection in self.detected_objects:
            label = detection.label
            masked_image = detection.pcd_mask.reshape(self.color_image.shape[:2])
            v, u = np.where(masked_image)
            rgb_mask = random_colored_mask(masked_image)
            image = cv2.addWeighted(image, 1, rgb_mask, 0.5, 0)
            cv2.rectangle(image, (u.min(), v.min()), (u.max(), v.max()), color=(0, 255, 0), thickness=rect_th)
            cv2.putText(image, label, (u.min(), v.min() + 10), cv2.FONT_HERSHEY_SIMPLEX, text_size, (0, 255, 0), thickness=text_th)

        plt.figure(figsize=(10, 10))
        plt.imshow(image)
        plt.tight_layout()
        plt.show()

    def visualize_objects_3d(self, name: str = 'mesh', overlay_background_pcd: bool = False):
        """Visualize the detected objects in 3D.

        Args:
            name: the name of the geometry to visualize. Can be one of the following:
                - 'mesh': the reconstructed mesh.
                - 'pcd': the reconstructed point cloud.
                - 'raw_pcd': the raw point cloud.
            overlay_background_pcd: whether to overlay the background point cloud.
        """
        print('Visualizing point cloud.')
        geometries = list()
        for obj in self.detected_objects:
            if name == 'mesh' and obj.reconstructed_mesh is not None:
                geometries.append(obj.reconstructed_mesh)
            elif name == 'pcd' and obj.reconstructed_pcd is not None:
                geometries.append(obj.reconstructed_pcd)
            elif name == 'raw_pcd':
                geometries.append(obj.pcd)

            # also add the bounding box
            bbox = geometries[-1].get_axis_aligned_bounding_box()
            bbox.color = (0, 1, 0)
            geometries.append(bbox)

        world_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.05, origin=[0, 0, 0])
        geometries.append(world_frame)
        if self.camera_transform is not None:
            camera = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.05).transform(np.linalg.inv(self.camera_transform.extrinsics))
            geometries.append(camera)

        # visualize a table plane
        if self.table_model is not None:
            # table_model was estimated by RANSAC. So it's represented as a plane equation.
            # We need to convert it to a plane mesh.
            table_plane = o3d.geometry.TriangleMesh.create_box(width=1, height=1, depth=0.001).translate(np.array([0, -0.5, 0]))
            table_plane.paint_uniform_color([0.5, 0.5, 0.5])
            T = compute_transformation_from_plane_equation(*self.table_model)
            table_plane = table_plane.transform(T)
            geometries.append(table_plane)

        if overlay_background_pcd:
            geometries.append(self.foreground_pcd)

        # o3d.visualization.draw_geometries(geometries)
        visualize_with_camera_matrix(geometries, np.linalg.inv(self.camera_transform.extrinsics))

    def run_default_pipeline(
        self, camera_transformation,
        x_range: Tuple[float, float] = (0.2, 1.0), y_range: Tuple[float, float] = (-0.5, 0.5), z_range: Tuple[float, float] = (-0.1, 0.5), depth_range: Tuple[float, float] = (0.05, 2.0),
        table_ransac_distance_threshold: float = 0.01, table_ransac_n: int = 3, table_ransac_num_iterations: int = 1000,
        image_segmentation_model: Optional[ImageBasedPCDSegmentationModel] = None, detection_min_points: int = 50, enforce_detected_points: bool = True,
        dbscan: bool = True, dbscan_eps: float = 0.01, dbscan_min_samples: int = 500, dbscan_min_points: int = 1500,
        dbscan_with_sam: bool = True, sam_model: Optional[PointGuidedImageSegmentationModel] = None, sam_max_points: int = 50000,
        outlier_filter: bool = True, outlier_filter_radius: float = 0.005, outlier_filter_min_neighbors: int = 50,
        min_height: float = 0.015,
        extrusion_min_height: float = 0.03,
        mesh_alpha: float = 0.1
    ):
        """Running the entire pipeline with the default parameters."""
        pbar = jacinle.tqdm_pbar(total=7, desc='Running object-centric vision pipeline', leave=False)
        with pbar:
            pbar.set_description('Reconstructing point cloud')
            self.construct_pointcloud(camera_transformation)
            pbar.update(); pbar.set_description('Segmenting the foreground')
            self.construct_foreground_by_cropping(x_range, y_range, z_range, depth_range)
            pbar.update(); pbar.set_description('Estimating the table plane')
            self.construct_table_segmentation(method='ransac', distance_threshold=table_ransac_distance_threshold, ransac_n=table_ransac_n, num_iterations=table_ransac_num_iterations)
            pbar.update(); pbar.set_description('Segmenting the objects')
            self.construct_object_segmentation(
                image_segmentation_model, detection_min_points=detection_min_points, enforce_detected_points=enforce_detected_points,
                dbscan=dbscan, dbscan_eps=dbscan_eps, dbscan_min_samples=dbscan_min_samples, dbscan_min_points=dbscan_min_points,
                dbscan_with_sam=dbscan_with_sam, sam_model=sam_model, sam_max_points=sam_max_points,
                outlier_filter=outlier_filter, outlier_filter_radius=outlier_filter_radius, outlier_filter_min_neighbors=outlier_filter_min_neighbors,
                min_height=min_height
            )
            pbar.update(); pbar.set_description('3D-reconstructing the objects')
            self.construct_object_completion_by_extrusion(min_height=extrusion_min_height)
            pbar.update(); pbar.set_description('Mesh-reconstructing the objects')
            self.construct_object_meshes(alpha=mesh_alpha)
            pbar.update(); pbar.set_description('Running NMS')
            self._detected_objects = mesh_nms(self.detected_objects, 0.1)


def get_pointcloud(color_image: np.ndarray, depth_image: np.ndarray, camera_transformation: Optional[CameraTransformation] = None) -> o3d.geometry.PointCloud:
    """Reconstruct a colored point cloud from a color and depth image.

    Args:
        color_image: a numpy array of shape (H, W, 3) representing the color image.
        depth_image: a numpy array of shape (H, W) representing the depth image.
        camera_transformation: a CameraTransformation object.

    Returns:
        a colored open3d point cloud.
    """

    xs = np.arange(0, color_image.shape[1])
    ys = np.arange(0, color_image.shape[0])
    xx, yy = np.meshgrid(xs, ys)
    xx = xx.flatten()
    yy = yy.flatten()
    zz = depth_image.flatten() / 1000.0
    colors = color_image[yy, xx, :].astype(np.float64) / 255.0

    if camera_transformation is not None:
        xx, yy, zz = camera_transformation.c2w(xx, yy, zz)

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(np.stack([xx, yy, zz], axis=1))
    pcd.colors = o3d.utility.Vector3dVector(colors)

    return pcd


def filter_pointcloud_range(
    pcd: o3d.geometry.PointCloud,
    x_range: Optional[Tuple[float, float]] = None, y_range: Optional[Tuple[float, float]] = None, z_range: Optional[Tuple[float, float]] = None,
    return_condition: bool = False
) -> Union[np.ndarray, o3d.geometry.PointCloud]:
    """Filter a point cloud by a range of x, y, z coordinates.

    Args:
        pcd: an open3d point cloud.
        x_range: a tuple of (min, max) x coordinates.
        y_range: a tuple of (min, max) y coordinates.
        z_range: a tuple of (min, max) z coordinates.
        return_condition: whether to return the condition used to filter the point cloud.

    Returns:
        an open3d point cloud or a numpy array of shape (N,) representing the condition.
    """

    condition = np.ones(len(pcd.points), dtype=np.bool_)
    xyz = np.asarray(pcd.points)
    if x_range is not None:
        condition &= (xyz[:, 0] >= x_range[0]) & (xyz[:, 0] <= x_range[1])
    if y_range is not None:
        condition &= (xyz[:, 1] >= y_range[0]) & (xyz[:, 1] <= y_range[1])
    if z_range is not None:
        condition &= (xyz[:, 2] >= z_range[0]) & (xyz[:, 2] <= z_range[1])

    if return_condition:
        return condition
    pcd = pcd.select_by_index(np.where(condition)[0])
    return pcd


def ransac_table_detection(
    pcd: o3d.geometry.PointCloud, distance_threshold: float = 0.005, ransac_n: int = 3, num_iterations: int = 1000, return_condition: bool = False
) -> Union[Tuple[np.ndarray, np.ndarray], Tuple[o3d.geometry.PointCloud, o3d.geometry.PointCloud]]:

    """Detect the table plane using RANSAC.

    Args:
        pcd: an open3d point cloud.
        distance_threshold: the distance threshold for RANSAC.
        ransac_n: the number of points to sample for RANSAC.
        num_iterations: the number of iterations for RANSAC.
        return_condition: whether to return the condition used to filter the point cloud.

    Returns:
        If return_condition is True, returns a tuple of (model, condition). Condition is a numpy array of shape (N,) representing the condition.
        Otherwise, returns a tuple of (inliers, outliers). Inliers and outliers are open3d point clouds.
    """

    model, inliers = pcd.segment_plane(distance_threshold, ransac_n, num_iterations)
    if return_condition:
        condition = np.zeros(len(pcd.points), dtype=np.bool_)
        condition[inliers] = True
        return model, condition

    inlier_cloud = pcd.select_by_index(inliers)
    outlier_cloud = pcd.select_by_index(inliers, invert=True)
    return inlier_cloud, outlier_cloud


def threshold_table_detection(
    pcd: o3d.geometry.PointCloud, z_threshold: float = 0.005, return_condition: bool = False
) -> Union[Tuple[np.ndarray, np.ndarray], Tuple[o3d.geometry.PointCloud, o3d.geometry.PointCloud]]:
    """Detect the table plane using a simple threshold.

    Args:
        pcd: an open3d point cloud.
        z_threshold: the threshold for the z coordinate.
        return_condition: whether to return the condition used to filter the point cloud.

    Returns:
        If return_condition is True, returns a numpy array of shape (N,) representing the condition.
        Otherwise, returns a tuple of (inliers, outliers). Inliers and outliers are open3d point clouds.
    """

    xyz = np.asarray(pcd.points)
    condition = np.abs(xyz[:, 2]) < z_threshold

    if return_condition:
        # TODO(Jiayuan Mao): check whether this should be + or -.
        model = np.array([0, 0, 1, -z_threshold])
        return model, condition

    inlier_cloud = pcd.select_by_index(np.where(condition)[0])
    outlier_cloud = pcd.select_by_index(np.where(condition)[0], invert=True)

    return inlier_cloud, outlier_cloud


def project_pointcloud_to_plane(pcd: o3d.geometry.PointCloud, model: np.ndarray) -> o3d.geometry.PointCloud:
    """Project a point cloud to a plane (into a 2D shape)

    Args:
        pcd: an open3d point cloud.
        model: a numpy array of shape (4,) representing the plane model.

    Returns:
        an open3d point cloud.
    """

    xyz = np.asarray(pcd.points)
    d = np.dot(xyz, model[:3]) + model[3]
    xyz -= d[:, None] * model[None, :3]
    new_pcd = o3d.geometry.PointCloud()
    new_pcd.points = o3d.utility.Vector3dVector(xyz)
    new_pcd.colors = pcd.colors
    return new_pcd


def _cdist(xyz1: np.ndarray, xyz2: np.ndarray) -> np.ndarray:
    """Compute the pairwise distance between two point clouds.

    Args:
        xyz1: a numpy array of shape (N, 3).
        xyz2: a numpy array of shape (M, 3).

    Returns:
        a numpy array of shape (N, M).
    """

    return np.sqrt(np.sum((xyz1[:, None, :] - xyz2[None, :, :]) ** 2, axis=-1))


def object_reconstruction_by_extrusion(
    pcd: o3d.geometry.PointCloud, model: np.ndarray, extrusion_distance: float = 0.005, min_height: Optional[float] = 0.05,
    include_inner_points: bool = False
) -> o3d.geometry.PointCloud:
    """Reconstruct an object by extrusion.

    Args:
        pcd: an open3d point cloud.
        model: a numpy array of shape (4,) representing the plane model.
        extrusion_distance: the distance to extrude.
        min_height: the minimum height of the extrusion. If an object is too thin, we will extrude it to this height.
        include_inner_points: whether to include the inner points of the object.
    """

    # Compute the max distance from the plane.
    xyz = np.asarray(pcd.points, dtype=np.float32)
    d = np.abs(model[:3].dot(xyz.T) + model[3]) / np.linalg.norm(model[:3])
    max_d = np.max(d)

    # Extrude the point cloud.
    max_extrusion_nr = int(max_d / extrusion_distance)

    projection = project_pointcloud_to_plane(pcd, model)
    # Subsample the projected point cloud.
    projection = projection.voxel_down_sample(voxel_size=0.005)
    # projection = projection.uniform_down_sample(every_k_points=10)
    projected_xyz = np.asarray(projection.points)
    projected_rgb = np.asarray(projection.colors)

    points = [(projected_xyz, projected_rgb)]
    if include_inner_points:
        for i in range(max_extrusion_nr):
            extruded_xyz = projected_xyz + (i + 1) * extrusion_distance * model[:3]
            # Compute the distance of all points to the extruded point cloud.
            distances = _cdist(extruded_xyz, xyz)
            # Find the closest point on the original pcd for each point.
            closest_points = np.min(distances, axis=1)
            # If the closest point is too close, we stop extruding for this point.
            mask = closest_points > extrusion_distance
            extruded_xyz = extruded_xyz[mask]
            extruded_rgb = projected_rgb[mask]

            points.append((extruded_xyz, extruded_rgb))

            projected_xyz = projected_xyz[mask]
            projected_rgb = projected_rgb[mask]

    if min_height is not None and np.median(xyz[:, 2]) < min_height:
        # if the object is just too thin, we add a top surface.
        points.append((projected_xyz + np.array([0, 0, 0.05]), projected_rgb))
    else:
        points.append((xyz, np.asarray(pcd.colors)))

    # Concatenate all points.
    xyz = np.concatenate([p[0] for p in points], axis=0)
    rgb = np.concatenate([p[1] for p in points], axis=0)
    pcd = o3d.geometry.PointCloud()
    pcd.points, pcd.colors = o3d.utility.Vector3dVector(xyz), o3d.utility.Vector3dVector(rgb)
    return pcd


def remove_outliers(pcd: o3d.geometry.PointCloud, min_neighbors: int, radius: float, verbose: bool = False) -> o3d.geometry.PointCloud:
    """Remove outliers from a point cloud.

    Note: this function is not used in the pipeline.

    Args:
        pcd: an open3d point cloud.
        min_neighbors: the minimum number of neighbors.
        radius: the radius to search for neighbors.

    Returns:
        an open3d point cloud.
    """

    if verbose:
        print("  Removing outliers... Before: ", len(pcd.points))
    pcd, _ = pcd.remove_radius_outlier(nb_points=min_neighbors, radius=radius)
    if verbose:
        print("  Removing outliers... After: ", len(pcd.points))
    return pcd


def mesh_reconstruction_alpha_shape(pcd: o3d.geometry.PointCloud, alpha: float = 0.1) -> o3d.geometry.TriangleMesh:
    """Reconstruct a mesh from a point cloud.

    Args:
        pcd: an open3d point cloud.
        alpha: the alpha value for the alpha shape.

    Returns:
        an open3d triangle mesh.
    """
    pcd = pcd.voxel_down_sample(voxel_size=0.0025)
    pcd.estimate_normals()
    # Project the object to the table plane
    mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_alpha_shape(pcd, alpha=alpha)
    # t_mesh = o3d.t.geometry.TriangleMesh.from_legacy(mesh)
    # t_filled = t_mesh.fill_holes(hole_size=1)
    # mesh = t_filled.to_legacy()
    mesh.compute_vertex_normals()
    mesh.compute_triangle_normals()
    return mesh


def mesh_nms(detected_objects: List[ObjectDetectionResult], bbox_overlapping_threshold: float = 0.1) -> List[ObjectDetectionResult]:
    """Perform non-maximum suppression on the detected objects.

    Args:
        detected_objects: a list of ObjectDetectionResult.
        bbox_overlapping_threshold: the threshold for overlapping.

    Returns:
        a list of ObjectDetectionResult.
    """

    # Sort the objects by their scores.
    # Compute the pairwise IoU.
    ious = np.zeros((len(detected_objects), len(detected_objects)))
    for i in range(len(detected_objects)):
        for j in range(i + 1, len(detected_objects)):
            ious[i, j] = ious[j, i] = compute_iou_3d(detected_objects[i].mesh_bbox, detected_objects[j].mesh_bbox)

    ious_visualization = [['' for _ in range(len(detected_objects) + 1)] for _ in range(len(detected_objects) + 1)]
    for i in range(len(detected_objects)):
        for j in range(len(detected_objects)):
            if i == j:
                ious_visualization[i + 1][j + 1] = 'X'
            else:
                ious_visualization[i + 1][j + 1] = f'{ious[i, j]:.2f}'
        ious_visualization[i + 1][0] = f'{detected_objects[i].label}'
        ious_visualization[0][i + 1] = f'{detected_objects[i].label}'
    print(jacinle.tabulate(ious_visualization, tablefmt='simple'))

    # Perform NMS.
    selected_objects_indices = []
    for i in range(len(detected_objects)):
        if len(selected_objects_indices) == 0:
            selected_objects_indices.append(i)
        else:
            if np.max(ious[i, selected_objects_indices]) > bbox_overlapping_threshold:
                continue
            selected_objects_indices.append(i)
    return [detected_objects[i] for i in selected_objects_indices]


def compute_iou_3d(bbox1: o3d.geometry.AxisAlignedBoundingBox, bbox2: o3d.geometry.AxisAlignedBoundingBox) -> float:
    """Compute the 3D IoU between two bounding boxes.

    Args:
        bbox1: an open3d axis aligned bounding box.
        bbox2: an open3d axis aligned bounding box.

    Returns:
        the 3D IoU.
    """

    # Compute the intersection.
    bbox1 = np.asarray(bbox1.get_box_points())
    bbox2 = np.asarray(bbox2.get_box_points())

    a, b = bbox1.min(axis=0), bbox1.max(axis=0)
    c, d = bbox2.min(axis=0), bbox2.max(axis=0)

    x, y = np.maximum(a, c), np.minimum(b, d)
    intersection = np.maximum(y - x, 0)
    intersection_volume = np.prod(intersection)

    a_volume = np.prod(b - a)
    b_volume = np.prod(d - c)

    union_volume = a_volume + b_volume - intersection_volume
    return intersection_volume / union_volume


def canonize_mesh_center_(mesh_: o3d.geometry.TriangleMesh) -> Tuple[o3d.geometry.TriangleMesh, np.ndarray]:
    """Canonize the mesh center. Note that this function modifies the mesh in place.

    Args:
        mesh_: an open3d triangle mesh.

    Returns:
        a tuple of an open3d triangle mesh, the center of the mesh.
    """

    mesh_copy = o3d.geometry.TriangleMesh(mesh_)
    # Compute the center of the mesh.
    center = mesh_copy.get_center()
    # Compute the transformation matrix.
    T = np.eye(4)
    T[:3, 3] = -center
    # Apply the transformation matrix.
    mesh_copy.transform(T)
    return mesh_copy, center


def compute_transformation_from_plane_equation(a, b, c, d):
    """Compute the transformation matrix from a plane equation.

    Args:
        a, b, c, d: the plane equation parameters.

    Returns:
        a 4x4 transformation matrix.
    """
    # Normal of the plane
    N = np.array([a, b, c])
    # Normalize the normal vector
    n = N / np.linalg.norm(N)
    # Z axis
    z_axis = np.array([0, 0, 1])
    # Axis around which to rotate - cross product between Z axis and n
    axis = np.cross(z_axis, n)
    axis = axis / np.linalg.norm(axis)  # normalize the axis
    # Angle between Z axis and n
    angle = np.arccos(np.dot(z_axis, n))
    # Compute rotation matrix using scipy's Rotation
    rotation_matrix = quat2mat(axisangle2quat(axis, angle))
    # Compute translation vector
    translation_vector = -d * n
    # Construct transformation matrix
    transformation_matrix = np.eye(4)
    transformation_matrix[:3, :3] = rotation_matrix
    transformation_matrix[:3, 3] = translation_vector
    return transformation_matrix


def visualize_with_camera_matrix(geometries, camera_extrinsics):
    """Visualize the geometries with a camera matrix.

    Args:
        geometries: a list of open3d geometries.
        camera_extrinsics: a 4x4 camera extrinsics matrix.
    """
    vis = o3d.visualization.VisualizerWithKeyCallback()
    vis.create_window()
    vis.get_render_option().background_color = np.asarray([0, 0, 0])

    view_ctl = vis.get_view_control()

    for geom in geometries:
        vis.add_geometry(geom)

    cam = view_ctl.convert_to_pinhole_camera_parameters()
    cam.extrinsic = camera_extrinsics
    view_ctl.convert_from_pinhole_camera_parameters(cam)

    vis.run()
    vis.destroy_window()


def load_scene_in_pybullet(client: BulletClient, scene_dir: str, verbose: bool = False, simplify_name: bool = True) -> Dict[str, Dict]:
    scene_metafile = osp.join(scene_dir, 'metainfo.json')
    scene_metainfo = jacinle.load(scene_metafile)

    unique_labels = set()
    non_unique_labels = set()
    non_unique_label_counts = dict()

    if simplify_name:
        for o in scene_metainfo['objects']:
            label = o['label'].replace(' ', '_')
            if label in unique_labels:
                non_unique_labels.add(label)
                non_unique_label_counts[label] = 0
                unique_labels.remove(label)
            else:
                unique_labels.add(label)

    load_metainfo = dict()
    for o in scene_metainfo['objects']:
        urdf_filename = osp.basename(o['urdf'])
        urdf_filename = osp.join(scene_dir, urdf_filename)
        pos = o['pos']
        pos = (pos[0], pos[1], pos[2] + 0.01)

        if simplify_name:
            name = o['label'].replace(' ', '_')
            if name in non_unique_labels:
                non_unique_label_counts[name] += 1
                name = name + '_' + str(non_unique_label_counts[name])
            else:
                assert name in unique_labels
        else:
            name = o['unique_label'].replace(' ', '_')

        index = client.load_urdf(urdf_filename, body_name=name, pos=pos)
        load_metainfo[name] = dict(id=index, label=o['label'])

    return load_metainfo


concepts/hw_interface/george_vision/object_centric_vision.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : segmentation_models.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/05/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import random
from typing import Tuple, List
from dataclasses import dataclass

import cv2
import numpy as np
import torch
import matplotlib.pyplot as plt


@dataclass
class InstanceSegmentationResult(object):
    masks: np.ndarray
    """A numpy array of shape (N, H, W) and dtype bool."""

    pred_cls: List[str]
    """A list of length N."""

    pred_boxes: List[Tuple[Tuple[int, int], Tuple[int, int]]]
    """A list of length N. Each element is a tuple of ((x1, y1), (x2, y2))."""

    @property
    def nr_objects(self) -> int:
        return len(self.pred_cls)


class ImageBasedPCDSegmentationModel(object):
    def __init__(self, model_name: str = 'maskrcnn_resnet50_fpn_v2', device: str = 'cpu', score_threshold: float = 0.5):
        from torchvision.models.detection import maskrcnn_resnet50_fpn, MaskRCNN_ResNet50_FPN_Weights
        from torchvision.models.detection import maskrcnn_resnet50_fpn_v2, MaskRCNN_ResNet50_FPN_V2_Weights

        self.model_name = model_name
        self.device = torch.device(device)

        if self.model_name == 'maskrcnn_resnet50_fpn':
            model_fn = maskrcnn_resnet50_fpn
            weights = MaskRCNN_ResNet50_FPN_Weights.COCO_V1
        elif self.model_name == 'maskrcnn_resnet50_fpn_v2':
            model_fn = maskrcnn_resnet50_fpn_v2
            weights = MaskRCNN_ResNet50_FPN_V2_Weights.COCO_V1
        else:
            raise ValueError('Unknown model name: {}.'.format(model_name))

        self.model = model_fn(weights=weights)
        self.model.eval()
        self.model.to(self.device)

        self.preprocess = weights.transforms()
        self.score_threshold = score_threshold

    def segment_image(self, image: np.ndarray) -> InstanceSegmentationResult:
        """Segment an image into objects and background.

        Args:
            image: a numpy array of shape (H, W, 3) and dtype uint8.

        Returns:
            A tuple of (segmented_image, object_names, object_bboxes).
            segmented_image is a numpy array of shape (H, W, 3) and dtype uint8.
            object_names is a list of strings.
            object_bboxes is a list of tuples of ((x1, y1), (x2, y2)).
        """

        image = torch.tensor(image.transpose(2, 0, 1) / 255.0, dtype=torch.float32)
        image= self.preprocess(image)
        pred = self.model([image])[0]

        pred_score = list(pred['scores'].detach().numpy())
        pred_t = [pred_score.index(x) for x in pred_score if x > self.score_threshold][-1]

        masks = (pred['masks'] > 0.5).squeeze().detach().cpu().numpy()
        pred_class = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in list(pred['labels'].numpy())]
        pred_boxes = [
            ((int(i[0]), int(i[1])), (int(i[2]), int(i[3])))
            for i in list(pred['boxes'].detach().numpy().astype(np.int64))
        ]

        masks = masks[:pred_t+1]
        pred_boxes = pred_boxes[:pred_t+1]
        pred_class = pred_class[:pred_t+1]
        return InstanceSegmentationResult(masks, pred_class, pred_boxes)


class PointGuidedImageSegmentationModel(object):
    def __init__(self, checkpoint_path, model='sam_default'):
        from segment_anything import SamPredictor, sam_model_registry

        if model == 'sam_default':
            model = sam_model_registry['default']
        else:
            raise ValueError('Unknown model name: {}.'.format(model))
        self.model = model(checkpoint=checkpoint_path)
        self.predicator = SamPredictor(self.model)
        self.last_image_id = None

    def segment_from_point(self, image, point):
        if self.last_image_id is None or self.last_image_id != id(image):

            self.predicator.set_image(image)
            self.last_image_id = id(image)
        masks, _, _ = self.predicator.predict(point_coords=np.array([point]), point_labels=np.array([1]))
        if len(masks) > 1:
            mask = masks[-2].astype(np.uint8)
        else:
            mask = masks[-1].astype(np.uint8)
        mask = remove_remains(mask, point)
        return mask


COCO_INSTANCE_CATEGORY_NAMES = [
    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
    'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A',
    'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
    'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
    'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
    'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
    'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
    'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
    'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]


def remove_remains(img, interest_point):
    """Remove remains which are not adjacent with interest_point."""
    img = img.copy().astype(np.uint8)
    h, w = img.shape[:2]
    mask = np.zeros((h + 2, w + 2), np.uint8)

    img_inv = img.copy()
    cv2.floodFill(img_inv, mask, tuple(interest_point), 0)
    img -= img_inv
    return img


def random_colored_mask(image):
    colours = [
        [0, 255, 0], [0, 0, 255], [255, 0, 0],
        [0, 255, 255], [255, 255, 0], [255, 0, 255],
        [80, 70, 180], [250, 80, 190], [245, 145, 50], [70, 150, 250], [50, 190, 190]
    ]
    rgb = np.array(random.choice(colours))
    rgb = np.reshape(rgb, (1, 1, 3))
    image = image[:, :, None]
    image = np.where(image != 0, image * rgb, image)
    return image.astype(np.uint8)


def visualize_instance_segmentation(image: np.ndarray, result: InstanceSegmentationResult, rect_th: int = 3, text_size: int = 1, text_th: int = 3):
    image = image.copy()
    for i in range(len(result.masks)):
        if 'table' in result.pred_cls[i]:
            continue
        rgb_mask = random_colored_mask(result.masks[i])
        image = cv2.addWeighted(image, 1, rgb_mask, 0.5, 0)
        cv2.rectangle(image, result.pred_boxes[i][0], result.pred_boxes[i][1], color=(0, 255, 0), thickness=rect_th)
        cv2.putText(image, result.pred_cls[i], (result.pred_boxes[i][0][0], result.pred_boxes[i][0][1] + 10), cv2.FONT_HERSHEY_SIMPLEX, text_size, (0, 255, 0), thickness=text_th)

    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    plt.xticks([])
    plt.yticks([])
    plt.tight_layout()
    plt.show()


concepts/hw_interface/george_vision/segmentation_models.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/12/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.



concepts/hw_interface/realsense/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : device.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/12/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""
Wrapper around pyrealsense2. Based on:
https://github.com/IntelRealSense/librealsense/issues/8388#issuecomment-782395443
"""

from typing import Dict, List, Optional, Tuple, ClassVar, Callable

import numpy as np
import pyrealsense2 as rs

__all__ = ['RealSenseDevice', 'start_pipelines', 'stop_pipelines']


# Note: change these to fit your use case. Assuming USB 3.2 connection.
_NAME_TO_STREAM_CONFIGURATIONS: Dict[str, List[Tuple]] = {
    # Mapping of camera name to a list of streams to enable
    # in the cfg.enable_stream format
    "Intel RealSense D435": [
        (rs.stream.depth, 1280, 720, rs.format.z16, 30),
        (rs.stream.color, 1920, 1080, rs.format.bgr8, 30),
    ],
    "Intel RealSense L515": [
        (rs.stream.depth, 1024, 768, rs.format.z16, 30),
        (rs.stream.color, 1920, 1080, rs.format.bgr8, 30),
    ],
}


class RealSenseDevice(object):
    align: ClassVar[Callable] = rs.align(rs.stream.color)
    ctx: ClassVar[rs.context] = rs.context()

    def __init__(self, name: str, serial_number: str):
        self.name = name
        self.serial_number = serial_number

        self.pipeline: Optional[rs.pipeline] = None
        self.registered_points: List[Tuple[float, float]] = list()
        self.color_intrinsics = None
        self.depth_intrinsics = None

    @classmethod
    def from_rs_device(cls, dev: rs.device) -> 'RealSenseDevice':
        name = dev.get_info(rs.camera_info.name)
        serial_number = dev.get_info(rs.camera_info.serial_number)
        return cls(name, serial_number)

    @classmethod
    def find_devices(cls, device_filter: str = "") -> List['RealSenseDevice']:
        """
        Get devices as detected by RealSense and filter devices that only
        contain the provided device_filter string in their name.
        e.g. to filter for D435 only you can call `find_devices("D435")`
        """
        devices = [cls.from_rs_device(dev) for dev in cls.ctx.devices]

        # Filter devices
        if device_filter:
            device_filter = device_filter.lower()
            devices = [d for d in devices if device_filter in d.name.lower()]
            print(f"Found devices (filter={device_filter}): {devices}")
        else:
            print(f"Found devices: {devices}")

        if not devices:
            raise RuntimeError("No devices connected!")
        return devices

    @property
    def stream_configurations(self) -> List[Tuple]:
        if self.name not in _NAME_TO_STREAM_CONFIGURATIONS:
            raise RuntimeError(f"Configuration not specified for {self.name}")

        return _NAME_TO_STREAM_CONFIGURATIONS[self.name]

    def start_pipeline(self) -> None:
        """Start RealSense pipeline"""
        if self.pipeline is not None:
            print(f"Pipeline already started for {self}")
            return

        # Setup pipeline and configuration
        pipeline = rs.pipeline(self.ctx)
        cfg = rs.config()
        cfg.enable_device(self.serial_number)
        for stream_configuration in self.stream_configurations:
            cfg.enable_stream(*stream_configuration)

        try:
            prof = pipeline.start(cfg)

            depth_stream = prof.get_stream(rs.stream.depth)
            color_stream = prof.get_stream(rs.stream.color)
            depth_intr = depth_stream.as_video_stream_profile().get_intrinsics()
            color_intr = color_stream.as_video_stream_profile().get_intrinsics()

            self.depth_intrinsics = get_intrinsics_matrix(depth_intr)
            self.color_intrinsics = get_intrinsics_matrix(color_intr)
        except RuntimeError as e:
            message = str(e)
            if message == "Couldn't resolve requests":
                # Something wrong with stream configurations probably
                raise RuntimeError(
                    f"{message} for {self}. Check stream configuration and USB connection."
                )
            else:
                raise e

        # Set the pipeline on the device
        self.pipeline = pipeline
        print(f"Started pipeline for {self}")

    def stop_pipeline(self) -> None:
        """Stop RealSense pipeline"""
        if not self.pipeline:
            print(f"Warning! Device {self} does not have a pipeline initialized")
        else:
            self.pipeline.stop()
            print(f"Stopped pipeline for {self}")

    def capture_images(self) -> Tuple[np.ndarray, np.ndarray]:
        """Capture color and depth images"""
        if self.pipeline is None:
            raise RuntimeError(f"Pipeline for {self} not started!")

        # Get frames and align
        frames = self.pipeline.wait_for_frames()
        aligned_frames = self.align.process(frames)

        # Get color and depth frame
        depth_frame = aligned_frames.get_depth_frame()
        color_frame = aligned_frames.get_color_frame()

        if not depth_frame or not color_frame:
            raise RuntimeError("Could not capture both color and depth frame.")

        # Convert to numpy arrays
        depth_image = np.asanyarray(depth_frame.get_data())
        color_image = np.asanyarray(color_frame.get_data())

        return color_image, depth_image

    def capture_color_image(self) -> np.ndarray:
        return self.capture_images()[0]

    def capture_depth_image(self) -> np.ndarray:
        return self.capture_images()[1]

    def __str__(self) -> str:
        return f"{self.name} ({self.serial_number})"

    def __repr__(self) -> str:
        return str(self)


def start_pipelines(devices: List[RealSenseDevice]) -> None:
    """Enable each device by starting a stream"""
    for device in devices:
        device.start_pipeline()


def stop_pipelines(devices: List[RealSenseDevice]) -> None:
    """Stop all the pipelines"""
    for device in devices:
        device.stop_pipeline()


def get_intrinsics_matrix(intr) -> np.ndarray:
    fx = float(intr.fx)
    fy = float(intr.fy)
    ppx = float(intr.ppx)
    ppy = float(intr.ppy)
    axs = 0.0
    return np.array([
        [fx, axs, ppx],
        [0.0, fy, ppy],
        [0.0, 0.0, 1.0]
    ])


concepts/hw_interface/realsense/device.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : visualizer.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/12/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import os
import numpy as np
import cv2
import jacinle.io as io

from datetime import datetime
from enum import Enum
from typing import TYPE_CHECKING, Optional, Tuple, List

if TYPE_CHECKING:
    from concepts.hw_interface.realsense.device import RealSenseDevice


class _WindowEvent(Enum):
    EXIT = "exit"
    SAVE = "save"
    NONE = "none"


class CV2Visualizer(object):
    def __init__(self, registered_points: Optional[List[Tuple[int, int]]] = None):
        if registered_points is None:
            self.registered_points = []
        else:
            self.registered_points = list(registered_points)

    def register_point(self, event, x, y, flag, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            found = None
            for i, (xx, yy) in enumerate(self.registered_points):
                if np.linalg.norm(np.array([xx, yy], dtype='float32') - np.array([x, y], dtype='float32')) < 10:
                    found = i
                    break

            if found is not None:
                self.registered_points = self.registered_points[:found] + self.registered_points[found + 1:]
            else:
                self.registered_points.append((x, y))

    def imshow(self, color_image: np.ndarray, depth_image: np.ndarray) -> None:
        color_image_new = color_image.copy()
        for x, y in self.registered_points:
            cv2.circle(color_image_new, (x, y), 5, (255, 0, 0), -1)

        # Form heatmap for depth and stack with color image
        depth_colormap = cv2.applyColorMap(
            cv2.convertScaleAbs(depth_image, alpha=0.03), cv2.COLORMAP_JET
        )
        images = np.hstack((color_image_new, depth_colormap))

        window_name = str(self)
        cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
        cv2.imshow(window_name, images)
        cv2.setMouseCallback(window_name, self.register_point)
        cv2.waitKey(1)

    def run(self, color_image: np.ndarray, depth_image: Optional[np.ndarray] = None, file_dirname: str = 'visualization', file_suffix: str = '', save: bool = True) -> None:
        while True:
            color_image_new = color_image.copy()
            for x, y in self.registered_points:
                cv2.circle(color_image_new, (x, y), 5, (255, 0, 0), -1)

            # Form heatmap for depth and stack with color image
            if depth_image is not None:
                depth_colormap = cv2.applyColorMap(
                    cv2.convertScaleAbs(depth_image, alpha=0.03), cv2.COLORMAP_JET
                )
                images = np.hstack((color_image_new, depth_colormap))
            else:
                images = color_image_new

            window_name = str(self)
            cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
            cv2.imshow(window_name, images)
            cv2.setMouseCallback(window_name, self.register_point)
            key = cv2.waitKey(1)

            # Exit on 'esc' or 'q'
            if key & 0xFF == ord("q") or key == 27:
                cv2.destroyWindow(window_name)
                break

            if key == ord('c'):
                print('Registered points:')
                print('-' * 80)
                print(self.registered_points)

                if save:
                    cv2.imwrite(f'{file_dirname}/calibration{file_suffix}.png', images)
                    io.dump(f'{file_dirname}/calibration_image_points{file_suffix}.pkl', self.registered_points)

        if save:
            cv2.imwrite(f'{file_dirname}/calibration{file_suffix}.png', images)
            io.dump(f'{file_dirname}/calibration_image_points{file_suffix}.pkl', self.registered_points)

        cv2.destroyAllWindows()
        return self.registered_points


class RealSenseVisualizer(object):
    def __init__(self, device):
        self.device: 'RealSenseDevice' = device

    def visualize(self, save_dir: str = "", save_image: bool = False) -> _WindowEvent:
        """
        Visualize color and depth images in a cv2 window.
        Terminates when 'esc' or 'q' key is pressed.
        Saves an image when the 's' key is pressed or if the 'save_image'
        flag is specified.
        Images are saved to the specified save_dir, which we
        assume to already exist.
        Returns a _WindowEvent indicating status of cv2.
        """
        color_image, depth_image = self.device.capture_images()

        # Form heatmap for depth and stack with color image
        depth_colormap = cv2.applyColorMap(
            cv2.convertScaleAbs(depth_image, alpha=0.03), cv2.COLORMAP_JET
        )
        images = np.hstack((color_image, depth_colormap))

        window_name = str(self)
        cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
        cv2.imshow(window_name, images)
        cv2.setMouseCallback(window_name, self.register_point)
        key = cv2.waitKey(1)

        # Exit on 'esc' or 'q'
        if key & 0xFF == ord("q") or key == 27:
            cv2.destroyWindow(window_name)
            return _WindowEvent.EXIT

        # Save images if 's' key is pressed
        if key == 115 or save_image:
            timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
            depth_fname = os.path.join(
                save_dir, f"{timestamp}-{self.device.serial_number}-depth.png"
            )
            color_fname = os.path.join(
                save_dir, f"{timestamp}-{self.device.serial_number}-color.png"
            )
            cv2.imwrite(depth_fname, depth_image)
            cv2.imwrite(color_fname, color_image)
            print(f"Saved depth image for {self} to {depth_fname}")
            print(f"Saved color image to {self} to {color_fname}")
            return _WindowEvent.SAVE

        if key == ord('c'):
            print('CALIBRATION registered points:')
            print('-' * 80)
            print(self.registered_points)

        return _WindowEvent.NONE


def visualize_devices(devices: List['RealSenseDevice'], save_dir: str = "") -> None:
    """
    Visualizes all the devices in a cv2 window. Press 'q' or 'esc' on any of
    the windows to exit the infinite loop.
    Press the 's' key in a specific window to save the color and depth image
    to disk.
    You can use a similar loop interface in other places where you need
    a live camera feed (e.g. collecting demonstrations).
    """
    while True:
        should_exit = False
        save = False

        for device in devices:
            window_event = device.visualize(save_dir, save_image=save)
            if window_event == _WindowEvent.SAVE:
                # We use this to propagate a save command across all windows
                save = True

            # Exit all windows
            if window_event == _WindowEvent.EXIT:
                should_exit = True
                break

        if should_exit:
            print("Exit key pressed.")
            cv2.destroyAllWindows()
            break


def run_4corner_calibration(devices: List['RealSenseDevice'], save_dir: str = "") -> None:
    while True:
        should_exit = False

        for device in devices:
            window_event = device.visualize(save_dir, save_image=False)
            # Exit all windows
            if window_event == _WindowEvent.EXIT:
                should_exit = True
                break

        if should_exit:
            print("Exit key pressed.")
            cv2.destroyAllWindows()
            break


if __name__ == "__main__":
    import jacinle
    from concepts.hw_interface.realsense.device import start_pipelines, stop_pipelines

    parser = jacinle.JacArgumentParser()
    parser.add_argument('--device', choices=['l515', 'd435'], default='d435')
    args = parser.parse_args()

    # Detect devices, start pipelines, visualize, and stop pipelines
    devices = RealSenseDevice.find_devices(args.device)
    start_pipelines(devices)

    try:
        visualize_devices(devices)
    finally:
        stop_pipelines(devices)


concepts/hw_interface/realsense/visualizer.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/13/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data structures for language semantics and tools for parsing.
Most of this package are designed for neuro-symbolic representations of
natrual language semantics and grammar-based language learning.

See the following papers for more details:

- The Neuro-Symbolic Concept Learner: Interpreting Scenes, Words, and Sentences From Natural Supervision (http://nscl.csail.mit.edu/)
- Grammar-Based Grounded Lexicon Learning (https://g2l2.csail.mit.edu/)

Here's a quick summary of the sub-modules:

- :mod:`concepts.language.ccg` contains tools for parsing sentences with (purely symbolic) Combinatory Categorial Grammars (CCGs).
"""

concepts/language/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/04/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Tools for parsing sentences with (purely symbolic) Combinatory Categorial Grammars."""

from . import learning
from .composition import CCGCompositionDirection, CCGCompositionType, CCGCompositionContext, get_ccg_composition_context, CCGCompositionResult, CCGComposable, CCGCompositionSystem
from .syntax import CCGSyntaxType, CCGPrimitiveSyntaxType, CCGConjSyntaxType, CCGComposedSyntaxType, CCGSyntaxSystem
from .semantics import CCGSemanticsConjFunction, CCGSemanticsSimpleConjFunction, CCGSemanticsLazyValue, CCGSemantics
from .grammar import Lexicon, LexiconUnion, CCGNode, compose_ccg_nodes, CCG

__all__ = [
    'CCGCompositionDirection', 'CCGCompositionType',
    'CCGCompositionContext', 'get_ccg_composition_context',
    'CCGCompositionResult', 'CCGComposable',
    'CCGCompositionSystem',
    'CCGSyntaxType', 'CCGPrimitiveSyntaxType', 'CCGConjSyntaxType', 'CCGComposedSyntaxType',
    'CCGSyntaxSystem',
    'CCGSemanticsConjFunction', 'CCGSemanticsSimpleConjFunction',
    'CCGSemanticsLazyValue', 'CCGSemantics',
    'Lexicon', 'LexiconUnion', 'CCGNode', 'compose_ccg_nodes', 'CCG',
    'learning'
]


concepts/language/ccg/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : composition.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/05/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Composition rules for CCG.

The two main classes are :class:`CCGComposable` and :class:`CCGCompositionSystem`.
"""

import contextlib
from typing import TYPE_CHECKING, Optional, Union, List, Dict, Callable
from dataclasses import dataclass
from jacinle.utils.cache import cached_property
from jacinle.utils.defaults import option_context
from jacinle.utils.enum import JacEnum
from jacinle.utils.printing import indent_text

if TYPE_CHECKING:
    from concepts.language.ccg.grammar import CCGNode

__all__ = [
    'CCGCompositionDirection', 'CCGCompositionType', 'CCGCompositionError',
    'CCGCompositionContext', 'get_ccg_composition_context',
    'CCGCompositionResult', 'CCGCoordinationImmNode', 'CCGComposable',
    'CCGCompositionSystem'
]


class CCGCompositionDirection(JacEnum):
    """Composition directions (left or right)."""
    LEFT = 'left'
    RIGHT = 'right'


class CCGCompositionType(JacEnum):
    """Composition types (e.g., application and coordination)."""

    LEXICON = 'lexicon'
    FORWARD_APPLICATION = 'forward_application'
    BACKWARD_APPLICATION = 'backward_application'
    COORDINATION = 'coordination'
    NONE = 'none'


class CCGCompositionError(Exception):
    """The error raised when composition fails."""
    pass


class CCGCompositionContext(option_context(
    '_CCGCompositionContext',
    syntax=True,
    semantics=True,
    semantics_lazy_composition=False,
    exc_verbose=True
)):
    """An option context for CCG composition."""

    syntax: bool
    """Whether to perform syntax composition."""

    semantics: bool
    """Whether to perform semantics composition."""

    semantics_lazy_composition: bool
    """Whether to perform lazy semantics composition."""

    exc_verbose: bool
    """Whether to raise verbose exceptions."""

    @contextlib.contextmanager
    def exc(self, exc_type: Optional[type] = None, from_: Optional[Exception] = None):
        """Context manager for handling composition errors. If `exc_verbose` is True, the error will be printed out.

        Example:
            >>> with get_ccg_composition_context().exc():
            >>>     raise CCGCompositionError('some error')

        Args:
            exc_type: the exception type to raise. If None, the original exception will be raised.
            from_: the original exception to raise.
        """
        if self.exc_verbose:
            yield
        else:
            if exc_type is None:
                exc_type = CCGCompositionError
            if from_ is not None:
                raise exc_type() from from_
            raise exc_type()


get_ccg_composition_context: Callable[[], CCGCompositionContext] = CCGCompositionContext.get_default


@dataclass
class CCGCoordinationImmNode(object):
    """An intermediate node for coordination."""

    conj: 'CCGComposable'
    """The conjunction node."""

    rhs: 'CCGComposable'
    """The right-hand side node."""

    # Adding a few properties to make it compatible with CCGComposable.

    @property
    def is_none(self) -> bool:
        """Whether this node is a None node. This property is always False."""
        return False

    @property
    def is_conj(self) -> bool:
        """Whether this node is a conjunction node. This property is always False."""
        return False

    @property
    def is_value(self) -> bool:
        """Whether this node is a value node. This property is always False."""
        return False

    @property
    def is_function(self) -> bool:
        """Whether this node is a function node. This property is always False."""
        return False


class CCGComposable(object):
    """The basic class for composable elements (including syntax and semantics) for CCG."""

    @property
    def is_none(self) -> bool:
        """Whether this element is None."""
        return False

    @property
    def is_conj(self) -> bool:
        """Whether this element is a conjunction."""
        return False

    def compose(self, rhs: Union['CCGComposable', CCGCoordinationImmNode], composition_type: CCGCompositionType) -> Union['CCGComposable', CCGCoordinationImmNode]:
        """Compose this element with another element. This function will call the corresponding composition function according to the composition type.
        Note that since the coordination composition has three arguments, this function will return a :class:`CCGCoordinationImmNode` for the first two arguments
        in coordination composition.

        Args:
            rhs: the right-hand side element.
            composition_type: the composition type.

        Returns:
            The composed element.
        """

        if isinstance(rhs, CCGCoordinationImmNode) and composition_type is not CCGCompositionType.COORDINATION:
            raise CCGCompositionError('Can not make non-coordination composition for CCGCoordinationImmNode.')
        if (self.is_none or (not isinstance(rhs, CCGCoordinationImmNode) and rhs.is_none)) and composition_type is not CCGCompositionType.NONE:
            raise CCGCompositionError('Can not make non-None composition with none elements.')

        if composition_type is CCGCompositionType.LEXICON:
            raise CCGCompositionError('Lexicon composition type is only used for leaf level nodes.')
        elif composition_type is CCGCompositionType.FORWARD_APPLICATION:
            return self.fapp(rhs)
        elif composition_type is CCGCompositionType.BACKWARD_APPLICATION:
            return rhs.bapp(self)
        elif composition_type is CCGCompositionType.COORDINATION:
            return self.coord(rhs)
        elif composition_type is CCGCompositionType.NONE:
            return self.none(rhs)

    def fapp(self, rhs: 'CCGComposable') -> 'CCGComposable':
        """Forward application composition."""
        assert not self.is_none and not rhs.is_none
        return self._fapp(rhs)

    def bapp(self, lhs: 'CCGComposable') -> 'CCGComposable':
        """Backward application composition."""
        assert not self.is_none and not lhs.is_none
        return self._bapp(lhs)

    def none(self, rhs: 'CCGComposable') -> 'CCGComposable':
        """None composition (composition with a None element)."""
        if rhs.is_none:
            return self
        elif self.is_none:
            return rhs
        with get_ccg_composition_context().exc(CCGCompositionError):
            raise CCGCompositionError(f'Invalid None composition: lhs={self} and rhs={rhs}.')

    def coord(self, other: Union['CCGComposable', CCGCoordinationImmNode]) -> Union['CCGComposable', CCGCoordinationImmNode]:
        """Coordination composition."""
        if isinstance(other, CCGCoordinationImmNode):
            return other.conj.coord3(self, other.rhs)
        elif self.is_conj:
            return CCGCoordinationImmNode(self, other)
        with get_ccg_composition_context().exc(CCGCompositionError):
            raise CCGCompositionError(f'Invalid coordination composition: lhs={self}, rhs={other}.')

    def coord3(self, lhs: 'CCGComposable', rhs: 'CCGComposable') -> 'CCGComposable':
        """Coordination composition with three elements."""
        assert not self.is_none and not lhs.is_none and not rhs.is_none
        return self._coord3(lhs, rhs)

    def _fapp(self, rhs: 'CCGComposable') -> 'CCGComposable':
        raise NotImplementedError()

    def _bapp(self, lhs: 'CCGComposable') -> 'CCGComposable':
        raise NotImplementedError()

    def _coord3(self, lhs: 'CCGComposable', rhs: 'CCGComposable') -> 'CCGComposable':
        raise NotImplementedError()


@dataclass
class CCGCompositionResult(object):
    """The result of a CCG composition."""

    composition_type: CCGCompositionType
    """The composition type applied at the current node."""

    result: Union[CCGComposable, CCGCoordinationImmNode, 'CCGNode']
    """The result of the composition."""


class CCGCompositionSystem(object):
    """The CCG composition system. It keeps track of the rules that can be used for composition."""

    def __init__(self, name: str, weights: Dict[CCGCompositionType, float]):
        """Initialize the CCG composition system.

        Args:
            name: the name of the composition system.
            weights: the weights of the composition types, which should be a dictionary mapping from :class:`CCGCompositionType` to float.
        """
        self.name = name
        self.weights = weights

    @cached_property
    def allowed_composition_types(self) -> List[CCGCompositionType]:
        """Get the list of allowed composition types.

        Returns:
            the list of allowed composition types.
        """
        return [c for c in CCGCompositionType.choice_objs() if c in self.weights and c is not CCGCompositionType.LEXICON]

    def try_compose(self, lhs: CCGComposable, rhs: CCGComposable) -> CCGCompositionResult:
        """Try to compose two elements.
        This function will try to compose the two elements with all allowed composition types, and return the result that works.

        Args:
            lhs: the left-hand side element.
            rhs: the right-hand side element.

        Returns:
            The composition result.
        """

        results = list()
        exceptions = list()
        for composition_type in self.allowed_composition_types:
            try:
                ret = lhs.compose(rhs, composition_type)
                results.append((composition_type, ret))
            except CCGCompositionError as e:
                exceptions.append(e)

        if len(results) == 1:
            return CCGCompositionResult(*results[0])
        elif len(results) == 0:
            with get_ccg_composition_context().exc():
                fmt = f'Failed to compose CCGNodes {lhs} and {rhs}.\n'
                fmt += 'Detailed messages are:\n'
                for t, e in zip(self.allowed_composition_types, exceptions):
                    fmt += indent_text('Trying CCGCompositionType.{}:\n{}'.format(t.name, str(e))) + '\n'
                raise CCGCompositionError(fmt.rstrip())
        else:
            with get_ccg_composition_context().exc():
                fmt = f'Got ambiguous composition for CCGNodes {lhs} and {rhs}.\n'
                fmt += 'Candidates are:\n'
                for r in results:
                    fmt += indent_text('CCGCompositionType.' + str(r[0].name)) + '\n'
                raise CCGCompositionError(fmt.rstrip())

    def __str__(self) -> str:
        return f'CCGCompositionSystem({self.name})'

    __repr__ = __str__

    def format_summary(self) -> str:
        """Format the summary of the composition system."""
        fmt = 'Allowed composition types:\n'
        for type, weight in self.weights.items():
            fmt += '  CCGCompositionType.' + type.name + ': ' + str(weight) + '\n'
        fmt = 'CCGCompositionSystem: {}\n'.format(self.name) + indent_text(fmt.rstrip())
        return fmt

    def print_summary(self):
        """Print the summary of the composition system."""
        print(self.format_summary())

    @classmethod
    def make_default(cls) -> 'CCGCompositionSystem':
        """Make the default CCG composition system."""
        return cls.make_function_application()

    @classmethod
    def make_function_application(cls) -> 'CCGCompositionSystem':
        """Make the CCG composition system that only allows function application."""
        return cls('function_application', {
            CCGCompositionType.LEXICON: 0,
            CCGCompositionType.FORWARD_APPLICATION: 0,
            CCGCompositionType.BACKWARD_APPLICATION: 0,
            CCGCompositionType.NONE: 0
        })

    @classmethod
    def make_categorial_grammar(cls) -> 'CCGCompositionSystem':
        """Make the CCG composition system that allows function application and coordination (i.e., categorial grammar)."""
        return cls('categorial_grammar', {
            CCGCompositionType.LEXICON: 0,
            CCGCompositionType.FORWARD_APPLICATION: 0,
            CCGCompositionType.BACKWARD_APPLICATION: 0,
            CCGCompositionType.COORDINATION: 0,
            CCGCompositionType.NONE: 0,
        })


concepts/language/ccg/composition.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : grammar.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/05/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data structures and algorithms for lexicon, grammar and parsing."""

import copy
import inspect
import contextlib
import itertools
from collections import defaultdict
from typing import Any, Optional, Union, Iterable, Tuple, List, Dict, Callable
from jacinle.utils.printing import indent_text, print_to_string

from concepts.dsl.dsl_types import FormatContext
from concepts.dsl.dsl_functions import Function
from concepts.dsl.value import Value
from concepts.dsl.expression import ConstantExpression, FunctionApplicationExpression
from concepts.dsl.function_domain import FunctionDomain
from concepts.language.ccg.composition import (
    CCGCompositionType, CCGCompositionSystem,
    CCGCompositionError, get_ccg_composition_context,
    CCGCompositionResult, CCGCoordinationImmNode
)
from concepts.language.ccg.syntax import CCGSyntaxType, CCGSyntaxSystem
from concepts.language.ccg.semantics import CCGSemantics, CCGSemanticsSugar


__all__ = [
    'CCGParsingError',
    'Lexicon', 'LexiconUnion',
    'CCGNode', 'compose_ccg_nodes', 'CCG'
]

_profile = getattr(__builtins__, 'profile', lambda x: x)


class CCGParsingError(Exception):
    """An exception raised when parsing fails."""
    pass


class Lexicon(object):
    """The data structure for lexicon entries."""

    def __init__(self, syntax: Optional[CCGSyntaxType], semantics: Optional[CCGSemantics], weight: float = 0, extra: Optional[Any] = None):
        """Initialize the lexicon entry.

        Args:
            syntax: the syntax of the lexicon entry.
            semantics: the semantics of the lexicon entry.
            weight: the weight of the lexicon entry.
            extra: the extra information stored at the lexicon entry. If not None, this will be a tuple of (word, word_index, lexicon_entry_index).
        """
        self.syntax = syntax
        self.semantics = semantics
        self.weight = weight
        self.extra = extra

    syntax: Optional[CCGSyntaxType]
    """The syntax of the lexicon entry."""

    semantics: Optional[CCGSemantics]
    """The semantics of the lexicon entry."""

    weight: float
    """The weight of the lexicon entry."""

    extra: Optional[Any]
    """The extra information stored at the lexicon entry."""

    def __str__(self) -> str:
        fmt = type(self).__name__ + '['
        fmt += 'syntax=' + str(self.syntax) + ', '
        fmt += 'semantics=' + indent_text(
            str(self.semantics.value)
        ).lstrip() + ', '
        fmt += 'weight=' + str(self.weight) + ''
        fmt += ']'
        return fmt

    def __repr__(self):
        return str(self)


class _LexiconUnionType(object):
    def __init__(self, *annotations):
        self.annotations = annotations


class _LexiconUnionTyping(object):
    def __getitem__(self, item):
        if isinstance(item, tuple) and not isinstance(item[0], tuple):
            return _LexiconUnionType(item)
        return _LexiconUnionType(*item)


LexiconUnion = _LexiconUnionTyping()
"""A syntax sugar that allows users to define several lexicon entries for a single word.
This is useful when the word has multiple meanings (see :meth:`CCG.define` for details)."""


class CCGNode(object):
    """A node in the CCG parsing tree."""

    def __init__(
        self,
        composition_system: CCGCompositionSystem,
        syntax: Optional[CCGSyntaxType], semantics: Optional[CCGSemantics], composition_type: CCGCompositionType,
        lexicon: Optional[Lexicon] = None,
        lhs: Optional['CCGNode'] = None, rhs: Optional['CCGNode'] = None,
        weight: Optional[float] = None
    ):
        """Construct a CCG node.

        Args:
            composition_system: the composition system.
            syntax: the syntax type.
            semantics: the semantics.
            composition_type: the composition type.
            lexicon: the lexicon (if the composition type is LEXICON).
            lhs: the left child of the node (for other types of composition).
            rhs: the right child of the node (for other types of composition).
            weight: the weight. If not given, it will be computed automatically (lhs + rhs + composition weight).
        """

        self.composition_system = composition_system
        self.syntax = syntax
        self.semantics = semantics
        self.composition_type = composition_type

        self.lexicon = lexicon
        self.lhs = lhs
        self.rhs = rhs

        self.weight = weight
        if self.weight is None:
            self.weight = self._compute_weight()

    def _compute_weight(self):
        if self.composition_type is CCGCompositionType.LEXICON:
            return self.lexicon.weight

        return self.lhs.weight + self.rhs.weight + self.composition_system.weights[self.composition_type]

    composition_system: CCGCompositionSystem
    """The composition system."""

    syntax: Optional[CCGSyntaxType]
    """The syntax type of the node."""

    semantics: Optional[CCGSemantics]
    """The semantics of the node."""

    composition_type: CCGCompositionType
    """The composition type of the node."""

    lexicon: Optional[Lexicon]
    """The lexicon (if the composition type is LEXICON)."""

    lhs: Optional['CCGNode']
    """The left child of the node (for other types of composition)."""

    rhs: Optional['CCGNode']
    """The right child of the node (for other types of composition)."""

    weight: float
    """The weight of the node."""

    @_profile
    def compose(self, rhs: 'CCGNode', composition_type: Optional[CCGCompositionType] = None) -> Union[CCGCompositionResult, 'CCGNode']:
        """Compose the current node with another node. This function will automatically guess the composition type.

        Args:
            rhs: the right node.
            composition_type: the composition type. If not given, it will be guessed automatically.

        Returns:
            The composition result. When the composition type is not given, it will be a :class:`CCGCompositionResult` object.
        """
        if composition_type is not None:
            try:
                ctx = get_ccg_composition_context()
                self.compose_check(rhs, composition_type)  # throws CCGCompositionError
                new_syntax, new_semantics = None, None
                if ctx.syntax:
                    new_syntax = self.syntax.compose(rhs.syntax, composition_type)
                if ctx.semantics:
                    new_semantics = self.semantics.compose(rhs.semantics, composition_type)
                node = self.__class__(
                    self.composition_system, new_syntax, new_semantics, composition_type,
                    lhs=self, rhs=rhs
                )
                return node
            except CCGCompositionError as e:
                raise e
        else:
            results = list()
            exceptions = list()

            composition_types = self.compose_guess(rhs)
            if composition_types is None:
                composition_types = self.composition_system.allowed_composition_types

            for composition_type in composition_types:
                try:
                    results.append((composition_type, self.compose(rhs, composition_type)))
                except CCGCompositionError as e:
                    exceptions.append(e)

            if len(results) == 1:
                return CCGCompositionResult(*results[0])
            elif len(results) == 0:
                with get_ccg_composition_context().exc():
                    fmt = f'Failed to compose CCGNodes {self} and {rhs}.\n'
                    fmt += 'Detailed messages are:\n'
                    for t, e in zip(composition_types, exceptions):
                        fmt += indent_text('Trying CCGCompositionType.{}:\n{}'.format(t.name, str(e))) + '\n'
                    raise CCGCompositionError(fmt.rstrip())
            else:
                with get_ccg_composition_context().exc():
                    fmt = f'Got ambiguous composition for CCGNodes {self} and {rhs}.\n'
                    fmt += 'Candidates are:\n'
                    for r in results:
                        fmt += indent_text('CCGCompositionType.' + str(r[0].name)) + '\n'
                    raise CCGCompositionError(fmt.rstrip())

    def compose_check(self, rhs: 'CCGNode', composition_type: CCGCompositionType):
        """Check if the current node can be composed with another node. If the check fails, a CCGCompositionError will be
        raised.

        Args:
            rhs: the right node.
            composition_type: the composition type.

        Raises:
            CCGCompositionError: if the check fails.
        """
        if (
            isinstance(self.syntax, CCGCoordinationImmNode) or
            isinstance(self.semantics, CCGCoordinationImmNode)
        ):
            raise CCGCompositionError('Can not make non-coordination composition for CCGCoordinationImmResult.')
        if (
            isinstance(self.syntax, CCGCoordinationImmNode) or
            isinstance(self.semantics, CCGCoordinationImmNode) or
            isinstance(rhs.syntax, CCGCoordinationImmNode) or
            isinstance(rhs.semantics, CCGCoordinationImmNode)
        ):
            if composition_type is not CCGCompositionType.COORDINATION:
                raise CCGCompositionError('Can not make non-coordination composition for CCGCoordinationImmResult.')

    def compose_guess(self, rhs: 'CCGNode') -> Optional[Tuple[CCGCompositionType]]:
        """Guess the composition type of the current node with another node.

        Args:
            rhs: the right node.

        Returns:
            The guessed composition type. If None, all composition types are allowed.
        """
        return None

    def linearize_lexicons(self) -> List[Lexicon]:
        """Linearize the lexicons of the node. It return a list of lexicons corresponding to the leaves of parsing tree.

        Returns:
            The list of lexicons.
        """
        if self.lexicon is not None:
            return [self.lexicon]
        return self.lhs.linearize_lexicons() + self.rhs.linearize_lexicons()

    def as_nltk_str(self) -> str:
        """Convert the node to a string in nltk format."""
        if self.composition_type is CCGCompositionType.LEXICON:
            if self.lexicon.extra is not None:
                meaning = str(self.lexicon.extra[0])
            else:
                with FormatContext(function_format_lambda=True).as_default():
                    meaning = str(self.semantics.value).replace('(', '{').replace(')', '}')
            return f'({str(self.syntax)} {meaning})'

        if self.composition_type is CCGCompositionType.COORDINATION:
            return f'({str(self.syntax)} {self.lhs.as_nltk_str()} {self.rhs.lhs.as_nltk_str()} {self.rhs.rhs.as_nltk_str()})'
        return f'({str(self.syntax)} {self.lhs.as_nltk_str()} {self.rhs.as_nltk_str()})'

    def format_nltk_tree(self) -> str:
        """Format the node as a nltk tree."""
        with print_to_string() as fmt:
            self.print_nltk_tree()
        return fmt.get()

    def print_nltk_tree(self):
        """Print the node as a nltk tree."""
        from nltk.tree import Tree
        parsing_nltk = Tree.fromstring(self.as_nltk_str())
        parsing_nltk.pretty_print()

    def __str__(self) -> str:
        fmt = type(self).__name__ + '[\n'
        fmt += '  syntax   : ' + str(self.syntax) + '\n'
        with FormatContext(function_format_lambda=True).as_default():
            fmt += '  semantics: ' + indent_text(str(self.semantics), indent_format=' ' * 13).lstrip() + '\n'
        fmt += '  weight   : ' + str(self.weight) + '\n'
        fmt += ']'
        return fmt

    def __repr__(self) -> str:
        return str(self)


def compose_ccg_nodes(lhs: CCGNode, rhs: CCGNode, composition_type: Optional[CCGCompositionType] = None) -> CCGCompositionResult:
    """Compose two CCG nodes. If the composition type is not specified, it will be guessed.

    Args:
        lhs: the left node.
        rhs: the right node.
        composition_type: the composition type.

    Returns:
        The composition result.
    """
    return lhs.compose(rhs, composition_type)


class CCG(object):
    """The data structure of a CCG grammar and the implementation for parsing."""

    def __init__(self, domain: FunctionDomain, syntax_system: CCGSyntaxSystem, composition_system: Optional[CCGCompositionSystem] = None):
        """Initialize the CCG parser.

        Args:
            domain: the function domain.
            syntax_system: the syntax system, containing primitive and conjunctive syntax types.
            composition_system: the composition system, containing allowed composition rules.
        """
        self.function_domain = domain
        self.syntax_system = syntax_system
        self.semantics_sugar = CCGSemanticsSugar(self.function_domain)
        self.composition_system = composition_system
        self.lexicon_entries = defaultdict(list)

        if self.composition_system is None:
            self.composition_system = CCGCompositionSystem.make_default()

    function_domain: FunctionDomain
    """The function domain."""

    syntax_system: CCGSyntaxSystem
    """The syntax system, containing primitive and conjunctive syntax types."""

    semantics_sugar: CCGSemanticsSugar
    """The semantics sugar, which is a helper function to convert lambda functions to a :class:`concepts.language.ccg.semantics.CCGSemantics` instance."""

    composition_system: CCGCompositionSystem
    """The composition system, containing allowed composition rules."""

    lexicons: Dict[str, List[Lexicon]]
    """The lexicons entries, which is a dictionary from word to a list of :class:`Lexicon` instances."""

    def clone(self, deep: bool = False) -> 'CCG':
        """Clone the CCG grammar.

        Args:
            deep: whether to make a deep copy of the class. If False, only the lexicons will be copied.

        Returns:
            The cloned CCG grammar.
        """
        if deep:
            return copy.deepcopy(self)
        new_obj = self.__class__(self.function_domain, self.syntax_system, self.composition_system)
        new_obj.lexicon_entries = defaultdict(list)
        for k, v in self.lexicon_entries.items():
            new_obj.lexicon_entries[k].extend(v)
        return new_obj

    def make_node(self, arg1: Union[str, Lexicon, CCGNode], arg2: Optional[CCGNode] = None, *, composition_type: Optional[CCGCompositionType] = None) -> CCGNode:
        """Make a CCG node by trying to compose two nodes or retriving an lexicon entry. There are three use cases:

        - ``make_node(lexicon)``: make a node from a :class:`Lexicon` instance.
        - ``make_node(word)``: make a node from a word, which will be looked up in the lexicon. When there are multiple lexicon entries, an error will be raised.
        - ``make_node(lhs, rhs)``: make a node by composing two nodes.

        Args:
            arg1: the first argument.
            arg2: the second argument.
            composition_type: the composition type. If not specified, it will be guessed.

        Returns:
            The composed CCG node.
        """
        if isinstance(arg1, str):
            if arg1 in self.lexicon_entries:
                tot_entries = len(self.lexicon_entries[arg1])
                if tot_entries == 1:
                    return self.make_node(self.lexicon_entries[arg1][0])
                raise CCGParsingError('Ambiguous lexicon entry for word: "{}" (n = {}).'.format(arg1, tot_entries))
            raise CCGParsingError('Out-of-vocab word: {}.'.format(arg1))

        if isinstance(arg1, Lexicon):
            return CCGNode(self.composition_system, arg1.syntax, arg1.semantics, CCGCompositionType.LEXICON, lexicon=arg1)
        else:
            assert arg2 is not None
            return compose_ccg_nodes(arg1, arg2, composition_type=composition_type).result

    @property
    def Syntax(self):
        """A syntax sugar for defining syntax types. For example, use ``self.Syntax['S/NP']`` to define a syntax type."""
        return self.syntax_system

    @property
    def Semantics(self):
        """A syntax sugar for defining semantics. For example, use ``self.Semantics[lambda x: f(x)]`` to define a semantics."""
        return self.semantics_sugar

    def add_entry(self, word: str, lexicon: Lexicon):
        """Add a lexicon entry.

        Args:
            word: the word.
            lexicon: the lexicon entry.
        """
        self.lexicon_entries[word].append(lexicon)

    def add_entry_simple(
        self, word: str,
        syntax: Union[None, str, CCGSyntaxType],
        semantics: Union[None, Value, Function, ConstantExpression, FunctionApplicationExpression, Callable],
        weight: float = 0
    ):
        """Add a lexicon entry with a syntax type and a semantic form.

        Args:
            word: the word.
            syntax: the syntax type. It can be a string, a :class:`CCGSyntaxType` instance, or None.
            semantics: the semantics. It can be a value, a function, a constant expression, a function application expression, or a callable Python function.
            weight: the weight of the lexicon entry.
        """
        self.lexicon_entries[word].append(Lexicon(self.Syntax[syntax], self.Semantics[semantics], weight=weight))

    def clear_entries(self, word: str):
        """Clear all lexicon entries for a word.

        Args:
            word: the word.
        """
        self.lexicon_entries[word].clear()

    def update_entries(self, entries_dict: Dict[str, Iterable[Lexicon]]):
        """Update the lexicon entries.

        Args:
            entries_dict: the lexicon entries, which is a dictionary from word to a list of :class:`Lexicon` instances.
        """
        for word, entries in entries_dict.items():
            for entry in entries:
                self.add_entry(word, entry)

    @contextlib.contextmanager
    def define(self):
        """A helper function to define a CCG grammar. For example:

        .. code-block:: python

            with ccg.define():
                red: 'N/N', lambda x: filter(x, 'red')
                blue: 'N/N', lambda x: filter(x, 'blue'), 0.5  # weight = 0.5
        """
        locals_before = inspect.stack()[2][0].f_locals.copy()
        annotations_before = locals_before.get('__annotations__', dict()).copy()
        yield self
        locals_after = inspect.stack()[2][0].f_locals.copy()
        annotations_after = locals_after.get('__annotations__', dict()).copy()

        new_annotations = {
            k: v for k, v in annotations_after.items()
            if k not in annotations_before or annotations_after[k] != annotations_before[k]
        }

        if len(new_annotations) == 0:
            raise ValueError('ccg.define() is only allowed at the global scope.')

        def add_entry(word, annotation):
            assert isinstance(annotation, tuple) and len(annotation) in (2, 3)
            assert isinstance(annotation[0], CCGSyntaxType)
            assert isinstance(annotation[1], CCGSemantics)

            weight = 0
            if len(annotation) == 2:
                syntax, semantics = annotation
            else:
                syntax, semantics, weight = annotation

            self.add_entry(word, Lexicon(syntax, semantics, weight))

        for const_name, raw_annotation in new_annotations.items():
            if isinstance(raw_annotation, _LexiconUnionType):
                for a in raw_annotation.annotations:
                    add_entry(const_name, a)
            else:
                add_entry(const_name, raw_annotation)

    def parse(self, sentence: Union[str, Iterable[str]], beam: Optional[int] = None, preserve_syntax_types: bool = True) -> List[CCGNode]:
        """Parse a sentence.

        Args:
            sentence: the sentence to parse. It can be a string or a list of words.
            beam: the beam size. If not specified, it will perform the full chart parsing (beam size = Infinity).
            preserve_syntax_types: whether to preserve the syntax types in beam searcm.
                When it is False, the algorithm will sort all nodes by their weights and return the top-k nodes.
                When it is True, the algorithm will sort all nodes by their weights and return the top-k nodes for each syntax type.

        Returns:
            A list of :class:`CCGNode` instances.
        """
        if isinstance(sentence, str):
            sentence = sentence.split()

        length = len(sentence)
        dp = [[list() for _ in range(length + 1)] for _ in range(length)]
        for i, word in enumerate(sentence):
            if word not in self.lexicon_entries:
                raise CCGParsingError('Out-of-vocab word: {}.'.format(word))

            dp[i][i + 1] = [self.make_node(length) for length in self.lexicon_entries[word]]

        def merge(list1, list2):
            output_list = list()
            for node1, node2 in itertools.product(list1, list2):
                try:
                    node = compose_ccg_nodes(node1, node2).result
                    output_list.append(node)
                except CCGCompositionError:
                    pass
            return output_list

        for length in range(2, length + 1):
            for i in range(0, length + 1 - length):
                j = i + length
                for k in range(i + 1, j):
                    dp[i][j].extend(merge(dp[i][k], dp[k][j]))
                if beam is not None:
                    if preserve_syntax_types:
                        dp[i][j] = _filter_beam_per_type(dp[i][j], beam)
                    else:
                        dp[i][j] = sorted(dp[i][j], key=lambda x: x.weight, reverse=True)[:beam]

        return sorted(dp[0][length], key=lambda x: x.weight)

    def _format_lexicon_entries(self) -> str:
        fmt = 'Lexicon Entries:\n'
        # max_words_len = max([len(x) for x in self.lexicons])
        for word, lexicons in self.lexicon_entries.items():
            for lexicon in lexicons:
                this_fmt = f'{word}: ' + str(lexicon)
                fmt += indent_text(this_fmt) + '\n'
        return fmt

    def format_summary(self) -> str:
        """Format the summary of the CCG grammar."""
        fmt = 'Combinatory Categorial Grammar\n'
        fmt += indent_text(str(self.function_domain)) + '\n'
        fmt += indent_text(str(self.syntax_system)) + '\n'
        fmt += indent_text(str(self.composition_system)) + '\n'
        fmt += indent_text(self._format_lexicon_entries())
        return fmt

    def print_summary(self):
        """Print the summary of the CCG grammar."""
        print(self.format_summary())


def _filter_beam_per_type(nodes: List[CCGNode], beam: int) -> List[CCGNode]:
    all_nodes_by_type = defaultdict(list)
    for node in nodes:
        if isinstance(node.syntax, CCGCoordinationImmNode):
            typename = f'COORDIMM({node.syntax.conj.typename}, {node.syntax.rhs.typename})'
        else:
            typename = node.syntax.typename
        all_nodes_by_type[typename].append(node)
    for typename, nodes in all_nodes_by_type.items():
        all_nodes_by_type[typename] = sorted(nodes, key=lambda x: x.weight, reverse=True)[:beam]
    return list(itertools.chain.from_iterable(all_nodes_by_type.values()))


concepts/language/ccg/grammar.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : learning.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/06/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Algorithms for learning CCG entries from exmaples."""

import itertools
import collections
from typing import Any, Optional, Union, Iterable, Tuple, List, Callable
from dataclasses import dataclass
from jacinle.utils.printing import indent_text

from concepts.dsl.dsl_types import ConstantType
from concepts.dsl.dsl_functions import Function
from concepts.dsl.value import Value
from concepts.dsl.expression import ConstantExpression, FunctionApplicationExpression
from concepts.dsl.executors.executor_base import DSLExecutionError
from concepts.language.ccg.composition import CCGCompositionContext, CCGCompositionError
from concepts.language.ccg.semantics import CCGSemantics
from concepts.language.ccg.grammar import CCG, Lexicon, CCGNode
from concepts.language.ccg.search import CCGSyntaxSearcherBase, CCGSyntaxEnumerativeSearcher, CCGSemanticsSearcherBase, CCGSemanticsEnumerativeSearcher
from concepts.language.ccg.search import CCGSyntaxSearchResult, CCGSemanticsSearchResult

__all__ = ['CCGLearningResult', 'by_parsing', 'by_parsing_with_lexicon_gen', 'by_grounding']


@dataclass
class CCGLearningResult(object):
    """The result of CCG learning."""

    words: Tuple[str, ...]
    """The list of words learned by the algorithm."""

    lexicons: Tuple[Lexicon, ...]
    """The list of lexicons learned by the algorithm. It is a tuple of lexicons that has the same length as `words`."""

    parsing_results: Tuple[CCGNode, ...]
    """The parsing results of the entire sentence based on the learned lexicons."""

    def format_summary(self):
        fmt = 'Learned lexicons:\n'
        for w, l in zip(self.words, self.lexicons):
            fmt += indent_text(w + ': ' + str(l)) + '\n'
        fmt += 'Parsing result:\n'
        for i, r in enumerate(self.parsing_results):
            fmt += f'{i}: ' + indent_text(str(r), indent_format='    ').lstrip() + '\n'
        return fmt

    def print_summary(self):
        print(self.format_summary())


def by_parsing(
    ccg: CCG, sentence: Union[str, Iterable[str]], *,
    novel_words: Optional[tuple] = None,
    candidate_syntax_types: Optional[List[CCGSyntaxSearchResult]] = None,
    syntax_searcher: Optional[CCGSyntaxSearcherBase] = None,
    syntax_searcher_kwargs: Optional[dict] = None,
    candidate_semantics: Optional[List[CCGSemanticsSearchResult]] = None,
    semantics_searcher: Optional[CCGSemanticsSearcherBase] = None,
    semantics_searcher_kwargs: Optional[dict] = None,
    bind_concepts: bool = True
) -> List[CCGLearningResult]:
    """Learn CCG lexicon entries from a sentence by trying to parse the sentence.

    Args:
        ccg: the CCG grammar.
        sentence: the sentence to be parsed.
        novel_words: the list of novel words to be learned. If not specified, the algorithm will detect all novel words in the sentence.
        candidate_syntax_types: the list of candidate syntax types to be used for parsing.
            If not specified, the algorithm will use the syntax searcher to generate the candidate syntax types.
        syntax_searcher: the syntax searcher to be used for generating candidate syntax types.
            If not specified, the algorithm will use the enumerative searcher.
        syntax_searcher_kwargs: the keyword arguments for the syntax searcher.
        candidate_semantics: the list of candidate semantics to be used for parsing.
            If not specified, the algorithm will use the semantics searcher to generate the candidate semantics.
        semantics_searcher: the semantics searcher to be used for generating candidate semantics.
            If not specified, the algorithm will use the enumerative searcher.
        semantics_searcher_kwargs: the keyword arguments for the semantics searcher.
        bind_concepts: whether to bind concepts in the semantics. This will allow algorithm to invent novel concepts while learning.

    Returns:
        The result of the learning process, as a list of :class:`CCGLearningResult`.
    """

    if isinstance(sentence, str):
        sentence = sentence.split()

    if candidate_syntax_types is None:
        if syntax_searcher is None:
            syntax_searcher = CCGSyntaxEnumerativeSearcher(ccg.syntax_system)
        if syntax_searcher_kwargs is None:
            syntax_searcher_kwargs = dict()
        candidate_syntax_types = syntax_searcher.gen(**syntax_searcher_kwargs)

    if candidate_semantics is None:
        if semantics_searcher is None:
            semantics_searcher = CCGSemanticsEnumerativeSearcher(ccg.function_domain)
        if semantics_searcher_kwargs is None:
            semantics_searcher_kwargs = dict()
        candidate_semantics = semantics_searcher.gen(**semantics_searcher_kwargs)

    semantics_search_results_arity = collections.defaultdict(list)
    for r in candidate_semantics:
        if not bind_concepts and r.nr_constant_arguments > 0:
            continue  # if we do not allow binding new concepts, we skip those semantics with constant arguments.
        semantics_search_results_arity[r.nr_variable_arguments].append(r.semantics)

    if novel_words is None:
        novel_words = set()
        for word in sentence:
            if word not in ccg.lexicon_entries or len(ccg.lexicon_entries[word]) == 0:
                novel_words.add(word)
        novel_words = tuple(novel_words)

    new_ccg = ccg.clone()

    success_syntax_types = list()
    for syntax_types in itertools.product(candidate_syntax_types, repeat=len(novel_words)):
        for word, syntax in zip(novel_words, syntax_types):
            new_ccg.add_entry(word, Lexicon(syntax.syntax, None))

        try:
            with CCGCompositionContext(semantics=False, exc_verbose=False).as_default():
                parsing_results = new_ccg.parse(sentence)  # TODO: maybe cache the results.
            if len(parsing_results) > 0:
                # TODO: add a flag to control this check.
                found_success = False
                for r in parsing_results:
                    if r.syntax.typename in syntax_searcher._starting_symbols:
                        found_success = True
                        break
                if found_success:
                    success_syntax_types.append(syntax_types)
        except CCGCompositionError:
            pass

        for word in novel_words:
            new_ccg.clear_entries(word)

    success_results = list()
    for syntax_types in success_syntax_types:
        for semantics_list in itertools.product(*[
            semantics_search_results_arity[syntax.syntax.arity]
            for syntax in syntax_types
        ]):
            this_lexicons = list()
            for word, syntax, semantics in zip(novel_words, syntax_types, semantics_list):
                syntax, semantics = syntax.syntax, semantics
                if bind_concepts:
                    semantics = _bind_concepts(semantics, word)
                lexicon = Lexicon(syntax, semantics)
                new_ccg.add_entry(word, lexicon)
                this_lexicons.append(lexicon)

            try:
                with CCGCompositionContext(exc_verbose=True).as_default():
                    parsing_results = new_ccg.parse(sentence)
                if len(parsing_results) > 0:
                    success_results.append(CCGLearningResult(novel_words, tuple(this_lexicons), parsing_results))
            except CCGCompositionError:
                pass

            for word in novel_words:
                new_ccg.clear_entries(word)

    return success_results


def _bind_concepts(semantics: CCGSemantics, word: str):
    if semantics.is_function:
        mapping = dict()
        for i, argument in enumerate(semantics.value.arguments):
            if isinstance(argument.dtype, ConstantType):
                mapping[f'#{i}'] = Value(argument.dtype, word)
        if len(mapping) > 0:
            semantics = CCGSemantics(semantics.value.partial(**mapping, execute_fully_bound_functions=True), is_conj=semantics.is_conj)
    return semantics


def by_parsing_with_lexicon_gen(
    ccg: CCG, sentence: Union[str, list, tuple],
    lexicon_generator: Callable[[str], Iterable[Lexicon]],
    novel_words: Optional[tuple] = None
):
    """Learn CCG lexicon entries from a sentence by trying to parse the sentence.
    Unlike :func:`by_parsing`, this function takes a lexicon generator instead of a list of candidate syntax types and semantics.

    Args:
        ccg: the CCG grammar.
        sentence: the sentence to be parsed.
        lexicon_generator: the lexicon generator to be used for generating candidate lexicons. It takes a word as input and returns a list of candidate lexicons.
        novel_words: the list of novel words to be learned. If not specified, the algorithm will detect all novel words in the sentence.

    Returns:
        The result of the learning process, as a list of :class:`CCGLearningResult`.
    """

    if isinstance(sentence, str):
        sentence = sentence.split()

    if novel_words is None:
        novel_words = set()
        for word in sentence:
            if word not in ccg.lexicon_entries or len(ccg.lexicon_entries[word]) == 0:
                novel_words.add(word)
    novel_words = tuple(novel_words)

    new_ccg = ccg.clone()
    candidate_lexicons = [lexicon_generator(word) for word in novel_words]

    success_results = list()
    # TODO(Jiayuan Mao @ 07/11): optimize by first filtering syntax types.
    for new_lexicon in itertools.product(*candidate_lexicons):
        for word, entry in zip(novel_words, new_lexicon):
            new_ccg.add_entry(word, entry)

        try:
            with CCGCompositionContext(exc_verbose=False).as_default():
                parsing_results = new_ccg.parse(sentence)
            if len(parsing_results) > 0:
                success_results.append(CCGLearningResult(novel_words, tuple(new_lexicon), parsing_results))
        except CCGCompositionError:
            pass

        for word in novel_words:
            new_ccg.clear_entries(word)

    return success_results


def by_grounding(
    by_parsing_learning_func: Callable,
    ccg: CCG, sentence: Union[str, Iterable[str]],
    executor: Callable[[Union[Function, ConstantExpression, FunctionApplicationExpression]], Any],
    criterion: Callable[[Any], bool],
    **by_parsing_kwargs
) -> List[CCGLearningResult]:
    """Learn a CCG lexicon from a sentence by both parsing the sentence and grounding the parsing result with a given executor.

    Args:
        by_parsing_learning_func: the function to be used for learning a CCG lexicon from a sentence by parsing the sentence.
            Can be either :func:`by_parsing` or :func:`by_parsing_with_lexicon_gen`.
        ccg: the CCG grammar.
        sentence: the sentence to be parsed.
        executor: the executor used to ground the parsing result.
        criterion: the criterion used to classify whether the execution result is correct.
        **by_parsing_kwargs: the keyword arguments to be passed to the ``by_parsing_learning_func``.

    Returns:
        The result of the learning process, as a list of :class:`CCGLearningResult`.
    """
    by_parsing_results = by_parsing_learning_func(ccg, sentence, **by_parsing_kwargs)

    success_results = list()
    for r in by_parsing_results:
        ccg_node = r.parsing_results[0]
        try:
            result = executor(ccg_node.expression.value)
            if criterion(result):
                success_results.append(r)
        except DSLExecutionError:
            pass
    return success_results


def auto_research_novel_words(
    base_learning_func: Callable,
    ccg: CCG,
    sentence: Union[str, Iterable[str]],
    max_research: int,
    **kwargs
) -> List[CCGLearningResult]:
    """A helper function that automatically perform re-search for known words in a sentence.

    Args:
        base_learning_func: the base learning function to be used for learning a CCG lexicon from a sentence.
            Can be :func:`by_parsing`, or :func:`by_parsing_with_lexicon_gen`, or :func:`by_grounding`.
        ccg: the CCG grammar.
        sentence: the sentence to be parsed.
        max_research: the maximum number of words whose lexicon will be re-searched.
        **kwargs: additional keyword arguments to be passed to the ``base_learning_func``.
    """
    if isinstance(sentence, str):
        sentence = sentence.split()

    novel_words = set()
    known_words = set()
    for word in sentence:
        if word not in ccg.lexicon_entries or len(ccg.lexicon_entries[word]) == 0:
            novel_words.add(word)
        else:
            known_words.add(word)
    novel_words_tuple = tuple(novel_words)
    known_words_tuple = tuple(known_words)

    rv = base_learning_func(ccg, sentence, **kwargs)
    if len(rv) > 0:
        return rv

    kwargs = kwargs.copy()
    for nr_research in range(1, max_research + 1):
        rv_accumulate = list()
        for comb in itertools.combinations(known_words_tuple, nr_research):
            new_novel_words = novel_words_tuple + tuple(comb)
            kwargs['novel_words'] = new_novel_words
            rv = base_learning_func(ccg, sentence, **kwargs)
            rv_accumulate.extend(rv)
        if len(rv_accumulate) > 0:
            return rv_accumulate

    return list()


concepts/language/ccg/learning.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : search.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/05/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Algorithms for enumerate possible syntax types and semantic forms in a domain."""

import itertools
from typing import Any, Optional, Union, Iterable, Tuple, List, Callable
from dataclasses import dataclass
from jacinle.utils.defaults import default_args, ARGDEF

from concepts.dsl.dsl_functions import FunctionType, Function
from concepts.dsl.expression import ConstantExpression, FunctionApplicationExpression
from concepts.dsl.function_domain import FunctionDomain
from concepts.dsl.learning.function_domain_search import FunctionDomainExpressionEnumerativeSearcher, stat_function
from concepts.language.ccg.syntax import CCGSyntaxSystem, CCGPrimitiveSyntaxType, CCGSyntaxType
from concepts.language.ccg.semantics import CCGSemantics

__all__ = [
    'CCGSyntaxSearchResult', 'CCGSyntaxSearcherBase', 'CCGSyntaxEnumerativeSearcher', 'gen_syntax_search_result_from_syntax_types',
    'CCGSemanticsSearchResult', 'CCGSemanticsSearcherBase', 'CCGSemanticsEnumerativeSearcher', 'gen_semantics_search_result_from_functions',
]


@dataclass
class CCGSyntaxSearchResult(object):
    """Search result for :class:`CCGSyntaxSearcherBase`."""

    syntax: CCGSyntaxType
    """The candidate syntax type."""

    depth: int
    """The depth of the candidate syntax type."""


class CCGSyntaxSearcherBase(object):
    """Base class for syntax searchers."""

    def __init__(self, syntax_system: CCGSyntaxSystem):
        """Initialize the searcher.

        Args:
            syntax_system: the syntax system, containing all primitive and conjunction syntax types.
        """
        self._syntax_system = syntax_system

    @property
    def syntax_system(self):
        """The syntax system."""
        return self._syntax_system

    def gen(self) -> List[CCGSyntaxSearchResult]:
        """Generate a list of candidate syntax types.

        Returns:
            A list of candidate syntax types.
        """
        raise NotImplementedError()


class CCGSyntaxEnumerativeSearcher(CCGSyntaxSearcherBase):
    """Enumerative searcher for syntax types."""

    def __init__(self, syntax_system: CCGSyntaxSystem, starting_symbols: Iterable[str] = ('S', )):
        """Initialize the searcher.

        Args:
            syntax_system: the syntax system, containing all primitive and conjunction syntax types.
            starting_symbols: the root primitive symbol for candidate syntax types.
        """
        super().__init__(syntax_system)
        self._starting_symbols = tuple(starting_symbols)

        for s in self._starting_symbols:
            assert isinstance(self._syntax_system[s], CCGPrimitiveSyntaxType)

    def gen(self, max_depth: int = 3, allow_functor_type: bool = False) -> List[CCGSyntaxSearchResult]:
        """Generate a list of candidate syntax types.

        Args:
            max_depth: The maximum depth of the syntax tree.
            allow_functor_type: Whether to allow functor types during composition. When this is set to False, the
                function will not generate any syntax types that contains functor-typed arguments. For example S/(NP/NP).

        Returns:
            A list of candidate syntax types.
        """
        current = {i: list() for i in range(1, max_depth + 1)}
        current_typenames = set()

        def add(depth, syntax):
            if syntax.typename not in current_typenames:
                current[depth].append(syntax)
                current_typenames.add(syntax.typename)

        for symbol in self._syntax_system.types.values():
            add(1, symbol)

        for depth in range(2, max_depth + 1):
            if allow_functor_type:
                for depth1 in range(1, depth):
                    for depth2 in range(1, depth - depth1 + 1):
                        for syntax1, syntax2 in itertools.product(current[depth1], current[depth2]):
                            if syntax2.typename not in self._starting_symbols:
                                add(depth, syntax1 / syntax2)
                                add(depth, syntax1 // syntax2)
            else:
                depth1 = depth - 1
                depth2 = depth - depth1
                for syntax1, syntax2 in itertools.product(current[depth1], current[depth2]):
                    if syntax2.typename not in self._starting_symbols:
                        add(depth, syntax1 / syntax2)
                        add(depth, syntax1 // syntax2)

        results = list()
        for k, vs in current.items():
            results.extend([CCGSyntaxSearchResult(v, k) for v in vs])
        return results


def gen_syntax_search_result_from_syntax_types(syntax_system: CCGSyntaxSystem, syntax_types: Iterable[CCGSyntaxType]) -> List[CCGSyntaxSearchResult]:
    """Generate a list of syntax search results from a list of syntax types.

    Args:
        syntax_system: the syntax system, containing all primitive and conjunction syntax types.
        syntax_types: the syntax types.

    Returns:
        A list of syntax search results.
    """
    results = list()
    for syntax_type in syntax_types:
        results.append(CCGSyntaxSearchResult(syntax_type, 0))
    return results


@dataclass
class CCGSemanticsSearchResult(object):
    """Search result for :class:`CCGSemanticsSearcherBase`."""

    semantics: CCGSemantics
    """The candidate semantic form."""

    depth: int
    """The depth of the candidate semantics."""

    nr_constant_arguments: int
    """The number of constant arguments in the semantic form."""

    nr_variable_arguments: int
    """The number of variable arguments in the semantic form."""

    nr_function_arguments: int
    """The number of function arguments in the semantic form."""


class CCGSemanticsSearcherBase(object):
    """Base class for semantics searchers."""

    def gen(self) -> List[CCGSemanticsSearchResult]:
        """Generate a list of candidate semantic forms.

        Returns:
            A list of candidate semantic forms.
        """
        raise NotImplementedError()


_Types = FunctionDomain.AllowedTypes


class CCGSemanticsEnumerativeSearcher(CCGSemanticsSearcherBase):
    """Enumerative searcher for semantics."""

    def __init__(self, domain: FunctionDomain):
        """Initialize the searcher.

        Args:
            domain: the domain of the semantics.
        """
        self._domain = domain
        self._enumerator = FunctionDomainExpressionEnumerativeSearcher(domain)

    def gen(
        self,
        max_depth: int = ARGDEF,
        return_type: Optional[Union[_Types, Tuple[_Types, ...], List[_Types]]] = None,
        max_variable_arguments: int = ARGDEF,
        max_constant_arguments: int = ARGDEF,
        max_function_arguments: int = ARGDEF,
        search_constants: bool = ARGDEF,
        hash_function: Callable[[Union[Function, FunctionApplicationExpression]], Any] = None,
        verbose: bool = False
    ) -> List[CCGSemanticsSearchResult]:
        """Generate a list of candidate semantic forms.

        Args:
            max_depth: the maximum depth of the semantics tree.
            return_type: the return type of the semantics.
            max_variable_arguments: the maximum number of variables in the semantics.
            max_constant_arguments: the maximum number of constants in the semantics.
            max_function_arguments: the maximum number of functions in the semantics.
            search_constants: whether to search for constants when generating semantic forms.
            hash_function: an optional hash function that will be used to filter out duplicate functions.
            verbose: whether to print out the progress.

        Returns:
            A list of candidate semantic forms.
        """
        return (
            (self.gen_constant_semantics(return_type) if search_constants else list()) +
            self.gen_function_semantics(
                return_type,
                max_depth=max_depth,
                max_variable_arguments=max_variable_arguments,
                max_constant_arguments=max_constant_arguments,
                max_function_arguments=max_function_arguments,
                search_constants=search_constants,
                hash_function=hash_function,
                verbose=verbose
            )
        )

    def gen_constant_semantics(
        self,
        return_type: Optional[Union[_Types, Tuple[_Types, ...], List[_Types]]] = None
    ) -> List[CCGSemanticsSearchResult]:
        """Generate a list of candidate semantic forms that are simply constants.

        Args:
            return_type: the return type of the semantics.

        Returns:
            A list of candidate semantic forms.
        """
        return [CCGSemanticsSearchResult(
            CCGSemantics(c.expression),
            depth=0, nr_constant_arguments=1, nr_variable_arguments=0, nr_function_arguments=0
        ) for c in self._enumerator.gen_constant_expressions(return_type)]

    @default_args
    def gen_function_semantics(
        self,
        return_type: Optional[Union[_Types, Tuple[_Types, ...], List[_Types]]] = None, *,
        max_depth: int = 3,
        max_variable_arguments: int = 2,
        max_constant_arguments: int = 1,
        max_function_arguments: int = 0,
        search_constants: bool = False,
        hash_function: Callable[[Union[Function, FunctionApplicationExpression]], Any] = None,
        verbose: bool = False
    ) -> List[CCGSemanticsSearchResult]:
        return [CCGSemanticsSearchResult(
            CCGSemantics(f.expression),
            depth=f.depth,
            nr_constant_arguments=f.nr_constant_arguments,
            nr_variable_arguments=f.nr_variable_arguments,
            nr_function_arguments=f.nr_function_arguments
        ) for f in self._enumerator.gen_function_application_expressions(
            return_type=return_type,
            max_depth=max_depth,
            max_variable_arguments=max_variable_arguments,
            max_constant_arguments=max_constant_arguments,
            max_function_arguments=max_function_arguments,
            search_constants=search_constants,
            hash_function=hash_function,
            verbose=verbose,
        )]

    def _gen_function_primitives(self, ret_type, max_function_arguments):
        def function_call(func, *args):
            return func(*args)

        def gen():
            types = tuple(self._domain.types.values())
            for repeat in range(1, max_function_arguments + 1):
                for arg_types in itertools.product(types, repeat=repeat):
                    yield Function(
                        '__lambda__',
                        FunctionType(
                            [FunctionType(arg_types, ret_type), ] + list(arg_types),
                            ret_type
                        ),
                        overridden_call=function_call,
                    )

        return tuple(gen())

    @staticmethod
    def gen_lambda(f1: Function, arg_index=None, f2: Optional[Function] = None):
        if arg_index is None:
            return Function(
                '__lambda__',
                FunctionType(f1.ftype.argument_types, f1.ftype.return_type), overridden_call=f1
            )
        else:
            f1_arg_types = f1.ftype.argument_types
            f2_arg_types = f2.ftype.argument_types
            arg_types = f2_arg_types + f1_arg_types[:arg_index] + f1_arg_types[arg_index + 1:]

            def new_function_call(*args):
                f2_args = args[:f2.nr_arguments]
                f1_args = list(args[f2.nr_arguments:])

                f2_ret = f2(*f2_args)
                f1_args.insert(arg_index, f2_ret)
                return f1(*f1_args)

            return Function(
                '__lambda__',
                FunctionType(arg_types, f1.ftype.return_type),
                overridden_call=new_function_call
            )


def gen_semantics_search_result_from_functions(functions: Iterable[Union[Function, FunctionApplicationExpression, ConstantExpression]]) -> List[CCGSemanticsSearchResult]:
    """
    Generate a list of semantics search results from functions.

    Args:
        functions: the functions.

    Returns:
        A list of :class:`CCGSemanticsSearchResult` instances.
    """
    results = list()
    for f in functions:
        if isinstance(f, Function):
            stat = stat_function(f)
            results.append(CCGSemanticsSearchResult(CCGSemantics(f), 1, nr_constant_arguments=stat.nr_constant_arguments, nr_variable_arguments=stat.nr_variable_arguments, nr_function_arguments=stat.nr_function_arguments))
        else:
            results.append(CCGSemanticsSearchResult(CCGSemantics(f), 1, nr_constant_arguments=0, nr_variable_arguments=0, nr_function_arguments=0))
    return results


concepts/language/ccg/search.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : semantics.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/04/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data structures for semantic forms in a linguistic CCG.

The basic class is :class:`CCGSemantics`, which is a wrapper of a semantic form (a functor or a value).
There is a
"""

from typing import Optional, Union, Tuple, Dict, Callable
from dataclasses import dataclass
from jacinle.utils.cache import cached_property
from jacinle.utils.printing import indent_text

from concepts.dsl.dsl_functions import Function, FunctionArgumentResolutionContext, FunctionArgumentResolutionError
from concepts.dsl.function_domain import FunctionDomain
from concepts.dsl.value import Value
from concepts.dsl.expression import ConstantExpression, FunctionApplicationExpression
from concepts.language.ccg.composition import CCGCompositionType, CCGComposable, CCGCompositionError, get_ccg_composition_context

__all__ = [
    'CCGSemanticsCompositionError',
    'CCGSemanticsConjFunction', 'CCGSemanticsSimpleConjFunction',
    'CCGSemanticsLazyValue', 'CCGSemantics',
    'CCGSemanticsSugar',
]


class CCGSemanticsCompositionError(CCGCompositionError):
    """Raised when the semantics composition fails."""

    def __init__(self, message=None, lhs=None, rhs=None, error=None, conj=None):
        """Initialize the error.

        Args:
            message: the error message.
            lhs: the lhs semantics.
            rhs: the rhs semantics.
            error: the error raised by the left/right semantics.
            conj: the conjunction.
        """

        if message is None:
            super().__init__(None)
        else:
            message += '\nLeft:' + indent_text(str(lhs)) + '\nRight:' + indent_text(str(rhs))
            if conj is not None:
                message += '\nConj: ' + indent_text(str(conj))
            message += '\nOriginal error message: ' + indent_text(str(error)).lstrip()
            super().__init__('(Semantics) ' + message)


@dataclass
class CCGSemanticsConjFunction(object):
    """CCGSemanticsConjFunction is a wrapper that represents the semantics form of a conjunction term (e.g., AND)."""

    impl: Callable[[Union[Value, Function], Union[Value, Function]], Union[Value, Function]]
    """The underlying implementation of the conjunction."""

    def __call__(self, lhs: Union[Value, Function], rhs: Union[Value, Function]) -> Union[Value, Function]:
        """Perform the conjunction.

        Args:
            lhs: the left-hand side semantics.
            rhs: the right-hand side semantics.

        Returns:
            the conjunction result.
        """
        return self.impl(lhs, rhs)


class CCGSemanticsSimpleConjFunction(CCGSemanticsConjFunction):
    """A simple implementation for :class:`CCGSemanticsConjValue`.

    This function takes a function that works for :class:`~concepts.dsl.value.Value` and automatically
    converts it to a function that works for :class:`~concepts.dsl.function.Function`.
    """
    impl: Callable[[Union[Value, Function], Union[Value, Function]], Union[Value, Function]]

    def __call__(self, lhs: Union[Value, Function], rhs: Union[Value, Function]) -> Union[Value, Function]:
        """Perform the conjunction.

        Args:
            lhs: the left-hand side semantics.
            rhs: the right-hand side semantics.

        Returns:
            the conjunction result.
        """
        if isinstance(lhs, Function) and isinstance(rhs, Function):
            def body(*args, **kwargs):
                return self.impl(lhs(*args, **kwargs), rhs(*args, **kwargs))
            return Function('__conj__', lhs.ftype, overridden_call=body)

        if isinstance(lhs, Value) and isinstance(rhs, Value):
            return self.impl(lhs, rhs)

        raise CCGSemanticsCompositionError(f'Cannot perform conjunction between {lhs} and {rhs}.')


@dataclass
class CCGSemanticsLazyValue(object):
    """A wrapper that represents the semantic form of a node that is not yet evaluated.
    Specifically, this function stores the composition type and the individual components."""

    composition_type: Optional[CCGCompositionType] = None
    """The composition type."""

    lhs: Optional[Union['CCGSemantics', 'CCGSemanticsLazyValue']] = None
    """The left-hand side semantics."""

    rhs: Optional[Union['CCGSemantics', 'CCGSemanticsLazyValue']] = None
    """The right-hand side semantics."""

    conj: Optional[Callable] = None
    """The conjunction function (optional)."""

    def execute(self) -> Union[Value, Function]:
        """Execute the lazy value recursively and return the semantic form."""
        lhs, rhs = self.lhs, self.rhs
        if isinstance(lhs, CCGSemanticsLazyValue):
            lhs = lhs.execute()
        if isinstance(rhs, CCGSemanticsLazyValue):
            rhs = rhs.execute()

        if self.composition_type in (CCGCompositionType.FORWARD_APPLICATION, CCGCompositionType.BACKWARD_APPLICATION):
            if self.composition_type is CCGCompositionType.FORWARD_APPLICATION:
                return _forward_application(lhs, rhs)
            else:
                return _backward_application(lhs, rhs)
        elif self.composition_type is CCGCompositionType.COORDINATION:
            return _coordination(lhs, self.conj, rhs)
        else:
            raise NotImplementedError(f'Unknown composition type: {self.composition_type}.')


class CCGSemantics(CCGComposable):
    """CCGSemantics is a wrapper of a semantic form (a functor or a value)."""

    def __init__(self, value: Union[None, Callable, Function, ConstantExpression, FunctionApplicationExpression, CCGSemanticsLazyValue], *, is_conj: bool = False):
        self._value = value
        self._is_conj = is_conj

        if isinstance(self._value, Function) and self._value.nr_arguments == 0:
            self._value = self._value()

    @property
    def value(self) -> Union[Callable, Function, ConstantExpression, FunctionApplicationExpression, CCGSemanticsLazyValue]:
        """The semantic form."""
        return self._value

    @property
    def is_none(self):
        """Whether the semantics is None."""
        return self.value is None

    @property
    def is_conj(self):
        """Whether the semantics is a conjunction."""
        return self._is_conj

    @property
    def is_py_function(self):
        """Whether the semantics is a Python function."""
        return callable(self.value) and not self.is_function

    @property
    def is_lazy(self):
        """Whether the semantics is a lazy value."""
        return isinstance(self.value, CCGSemanticsLazyValue)

    @property
    def is_function(self):
        """Whether the semantics is a function."""
        if self.is_lazy:
            raise ValueError('Cannot check is_function for CCGSemanticsLazyValue')
        return isinstance(self.value, Function)

    @property
    def is_constant(self):
        """Whether the semantics is a constant."""
        if self.is_lazy:
            raise ValueError('Cannot check is_value for CCGSemanticsLazyValue')
        return isinstance(self.value, ConstantExpression)

    @property
    def is_function_application(self):
        """Whether the semantics is a function application expression."""
        if self.is_lazy:
            raise ValueError('Cannot check is_function_application for CCGSemanticsLazyValue')
        return isinstance(self.value, FunctionApplicationExpression)

    @property
    def is_value(self):
        """Whether the semantics is a value (either a constant or a function application expression)."""
        return self.is_constant or self.is_function_application

    @property
    def flags(self) -> Dict[str, bool]:
        """A set of flags that indicates the type of the semantics."""
        return {
            'is_none': self.is_none,
            'is_conj': self.is_conj,
            'is_py_function': self.is_py_function,
            'is_lazy': self.is_lazy,
            'is_function': self.is_function if not self.is_lazy else None,
            'is_constant': self.is_constant if not self.is_lazy else None,
            'is_function_application': self.is_function_application if not self.is_lazy else None,
            'is_value': self.is_value if not self.is_lazy else None,
        }

    @property
    def return_type(self):
        """The return type of the semantics. If the semantics is a function or a value, the return type is the type of the
        function or the value. Otherwise, this function will raise an error."""
        if self.is_value:
            return self.value.return_type
        elif self.is_function:
            return self.value.ftype.return_type
        else:
            raise AttributeError('Cannot get the return type of None, PyFunction, or Lazy semantics.')

    @cached_property
    def arity(self):
        """The arity of the semantics. If the semantics is a value, this function will return 0. If the semantics
        is a function, this function will return the arity of the function. Otherwise, this function will raise an error."""
        if self.is_value:
            return 0
        elif self.is_function:
            return self.value.ftype.nr_variable_arguments
        else:
            raise AttributeError('Cannot get the arity of None, PyFunction, or Lazy semantics.')

    def __str__(self) -> str:
        if self.value is None:
            return type(self).__name__ + '[None]'
        if self.is_conj:
            return type(self).__name__ + '[' + str(self.value) + ', CONJ]'
        return type(self).__name__ + '[' + str(self.value) + ']'

    def __repr__(self) -> str:
        return str(self)

    @cached_property
    def hash(self):  # for set/dict indexing.
        return hash(str(self))

    def _fapp(self, rhs: 'CCGSemantics') -> 'CCGSemantics':
        if get_ccg_composition_context().semantics_lazy_composition:
            return type(self)(CCGSemanticsLazyValue(CCGCompositionType.FORWARD_APPLICATION, self.value, rhs.value))
        else:
            return type(self)(_forward_application(self.value, rhs.value))

    def _bapp(self, lhs: 'CCGSemantics') -> 'CCGSemantics':
        if get_ccg_composition_context().semantics_lazy_composition:
            return type(self)(CCGSemanticsLazyValue(CCGCompositionType.BACKWARD_APPLICATION, lhs.value, self.value))
        else:
            return type(self)(_backward_application(lhs.value, self.value))

    def _coord3(self, lhs: 'CCGSemantics', rhs: 'CCGSemantics') -> 'CCGSemantics':
        if get_ccg_composition_context().semantics_lazy_composition:
            return type(self)(CCGSemanticsLazyValue(CCGCompositionType.COORDINATION, lhs.value, rhs.value, conj=self.value))
        else:
            return type(lhs)(_coordination(lhs.value, self.value, rhs.value))


def _forward_application(lhs, rhs):
    ctx = get_ccg_composition_context()

    if isinstance(lhs, Function):
        try:
            with FunctionArgumentResolutionContext(exc_verbose=ctx.exc_verbose).as_default():
                return lhs.partial(rhs, execute_fully_bound_functions=True)
        except FunctionArgumentResolutionError as e:
            with ctx.exc(CCGSemanticsCompositionError, e):
                raise CCGSemanticsCompositionError('Cannot make forward application.', lhs, rhs, e) from e
    with ctx.exc(CCGSemanticsCompositionError):
        raise CCGSemanticsCompositionError('Cannot make forward application.', lhs, rhs, 'Functor/Value types do not match.')


def _backward_application(lhs, rhs):
    ctx = get_ccg_composition_context()

    if isinstance(rhs, Function):
        try:
            with FunctionArgumentResolutionContext(exc_verbose=ctx.exc_verbose).as_default():
                return rhs.partial(lhs, execute_fully_bound_functions=True)
        except FunctionArgumentResolutionError as e:
            with ctx.exc(CCGSemanticsCompositionError, e):
                raise CCGSemanticsCompositionError('Cannot make backward application.', lhs, rhs, e) from e
    with ctx.exc(CCGSemanticsCompositionError):
        raise CCGSemanticsCompositionError('Cannot make backward application.', lhs, rhs, 'Functor/Value types do not match.')


def _coordination(lhs, conj, rhs):
    ctx = get_ccg_composition_context()

    if isinstance(lhs, Function) and isinstance(rhs, Function) and callable(conj):
        if lhs.ftype.nr_variable_arguments == rhs.ftype.nr_variable_arguments:
            return conj(lhs, rhs)
    if isinstance(lhs, Value) and isinstance(rhs, Value) and callable(conj):
        return conj(lhs, rhs)

    with ctx.exc(CCGSemanticsCompositionError):
        raise CCGSemanticsCompositionError(
            'Cannot make coordination.',
            lhs, rhs, conj=conj,
            error='Functor arity does not match.'
        )


class CCGSemanticsSugar(object):
    """A syntax sugar that allows users to write CCG semantics in a more natural way.
    This class is initialized in the :class:`~concepts.language.ccg.grammar.CCG` class.
    """

    domain: FunctionDomain

    def __init__(self, domain: FunctionDomain):
        """Initialize the syntax sugar.

        Args:
            domain: the function domain.
        """
        self.domain = domain

    domain: FunctionDomain
    """The underlying function domain."""

    def __getitem__(self, item: Optional[Union[Value, Function, ConstantExpression, FunctionApplicationExpression, Callable, Tuple[Callable, Dict]]]) -> CCGSemantics:
        """Create a :class:`CCGSemantics` instance from a function or a value."""
        if item is None:
            return CCGSemantics(None)
        if isinstance(item, CCGSemantics):
            return item
        if isinstance(item, Value):
            return CCGSemantics(ConstantExpression(item))
        if isinstance(item, (Function, ConstantExpression, FunctionApplicationExpression)):
            return CCGSemantics(item)
        if isinstance(item, tuple):
            assert len(item) == 2
            return CCGSemantics(self.domain.lam(item[0], typing_cues=item[1]))
        assert callable(item)
        return CCGSemantics(self.domain.lam(item))


concepts/language/ccg/semantics.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : syntax.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/05/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data structures for syntax types in a linguistic CCG."""

from typing import Optional, Union, Tuple, List
from jacinle.utils.cache import cached_property
from jacinle.utils.meta import repr_from_str
from jacinle.utils.printing import indent_text

from concepts.language.ccg.composition import CCGCompositionDirection, CCGCompositionError, get_ccg_composition_context, CCGComposable

__all__ = [
    'CCGSyntaxCompositionError', 'CCGSyntaxTypeParsingError',
    'CCGSyntaxType', 'CCGPrimitiveSyntaxType', 'CCGConjSyntaxType', 'CCGComposedSyntaxType',
    'CCGSyntaxSystem', 'parse_syntax_type'
]


class CCGSyntaxCompositionError(CCGCompositionError):
    """Raised when the composition of two syntax types is not allowed."""

    def __init__(self, message: Optional[str] = None):
        if message is None:
            super().__init__(None)
        else:
            super().__init__('(Syntax) ' + message)


class CCGSyntaxTypeParsingError(Exception):
    """Raised when the parsing of a syntax type string fails."""


class CCGSyntaxType(CCGComposable):
    """Syntax types for CCG.

    There are three types of syntax types:

        - Primitive syntax types: `N`, `S`, `NP`, `VP`, etc.
        - Composed syntax types: `S/NP`, `S\\NP`, etc.
        - Conjunction syntax types: `CONJ`, etc.
    """

    def __init__(self, typename: Optional[str] = None):
        self.typename = typename

    # This property is inherited from CCGComposable.
    @property
    def is_none(self) -> bool:
        return self.typename is None

    # This property is inherited from CCGComposable.
    @property
    def is_conj(self) -> bool:
        return False

    @property
    def arity(self) -> int:
        """The arity of the syntax type. That is, the number of arguments it needs to combine before it becomes a primitive syntax type."""
        return 0

    @property
    def is_function(self) -> bool:
        """Whether the syntax type is a function type. That is, whether it can do function application with another syntax type."""
        return False

    @property
    def is_value(self) -> bool:
        """Whether the syntax type is a value type. That is, whether it is a primitive syntax type."""
        return False

    @property
    def parenthesis_typename(self) -> str:
        """Return the typename with parenthesis."""
        return self.typename

    def _fapp(self, right: 'CCGSyntaxType') -> 'CCGSyntaxType':
        return _forward_application(self, right)

    def _bapp(self, lhs: 'CCGSyntaxType') -> 'CCGSyntaxType':
        return _backward_application(lhs, self)

    def _coord3(self, lhs: 'CCGSyntaxType', rhs: 'CCGSyntaxType') -> 'CCGSyntaxType':
        return _coordination(lhs, self, rhs)

    def __str__(self) -> str:
        return str(self.typename)

    __repr__ = repr_from_str

    def __truediv__(self, other: 'CCGSyntaxType') -> 'CCGSyntaxType':
        """Construct a `A/B` syntax type."""
        return CCGComposedSyntaxType(self, other, direction=CCGCompositionDirection.RIGHT)

    def __floordiv__(self, other: 'CCGSyntaxType') -> 'CCGSyntaxType':
        """Construct a `A\\B` syntax type."""
        return CCGComposedSyntaxType(self, other, direction=CCGCompositionDirection.LEFT)

    def __eq__(self, other: 'CCGSyntaxType') -> bool:
        """Return whether two syntax types are equal."""
        return self.typename == other.typename

    def __ne__(self, other: 'CCGSyntaxType') -> bool:
        return self.typename != other.typename

    def __hash__(self):
        return str(self)

    def __lt__(self, other: 'CCGSyntaxType') -> bool:
        """Customized comparison function for sorting a list of syntax types."""
        a, b = str(self), str(other)
        return (a.count('/') + a.count('\\'), a) < (b.count('/') + b.count('\\'), b)

    def flatten(self) -> List[Union['CCGSyntaxType', Tuple['CCGSyntaxType', CCGCompositionDirection]]]:
        """Flatten the recursive definition of a syntax type into a list of lower-level syntax types. For example,
        the syntax type ``S/NP`` will be flattened into ``[S, NP, (S/NP, RIGHT)]``.

        Returns:
            the list of flattened lower-level syntax types.
        """
        raise NotImplementedError()


class CCGPrimitiveSyntaxType(CCGSyntaxType):
    """The primitive syntax types (e.g., NP)."""

    @property
    def is_value(self) -> bool:
        return True

    def flatten(self) -> List[CCGSyntaxType]:
        return [self]


class CCGConjSyntaxType(CCGSyntaxType):
    """A conjunction syntax type."""

    @property
    def is_conj(self):
        return True

    def __call__(self, lhs: CCGSyntaxType, rhs: CCGSyntaxType) -> CCGSyntaxType:
        """Construct the resulting syntax type for `A CONJ B` given A and B.

        Args:
            lhs: The left syntax type (A).
            rhs: The right syntax type (B).

        Returns:
            CCGSyntaxType: The resulting syntax type.
        """
        return lhs

    def flatten(self) -> List[CCGSyntaxType]:
        return [self]


class CCGComposedSyntaxType(CCGSyntaxType):
    """A composed syntax type (e.g., S/NP)."""

    def __init__(self, main: CCGSyntaxType, sub: CCGSyntaxType, direction: CCGCompositionDirection):
        """Initialize the composed syntax type.

        Args:
            main: the main syntax type (e.g., S).
            sub: the sub syntax type (e.g., NP).
            direction: the composition direction (e.g., RIGHT).
        """

        self.main = main
        self.sub = sub
        self.direction = CCGCompositionDirection.from_string(direction)

        if self.direction is CCGCompositionDirection.RIGHT:
            typename = self.main.typename + '/' + self.sub.parenthesis_typename
        else:
            typename = self.main.typename + '\\' + self.sub.parenthesis_typename
        super().__init__(typename)

    @cached_property
    def arity(self) -> int:
        return self.main.arity + 1

    @property
    def is_function(self) -> bool:
        return True

    @property
    def parenthesis_typename(self) -> str:
        return '{' + f'{self.typename}' + '}'

    def flatten(self) -> List[Union[CCGSyntaxType, Tuple[CCGSyntaxType, CCGCompositionDirection]]]:
        ret = self.main.flatten()
        ret.append((self.sub, self.direction))
        return ret


def _forward_application(lhs, rhs):
    if isinstance(lhs, CCGComposedSyntaxType):
        if lhs.direction == CCGCompositionDirection.RIGHT:
            if lhs.sub == rhs:
                return lhs.main
    with get_ccg_composition_context().exc(CCGSyntaxCompositionError):
        raise CCGSyntaxCompositionError(f'Cannot make forward application of {lhs} and {rhs}.')


def _backward_application(lhs, rhs):
    if isinstance(rhs, CCGComposedSyntaxType):
        if rhs.direction == CCGCompositionDirection.LEFT:
            if rhs.sub == lhs:
                return rhs.main
    with get_ccg_composition_context().exc(CCGSyntaxCompositionError):
        raise CCGSyntaxCompositionError(f'Cannot make backward application of {lhs} and {rhs}.')


def _coordination(lhs, conj, rhs):
    if lhs == rhs and isinstance(conj, CCGConjSyntaxType):
        return conj(lhs, rhs)
    with get_ccg_composition_context().exc(CCGSyntaxCompositionError):
        raise CCGSyntaxCompositionError(f'Cannot make coordination of {lhs} {conj} {rhs}.')


class CCGSyntaxSystem(object):
    """A data structure that keeps track of a set of primitive and conjunction syntax types allowed in a grammar."""

    def __init__(self):
        self.types = dict()

    def define_primitive_type(self, stype: Union[CCGSyntaxType, str]):
        """Define a primitive syntax type.

        Args:
            stype: The syntax type to be defined.
        """
        if isinstance(stype, CCGSyntaxType):
            self.types[stype.typename] = stype
        elif isinstance(stype, str):
            self.types[stype] = CCGPrimitiveSyntaxType(stype)
        else:
            raise TypeError(f'Invalid type: {stype}.')

    def define_conj_type(self, stype: Union[CCGSyntaxType, str]):
        """Define a conj syntax type.

        Args:
            stype: The syntax type to be defined.
        """
        if isinstance(stype, CCGSyntaxType):
            self.types[stype.typename] = stype
        elif isinstance(stype, str):
            self.types[stype] = CCGConjSyntaxType(stype)
        else:
            raise TypeError(f'Invalid type: {stype}.')

    def __getitem__(self, item: Optional[Union[CCGSyntaxType, str]]) -> CCGSyntaxType:
        """A syntax sugar for `parse_syntax_type`.

        - When the string is `None`, return `None`.
        - When the string is a `CCGSyntaxType`, return the type itself.

        Args:
            item: The string to be parsed.

        Returns:
            CCGSyntaxType: The parsed syntax type.
        """
        if item is None:
            return CCGSyntaxType(None)
        if isinstance(item, CCGSyntaxType):
            return item
        return parse_syntax_type(item, syntax_system=self)

    def __str__(self) -> str:
        return 'CCGSyntaxSystem(' + ', '.join([str(x) for x in self.types.keys()]) + ')'

    __repr__ = __str__

    def format_sumamry(self) -> str:
        fmt = 'Primitive and Conjunction types:\n'
        for type in self.types.values():
            fmt += '  ' + str(type) + '\n'
        fmt = 'CCGSyntaxSystem:\n' + indent_text(fmt.rstrip())
        return fmt

    def print_summary(self):
        print(self.format_sumamry())


def parse_syntax_type(string: str, syntax_system: Optional[CCGSyntaxSystem] = None) -> CCGSyntaxType:
    """Parse a string to a syntax type.

    Args:
        string: The string to be parsed.
        syntax_system: The syntax system to be used. Defaults to None.

    Returns:
        CCGSyntaxType: The parsed syntax type.
    """

    def parse_inner(current):
        if current == '':
            raise CCGSyntaxTypeParsingError('Invalid syntax type string (got empty type): {}.'.format(string))
        nr_parenthesis = 0
        last_op = None
        for i, c in enumerate(current):
            if c in r'\/':
                if nr_parenthesis == 0:
                    last_op = i
            if c == '(':
                nr_parenthesis += 1
            elif c == ')':
                nr_parenthesis -= 1
                if nr_parenthesis < 0:
                    raise CCGSyntaxTypeParsingError('Invalid parenthesis (extra ")"): {}.'.format(string))
        if nr_parenthesis != 0:
            raise CCGSyntaxTypeParsingError('Invalid parenthesis (extra "("): {}.'.format(string))

        if last_op is None:
            if current[0] == '(' and current[-1] == ')':
                return parse_inner(current[1:-1])
            else:
                if syntax_system is None:
                    return CCGSyntaxType(current)
                else:
                    if current in syntax_system.types:
                        return syntax_system.types[current]
                    else:
                        raise CCGSyntaxTypeParsingError('Unknown primitive syntax type {} during parsing {}.'.format(current, string))

        last_op_value = CCGCompositionDirection.RIGHT if current[last_op] == '/' else CCGCompositionDirection.LEFT
        return CCGComposedSyntaxType(
            parse_inner(current[:last_op]),
            parse_inner(current[last_op + 1:]),
            direction=last_op_value
        )

    return parse_inner(string)


concepts/language/ccg/syntax.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 04/23/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/language/gpt_parsing/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : caption_sng.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 04/23/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import openai
from dataclasses import dataclass
from typing import Any, Tuple, Dict
from concepts.language.gpt_parsing.utils import TagNotUniqueError, load_prompt, extract_tag


@dataclass
class CaptionSceneGraph(object):
    objects: Tuple[str]
    relations: Tuple[Tuple[str, str, str]]


class CaptionSNGParser(object):
    def __init__(self, max_tokens: int = 1024):
        self.prompt = load_prompt('gpt-35-turbo-chat-caption')
        self.max_tokens = max_tokens

    def parse(self, sentence: str) -> Dict[str, Any]:
        response = openai.ChatCompletion.create(
            model='gpt-3.5-turbo',
            messages=[
                self.prompt[0],  # the system prmopt
                {'role': 'user', 'content': sentence}
            ],
            max_tokens=self.max_tokens,
        )

        parsing = None
        exception = None
        try:
            parsing = self.extract(response['choices'][0]['message']['content'])
        except TagNotUniqueError as e:
            exception = e

        return {
            'sentence': sentence,
            'raw_response': response,
            'parsing': parsing,
            'exception': exception,
        }

    def parse_batch(self, sentences):
        # TODO(Jiayuan Mao @ 2023/04/23): support batchified parsing.
        raise NotImplementedError()

    def extract(self, raw_response):
        # extract all objects
        objects = extract_tag(raw_response, 'object', unique=False)

        # extract all relations
        raw_relations = extract_tag(raw_response, 'relation', unique=False)

        objects = set(objects)
        relations = list()
        for relation in raw_relations:
            this_subject = extract_tag(relation, 'subject', unique=True)
            this_predicate = extract_tag(relation, 'predicate', unique=True)
            this_object = extract_tag(relation, 'object', unique=True)

            if this_subject not in objects:
                objects.add(this_subject)
            if this_object not in objects:
                objects.add(this_object)
            relations.append((this_subject, this_predicate, this_object))

        objects = tuple(sorted(objects))
        relations = tuple(sorted(relations))
        return CaptionSceneGraph(objects, relations)


default_caption_sng_parser = CaptionSNGParser()


def parse_caption(sentence: str) -> Dict[str, Any]:
    return default_caption_sng_parser.parse(sentence)


if __name__ == '__main__':
    print(parse_caption('A little girl and a woman are having their picture taken in front of a desert.'))


concepts/language/gpt_parsing/caption_sng.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : utils.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 04/23/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import re
import os.path as osp
from typing import Union, List, Dict


class ParsingFailedError(Exception):
    pass


def load_prompt(identifier) -> List[Dict[str, str]]:
    prompt_filename = osp.join(osp.dirname(__file__), 'prompts', identifier + '.txt')
    with open(prompt_filename, 'r') as f:
        content = f.read()
        return [
            {'role': 'system', 'content': content},
        ]


class TagNotUniqueError(ParsingFailedError):
    pass


def extract_tag(content: str, tag: str, unique: bool = False) -> Union[str, List[str]]:
    """Extract all matched content inside <tag></tag>.

    Args:
        content: the input string.
        tag: the tag name.
        unique: if True, only return the first matched content and raises an error.
    """

    pattern = r'<{}>(.*?)</{}>'.format(tag, tag)
    matches = re.findall(pattern, content, re.DOTALL)

    if unique:
        if len(matches) != 1:
            raise TagNotUniqueError('Tag "{}" is not unique. Content: {}'.format(tag, content))
        return matches[0]
    else:
        return matches


concepts/language/gpt_parsing/utils.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/07/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data structures and algorithms for Neural CCG.
This module is the core of paper [Grammar-Based Grounded Lexicon Learning](https://arxiv.org/abs/2202.08806).

See the documentation for file ``grammar.py`` for more details.
"""

concepts/language/neural_ccg/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : ckyee.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/11/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""NeuralCCG with CKY-EE (CKY with Expected Execution)."""

from collections import defaultdict
from dataclasses import dataclass
from typing import Optional, Union, Iterable, Tuple, List, Dict, Callable

import torch
import torch.nn.functional as F

import jacinle.random as jacrandom
import jactorch

from concepts.dsl.function_domain import FunctionDomain
from concepts.dsl.value import Value
from concepts.dsl.executors.function_domain_executor import FunctionDomainExecutor
from concepts.language.ccg.composition import CCGCompositionSystem, CCGCoordinationImmNode
from concepts.language.neural_ccg.grammar import NeuralCCGSemanticsPartialTypeLex, NeuralCCGSemantics, NeuralCCGNode, NeuralCCG
from concepts.language.neural_ccg.search import NeuralCCGLexiconSearchResult


__all__ = ['CKYEEExpectationFunction', 'CKYEEExpectationConfig', 'NeuralCKYEE']

_profile = getattr(__builtins__, 'profile', lambda x: x)


class CKYEEExpectationFunction(object):
    """A collection of functions to perform expected execution. This class maintains a set of functions of the following form:

    .. code-block:: python

        def expectation_<typename>(self, values: List[Value], weights: List[torch.Tensor]) -> Tuple[Optional[Value], Optional[Tuple[torch.Tensor, float]]]:
            ...

    where ``typename`` is the name of a type, and ``values`` is a list of values of the type, and ``weights`` is a list of corresponding weights.
    The function should return a tuple of (expectation, (sum_weight, max_weight)), where ``expectation`` is the expectation of the values,
    as a :class:`~concepts.dsl.value.Value` instance, and ``sum_weight`` is the sum of the weights (as a PyTorch tensor),
    and ``max_weight`` is the maximum of the weights (as a float).
    """

    def __init__(self, domain: FunctionDomain):
        """Initialize the expectation function.

        Args:
            domain: the function domain.
        """
        self._domain = domain
        self._expectation_functions = dict()
        self._init_expectation_functions()

    _domain: FunctionDomain
    _expectation_functions: Dict[str, Callable[[List[Value], List[torch.Tensor]], Tuple[torch.Tensor, Tuple[torch.Tensor, float]]]]

    def register_function(self, typename: str, function: Callable[[List[Value], List[torch.Tensor]], Tuple[torch.Tensor, Tuple[torch.Tensor, float]]]):
        """Register a function to compute the expectation of a list of values.

        Args:
            typename: the name of the type.
            function: the function to compute the expectation of a list of values.
        """
        self._expectation_functions[typename] = function

    def get_function(self, typename: str) -> Optional[Callable[[List[Value], List[torch.Tensor]], Tuple[torch.Tensor, Tuple[torch.Tensor, float]]]]:
        """Get the registered function for a type.

        Args:
            typename: the name of the type.

        Returns:
            The registered function.
        """
        return self._expectation_functions.get(typename, None)

    def _init_expectation_functions(self):
        for typename in self._domain.types:
            funcname = 'expectation_' + typename
            if hasattr(self, funcname):
                self.register_function(typename, getattr(self, funcname))

    def expectation(self, values: List[Value], weights: List[torch.Tensor]) -> Tuple[Optional[Value], Optional[Tuple[torch.Tensor, float]]]:
        """Compute the expectation of a list of values.

        Args:
            values: a list of values.
            weights: a list of corresponding weights.

        Returns:
            A tuple of (expectation, (sum_weight, max_weight)).
            The ``expectation`` is the expectation of the values, as a :class:`~concepts.dsl.value.Value` instance.
            The second element is a tuple of (sum_weight, max_weight), where ``sum_weight`` is the sum of the weights (as a PyTorch tensor),
            and ``max_weight`` is the maximum of the weights (as a float).
        """
        if not isinstance(values[0], Value):
            return None, None
        function = self.get_function(values[0].dtype.typename)
        if function is None:
            return None, None
        return function(values, weights)

    def _expectation_set_tensors(self, sets: Union[Tuple[torch.Tensor, ...], List[torch.Tensor]], weights: Union[Tuple[torch.Tensor, ...], List[torch.Tensor]], log: bool = False) -> Tuple[torch.Tensor, Tuple[torch.Tensor, float]]:
        """A helper function to compute the expectation of a list of tensors representing sets.
        Typically, the tensor representation of a set is a one-hot vector, where each entry represents the probabbility of the corresponding element in the set.

        Args:
            sets: a list of tensors representing sets.
            weights: a list of corresponding weights.
            log: whether the probabilities in sets are in log-space.

        Returns:
            A tuple of (set, (sum_weight, max_weight)).
            The ``set`` is a tensor representing the expectation of the list of sets, as a PyTorch tensor.
            The second element is a tuple of (sum_weight, max_weight), where ``sum_weight`` is the sum of the weights (as a PyTorch tensor),
            and ``max_weight`` is the maximum of the weights (as a float).
        """
        sets_tensor = torch.stack(sets, dim=-1)
        weights_tensor = torch.stack(weights, dim=0)

        weights_sum = jactorch.logsumexp(weights_tensor)
        weights_max = weights_tensor.argmax().item()

        if log:
            weights_tensor = F.log_softmax(weights_tensor, dim=-1)
            weights_tensor = jactorch.add_dim_as_except(weights_tensor, sets_tensor, -1)
            output = jactorch.logsumexp(sets_tensor + weights_tensor, dim=-1)
        else:
            weights_tensor = F.softmax(weights_tensor, dim=-1)
            output = torch.matmul(sets_tensor, weights_tensor)

        return output, (weights_sum, weights_max)


def aggregate_weights(weights: Union[torch.Tensor, Tuple[torch.Tensor, ...], List[torch.Tensor]], log=True) -> Tuple[torch.Tensor, float]:
    """Sum up a list of parsing weights (probabilities).

    Args:
        weights: a list of weights. It can be a single tensor, or a list of tensors.
        log: whether the weights are in log space. When this is True, the output will be computed as ``logsumexp``.

    Returns:
        A tuple of (sum, max).
    """
    if torch.is_tensor(weights):
        weights_tensor = weights
    elif isinstance(weights, (tuple, list)):
        weights_tensor = torch.stack(weights, dim=0)
    else:
        weights_tensor = torch.tensor(weights)

    if log:
        return jactorch.logsumexp(weights_tensor), weights_tensor.argmax().item()
    else:
        return weights_tensor.sum(), weights_tensor.argmax().item()


@dataclass
class CKYEEExpectationConfig(object):
    """Configurations for CKYEE expectation computation."""
    compress_values: bool = True
    compress_0varbinding_functions: bool = True
    compress_1varbinding_functions: bool = True
    sample: bool = False


def _get_typename(x):
    return x.typename if x is not None else None


class NeuralCKYEE(NeuralCCG):
    """The neural CCG grammar with CKY-EE for chart parsing."""

    def __init__(
        self,
        domain: FunctionDomain,
        executor: FunctionDomainExecutor,
        candidate_lexicon_entries: Iterable[NeuralCCGLexiconSearchResult],
        expectation_function: CKYEEExpectationFunction,
        composition_system: Optional[CCGCompositionSystem] = None,
        *,
        expectation_config: Optional[CKYEEExpectationConfig] = None,
        joint_execution: bool = True,
        allow_none_lexicon: bool = False,
        reweight_meaning_lex: bool = False
    ):
        """Initialize the neural CCG grammar with CKY-EE.

        Args:
            domain: the function domain of the grammar.
            executor: the executor for the function domain.
            candidate_lexicon_entries:  a list of candidate lexicon entries.
            exepctation_function: the expectation function for CKY-EE. See :class:`CKYEEExpectationFunction`.
            composition_system: the composition system. If None, the default composition system will be used.
            expectation_config: the configuration for CKY-EE expectation computation.
            joint_execution: whether to execute the partial programs during CKY.
            allow_none_lexicon: whether to allow None lexicon.
            reweight_meaning_lex: whether to reweight the meaning lexicon entries. Specifically, if there are two parsings
                share the same set of lexicon entries (i.e., caused by ambiguities in combination), specifying this flag
                will reweight both of them to be 1 / (number of parsings that used this set of lexicon entries).
        """
        super().__init__(
            domain, executor, candidate_lexicon_entries, composition_system,
            joint_execution=joint_execution, allow_none_lexicon=allow_none_lexicon, reweight_meaning_lex=reweight_meaning_lex
        )
        self.expectation_function = expectation_function
        self.expectation_config = expectation_config if expectation_config is not None else CKYEEExpectationConfig()

    training: bool
    domain: FunctionDomain
    executor: FunctionDomainExecutor
    composition_system: CCGCompositionSystem
    candidate_lexicon_entries: Tuple['NeuralCCGLexiconSearchResult', ...]
    joint_execution: bool
    allow_none_lexicon: bool
    reweight_meaning_lex: bool

    def _unique_lexicon(self, nodes: List[NeuralCCGNode], syntax_only: bool) -> List[NeuralCCGNode]:
        if syntax_only:
            return self._unique_syntax_only(nodes)
        return self._collect_nodes_by_type(nodes)

    def _unique(self, nodes: List[NeuralCCGNode], syntax_only: bool) -> List[NeuralCCGNode]:
        nodes = super()._unique(nodes, syntax_only)
        if syntax_only:
            return nodes
        nodes = self._collect_nodes_by_type(nodes)
        return nodes

    def _reweight_parse_trees(self, parsings: List[NeuralCCGNode]) -> List[NeuralCCGNode]:
        raise NotImplementedError('Reweight pr[meaning | lex_i] is not supported for NeuralSoftCCG.')

    def _collect_nodes_by_type(self, nodes: List[NeuralCCGNode]) -> List[NeuralCCGNode]:
        output_nodes = nodes

        if self.expectation_config.compress_values:
            output_nodes = self._collect_values(output_nodes)
        if self.expectation_config.compress_0varbinding_functions:
            output_nodes = self._collect_0varbinding_functions(output_nodes)
        if self.expectation_config.compress_1varbinding_functions:
            output_nodes = self._collect_1varbinding_functions(output_nodes)
        if self.expectation_config.sample:
            output_nodes = self._collect_by_sample(output_nodes)

        return output_nodes

    @_profile
    def _collect_values(self, nodes):
        output_nodes = list()
        value_nodes = defaultdict(lambda: (list(), list(), list()))
        for node in nodes:
            if not isinstance(node.syntax, CCGCoordinationImmNode) and node.syntax.is_value:
                rec = value_nodes[(node.syntax.return_type.typename, _get_typename(node.syntax.lang_syntax_type))]
                rec[0].append(node)
                if self.joint_execution:
                    rec[1].append(node.execution_result)
                rec[2].append(node.weight)
            else:
                output_nodes.append(node)

        for typename, rec in value_nodes.items():
            output_nodes.extend(self._collect_values_once(*rec))
        return output_nodes

    @_profile
    def _collect_0varbinding_functions(self, nodes):
        output_nodes = list()

        function_nodes = defaultdict(lambda: (list(), list(), list()))
        for node in nodes:
            sem = node.semantics
            if not isinstance(node.syntax, CCGCoordinationImmNode) and sem.partial_type is not None and len(sem.partial_type) == 1 and isinstance(sem.partial_type[0], NeuralCCGSemanticsPartialTypeLex):
                rec = function_nodes[(sem.partial_type[0], _get_typename(node.syntax.lang_syntax_type))]
                rec[0].append(node)
                rec[2].append(node.weight)
            else:
                output_nodes.append(node)

        for typename, rec in function_nodes.items():
            output_nodes.extend(self._collect_0varbinding_functions_once(*rec))

        return output_nodes

    @_profile
    def _collect_1varbinding_functions(self, nodes):
        output_nodes = list()

        function_nodes = defaultdict(lambda: (list(), list(), list()))
        for node in nodes:
            sem = node.semantics
            if not isinstance(node.syntax, CCGCoordinationImmNode) and sem.partial_type is not None and len(sem.partial_type) == 2 and isinstance(sem.partial_type[0], NeuralCCGSemanticsPartialTypeLex):
                rec = function_nodes[(sem.partial_type[0], _get_typename(node.syntax.lang_syntax_type))]
                rec[0].append(node)
                if self.joint_execution:
                    rec[1].append(node.semantics.execution_buffer[1])
                rec[2].append(node.weight)
            else:
                output_nodes.append(node)

        for typename, rec in function_nodes.items():
            output_nodes.extend(self._collect_1varbinding_functions_once(*rec))

        return output_nodes

    @_profile
    def _collect_by_sample(self, nodes):
        output_nodes = list()
        value_nodes = defaultdict(lambda: (list(), list()))
        for node in nodes:
            rec = value_nodes[str(node.syntax)]
            rec[0].append(node)
            rec[1].append(node.weight)

        for typename, rec in value_nodes.items():
            output_nodes.extend(self._collect_by_sample_once(*rec))
        return output_nodes

    def _collect_values_once(self, nodes: List[NeuralCCGNode], all_results: List[Value], all_weights: List[torch.Tensor]) -> List[NeuralCCGNode]:
        if len(nodes) > 1:
            if self.joint_execution:
                compressed_node, weights = self.expectation_function.expectation(all_results, all_weights)
            else:
                compressed_node = None
                weights = aggregate_weights(all_weights, log=True)

            if self.joint_execution and compressed_node is None:
                return nodes
            else:
                sample_node = nodes[weights[1]]
                new_node = NeuralCCGNode(
                    composition_system=sample_node.composition_system,
                    syntax=sample_node.syntax,
                    semantics=NeuralCCGSemantics(
                        sample_node.semantics.value,
                        execution_buffer=[compressed_node] if compressed_node is not None else None,
                        partial_type=sample_node.semantics.partial_type,
                        nr_execution_steps=self.count_nr_execution_steps(nodes)
                    ),
                    composition_type=sample_node.composition_type,
                    lexicon=sample_node.lexicon, lhs=sample_node.lhs, rhs=sample_node.rhs,
                    composition_str=sample_node.composition_str,
                    weight=weights[0]
                )
                new_node.set_used_lexicon_entries(self.gen_valid_lexicons(nodes))
                return [new_node]
        else:
            return nodes

    def _collect_0varbinding_functions_once(self, nodes: List[NeuralCCGNode], all_results: List[Value], all_weights: List[torch.Tensor]) -> List[NeuralCCGNode]:
        if len(nodes) > 1:
            weights = aggregate_weights(all_weights, log=True)
            sample_node = nodes[0]
            new_node = NeuralCCGNode(
                composition_system=sample_node.composition_system,
                syntax=sample_node.syntax,
                semantics=sample_node.semantics,
                composition_type=sample_node.composition_type,
                lexicon=sample_node.lexicon, lhs=sample_node.lhs, rhs=sample_node.rhs,
                composition_str=sample_node.composition_str,
                weight=weights[0]
            )
            new_node.set_used_lexicon_entries(self.gen_valid_lexicons(nodes))
            return [new_node]
        else:
            return nodes

    def _collect_1varbinding_functions_once(self, nodes: List[NeuralCCGNode], all_results: List[Value], all_weights: List[torch.Tensor]) -> List[NeuralCCGNode]:
        if len(nodes) > 1:
            if self.joint_execution:
                compressed_node, weights = self.expectation_function.expectation(all_results, all_weights)
            else:
                compressed_node = None
                weights = aggregate_weights(all_weights, log=True)

            if self.joint_execution and compressed_node is None:
                return nodes
            else:
                sample_node = nodes[weights[1]]
                new_node = NeuralCCGNode(
                    composition_system=sample_node.composition_system,
                    syntax=sample_node.syntax,
                    semantics=NeuralCCGSemantics(
                        sample_node.semantics.value,
                        execution_buffer=[
                            sample_node.semantics.execution_buffer[0],
                            compressed_node
                        ] if compressed_node is not None else None,
                        partial_type=sample_node.semantics.partial_type,
                        nr_execution_steps=self.count_nr_execution_steps(nodes)
                    ),
                    composition_type=sample_node.composition_type,
                    lexicon=sample_node.lexicon, lhs=sample_node.lhs, rhs=sample_node.rhs,
                    composition_str=sample_node.composition_str,
                    weight=weights[0]
                )
                new_node.set_used_lexicon_entries(self.gen_valid_lexicons(nodes))
                return [new_node]
        else:
            return nodes

    def _collect_by_sample_once(self, nodes: List[NeuralCCGNode], all_weights: List[torch.Tensor]) -> List[NeuralCCGNode]:
        if len(nodes) > 1:
            weights_tensor = torch.stack(all_weights, dim=0).detach()
            index = jacrandom.choice(weights_tensor.shape[0], p=jactorch.as_numpy(F.softmax(weights_tensor)))
            weights_tensor.data[index] = -1e9
            total_weights = jactorch.logaddexp(all_weights[index], jactorch.logsumexp(weights_tensor))

            sample_node = nodes[index]
            new_node = NeuralCCGNode(
                composition_system=sample_node.composition_system,
                syntax=sample_node.syntax,
                semantics=NeuralCCGSemantics(
                    sample_node.semantics.value,
                    execution_buffer=sample_node.semantics.execution_buffer,
                    partial_type=sample_node.semantics.partial_type,
                    nr_execution_steps=self.count_nr_execution_steps(nodes)
                ),
                composition_type=sample_node.composition_type,
                lexicon=sample_node.lexicon, lhs=sample_node.lhs, rhs=sample_node.rhs,
                composition_str=sample_node.composition_str,
                weight=total_weights
            )
            new_node.set_used_lexicon_entries(self.gen_valid_lexicons(nodes))
            return [new_node]
        else:
            return nodes

concepts/language/neural_ccg/ckyee.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : grammar.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/07/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""The grammar components for neural CCGs, including syntax types, semantic forms, lexicon, and the parser.
There are a few debug options that can be enabled by setting the corresponding flags to True, see the
:class:`foptions` variable for details.
"""

import math
import heapq
import itertools

from typing import TYPE_CHECKING, TypeVar, Any, Optional, Union, Iterable, Sequence, Tuple, List, Set, Dict, Callable
from collections import defaultdict, Counter
from dataclasses import dataclass
from tabulate import tabulate

import torch
import torch.nn as nn
import torch.nn.functional as F
import jacinle
import jactorch

from jacinle.config.environ_v2 import configs, def_configs
from jacinle.utils.cache import cached_property
from jacinle.utils.meta import repr_from_str

from concepts.dsl.dsl_types import FormatContext, TypeBase, ObjectType, ValueType, ConstantType
from concepts.dsl.dsl_functions import FunctionType, Function
from concepts.dsl.value import Value
from concepts.dsl.expression import ConstantExpression, FunctionApplicationExpression, get_type
from concepts.dsl.function_domain import FunctionDomain
from concepts.dsl.executors.function_domain_executor import FunctionDomainExecutor
from concepts.language.ccg.composition import CCGCompositionType, CCGCompositionDirection, CCGCompositionSystem, CCGComposable, CCGCoordinationImmNode, CCGCompositionResult
from concepts.language.ccg.composition import CCGCompositionError, get_ccg_composition_context, CCGCompositionContext
from concepts.language.ccg.syntax import CCGSyntaxType, CCGSyntaxCompositionError
from concepts.language.ccg.semantics import CCGSemantics, CCGSemanticsLazyValue
from concepts.language.ccg.grammar import CCGNode, Lexicon

if TYPE_CHECKING:
    from concepts.language.neural_ccg.search import NeuralCCGLexiconSearchResult

T = TypeVar('T')


__all__ = [
    'foptions',
    'LinearizationTuple', 'parse_linearization_string', 'NeuralCCGSyntaxType', 'NeuralCCGConjSyntaxType',
    'NeuralCCGSemanticsPartialTypeLex', 'NeuralCCGSemantics',
    'NeuralCCGGroundingFunction', 'NeuralCCGSemanticsExecutionBuffer', 'NeuralCCGConjGroundingFunction', 'NeuralCCGSimpleConjFunction',
    'NeuralCCGNode', 'compose_neural_ccg_node', 'NeuralCCG',
]


logger = jacinle.get_logger(__file__)
_profile = getattr(__builtins__, 'profile', lambda x: x)


with def_configs():
    configs.neural_ccg.debug_print = False
    configs.neural_ccg.debug_stat_types = False
    configs.neural_ccg.training_use_lazy_semantics = True
    configs.neural_ccg.compose_function_note = False


foptions = configs.neural_ccg
"""Options for the CCG parser. Use ``foptions.xxx`` to access the options.

- ``debug_print``: print debug information.
- ``debug_stat_types``: print statistics about the types.
- ``training_use_lazy_semantics``: whether to use lazy semantics during training.
- ``compose_function_note``: whether to compose the function notes of semantics.
    Specifically, if this is set to True, then every :class:`NeuralCCGGroundingFunction` instance will compose a note
    that records the trace of the fucntions that has been called.
"""


@dataclass
class LinearizationTuple(object):
    """A pair of (index, direction) for representing linearization of function arguments."""

    index: int
    """The index of the argument."""

    direction: CCGCompositionDirection
    """The direction of the linearization (left or right)."""

    def __str__(self) -> str:
        if self.direction is CCGCompositionDirection.LEFT:
            return f'{self.index}<'
        else:
            return f'{self.index}>'


def parse_linearization_string(string: str) -> List[LinearizationTuple]:
    r"""Parse a linearization string (e.g., ``/0\1``) into a list of tuples.

    Args:
        string: the linearization string.

    Returns:
        A list of tuples, each tuple is a pair of ``(index, direction)``.
    """
    if string == '':
        return []
    string = string.replace('/', '@/').replace('\\', '@\\')
    assert string.startswith('@')
    parts = string.split('@')[1:]
    linearization = list()
    for p in parts:
        linearization.append(LinearizationTuple(
            int(p[1:]),
            CCGCompositionDirection.RIGHT if p[0] == '/' else CCGCompositionDirection.LEFT
        ))
    return linearization


class NeuralCCGSyntaxType(CCGComposable):
    """Data structure for neural CCG syntax types.
    NeuralCCGSyntaxType implements a different version of the syntax type than the CCGSyntaxType.

    Specifically, the :class:`NeuralCCGSyntaxType` define syntaxes based on the signature of its meaning. It has two parts:

    1. the signature of the semantic form, including the argument types and the return type.
    2. the linearization of the semantics.

    IMPORTANT: Arguments are combined from right to left. Thus, argument_types[-1] is the first argument.
    There is one possible confusion about linearization:

    Although LinearizationTuple has two fields: index and direction, only type will be used during computation.
    LinearizationTuple.index is used to record the mapping from the original "function" definition, most likely
    produced by iter_from_function.

    Example: Suppose we have a function definition def func(string, int) -> string.
    Running iter_from_function(func) will get you six possible syntax types:

        string/0:string/1:int, string\\0:string/1:int, string\\0:string\\1:int,
        string/1:int/0:string, string\\1:int/0:string, string\\1:int\\0:string

    Consider a specifc type string\\0:int/1:string. It has the following fields:

    - return_type: string
    - argument_types: [int, string]
    - linearization: [LinearizationTuple(0, CCGCompositionDirection.LEFT), LinearizationTuple(1, CCGCompositionDirection.RIGHT)]

    Here, the first argument should come from fapp, of type string. The second argument should come from bapp,
    of type int. The linearization[i].index is never used in the computation. However, it does record the argument
    order mapping from the original "func(string, int) -> string".

    See also :meth:`NeuralCCGSyntaxType.iter_from_function`.

    See also :meth:`concepts.language.neural_ccg.search.NeuralCCGLexiconSearcher.gen_lexicons`.
    """

    return_type: Union[None, str, ObjectType, ValueType]
    """The return type of the syntax."""

    argument_types: Tuple[Union[ObjectType, ValueType, FunctionType], ...]
    """The argument types of the syntax."""

    linearization: Tuple[LinearizationTuple, ...]
    """The linearization of arguments. See the docstring for details."""

    lang_syntax_type: Optional[CCGSyntaxType]
    """Optionally, a :class:`~concepts.language.ccg.CCGSyntaxType` can be associated with the :class:`NeuralCCGSyntaxType`."""

    def __init__(
        self,
        return_type: Union[None, str, ObjectType, ValueType],
        argument_types: Optional[Iterable[Union[ObjectType, ValueType, FunctionType]]] = None,
        linearization: Optional[Iterable[LinearizationTuple]] = None,
        lang_syntax_type: Optional[CCGSyntaxType] = None,
        function_typename: Optional[str] = None
    ):
        self.return_type = return_type
        self.argument_types = tuple(argument_types) if argument_types is not None else tuple()
        self.linearization = tuple(linearization) if linearization is not None else tuple()
        self.lang_syntax_type = lang_syntax_type
        self._function_typename = function_typename

        if self.is_function and self._function_typename is None:
            self._function_typename = FunctionType(self.argument_types, self.return_type).typename

    def derive_lang_syntax_type(self, lang_syntax_type: CCGSyntaxType) -> 'NeuralCCGSyntaxType':
        """Create a new syntax type that has an associated standard linguistic syntax type.

        Args:
            lang_syntax_type: the syntax type.

        Returns:
            A new :class:`NeuralCCGSyntaxType` with the same signature and linearization, but with an associated
            standard linguistic syntax type.
        """
        return NeuralCCGSyntaxType(self.return_type, self.argument_types, self.linearization, lang_syntax_type)

    @classmethod
    def iter_from_function(
        cls, function: Function, nr_used_arguments: Optional[int] = None
    ) -> Iterable[Tuple['NeuralCCGSyntaxType', Function]]:
        """
        Create all possible linearization of the corresponding function.

        Args:
            function: the function.
            nr_used_arguments: do enumeration only for the first N arguments. Useful when the function has constant
                arguments that will be bound with word types.
        Yield:
            all possible :class:`NeuralCCGSyntaxType` derived from the given function signature.
        """
        rtype = function.ftype.return_type
        if nr_used_arguments is None:
            nr_used_arguments = len(function.ftype.argument_types)
        argtypes = function.ftype.argument_types[:nr_used_arguments]

        for indices in itertools.permutations(range(len(argtypes))):
            for nr_left in range(0, 1 + len(argtypes)):
                linearization = list()
                linearization.extend([
                    LinearizationTuple(x, CCGCompositionDirection.LEFT)
                    for x in indices[:nr_left]
                ])
                linearization.extend([
                    LinearizationTuple(x, CCGCompositionDirection.RIGHT)
                    for x in indices[nr_left:]
                ])
                permuted_argtypes = [argtypes[i] for i in indices]
                new_func = function.remap_arguments(
                    [x.index for x in linearization[::-1]] +
                    list(range(nr_used_arguments, function.ftype.nr_arguments))
                )
                syntax_type = NeuralCCGSyntaxType(rtype, permuted_argtypes, linearization, function_typename=new_func.ftype.typename)
                yield syntax_type, new_func

    @cached_property
    def typename(self) -> str:
        """The functional typename. Does NOT include the linearization.

        Returns:
            The functional typename.
        """
        if self.is_none:
            return '<None>'
        elif self.is_function:
            # NB(Jiayuan Mao @ 04/05): this line should be synced up with typing.py FunctionType::_gen_typename().
            return self._function_typename
        elif self.is_value:
            return self.return_type.typename
        elif self.is_conj:
            return self.return_type

    # This property is inherited from CCGComposable.
    @property
    def is_none(self) -> bool:
        return self.return_type is None

    # This property is inherited from CCGComposable.
    @property
    def is_conj(self) -> bool:
        return False

    @property
    def arity(self) -> int:
        """The arity of the syntax type. That is, the number of arguments it needs to combine before it becomes a primitive syntax type."""
        if self.is_function:
            return len(self.argument_types)
        elif self.is_value:
            return 0
        else:
            raise AttributeError('Cannot get the arity of None syntax.')

    @property
    def is_function(self) -> bool:
        """Whether the syntax type is a function type. That is, whether it can do function application with another syntax type."""
        return not self.is_none and not self.is_conj and self.argument_types is not None and len(self.argument_types) > 0

    @property
    def is_value(self) -> bool:
        """Whether the syntax type is a value type. That is, whether it is a primitive syntax type."""
        return not self.is_none and not self.is_conj and (self.argument_types is None or len(self.argument_types) == 0)

    @_profile
    def _fapp(self, rhs: 'NeuralCCGSyntaxType') -> 'NeuralCCGSyntaxType':
        new_syntax_type = None
        if self.lang_syntax_type is not None:
            assert rhs.lang_syntax_type is not None
            new_syntax_type = self.lang_syntax_type.fapp(rhs.lang_syntax_type)

        if self.is_function:
            if self.argument_types[-1].typename == rhs.typename and self.linearization[-1].direction is CCGCompositionDirection.RIGHT:
                return NeuralCCGSyntaxType(self.return_type, self.argument_types[:-1], self.linearization[:-1], lang_syntax_type=new_syntax_type)
        raise CCGCompositionError()

    @_profile
    def _bapp(self, lhs: 'NeuralCCGSyntaxType') -> 'NeuralCCGSyntaxType':
        new_syntax_type = None
        if self.lang_syntax_type is not None:
            assert lhs.lang_syntax_type is not None
            new_syntax_type = self.lang_syntax_type.bapp(lhs.lang_syntax_type)

        if self.is_function:
            if self.argument_types[-1].typename == lhs.typename and self.linearization[-1].direction is CCGCompositionDirection.LEFT:
                return NeuralCCGSyntaxType(self.return_type, self.argument_types[:-1], self.linearization[:-1], lang_syntax_type=new_syntax_type)
        raise CCGCompositionError()

    def _coord3(self, lhs: 'NeuralCCGSyntaxType', rhs: 'NeuralCCGSyntaxType') -> 'NeuralCCGSyntaxType':
        raise NotImplementedError('Coordination is not supported for primitive neural CCG syntax types. Use NeuralCCGConjSyntaxType instead.')

    def __str__(self) -> str:
        if self.is_none:
            return '<None>'
        if isinstance(self.return_type, str):
            fmt = self.return_type
        else:
            fmt = self.return_type.short_str()
        if self.is_function:
            fmt = self.return_type.short_str()
            for arg, lin in zip(self.argument_types, self.linearization):
                if lin.direction is CCGCompositionDirection.LEFT:
                    fmt = fmt + '\\' + arg.short_str()
                else:
                    fmt = fmt + '/' + arg.short_str()
        if self.lang_syntax_type is not None:
            fmt += '[' + str(self.lang_syntax_type) + ']'
        return fmt

    __repr__ = repr_from_str

    def __eq__(self, other: 'NeuralCCGSyntaxType') -> bool:
        return (
            self.return_type == other.return_type and
            self.argument_types == other.argument_types and
            self.linearization == other.linearization
        )

    def __hash__(self):
        return hash(str(self))

    def __lt__(self, other: 'NeuralCCGSyntaxType') -> bool:
        a, b = str(self), str(other)
        return (a.count('/') + a.count('\\'), a) < (b.count('/') + b.count('\\'), b)


class NeuralCCGConjSyntaxType(NeuralCCGSyntaxType):
    """Data structure for conjunction syntax types. Conjunction syntax types are represented as a single string (stored in ``return_type``."""

    return_type: str
    """The return type of the syntax type. For conjunction syntax types, this is the string name of the conjunction."""

    argument_types: Tuple[TypeBase, ...]

    linearization: Tuple[LinearizationTuple, ...]

    lang_syntax_type: Optional[CCGSyntaxType]

    def __init__(self, typename: str):
        super().__init__(typename)

    @property
    def is_conj(self):
        return True

    def __call__(self, lhs: NeuralCCGSyntaxType, rhs: NeuralCCGSyntaxType) -> NeuralCCGSyntaxType:
        """Perform conjunction of two syntax types."""
        return lhs

    @_profile
    def _coord3(self, lhs: NeuralCCGSyntaxType, rhs: NeuralCCGSyntaxType) -> NeuralCCGSyntaxType:
        if lhs == rhs:
            return self(lhs, rhs)
        raise CCGSyntaxCompositionError()


@dataclass
class NeuralCCGSemanticsPartialTypeLex(object):
    """A data structure that will be used in representing the semantic type of a partially applied function."""

    candidate_lexicon_index: int
    """The index of the candidate lexicon entry."""

    word_index: Optional[int]
    """The index of the word in the sentence. This will only be set if the lexicon entry contains constants."""

    def __str__(self):
        return f'Lex({self.candidate_lexicon_index}, word_binding={self.word_index})'

    def __repr__(self):
        return str(self)


class NeuralCCGSemantics(CCGSemantics):
    """
    We implement the execution result of the current subtree in the following way:
    For primitive types (syntax.is_value == True), the execution_buffer field is a list containing a single element
    recording the execution result.
    For function types (syntax.is_function == True), the execution_buffer stores a partially bound function application.
    It is stored as a list, where the first element is the function and the rest are bound arguments (the order of
    these elements is decided by syntax.linearization).

    There are two auxiliary properties of this class:

    - `partial_type`: It summarizes the intermediate execution result of the composition. Specifically, if there is only one string in the list,
        the string is the type of the semantics LF. Otherwise, the first element is of type NeuralCCGSemanticsPartialTypeLex,
        which records the candidate lexicon index and the word index (which can be None).
    - `nr_execution_steps`: the number of forwarding steps that have been executed during the construction of this subtree.
    """

    def __init__(
        self,
        value: Union[None, Callable, Function, ConstantExpression, FunctionApplicationExpression, CCGSemanticsLazyValue],
        execution_buffer: Optional[List[Union[None, Callable, Value]]] = None,
        partial_type: Optional[List[str]] = None,
        nr_execution_steps: int = 0,
        is_conj: bool = False
    ):
        """Initialize a neural CCG semantics object.

        Args:
            value: the value of the semantics.
            execution_buffer: the buffer for partially executed values. See the docstring for details.
            partial_type: the partial type of the semantics. See the docstring for details.
            nr_execution_steps: the number of execution steps that have been executed.
            is_conj: whether the semantics is a conjunction.
        """
        super().__init__(value, is_conj=is_conj)

        self.execution_buffer = execution_buffer
        self.partial_type = partial_type
        self.nr_execution_steps = nr_execution_steps

    execution_buffer: Optional[List[Union[None, Callable, Value]]]
    """The execution buffer for partially executed values."""

    partial_type: Optional[List[str]]
    """The partial type of the semantics."""

    nr_execution_steps: int
    """The number of execution steps that have been executed."""

    @_profile
    def _fapp(self, rhs: 'NeuralCCGSemantics') -> 'NeuralCCGSemantics':
        output: NeuralCCGSemantics = super()._fapp(rhs)
        output.execution_buffer = _merge_list(self.execution_buffer, rhs.execution_buffer)
        output.partial_type = _merge_list(self.partial_type, rhs.partial_type)
        output.nr_execution_steps = self.nr_execution_steps + rhs.nr_execution_steps
        return output.canonize_execution_buffer()

    @_profile
    def _bapp(self, lhs: 'NeuralCCGSemantics') -> 'NeuralCCGSemantics':
        output: NeuralCCGSemantics = super()._fapp(lhs)
        output.execution_buffer = _merge_list(self.execution_buffer, lhs.execution_buffer)
        output.partial_type = _merge_list(self.partial_type, lhs.partial_type)
        output.nr_execution_steps = self.nr_execution_steps + lhs.nr_execution_steps
        return output.canonize_execution_buffer()

    @_profile
    def _coord3(self, lhs: 'NeuralCCGSemantics', rhs: 'NeuralCCGSemantics') -> 'NeuralCCGSemantics':
        output: NeuralCCGSemantics = super()._coord3(lhs, rhs)
        assert isinstance(self.execution_buffer[0], NeuralCCGConjGroundingFunction)
        ret = self.execution_buffer[0](
            lhs.execution_buffer,
            lhs.partial_type,
            rhs.execution_buffer,
            rhs.partial_type
        )
        output.execution_buffer = ret.execution_buffer
        output.partial_type = ret.partial_type
        output.nr_execution_steps = lhs.nr_execution_steps + rhs.nr_execution_steps + ret.nr_execution_steps
        return output

    def canonize_execution_buffer(self) -> 'NeuralCCGSemantics':
        """Canonize the execution buffer. Specifically, if the execution buffer has all argument values to the function,
        we execute the function and replace the execution buffer with the result."""
        if self.execution_buffer is not None and len(self.execution_buffer) > 0:
            r0 = self.execution_buffer[0]
            if isinstance(r0, NeuralCCGGroundingFunction):
                if len(self.execution_buffer) - 1 == r0.nr_arguments:
                    self.execution_buffer = [r0(*self.execution_buffer[1:])]
                    self.partial_type = [get_type(self.execution_result).typename]
                    self.nr_execution_steps += 1

        return self

    def has_execution_result(self):
        """Check whether the semantics has a fully-executed result."""
        if self.execution_buffer is None:
            return False
        if len(self.execution_buffer) != 1:
            return False
        return True

    @property
    def execution_result(self):
        """Get the fully-executed result. If the result is not fully executed, raise an exception."""
        if self.execution_buffer is None:
            return None
        if len(self.execution_buffer) != 1:
            raise ValueError('The result being queried is not a single output.')
        return self.execution_buffer[0]

    def set_execution_result(self, value):
        """Set the fully-executed result.

        Args:
            value: the value to be set.
        """
        assert self.execution_buffer is not None
        if len(self.execution_buffer) != 1:
            raise ValueError('The result being queried is not a single output.')
        self.execution_buffer[0] = value


class NeuralCCGGroundingFunction(object):
    """A data structure contains the implementation of a function (for neural CCG semantics)."""

    def __init__(
        self,
        expression: Union[Callable, ConstantExpression, Function, FunctionApplicationExpression],
        executor: FunctionDomainExecutor,
        nr_arguments: Optional[int] = None,
        constant_arg_types: Optional[Iterable[ConstantType]] = None,
        bound_constant: Optional[str] = None,
        note: Any = '<anonymous>'
    ):
        """Initialize a neural CCG function.

        Args:
            expression: the function application expression, the function, or a python Callable function representing the semantic form of the function.
            executor: the executor for the domain.
            nr_arguments: the number of arguments to the function.
            constant_arg_types: the types of the constant arguments. Note that constant-typed arguments always come last.
            bound_constant: the constant that is bound to the function. During execution, constant-typed arguments will be filled with this constant.
            note: the note for the function. This is used for debugging.
        """
        self.expression = expression
        self.executor = executor
        if nr_arguments is None:
            if isinstance(expression, (ConstantExpression, FunctionApplicationExpression)):
                self.nr_arguments = 0
            elif isinstance(expression, Function):
                self.nr_arguments = expression.ftype.nr_arguments
            else:
                raise ValueError('The number of arguments must be specified when the expression is a Python callable.')
        else:
            self.nr_arguments = nr_arguments
        self.constant_arg_types = tuple(constant_arg_types) if constant_arg_types is not None else tuple()
        self.bound_constant = bound_constant
        self.note = note

    expression: Union[Callable, Function, FunctionApplicationExpression]
    """The function application expression, the function, or a python Callable function representing the semantic form of the function."""

    executor: FunctionDomainExecutor
    """The executor for the domain."""

    nr_arguments: int
    """The number of arguments to the function."""

    constant_arg_types: Tuple[ConstantType]
    """The types of the constant arguments. Note that constant-typed arguments always come last."""

    bound_constant: Optional[str]
    """The constant that is bound to the function. During execution, constant-typed arguments will be filled with this constant."""

    note: Any
    """The note for the function. This is used for debugging."""

    def bind_constants(self, constant: str) -> 'NeuralCCGGroundingFunction':
        """Bind a constant to the function. This function derives a new function with the same underlying expression.

        Args:
            constant: the constant to be bound.

        Returns:
            the new function with the constant bound.
        """
        return type(self)(
            self.expression, self.executor,
            self.nr_arguments, self.constant_arg_types,
            bound_constant=constant, note=self.note
        )

    def __call__(self, *args: Union[Callable, Value]) -> Value:
        """Execute the function with the given arguments.

        Args:
            *args: the arguments to the function.

        Returns:
            the result of the function.
        """
        args = list(args)
        if len(self.constant_arg_types) > 0:
            assert self.bound_constant is not None
        for arg in self.constant_arg_types:
            args.append(Value(arg, self.bound_constant))

        for i, arg in enumerate(args):
            if isinstance(arg, NeuralCCGGroundingFunction):
                args[i] = arg.expression

        if isinstance(self.expression, Function):
            retval = self.executor.execute_function(self.expression, *args)
        elif isinstance(self.expression, Callable):
            expression = self.expression(*args)
            retval = self.executor.execute(expression)
        else:
            assert len(args) == 0
            retval = self.executor.execute(self.expression)

        if foptions.compose_function_note:
            retval.note = [self.note]
            if len(self.constant_arg_types) > 0:
                retval.note.append('const:' + self.bound_constant)
            for arg in args:
                if hasattr(arg, 'note'):
                    retval.note.append(arg.note)

        return retval


@dataclass
class NeuralCCGSemanticsExecutionBuffer(object):
    """A small data structure that contains the execution buffer, partial type, and number of execution steps.
    This is used in the underlying implementation of the conjunction-typed semantics."""

    execution_buffer: Optional[List[Union[None, Callable, Value]]] = None
    partial_type: Optional[List[str]] = None
    nr_execution_steps: int = 0


class NeuralCCGConjGroundingFunction(object):
    """A wrapper for the underlying implementation of a conjunction-typed semantic form (e.g., AND)."""

    def __call__(
        self,
        left_buffer: List[Union[None, Callable, Value]],
        left_type: List[str],
        right_buffer: List[Union[None, Callable, Value]],
        right_type: List[str],
    ) -> NeuralCCGSemanticsExecutionBuffer:
        """Execute the conjunction-typed semantic form."""
        raise NotImplementedError()


class NeuralCCGSimpleConjFunction(NeuralCCGConjGroundingFunction):
    """A simple implementation for conjunction-typed semantics.

    This function takes a function that works for :class:`~concepts.dsl.value.Value` and automatically
    converts it to a function that works for :class:`NeuralCCGGroundingFunction`.

    The initializer takes a function that takes two arguments and returns a :class:`~concepts.dsl.value.Value`.
    When the conjunction lexicon is combined with two semantic forms which has a value type, this combination function
    will be called to produce the result. When the conjunction lexicon is combined with two semantic forms which has
    a function type, this combination function will be called to produce a new function:

    .. code-block:: python

        def new_function(*args):
            return combination_function(left_function(*args), right_function(*args))
    """

    def __init__(self, executor: FunctionDomainExecutor, primitive_function: Function, note='<anonymous conj>'):
        """Initialize the conjunction-typed semantics.

        Args:
            executor: the executor for the domain.
            primitive_function: the primitive function that is used to implement the conjunction-typed semantics.
            note: the note for the function. This is used for debugging.
        """

        self.executor = executor
        self.primitive_function = primitive_function
        self.note = note

    executor: FunctionDomainExecutor
    """The executor for the domain."""

    primitive_function: Function
    """The primitive function that is used to execute the conjunction-typed semantics at the primitive level."""

    note: Any
    """The note for the function. This is used for debugging."""

    def __call__(
        self,
        left_buffer: List[Union[None, Callable, Value]],
        left_type: List[str],
        right_buffer: List[Union[None, Callable, Value]],
        right_type: List[str],
    ) -> NeuralCCGSemanticsExecutionBuffer:

        lfunc = left_buffer[0]
        rfunc = right_buffer[0]
        if isinstance(lfunc, NeuralCCGGroundingFunction) and isinstance(rfunc, NeuralCCGGroundingFunction):
            nr_left = lfunc.nr_arguments - len(left_buffer) + 1
            nr_right = rfunc.nr_arguments - len(right_buffer) + 1
            assert nr_left == nr_right

            def body(*args):
                assert len(args) == nr_left == nr_right
                lval = lfunc(*left_buffer[1:], *args)
                rval = rfunc(*right_buffer[1:], *args)
                return self.executor.execute_function(self.primitive_function, lval, rval)

            return NeuralCCGSemanticsExecutionBuffer(
                [NeuralCCGGroundingFunction(body, self.executor, nr_left, note=self.note)],
                [self.note],
                0
            )

        if (
            len(left_buffer) == 1 and len(right_buffer) == 1 and
            isinstance(left_buffer[0], (Value, FunctionApplicationExpression)) and
            isinstance(right_buffer[0], (Value, FunctionApplicationExpression))
        ):
            lval, rval = left_buffer[0], right_buffer[0]
            ret = self.executor.execute_function(self.primitive_function, lval, rval)
            return NeuralCCGSemanticsExecutionBuffer([ret], [get_type(ret).typename], 1)

        raise CCGCompositionError('NeuralCCGConjunction failed.')


class NeuralCCGNode(CCGNode):
    """A node in the Neural CCG parsing tree."""

    def __init__(
        self, composition_system: CCGCompositionSystem,
        syntax: NeuralCCGSyntaxType, semantics: NeuralCCGSemantics, composition_type: CCGCompositionType,
        lexicon: Optional[Lexicon] = None,
        lhs: Optional['NeuralCCGNode'] = None, rhs: Optional['NeuralCCGNode'] = None,
        weight: Optional[Union[float, torch.Tensor]] = None,
        composition_str: str = None,
    ):
        """Initialize a CCG node.

        Args:
            composition_system: the composition system.
            syntax: the syntax of the node.
            semantics: the semantics of the node.
            composition_type: the composition type of the node.
            lexicon: the lexicon of the node (only if the composition type is LEXICON).
            lhs: the left child of the node (when available).
            rhs: the right child of the node (when available).
            weight: the weight of the node.
            composition_str: a string representation of the composition, which can be used to remove duplicated trees.
        """
        super().__init__(
            composition_system, syntax, semantics, composition_type,
            lexicon=lexicon, lhs=lhs, rhs=rhs,
            weight=weight
        )

        self.composition_str = composition_str
        if self.composition_str is None:
            self.composition_str = self._build_composition_str()
        self._used_lexicon_entries: Optional[Dict[str, Set[Lexicon]]] = None

    def _build_composition_str(self):
        if self.composition_type is CCGCompositionType.LEXICON:
            if self.is_none:
                return ''
            else:
                return f'{self.lexicon.extra[1]}:{self.lexicon.extra[2]}'

        if isinstance(self.rhs.syntax, CCGCoordinationImmNode):
            return f'({self.lhs.composition_str} {self.rhs.lhs.composition_str}, {self.rhs.rhs.composition_str})'

        if self.lhs.syntax.is_none:
            return self.rhs.composition_str
        if self.rhs.syntax.is_none:
            return self.lhs.composition_str
        return f'({self.lhs.composition_str} {self.rhs.composition_str})'

    composition_system: CCGCompositionSystem
    syntax: Optional[NeuralCCGSyntaxType]
    semantics: Optional[NeuralCCGSemantics]
    composition_type: CCGCompositionType
    lexicon: Optional[Lexicon]
    lhs: Optional['NeuralCCGNode']
    rhs: Optional['NeuralCCGNode']
    weight: Union[float, torch.Tensor]
    composition_str: str
    """A string representation of the composition inside the tree."""

    @property
    def is_none(self):
        """Whether the node is a None node."""
        return self.syntax.is_none

    @property
    def execution_result(self):
        """The execution result of the node."""
        return self.semantics.execution_result

    # NB(Jiayuan Mao @ 04/11) interface for the neural_ccg_soft implementation.
    @property
    def used_lexicon_entries(self) -> Dict[str, Set[Lexicon]]:
        """A set of lexicon entries used in the derivation."""
        return self._used_lexicon_entries

    def set_used_lexicon_entries(self, used_lexicon_entries: Dict[str, Set[Lexicon]]):
        """Set the used lexicon entries."""
        assert self._used_lexicon_entries is None
        self._used_lexicon_entries = used_lexicon_entries

    def compose_check(self, rhs: 'NeuralCCGNode', composition_type: CCGCompositionType):
        super().compose_check(rhs, composition_type)
        if composition_type is CCGCompositionType.FORWARD_APPLICATION:
            if self.syntax.is_function and isinstance(self.syntax.argument_types[-1], FunctionType):
                if rhs.composition_type is not CCGCompositionType.LEXICON:
                    raise CCGCompositionError('Does not support combination with non-lexical functor type.')
        elif composition_type is CCGCompositionType.BACKWARD_APPLICATION:
            if rhs.syntax.is_function and isinstance(rhs.syntax.argument_types[-1], FunctionType):
                if self.composition_type is not CCGCompositionType.LEXICON:
                    raise CCGCompositionError('Does not support combination with non-lexical functor type.')

    def compose_guess(self, rhs: 'NeuralCCGNode') -> Tuple[CCGCompositionType]:
        if isinstance(self.syntax, CCGCoordinationImmNode):
            return tuple()
        if isinstance(rhs.syntax, CCGCoordinationImmNode):
            return (CCGCompositionType.COORDINATION, )
        if self.syntax.is_conj:
            return (CCGCompositionType.COORDINATION, )
        if rhs.syntax.is_conj:
            return tuple()

        if self.syntax.is_none or rhs.syntax.is_none:
            return (CCGCompositionType.NONE, )
        if self.syntax.is_function:
            if self.syntax.argument_types[-1].typename == rhs.syntax.typename and self.syntax.linearization[-1].direction is CCGCompositionDirection.RIGHT:
                return (CCGCompositionType.FORWARD_APPLICATION, )
        if rhs.syntax.is_function:
            if rhs.syntax.argument_types[-1].typename == self.syntax.typename and rhs.syntax.linearization[-1].direction is CCGCompositionDirection.LEFT:
                return (CCGCompositionType.BACKWARD_APPLICATION, )

        return tuple()


def compose_neural_ccg_node(lhs: NeuralCCGNode, rhs: NeuralCCGNode, composition_type: Optional[CCGCompositionType] = None) -> Union[NeuralCCGNode, CCGCompositionResult]:
    """Compose two neural CCG nodes.

    Args:
        lhs: the left node.
        rhs: the right node.

    Returns:
        The composed node.
    """
    return lhs.compose(rhs, composition_type=composition_type)


class NeuralCCG(nn.Module):
    """The neural CCG grammar and the implementation for parsing."""

    training: bool

    def __init__(
        self,
        domain: FunctionDomain,
        executor: FunctionDomainExecutor,
        candidate_lexicon_entries: Iterable['NeuralCCGLexiconSearchResult'],
        composition_system: Optional[CCGCompositionSystem] = None,
        joint_execution: bool = True,  # execute the partial programs during CKY.
        allow_none_lexicon: bool = False,
        reweight_meaning_lex: bool = False
    ):
        """Initialize the neural CCG grammar.

        Args:
            domain: the function domain of the grammar.
            executor: the executor for the function domain.
            candidate_lexicon_entries:  a list of candidate lexicon entries.
            composition_system: the composition system. If None, the default composition system will be used.
            joint_execution: whether to execute the partial programs during CKY.
            allow_none_lexicon: whether to allow None lexicon.
            reweight_meaning_lex: whether to reweight the meaning lexicon entries. Specifically, if there are two parsings
                share the same set of lexicon entries (i.e., caused by ambiguities in combination), specifying this flag
                will reweight both of them to be 1 / (number of parsings that used this set of lexicon entries).
        """
        super().__init__()

        if composition_system is None:
            composition_system = CCGCompositionSystem.make_default()

        self.domain = domain
        self.executor = executor
        self.composition_system = composition_system
        self.candidate_lexicon_entries = tuple(candidate_lexicon_entries)
        self.joint_execution = joint_execution
        self.allow_none_lexicon = allow_none_lexicon
        self.reweight_meaning_lex = reweight_meaning_lex

    domain: FunctionDomain
    """The function domain."""

    executor: FunctionDomainExecutor
    """The executor for the domain."""

    composition_system: CCGCompositionSystem
    """The composition system."""

    candidate_lexicon_entries: Tuple['NeuralCCGLexiconSearchResult', ...]
    """The candidate lexicon entries."""

    joint_execution: bool
    """Whether to execute the partial programs during CKY."""

    allow_none_lexicon: bool
    """Whether to allow None lexicon."""

    reweight_meaning_lex: bool
    """Whether to reweight the meaning lexicon entries."""

    @property
    def nr_candidate_lexicon_entries(self):
        return len(self.candidate_lexicon_entries)

    @_profile
    def parse(
        self,
        words: Union[Sequence[str], str],
        distribution_over_lexicon_entries: torch.Tensor,
        used_lexicon_entries: Optional[Dict[str, Set[int]]] = None,
        acceptable_rtypes: Optional[Sequence[TypeBase]] = None,
        max_research: int = 0
    ) -> List[NeuralCCGNode]:
        """Parse the sentence using the CKY algorithm. The function maintains a "chart" for all possible spans of the
        sentence, and for each span, it maintains a list of possible CCG nodes. When ``max_research`` is set to a positive
        integer, for each span, the function will keep a tuple of lists of CCG nodes, where ``dp[i][j][k]`` corresponds to
        candidate CCG nodes for span ``[i, j]`` with ``k`` words being re-searched.

        Args:
            words: the words of the sentence. If the input is a string, it will be tokenized using the default ``.split()`` method.
            distribution_over_lexicon_entries: the distribution over the lexicon entries for each word.
            used_lexicon_entries: the used lexicon entries for each word. If None, it will be computed from the distribution.
                This is a dictionary mapping from the word to the set of lexicon entry indices.
            acceptable_rtypes: the acceptable return types. If None, it will accept all return types.
            max_research: the maximum number of words whose lexicon entries can be re-searched (i.e., the ``used_lexicon_entries`` for this word will be ignored).

        Returns:
            The parsing result.
        """

        if isinstance(words, str):
            words = words.split()

        if not self.training:
            max_research = 0

        length = distribution_over_lexicon_entries.size(0)
        assert length == len(words)

        # Initialize the chart.
        if max_research == 0:
            dp = [[list() for _ in range(length + 1)] for _ in range(length)]
        else:
            dp = [[tuple([list() for _ in range(max_research + 1)] for _ in range(length + 1)) for _ in range(length)]]

        if foptions.debug_print:
            print('Parsing: ', words)
            print('-' * 120)

        use_lazy_semantics = False
        if self.training:
            use_lazy_semantics = foptions.training_use_lazy_semantics

        syntax_only = False
        if get_ccg_composition_context().semantics is False:
            syntax_only = True

        for i in range(length):
            dp[i][i + 1] = self._gen_lexicons(words, i, distribution_over_lexicon_entries, used_lexicon_entries, max_research=max_research)

            if self.training:
                if max_research == 0:
                    dp[i][i + 1] = self._unique_lexicon(dp[i][i + 1], syntax_only=syntax_only)
                else:
                    dp[i][i + 1] = tuple(self._unique_lexicon(x, syntax_only=syntax_only) for x in dp[i][i + 1])

            # NB(Jiayuan Mao @ 03/01): Remove "select max" at the lexical level to implement SFINAE.
            # if not self.training:
            #     dp[i][i+1] = [max(dp[i][i+1], key=lambda node: node.weight.item())] if len(dp[i][i+1]) > 0 else []
            # dp[i][i+1] = [max(dp[i][i+1], key=lambda node: node.weight.item())] if len(dp[i][i+1]) > 0 else []

            # Block :: Debug Print {{{
            if foptions.debug_print:
                print(f'Span: [{i}, {i+1})', words[i])
                with FormatContext(expr_max_length=-1).as_default():
                    for node in sorted(_maybe_flatten_list(dp[i][i + 1]), key=lambda x: x.syntax):
                        print(jacinle.indent_text(
                            f'Weight: {node.weight.item():.4f}; Semantics: {str(node.expression)}; Syntax: {str(node.syntax)}'
                        ))
            # }}} End Block :: Debug Print

        with CCGCompositionContext(semantics_lazy_composition=use_lazy_semantics, exc_verbose=False).as_default():
            for phrase_length in range(2, length + 1):
                for i in range(0, length + 1 - phrase_length):
                    j = i + phrase_length
                    for k in range(i + 1, j):
                        if max_research == 0:
                            dp[i][j].extend(self._merge(dp[i][k], dp[k][j]))
                        else:
                            for mri in range(max_research + 1):
                                for mrj in range(max_research + 1):
                                    if mri + mrj <= max_research:
                                        dp[i][j][mri + mrj].extend(self._merge(dp[i][k][mri], dp[k][j][mrj]))

                    # Block :: Debug Print {{{
                    if foptions.debug_print:
                        print(f'Span: [{i}, {j}) :: Before Unique', words[i:j])
                        with FormatContext(expr_max_length=-1).as_default():
                            for node in sorted(_maybe_flatten_list(dp[i][j]), key=lambda x: x.syntax):
                                print(jacinle.indent_text(
                                    f'Weight: {node.weight.item():.4f}; Semantics: {str(node.expression)}; Syntax: {str(node.syntax)}'
                                ))
                    # }}} End Block :: Debug Print

                    if not self.training:
                        assert max_research == 0
                        dp[i][j] = [max(dp[i][j], key=lambda node: node.weight.item())] if len(dp[i][j]) > 0 else []
                    else:
                        if max_research == 0:
                            dp[i][j] = self._unique(dp[i][j], syntax_only=syntax_only)
                        else:
                            dp[i][j] = tuple(self._unique(x, syntax_only=syntax_only) for x in dp[i][j])

                    # Block :: Debug Print {{{
                    if foptions.debug_print:
                        print(f'Span: [{i}, {j}) :: After Unique', words[i:j])
                        with FormatContext(expr_max_length=-1).as_default():
                            for node in sorted(_maybe_flatten_list(dp[i][j]), key=lambda x: x.syntax):
                                print(jacinle.indent_text(
                                    f'Weight: {node.weight.item():.4f}; Semantics: {str(node.expression)}; Syntax: {str(node.syntax)}'
                                ))
                    # }}} End Block :: Debug Print

                    if foptions.debug_stat_types:
                        count = defaultdict(int)
                        for node in _maybe_flatten_list(dp[i][j]):
                            count[str(node.syntax)] += 1
                        print(f'Span: [{i}, {j}) ::', count)

        ret = _maybe_flatten_list(dp[0][length])
        ret = self._select_acceptable_parsings(ret, acceptable_rtypes)
        if self.reweight_meaning_lex:
            ret = self._reweight_parse_trees(ret)
        ret = sorted(ret, key=lambda node: node.weight.item(), reverse=True)

        if foptions.debug_print:
            input('Press enter to continue...')
            import ipdb; ipdb.set_trace()

        return ret

    @_profile
    def parse_beamsearch(
        self,
        words: Union[Sequence[str], str],
        distribution_over_lexicon_entries: torch.Tensor,
        used_lexicon_entries: Optional[Dict[str, Set[int]]] = None,
        acceptable_rtypes: Optional[List[TypeBase]] = None,
        beam: int = 5, lexical_beam: int = 0,
        out_of_vocab_weight: Optional[float] = None
    ) -> List[NeuralCCGNode]:
        """Parse the sentence with the CKY algorithm, but with beam search. Note that this function does not support re-search for used lexicon entries,
        and it can only be used when the model is in ``eval`` mode, because the beam-search will break the correctness of the gradient computation.

        Args:
            words: the words of the sentence. If the input is a string, it will be tokenized using the default ``.split()`` method.
            distribution_over_lexicon_entries: the distribution over the lexicon entries for each word.
            used_lexicon_entries: the used lexicon entries for each word. If None, it will be computed from the distribution.
                This is a dictionary mapping from the word to the set of lexicon entry indices.
            acceptable_rtypes: the acceptable return types. If None, it will accept all return types.
            beam: the beam size.
            lexical_beam: the lexical beam size. If 0, we will keep all the lexical entries.
            out_of_vocab_weight: the weight for out-of-vocabulary words. Specifically, if this value is set, we will consider all lexical entries
                (instead of only those in the ``used_lexicon_entries``). However, the corresponding weights for those entries will be set to
                ``out_of_vocab_weight``.

        Returns:
            The parsing result, as a list of parsing trees.
        """
        assert not self.training

        length = distribution_over_lexicon_entries.size(0)
        dp = [[list() for _ in range(length + 1)] for _ in range(length)]

        if lexical_beam is None:
            lexical_beam = beam

        for i in range(length):
            dp[i][i + 1] = self._gen_lexicons(words, i, distribution_over_lexicon_entries, used_lexicon_entries, out_of_vocab_weight=out_of_vocab_weight)
            if lexical_beam > 0:
                dp[i][i + 1] = heapq.nlargest(lexical_beam, dp[i][i + 1], key=lambda node: node.weight.item()) if len(dp[i][i + 1]) > 0 else list()

        with CCGCompositionContext(semantics_lf=False, exc_verbose=False).as_default():
            for span_length in range(2, length + 1):
                for i in range(0, length + 1 - span_length):
                    j = i + span_length
                    for k in range(i + 1, j):
                        dp[i][j].extend(self._merge(dp[i][k], dp[k][j]))

                    if span_length != length:
                        if len(dp[i][j]) > beam:
                            dp[i][j] = heapq.nlargest(beam, dp[i][j], key=lambda node: node.weight.item()) if len(dp[i][j]) > 0 else []

        ret = self._select_acceptable_parsings(dp[0][length], acceptable_rtypes)
        ret = sorted(ret, key=lambda node: node.weight.item(), reverse=True)
        return ret

    def _gen_lexicons(
        self,
        words: Sequence[str],
        word_index: int,
        distribution_over_lexicon_entries: torch.Tensor,
        used_lexicon_entries: Optional[Dict[str, Set[int]]] = None,
        out_of_vocab_weight: Optional[float] = None,
        max_research: int = 0
    ) -> Union[List[NeuralCCGNode], Tuple[List[NeuralCCGNode], ...]]:
        """Generate the lexicon nodes for the given word.

        Args:
            words: the words of the sentence.
            word_index: the index of the word.
            distribution_over_lexicon_entries: the distribution over the lexicon entries for each word.
            used_lexicon_entries: the used lexicon entries for each word. It should be a dictionary mapping from the word to the set of lexicon entry indices.
            out_of_vocab_weight: the weight for out-of-vocab words. This will only take effect in the evaluation mode. See :meth:`parse_beamsearch` for more details.
            max_research: the maximum number of re-search for used lexicon entries. See :meth:`parse` for more details.

        Returns:
            When ``max_research`` is 0, it will return a list of lexicon nodes. Otherwise, it will return a tuple of lists of lexicon nodes.
        """
        word = words[word_index]
        use_used_lexicon_record = used_lexicon_entries is not None and word in used_lexicon_entries

        # When ``out_of_vocab_weight`` is None, we will still consider using all candidate lexicon entries, but
        # for candidates that are not in the ``used_lexicon_entries`` record, we will set their weight to be
        # ``out_of_vocab_weight``.
        if not self.training and out_of_vocab_weight is not None:
            use_used_lexicon_record = False

        # When ``max_research`` is not 0, we will still consider using all candidate lexicon entries.
        if self.training and max_research != 0:
            use_used_lexicon_record = False

        if use_used_lexicon_record:
            lexicon_entry_indices = used_lexicon_entries[word]
        else:
            lexicon_entry_indices = range(distribution_over_lexicon_entries.size(1))

        if max_research == 0:
            lexicon_entries = list()
        else:
            lexicon_entries = tuple(list() for _ in range(max_research + 1))

        ctx = get_ccg_composition_context()

        for j in lexicon_entry_indices:
            candidate = self.candidate_lexicon_entries[j]
            syn, sem, exe = candidate.syntax, candidate.semantics, candidate.executor

            if ctx.semantics:
                is_conj = syn.is_conj
                sem = self._bind_constants(sem, exe, word, word_index=word_index, candidate_lexicon_entry_index=j, is_conj=is_conj)
                sem = sem.canonize_execution_buffer()
            else:
                sem = None

            weight = distribution_over_lexicon_entries[word_index, j]
            if not self.training and out_of_vocab_weight is not None:
                if used_lexicon_entries is not None and word in used_lexicon_entries:
                    if j not in used_lexicon_entries[word]:
                        weight = out_of_vocab_weight

            lexicon = Lexicon(syn, sem, weight, extra=(words[word_index], word_index, j))
            node = NeuralCCGNode(
                self.composition_system,
                lexicon.syntax, lexicon.semantics, CCGCompositionType.LEXICON, lexicon=lexicon
            )

            if max_research == 0:
                lexicon_entries.append(node)
            else:
                research_index = 0
                if used_lexicon_entries is not None and words[word_index] in used_lexicon_entries:
                    if j not in used_lexicon_entries[words[word_index]]:
                        research_index = 1
                lexicon_entries[research_index].append(node)

        return lexicon_entries

    def _bind_constants(
        self,
        expression: Union[None, Callable, ConstantExpression, Function, FunctionApplicationExpression],
        compiled_executor: Optional[Callable],
        word: str,
        word_index: int,
        candidate_lexicon_entry_index: int,
        is_conj: bool
    ) -> NeuralCCGSemantics:
        """Bind a constant to a semantic form.

        Args:
            expression: the semantic form.
            compiled_executor: the compiled executor. When ``self.joint_execution`` is set to False, this argument will be None.
            word: the word.
            word_index: the word index in the sentence. This will only be used in compiling the partial type.
            candidate_lexicon_entry_index: the index of the candidate lexicon entry. This will only be used in compiling the partial type.
            is_conj: whether the semantic form is a conjunction.

        Returns:
            The semantic form with the constant bound.
        """
        if expression is None:
            return NeuralCCGSemantics(None)

        new_expression = expression
        new_buffer = compiled_executor
        new_partial_type = None
        nr_execution_steps = 0

        # NB(Jiayuan Mao @ 2022/12/11): this function does need to handle "execution" because a ``canonize_execution_buffer`` function
        # will be called in ``_gen_lexicons`` immediately after this function.
        if isinstance(expression, (ConstantExpression, FunctionApplicationExpression)):
            pass
        elif isinstance(expression, Function):
            bindings = dict()
            for denotation, arg in enumerate(expression.ftype.argument_types):
                if isinstance(arg, ConstantType):
                    bindings['#' + str(denotation)] = Value(arg, word)

            if len(bindings) == 0:
                new_partial_type = [NeuralCCGSemanticsPartialTypeLex(candidate_lexicon_entry_index, None)]
            else:
                new_expression = expression.partial(**bindings, execute_fully_bound_functions=True)
                new_buffer = compiled_executor.bind_constants(word) if compiled_executor is not None else None
                new_partial_type = [NeuralCCGSemanticsPartialTypeLex(candidate_lexicon_entry_index, new_partial_type)]
        else:
            new_partial_type = [NeuralCCGSemanticsPartialTypeLex(candidate_lexicon_entry_index, None)]

        return NeuralCCGSemantics(
            new_expression,
            execution_buffer=[new_buffer],
            partial_type=new_partial_type,
            nr_execution_steps=nr_execution_steps,
            is_conj=is_conj
        )

    def _merge(self, list1, list2):
        """The merge function for chart parsing."""
        # TODO(Jiayuan Mao @ 2020/01/11): implement the merge sort.
        output_list = list()
        for node1, node2 in itertools.product(list1, list2):
            try:
                node = compose_neural_ccg_node(node1, node2).result
                output_list.append(node)
            except CCGCompositionError:
                pass

        return output_list

    def _unique_lexicon(self, nodes: List[NeuralCCGNode], syntax_only: bool) -> List[NeuralCCGNode]:
        """Unique the lexicon entries during the chart parsing."""
        if syntax_only:
            return self._unique_syntax_only(nodes)
        return nodes

    def _unique(self, nodes: List[NeuralCCGNode], syntax_only: bool) -> List[NeuralCCGNode]:
        """Unique the nodes for intermediate results during the chart parsing. This will only be triggered when the system allows None lexicon entriese, in which
        case different set of leaf lexicon entries will lead to the same parsing tree."""
        # NB(Jiayuan Mao @ 2020/02/10): perform a unique() operation.
        # We do not need to sum up the weights here, since the duplicates occur because of the empty semantic forms.
        if syntax_only:
            return self._unique_syntax_only(nodes)

        if self.allow_none_lexicon:
            nodes_by_composition_string = defaultdict(list)
            for v in nodes:
                nodes_by_composition_string[v.composition_str].append(v)
            nodes = [v[0] for v in nodes_by_composition_string.values()]
        return nodes

    def _unique_syntax_only(self, nodes: List[NeuralCCGNode]) -> List[NeuralCCGNode]:
        """Perform unique operation based on the syntax only. When there are multiple nodes with the same syntax, we will only keep the semantic forms for the one with the highest weight."""
        output_nodes = list()

        # a dictionary mapping from syntax types to (semantics, weight) pairs.
        nodes_by_syntax = defaultdict(lambda: (list(), list()))
        for node in nodes:
            rec = nodes_by_syntax[str(node.syntax)]
            rec[0].append(node)
            rec[1].append(node.weight)

        for typename, (this_nodes, this_weights) in nodes_by_syntax.items():
            if len(this_nodes) > 1:
                sample_node = this_nodes[0]
                weights_tensor = torch.stack(this_weights, dim=0)
                total_weights = jactorch.logsumexp(weights_tensor)
                output_nodes.append(NeuralCCGNode(
                    composition_system=sample_node.composition_system,
                    syntax=sample_node.syntax,
                    semantics=sample_node.expression,
                    composition_type=sample_node.composition_type,
                    lexicon=sample_node.lexicon, lhs=sample_node.lhs, rhs=sample_node.rhs,
                    composition_str=sample_node.composition_str,
                    weight=total_weights
                ))
                output_nodes[-1].set_used_lexicon_entries(gen_used_lexicon_entries(this_nodes))
            else:
                output_nodes.append(this_nodes[0])
        return output_nodes

    def _select_acceptable_parsings(self, parsings: List[NeuralCCGNode], acceptable_rtypes: Sequence[TypeBase]) -> List[NeuralCCGNode]:
        """Select a subset of the parsing trees that are whose final types are in the acceptable_rtypes."""
        if acceptable_rtypes is None:
            return [node for node in parsings if not node.syntax.is_none]
        parsings_filtered = list()
        for node in parsings:
            if not isinstance(node.syntax, CCGCoordinationImmNode) and not node.syntax.is_none and node.syntax.is_value and node.syntax.return_type in acceptable_rtypes:
                parsings_filtered.append(node)
        return parsings_filtered

    def _reweight_parse_trees(self, parsings: List[NeuralCCGNode]) -> List[NeuralCCGNode]:
        """Reweight the parse trees based on lexicon entries. Specifically, if there are multiple parsing trees with the same leaf-level lexicon
        entries, their weight will be divided by the number of such trees."""
        lexicon_strs = list()
        for p in parsings:
            lexicon_strs.append(p.composition_str.replace('(', '').replace(')', ''))
        counter = Counter(lexicon_strs)
        for p, ls in zip(parsings, lexicon_strs):
            if counter[ls] > 1:
                # p.weight = p.weight * 1/(counter[ls])
                p.weight = p.weight - math.log(counter[ls])
        return parsings

    def format_lexicon_table_sentence(self, words, distribution_over_lexicon_entries, max_entries=None):
        return self.format_lexicons_table(None, distribution_over_lexicon_entries_words=distribution_over_lexicon_entries, words=words, max_entries=max_entries)

    def format_lexicons_table(
        self,
        vocab,
        distribution_over_lexicon_entries=None,
        distribution_over_lexicon_entries_words=None,
        *,
        used_lexicon_weights=None, words=None, full=False,
        max_entries=None, print_grad=True
    ):
        if words is None:
            if full or used_lexicon_weights is None:
                words = vocab
            else:
                words = used_lexicon_weights.keys()

        table = list()

        if distribution_over_lexicon_entries is not None:
            has_grad = print_grad and distribution_over_lexicon_entries.grad is not None
        else:
            assert distribution_over_lexicon_entries_words is not None
            has_grad = print_grad and distribution_over_lexicon_entries_words.grad is not None

        # if print_grad and not has_grad:
        #     logger.warning('print_grad is specified for lexicon table but gradient information is not available.')

        for word_idx, word in enumerate(words):
            if word[0] == '<' and word[-1] == '>':
                continue

            if used_lexicon_weights is not None and word in used_lexicon_weights:
                lexicon_indices = list(used_lexicon_weights[word])
            else:
                lexicon_indices = range(self.nr_candidate_lexicon_entries)

            if distribution_over_lexicon_entries is not None:
                lexicon_weights = F.softmax(torch.tensor([
                    distribution_over_lexicon_entries[vocab.map(word), i].item()
                    for i in lexicon_indices
                ]), dim=-1)

                # TODO(Jiayuan Mao @ 02/10): take a softmax for visualization?
                all_lexicons = [(
                    i,
                    self.candidate_lexicon_entries[i],
                    lexicon_weights[j],
                    distribution_over_lexicon_entries.grad[vocab.map(word), i].item() if has_grad else None
                ) for j, i in enumerate(lexicon_indices)]
            else:
                lexicon_weights = F.softmax(torch.tensor([
                    distribution_over_lexicon_entries_words[word_idx, i].item()
                    for i in lexicon_indices
                ]), dim=-1)

                # TODO(Jiayuan Mao @ 02/10): take a softmax for visualization?
                all_lexicons = [(
                    i,
                    self.candidate_lexicon_entries[i],
                    lexicon_weights[j],
                    distribution_over_lexicon_entries_words.grad[word_idx, i].item() if has_grad else None
                ) for j, i in enumerate(lexicon_indices)]

            all_lexicons.sort(key=lambda x: x[2], reverse=True)
            if max_entries is not None:
                all_lexicons = all_lexicons[:max_entries]

            first = word
            for idx, lexicon, score, grad in all_lexicons:
                if print_grad:
                    table.append((first, idx, '{:3.4f}'.format(score), ('g: {:3.4f}'.format(grad) if grad is not None else 'g: None'), str(lexicon.syntax), str(lexicon.semantics)))
                else:
                    table.append((first, idx, '{:3.4f}'.format(score), str(lexicon.syntax), str(lexicon.semantics)))

                first = ''

        if print_grad:
            return tabulate(table, headers=['', 'i', 'weight', 'w_grad', 'syntax', 'semantics'])
        else:
            return tabulate(table, headers=['', 'i', 'weight', 'syntax', 'semantics'])


def _merge_list(*lists: Optional[List[T]]) -> List[T]:
    output = list()
    for x in lists:
        if x is not None:
            output.extend(x)
    return output


def _maybe_flatten_list(tuple_of_list: Union[List[T], Tuple[List[T], ...]]) -> List[T]:
    if isinstance(tuple_of_list, list):
        return tuple_of_list
    output = list()
    for x in tuple_of_list:
        output.extend(x)
    return output


def gen_used_lexicon_entries(parsings: Iterable[NeuralCCGNode]) -> Dict[str, Set[int]]:
    """Generate a dictionary of used lexicon entries from a list of parsing results. The result is a mapping from words to lexicon entry indices.

    Args:
        parsings: a list of parsing results.

    Returns:
        A dictionary of used lexicon entries.
    """
    used_lexicon_entries = defaultdict(set)

    def dfs(ccg_node: NeuralCCGNode):
        nonlocal used_lexicon_entries

        # NB(Jiayuan Mao @ 04/11): handles compressed nodes generated by neural soft ccg.
        if ccg_node.used_lexicon_entries is not None:
            for k, v in ccg_node.used_lexicon_entries.items():
                used_lexicon_entries[k].update(v)
        else:
            if ccg_node.composition_type is CCGCompositionType.LEXICON:
                word, _, index = ccg_node.lexicon.extra
                used_lexicon_entries[word].add(index)
            else:
                dfs(ccg_node.lhs)
                dfs(ccg_node.rhs)

    for p in parsings:
        dfs(p)
    return used_lexicon_entries


def count_nr_execution_steps(parsings: Iterable[NeuralCCGNode]) -> int:
    """Count the number of execution steps from a list of parsing results."""
    answer = 0

    for p in parsings:
        answer += p.expression.nr_execution_steps
    return answer


concepts/language/neural_ccg/grammar.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : search.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 04/13/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Algorithms for enumerate possible lexicon entries in Neural CCG."""

from dataclasses import dataclass
from typing import Any, Optional, Union, Iterable, Tuple, List, Callable

import jacinle
import numpy as np
from scipy.optimize import linear_sum_assignment

from concepts.dsl.dsl_types import TypeBase
from concepts.dsl.dsl_functions import FunctionType, Function
from concepts.dsl.expression import ConstantExpression, FunctionApplicationExpression
from concepts.dsl.value import Value
from concepts.dsl.executors.function_domain_executor import FunctionDomainExecutor
from concepts.dsl.learning.function_domain_search import FunctionDomainExpressionSearchResult
from concepts.language.ccg.syntax import CCGSyntaxType, CCGPrimitiveSyntaxType, CCGComposedSyntaxType
from concepts.language.ccg.search import stat_function
from concepts.language.neural_ccg.grammar import parse_linearization_string, NeuralCCGSyntaxType
from concepts.language.neural_ccg.grammar import NeuralCCGGroundingFunction

__all__ = [
    'NeuralCCGLexiconSearchResult',
    'NeuralCCGLexiconSearcherBase',
    'NeuralCCGLexiconEnumerativeSearcher',
    'NeuralCCGLexiconEnumerativeSearcherWithSyntax',
    'gen_lexicon_search_results_from_functions',
    'gen_lexicon_search_results_from_syntax_and_semantics'
]

logger = jacinle.get_logger(__file__)


@dataclass
class NeuralCCGLexiconSearchResult(object):
    """A data structure representing the search result of a neural CCG lexicon searcher."""

    syntax: NeuralCCGSyntaxType
    """The syntax type of the lexicon."""

    semantics: Union[None, ConstantExpression, Function, FunctionApplicationExpression]
    """The semantic form of the lexicon."""

    executor: Optional[NeuralCCGGroundingFunction]
    """Optionally, the executor for the semantic form."""


class NeuralCCGLexiconSearcherBase(object):
    """The base class for neural CCG lexicon searchers."""

    def gen(self) -> List[NeuralCCGLexiconSearchResult]:
        """Generate a list of candidate lexicons.

        Returns:
            A list of candidate lexicons, as a list of :class:`NeuralCCGLexiconSearchResult`.
        """
        raise NotImplementedError()


class NeuralCCGLexiconEnumerativeSearcher(NeuralCCGLexiconSearcherBase):
    """An enumerative neural CCG lexicon searcher. This function takes a list of candidate function expressions, and generates a list of candidate lexicon entries."""

    def __init__(self, candidate_expressions: List[FunctionDomainExpressionSearchResult], executor: FunctionDomainExecutor):
        """Initialize the searcher.

        Args:
            candidate_expressions: a list of candidate function expressions (including constants, functions, and function applications).
            executor: the executor for the semantic forms.
        """
        self._candidate_expressions = candidate_expressions
        self._executor = executor

    def gen(
        self,
        init_executor: bool = True,
        allow_none_lexicon: bool = False,
        permute_arguments: bool = False,
        semantics_hash: Callable[[Union[Value, Function, FunctionApplicationExpression]], Any] = None
    ) -> List[NeuralCCGLexiconSearchResult]:
        def gen() -> Iterable[Tuple[NeuralCCGSyntaxType, Union[None, ConstantExpression, Function, FunctionApplicationExpression]]]:
            for result in self._candidate_expressions:
                expression = result.expression
                if isinstance(expression, ConstantExpression):
                    syntax_type = NeuralCCGSyntaxType(expression.return_type)
                    yield syntax_type, expression
                if isinstance(expression, FunctionApplicationExpression):
                    syntax_type = NeuralCCGSyntaxType(expression.return_type)
                    yield syntax_type, expression
                elif isinstance(expression, Function):
                    for syntax_type, sub_function in NeuralCCGSyntaxType.iter_from_function(expression, result.nr_variable_arguments):
                        yield syntax_type, sub_function
                else:
                    raise TypeError(f'Invalid semantics: {expression}.')

            if allow_none_lexicon:
                yield NeuralCCGSyntaxType(None), None

        results = gen_lexicon_search_results_from_syntax_and_semantics(
            gen(), self._executor, init_executor=init_executor
        )
        results = self._unique_semantics(results, semantics_hash)
        return results

    def _unique_semantics(self, input_list: List[NeuralCCGLexiconSearchResult], hash_func: Optional[Callable[[Union[Value, Function, FunctionApplicationExpression]], Any]] = None) -> List[NeuralCCGLexiconSearchResult]:
        if hash_func is None:
            hash_func = str

        output_list = list()
        hash_set = set()
        for r in input_list:
            try:
                lex_hash = (str(r.syntax), hash_func(r.semantics))
                if lex_hash not in hash_set:
                    hash_set.add(lex_hash)
                    output_list.append(r)
            except TypeError:
                output_list.append(r)
        return output_list


class NeuralCCGLexiconEnumerativeSearcherWithSyntax(NeuralCCGLexiconSearcherBase):
    """An enumerative neural CCG lexicon searcher. This function also searches for the linguistic syntax type of the semantic form."""

    def __init__(
        self,
        candidate_expressions: List[FunctionDomainExpressionSearchResult],
        allowed_syntax_types: List[CCGSyntaxType],
        executor: FunctionDomainExecutor,
        primitive_syntax_compatible=Callable[[TypeBase, CCGPrimitiveSyntaxType], bool],
    ):
        """Initialize the searcher.

        Args:
            candidate_semantics (Union[list, CCGSemanticsSearcherBase]): the semantics searcher.
            allowed_syntax_types (List[CCGSyntaxType]): a list of allowed syntax types.
            executor (QSImplementation): executor.
            primitive_syntax_compatible ((TypeBase, CCGBasicSyntaxType) -> bool): an optional function.
        """

        self._basic_searcher = NeuralCCGLexiconEnumerativeSearcher(candidate_expressions, executor=executor)
        self._allowed_syntax_types = allowed_syntax_types
        self._primitive_syntax_compatible = primitive_syntax_compatible

        if self._primitive_syntax_compatible is None:
            self._primitive_syntax_compatible = lambda x, y: True

    def is_primitive_compatible(self, dsl_type: TypeBase, syntax_type: CCGSyntaxType):
        if isinstance(dsl_type, FunctionType):
            syntax_type = syntax_type.flatten()
            if len(syntax_type) - 1 != dsl_type.nr_arguments:
                return False
            if not self.is_primitive_compatible(dsl_type.return_type, syntax_type[0]):
                return False

            align = np.zeros((len(dsl_type.argument_types), len(syntax_type[1:])), dtype=np.int32)
            for i, dsl_type_arg in enumerate(dsl_type.argument_types):
                for j, syntax_arg in enumerate(syntax_type[1:]):
                    align[i, j] = int(self.is_primitive_compatible(dsl_type_arg, syntax_arg[0]))
            if align[linear_sum_assignment(align, maximize=True)].sum() == len(dsl_type.argument_types):  # max matching exists
                return True
        else:
            if not isinstance(syntax_type, CCGPrimitiveSyntaxType):
                return False
            return self._primitive_syntax_compatible(dsl_type, syntax_type)

    def is_compatible(self, semantics_type: NeuralCCGSyntaxType, syntax_type: CCGSyntaxType):
        if isinstance(syntax_type, CCGPrimitiveSyntaxType):
            if semantics_type.is_function:
                return False
            return self.is_primitive_compatible(semantics_type.return_type, syntax_type)
        else:
            assert isinstance(syntax_type, CCGComposedSyntaxType)
            if not semantics_type.is_function:
                return False
            if syntax_type.direction != semantics_type.linearization[-1].direction:
                return False
            if not self.is_primitive_compatible(semantics_type.argument_types[-1], syntax_type.sub):
                return False
            new_semantics_type = NeuralCCGSyntaxType(
                semantics_type.return_type,
                argument_types=semantics_type.argument_types[:-1], linearization=semantics_type.linearization[:-1]
            )
            return self.is_compatible(new_semantics_type, syntax_type.main)

    def gen(
        self,
        init_executor: bool = True,
        allow_none_lexicon: bool = False,
        semantics_hash: Callable[[Union[Value, Function, FunctionApplicationExpression]], Any] = None
    ) -> List[NeuralCCGLexiconSearchResult]:
        without_syntax = self._basic_searcher.gen(init_executor=init_executor, allow_none_lexicon=allow_none_lexicon, semantics_hash=semantics_hash)
        with_syntax = list()
        for x in without_syntax:
            for y in self._allowed_syntax_types:
                if self.is_compatible(x.syntax, y):
                    with_syntax.append(NeuralCCGLexiconSearchResult(x.syntax.derive_lang_syntax_type(y), x.semantics, x.executor))
        return with_syntax


def gen_lexicon_search_results_from_functions(
    functions: Iterable[Tuple[Union[ConstantExpression, Function, FunctionApplicationExpression], str]],
    executor: FunctionDomainExecutor,
    init_executor: bool = True,
    allow_none_lexicon: bool = False,
) -> List[NeuralCCGLexiconSearchResult]:
    """Generate lexicon search results from a given list of functions and the corresponding linearization strings.

    Args:
        functions: a list of functions and the corresponding linearization strings.
        executor: the executor for the domain.
        init_executor: whether to initialize the executor.
        allow_none_lexicon: whether to allow None lexicon.

    Returns:
        a list of lexicon search results.
    """
    def gen() -> Iterable[Tuple[NeuralCCGSyntaxType, Union[None, ConstantExpression, Function, FunctionApplicationExpression]]]:
        for f, linearization_string in functions:
            if isinstance(f, ConstantExpression):
                yield f.return_type, f
            if isinstance(f, FunctionApplicationExpression):
                yield NeuralCCGSyntaxType(f.return_type), f
            elif isinstance(f, Function):
                f_stat = stat_function(f)
                linearization = parse_linearization_string(linearization_string)
                syntax_type = NeuralCCGSyntaxType(
                    f.ftype.return_type,
                    f.ftype.argument_types[:f_stat.nr_variable_arguments],
                    linearization
                )
                sub_f = f.remap_arguments(
                    [x.index for x in syntax_type.linearization[::-1]] +
                    list(range(f_stat.nr_variable_arguments, f_stat.nr_variable_arguments + f_stat.nr_constant_arguments))
                )
                yield syntax_type, sub_f
        if allow_none_lexicon:
            yield NeuralCCGSyntaxType(None), None

    return gen_lexicon_search_results_from_syntax_and_semantics(
        gen(), executor, init_executor=init_executor
    )


def gen_lexicon_search_results_from_syntax_and_semantics(
    results: Iterable[Tuple[NeuralCCGSyntaxType, Union[None, ConstantExpression, Function, FunctionApplicationExpression]]],
    executor: FunctionDomainExecutor,
    init_executor: bool = True,
) -> List[NeuralCCGLexiconSearchResult]:
    """Generate lexicon search results from a given list of syntax and semantic forms.

    Args:
        results: a list of syntax and semantic forms.
        executor: the executor for the domain.
        init_executor: whether to initialize the executor.

    Returns:
        A list of lexicon search results.
    """
    outputs = list()
    for syntax_type, expression in results:
        if expression is None:
            outputs.append(NeuralCCGLexiconSearchResult(syntax_type, None, None))
        elif isinstance(expression, ConstantExpression):
            exe = expression.constant if init_executor else None
            outputs.append(NeuralCCGLexiconSearchResult(syntax_type, expression, exe))
        if isinstance(expression, FunctionApplicationExpression):
            exe = NeuralCCGGroundingFunction(expression, executor, nr_arguments=0, constant_arg_types=tuple(), note=str(expression)) if init_executor else None
            outputs.append(NeuralCCGLexiconSearchResult(syntax_type, expression, exe))
        elif isinstance(expression, Function):
            exe = NeuralCCGGroundingFunction(
                expression, executor, syntax_type.arity,
                expression.ftype.argument_types[syntax_type.arity:],
                note=str(expression)
            ) if init_executor else None
            outputs.append(NeuralCCGLexiconSearchResult(syntax_type, expression, exe))
        else:
            raise TypeError(f'Unknown expression type: {type(expression)}.')

    return outputs

concepts/language/neural_ccg/search.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : tensorized_syntax.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 06/21/2021
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import heapq
import collections
from typing import Optional, Union, Iterable, Tuple, List, Dict, Callable

import torch
import torch.nn as nn

import jactorch
from jacinle.utils.cache import cached_property

from concepts.language.ccg.composition import CCGCompositionType, CCGCompositionContext, CCGCompositionError, CCGCompositionSystem
from concepts.language.ccg.syntax import CCGSyntaxType
from concepts.language.ccg.grammar import Lexicon, CCGNode

__all__ = ['TensorizedSyntaxOnlyCCG']


class TensorizedSyntaxOnlyCCG(nn.Module):
    """Data structures and parsing implementation for a syntax-only CCG grammar. It uses a tensor-based representation for the chart, therefore is highly parallelizable."""

    def __init__(self, candidate_syntax_types: Iterable[CCGSyntaxType], composition_system: Optional[CCGCompositionSystem] = None):
        """Initialize the CCG grammar.

        Args:
            candidate_syntax_types: the candidate syntax types.
            composition_system: the composition system. If None, the default composition system will be used.
        """
        super().__init__()

        self.candidate_syntax_types = tuple(candidate_syntax_types)
        self.composition_system = composition_system
        if self.composition_system is None:
            self.composition_system = CCGCompositionSystem.make_default()

    training: bool

    candidate_syntax_types: Tuple[CCGSyntaxType, ...]
    """A list of candidate syntax types."""

    composition_system: CCGCompositionSystem
    """The composition system."""

    @cached_property
    def syntax_name2index(self) -> Dict[str, int]:
        """A mapping from syntax type name to the corresponding index."""
        return {
            (s.typename if isinstance(s, CCGSyntaxType) else s): i
            for i, s in enumerate(self.candidate_syntax_types)
        }

    def parse(
        self,
        words: torch.Tensor,  # [batch_size, max_seqlen]
        words_length: torch.Tensor,  # [batch_size]
        distribution_over_syntax_types: torch.Tensor,  # [batch_size, max_seqlen, nr_types],
        transition_matrix: torch.Tensor,  # [A: nr_types, B: nr_types, C: nr_types]  AB -> C
        allowed_root_type_indices: Union[int, List[int]]
    ) -> torch.Tensor:
        """Parse a batch of sentences.

        Args:
            words: the word indices, should have shape ``[batch_size, max_seqlen]``.
            words_length: the length of each sentence, should have shape ``[batch_size]``.
            distribution_over_syntax_types: the distribution over syntax types for each word, should have shape ``[batch_size, max_seqlen, nr_types]``.
            transition_matrix: the transition matrix, should have shape [A: nr_types, B: nr_types, C: nr_types]. Basically, each entry [A, B, C] means the probability of composing a node of type A and a node of type B to a node of type C.
            allowed_root_type_indices: the indices of the syntax types that are allowed to be the root of the tree. It can be either a single integer or a list of integers.

        Returns:
            The root nodes of the trees, with shape ``[batch_size]`` (when ``allowed_root_type_indices`` is a single integer) or ``[batch_size, nr_allowed_root_types]`` (when ``allowed_root_type_indices`` is a list of integers).
            Each entry i can be interpreted as the log-probability of the root node being syntax type ``allowed_root_type_indices[i]``.
        """
        batch_size, length = words.size()

        dp = [[None for _ in range(length + 1)] for _ in range(length)]

        for i in range(length):
            dp[i][i + 1] = distribution_over_syntax_types[:, i]

        for span_length in range(2, length + 1):
            for i in range(0, length + 1 - span_length):
                j = i + span_length

                scores = list()
                for k in range(i + 1, j):
                    scores.append(self.merge(dp[i][k], dp[k][j], transition_matrix))

                scores = torch.stack(scores, dim=-1)  # [batch_size, nr_types, K]
                scores = jactorch.logsumexp(scores, dim=-1, keepdim=False)
                dp[i][j] = scores

        ret = list()
        for i in range(words_length.size(0)):
            ret.append(dp[0][words_length[i].item()][i])
        ret = torch.stack(ret, dim=0)

        if isinstance(allowed_root_type_indices, int):
            ret = ret[:, allowed_root_type_indices]
        else:
            ret = jactorch.logsumexp(ret[:, allowed_root_type_indices], dim=-1)

        return ret

    @jactorch.no_grad_func
    def parse_beamsearch(
        self,
        words: torch.Tensor,
        words_length: torch.Tensor,  # [batch_size]
        words_tokenized: List[List[str]],
        distribution_over_syntax_types: torch.Tensor,  # [batch_size, max_seqlen, nr_types],
        transition_matrix: torch.Tensor,  # [A: nr_types, B: nr_types, C: nr_types]  AB -> C
        allowed_root_types: List[str],
        beam_size: Optional[int] = 3
    ) -> List[List[CCGNode]]:
        """Parse a batch of sentences using beam search. This function will return candidate parsing trees (instead of just the probability
        distribution of the root note). Therefore, this function is significantly slower than :meth:`parse` and should not be used for training.

        Args:
            words: the word indices, should have shape ``[batch_size, max_seqlen]``.
            words_length: the length of each sentence, should have shape ``[batch_size]``.
            words_tokenized: the tokenized words.
            distribution_over_syntax_types: the distribution over syntax types for each word, should have shape ``[batch_size, max_seqlen, nr_types]``.
            transition_matrix: the transition matrix, should have shape [A: nr_types, B: nr_types, C: nr_types]. Basically, each entry [A, B, C] means the probability of composing a node of type A and a node of type B to a node of type C.
            allowed_root_types: the syntax types that are allowed to be the root of the tree.
            beam_size: the beam size.

        Returns:
            A list of candidate parsing trees for each sentence. Each entry is a list of trees, where each tree is a :class:`CCGNode` object.
        """
        assert not self.training

        batch_size = words.size(0)
        batch_ret = list()

        with CCGCompositionContext(semantics=False).as_default():
            for batch_index in range(batch_size):
                length = words_length[batch_index].item()
                dp = [[None for _ in range(length + 1)] for _ in range(length)]

                for i in range(length):
                    dp[i][i + 1] = list()
                    for j in range(1, distribution_over_syntax_types.size(2)):  # skip <PAD>
                        weight = distribution_over_syntax_types[batch_index, i, j]
                        dp[i][i + 1].append(CCGNode(
                            self.composition_system,
                            syntax=self.candidate_syntax_types[j],
                            semantics=None,
                            composition_type=CCGCompositionType.LEXICON,
                            lexicon=Lexicon(
                                self.candidate_syntax_types[j],
                                None,
                                weight=weight,
                                extra=(words_tokenized[batch_index][i], j)
                            ),
                            weight=weight
                        ))

                    dp[i][i + 1] = self._beamsearch_nlargest(beam_size, dp[i][i + 1], key_func=lambda node: node.weight.item())

                for span_length in range(2, length + 1):
                    for i in range(0, length + 1 - span_length):
                        j = i + span_length

                        dp[i][j] = list()
                        for k in range(i + 1, j):
                            dp[i][j].extend(self.merge_beamsearch(dp[i][k], dp[k][j], transition_matrix))
                        dp[i][j] = self._beamsearch_nlargest(beam_size, dp[i][j], key_func=lambda node: node.weight.item())

                ret = dp[0][length]
                ret = list(filter(lambda node: node.syntax.typename in allowed_root_types, ret))
                batch_ret.append(ret)

        return batch_ret

    def _beamsearch_nlargest(self, beam_size: int, input_list: List[CCGNode], key_func: Callable[[CCGNode], float]) -> List[CCGNode]:
        """Return the nlargest elements in the list, using the key function to determine the order."""
        input_by_type = collections.defaultdict(list)
        for node in input_list:
            input_by_type[node.syntax.typename].append(node)
        for key, value in input_by_type.items():
            input_by_type[key] = heapq.nlargest(beam_size, value, key_func)
        return [v for value in input_by_type.values() for v in value]

    def merge(self, A: torch.Tensor, B: torch.Tensor, transition_matrix: torch.Tensor):
        """The underlying merge function for chart parsing.

        Args:
            A: the distribution over syntax types for the left child, should have shape ``[batch_size, nr_types]``.
            B: the distribution over syntax types for the right child, should have shape ``[batch_size, nr_types]``.
            transition_matrix: the transition matrix, should have shape [A: nr_types, B: nr_types, C: nr_types]. Basically, each entry [A, B, C] means the probability of composing a node of type A and a node of type B to a node of type C.

        Returns:
            The distribution over syntax types for the merged node, should have shape ``[batch_size, nr_types]``.
        """
        # A : [batch_size, nr_types]
        # B : [batch_size, nr_types]
        # transition_matrix: [A: nr_types, B: nr_types, C: nr_types]

        # return _cky_merge(A, B, transition_matrix)

        T = transition_matrix
        T = jactorch.logmatmulexp(A, T)
        T = jactorch.batch_logmatmulexp(B.unsqueeze(1), T).squeeze(1)
        return T

    def merge_beamsearch(self, a_list: List[CCGNode], b_list: List[CCGNode], transition_matrix: torch.Tensor) -> List[CCGNode]:
        """The underlying merge function for chart parsing with beam search.

        Args:
            a_list: the list of nodes for the left child.
            b_list: the list of nodes for the right child.
            transition_matrix: the transition matrix, should have shape [A: nr_types, B: nr_types, C: nr_types]. Basically, each entry [A, B, C] means the probability of composing a node of type A and a node of type B to a node of type C.

        Returns:
            The list of merged nodes.
        """
        for x in a_list:
            for y in b_list:
                try:
                    _, z = x.compose(y)
                except CCGCompositionError:
                    continue

                z.weight = x.weight + y.weight + transition_matrix[
                    self.syntax_name2index[x.syntax.typename],
                    self.syntax_name2index[y.syntax.typename],
                    self.syntax_name2index[z.syntax.typename]
                ]

                yield z


@torch.jit.script
def _cky_logsumexp(tensor: torch.Tensor, dim: int = 0) -> torch.Tensor:
    tensor_max = tensor.max(dim=dim, keepdim=True)[0]
    tensor = tensor - tensor_max
    tensor_max = tensor_max.squeeze(dim)
    tensor = tensor.exp().sum(dim=dim).log() + tensor_max
    return tensor


@torch.jit.script
def _cky_merge(A: torch.Tensor, B: torch.Tensor, T: torch.Tensor) -> torch.Tensor:
    batch_size, nr_types = A.size()
    C = torch.zeros_like(A)
    for bs in range(batch_size):
        c = A[bs].unsqueeze(-1).unsqueeze(-1) + B[bs].unsqueeze(-1).unsqueeze(0) + T
        c = c.reshape(nr_types * nr_types, nr_types)
        c = _cky_logsumexp(c)
        C[bs, :] = c
    return C


concepts/language/neural_ccg/tensorized_syntax.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/13/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Neural network modules that are used in concept learning, such as logic rule modules."""

concepts/nn/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Data structures and algorithms for model learning and task and motion planning.

The PDSketch is a sketch extension of PDDL.
"""

__all__ = [
    'Predicate', 'flatten_expression', 'get_used_state_variables',
    'Precondition', 'Effect', 'Operator', 'OperatorApplier', 'gen_all_grounded_actions', 'gen_all_partially_grounded_actions',
    'Domain', 'Problem', 'State',
    'load_domain_file', 'load_domain_string', 'load_problem_file', 'load_problem_string', 'parse_expression',
    'PDSketchExecutor', 'config_function_implementation', 'StateDefinitionHelper',
    'csp_solvers', 'planners', 'strips'
]

from .predicate import Predicate, flatten_expression, get_used_state_variables
from .operator import Precondition, Effect, Operator, OperatorApplier, gen_all_grounded_actions, gen_all_partially_grounded_actions
from .domain import Domain, Problem, State
from .parsers.pdsketch_parser import PDSketchParser, load_domain_file, load_domain_string, load_problem_file, load_problem_string, parse_expression
from .executor import PDSketchExecutor, config_function_implementation, StateDefinitionHelper, GeneratorManager
from .simulator_interface import PDSketchSimulatorInterface

from . import csp_solvers
from . import planners
from . import strips

# from .value import *
# from .optimistic import *
# from .state import *
# from .expr import *
# from .operator import *
# from .domain import *
# from .session import *
# from .ao_discretization import *
#
# from .parser import *
# from .planners import *
# from .csp_solvers import *
#
# from . import strips
# from . import nn
# from . import rl

concepts/pdsketch/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/24/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""This module contains CSP solvers for PDSketch.
These solvers will be primarily used in search-and-optimization problems in the task and motion planning problem,
although some strategies and solvers can be used in general domains as well.

Current solvers include:

- Brute-force sampling-based CSP solver. See :mod:`concepts.pdsketch.csp_solvers.brute_force_sampling`.
- DPLL-sampling-based CSP solver (a generalization of the DPLL algorithm to mixed discrete-continuous CSPs). See :mod:`concepts.pdsketch.csp_solvers.dpll_sampling`.
"""

from .dpll_sampling import csp_dpll_sampling_solve, csp_dpll_simplify


concepts/pdsketch/csp_solvers/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : brute_force_sampling.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 11/09/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""
This file corresponds to a brute-force sampling-based strategy for CSP solving.

Given a set of generators, and a CSP problem, we randomly sample a lot of values from the generators,
and then assign variables in the CSP problem to the sampled values (enumeratively). If the CSP problem
is satisfied, we return the assignment. Otherwise, we continue to sample and assign.

In practice, we do not directly implement this algorithm. When we are solving a task and motion planning problem
with continuous-parameterized operators, we first sample a large collection of possible continuous parameters,
and then reduce the problem into a basic discrete search problem.

The main function of this file is `generate_continuous_values`, which implements the sampling for continuous parameters
given the initial state and a collection of generators.
"""

import itertools
from typing import Tuple, List, Dict

import torch

from concepts.dsl.dsl_types import NamedTensorValueType
from concepts.dsl.tensor_value import TensorValue
from concepts.dsl.tensor_state import TensorState
from concepts.pdsketch.generator import Generator
from concepts.pdsketch.executor import PDSketchExecutor

__all__ = ['extract_generator_data', 'ContinuousValueDict', 'generate_continuous_values', 'expand_continuous_values_']


def extract_generator_data(executor: PDSketchExecutor, state: TensorState, generator: Generator) -> Tuple[List[TensorValue], List[TensorValue]]:
    """Extract the positive data from the state for a given generator. Specifically, it will test the `certifies` condition
    of the generator on the input state. If the condition is satisfied, it will extract the inputs (context) and outputs (generates)
    of the generator.

    Args:
        executor: a pdsketch expression executor.
        state: the input state.
        generator: the generator.

    Returns:
        a tuple of (inputs, outputs) for the generator.
    """
    result = executor.execute(generator.certifies, state, generator.arguments)
    result.tensor = torch.ge(result.tensor, 0.5)
    if result.tensor_mask is not None:
        result.tensor = torch.logical_and(result.tensor, torch.ge(result.tensor_mask, 0.5))

    def _index(value, mask):
        value = value.expand_as(mask)
        return value.tensor[mask.tensor]

    contexts = [_index(executor.execute(c, state, generator.arguments), result) for c in generator.context]
    generates = [_index(executor.execute(c, state, generator.arguments), result) for c in generator.generates]
    return contexts, generates


ContinuousValueDict = Dict[str, List[TensorValue]]


def generate_continuous_values(executor: PDSketchExecutor, state: TensorState, nr_iterations: int = 1, nr_samples: int = 5) -> ContinuousValueDict:
    """The function generate_continuous_values and expand_continuous_values jointly implements the incremental search
    algorithm for Task and Motion Planning.

    Basically, the algorithm starts from generating a large collection of possible continuous parameters by "expanding"
    from the continuous parameters in the input state. Next, it reduces the TAMP problem into a basic discrete search
    problem. The downside of this approach is that it requires grounding a large collection of possible values,
    but it is in theory probabilistically complete.

    Args:
        executor: a pdsketch expression executor.
        state: the input state.
        nr_iterations: the number of iterations.
        nr_samples: the number of samples for each generator.

    Returns:
        a dictionary of continuous values.
    """

    domain = executor.domain
    continuous_values = dict()
    for dtype in domain.types.values():
        if isinstance(dtype, NamedTensorValueType):
            continuous_values[dtype.typename] = list()

    for key, value in domain.functions.items():
        if key in state.features.all_feature_names and isinstance(value.return_type, NamedTensorValueType):
            dtype = value.return_type
            assert isinstance(dtype, NamedTensorValueType)
            feat = state.features[key].tensor
            feat = feat.reshape((-1, ) + feat.shape[-dtype.ndim():])
            continuous_values[dtype.typename].extend([TensorValue.from_tensor(x, dtype) for x in feat])

    for i in range(nr_iterations):
        expand_continuous_values_(executor, continuous_values, nr_samples=nr_samples)
    return continuous_values


def expand_continuous_values_(executor: PDSketchExecutor, current: ContinuousValueDict, nr_samples: int = 5):
    """Internal function used by :func:`generate_continuous_values`. Given the current set of continuous values,
    it tries to apply all generators using current values as their input. It inplace updates the `current` dictionary.
    with the new values.

    Args:
        executor: a pdsketch expression executor.
        current: the current continuous values.
        nr_samples: the number of samples for each generator.
    """

    domain = executor.domain
    for gen_name, gen_def in domain.generators.items():
        arguments = list()
        for arg in gen_def.context:
            assert isinstance(arg.return_type, NamedTensorValueType)
            arguments.append(current[arg.return_type.typename])
        for comb in itertools.product(*arguments):
            generator = executor.get_function_implementation(f'generator::{gen_name}').iter_from(*comb)
            for outputs in itertools.islice(generator, nr_samples):
                for output, output_def in zip(outputs, gen_def.generates):
                    assert isinstance(output_def.return_type, NamedTensorValueType)
                    current[output_def.return_type.typename].append(output)


concepts/pdsketch/csp_solvers/brute_force_sampling.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : dpll_sampling.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 11/20/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""The DPLL-Sampling algorithm for solving CSPs. This algorithm is specifically designed
for solving CSPs with mixed Boolean and continuous variables. At a high level, the algorithm
uses DPLL-style search to find a solution to Boolean variables. After the value for all Boolean
variables are fixed, the algorithm uses a sampling-based method to find a solution to the continuous variables.
"""

from typing import Optional, Union, Tuple, List, Dict

import itertools
import collections
import jacinle

from concepts.dsl.dsl_types import BOOL, ValueType, NamedTensorValueType, TensorValueTypeBase, PyObjValueType
from concepts.dsl.dsl_functions import Function
from concepts.dsl.tensor_value import TensorValue
from concepts.dsl.constraint import Constraint, ConstraintSatisfactionProblem, OptimisticValue, Assignment, AssignmentType, AssignmentDict, SimulationFluentConstraintFunction
from concepts.dsl.expression import VariableExpression, ValueOutputExpression, BoolOpType, QuantificationOpType, BoolExpression, PredicateEqualExpression, FunctionApplicationExpression
from concepts.dsl.expression import is_and_expr, iter_exprs

from concepts.pdsketch.executor import PDSketchExecutor, GeneratorManager
from concepts.pdsketch.predicate import Predicate
from concepts.pdsketch.generator import Generator, FancyGenerator

__all__ = ['CSPNotSolvable', 'CSPNoGenerator', 'csp_dpll_sampling_solve', 'csp_dpll_simplify']

ConstraintList = List[Optional[Constraint]]


class CSPNotSolvable(Exception):
    """An exception raised when the CSP is not solvable."""
    pass


class CSPNoGenerator(Exception):
    """An exception raised when there is no generator that can be matched in order to solve the CSP.
    Note that this does not mean that the CSP is not solvable."""
    pass


def _determined(*args) -> bool:
    """Helper function: if all arguments are determined.

    Args:
        *args: the arguments.

    Returns:
        True if all arguments are determined.
    """
    for x in args:
        if isinstance(x, OptimisticValue):
            return False
    return True


def _ground_assignment_value_partial(assignments: Dict[int, Assignment], dtype: ValueType, identifier: int) -> Union[TensorValue, OptimisticValue]:
    """Get the value of a variable based on the assignment dictionary. It will follow the EQUAL assignment types.

    The key difference between the :meth:`~concepts.dsl.constraint.ground_assignment_value` is the return type of this function.
    Specifically, the :meth:`concepts.dsl.constraint.ground_assignment_value` (exported) function returns the actual Value object.
    This function returns a wrapped Value object. When the value is not determined, it will return an OptimisticValue.

    Args:
        assignments: the assignment dictionary.
        dtype: the type of the variable.

    Returns:
        the value of the variable, wrapped in either :class:`~concepts.dsl.constraint.DeterminedValue` or :class:`~concepts.dsl.constraint.OptimisticValue`.
    """
    while identifier in assignments and assignments[identifier].t is AssignmentType.EQUAL:
        identifier = assignments[identifier].d
    if identifier in assignments:
        return assignments[identifier].d
    return OptimisticValue(dtype, identifier)


def _apply_assignments(executor: PDSketchExecutor, constraints: ConstraintList, assignments: Dict[int, Assignment]) -> ConstraintList:
    """Apply the assignments to the constraints. Essentially, it replaces all variables that have been determined in the assignment dictionary with the value.

    This function will also check all the constraints to make sure that the assignments are valid.
    When a constraint is invalid, this function will raises CSPNotSolvable. Otherwise, constraints that have been satisfied will be removed from the list.

    Args:
        executor: the executor.
        constraints: the list of constraints.
        assignments: the dictionary of assignments.

    Returns:
        the list of constraints that have not been satisfied.
    """
    new_constraints = list()
    for c in constraints:
        if c is None:
            continue

        # If the return value of the constraint is ignored, simply ignore the entire constraint.
        if isinstance(c.rv, OptimisticValue) and c.rv.identifier in assignments and assignments[c.rv.identifier].t is AssignmentType.IGNORE:
            continue

        # Ground the arguments and the return value.
        new_args = list(c.arguments)
        for i, x in enumerate(c.arguments):
            if isinstance(x, OptimisticValue) and x.identifier in assignments:
                new_args[i] = _ground_assignment_value_partial(assignments, x.dtype, x.identifier)
        new_rv = c.rv
        if isinstance(c.rv, OptimisticValue) and c.rv.identifier in assignments:
            new_rv = _ground_assignment_value_partial(assignments, c.rv.dtype, c.rv.identifier)

        # Evaluate the constraint.
        nc = Constraint(c.function, new_args, new_rv, note=c.note)
        if _determined(nc.rv) and _determined(*nc.arguments) and not isinstance(nc.function, SimulationFluentConstraintFunction):
            if _check_constraint(executor, nc):
                continue
            else:
                raise CSPNotSolvable(f'Constraint {c} is not satisfied.')
        new_constraints.append(nc)
    return new_constraints


def _check_constraint(executor: PDSketchExecutor, c: Constraint) -> bool:
    """Helper function: check if a constraint has been satisfied based on the current assignments.

    Args:
        executor: the executor.
        c: the constraint.

    Returns:
        True if the constraint has been satisfied.
    """
    if c.function is BoolOpType.NOT:
        return c.arguments[0].item() == (not c.rv.item())
    elif c.function in (QuantificationOpType.FORALL, BoolOpType.AND):
        return all([x.item() for x in c.arguments]) == c.rv.item()
    elif c.function in (QuantificationOpType.EXISTS, BoolOpType.OR):
        return any([x.item() for x in c.arguments]) == c.rv.item()
    elif c.is_equal_constraint:
        if c.arguments[0].dtype == BOOL:
            return (c.arguments[0].item() == c.arguments[1].item()) == c.rv.item()
        else:
            return _check_eq(executor, c.arguments[0].dtype, c.arguments[0], c.arguments[1]) == c.rv.item()
    elif isinstance(c.function, SimulationFluentConstraintFunction):
        return False
    else:
        assert isinstance(c.function, Predicate)
        # NB(Jiayuan Mao @ 09/05): for generator placeholders, they can only be set true through the corresponding generators.
        if c.function.is_generator_placeholder:
            return False
        func = executor.get_function_implementation(f'predicate::{c.function.name}')
        rv = func(*c.arguments, return_type=c.function.return_type)
        if rv.dtype == BOOL:
            return (rv.item() > 0.5) == c.rv.item()
        else:
            return _check_eq(executor, c.function.return_type, rv, c.rv.item())


def _check_eq(executor: PDSketchExecutor, dtype: Union[TensorValueTypeBase, PyObjValueType], v1: TensorValue, v2: TensorValue) -> bool:
    """Helper function: check if two values are equal. Internally used by :meth:`_check_constraint`.

    Args:
        executor: the executor.
        dtype: the type of the values.
        v1: the first value.
        v2: the second value.

    Returns:
        True if the two values are equal.
    """
    if isinstance(dtype, TensorValueTypeBase) and dtype.is_intrinsically_quantized():
        return (v1.tensor == v2.tensor).item()
    assert isinstance(dtype, NamedTensorValueType)
    eq_function = executor.get_function_implementation('type::' + dtype.typename + '::equal')
    return bool(eq_function(v1, v2).item())


def _filter_deterministic_equal(executor: PDSketchExecutor, constraints: ConstraintList, assignments: Dict[int, Assignment]) -> Tuple[bool, ConstraintList]:
    """Filter the constraints to remove the ones that are determined to be equal.

    Args:
        executor: the executor.
        constraints: the list of constraints.
        assignments: the dictionary of assignments.

    Returns:
        a tuple of (whether we have made progress, the list of constraints that have not been satisfied).
    """

    progress = False
    for i, c in enumerate(constraints):
        if c.is_equal_constraint and isinstance(c.rv, TensorValue):
            # If the constraint looks like: (x == y) == True, then we can set x = y.
            if c.rv.item():
                if isinstance(c.arguments[0], OptimisticValue):
                    if isinstance(c.arguments[1], OptimisticValue):
                        assignments[c.arguments[0].identifier] = Assignment(AssignmentType.EQUAL, c.arguments[1].identifier)
                    else:
                        assignments[c.arguments[0].identifier] = Assignment(AssignmentType.VALUE, c.arguments[1])
                    constraints[i] = None
                elif isinstance(c.arguments[1], OptimisticValue):
                    if isinstance(c.arguments[0], OptimisticValue):
                        assignments[c.arguments[1].identifier] = Assignment(AssignmentType.EQUAL, c.arguments[0].identifier)
                    else:
                        assignments[c.arguments[1].identifier] = Assignment(AssignmentType.VALUE, c.arguments[0])
                    constraints[i] = None
                else:
                    raise AssertionError('Sanity check failed.')
                progress = True
            else:
                if c.arguments[0].dtype == BOOL:
                    constraints[i] = Constraint(BoolOpType.NOT, [c.arguments[0]], c.arguments[1], note=c.note)
                    progress = True
    if progress:
        return progress, _apply_assignments(executor, constraints, assignments)
    return progress, constraints


def _filter_unused_rhs(executor: PDSketchExecutor, constraints: ConstraintList, assignments: Dict[int, Assignment], index2actionable: Dict[int, bool]) -> ConstraintList:
    """Filter out constraints that only appear once in the RHS of the constraints. In this case, the variable can be ignored and the related constraints can be removed.

    Args:
        executor: the executor.
        constraints: the list of constraints.
        assignments: the dictionary of assignments.

    Returns:
        the list of constraints that have not been satisfied, after removing all unused variables.
    """
    used: Dict[int, int] = collections.defaultdict(int)
    for i, actionable in index2actionable.items():
        if actionable:
            used[i] += 100

    for c in constraints:
        for x in c.arguments:
            if isinstance(x, OptimisticValue):
                used[x.identifier] += 100  # as long as a variable has appeared in the lhs of a constraint, it is used.
        if isinstance(c.rv, OptimisticValue):
            used[c.rv.identifier] += 1  # if the variable has only appeared in the rhs of a constraint for once, it is not used.
    for k, v in used.items():
        if v == 1:
            assignments[k] = Assignment(AssignmentType.IGNORE, None)
    return _apply_assignments(executor, constraints, assignments)


def _filter_deterministic_clauses(executor: PDSketchExecutor, constraints: ConstraintList, assignments: Dict[int, Assignment]) -> Tuple[bool, ConstraintList]:
    """Filter out Boolean constraints that have been determined. For example, AND(x, y, z) == true, then
        x == true, y == true, z == true. This function will remove the AND(x, y, z) constraint.

    There is another case that this function handles: for Boolean constraints, if everything on the LHS is determined, then
    we can determine the RHS.

    Args:
        executor: the executor.
        constraints: the list of constraints.
        assignments: the dictionary of assignments.

    Returns:
        a tuple of (whether we have made progress, the list of constraints that have not been satisfied).
    """
    progress = False
    for i, c in enumerate(constraints):
        if isinstance(c.function, (QuantificationOpType, BoolOpType)):
            if _determined(c.rv):
                if (
                    (c.function in (QuantificationOpType.FORALL, BoolOpType.AND)) or
                    (c.function in (QuantificationOpType.EXISTS, BoolOpType.OR) and len(c.arguments) <= 1)
                ):
                    if c.rv.item():
                        for x in c.arguments:
                            if isinstance(x, OptimisticValue):
                                assignments[x.identifier] = Assignment(AssignmentType.VALUE, True)
                                progress = True
                                # print('assign', optimistic_value_id(x), True)
                            elif not x:
                                raise CSPNotSolvable()
                    else:
                        # AND(x, y, z) == false
                        determined_values = [x for x in c.arguments if _determined(x)]
                        if False in determined_values:
                            progress = True
                            constraints[i] = None
                        elif len(determined_values) == len(c.arguments):
                            raise CSPNotSolvable()
                        elif len(determined_values) == len(c.arguments) - 1:
                            for x in c.arguments:
                                if not _determined(x):
                                    progress = True
                                    assignments[x.identifier] = Assignment(AssignmentType.VALUE, False)
                elif c.function in (QuantificationOpType.EXISTS, BoolOpType.OR):
                    if not c.rv.item():
                        for x in c.arguments:
                            if isinstance(x, OptimisticValue):
                                progress = True
                                assignments[x.identifier] = Assignment(AssignmentType.VALUE, False)
                                # print('assign', optimistic_value_id(x), False)
                            elif x:
                                raise CSPNotSolvable()
                    else:
                        # OR(x, y, z) == TRUE
                        determined_values = [x.item() for x in c.arguments if _determined(x)]
                        if True in determined_values:
                            progress = True
                            constraints[i] = None
                        elif len(determined_values) == len(c.arguments):
                            raise CSPNotSolvable()
                        elif len(determined_values) == len(c.arguments) - 1:
                            for x in c.arguments:
                                if not _determined(x):
                                    progress = True
                                    assignments[x.identifier] = Assignment(AssignmentType.VALUE, True)
                elif c.function is BoolOpType.NOT:
                    progress = True
                    assignments[c.arguments[0].identifier] = Assignment(AssignmentType.VALUE, not c.rv.item())
            elif _determined(*c.arguments):
                progress = True
                if c.function in (QuantificationOpType.FORALL, BoolOpType.AND):
                    assignments[c.rv.identifier] = Assignment(AssignmentType.VALUE, all(x.item() for x in c.arguments))
                elif c.function in (QuantificationOpType.EXISTS, BoolOpType.OR):
                    assignments[c.rv.identifier] = Assignment(AssignmentType.VALUE, any(x.item() for x in c.arguments))
                elif c.function is BoolOpType.NOT:
                    assignments[c.rv.identifier] = Assignment(AssignmentType.VALUE, not c.arguments[0].item())
        elif c.is_equal_constraint and _determined(*c.arguments):
            progress = True
            assignments[c.rv.identifier] = Assignment(AssignmentType.VALUE, c.arguments[0].item() == c.arguments[1].item())
    if progress:
        return progress, _apply_assignments(executor, constraints, assignments)
    return progress, constraints


def _find_bool_variable(executor: PDSketchExecutor, constraints: ConstraintList, assignments: Dict[int, Assignment]) -> Optional[int]:
    """Find a Boolean variable that is not determined. As a heuristic, we will look for the variable that appear in the maximum number of constraints.

    Args:
        constraints: the list of constraints.
        assignments: the dictionary of assignments.

    Returns:
        the variable that is not determined.
    """
    count: Dict[int, int] = collections.defaultdict(int)
    for c in constraints:
        for x in itertools.chain(c.arguments, [c.rv]):
            if isinstance(x, OptimisticValue) and x.identifier not in assignments and x.dtype == BOOL:
                count[x.identifier] += 1
    if len(count) == 0:
        return None
    return max(count, key=count.get)


def _find_grounded_function_application(executor: PDSketchExecutor, constraints: ConstraintList) -> Optional[Constraint]:
    """Find a function application whose arguments are all determined.

    Args:
        executor: the executor.
        constraints: the list of constraints.

    Returns:
        the function application that is not determined.
    """
    for c in constraints:
        if _determined(*c.arguments) and isinstance(c.function, Function):
            return c

    return None


GeneratorMatchingInputType = List[Optional[TensorValue]]
GeneratorMatchingOutputType = List[Optional[OptimisticValue]]
GeneratorMatchingIOReturnType = Tuple[Optional[GeneratorMatchingInputType], Optional[GeneratorMatchingOutputType]]


def _match_generator(c: Constraint, g: Generator, certifies_expr: Optional[ValueOutputExpression] = None, allow_star_matching: bool = False) -> GeneratorMatchingIOReturnType:
    def gen_input_output(func_arguments, rv_variable=None):
        inputs: GeneratorMatchingInputType = [None for _ in range(len(g.input_vars))]
        outputs: GeneratorMatchingOutputType = [None for _ in range(len(g.output_vars))]
        for argc, argg in zip(c.arguments, func_arguments):
            if isinstance(argc, OptimisticValue):
                if argg.name.startswith('?g'):
                    outputs[int(argg.name[2:])] = argc
                else:
                    return None, None
            else:
                if argg.name.startswith('?c'):
                    inputs[int(argg.name[2:])] = argc
                elif allow_star_matching and argg.name == '??':
                    continue
                else:
                    return None, None
        if rv_variable is not None:
            if rv_variable.name.startswith('?c'):
                inputs[int(rv_variable.name[2:])] = c.rv
            else:
                return None, None
        return inputs, outputs

    if certifies_expr is None:
        certifies_expr = g.flatten_certifies

    if isinstance(c.rv, TensorValue) and c.rv.dtype == BOOL:
        if c.rv.item():  # match (pred ?x ?y) == True
            if isinstance(certifies_expr, FunctionApplicationExpression):
                if c.function.name == certifies_expr.function.name:
                    return gen_input_output(certifies_expr.arguments)
        else:
            if isinstance(certifies_expr, BoolExpression) and certifies_expr.bool_op is BoolOpType.NOT:  # match (pred ?x ?y) == False
                inner_expr = certifies_expr.arguments[0]
                if isinstance(inner_expr, FunctionApplicationExpression):
                    if c.function.name == inner_expr.function.name:
                        return gen_input_output(inner_expr.arguments)
                elif c.is_equal_constraint and isinstance(inner_expr, PredicateEqualExpression):  # match (equal ?x ?y) == False
                    if c.arguments[0].dtype == inner_expr.predicate.return_type:
                        return gen_input_output([inner_expr.predicate, inner_expr.value])
    if isinstance(c.rv, TensorValue) and isinstance(c.function, Function):  # match (pred ?x ?y) == ?z
        if isinstance(certifies_expr, PredicateEqualExpression) and c.function.name == certifies_expr.predicate.function.name:
            if isinstance(certifies_expr.value, VariableExpression):
                return gen_input_output(certifies_expr.predicate.arguments, certifies_expr.value)
    return None, None


GeneratorMatchingReturnType = Optional[List[Tuple[Constraint, Generator, Optional[GeneratorMatchingInputType], Optional[GeneratorMatchingOutputType]]]]


def _find_gen_variable(executor: PDSketchExecutor, constraints: ConstraintList) -> GeneratorMatchingReturnType:
    # Step 1: find all applicable generators.
    all_generators = list()
    for c in constraints:
        for g in executor.domain.generators.values():
            i, o = _match_generator(c, g)
            # jacinle.log_function.print('matching', c, g, i, o)
            if i is not None:
                all_generators.append((c, g, i, o))

    # Step 2: find if there is any variable with only one generator.
    target_to_generator: Dict[int, list] = collections.defaultdict(list)
    for c, g, i, o in all_generators:
        for target in o:
            target_to_generator[target.identifier].append((c, g, i, o))

    for target, generators in target_to_generator.items():
        if len(target_to_generator[target]) == 1:
            return target_to_generator[target]

    if len(all_generators) > 0:
        max_priority = max([r[1].priority for r in all_generators])
        all_generators = [r for r in all_generators if r[1].priority == max_priority]
        return all_generators

    return None


def _find_gen_variable_advanced(executor: PDSketchExecutor, constraints: ConstraintList) -> GeneratorMatchingReturnType:
    def match_io(list1, list2):
        for x, y in zip(list1, list2):
            if x is None or y is None:
                continue
            if isinstance(x, OptimisticValue):
                if isinstance(y, OptimisticValue):
                    if x.identifier != y.identifier:
                        return False
                else:
                    return False
            elif isinstance(x, TensorValue):
                if isinstance(y, TensorValue):
                    if not _check_eq(executor, x.dtype, x, y):
                        return False
                else:
                    return False
            else:
                raise TypeError(f'Invalid type: {type(x)}.')
        return True

    def is_star_expression(sub_certifies):
        for x in iter_exprs(sub_certifies):
            if isinstance(x, VariableExpression) and x.name == '??':
                return True
        return False

    generator2matched = collections.defaultdict(list)
    for c in constraints:
        for g in executor.domain.generators.values():
            if is_and_expr(g.flatten_certifies):
                for sub_certifies_index, sub_certifies in enumerate(g.flatten_certifies.arguments):
                    i, o = _match_generator(c, g, sub_certifies, allow_star_matching=True)
                    if i is not None:
                        generator2matched[g].append((c, i, o, sub_certifies_index))

    all_generators = list()
    # TODO: implement the exact matching algorithm.
    for g, matched in generator2matched.items():
        all_matches = list()
        for result_index in range(len(matched)):
            c, i, o, sub_certifies_index = matched[result_index]

            used = False
            # constraints, g, i, o, matched_sub_certifies_index
            for mcs, mg, mi, mo, matched_sci in all_matches:
                if match_io(i, mi) and match_io(o, mo):
                    for j in range(len(mi)):
                        if mi[j] is None:
                            mi[j] = i[j]
                    for j in range(len(mo)):
                        if mo[j] is None:
                            mo[j] = o[j]

                    mcs.append(c)
                    matched_sci.add(sub_certifies_index)
                    used = True
                    break
            if not used:
                all_matches.append(([c], g, i, o, {sub_certifies_index}))

        this_is_star_expression = list()
        for sub_certifies in g.flatten_certifies.arguments:
            this_is_star_expression.append(is_star_expression(sub_certifies))

        for mcs, mg, mi, mo, matched_sci in all_matches:
            # the matched constraints should cover all sentences in the flatten_expression.
            match_succ = True
            for sub_certifies_index, sub_certifies in enumerate(g.flatten_certifies.arguments):
                if not this_is_star_expression[sub_certifies_index] and sub_certifies_index not in matched_sci:
                    match_succ = False
                    break
            if not match_succ:
                continue
            if None not in mi and None not in mo:
                all_generators.append((mcs, mg, mi, mo))

    return all_generators if len(all_generators) > 0 else None


def _find_typegen_variable(executor: PDSketchExecutor, dtype: ValueType) -> Optional[Generator]:
    assert isinstance(dtype, NamedTensorValueType)
    for g in executor.domain.generators.values():
        if len(g.function.arguments) == 0 and g.function.return_type[0] == dtype:
            if isinstance(g.certifies, BoolExpression) and g.certifies.bool_op is BoolOpType.AND:
                return g
    return None


def _find_fancy_gen_variable(executor: PDSketchExecutor, csp: ConstraintSatisfactionProblem, constraints: ConstraintList, assignments: AssignmentDict) -> Optional[List[Tuple[List[Constraint], Dict[int, Union[TensorValueTypeBase, PyObjValueType]], FancyGenerator]]]:
    results = list()
    for g in sorted(executor.domain.fancy_generators.values(), key=lambda generator: generator.priority, reverse=True):
        g: FancyGenerator
        this_constraints = list()
        this_variable_dtypes = dict()

        assert is_and_expr(g.flatten_certifies)
        for certifies_expr in g.flatten_certifies.arguments:
            assert isinstance(certifies_expr, FunctionApplicationExpression)
            for arg in certifies_expr.arguments:
                assert isinstance(arg, VariableExpression) and arg.name == '??'
            for c in constraints:
                if c.function.name == certifies_expr.function.name:
                    this_constraints.append(c)
                for arg in itertools.chain(c.arguments, [c.rv]):
                    if isinstance(arg, OptimisticValue):
                        this_variable_dtypes[arg.identifier] = arg.dtype
        if len(this_constraints) == 0:
            continue
        results.append((this_constraints, this_variable_dtypes, g))
    if len(results) > 0:
        return results
    return None


def _find_gen_variable_combined(executor: PDSketchExecutor, csp: ConstraintSatisfactionProblem, constraints: ConstraintList, assignments: AssignmentDict) -> GeneratorMatchingReturnType:
    """Combine the generator matching in the following order:

    1. Use :func:`_find_gen_variable` to find the generator with the highest priority.
    2. Use :func:`_find_gen_variable_advanced` to find the generator with the highest priority, using star-matching.
    3. Use :func:`_find_typegen_variable` to find the generator with the highest priority, using type-matching.
    """
    rv = _find_gen_variable(executor, constraints)
    if rv is not None:
        return rv
    rv = _find_gen_variable_advanced(executor, constraints)
    if rv is not None:
        return rv
    for name, dtype in csp.index2type.items():
        if name not in assignments and isinstance(dtype, NamedTensorValueType):
            g = _find_typegen_variable(executor, dtype)
            if g is not None:
                rv = [(None, g, [], [OptimisticValue(dtype, name)])]
                return rv
    return None


def csp_dpll_sampling_solve(
    executor: PDSketchExecutor, csp: ConstraintSatisfactionProblem, *,
    generator_manager: Optional[GeneratorManager] = None,
    max_generator_trials: int = 3,
    enable_ignore: bool = False, solvable_only: bool = False,
    verbose: bool = False
) -> Optional[Union[bool, AssignmentDict]]:
    """Solve the constraint satisfaction problem using the DPLL-sampling algorithm.

    Args:
        executor: the executor.
        csp: the constraint satisfaction problem.
        max_generator_trials: the maximum number of trials for each generator.
        enable_ignore: whether to ignore constraints whose RHS value is not determined.
        solvable_only: whether to only return whether the problem is solvable, without returning the solution.

    Returns:
        When `solvable_only` is True, return a single Boolean value indicating whether the problem is solvable.
        When `solvable_only` is False, return an assignment dictionary.
        When the problem is not solvable, return None.

    Raises:
        CSPNotSolvable: when the problem is not solvable.
        CSPNoGenerator: when no generator can be found to solve the problem. However, the problem may still be solvable.
    """
    if generator_manager is None:
        generator_manager = GeneratorManager(executor, store_history=False)

    constraints = csp.constraints.copy()

    @jacinle.log_function(verbose=False)
    def dfs(constraints, assignments):
        if len(constraints) == 0:
            return assignments

        progress = True
        while progress:
            progress, constraints = _filter_deterministic_equal(executor, constraints, assignments)
        if enable_ignore:
            constraints = _filter_unused_rhs(executor, constraints, assignments, csp.index2actionable)
        progress = True
        while progress:
            progress, constraints = _filter_deterministic_clauses(executor, constraints, assignments)

        if verbose:
            jacinle.log_function.print('Remaining constraints:', len(constraints))
            jacinle.log_function.print(*constraints, sep='\n')

        if len(constraints) == 0:
            return assignments

        if (next_bool_var := _find_bool_variable(executor, constraints, assignments)) is not None:
            assignments_true = assignments.copy()
            assignments_true[next_bool_var] = Assignment(AssignmentType.VALUE, True)
            try:
                constraints_true = _apply_assignments(executor, constraints, assignments_true)
                return dfs(constraints_true, assignments_true)
            except CSPNotSolvable:
                pass

            assignments_false = assignments.copy()
            assignments_false[next_bool_var] = Assignment(AssignmentType.VALUE, False)
            try:
                constraints_false = _apply_assignments(executor, constraints, assignments_false)
                return dfs(constraints_false, assignments_false)
            except CSPNotSolvable:
                pass

            raise CSPNotSolvable()
        elif (next_fapp := _find_grounded_function_application(executor, constraints)) is not None:
            function: Predicate = next_fapp.function
            arguments = next_fapp.arguments

            external_function = executor.get_function_implementation(f'predicate::{function.name}')
            output = external_function(*arguments, auto_broadcast=False)

            target = next_fapp.rv
            new_assignments = assignments.copy()
            new_assignments[target.identifier] = Assignment(AssignmentType.VALUE, output)
            try:
                new_constraints = constraints.copy()
                new_constraints[new_constraints.index(next_fapp)] = None
                new_constraints = _apply_assignments(executor, new_constraints, new_assignments)
                return dfs(new_constraints, new_assignments)
            except CSPNotSolvable:
                pass

            raise CSPNotSolvable()
        elif (next_gen_vars := _find_fancy_gen_variable(executor, csp, constraints, assignments)) is not None:
            if len(next_gen_vars) > 0 and next_gen_vars[0][1].unsolvable:
                raise CSPNotSolvable()

            for vv in next_gen_vars:
                c, dtype_mapping, g = vv

                generator = generator_manager.call(g, max_generator_trials, tuple(), c)
                generator = iter(generator)
                for j in range(max_generator_trials):
                    try:
                        output_index, outputs = next(generator)
                    except StopIteration:
                        break
                    if outputs is None:
                        break

                    assert isinstance(outputs, dict)
                    new_assignments = assignments.copy()
                    for k, v in outputs.items():
                        # v = TensorValue.from_scalar(v, dtype_mapping[k])
                        new_assignments[k] = Assignment(AssignmentType.VALUE, v)
                    try:
                        new_constraints = constraints.copy()
                        for cc in c:
                            new_constraints[new_constraints.index(cc)] = None

                        new_constraints = _apply_assignments(executor, new_constraints, new_assignments)
                        return dfs(new_constraints, new_assignments)
                    except CSPNotSolvable:
                        pass

            raise CSPNotSolvable()
        elif (next_gen_vars := _find_gen_variable_combined(executor, csp, constraints, assignments)) is not None:
            if len(next_gen_vars) > 1:
                # jacinle.log_function.print('generator orders', *[str(vv[1]).split('\n')[0] for vv in next_gen_vars], sep='\n  ')
                pass

            if len(next_gen_vars) > 0 and next_gen_vars[0][1].unsolvable:
                raise CSPNotSolvable()

            for vv in next_gen_vars:
                c, g, args, outputs_target = vv

                generator = generator_manager.call(g, max_generator_trials, tuple(args), c)
                generator = iter(generator)
                for j in range(max_generator_trials):
                    try:
                        output_index, outputs = next(generator)
                    except StopIteration:
                        break
                    if outputs is None:
                        break
                    if not isinstance(outputs, tuple) and g.function.ftype.is_singular_return:
                        outputs = (outputs, )

                    if not g.function.ftype.is_singular_return:
                        assert len(outputs) == len(g.function.return_type)

                    # jacinle.log_function.print('running generator', g, f'count = {j}')
                    new_assignments = assignments.copy()
                    for output, target in zip(outputs, outputs_target):
                        # output = TensorValue.from_scalar(output, target.dtype)
                        new_assignments[target.identifier] = Assignment(AssignmentType.VALUE, output)
                        # jacinle.log_function.print('assigned', target, output)
                    try:
                        new_constraints = constraints.copy()
                        if isinstance(c, list):
                            for cc in c:
                                new_constraints[new_constraints.index(cc)] = None
                        else:
                            new_constraints[new_constraints.index(c)] = None

                        new_constraints = _apply_assignments(executor, new_constraints, new_assignments)
                        return dfs(new_constraints, new_assignments)
                    except CSPNotSolvable:
                        pass

            raise CSPNotSolvable()
        else:
            # jacinle.log_function.print('Can not find a generator. Constraints:\n  ' + '\n  '.join([str(x) for x in constraints]))
            raise CSPNoGenerator('Can not find a generator. Constraints:\n  ' + '\n  '.join([str(x) for x in constraints]))

    try:
        assignments = dfs(constraints, {})
        if solvable_only:
            return True
        for name, dtype in csp.index2type.items():
            if name not in assignments:
                g = _find_typegen_variable(executor, dtype)
                if g is None:
                    raise NotImplementedError('Can not find a generator for unbounded variable {}, type {}.'.format(name, dtype))
                else:
                    output, = executor.get_function_implementation(g.function.name)()
                    assignments[name] = Assignment(AssignmentType.VALUE, TensorValue.from_scalar(output, dtype))
        return assignments
    except CSPNotSolvable:
        return None
    except CSPNoGenerator:
        raise


def csp_dpll_simplify(
    executor: PDSketchExecutor,
    csp: ConstraintSatisfactionProblem,
    enable_ignore: bool = True, return_assignments: bool = False
) -> Union[ConstraintSatisfactionProblem, Tuple[ConstraintSatisfactionProblem, AssignmentDict]]:
    """Simplify the CSP using DPLL algorithm.

    Args:
        executor: the executor.
        csp: the CSP.
        enable_ignore: whether to ignore constraints whose RHS value is not determined.
        return_assignments: whether to return the assignments.

    Returns:
        the simplified CSP.
    """

    constraints = csp.constraints.copy()
    assignments = dict()

    if len(constraints) == 0:
        return csp

    while True:
        nr_constraints = len(constraints)

        progress = True
        while progress:
            progress, constraints = _filter_deterministic_equal(executor, constraints, assignments)
        if enable_ignore:
            constraints = _filter_unused_rhs(executor, constraints, assignments, csp.index2actionable)
        progress = True
        while progress:
            progress, constraints = _filter_deterministic_clauses(executor, constraints, assignments)

        if len(constraints) == nr_constraints:
            break

    if return_assignments:
        return csp.clone(constraints), assignments
    return csp.clone(constraints)


concepts/pdsketch/csp_solvers/dpll_sampling.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : domain.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/30/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import itertools
from typing import TYPE_CHECKING, Any, Optional, Union, Sequence, Tuple, List, Dict

import jacinle
import torch
from jacinle.utils.printing import indent_text, stprint

from concepts.dsl.dsl_types import BOOL, FLOAT32, INT64, ObjectType, TensorValueTypeBase, PyObjValueType, ScalarValueType, VectorValueType, NamedTensorValueType, Variable, ObjectConstant
from concepts.dsl.dsl_functions import FunctionArgumentType, FunctionReturnType, FunctionType, Function
from concepts.dsl.dsl_domain import DSLDomainBase
from concepts.dsl.expression import Expression, ExpressionDefinitionContext, VariableExpression, ValueOutputExpression, cvt_expression_list
from concepts.dsl.expression import FunctionApplicationExpression, AssignExpression, ConditionalAssignExpression, DeicticAssignExpression
from concepts.dsl.constraint import is_optimistic_value
from concepts.dsl.tensor_state import NamedObjectTensorState

from concepts.pdsketch.predicate import Predicate, flatten_expression, get_used_state_variables
from concepts.pdsketch.operator import Precondition, Effect, Controller, Operator, MacroOperator, OperatorApplier, OperatorApplicationExpression, AchieveExpression, RegressionRule
from concepts.pdsketch.generator import Generator, FancyGenerator

if TYPE_CHECKING:
    from concepts.pdsketch.executor import PDSketchExecutor

logger = jacinle.get_logger(__file__)

__all__ = ['Domain', 'Problem', 'State']


class _TypedVariableView(object):
    """Use `domain.typed_variable['type_name']('variable_name')`"""

    def __init__(self, domain):
        self.domain = domain

    def __getitem__(self, typename):
        def function(string):
            return Variable(string, self.domain.types[typename])
        return function


class Domain(DSLDomainBase):
    """The planning domain definition."""

    def __init__(self, name: Optional[str] = None):
        """Initialize a planning domain.

        Args:
            name: The name of the domain.
        """
        super().__init__(name)

        self.operators = dict()
        self.operator_templates = dict()
        self.regression_rules = dict()
        self.axioms = dict()
        self.generators = dict()
        self.fancy_generators = dict()

        self.external_functions = dict()
        self.external_function_crossrefs = dict()

        self.tv = self.typed_variable = _TypedVariableView(self)

    name: str

    types: Dict[str, Union[ObjectType, PyObjValueType, TensorValueTypeBase]]
    """The types defined in the domain, as a dictionary from type names to types."""

    functions: Dict[str, Predicate]
    """A mapping from function name to the corresponding :class:`~concepts.pdsketch.predicate.Predicate` class. Note that,
    unlike the basic :class:`~concepts.dsl.dsl_domain.DSLDomainBase`, in planning domain, all functions should be of type :class:`~concepts.pdsketch.predicate.Predicate`."""

    constants: Dict[str, ObjectConstant]
    """The constants defined in the domain, as a dictionary from constant names to values."""

    operators: Dict[str, Union[Operator, MacroOperator]]
    """A mapping of operators: from operator name to the corresponding :class:`~concepts.pdsketch.operator.Operator` class."""

    operator_templates: Dict[str, Operator]
    """A mapping of operator templates: from operator name to the corresponding :class:`~concepts.pdsketch.operator.Operator` class."""

    regression_rules: Dict[str, RegressionRule]
    """A mapping of regression rules: from regression rule name to the corresponding :class:`~concepts.pdsketch.operator.RegressionRule` class."""

    axioms: Dict[str, Operator]
    """A mapping of axioms: from axiom name to the corresponding :class:`~concepts.pdsketch.operator.Operator` class."""

    generators: Dict[str, Generator]
    """A mapping of generators: from generator name to the corresponding :class:`~concepts.pdsketch.generator.Generator` class."""

    fancy_generators: Dict[str, FancyGenerator]
    """A mapping of fancy generators: from fancy generator name to the corresponding :class:`~concepts.pdsketch.generator.FancyGenerator` class."""

    external_functions: Dict[str, Function]
    """A mapping of external functions: from function name to the corresponding :class:`~concepts.dsl.dsl_functions.Function` class."""

    external_function_crossrefs: Dict[str, str]
    """A mapping from function name to another function name. This is useful when defining one function as an derived function of another function."""

    tv: _TypedVariableView
    """A helper function that returns a variable with the given type.
    For example, `domain.tv['object']('x')` returns a variable of type `object` with name `x`."""

    def __getattr__(self, item):
        if item.startswith('__') and item.endswith('__'):
            raise AttributeError

        # NB(Jiayuan Mao @ 09/03): PDDL definition convention.
        item = item.replace('_', '-')

        if item.startswith('t-'):
            return self.types[item[2:]]
        elif item.startswith('p-') or item.startswith('f-'):
            return self.functions[item[2:]]
        elif item.startswith('op-'):
            return self.operators[item[3:]]
        elif item.startswith('gen-'):
            return self.generators[item[4:]]
        raise NameError('Unknown attribute: {}.'.format(item))

    def set_name(self, name: str):
        """Set the name of the domain.

        Args:
            name: the new name of the domain.
        """
        self.name = name

    BUILTIN_TYPES = ['object', 'pyobject', 'bool', 'int64', 'float32']
    BUILTIN_NUMERIC_TYPES = {
        'bool': BOOL,
        'int64': INT64,
        'float32': FLOAT32
    }

    def define_type(self, typename, parent_name: Optional[Union[VectorValueType, ScalarValueType, str]] = 'object') -> Union[ObjectType, PyObjValueType, VectorValueType, ScalarValueType]:
        """Define a new type.

        Args:
            typename: the name of the new type.
            parent_name: the parent type of the new type, default to 'object'.

        Returns:
            the newly defined type.
        """

        if typename == 'object':
            logger.warning_once('Shadowing built-in type name "object".')
        elif typename in type(self).BUILTIN_TYPES:
            raise ValueError('Typename {} is a built-in type.'.format(typename))

        assert isinstance(parent_name, (str, VectorValueType)), f'Currently only support inheritance from builtin types: {type(self).BUILTIN_TYPES}.'

        if isinstance(parent_name, str):
            if parent_name == 'object':
                self.types[typename] = ObjectType(typename)
            elif parent_name == 'pyobject':
                dtype = PyObjValueType(typename)
                self.types[typename] = dtype
                self.declare_external_function(f'type::{typename}::equal', [dtype, dtype], BOOL)
            elif parent_name == 'int64':
                dtype = NamedTensorValueType(typename, INT64)
                self.types[typename] = dtype
                self.declare_external_function(f'type::{typename}::equal', [dtype, dtype], BOOL)
            elif parent_name == 'float32':
                dtype = NamedTensorValueType(typename, FLOAT32)
                self.types[typename] = dtype
                self.declare_external_function(f'type::{typename}::equal', [dtype, dtype], BOOL)
            else:
                raise ValueError(f'Unknown parent type: {parent_name}.')
        elif isinstance(parent_name, VectorValueType):
            dtype = NamedTensorValueType(typename, parent_name)
            self.types[typename] = dtype
            self.declare_external_function(f'type::{typename}::equal', [dtype, dtype], BOOL)
        else:
            raise ValueError(f'Unknown parent type: {parent_name}.')

        return self.types[typename]

    def get_type(self, typename: str) -> Union[ObjectType, PyObjValueType, VectorValueType, ScalarValueType, NamedTensorValueType]:
        """Get a type by name.

        Args:
            typename: the name of the type.

        Returns:
            the type with the given name.
        """
        if typename in type(self).BUILTIN_NUMERIC_TYPES:
            return type(self).BUILTIN_NUMERIC_TYPES[typename]
        if typename not in self.types:
            raise ValueError(f'Unknown type: {typename}.')
        return self.types[typename]

    def define_predicate(
        self, name: str, arguments: Sequence[FunctionArgumentType], return_type: FunctionReturnType = BOOL,
        observation: Optional[bool] = None, state: Optional[bool] = None,
        generator_placeholder: bool = False, inplace_generators: Optional[Sequence[str]] = None,
    ):
        """Define a new predicate.

        Args:
            name: the name of the new predicate.
            arguments: the arguments of the new predicate.
            return_type: the return type of the new predicate.
            observation: whether the new predicate is an observation variable.
            state: whether the new predicate is a state variable.
            generator_placeholder: whether the new predicate is a generator placeholder.

        Returns:
            the newly defined predicate.
        """
        predicate = Predicate(name, FunctionType(arguments, return_type), observation=observation, state=state, generator_placeholder=generator_placeholder, inplace_generators=inplace_generators)
        self.define_predicate_inner(name, predicate)
        return predicate

    def define_derived(self, name: str, arguments: Sequence[FunctionArgumentType], return_type: Optional[FunctionReturnType] = None, expr: ValueOutputExpression = None, state: bool = False, generator_placeholder: bool = False):
        """Define a new derived predicate. Note that a derived predicate can not be an observation variable.

        Args:
            name: the name of the new derived predicate.
            arguments: the arguments of the new derived predicate.
            return_type: the return type of the new derived predicate.
            expr: the expression of the new derived predicate.
            state: whether the new derived predicate is a state variable.
            generator_placeholder: whether the new derived predicate is a generator placeholder.

        Returns:
            the newly defined derived predicate.
        """
        predicate_def = Predicate(name, FunctionType(arguments, return_type), observation=False, state=state, generator_placeholder=generator_placeholder, derived_expression=expr)
        return self.define_predicate_inner(name, predicate_def)

    def define_predicate_inner(self, name, predicate_def: Predicate):
        self.functions[name] = predicate_def

        # NB(Jiayuan Mao @ 07/21): a non-cacheable function is basically an external function.
        if not predicate_def.is_cacheable and predicate_def.derived_expression is None:
            identifier = f'predicate::{name}'
            self.external_functions[identifier] = predicate_def

        return predicate_def

    def define_operator(self, name: str, parameters: Sequence[Variable], preconditions: Sequence[Precondition], effects: Sequence[Effect], controller: Controller, template=False, extends=None) -> Operator:
        """Define a new operator.

        Args:
            name: the name of the new operator.
            parameters: the parameters of the new operator.
            preconditions: the preconditions of the new operator.
            effects: the effects of the new operator.
            controller: the controller of the new operator.
            template: whether the new operator is a template.
            extends: the parent operator of the new operator.

        Returns:
            the newly defined operator.
        """
        self.operators[name] = op = Operator(name, parameters, preconditions, effects, controller, extends=extends, is_template=template)
        return op

    def define_operator_inner(self, name: str, operator: Operator) -> Operator:
        assert name not in self.operators
        self.operators[name] = operator
        return operator

    def define_regression_rule(
        self, name: str, parameters: Sequence[Variable],
        preconditions: Sequence[Precondition], preconstraints: Sequence[Precondition],
        goal: Sequence[Effect], goal_expression: ValueOutputExpression, body: Sequence[Union[OperatorApplicationExpression, AchieveExpression]],
        always: bool = False
    ):
        """Define a new regression rule.

        Args:
            name: the name of the new regression rule.
            parameters: the parameters of the new regression rule.
            preconditions: the preconditions of the new regression rule.
            preconstraints: the preconstraints of the new regression rule.
            goal: the goal of the new regression rule.
            body: the body of the new regression rule.
            always: whether the new regression rule is always applicable.

        Returns:
            the newly defined regression rule.
        """
        self.regression_rules[name] = rule = RegressionRule(name, parameters, preconditions, preconstraints, goal, goal_expression, body, always=always)
        return rule

    def define_axiom(self, name: Optional[str], parameters: Sequence[Variable], preconditions: Sequence[Precondition], effects: Sequence[Effect]) -> Operator:
        """Define a new axiom.

        Args:
            name: the name of the new axiom. If None, a unique name will be generated.
            parameters: the parameters of the new axiom.
            preconditions: the preconditions of the new axiom.
            effects: the effects of the new axiom.

        Returns:
            the newly defined axiom.
        """

        if name is None:
            name = f'axiom_{len(self.axioms)}'
        self.axioms[name] = op = Operator(name, parameters, preconditions, effects, is_axiom=True)
        return op

    def define_macro(self, name: str, parameters: Sequence[Variable], sub_operators: Sequence[OperatorApplier]):
        """Define a new macro.

        Args:
            name: the name of the new macro.
            parameters: the parameters of the new macro.
            sub_operators: the sub operators of the new macro.

        Returns:
            the newly defined macro.
        """
        self.operators[name] = op = MacroOperator(name, parameters, sub_operators)
        return op

    def define_generator(
        self, name: str, parameters: Sequence[Variable], certifies: ValueOutputExpression,
        context: Sequence[Union[VariableExpression, ValueOutputExpression]], generates: Sequence[Union[VariableExpression, ValueOutputExpression]],
        priority: int = 0, unsolvable: bool = False
    ) -> Generator:
        """Define a new generator.

        Args:
            name: the name of the new generator.
            parameters: the parameters of the new generator.
            certifies: the certified condition of the new generator.
            context: the context of the new generator.
            generates: the generates of the new generator.
            priority: the priority of the new generator.
            unsolvable: whether the new generator is unsolvable.

        Returns:
            the newly defined generator.
        """

        if unsolvable:
            priority = int(1e9)

        identifier = f'generator::{name}'
        context: List[Union[VariableExpression, ValueOutputExpression]] = cvt_expression_list(context)
        generates: List[Union[VariableExpression, ValueOutputExpression]] = cvt_expression_list(generates)

        arguments = [Variable(f'?c{i}', c.return_type) for i, c in enumerate(context)]
        return_type = [target.return_type for target in generates]
        output_vars = [Variable(f'?g{i}', g.return_type) for i, g in enumerate(generates)]
        function = Function(identifier, FunctionType(arguments, return_type))

        all_variables = {c: cv for c, cv in zip(context, arguments)}
        all_variables.update({g: gv for g, gv in zip(generates, output_vars)})
        ctx = ExpressionDefinitionContext(*arguments, *output_vars, domain=self)
        flatten_certifies = flatten_expression(certifies, all_variables, ctx, flatten_cacheable_expression=True)

        if not unsolvable:
            self.external_functions[identifier] = function

        if name in self.generators:
            raise ValueError(f'Duplicate generator: {name}.')
        self.generators[name] = generator = Generator(
            name, parameters, certifies,
            context=context, generates=generates,
            function=function, output_vars=output_vars, flatten_certifies=flatten_certifies,
            priority=priority, unsolvable=unsolvable
        )
        return generator

    def define_fancy_generator(
        self, name: str, certifies: ValueOutputExpression,
        priority: int = 10, unsolvable: bool = False
    ) -> Generator:
        """Declare a new fancy generator. The difference between a fancy generator and a normal generator is that
        a fancy generator is not directional. That is, it can generate a set of variables satisfies the constraints,
        without requring specific `contexts` to `genereates` directions. Therefore, we don't need to specify the
        `context` and `generates` of a fancy generator.

        Args:
            name: the name of the new fancy generator.
            certifies: the certified condition of the new fancy generator.
            priority: the priority of the new fancy generator.
            unsolvable: whether the new fancy generator is unsolvable.

        Returns:
            the newly declared fancy generator.
        """
        if unsolvable:
            priority = int(1e9)

        identifier = f'generator::{name}'
        # TODO(Jiayuan Mao @ 2023/04/04): fix the typing for this.
        function = Function(identifier, FunctionType([], []))
        flatten_certifies = certifies

        if not unsolvable:
            self.external_functions[identifier] = function

        if name in self.generators:
            raise ValueError(f'Duplicate generator: {name}.')
        self.fancy_generators[name] = generator = FancyGenerator(name, certifies, function=function, flatten_certifies=flatten_certifies, priority=priority, unsolvable=unsolvable)
        return generator

    def declare_external_function(self, function_name: str, argument_types: Sequence[FunctionArgumentType], return_type: FunctionArgumentType, kwargs: Optional[Dict[str, Any]] = None) -> Function:
        """Declare an external function.

        Args:
            function_name: the name of the external function.
            argument_types: the argument types of the external function.
            return_type: the return type of the external function.
            kwargs: the keyword arguments of the external function. Supported keyword arguments are:
                - ``observation``: whether the external function is an observation variable.
                - ``state``: whether the external function is a state variable.
        """
        if kwargs is None:
            kwargs = dict()

        self.external_functions[function_name] = Predicate(function_name, FunctionType(argument_types, return_type), **kwargs)
        return self.external_functions[function_name]

    def declare_external_function_crossref(self, function_name: str, cross_ref_name: str):
        """Declare a cross reference to an external function.
        This is useful when one function is an derived function of another function.

        Args:
            function_name: the name of the external function.
            cross_ref_name: the name of the cross reference.
        """
        self.external_function_crossrefs[function_name] = cross_ref_name

    def parse(self, string: Union[str, Expression], variables: Optional[Sequence[Variable]] = None) -> Expression:
        """Parse a string into an expression.

        Args:
            string: the string to be parsed.
            variables: the variables to be used in the expression.

        Returns:
            the parsed expression.
        """
        if isinstance(string, Expression):
            return string

        from concepts.pdsketch.parsers.pdsketch_parser import parse_expression
        return parse_expression(self, string, variables)

    def incremental_define(self, string: str):
        """Incrementally define new parts of the domain.

        Args:
            string: the string to be parsed and defined.
        """
        from concepts.pdsketch.parsers.pdsketch_parser import load_domain_string_incremental
        load_domain_string_incremental(self, string)

    def print_summary(self):
        """Print a summary of the domain."""
        print(f'Domain {self.name}')
        stprint(key='Types: ', data=self.types, indent=1)
        stprint(key='Functions: ', data=self.functions, indent=1)
        stprint(key='External Functions: ', data=self.external_functions, indent=1)
        stprint(key='Generators: ', data=self.generators, indent=1)
        stprint(key='Fancy Generators: ', data=self.fancy_generators, indent=1)
        print('  Operators:')
        if len(self.operators) > 0:
            for op in self.operators.values():
                print(indent_text(op.pddl_str(), level=2))
        else:
            print('    <Empty>')
        print('  Axioms:')
        if len(self.axioms) > 0:
            for op in self.axioms.values():
                print(indent_text(op.pddl_str(), level=2))
        else:
            print('    <Empty>')
        print('  Regression Rules:')
        if len(self.regression_rules) > 0:
            for op in self.regression_rules.values():
                print(indent_text(op.pddl_str(), level=2))
        else:
            print('    <Empty>')

    def post_init(self):
        """Post-initialization of the domain.
        This function should be called by the domain generator after all the domain definitions (predicates and operators) are done.
        Currently, the following post-initialization steps are performed:

        1. Analyze the static predicates.
        """
        self._analyze_static_predicates()

    def _analyze_static_predicates(self):
        """Run static analysis on the predicates to determine which predicates are static."""
        dynamic = set()
        for op in itertools.chain(self.operators.values(), self.axioms.values()):
            if isinstance(op, MacroOperator):
                continue
            for eff in op.effects:
                if isinstance(eff.assign_expr, (AssignExpression, ConditionalAssignExpression)):
                    dynamic.add(eff.assign_expr.predicate.function.name)
                elif isinstance(eff.assign_expr, DeicticAssignExpression):
                    expr = eff.unwrapped_assign_expr
                    assert isinstance(expr, (AssignExpression, ConditionalAssignExpression))
                    dynamic.add(expr.predicate.function.name)
                else:
                    raise TypeError(f'Unknown effect type: {eff.assign_expr}.')

        # propagate the static predicates.
        for p in self.functions.values():
            if p.is_state_variable:
                p.mark_static(p.name not in dynamic)
            else:
                if p.is_cacheable and p.derived_expression is not None:
                    used_predicates = get_used_state_variables(p.derived_expression)
                    static = True
                    for predicate_def in used_predicates:
                        if not predicate_def.is_static:
                            static = False
                            break
                    p.mark_static(static)


class Problem(object):
    """The representation for a planning problem. It contains the set of objects, a inital state (a set of propositions), and a goal expression."""

    def __init__(self):
        """Initialize the problem."""
        self.objects = dict()
        self.predicates = list()
        self.goal = None

    objects: Dict[str, str]
    """The set of objects, which are mappings from object names to object type names."""

    predicates: List[FunctionApplicationExpression]
    """The initial state, which is a set of propositions."""

    goal: Optional[ValueOutputExpression]
    """The goal expression."""

    def add_object(self, name: str, typename: str):
        """Add an object to the problem.

        Args:
            name: the name of the object.
            typename: the type of the object.
        """
        self.objects[name] = typename

    def add_proposition(self, proposition: FunctionApplicationExpression):
        """Add a proposition to the initial problem.

        Args:
            proposition: the proposition to add.
        """
        self.predicates.append(proposition)

    def set_goal(self, goal: ValueOutputExpression):
        """Set the goal of the problem.

        Args:
            goal: the goal expression.
        """
        self.goal = goal

    def to_state(self, executor: 'PDSketchExecutor') -> 'State':
        """Convert the problem to a :class:`State` object.

        Args:
            executor: the executor to use to instantiate the state.

        Returns:
            the state object.
        """

        domain = executor.domain
        object_names = list(self.objects.keys())
        object_types = [executor.domain.types[self.objects[name]] for name in object_names]

        for constant in domain.constants.values():
            object_names.append(constant.name)
            object_types.append(constant.dtype)

        state = State(None, object_names, object_types)

        from concepts.pdsketch.executor import StateDefinitionHelper
        ctx = StateDefinitionHelper(executor, state)
        predicates = list()
        for p in self.predicates:
            predicates.append(ctx.get_predicate(p.function.name)(*[arg.constant.name for arg in p.arguments]))
        ctx.define_predicates(predicates)

        return state


class State(NamedObjectTensorState):
    """Planning domain state."""

    def make_quantized(self, domain: Domain, features=None) -> 'State':
        """Make a quantized version of the state.

        Args:
            domain: the planning domain.
            features: the features to use for quantization. If None, use all state variables.

        Returns:
            the quantized state.
        """
        if features is None:
            features = [name for name in self.features.all_feature_names if domain.functions[name].is_state_variable]

        new_tensor_dict = dict()
        for feature_name in features:
            new_tensor_dict[feature_name] = self.features[feature_name].simple_quantize()
        return type(self)(self.object_types, new_tensor_dict, self.object_names)

    def generate_tuple_description(self, domain: Domain) -> Tuple[int, ...]:
        """Generate a tuple description of the state.

        Args:
            domain: the planning domain.

        Returns:
            the tuple description of the state.
        """
        rv = list()
        for feature_name in sorted(self.features.all_feature_names):
            if domain.functions[feature_name].is_state_variable:
                feature = self.features[feature_name]
                if isinstance(feature.dtype, TensorValueTypeBase) and feature.dtype.is_intrinsically_quantized():
                    rv.extend(_maybe_apply_optimistic_mask(feature.tensor, feature.tensor_optimistic_values).flatten().tolist())
                elif feature.tensor_quantized_values is not None:
                    rv.extend(_maybe_apply_optimistic_mask(feature.tensor_quantized_values, feature.tensor_optimistic_values).flatten().tolist())
                else:
                    raise RuntimeError(f'Cannot generate tuple description for feature {feature_name}.')
        return tuple(rv)


def _maybe_apply_optimistic_mask(tensor, optimistic_values):
    if optimistic_values is None:
        return tensor
    assert tensor.shape == optimistic_values.shape
    optimistic_mask = is_optimistic_value(optimistic_values)
    return torch.where(optimistic_mask, optimistic_values, tensor.to(torch.int64))


concepts/pdsketch/domain.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : executor.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 11/18/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""This file contains the executor classes for the PDSketch framework.

All the executors are based on :class:`~concepts.dsl.tensor_value.TensorValue` classes.
It supports all expressions defined in :mod:`~concepts.dsl.expressions`, including the basic
function application expressions and a few first-order logic expression types. The executors
are designed to be "differentiable", which means we can directly do backpropagation on the
computed values.

The main entry for the executor is the :class:`PDSketchExecutor` class.
Internally it contains two executor implementations: the basic one, and an "optimistic" one,
which handles the case where the value of a variable can be unknown and "optimistic".
"""

import itertools
import functools
from typing import Any, Optional, Union, Iterator, Sequence, Tuple, List, Mapping, Dict, Callable
from collections import defaultdict
from tabulate import tabulate

import torch
import jactorch
from jacinle.logging import get_logger

import concepts.dsl.expression as E
from concepts.dsl.dsl_types import BOOL, TensorValueTypeBase, ScalarValueType, VectorValueType, NamedTensorValueType, PyObjValueType, QINDEX, Variable
from concepts.dsl.expression import Expression, BoolOpType, QuantificationOpType
from concepts.dsl.expression_visitor import ExpressionVisitor
from concepts.dsl.tensor_value import TensorizedPyObjValues, TensorValue, MaskedTensorStorage
from concepts.dsl.tensor_state import StateObjectReference, TensorState, NamedObjectTensorState, ObjectNameArgument, ObjectTypeArgument
from concepts.dsl.constraint import OPTIM_MAGIC_NUMBER_MAGIC, is_optimistic_value, OptimisticValue, cvt_opt_value, Constraint, EqualityConstraint, ConstraintSatisfactionProblem, AssignmentDict, SimulationFluentConstraintFunction
from concepts.dsl.parsers.parser_base import ParserBase
from concepts.dsl.executors.tensor_value_executor import TensorValueExecutorBase, TensorValueExecutorReturnType, compose_bvdict_args
from concepts.pdsketch.predicate import Predicate
from concepts.pdsketch.operator import Operator, MacroOperator, OperatorApplier
from concepts.pdsketch.generator import Generator
from concepts.pdsketch.domain import Domain, State

logger = get_logger(__file__)

__all__ = [
    'PythonFunctionRef', 'PythonFunctionCrossRef', 'config_function_implementation',
    'PDSketchExecutor', 'PDSketchExecutionDefaultVisitor', 'PDSketchExecutionCSPVisitor',
    'StateDefinitionHelper',
    'GeneratorManager', 'wrap_singletime_function_to_iterator'
]


class PythonFunctionRef(object):
    """A reference to a Python function.

    This class is used to wrap external function implementations in domains.
    """

    def __init__(
        self, function: Callable, function_quantized: Optional[Callable] = None,
        *,
        return_type: Optional[Union[TensorValueTypeBase, PyObjValueType, Tuple[Union[TensorValueTypeBase, PyObjValueType], ...]]] = None,
        support_batch: bool = False, auto_broadcast: bool = False,
        use_object_names: bool = True, unwrap_values: Optional[bool] = None,
        include_executor_args: bool = False, is_iterator: bool = False,
        executor: Optional['PDSketchExecutor'] = None
    ):
        """Initialize a Python function reference.

        Args:
            function: the function to be wrapped.
            function_quantized: the quantized version of the function (can be None).
            support_batch: whether the function supports batched inputs.
            auto_broadcast: whether the executor should automatically broadcast the arguments before calling the function.
            use_object_names: whether the executor should use the object names in the state (instead of the index).
            unwrap_values: whether the executor should unwrap the tensor values before calling the function.
            include_executor_args: whether the caller should include the executor as the first argument.
            is_iterator: whether the function is an iterator.
            executor: the executor that is using this function reference.
        """

        self.function = function
        self.function_quantized = function_quantized
        self.return_type = return_type
        self.support_batch = support_batch
        self.auto_broadcast = auto_broadcast
        self.use_object_names = use_object_names
        if unwrap_values is None:
            unwrap_values = not support_batch or auto_broadcast
        self.unwrap_values = unwrap_values
        self.include_executor_args = include_executor_args
        self.is_iterator = is_iterator
        self._executor = executor

    function: Callable
    """The internal implementation of the function."""

    function_quantized: Optional[Callable]
    """The quantized version of the function (can be None)."""

    return_type: Optional[Union[TensorValueTypeBase, PyObjValueType, Tuple[Union[TensorValueTypeBase, PyObjValueType], ...]]]
    """The return type of the function."""

    support_batch: bool
    """Whether the function supports batched inputs."""

    auto_broadcast: bool
    """Whether the executor should automatically broadcast the arguments before calling the function."""

    use_object_names: bool
    """Whether the executor should use the object names in the state (instead of the index)."""

    unwrap_values: bool
    """Whether the executor should unwrap the tensor values before calling the function."""

    include_executor_args: bool
    """Whether the caller should include the executor as the first argument."""

    is_iterator = False
    """Whether the function is an iterator."""

    def set_executor(self, executor: 'PDSketchExecutor') -> 'PythonFunctionRef':
        """Set the executor that is using this function reference.

        Args:
            executor: the executor that is using this function reference.

        Returns:
            the function reference itself.
        """
        self._executor = executor
        return self

    def __str__(self) -> str:
        return (
            'PythonFunctionRef('
            f'{self.function}, support_batch={self.support_batch}, auto_broadcast={self.auto_broadcast}, '
            f'use_object_names={self.use_object_names}, unwrap_values={self.unwrap_values}, include_executor_args={self.include_executor_args}, '
            f'is_iterator={self.is_iterator})'
        )

    def __repr__(self) -> str:
        return self.__str__()

    @property
    def flags(self) -> Dict[str, bool]:
        return {
            'support_batch': self.support_batch,
            'auto_broadcast': self.auto_broadcast,
            'use_object_names': self.use_object_names,
            'unwrap_values': self.unwrap_values,
            'include_executor_args': self.include_executor_args,
            'is_iterator': self.is_iterator,
        }

    def forward(
        self, argument_values: Sequence[TensorValueExecutorReturnType], return_type: Optional[Union[TensorValueTypeBase, PyObjValueType]] = None,
        auto_broadcast: bool = True, wrap_rv: bool = True
    ) -> Union[TensorValue, Tuple[TensorValue, ...]]:
        """Call the function.

        Args:
            argument_values: the arguments to the function.
            return_type: the type of the return value.
            auto_broadcast: whether the executor should automatically broadcast the arguments before calling the function.
            wrap_rv: whether the executor should wrap the return value.

        Returns:
            the result of the function.
        """
        function = self.function

        if self.use_object_names:
            argument_values = [v.name if isinstance(v, StateObjectReference) else v for v in argument_values]
        else:
            argument_values = [v.index if isinstance(v, StateObjectReference) else v for v in argument_values]

        if self.support_batch:
            if self.auto_broadcast and auto_broadcast:
                argument_values = expand_argument_values(argument_values)

        argument_values_flat = argument_values
        if self.unwrap_values:
            argument_values_flat = [v.tensor if isinstance(v, TensorValue) else v for v in argument_values_flat]
            if not self.support_batch:
                argument_values_flat = [v.item() if isinstance(v, TensorizedPyObjValues) else v for v in argument_values_flat]

        if self.include_executor_args:
            assert self._executor is not None, 'Executor is None.'
            rv = function(self._executor, *argument_values_flat)
        else:
            rv = function(*argument_values_flat)

        if not wrap_rv:
            return rv
        return self._wrap_rv(rv, return_type, argument_values, auto_broadcast)

    def forward_generator(
        self, argument_values: Sequence[TensorValueExecutorReturnType], return_type: Optional[Union[TensorValueTypeBase, PyObjValueType]] = None,
        auto_broadcast: bool = True, wrap_rv: bool = True
    ) -> Union[Iterator[TensorValue], Iterator[Tuple[TensorValue, ...]]]:
        """Call the function and return a generator.

        Args:
            argument_values: the arguments to the function.
            return_type: the type of the return value.
            auto_broadcast: whether the executor should automatically broadcast the arguments before calling the function.
            wrap_rv: whether the executor should wrap the
        """

        generator = self.forward(argument_values, return_type=return_type, auto_broadcast=auto_broadcast, wrap_rv=False)
        if not wrap_rv:
            yield from generator
        else:
            for v in generator:
                yield self._wrap_rv(v, return_type, argument_values, auto_broadcast)

    def __call__(self, *args, return_type: Optional[Union[TensorValueTypeBase, PyObjValueType]] = None, auto_broadcast: bool = True, wrap_rv: bool = True) -> Union[TensorValue, Tuple[TensorValue, ...]]:
        assert not self.is_iterator, 'Use iter_from to call an iterator function.'
        return self.forward(args, return_type=return_type, auto_broadcast=auto_broadcast, wrap_rv=wrap_rv)

    def iter_from(self, *args, return_type: Optional[Union[TensorValueTypeBase, PyObjValueType]] = None, auto_broadcast: bool = True, wrap_rv: bool = True) -> Union[Iterator[TensorValue], Iterator[Tuple[TensorValue, ...]]]:
        assert self.is_iterator, 'Use __call__ to call a non-iterator function.'
        return self.forward_generator(args, return_type=return_type, auto_broadcast=auto_broadcast, wrap_rv=wrap_rv)

    def _wrap_rv(self, rv, return_type, argument_values, auto_broadcast):
        if isinstance(rv, TensorValue):
            return rv
        elif isinstance(rv, tuple) and all(isinstance(v, TensorValue) for v in rv):
            return rv

        if return_type is None:
            return_type = self.return_type
        if return_type is None:
            raise RuntimeError('Return type can not be None if the function return is not a TensorValue.')

        if isinstance(return_type, tuple):
            if not isinstance(rv, tuple) and len(return_type) == 1:
                rv = (rv, )
            return tuple(self._wrap_single_rv(v, t, argument_values, auto_broadcast) for v, t in zip(rv, return_type))
        else:
            return self._wrap_single_rv(rv, return_type, argument_values, auto_broadcast)

    def _wrap_single_rv(self, rv, return_type, argument_values, auto_broadcast):
        if not self.support_batch:
            if isinstance(return_type, PyObjValueType):
                if isinstance(rv, TensorizedPyObjValues):
                    return TensorValue.from_tensorized_pyobj(rv, return_type)
                return TensorValue.from_scalar(rv, return_type)
            else:
                if isinstance(rv, torch.Tensor):
                    return TensorValue.from_tensor(rv, return_type)
                elif isinstance(rv, (bool, int, float)):
                    return TensorValue.from_scalar(rv, return_type)
                else:
                    raise ValueError(f'Unsupported return type: {type(rv)}')
        else:
            if isinstance(return_type, PyObjValueType):
                raise TypeError('Cannot return a PyObjValueType for a batched function.')
            else:
                if isinstance(rv, torch.Tensor):
                    first_tensor_arg = None
                    for arg in argument_values:
                        if isinstance(arg, TensorValue):
                            first_tensor_arg = arg
                            break
                    if not self.auto_broadcast or not auto_broadcast or first_tensor_arg is None:
                        raise ValueError('Cannot return a raw PyTorch tensor for a batched function without auto_broadcast.')
                    return TensorValue.from_tensor(rv, return_type, batch_variables=first_tensor_arg.batch_variables)
                else:
                    raise ValueError(f'Unsupported return type: {type(rv)}')


class PythonFunctionCrossRef(object):
    def __init__(self, cross_ref_name: str):
        self.cross_ref_name = cross_ref_name


def config_function_implementation(function: Optional[Callable] = None, *, function_quantized: Optional[Callable] = None, support_batch: bool = False, auto_broadcast: bool = True, use_object_names: bool = True, unwrap_values: Optional[bool] = None, include_executor_args: bool = False, is_iterator: bool = False) -> Callable:
    """Configure the implementation of a function in a domain.

    Args:
        function: the function to be wrapped.
        function_quantized: the quantized version of the function (can be None).
        support_batch: whether the function supports batched inputs.
        auto_broadcast: whether the executor should automatically broadcast the arguments before calling the function.
        use_object_names: whether the executor should use object names instead of indices.
        unwrap_values: whether the executor should unwrap the values before calling the function.
        include_executor_args: whether the executor should include itself as the first argument.
        is_iterator: whether the function is an iterator.

    Returns:
        the decorator.
    """

    function_implementation_configs = {
        'function_quantized': function_quantized,
        'support_batch': support_batch,
        'auto_broadcast': auto_broadcast,
        'use_object_names': use_object_names,
        'unwrap_values': unwrap_values,
        'include_executor_args': include_executor_args,
        'is_iterator': is_iterator
    }

    def wrapper(function: Callable, configs=function_implementation_configs):
        return PythonFunctionRef(function, **configs)

    if function is None:
        return wrapper
    return wrapper(function)


def _check_no_quantized_arguments(arguments):
    """A helper function to check that there are no quantized arguments. This function handles the migration from quantized tensor CSP computation to non-quantized tensor CSP computation."""
    # TODO(Jiayuan Mao @ 2023/08/15): remove this after the migration.
    for arg in arguments:
        if isinstance(arg, TensorValue):
            if isinstance(arg.dtype, TensorValueTypeBase):
                if arg.quantized and not arg.dtype.is_intrinsically_quantized():
                    raise RuntimeError('Quantized arguments are not supported.')


class PDSketchExecutor(TensorValueExecutorBase):
    """Planning domain expression executor. The basic concept in planning domain execution is the "state"."""

    def __init__(self, domain: Domain, parser: Optional[ParserBase] = None):
        """Initialize a PDSketch expression executor.

        Args:
            domain: the domain of this executor.
            parser: the parser to be used. This argument is optional. If provided, the execute function can take strings as input.
        """
        super().__init__(domain, parser)
        self._default_visitor = PDSketchExecutionDefaultVisitor(self)
        self._csp_visitor = PDSketchExecutionCSPVisitor(self)
        self._register_default_function_implementations()
        self._effect_update_from_simulation = False
        self._effect_action_index = None

    def _register_default_function_implementations(self):
        for t in self.domain.types.values():
            if isinstance(t, NamedTensorValueType):
                if isinstance(t.base_type, ScalarValueType):
                    self.register_function_implementation(
                        f'type::{t.typename}::equal',
                        PythonFunctionRef(
                            lambda x, y: TensorValue(BOOL, x.batch_variables, torch.eq(x.tensor, y.tensor), x.batch_dims),
                            support_batch=True, auto_broadcast=True, unwrap_values=False
                        )
                    )
                    self.register_function_implementation(
                        f'type::{t.typename}::hash',
                        PythonFunctionRef(lambda x: x.tensor.item(), support_batch=False, unwrap_values=False)
                    )
                elif isinstance(t.base_type, VectorValueType):
                    self.register_function_implementation(
                        f'type::{t.typename}::equal',
                        PythonFunctionRef(
                            lambda x, y: TensorValue(BOOL, x.batch_variables, torch.eq(x.tensor, y.tensor).all(dim=-1), x.batch_dims),
                            support_batch=True, auto_broadcast=True, unwrap_values=False
                        )
                    )

        for fname, cross_ref_name in self.domain.external_function_crossrefs.items():
            self.register_function_implementation(fname, PythonFunctionCrossRef(cross_ref_name))

    _domain: Domain

    @property
    def domain(self) -> Domain:
        return self._domain

    @property
    def effect_update_from_simulation(self) -> bool:
        """A context variable indicating whether the current effect should be updated from simulation, instead of the evaluation of expressions."""
        return self._effect_update_from_simulation

    @property
    def effect_action_index(self) -> Optional[int]:
        return self._effect_action_index

    def parse(self, string: Union[str, Expression], variables: Optional[Sequence[Variable]] = None) -> Expression:
        if isinstance(string, Expression):
            return string
        if self.parser is not None:
            return self.parser.parse_expression(string)
        return self._domain.parse(string, variables)

    _function_implementations: Dict[str, Union[PythonFunctionRef, PythonFunctionCrossRef]]

    @property
    def function_implementations(self) -> Dict[str, Union[PythonFunctionRef, PythonFunctionCrossRef]]:
        return self._function_implementations

    def register_function_implementation(self, name: str, func: Union[Callable, PythonFunctionRef, PythonFunctionCrossRef]):
        if isinstance(func, PythonFunctionRef):
            self._function_implementations[name] = func.set_executor(self)
        elif isinstance(func, PythonFunctionCrossRef):
            self._function_implementations[name] = func
        else:
            self._function_implementations[name] = PythonFunctionRef(func)

    def get_function_implementation(self, name: str) -> PythonFunctionRef:
        while name in self._function_implementations:
            func = self._function_implementations[name]
            if isinstance(func, PythonFunctionCrossRef):
                name = func.cross_ref_name
            else:
                return func
        raise KeyError(f'Function {name} not found.')

    def _execute(self, expression: Expression) -> TensorValueExecutorReturnType:
        if self.csp is not None:
            return self._csp_visitor.visit(expression)
        return self._default_visitor.visit(expression)

    def apply_precondition(self, operator: Union[Operator, OperatorApplier], state: TensorState, *args, csp: Optional[ConstraintSatisfactionProblem] = None) -> bool:
        """Apply the precondition of this operator to the given state.

        Args:
            operator: the operator to be applied.
            state: the state to be applied to.
            args: the arguments of the operator.
            csp: the CSP to be used for optimistic evaluation.

        Returns:
            bool: whether the precondition is satisfied.
        """

        if isinstance(operator, OperatorApplier):
            assert not operator.is_macro
            assert len(args) == 0, 'Operator applier does not support arguments.'
            operator, args = operator.operator, operator.arguments
        else:
            # NB(Jiayuan Mao @ 2023/03/17): sanity check due to the addition of macro operators.
            assert isinstance(operator, Operator)

        bounded_variables = compose_bvdict_args(operator.arguments, args, state=state)
        try:
            with self.with_state(state), self.with_csp(csp), self.with_bounded_variables(bounded_variables):
                for precondition in operator.preconditions:
                    pred_value = self._execute(precondition.bool_expr)
                    rv = pred_value.item()

                    if isinstance(rv, OptimisticValue):
                        if self.csp is not None:
                            self.csp.add_constraint(EqualityConstraint.from_bool(rv, True), note=f'precondition::{precondition.bool_expr.cached_string(-1)}')
                    else:
                        if rv < 0.5:
                            return False
            return True
        except Exception:
            logger.warning(f'Precondition evaluation failed: {precondition.bool_expr}.')
            raise

    def apply_precondition_debug(self, operator: Union[Operator, OperatorApplier], state: TensorState, *args, csp: Optional[ConstraintSatisfactionProblem] = None):
        """Apply the precondition of this operator to the given state, but in a dry-run mode. It will print out the evaluation results of each precondition.

        Example:
            .. code-block:: python

                (succ, state), csp = executor.apply(action, state, csp=csp)
                if succ:
                    pass
                else:
                    executor.apply_precondition_debug(action, state, csp=csp)  # will print out detailed infomation why the precondition is not satisfied.

        Args:
            operator: the operator to be applied.
            state: the state to be applied to.
            args: the arguments of the operator.
            csp: the CSP to be used for optimistic evaluation.
        """
        if isinstance(operator, OperatorApplier):
            assert not operator.is_macro
            assert len(args) == 0, 'Operator applier does not support arguments.'
            operator, args = operator.operator, operator.arguments
        else:
            # NB(Jiayuan Mao @ 2023/03/17): sanity check due to the addition of macro operators.
            assert isinstance(operator, Operator)

        bounded_variables = compose_bvdict_args(operator.arguments, args, state=state)
        try:
            with self.with_state(state), self.with_csp(csp), self.with_bounded_variables(bounded_variables):
                for precondition in operator.preconditions:
                    logger.info(f'Evaluate precondition: {precondition.bool_expr}')
                    pred_value = self._execute(precondition.bool_expr)
                    rv = pred_value.item()
                    logger.info(f'  Result = {rv}')

                    if isinstance(rv, OptimisticValue):
                        if self.csp is not None:
                            constraint = EqualityConstraint.from_bool(rv, True)
                            logger.info(f'  Add constraint: {constraint}')
        except Exception:
            logger.warning(f'Precondition evaluation failed: {precondition.bool_expr}.')
            raise

    def apply_effect(self, operator: Union[Operator, OperatorApplier], state: TensorState, *args, csp: Optional[ConstraintSatisfactionProblem] = None, action_index: Optional[int] = None, clone: bool = True) -> TensorState:
        """Apply the effect of this operator to the given state.

        Args:
            operator: the operator to be applied.
            state: the state to be applied to.
            args: the arguments of the operator.
            csp: the CSP to be used for optimistic evaluation.
            action_index: the index of the action in the trajectory (only effective when ``effect_update_from_simulation`` is active).
            clone: whether to clone the state before applying the effect.

        Returns:
            the new state after applying the effect.
        """
        if isinstance(operator, OperatorApplier):
            assert not operator.is_macro
            assert len(args) == 0, 'Operator applier does not support arguments.'
            operator, args = operator.operator, operator.arguments
        else:
            # NB(Jiayuan Mao @ 2023/03/17): sanity check due to the addition of macro operators.
            assert isinstance(operator, Operator)

        if clone:
            state = state.clone()

        bounded_variables = compose_bvdict_args(operator.arguments, args, state=state)

        try:
            with self.with_state(state), self.with_csp(csp), self.with_bounded_variables(bounded_variables):
                self._effect_action_index = action_index
                for effect in operator.effects:
                    self._effect_update_from_simulation = effect.update_from_simulation
                    self._execute(effect.assign_expr)
            return state
        except Exception:
            self._effect_action_index = None
            self._effect_update_from_simulation = False
            logger.warning(f'Effect application failed: {effect.assign_expr}.')
            raise

    def apply(self, operator: Union[Operator, OperatorApplier], state: TensorState, *args, clone: bool = True, csp: Optional[ConstraintSatisfactionProblem] = None, action_index: Optional[int] = None) -> Tuple[bool, TensorState]:
        """Apply an operator to a state.

        Args:
            operator: the operator to be applied.
            state: the state to be applied to.
            args: the arguments of the operator.
            clone: whether to clone the state before applying the effect.
            csp: the CSP to be used for optimistic evaluation.
            action_index: the index of the action in the trajectory (only effective when ``effect_update_from_simulation`` is active).

        Returns:
            a tuple of (whether the precondition is satisfied, the new state after applying the effect).
        """

        if isinstance(operator, (Operator, MacroOperator)):
            assert not operator.is_macro, 'Macro operators are not supported.'
            operator = operator(*args)
        else:
            assert not operator.operator.is_macro, 'Macro operators are not supported.'

        if self.apply_precondition(operator, state, csp=csp):
            return True, self.apply_effect(operator, state, csp=csp, clone=clone, action_index=action_index)
        return False, state

    def get_controller_args(self, operator: Union[Operator, OperatorApplier], state: TensorState, *args) -> Tuple[TensorValueExecutorReturnType, ...]:
        """Get the arguments of the controller of the given operator.

        Args:
            operator: the operator to be applied.
            state: the state to be applied to.
            args: the arguments of the operator.

        Returns:
            the arguments to the controller.
        """
        if isinstance(operator, OperatorApplier):
            assert not operator.is_macro
            assert len(args) == 0, 'Operator applier does not support arguments.'
            operator, args = operator.operator, operator.arguments
        else:
            # NB(Jiayuan Mao @ 2023/03/17): sanity check due to the addition of macro operators.
            assert isinstance(operator, Operator)

        bounded_variables = compose_bvdict_args(operator.arguments, args, state=state)

        try:
            with self.with_state(state), self.with_bounded_variables(bounded_variables):
                arguments = list()
                for arg in operator.controller.arguments:
                    arguments.append(self._execute(arg))
                return tuple(arguments)
        except Exception:
            logger.warning(f'Controller argument application failed: {arg}.')
            raise

    def new_state(self, object_names: ObjectNameArgument, object_types: ObjectTypeArgument = None, create_context: bool = False) -> Union[TensorState, Tuple[TensorState, 'StateDefinitionHelper']]:
        """Create a new state. This function also creates the state definition helper if the `create_context` argument is True. See the documentation of `StateDefinitionHelper` for more details.

        Args:
            object_names: the object names. It can be a list of strings (the names of the objects), or a dictionary mapping object names to their types. In this case, the `object_types` argument should be None.
            object_types: the object types, which is a list of :class:`~concepts.dsl.dsl_types.ObjectType` instances.
            create_context: whether to create the state definition helper.

        Returns:
            the new state if `create_context` is False, otherwise a tuple of the new state and the state definition helper.
        """

        if isinstance(object_names, dict):
            assert object_types is None
            object_types = tuple(object_names.values())
            object_names = tuple(object_names.keys())

        object_types_list = list()
        for object_type in object_types:
            if isinstance(object_type, str):
                object_types_list.append(self.domain.types[object_type])
            else:
                object_types_list.append(object_type)
        object_types = tuple(object_types_list)

        state = State({}, object_names, object_types)
        if create_context:
            return state, StateDefinitionHelper(self, state)
        return state


class PDSketchExecutionDefaultVisitor(ExpressionVisitor):
    """The underlying default implementation for :class:`PDSketchExecutor`. This function does not handle CSPs (a.k.a. optimistic execution)."""

    def __init__(self, executor: PDSketchExecutor):
        """Initialize a PDExpressionExecutionDefaultVisitor.

        Args:
            executor: the executor that uses this visitor.
        """
        self.executor = executor

    @property
    def csp(self) -> ConstraintSatisfactionProblem:
        return self.executor.csp

    def visit_variable_expression(self, expr: E.VariableExpression) -> TensorValueExecutorReturnType:
        variable = expr.variable
        return self.executor.bounded_variables[variable.dtype.typename][variable.name]

    def visit_object_constant_expression(self, expr: E.ObjectConstantExpression) -> StateObjectReference:
        const = expr.constant
        state = self.executor.state
        assert isinstance(state, NamedObjectTensorState)
        return StateObjectReference(
            const.name,
            state.get_typed_index(const.name, const.dtype.typename)
        )

    def visit_constant_expression(self, expr: E.ConstantExpression) -> TensorValueExecutorReturnType:
        value = expr.constant
        assert isinstance(value, TensorValue)
        return value

    def visit_function_application_expression(self, expr: E.FunctionApplicationExpression) -> TensorValueExecutorReturnType:
        function = expr.function
        return_type = function.return_type
        state = self.executor.state
        assert isinstance(function, Predicate)

        if function.is_generator_placeholder:  # always true branch
            argument_values = self.forward_args(*expr.arguments, force_tuple=True)
            argument_values = expand_argument_values(argument_values)
            _check_no_quantized_arguments(argument_values)

            batched_value = None
            for argv in argument_values:
                if isinstance(argv, TensorValue):
                    batched_value = argv
                    break
            assert batched_value is not None

            rv = torch.ones(
                batched_value.tensor.shape[:batched_value.total_batch_dims],
                dtype=torch.bool, device=batched_value.tensor.device
            )
            assert return_type == BOOL

            # Add "true" asserts to the csp.
            if self.csp is not None and not self.executor.optimistic_execution:
                expr_string = expr.cached_string(-1)
                for ind in _iter_tensor_indices(rv):
                    self.csp.add_constraint(Constraint.from_function(
                        function,
                        [argv.fast_index(tuple(ind)) for argv in argument_values],
                        True
                    ), note=f'{expr_string}::{ind}' if len(ind) > 0 else expr_string)

            return TensorValue(
                BOOL, batched_value.batch_variables,
                rv, batch_dims=state.batch_dims
            )
        elif function.is_cacheable and function.name in state.features:
            argument_values = self.forward_args(*expr.arguments, force_tuple=True)
            argument_values = [v.index if isinstance(v, StateObjectReference) else v for v in argument_values]
            batch_variables = [arg.name for arg, value in zip(expr.arguments, argument_values) if value == QINDEX]
            value = state.features[function.name][tuple(argument_values)]
            return value.rename_batch_variables(batch_variables)
        elif function.is_derived:
            argument_values = self.forward_args(*expr.arguments, force_tuple=True)
            with self.executor.with_bounded_variables({k: v for k, v in zip(function.arguments, argument_values)}):
                return self.visit(function.derived_expression)
        else:
            # dynamic predicate is exactly the same thing as a pre-defined external function.
            argument_values = self.forward_args(*expr.arguments, force_tuple=True)
            # only supports external functions with a single return value.
            return self.forward_external_function('predicate::' + function.name, argument_values, return_type, expression=expr)

    def visit_bool_expression(self, expr: E.BoolExpression, argument_values: Optional[Tuple[TensorValueExecutorReturnType, ...]] = None) -> TensorValueExecutorReturnType:
        if argument_values is None:
            argument_values = self.forward_args(*expr.arguments, force_tuple=True)
            argument_values = expand_argument_values(argument_values)

        _check_no_quantized_arguments(argument_values)

        assert len(argument_values) > 0
        assert all(isinstance(v, TensorValue) for v in argument_values)

        dtype = argument_values[0].dtype
        batch_variables = argument_values[0].batch_variables

        if expr.bool_op is BoolOpType.NOT:
            assert len(argument_values) == 1
            return TensorValue(
                dtype, batch_variables,
                1 - argument_values[0].tensor,
                batch_dims=self.executor.state.batch_dims
            )
        elif expr.bool_op is BoolOpType.AND:
            if len(argument_values) == 1:
                return argument_values[0]
            return TensorValue(
                dtype, batch_variables,
                MaskedTensorStorage(torch.stack([argv.tensor for argv in argument_values], dim=-1).amin(dim=-1), None, argument_values[0].tensor_mask),
                batch_dims=self.executor.state.batch_dims
            )
        elif expr.bool_op is BoolOpType.OR:
            if len(argument_values) == 1:
                return argument_values[0]
            return TensorValue(
                dtype, batch_variables,
                MaskedTensorStorage(torch.stack([argv.tensor for argv in argument_values], dim=-1).amax(dim=-1), None, argument_values[0].tensor_mask),
                batch_dims=self.executor.state.batch_dims
            )
        else:
            raise ValueError(f'Unknown bool op type: {expr.bool_op}.')

    def visit_quantification_expression(self, expr: E.QuantificationExpression, value: Optional[TensorValue] = None) -> TensorValueExecutorReturnType:
        if value is None:
            with self.executor.new_bounded_variables({expr.variable: QINDEX}):
                value = self.forward_args(expr.expression)
            assert isinstance(value, TensorValue)
        _check_no_quantized_arguments([value])

        dtype = value.dtype
        batch_variables = value.batch_variables
        variable_index = batch_variables.index(expr.variable.name)
        batch_variables = batch_variables[:variable_index] + batch_variables[variable_index + 1:]

        if value.tensor_mask is None:
            tensor = value.tensor
            mask = None
        else:
            tensor, mask = value.tensor, value.tensor_mask
            if expr.quantification_op is QuantificationOpType.FORALL:
                tensor = (tensor * mask + (1 - mask)).to(tensor.dtype)
            elif expr.quantification_op is QuantificationOpType.EXISTS:
                tensor = (tensor * mask).to(tensor.dtype)
            else:
                raise ValueError(f'Unknown quantification op type: {expr.quantification_op}.')

        if expr.quantification_op is QuantificationOpType.FORALL:
            return TensorValue(
                dtype, batch_variables,
                MaskedTensorStorage(tensor.amin(dim=variable_index), None, mask),
                batch_dims=self.executor.state.batch_dims
            )
        elif expr.quantification_op is QuantificationOpType.EXISTS:
            return TensorValue(
                dtype, batch_variables,
                MaskedTensorStorage(tensor.amax(dim=variable_index), None, mask),
                batch_dims=self.executor.state.batch_dims
            )
        else:
            raise ValueError(f'Unknown quantifier type: {expr.quantification_op}.')

    def visit_predicate_equal_expression(self, expr: E.PredicateEqualExpression, feature: Optional[TensorValue] = None, value: Optional[TensorValue] = None) -> TensorValueExecutorReturnType:
        if feature is None or value is None:
            feature, value = self.forward_args(expr.predicate, expr.value)
            feature, value = expand_argument_values([feature, value])
        _check_no_quantized_arguments([feature, value])

        if isinstance(feature.dtype, PyObjValueType):
            rv = self.forward_external_function(f'type::{feature.dtype.typename}::equal', [feature, value], BOOL)
        elif isinstance(feature.dtype, NamedTensorValueType):
            rv = self.forward_external_function(f'type::{feature.dtype.typename}::equal', [feature, value], BOOL)
        else:
            raise NotImplementedError(f'Unsupported FeatureEqual computation for dtype {feature.dtype} and {value.dtype}.')

        return rv

    def visit_assign_expression(self, expr: E.AssignExpression):
        state = self.executor.state
        argument_values = self.forward_args(*expr.predicate.arguments, force_tuple=True)
        value = self.forward_args(expr.value)
        argument_values = [v.index if isinstance(v, StateObjectReference) else v for v in argument_values]

        # if state.features[expr.predicate.function.name].quantized:
        #     if not value.quantized:
        #         value = self.executor.value_quantizer.quantize_value(value)
        # else:
        #     if value.quantized:
        #         value = self.executor.value_quantizer.unquantize_value(value)
        _check_no_quantized_arguments([value, *argument_values, state.features[expr.predicate.function.name]])

        state.features[expr.predicate.function.name][tuple(argument_values)] = value

    def visit_conditional_select_expression(self, expr: E.ConditionalSelectExpression) -> TensorValueExecutorReturnType:
        value, condition = self.forward_args(expr.predicate, expr.condition)
        value = value.clone()
        if value.tensor_mask is None:
            value.tensor_mask = condition.tensor
        else:
            value.tensor_mask = torch.min(value.tensor_mask, condition.tensor)
        return value

    def visit_deictic_select_expression(self, expr: E.DeicticSelectExpression) -> TensorValueExecutorReturnType:
        with self.executor.new_bounded_variables({expr.variable: QINDEX}):
            return self.visit(expr.expression)

    CONDITIONAL_ASSIGN_QUANTIZE = False

    def visit_conditional_assign_expression(self, expr: E.ConditionalAssignExpression):
        state = self.executor.state
        argument_values = self.forward_args(*expr.predicate.arguments, force_tuple=True)
        value = self.forward_args(expr.value)
        condition = self.forward_args(expr.condition)
        _check_no_quantized_arguments([value, condition, state.features[expr.predicate.function.name]])

        condition_tensor = jactorch.quantize(condition.tensor) if self.CONDITIONAL_ASSIGN_QUANTIZE else condition.tensor

        feature = state.features[expr.predicate.function.name]
        origin_tensor = feature[argument_values].tensor  # I am not using feature.tensor[argument_values] because the current code will handle TensorizedPyObjValues too.
        # assert value.tensor.dim() == condition_tensor.dim() or value.tensor.dim() == 0

        if value.is_tensorized_pyobj:
            raise NotImplementedError('Cannot make conditional assignments for tensorized pyobj.')
        else:
            if condition_tensor.dim() < value.tensor.dim():
                condition_tensor = condition_tensor.unsqueeze(-1)
            state.features[expr.predicate.function.name].tensor[argument_values] = (
                condition_tensor.to(origin_tensor.dtype) * value.tensor + (1 - condition_tensor).to(origin_tensor.dtype) * origin_tensor
            )

    def visit_deictic_assign_expression(self, expr: E.DeicticAssignExpression):
        with self.executor.new_bounded_variables({expr.variable: QINDEX}):
            self.visit(expr.expression)

    def forward_args(self, *args, force_tuple: bool = False) -> Union[TensorValueExecutorReturnType, Tuple[TensorValueExecutorReturnType, ...]]:
        if len(args) == 1 and not force_tuple:
            return self.visit(args[0])
        else:
            return tuple(self.visit(arg) for arg in args)

    def forward_external_function(
        self, function_name: str, argument_values: Sequence[TensorValueExecutorReturnType],
        return_type: Union[TensorValueTypeBase, PyObjValueType], auto_broadcast: bool = False, expression: Optional[Expression] = None
    ) -> TensorValue:
        external_function = self.executor.get_function_implementation(function_name)
        assert isinstance(external_function, PythonFunctionRef)
        return external_function.forward(argument_values, return_type=return_type, auto_broadcast=auto_broadcast)


class PDSketchExecutionCSPVisitor(PDSketchExecutionDefaultVisitor):
    def __init__(self, executor: PDSketchExecutor):
        super().__init__(executor)

    def forward_external_function(
        self, function_name: str, argument_values: Sequence[TensorValueExecutorReturnType],
        return_type: Union[TensorValueTypeBase, PyObjValueType], auto_broadcast: bool = True, expression: Optional[E.FunctionApplicationExpression] = None
    ) -> TensorValue:
        argument_values = expand_argument_values(argument_values)
        optimistic_masks = [is_optimistic_value(argv.tensor_optimistic_values) for argv in argument_values if argv.tensor_optimistic_values is not None]
        if len(optimistic_masks) > 0:
            optimistic_mask = torch.stack(optimistic_masks, dim=-1).any(dim=-1)

            # TODO(Jiayuan Mao @ 2023/01/26): fix this hack, add a proper flag to indicate whether a function has to be executed in the simulator.
            if not self.executor.has_function_implementation('predicate::' + function_name):
                optimistic_mask[...] = True

            rv = super().forward_external_function(function_name, argument_values, return_type=return_type, auto_broadcast=auto_broadcast, expression=expression)

            if optimistic_mask.sum().item() == 0:
                return rv

            rv.init_tensor_optimistic_values()

            if self.executor.optimistic_execution:
                rv.tensor_optimistic_values[optimistic_mask.nonzero(as_tuple=True)] = OPTIM_MAGIC_NUMBER_MAGIC
            else:
                expr_string = expression.cached_string(-1)
                for ind in optimistic_mask.nonzero().tolist():
                    ind = tuple(ind)
                    new_identifier = self.executor.csp.new_var(return_type)
                    rv.tensor_optimistic_values[ind] = new_identifier
                    self.csp.add_constraint(Constraint.from_function(
                        expression.function,
                        [argv.fast_index(ind) for argv in argument_values],
                        new_identifier
                    ), note=f'{expr_string}::{ind}' if len(ind) > 0 else expr_string)
            return rv

            # TODO(Jiayuan Mao @ 2023/08/16): implement a faster verion of the computation.
            # if optimistic_mask.sum().item() == 0:
            #     pass  # just do the standard execution.
            # else:
            #     retain_mask = torch.logical_not(optimistic_mask)
            #     rv = torch.zeros(
            #         argument_values[0].tensor.shape,
            #         dtype=torch.int64,
            #         device=argument_values[0].tensor.device
            #     )

            #     if retain_mask.sum().item() > 0:
            #         argument_values_r = [TensorValue(argv.dtype, ['?x'], argv.tensor[retain_mask], 0, quantized=argv.quantized) for argv in argument_values]
            #         # No need to broadcast again.
            #         rv_r = super().forward_external_function(function_name, argument_values_r, return_type, auto_broadcast=False)
            #         if not rv_r.quantized:
            #             rv_r = executor.value_quantizer.quantize_value(rv_r)
            #         rv[retain_mask] = rv_r.tensor

            #     expr_string = expression.cached_string(-1)
            #     for ind in optimistic_mask.nonzero().tolist():
            #         ind = tuple(ind)
            #         new_identifier = self.executor.csp.new_var(return_type)
            #         rv[ind] = new_identifier
            #         self.csp.add_constraint(Constraint.from_function(
            #             expression.function,
            #             [argv.tensor[ind].item() if argv.quantized else argv.tensor[ind] for argv in argument_values],
            #             new_identifier
            #         ), note=f'{expr_string}::{ind}' if len(ind) > 0 else expr_string)

            #     return TensorValue(
            #         expression.function.return_type, argument_values[0].batch_variables if len(argument_values) > 0 else [],
            #         rv,
            #         batch_dims=executor.state.batch_dims, quantized=True
            #     )

        return super().forward_external_function(function_name, argument_values, return_type=return_type, auto_broadcast=auto_broadcast, expression=expression)

    def visit_function_application_expression(self, expr: E.FunctionApplicationExpression) -> TensorValueExecutorReturnType:
        return super().visit_function_application_expression(expr)

    def visit_bool_expression(self, expr: E.BoolExpression, argument_values: Optional[Tuple[TensorValueExecutorReturnType, ...]] = None) -> TensorValueExecutorReturnType:
        if argument_values is None:
            argument_values = self.forward_args(*expr.arguments, force_tuple=True)
            argument_values = list(expand_argument_values(argument_values))

        for argv in argument_values:
            assert argv.dtype == BOOL, 'Boolean expression only supports boolean values in CSP mode.'

        optimistic_masks = [is_optimistic_value(argv.tensor_optimistic_values) for argv in argument_values if argv.tensor_optimistic_values is not None]
        if len(optimistic_masks) > 0:
            optimistic_mask = torch.stack(optimistic_masks, dim=-1).any(dim=-1)
            if optimistic_mask.sum().item() > 0:
                rv = super().visit_bool_expression(expr, argument_values)
                rv.init_tensor_optimistic_values()

                if self.executor.optimistic_execution:
                    rv.tensor_optimistic_values[optimistic_mask.nonzero(as_tuple=True)] = OPTIM_MAGIC_NUMBER_MAGIC
                else:
                    expr_string = expr.cached_string(-1)
                    for ind in optimistic_mask.nonzero().tolist():
                        ind = tuple(ind)

                        this_argv = [argv.fast_index(ind, wrap=False) for argv in argument_values]
                        determined = None
                        if expr.return_type == BOOL:
                            if expr.bool_op is BoolOpType.NOT:
                                pass  # nothing we can do.
                            elif expr.bool_op is BoolOpType.AND:
                                if 0 in this_argv or False in this_argv:
                                    determined = False
                            elif expr.bool_op is BoolOpType.OR:
                                if 1 in this_argv or True in this_argv:
                                    determined = True
                            this_argv = [v for v in this_argv if isinstance(v, OptimisticValue)]
                        else:  # generalized boolean operations.
                            pass

                        if determined is None:
                            new_identifier = self.csp.new_var(BOOL)
                            rv.tensor_optimistic_values[ind] = new_identifier
                            self.csp.add_constraint(Constraint(
                                expr.bool_op,
                                this_argv,
                                cvt_opt_value(new_identifier, BOOL),
                            ), note=f'{expr_string}::{ind}' if len(ind) > 0 else expr_string)
                        else:
                            rv[ind] = determined
                return rv
            else:
                return super().visit_bool_expression(expr, argument_values)
        else:  # if len(optimistic_masks) == 0
            return super().visit_bool_expression(expr, argument_values)

    def visit_quantification_expression(self, expr: E.QuantificationExpression, value: Optional[TensorValue] = None) -> Any:
        if value is None:
            with self.executor.new_bounded_variables({expr.variable: QINDEX}):
                value = self.forward_args(expr.expression)
            assert isinstance(value, TensorValue)

        assert value.dtype == BOOL, 'Quantification expression only supports boolean values in CSP mode.'

        value.init_tensor_optimistic_values()
        rv = super().visit_quantification_expression(expr, value)

        dim = value.batch_variables.index(expr.variable.name) + value.batch_dims
        value_transposed = value.tensor
        optimistic_values_transposed = value.tensor_optimistic_values
        if dim != value.tensor.ndim - 1:
            value_transposed = value_transposed.transpose(dim, -1)  # put the target dimension last.
            optimistic_values_transposed = optimistic_values_transposed.transpose(dim, -1)
        optimistic_mask_transposed = is_optimistic_value(optimistic_values_transposed)

        value_transposed = torch.where(
            optimistic_mask_transposed,
            optimistic_values_transposed,
            value_transposed
        )
        optimistic_mask = optimistic_mask_transposed.any(dim=-1)

        if optimistic_mask.sum().item() == 0:
            return rv

        rv.init_tensor_optimistic_values()
        if self.executor.optimistic_execution:
            rv.tensor_optimistic_values[optimistic_mask.nonzero(as_tuple=True)] = OPTIM_MAGIC_NUMBER_MAGIC
        else:
            expr_string = expr.cached_string(-1)
            for ind in optimistic_mask.nonzero().tolist():
                ind = tuple(ind)

                this_argv = value_transposed[ind].tolist()
                determined = None
                if expr.quantification_op is QuantificationOpType.FORALL:
                    if 0 in this_argv or False in this_argv:
                        determined = False
                else:
                    if 1 in this_argv or True in this_argv:
                        determined = True
                this_argv = list(filter(is_optimistic_value, this_argv))

                if determined is None:
                    new_identifier = self.csp.new_var(BOOL)
                    rv.tensor_optimistic_values[ind] = new_identifier
                    self.csp.add_constraint(Constraint(
                        expr.quantification_op,
                        [OptimisticValue(value.dtype, int(v)) for v in this_argv],
                        OptimisticValue(value.dtype, new_identifier),
                    ), note=f'{expr_string}::{ind}' if len(ind) > 0 else expr_string)
                else:
                    rv.tensor[ind] = determined
        return rv

    def visit_predicate_equal_expression(self, expr: E.PredicateEqualExpression, feature: Optional[TensorValue] = None, value: Optional[TensorValue] = None) -> Any:
        if feature is None or value is None:
            feature, value = self.forward_args(expr.predicate, expr.value)
            feature, value = expand_argument_values([feature, value])

        feature.init_tensor_optimistic_values()
        value.init_tensor_optimistic_values()

        rv = super().visit_predicate_equal_expression(expr, feature, value)
        optimistic_mask = torch.logical_or(is_optimistic_value(feature.tensor_optimistic_values), is_optimistic_value(value.tensor_optimistic_values))

        if optimistic_mask.sum().item() == 0:
            return rv

        rv.init_tensor_optimistic_values()
        if self.executor.optimistic_execution:
            rv.tensor_optimistic_values[optimistic_mask.nonzero(as_tuple=True)] = OPTIM_MAGIC_NUMBER_MAGIC
        else:
            expr_string = expr.cached_string(-1)
            for ind in optimistic_mask.nonzero().tolist():
                ind = tuple(ind)
                this_argv = feature.fast_index(ind), value.fast_index(ind)
                new_identifier = self.csp.new_var(BOOL)
                rv.tensor_optimistic_values[ind] = new_identifier
                self.csp.add_constraint(EqualityConstraint(
                    *[cvt_opt_value(v, feature.dtype) for v in this_argv],
                    OptimisticValue(BOOL, new_identifier)
                ), note=f'{expr_string}::{ind}' if len(ind) > 0 else expr_string)
        return rv

    def visit_assign_expression(self, expr: E.AssignExpression) -> Any:
        if self.executor.effect_update_from_simulation:
            feature = self.executor.state.features[expr.predicate.function.name]
            feature.init_tensor_optimistic_values()
            argument_values = self.forward_args(*expr.predicate.arguments, force_tuple=True)

            assert self.executor.effect_action_index is not None, 'Effect action index must be set if the target predicate will be updated from simulation.'
            expr_string = expr.cached_string(-1)
            for entry_values in _expand_tensor_indices(feature, argument_values):
                if self.executor.optimistic_execution:
                    raise RuntimeError('Optimistic execution is not supported for effect update from simulation.')
                else:
                    opt_identifier = self.csp.new_var(feature.dtype, wrap=True)
                    feature.tensor_optimistic_values[entry_values] = opt_identifier.identifier
                    self.csp.add_constraint(Constraint(
                        SimulationFluentConstraintFunction(self.executor.effect_action_index, expr.predicate.function, entry_values),
                        [],
                        opt_identifier,
                        note=f'{expr_string}::{entry_values}' if len(entry_values) > 0 else expr_string
                    ))
        else:
            return super().visit_assign_expression(expr)

    def visit_conditional_select_expression(self, expr: E.ConditionalSelectExpression) -> TensorValueExecutorReturnType:
        return super().visit_conditional_select_expression(expr)

    def visit_deictic_select_expression(self, expr: E.DeicticSelectExpression) -> Any:
        return super().visit_deictic_select_expression(expr)

    def visit_conditional_assign_expression(self, expr: E.ConditionalAssignExpression) -> Any:
        if self.executor.effect_update_from_simulation:
            raise NotImplementedError('Conditional assign is not supported in simulation mode.')
        if self.executor.optimistic_execution:
            raise RuntimeError('Optimistic execution is not supported for conditional assign.')

        state = self.executor.state
        argument_values = self.forward_args(*expr.predicate.arguments, force_tuple=True)
        value = self.forward_args(expr.value)
        condition = self.forward_args(expr.condition)
        _check_no_quantized_arguments([value, condition, state.features[expr.predicate.function.name]])

        condition_tensor = jactorch.quantize(condition.tensor) if self.CONDITIONAL_ASSIGN_QUANTIZE else condition.tensor
        condition_tensor = (condition_tensor > 0.5).to(torch.bool)

        feature = state.features[expr.predicate.function.name]
        origin_tensor = feature[argument_values].tensor  # I am not using feature.tensor[argument_values] because the current code will handle TensorizedPyObjValues too.
        # assert value.tensor.dim() == condition_tensor.dim() or value.tensor.dim() == 0

        # NB(Jiayuan Mao @ 2023/08/15): conditional assignment does not support "soft" assignment.
        if value.is_tensorized_pyobj:
            raise NotImplementedError('Cannot make conditional assignments for tensorized pyobj.')
        else:
            if condition_tensor.dim() < value.tensor.dim():
                condition_tensor_expanded = condition_tensor.unsqueeze(-1)
            else:
                condition_tensor_expanded = condition_tensor

            feature.tensor[argument_values] = (
                condition_tensor_expanded.to(origin_tensor.dtype) * value.tensor + (1 - condition_tensor_expanded).to(origin_tensor.dtype) * origin_tensor
            )
            feature.init_tensor_optimistic_values()
            if value.tensor_optimistic_values is not None:
                feature.tensor_optimistic_values[argument_values] = (
                    condition_tensor.to(torch.int64) * value.tensor_optimistic_values +
                    (1 - condition_tensor).to(torch.int64) * state.features[expr.predicate.function.name].tensor_optimistic_values[argument_values]
                )

            if condition.tensor_optimistic_values is None:
                pass
            else:
                optimistic_mask = is_optimistic_value(condition.tensor_optimistic_values)
                if optimistic_mask.sum().item() == 0:
                    pass
                else:
                    expr_string = expr.cached_string(-1)
                    dtype = expr.predicate.function.return_type
                    for ind in optimistic_mask.nonzero().tolist():
                        ind = tuple(ind)

                        new_identifier = self.csp.new_var(dtype, wrap=True)
                        neg_condition_identifier = self.csp.new_var(BOOL, wrap=True)
                        eq_1_identifier = self.csp.new_var(BOOL, wrap=True)
                        eq_2_identifier = self.csp.new_var(BOOL, wrap=True)
                        condition_identifier = condition.tensor_optimistic_values[ind].item()

                        self.csp.add_constraint(EqualityConstraint(
                            new_identifier,
                            cvt_opt_value(value.fast_index(ind), dtype),
                            eq_1_identifier,
                        ), note=f'{expr_string}::{ind}::eq-1')
                        self.csp.add_constraint(EqualityConstraint(
                            new_identifier,
                            cvt_opt_value(origin_tensor.fast_index(ind) if isinstance(origin_tensor, TensorizedPyObjValues) else origin_tensor[ind], dtype),
                            eq_2_identifier
                        ), note=f'{expr_string}::{ind}::eq-2')

                        self.csp.add_constraint(Constraint(
                            BoolOpType.NOT,
                            [OptimisticValue(BOOL, condition_identifier)],
                            neg_condition_identifier
                        ), note=f'{expr_string}::{ind}::neg-cond')

                        self.csp.add_constraint(Constraint(
                            BoolOpType.OR,
                            [neg_condition_identifier, eq_1_identifier],
                            cvt_opt_value(True, BOOL)
                        ), note=f'{expr_string}::{ind}::implies-new')

                        self.csp.add_constraint(Constraint(
                            BoolOpType.OR,
                            [OptimisticValue(BOOL, condition_identifier), eq_2_identifier],
                            cvt_opt_value(True, BOOL)
                        ), note=f'{expr_string}::{ind}::implies-old')

                        feature.tensor_optimistic_values[ind] = new_identifier.identifier

    def visit_deictic_assign_expression(self, expr: E.DeicticAssignExpression) -> Any:
        return super().visit_deictic_assign_expression(expr)


def _iter_tensor_indices(target_tensor: torch.Tensor) -> Iterator[Tuple[int, ...]]:
    """Iterate from the indices of a tensor. """
    for ind in torch.nonzero(torch.ones_like(target_tensor)):
        yield tuple(ind.tolist())


def _expand_tensor_indices(target_value: TensorValue, input_indices: Tuple[Union[int, slice, StateObjectReference], ...]) -> Iterator[Tuple[int, ...]]:
    """Iterate over entry indices based on the input indices. Supported indices are int and QINDEX (:).

    Args:
        target_value: the target value, used to determine the size of ``QINDEX``.
        input_indices: the indices to iterate over.

    Yields:
        the entry indices.
    """
    indices = list()
    for i, ind in enumerate(input_indices):
        if isinstance(ind, int):
            indices.append(torch.tensor([ind], dtype=torch.int64))
        elif isinstance(ind, slice):
            assert ind.step is None and ind.start is None and ind.stop is None  # == ':'
            indices.append(torch.arange(target_value.tensor.shape[i], dtype=torch.int64))
        elif isinstance(ind, StateObjectReference):
            indices.append(torch.tensor([ind.index], dtype=torch.int64))
        else:
            raise ValueError(f'Invalid index type: {ind}')

    if len(indices) == 0:
        yield tuple()
        return

    if len(indices) == 1:
        for x in indices[0]:
            yield tuple([x.item()], )
        return

    indices = torch.meshgrid(*indices, indexing='ij')
    indices = [i.flatten() for i in indices]
    for i in range(len(indices[0])):
        yield tuple(indices[j][i].item() for j in range(len(indices)))


def expand_argument_values(argument_values: Sequence[Union[TensorValue, int, str, slice, StateObjectReference]]) -> List[TensorValue]:
    """Expand a list of argument values to the same batch size.

    Args:
        argument_values: a list of argument values.

    Returns:
        the result list of argument values. All return values will have the same batch size.
    """
    has_slot_var = False
    for arg in argument_values:
        if isinstance(arg, TensorValue):
            for var in arg.batch_variables:
                if var == '??':
                    has_slot_var = True
                    break
    if has_slot_var:
        return list(argument_values)

    if len(argument_values) < 2:
        return list(argument_values)

    argument_values = list(argument_values)
    batch_variables = list()
    batch_sizes = list()
    for arg in argument_values:
        if isinstance(arg, TensorValue):
            for var in arg.batch_variables:
                if var not in batch_variables:
                    batch_variables.append(var)
                    batch_sizes.append(arg.get_variable_size(var))
        else:
            assert isinstance(arg, (int, str, slice, StateObjectReference)), arg

    masks = list()
    for i, arg in enumerate(argument_values):
        if isinstance(arg, TensorValue):
            argument_values[i] = arg.expand(batch_variables, batch_sizes)
            if argument_values[i].tensor_mask is not None:
                masks.append(argument_values[i].tensor_mask)

    if len(masks) > 0:
        final_mask = torch.stack(masks, dim=-1).amin(dim=-1)
        for arg in argument_values:
            if isinstance(arg, TensorValue):
                arg.tensor_mask = final_mask
                arg._mask_certified_flag = True  # now we have corrected the mask.
    return argument_values


class _StateDefPredicate(object):
    def __init__(self, function, arguments):
        self.function = function
        self.arguments = arguments


class _StateDefPredicateApplier(object):
    def __init__(self, function):
        self.function = function

    def __call__(self, *args):
        return _StateDefPredicate(self.function, args)


class StateDefinitionHelper(object):
    """A helper class to define the planning domain state (:class:`~concepts.pdsketch.domain.State`). Typically you should use :meth:`~PDSketchExecutor.new_state` to create an instance of this helper class.

    The most important feature of this class is that it supports a human-friendly definition of Boolean predicates. For example, you can define the grounding values for a predicate ``p`` as follows:

    .. code-block:: python

        state, ctx = executor.new_state({'a': ObjectType('person'), 'b': ObjectType('person')})
        state.define_predicates([
            ctx.p('a', 'b')
        ])  # define the predicate p(a, b) as True. All other values are False by convention.

    """

    def __init__(self, executor: PDSketchExecutor, state: State):
        """Initialize a new instance of the state definition helper.

        Args:
            executor: the executor.
            state: a (possibly empty) state to be modified.
        """
        self.executor = executor
        self.domain = executor.domain
        self.state = state

    def get_predicate(self, name: str) -> _StateDefPredicateApplier:
        """Get a predicate definition helper by name. This can be used to define the grounding values of a predicate, e.g., ``ctx.get_predicate('p')('a', 'b')``.

        Args:
            name: the name of the predicate.

        Returns:
            a predicate definition helper.
        """
        if name in self.domain.functions:
            return _StateDefPredicateApplier(self.domain.functions[name])
        else:
            name = name.replace('_', '-')
            if name in self.domain.functions:
                return _StateDefPredicateApplier(self.domain.functions[name])
            else:
                raise NameError(f'Predicate {name} is not defined in the domain.')

    def __getattr__(self, name):
        return self.get_predicate(name)

    def define_predicates(self, predicates: Sequence[_StateDefPredicate]):
        """Define a list of grounding values of predicates. See the example in :class:`StateDefinitionHelper`.

        Args:
            predicates: a list of grounding values, created by ``ctx.p('a', 'b')``.
        """
        for function in self.domain.functions.values():
            if function.return_type != BOOL:
                continue
            if function.name in self.state.features.all_feature_names:
                continue

            if function.is_state_variable:
                sizes = list()
                for arg_def in function.arguments:
                    sizes.append(len(self.state.object_type2name[arg_def.typename]) if arg_def.typename in self.state.object_type2name else 0)
                self.state.features[function.name] = TensorValue.make_empty(BOOL, [var.name for var in function.arguments], tuple(sizes))

        for pred in predicates:
            assert isinstance(pred, _StateDefPredicate)
            assert pred.function.return_type == BOOL
            name = pred.function.name
            arguments = [self.state.get_typed_index(arg) for arg in pred.arguments]
            self.state.features[name].tensor[tuple(arguments)] = 1

    def define_feature(self, feature_name: str, tensor_or_mapping: Union[torch.Tensor, TensorizedPyObjValues, Mapping[Union[str, Tuple[str, ...]], Union[bool, int, float, torch.Tensor, TensorizedPyObjValues]]]):
        """Define a relational feature directly with :class:`torch.Tensor` objects. For example,

        .. code-block:: python

            state, ctx = executor.new_state({'a': ObjectType('person'), 'b': ObjectType('person')})
            # assume p is a Boolean predicate with the signature person x person -> bool
            state.define_feature('p', torch.tensor([[1, 0], [0, 1]]), quantized=True)

        Args:
            feature_name: the name of the feature.
            tensor_or_mapping: a tensor or a mapping from argument values to tensors. If a tensor is given, it is assumed that the tensor has exactly the same shape as the tensor of the feature.
                If the input is a mapping, the keys of the mapping are tuples (entry indices), and the values are tensors. The tensors will be filled into an all-zero tensor of the same shape as the feature tensor.
        """
        function = self.domain.functions[feature_name]
        sizes = list()
        for arg_def in function.arguments:
            sizes.append(len(self.state.object_type2name[arg_def.typename]) if arg_def.typename in self.state.object_type2name else 0)
        sizes = tuple(sizes)
        batch_variables = [var.name for var in function.arguments]
        return_type = function.return_type

        if isinstance(return_type, PyObjValueType):
            if isinstance(tensor_or_mapping, TensorizedPyObjValues):
                self.state.features[feature_name] = TensorValue.from_tensorized_pyobj(tensor_or_mapping, return_type, batch_variables)
                return
        else:
            if isinstance(tensor_or_mapping, torch.Tensor):
                self.state.features[feature_name] = TensorValue.from_tensor(tensor_or_mapping, return_type, batch_variables)
                return

        self.state.features[feature_name] = feature = TensorValue.make_empty(return_type, [var.name for var in function.arguments], sizes)
        for key, value in tensor_or_mapping.items():
            if isinstance(key, tuple):
                args = [self.state.get_typed_index(arg) for arg in key]
            else:
                assert isinstance(key, str)
                args = [self.state.get_typed_index(key)]
            feature[tuple(args)] = value

    def define_pyobj_feature(self, feature_name: str, pyobj_list: List[Any]):
        """Define a feature with a list of Python objects. The objects will be converted to tensors using the underlying :class:`~concepts.dsl.executors.tensor_value_executor.PyObjStore` of the executor.

        Args:
            feature_name: the name of the feature.
            pyobj_list: a list of Python objects (they can also be nested lists).
        """

        function = self.domain.functions[feature_name]
        sizes = list()
        for arg_def in function.arguments:
            sizes.append(len(self.state.object_type2name[arg_def.typename]) if arg_def.typename in self.state.object_type2name else 0)
        sizes = tuple(sizes)
        batch_variables = [var.name for var in function.arguments]
        return_type = function.return_type

        value = TensorValue.from_tensorized_pyobj(
            TensorizedPyObjValues(return_type, pyobj_list, sizes),
            dtype=return_type,
            batch_variables=batch_variables
        )

        self.state.features[feature_name] = value

    def set_value(self, feature_name, arguments: Union[Sequence[str], Sequence[int]], value: Union[torch.Tensor, bool, int, float, Any]):
        """Set a single entry in the feature representation. When the feature tensor has not been created, it will be created automatically.
        Note that, when creating the feature tensor, we will use ``quantized`` to determine whether the feature should be quantized or not in the state.

        .. code-block:: python

            state, ctx = executor.new_state({'a': ObjectType('person'), 'b': ObjectType('person')})
            # assume p is a Boolean predicate with the signature person x person -> bool
            state.set_value('p', ('a', 'b'), 1, quantized=True)

        Args:
            feature_name: the name of the feature.
            arguments: the argument values of the entry.
            value: the value of the entry.
        """

        function = self.domain.functions[feature_name]
        return_type = function.return_type

        if feature_name not in self.state.features:
            sizes = list()
            for arg_def in function.arguments:
                sizes.append(len(self.state.object_type2name[arg_def.typename]) if arg_def.typename in self.state.object_type2name else 0)
            sizes = tuple(sizes)
            self.state.features[feature_name] = TensorValue.make_empty(return_type, [var.name for var in function.arguments], sizes)

        assert len(arguments) == len(function.arguments)
        arguments = [self.state.get_typed_index(arg, arg_def.dtype.typename) if isinstance(arg, str) else arg for arg, arg_def in zip(arguments, function.arguments)]
        self.state.features[feature_name][tuple(arguments)] = value


class GeneratorManager(object):
    """The :class:`GeneratorManager` is used to manage calls to generators in the function dmoain. It is particularly useful for
    keep tracking of historical values generated by the generators."""

    def __init__(self, executor: PDSketchExecutor, store_history: bool = True):
        """Initialize the generator manager.

        Args:
            executor: the executor.
            store_history: whether to store the historical values generated by the generators.
        """

        self.executor = executor
        self.generator_calls = defaultdict(list)
        self.generator_calls_successful = defaultdict(list)
        self.generator_calls_count = defaultdict(int)

        self._store_history = store_history

    executor: PDSketchExecutor
    """The executor."""

    generator_calls: Dict[str, List[Tuple[Tuple[Any, ...], Tuple[Any, ...]]]]
    """Mappings from generator names to the list of calls made to the generator, including a tuple of the arguments and a tuple of the return values."""

    generator_calls_successful: Dict[str, List[bool]]
    """Mappings from generator names to the list of Boolean values indicating whether the generated values lead to successful solution."""

    @property
    def store_history(self) -> bool:
        """Whether to store the historical values generated by the generators."""
        return self._store_history

    def call(self, g: Generator, max_generator_trials: int, args: Tuple[Any, ...], constraint_list: Optional[List[Constraint]] = None) -> Iterator[Tuple[Tuple[str, int], Any]]:
        """Call a generator.

        Args:
            g: the generator.
            max_generator_trials: the maximum number of trials to generate values.
            args: the arguments of the generator.
            constraint_list: the list of constraints to be satisfied by the generated values. This will be passed to the generator function if the list contains more than one constraint.

        Yields:
            A tuple of (index, generated value). The index is a tuple of (generator_name, value_index).
        """

        generator_name = g.function.name
        generator = wrap_singletime_function_to_iterator(
            self.executor.get_function_implementation(generator_name),
            max_generator_trials
        )
        if constraint_list is not None or not isinstance(constraint_list, list):
            generator = generator(*args, return_type=g.output_type)
        else:
            generator = generator(*args, constraint_list, return_type=g.output_type)

        self.generator_calls_count[generator_name] += 1
        first = True
        for result in generator:
            if self._store_history:
                self.generator_calls[generator_name].append((args, result))
                self.generator_calls_successful[generator_name].append(False)
            if not first:
                self.generator_calls_count[generator_name] += 1
            else:
                first = False
            index = generator_name, len(self.generator_calls[generator_name]) - 1

            if not isinstance(result, tuple) and g.function.ftype.is_singular_return:
                result = (result, )

            if not g.function.ftype.is_singular_return:
                assert len(result) == len(generator.function.return_type)

            yield index, result

    def mark_success(self, assignment_dict: AssignmentDict):
        """Mark the values in an assignment dictionary as successful.

        Args:
            assignment_dict: the assignment dictionary.
        """
        assert self._store_history, 'Cannot mark success if history is not stored.'
        for _, value in assignment_dict.items():
            if value.generator_index is not None:
                name, index = value.generator_index
                self.generator_calls_successful[name][index] = True

    def export_generator_calls(self) -> Dict[str, List[Tuple[Tuple[Any, ...], Tuple[Any, ...], bool]]]:
        """Export the generator calls.

        Returns:
            a dictionary mapping from generator names to the list of calls made to the generator, including a tuple of the arguments
            and a tuple of the return values, and a Boolean value indicating whether the generated values lead to successful solution.
        """

        output_dict = defaultdict(list)
        for name, calls in self.generator_calls.items():
            for index, (args, result) in enumerate(calls):
                output_dict[name].append((args, result, self.generator_calls_successful[name][index]))
        return output_dict

    def export_generator_stats(self, divide_by: float = 1) -> str:
        """Export the generator statistics.

        Returns:
            a string containing the generator statistics.
        """

        rows = list()
        for name, count in self.generator_calls_count.items():
            rows.append((name, count / divide_by))
        rows.append(('Total', sum(count / divide_by for count in self.generator_calls_count.values())))
        return tabulate(rows, headers=['Generator', 'Calls'])


def wrap_singletime_function_to_iterator(function: PythonFunctionRef, max_examples: int) -> Callable[..., Iterator[Any]]:
    """Wrap a function that returns a single value to an iterator function.

    Args:
        function: the function.
        max_examples: the maximum number of examples.

    Returns:
        the iterator function.
    """

    if function.is_iterator:
        @functools.wraps(function)
        def wrapped(*args, **kwargs) -> Iterator[Any]:
            try:
                yield from itertools.islice(function.iter_from(*args, **kwargs), max_examples)
            except Exception as e:
                logger.warning(f'Exception raised when calling generator {function}: {e}')

        return wrapped

    @functools.wraps(function)
    def wrapped(*args, **kwargs) -> Iterator[Any]:
        for _ in range(max_examples):
            rv = function(*args, **kwargs)
            if rv is None:
                break
            yield rv

    return wrapped


concepts/pdsketch/executor.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : generator.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/04/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from typing import Union, Sequence, Tuple

from jacinle.utils.meta import repr_from_str

from concepts.dsl.dsl_types import Variable, TensorValueTypeBase, PyObjValueType
from concepts.dsl.dsl_functions import Function
from concepts.dsl.expression import VariableExpression, ValueOutputExpression

__all__ = ['Generator', 'FancyGenerator']


class Generator(object):
    """A generator is function that generates a set of values from a set of given values. Semantically, it certifies that
    the generated values satisfy a given condition."""

    def __init__(
        self,
        name: str,
        arguments: Sequence[Variable],
        certifies: ValueOutputExpression,
        context: Sequence[Union[VariableExpression, ValueOutputExpression]],
        generates: Sequence[Union[VariableExpression, ValueOutputExpression]],
        function: Function,
        output_vars: Sequence[Variable],
        flatten_certifies: ValueOutputExpression,
        priority: int = 0,
        unsolvable: bool = False
    ):
        self.name = name
        self.arguments = tuple(arguments)
        self.certifies = certifies
        self.context = tuple(context)
        self.generates = tuple(generates)
        self.function = function
        self.output_vars = tuple(output_vars)
        self.output_type = tuple(v.dtype for v in output_vars)
        self.flatten_certifies = flatten_certifies
        self.priority = priority
        self.unsolvable = unsolvable

    name: str
    """The name of the generator."""

    arguments: Tuple[Variable, ...]
    """The arguments of the generator."""

    certifies: ValueOutputExpression
    """The condition that the generated values should satisfy."""

    context: Tuple[Union[VariableExpression, ValueOutputExpression], ...]
    """The context values that the generator depends on."""

    generates: Tuple[Union[VariableExpression, ValueOutputExpression], ...]
    """The values that the generator generates."""

    function: Function
    """The declaration of the underlying function that generates the values."""

    output_vars: Tuple[Variable, ...]
    """The output variables of the function."""

    output_type: Tuple[Union[TensorValueTypeBase, PyObjValueType], ...]
    """The output type of the function."""

    flatten_certifies: ValueOutputExpression
    """The condition that the generated values should satisfy, flattened."""

    priority: int
    """The priority of the generator."""

    unsolvable: bool
    """Whether the generator is unsolvable."""

    @property
    def input_vars(self) -> Tuple[Variable, ...]:
        """The input variables of the function."""
        return self.function.arguments

    def __str__(self):
        arg_string = ', '.join([str(c) for c in self.context])
        gen_string = ', '.join([str(c) for c in self.generates])
        return (
            f'{self.name}({arg_string}) -> {gen_string}' + ' {\n'
            '  ' + str(self.function) + '\n'
            '  parameters: ' + str(self.arguments) + '\n'
            '  certifies:  ' + str(self.flatten_certifies) + '\n'
            '  context:    ' + str(self.context) + '\n'
            '  generates:  ' + str(self.generates) + '\n'
            '}'
        )

    __repr__ = repr_from_str


class FancyGenerator(object):
    def __init__(
        self,
        name: str,
        certifies: ValueOutputExpression,
        function: Function,
        flatten_certifies: ValueOutputExpression,
        priority: int = 0,
        unsolvable: bool = False
    ):
        self.name = name
        self.certifies = certifies
        self.function = function
        self.flatten_certifies = flatten_certifies
        self.priority = priority
        self.unsolvable = unsolvable

    name: str
    """The name of the generator."""

    certifies: ValueOutputExpression
    """The condition that the generated values should satisfy."""

    function: Function
    """The declaration of the underlying function that generates the values."""

    flatten_certifies: ValueOutputExpression
    """The condition that the generated values should satisfy, flattened."""

    priority: int
    """The priority of the generator."""

    unsolvable: bool
    """Whether the generator is unsolvable."""

    def __str__(self):
        return (
            f'{self.name}() -> ' + ' {\n'
            '  ' + str(self.function) + '\n'
            '  certifies: ' + str(self.flatten_certifies) + '\n'
            '}'
        )

    __repr__ = repr_from_str


concepts/pdsketch/generator.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : operator.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/04/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import itertools
from typing import TYPE_CHECKING, Any, Optional, Union, Iterable, Iterator, Sequence, Tuple, List, Mapping, Callable

from jacinle.logging import get_logger
from jacinle.utils.enum import JacEnum
from jacinle.utils.printing import indent_text
from jacinle.utils.meta import repr_from_str

from concepts.dsl.executors.tensor_value_executor import compose_bvdict_args
from concepts.dsl.dsl_types import ObjectType, NamedTensorValueType, PyObjValueType, Variable, UnnamedPlaceholder
from concepts.dsl.expression import (
    AndExpression, AssignExpression, ConditionalAssignExpression, DeicticAssignExpression,
    ValueOutputExpression, VariableExpression, VariableAssignmentExpression,
    find_free_variables
)
from concepts.dsl.tensor_value import TensorValue
from concepts.pdsketch.predicate import is_simple_bool, get_simple_bool_predicate

if TYPE_CHECKING:
    from concepts.pdsketch.executor import PDSketchExecutor
    from concepts.pdsketch.domain import State

logger = get_logger(__file__)

__all__ = [
    'Precondition', 'Effect', 'Operator', 'MacroOperator', 'OperatorApplier', 'OperatorApplicationExpression',
    'SubgoalSerializability', 'SubgoalContSerializability', 'AchieveExpression',
    'RegressionRule', 'RegressionRuleApplier',
    'gen_all_grounded_actions', 'filter_static_grounding', 'gen_all_partially_grounded_actions',
    'gen_all_grounded_regression_rules'
]


class Precondition(object):
    """The precondition of an operator. It is basically a wrapper around :class:`~concepts.dsl.expression.ValueOutputExpression`."""

    def __init__(self, bool_expr: ValueOutputExpression):
        self.bool_expr = bool_expr
        self.ao_discretization = None

    bool_expr: ValueOutputExpression
    """The underlying Boolean expression."""

    ao_discretization: Optional[Any]
    """The And-Or discretization of the precondition."""

    def __str__(self) -> str:
        return str(self.bool_expr)

    __repr__ = repr_from_str


class Effect(object):
    """The effect of an operator. It is basically a wrapper around :class:`~concepts.dsl.expression.VariableAssignmentExpression`."""

    def __init__(self, assign_expr: VariableAssignmentExpression):
        self.assign_expr = assign_expr
        self.update_from_simulation = False
        self.ao_discretization = None

    assign_expr: VariableAssignmentExpression
    """The underlying assign expression."""

    update_from_simulation: bool
    """Whether the effect should be updated from simulation, instead of the evaluation of the expression."""

    ao_discretization: Optional[Any]
    """The And-Or discretization of the effect."""

    def set_update_from_simulation(self, update_from_simulation: bool = True):
        self.update_from_simulation = update_from_simulation

    @property
    def unwrapped_assign_expr(self) -> Union[AssignExpression, ConditionalAssignExpression]:
        """Unwrap the DeicticAssignExpression and return the innermost AssignExpression."""
        expr = self.assign_expr
        if isinstance(expr, DeicticAssignExpression):
            expr = expr.expression
        assert isinstance(expr, (AssignExpression, ConditionalAssignExpression))
        return expr

    def __str__(self) -> str:
        return str(self.assign_expr)

    __repr__ = repr_from_str


class Controller(object):
    def __init__(self, name: str, arguments: Sequence[ValueOutputExpression]):
        self.name = name
        self.arguments = tuple(arguments)

    name: str
    """The identifier of the controller."""

    arguments: Tuple[ValueOutputExpression, ...]
    """The argument expressions to the controller function."""

    def __str__(self) -> str:
        return '{}({})'.format(self.name, ', '.join(map(str, self.arguments)))

    __repr__ = repr_from_str

    def pddl_str(self) -> str:
        return f'({self.name} {" ".join(map(str, self.arguments))})'


class OperatorBase(object):
    @property
    def is_primitive(self) -> bool:
        """Whether this operator is a primitive operator (instead of a macro operator or a sub operator in a macro)."""
        raise True

    @property
    def is_macro(self):
        """Whether the operator is a macro operator."""
        return False

    @property
    def is_axiom(self) -> bool:
        """Whether the operator is an axiom."""
        return False


class Operator(OperatorBase):
    """The operator definition in a planning domain."""

    def __init__(
        self,
        name: str,
        arguments: Sequence[Variable],
        preconditions: Sequence[Precondition],
        effects: Sequence[Effect],
        controller: Optional[Controller] = None,
        is_primitive: bool = True,
        is_axiom: bool = False,
        is_template: bool = False,
        extends: Optional[str] = None,
    ):
        self.name = name
        self.arguments = tuple(arguments)
        self.preconditions = tuple(preconditions)
        self.effects = tuple(effects)
        self.controller = controller
        self.is_template = is_template
        self.extends = extends

        self._is_primitive = is_primitive
        self._is_axiom = is_axiom

    name: str
    """The name of the operator."""

    arguments: Tuple[Variable]
    """The list of arguments of the operator."""

    preconditions: Tuple[Precondition]
    """The list of preconditions of the operator."""

    effects: Tuple[Effect]
    """The list of effects of the operator."""

    controller: Optional[Controller]
    """The controller function of the operator."""

    is_axiom: bool
    """Whether this operator is an axiom."""
    is_template: bool
    """Whether this operator is a template."""

    extends: Optional[str]
    """The name of the operator that this operator extends."""

    @property
    def is_primitive(self) -> bool:
        return self._is_primitive

    @property
    def is_axiom(self) -> bool:
        return self._is_axiom

    @property
    def nr_arguments(self) -> int:
        """The number of arguments of the operator."""
        return len(self.arguments)

    def rename(self, new_name: str, is_primitive: Optional[bool] = None) -> 'Operator':
        """Rename the operator."""

        if is_primitive is None:
            is_primitive = self.is_primitive

        return Operator(
            new_name,
            self.arguments,
            self.preconditions,
            self.effects,
            controller=self.controller,
            is_primitive=is_primitive,
            is_axiom=self.is_axiom,
            is_template=self.is_template,
            extends=self.extends,
        )

    def __call__(self, *args) -> 'OperatorApplier':
        """Ground the operator with a list of arguments."""
        output_args = list()
        if args[-1] is Ellipsis:
            args = args[:-1] + ('??', ) * (self.nr_arguments - len(args) + 1)
        for i, arg in enumerate(args):
            if isinstance(arg, str) and arg == '??':
                output_args.append(UnnamedPlaceholder(self.arguments[i].dtype))
            else:
                output_args.append(arg)
        return OperatorApplier(self, tuple(output_args))

    def __str__(self) -> str:
        arg_string = ', '.join([str(arg) for arg in self.arguments])
        return f'{self.name}({arg_string})'

    __repr__ = repr_from_str

    def pddl_str(self) -> str:
        """Return the PDDL representation of the operator."""

        if not self.is_axiom:
            def_name, def_name_a, def_name_p, def_name_e = f'action {self.name}', 'parameters', 'precondition', 'effect'
        else:
            def_name, def_name_a, def_name_p, def_name_e = 'axiom', 'vars', 'context', 'implies'
        arg_string = ' '.join([str(arg) for arg in self.arguments])
        pre_string = '\n'.join([indent_text(str(pre), 1, tabsize=3) for pre in self.preconditions])
        eff_string = '\n'.join([indent_text(str(eff), 1, tabsize=3) for eff in self.effects])
        controller_string = ''
        if self.controller is not None:
            controller_string = f'\n  :controller {self.controller.pddl_str()}'
        return f'''(:{def_name}
 :{def_name_a} ({arg_string})
 :{def_name_p} (and
   {pre_string.lstrip()}
 )
 :{def_name_e} (and
   {eff_string.lstrip()}
 ){controller_string}
)'''


class MacroOperator(OperatorBase):
    def __init__(
        self,
        name: str,
        arguments: Sequence[Variable],
        sub_operators: Sequence['OperatorApplier']
    ):
        self.name = name
        self.arguments = tuple(arguments)
        self.sub_operators = tuple(sub_operators)
        self._check_sub_operator_arguments()

    name: str
    """The name of the macro operator."""

    arguments: Tuple[Variable]
    """The list of arguments of the macro operator."""

    sub_operators: Tuple['OperatorApplier']
    """The list of sub operators of the macro operator."""

    def _check_sub_operator_arguments(self):
        for sub_operator in self.sub_operators:
            for argument in sub_operator.arguments:
                if not isinstance(argument, Variable):
                    raise ValueError('The arguments of a sub operator must be variables.')

    @property
    def nr_arguments(self) -> int:
        """The number of arguments of the operator."""
        return len(self.arguments)

    @property
    def is_macro(self) -> bool:
        """Whether this operator is a macro operator."""
        return True

    @property
    def is_primitive(self) -> bool:
        return False

    def __call__(self, *args) -> 'OperatorApplier':
        """Ground the operator with a list of arguments."""
        output_args = list()
        if args[-1] is Ellipsis:
            args = args[:-1] + ('??', ) * (self.nr_arguments - len(args) + 1)
        for i, arg in enumerate(args):
            if isinstance(arg, str) and arg == '??':
                output_args.append(UnnamedPlaceholder(self.arguments[i].dtype))
            else:
                output_args.append(arg)
        return OperatorApplier(self, tuple(output_args))

    def __str__(self) -> str:
        def_name = 'macro'
        arg_string = ', '.join([str(arg) for arg in self.arguments])
        return f'{def_name}::{self.name}({arg_string})'

    __repr__ = repr_from_str

    def pddl_str(self) -> str:
        """Return the PDDL representation of the operator."""

        arg_string = ' '.join([str(arg) for arg in self.arguments])
        action_string = '\n'.join([indent_text(str(action), 1, tabsize=3) for action in self.sub_operators])
        return f'''(:macro
 ({self.name} {arg_string})
 (and
   {action_string.lstrip()}
 )
)'''


class OperatorApplier(object):
    """An operator applier is essentially a grounded operator, composed of an operator and its arguments."""

    def __init__(self, operator: Union[Operator, MacroOperator], arguments: Sequence[Union[str, Variable, UnnamedPlaceholder, TensorValue]]):
        """Initialize an operator applier."""
        self.operator = operator
        self.arguments = tuple(arguments)

        if len(self.arguments) != len(self.operator.arguments):
            raise ValueError(f'The number of arguments does not match the operator: {self.operator}, arguments: {self.arguments}.')

    operator: Union[Operator, MacroOperator]
    """The operator."""

    arguments: Tuple[Union[str, Variable, UnnamedPlaceholder, TensorValue], ...]
    """The arguments of the grounded operator."""

    @property
    def is_macro(self) -> bool:
        """Whether this operator is a macro operator."""
        return isinstance(self.operator, MacroOperator)

    def iter_sub_operator_appliers(self) -> Iterator['OperatorApplier']:
        if not self.is_macro:
            yield self
            return

        operator: MacroOperator = self.operator
        for sub_operator in operator.sub_operators:
            new_sub_operator = sub_operator.replace_arguments({arg_def.name: arg for arg_def, arg in zip(operator.arguments, self.arguments)})
            if new_sub_operator.is_macro:
                yield from new_sub_operator.iter_sub_operator_appliers()
            else:
                yield new_sub_operator

    def replace_arguments(self, argument_map: Mapping[str, Union[Variable, UnnamedPlaceholder, TensorValue]]) -> 'OperatorApplier':
        """Replace the arguments of the operator applier with a map from argument names to new arguments."""
        new_arguments = list()
        for arg in self.arguments:
            if isinstance(arg, Variable) and arg.name in argument_map:
                new_arguments.append(argument_map[arg.name])
            else:
                new_arguments.append(arg)
        return OperatorApplier(self.operator, new_arguments)

    @property
    def name(self) -> str:
        """The name of the operator."""
        return self.operator.name

    def __str__(self) -> str:
        if not self.operator.is_axiom:
            def_name = 'action'
        else:
            def_name = 'axiom'

        def arg_str(arg_def, arg):
            if isinstance(arg, TensorValue):
                return f'{arg_def.name}={arg.short_str()}'
            elif isinstance(arg, UnnamedPlaceholder):
                return f'{arg_def.name}=??'
            else:
                return f'{arg_def.name}={arg}'

        arg_string = ', '.join([
            arg_str(arg_def, arg)
            for arg_def, arg in zip(self.operator.arguments, self.arguments)
        ])
        return f'{def_name}::{self.operator.name}({arg_string})'

    __repr__ = repr_from_str

    def pddl_str(self) -> str:
        arg_str = ' '.join([str(arg) for arg in self.arguments])
        return f'({self.operator.name} {arg_str})'


class OperatorApplicationExpression(object):
    """An abstract operator grounding. For example :code:`(move ?x ?y)` where :code:`?x` and :code:`?y` are variables in the context."""

    def __init__(self, operator: Operator, arguments: Sequence[Union[VariableExpression, UnnamedPlaceholder, ValueOutputExpression]]):
        self.operator = operator
        self.arguments = tuple(arguments)

    operator: Operator
    """The operator that is applied."""

    arguments: Tuple[Union[VariableExpression, UnnamedPlaceholder, ValueOutputExpression], ...]
    """The arguments of the operator."""

    def ground(self, executor: 'PDSketchExecutor') -> 'OperatorApplier':
        """Ground the operator statement."""
        return OperatorApplier(self.operator, tuple(executor.execute(arg) for arg in self.arguments))

    @property
    def name(self) -> str:
        """The name of the operator."""
        return self.operator.name

    def __str__(self) -> str:
        def_name = 'action'
        arg_string = ', '.join([
            arg_def.name + '=' + str(arg)
            for arg_def, arg in zip(self.operator.arguments, self.arguments)
        ])
        return f'{def_name}::{self.operator.name}({arg_string})'

    __repr__ = repr_from_str

    def pddl_str(self) -> str:
        arg_str = ' '.join([str(arg) for arg in self.arguments])
        return f'({self.operator.name} {arg_str})'


class SubgoalSerializability(JacEnum):
    STRONG = 'strong'
    RULE = 'rule'
    ORDER = 'order'


class SubgoalContSerializability(JacEnum):
    FORALL = 'forall'
    SOME = 'some'


class AchieveExpression(object):
    """An achieve expression. This is used in the definition of regression rules. Each achieve expression has two properties: serializability and cont_serializability.

        - serializability: 'strong', 'rule', or 'order'. Strong serializability means that the achieve expression is serializable for any state that satisfies the preconditions
            of the achieve expression, i.e., all achieve expressions prior to the current achieve expression. Rule serializability means that the achieve expression is serializable
            only for some refinement of the achieve expressions prior to the current achieve expression. However, there exists some refinement for previous expressions such that
            this new achieve expression can be planned based on the results of previous subgoals. Order serializability means that this subgoal is only serializable with respect
            to the order of the subgoals. Therefore, when plan for the sequence of subgoals, the order of the subgoals are preserved, but the preconditions for the subsequent
            rules might need to be promoted to a previous point in the plan.
        - cont_serializability: 'some' or 'forall'. Some continuous serializability means that the achieve expression is serializable for some continuous values of the action parameters.
            Forall continuous serializability means that the achieve expression is serializable for all continuous values of the action parameters.
    """

    def __init__(self, goal: ValueOutputExpression, maintains: Sequence[ValueOutputExpression], serializability: str = 'strong', cont_serializability: str = 'some'):
        self.goal = goal
        self.maintains = tuple(maintains)
        self.serializability = SubgoalSerializability.from_string(serializability)
        self.cont_serializability = SubgoalContSerializability.from_string(cont_serializability)

    goal: ValueOutputExpression
    """The goal of the achieve expression."""

    maintains: Tuple[ValueOutputExpression]
    """The list of maintain expressions."""

    serializability: SubgoalSerializability
    """The serializability of the achieve expression: 'strong', 'rule', or 'order'."""

    cont_serializability: SubgoalContSerializability
    """The continuous serializability of the achieve expression: 'some' or 'forall'."""

    def __str__(self):
        return f'achieve({self.goal}, serializability={self.serializability}, maintain={{{" ".join([str(m) for m in self.maintains])}}})'

    __repr__ = repr_from_str


class RegressionRule(object):
    def __init__(
        self,
        name: str,
        parameters: Sequence[Variable],
        preconditions: Sequence[Precondition],
        preconstraints: Sequence[Precondition],
        goal: Sequence[Effect],
        goal_expression: ValueOutputExpression,
        body: Sequence[Union[OperatorApplicationExpression, AchieveExpression]],
        always: bool = False
    ):
        """Initialize a regression rule."""

        self.name = name
        self.arguments = tuple(parameters)
        self.preconditions = tuple(preconditions)
        self.preconditions_conjunction = AndExpression(*[p.bool_expr for p in self.preconditions])
        self.preconstraints = tuple(preconstraints)
        self.goal = tuple(goal)
        self.goal_expression = goal_expression
        self.body = tuple(body)
        self.always = always

        self.goal_arguments, self.binding_arguments = self._split_arguments()
        self.max_reorder_prefix_length = self._compute_max_reorder_prefix_length()
        self.max_rule_prefix_length = self._compute_max_rule_prefix_length()

        assert self.max_reorder_prefix_length <= self.max_rule_prefix_length or self.max_rule_prefix_length == 0

    name: str
    """The name of the regression rule."""

    arguments: Tuple[Variable, ...]
    """The arguments to the regression rule."""

    goal_arguments: Tuple[Variable, ...]
    """The arguments that appear in the goal of the regression rule."""

    binding_arguments: Tuple[Variable, ...]
    """The arguments that do not appear in the goal of the regression rule."""

    preconditions: Tuple[Precondition, ...]
    """The preconditions of the regression rule."""

    preconditions_conjunction: AndExpression
    """The conjunction of the preconditions of the regression rule."""

    preconstraints: Tuple[Precondition, ...]
    """The preconstraints of the regression rule (the constraints in the maintain set)."""

    goal: Tuple[Effect, ...]
    """The goal of the regression rule."""

    goal_expression: ValueOutputExpression
    """The goal expression of the regression rule."""

    body: Tuple[Union[OperatorApplicationExpression, AchieveExpression], ...]
    """The body of the regression rule, including operator applications and achieve expressions."""

    always: bool
    """Whether the regression rule is always applicable."""

    max_reorder_prefix_length: int
    """The maximum length of the prefix of the body that might need to be reordered."""

    max_rule_prefix_length: int
    """The maximum length of the prefix of the body that need tracking of all possible refinements."""

    @property
    def nr_arguments(self) -> int:
        """The number of arguments of the operator."""
        return len(self.arguments)

    @property
    def serializability(self):
        return 'always' if self.always else 'sometimes'

    def _split_arguments(self) -> Tuple[Tuple[Variable, ...], Tuple[Variable, ...]]:
        goal_variables = find_free_variables(self.goal_expression)
        goal_variable_names = set(v.name for v in goal_variables)
        binding_variables = [arg for arg in self.arguments if arg.name not in goal_variable_names]
        return tuple(goal_variables), tuple(binding_variables)

    def _compute_max_reorder_prefix_length(self) -> int:
        max_reorder_prefix_length = len(self.body)
        while max_reorder_prefix_length > 0:
            item = self.body[max_reorder_prefix_length - 1]
            if isinstance(item, AchieveExpression) and item.serializability is SubgoalSerializability.ORDER:
                break
            max_reorder_prefix_length -= 1
        return max_reorder_prefix_length

    def _compute_max_rule_prefix_length(self) -> int:
        max_rule_prefix_length = len(self.body)
        while max_rule_prefix_length > 0:
            item = self.body[max_rule_prefix_length - 1]
            if isinstance(item, AchieveExpression) and item.serializability is SubgoalSerializability.RULE:
                break
            max_rule_prefix_length -= 1
        return max_rule_prefix_length

    def __str__(self) -> str:
        return f'{self.name}({", ".join(str(param) for param in self.arguments)}, always={self.always})'

    __repr__ = repr_from_str

    def pddl_str(self) -> str:
        body_str = '\n'.join([indent_text(str(action), 1, tabsize=3) for action in self.body])
        return f'''(:regression {self.name} [always={self.always}]
 :parameters ({' '.join(str(param) for param in self.arguments)})
 :preconditions (and {' '.join(str(precondition) for precondition in self.preconditions)})
 :preconstraints (and {' '.join(str(preconstraint) for preconstraint in self.preconstraints)})
 :goal (and {' '.join(str(effect) for effect in self.goal)})
 :rule (then
   {body_str.lstrip()}
 )
)'''

    def __call__(self, *args) -> 'RegressionRuleApplier':
        """Ground the operator with a list of arguments."""
        output_args = list()
        if args[-1] is Ellipsis:
            args = args[:-1] + ('??', ) * (self.nr_arguments - len(args) + 1)
        for i, arg in enumerate(args):
            if isinstance(arg, str) and arg == '??':
                output_args.append(UnnamedPlaceholder(self.arguments[i].dtype))
            else:
                output_args.append(arg)
        return RegressionRuleApplier(self, tuple(output_args))


class RegressionRuleApplier(object):
    def __init__(self, regression_rule: RegressionRule, arguments: Sequence[Union[Variable, TensorValue, Any]]):
        self.regression_rule = regression_rule
        self.arguments = tuple(arguments)

    regression_rule: RegressionRule
    """The regression rule that is applied."""

    arguments: Tuple[Union[Variable, TensorValue, Any], ...]
    """The arguments of the regression rule."""

    def __str__(self) -> str:
        return f'{self.regression_rule.name}({", ".join(str(arg) for arg in self.arguments)})'

    __repr__ = repr_from_str

    def pddl_str(self) -> str:
        arg_str = ' '.join([str(arg) for arg in self.arguments])
        return f'({self.regression_rule.name} {arg_str})'


def gen_all_grounded_actions(
    executor: 'PDSketchExecutor', state: 'State', continuous_values: Optional[Mapping[str, Iterable[TensorValue]]] = None,
    action_names: Optional[Sequence[str]] = None, action_filter: Optional[Callable[[OperatorApplier], bool]] = None, filter_static: bool = True,
    allow_macro_operator: bool = False
) -> List[OperatorApplier]:
    """Generate all grounded actions in a state. Note that this function does not check if the action is applicable at the current state.

    Args:
        executor: a :class:`~concepts.pdsketch.executor.PDSketchExecutor` object.
        state: the current state.
        continuous_values: a dictionary mapping the typename of continuous types to a list of possible values.
        action_names: a list of action names to generate. If None, all actions will be generated.
        action_filter: a function that takes an :class:`OperatorApplier` object and returns a boolean value indicating whether the action should be included in the result.
        filter_static: whether to use the :func:`filter_static_actions` function to filter out static actions. The function will check all static predicates in the
            domain and remove actions that will never been applicable.
        allow_macro_operator: whether to allow macro operators.
    """
    if action_names is not None:
        action_ops = [executor.domain.operators[x] for x in action_names]
    else:
        action_ops = list(executor.domain.operators.values())

    if not allow_macro_operator:
        action_ops = [x for x in action_ops if not x.is_macro]
    else:
        raise NotImplementedError('Macro operators are not supported yet.')

    actions = list()
    for op in action_ops:
        argument_candidates = list()
        for arg in op.arguments:
            if isinstance(arg.dtype, ObjectType):
                argument_candidates.append(state.object_type2name[arg.dtype.typename])
            else:
                assert isinstance(arg.dtype, NamedTensorValueType)
                argument_candidates.append(continuous_values[arg.dtype.typename])
        for comb in itertools.product(*argument_candidates):
            actions.append(op(*comb))
    if filter_static:
        actions = filter_static_grounding(executor, state, actions)
    if action_filter is not None:
        actions = list(filter(action_filter, actions))
    return actions


def filter_static_grounding(executor: 'PDSketchExecutor', state: 'State', actions: Sequence[Union[OperatorApplier, RegressionRuleApplier]]) -> List[Union[OperatorApplier, RegressionRuleApplier]]:
    """Filter out grounded actions or regression rules that do not satisfy static preconditions.

    Args:
        executor: a :class:`~concepts.pdsketch.executor.PDSketchExecutor` object.
        state: the current state.
        actions: a list of grounded actions to be filtered.

    Returns:
        a list of grounded actions that satisfy static preconditions.
    """

    output_actions = list()
    for action in actions:
        if isinstance(action, OperatorApplier):
            bounded_variables = compose_bvdict_args(action.operator.arguments, action.arguments, state=state)
            preconditions = action.operator.preconditions
        elif isinstance(action, RegressionRuleApplier):
            bounded_variables = compose_bvdict_args(action.regression_rule.arguments, action.arguments, state=state)
            preconditions = action.regression_rule.preconditions
        else:
            raise TypeError(f'filter_static_grounding only accepts OperatorApplier or RegressionRuleApplier sequences； got {type(action)} in the input list.')
        flag = True
        for pre in preconditions:
            if is_simple_bool(pre.bool_expr) and get_simple_bool_predicate(pre.bool_expr).is_static:
                rv = executor.execute(pre.bool_expr, state=state, bounded_variables=bounded_variables).item()
                if rv < 0.5:
                    flag = False
                    break
        if flag:
            output_actions.append(action)
    return output_actions


def gen_all_partially_grounded_actions(
    executor: 'PDSketchExecutor', state: 'State',
    action_names: Optional[Sequence[str]] = None, action_filter: Optional[Callable[[OperatorApplier], bool]] = None, filter_static: bool = True,
    allow_macro_operator: bool = False
) -> List[OperatorApplier]:
    """Generate all partially grounded actions in a state. Partially grounded actions are actions with only object-typed arguments grounded.
    For Value-typed arguments, the resulting :class:`~concepts.pdsketch.operators.OperatorApplier` object will have arguments that are
    placeholders (i.e., :class:`~concepts.dsl.dsl_types.UnnamedPlaceholder` objects).
    Note that, as :func:`generate_all_grounded_actions`, this function does not check if the action is applicable at the current state.

    Args:
        executor: a :class:`~concepts.pdsketch.executor.PDSketchExecutor` object.
        state: the current state.
        action_names: a list of action names to generate. If None, all actions will be generated.
        action_filter: a function that takes an :class:`OperatorApplier` object and returns a boolean value indicating whether the action should be included in the result.
        filter_static: whether to use the :func:`filter_static_actions` function to filter out static actions. The function will check all static predicates in the
            domain and remove actions that will never been applicable.
        allow_macro_operator: whether to allow macro operators.

    Returns:
        a list of partially grounded actions.
    """

    if action_names is not None:
        action_ops = [executor.domain.operators[x] for x in action_names]
    else:
        action_ops = list(executor.domain.operators.values())

    if not allow_macro_operator:
        action_ops = [x for x in action_ops if not x.is_macro]
    else:
        raise NotImplementedError('Macro operators are not supported yet.')

    actions = list()
    for op in action_ops:
        argument_candidates = list()
        for arg in op.arguments:
            if isinstance(arg.dtype, ObjectType):
                argument_candidates.append(state.object_type2name[arg.dtype.typename])
            else:
                assert isinstance(arg.dtype, (NamedTensorValueType, PyObjValueType))
                argument_candidates.append([UnnamedPlaceholder(arg.dtype)])
        for comb in itertools.product(*argument_candidates):
            actions.append(op(*comb))

    if filter_static:
        actions = filter_static_grounding(executor, state, actions)
    if action_filter is not None:
        actions = list(filter(action_filter, actions))
    return actions


def gen_all_grounded_regression_rules(
    executor: 'PDSketchExecutor', state: 'State', continuous_values: Optional[Mapping[str, Iterable[TensorValue]]] = None,
    regression_rule_names: Optional[Sequence[str]] = None, regression_rule_filter: Optional[Callable[[RegressionRuleApplier], bool]] = None, filter_static: bool = True,
) -> List[RegressionRuleApplier]:
    """Generate all grounded regression rules applicable in an environment, given the initial state.
    Note that this function does not check if the action is applicable at the current state.

    Args:
        executor: a :class:`~concepts.pdsketch.executor.PDSketchExecutor` object.
        state: the current state.
        continuous_values: a dictionary mapping the typename of continuous types to a list of possible values.
        regression_rule_names: a list of regression rule names to generate. If None, all regression rules will be generated.
        regression_rule_filter: a function that takes an :class:`RegressionRuleApplier` object and returns a boolean value indicating whether the action should be included in the result.
        filter_static: whether to use the :func:`filter_static_actions` function to filter out static regression rules. The function will check all static predicates in the precondition list.
    """
    if regression_rule_names is not None:
        lifted_regression_rules = [executor.domain.regression_rules[x] for x in regression_rule_names]
    else:
        lifted_regression_rules = list(executor.domain.regression_rules.values())

    regression_rules = list()
    for op in lifted_regression_rules:
        argument_candidates = list()
        for arg in op.arguments:
            if isinstance(arg.dtype, ObjectType):
                argument_candidates.append(state.object_type2name[arg.dtype.typename])
            else:
                assert isinstance(arg.dtype, NamedTensorValueType)
                argument_candidates.append(continuous_values[arg.dtype.typename])
        for comb in itertools.product(*argument_candidates):
            regression_rules.append(op(*comb))
    if filter_static:
        regression_rules = filter_static_grounding(executor, state, regression_rules)
    if regression_rule_filter is not None:
        regression_rules = list(filter(regression_rule_filter, regression_rules))
    return regression_rules


concepts/pdsketch/operator.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/31/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/pdsketch/parsers/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : pdsketch_parser.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/30/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import os.path as osp
import itertools
import collections
from typing import Optional, Union, Sequence, Tuple, Set, List
from lark import Lark, Tree, Transformer, v_args

import jacinle

import concepts.dsl.expression as E
from concepts.dsl.dsl_types import ObjectType, TensorValueTypeBase, VectorValueType, AutoType, BOOL, Variable, ObjectConstant, UnnamedPlaceholder
from concepts.dsl.expression import ExpressionDefinitionContext, get_expression_definition_context, FunctionApplicationExpression
from concepts.dsl.tensor_state import TensorState
from concepts.pdsketch.predicate import Predicate
from concepts.pdsketch.operator import Precondition, Effect, Controller, Operator, OperatorApplicationExpression, AchieveExpression
from concepts.pdsketch.generator import Generator
from concepts.pdsketch.domain import Domain, Problem
from concepts.pdsketch.executor import PDSketchExecutor

__all__ = [
    'PDSketchParser', 'PDSketchTransformer',
    'load_domain_file', 'load_domain_string', 'parse_expression', 'load_problem_file', 'load_problem_string', 'parse_expression', 'load_domain_string_incremental'
]

logger = jacinle.get_logger(__file__)

# lark.v_args
inline_args = v_args(inline=True)
DEBUG_LOG_COMPOSE = False


def _log_function(func):
    if DEBUG_LOG_COMPOSE:
        return jacinle.log_function(func)
    return func


class PDSketchParser(object):
    """Parser for PDSketch domain and problem files. Users should not use this class directly.
    Instead, use the following functions:

    - :func:`load_domain_file`
    - :func:`load_domain_string`
    - :func:`load_problem_file`
    - :func:`load_problem_string`
    - :func:`parse_expression`
    """

    grammar_file = osp.join(osp.dirname(__file__), 'pdsketch-v2.grammar')
    """The grammar definition for PDSketch."""

    def __init__(self):
        with open(type(self).grammar_file) as f:
            self.lark = Lark(f)

    def load(self, file) -> Tree:
        """Load a domain or problem file and return the corresponding tree."""
        with open(file) as f:
            return self.lark.parse(f.read())

    def loads(self, string) -> Tree:
        """Load a domain or problem string and return the corresponding tree."""
        return self.lark.parse(string)

    def make_domain(self, tree: Tree) -> Domain:
        """Construct a PDSketch domain from a tree."""
        assert tree.children[0].data == 'definition'
        transformer = PDSketchTransformer(Domain())
        transformer.transform(tree)
        domain = transformer.domain
        domain.post_init()
        return domain

    def make_problem(self, tree: Tree, domain: Domain, ignore_unknown_predicates: bool = False) -> Problem:
        """Construct a PDSketch problem from a tree."""
        assert tree.children[0].data == 'definition'
        transformer = PDSketchTransformer(domain, ignore_unknown_predicates=ignore_unknown_predicates)
        transformer.transform(tree)
        problem = transformer.problem
        return problem

    def make_expression(self, domain: Domain, tree: Tree, variables: Optional[Sequence[Variable]] = None) -> E.Expression:
        """Construct a PDSketch expression from a tree."""
        if variables is None:
            variables = list()
        transformer = PDSketchTransformer(domain, allow_object_constants=True)
        node = transformer.transform(tree).children[0]
        assert isinstance(node, (_FunctionApplicationImm, _QuantifierApplicationImm))
        with ExpressionDefinitionContext(*variables, domain=domain).as_default():
            return node.compose()

    def incremental_define_domain(self, domain: Domain, tree: Tree) -> None:
        """Incrementally define a PDSketch domain from a tree."""
        transformer = PDSketchTransformer(domain)
        transformer.transform(tree)
        domain.post_init()
        return domain


_parser = PDSketchParser()


def load_domain_file(filename: str) -> Domain:
    """Load a domain from a file.

    Args:
        filename: the filename of the domain.

    Returns:
        the domain.
    """
    tree = _parser.load(filename)
    domain = _parser.make_domain(tree)
    return domain


def load_domain_string(domain_string: str) -> Domain:
    """Load a domain from a string.

    Args:
        domain_string: the string of the domain.

    Returns:
        the domain.
    """
    tree = _parser.loads(domain_string)
    domain = _parser.make_domain(tree)
    return domain


def load_problem_file(filename: str, domain: Domain, ignore_unknown_predicates: bool = False, return_tensor_state: bool = True, executor: Optional[PDSketchExecutor] = None) -> Union[Tuple[TensorState, E.ValueOutputExpression], Problem]:
    """Load a problem from a file.

    Args:
        filename: the filename of the problem.
        domain: the domain.
        ignore_unknown_predicates: whether to ignore unknown predicates. If set to True, unknown predicates will be ignored with a warning.
            Otherwise, an error will be raised while parsing the problem.
        return_tensor_state: whether to return the initial state as a tensor state. If set to True, the return value wille be a tuple of (initial state, goal).
            Otherwise, the return value will be the :class:`~concepts.pdsketch.domain.Problem` object.
        executor: the executor used to create initial states.

    Returns:
        the problem, as a tuple of (initial state, goal).
    """
    tree = _parser.load(filename)
    with ExpressionDefinitionContext(domain=domain).as_default():
        problem = _parser.make_problem(tree, domain, ignore_unknown_predicates=ignore_unknown_predicates)

    if return_tensor_state:
        assert executor is not None
        return problem.to_state(executor), problem.goal
    return problem


def load_problem_string(problem_string: str, domain: Domain, ignore_unknown_predicates: bool = False, return_tensor_state: bool = True, executor: Optional[PDSketchExecutor] = None) -> Union[Tuple[TensorState, E.ValueOutputExpression], Problem]:
    """Load a problem from a string.

    Args:
        problem_string: the string of the problem.
        domain: the domain.
        ignore_unknown_predicates: whether to ignore unknown predicates. If set to True, unknown predicates will be ignored with a warning.
            Otherwise, an error will be raised while parsing the problem.
        return_tensor_state: whether to return the initial state as a tensor state. If set to True, the return value wille be a tuple of (initial state, goal).
            Otherwise, the return value will be the :class:`~concepts.pdsketch.domain.Problem` object.
        executor: the executor used to create initial states.

    Returns:
        the problem, as a tuple of (initial state, goal) if `return_tensor_state` is True, otherwise the :class:`~concepts.pdsketch.domain.Problem` object.
    """
    tree = _parser.loads(problem_string)
    with ExpressionDefinitionContext(domain=domain).as_default():
        problem = _parser.make_problem(tree, domain, ignore_unknown_predicates=ignore_unknown_predicates)

    if return_tensor_state:
        assert executor is not None
        return problem.to_state(executor), problem.goal
    return problem


def parse_expression(domain: Domain, string: str, variables: Optional[Sequence[Variable]] = None) -> E.Expression:
    """Parse an expression from a string.

    Args:
        domain: the domain.
        string: the string of the expression.
        variables: a list of variables that are used in the expression.

    Returns:
        the expression.
    """
    tree = _parser.loads(string)
    expr = _parser.make_expression(domain, tree, variables)
    return expr


def load_domain_string_incremental(domain: Domain, string: str) -> Domain:
    """Incrementally load a domain from a string.

    Args:
        domain: the domain.
        string: the string of the domain.

    Returns:
        the domain.
    """
    tree = _parser.loads(string)
    _parser.incremental_define_domain(domain, tree)
    return domain


class PDSketchTransformer(Transformer):
    """The tree-to-object transformer for PDSketch domain and problem files. Users should not use this class directly."""

    def __init__(self, init_domain: Domain = None, allow_object_constants: bool = True, ignore_unknown_predicates: bool = False):
        super().__init__()

        self.domain = init_domain
        self.problem = Problem()
        self.allow_object_constants = allow_object_constants
        self.ignore_unknown_predicates = ignore_unknown_predicates
        self.ignored_predicates: Set[str] = set()

    @inline_args
    def definition_decl(self, definition_type, definition_name):
        if definition_type.value == 'domain':
            self.domain.name = definition_name.value

    def type_definition(self, args):
        """Parse a type definition.

        Very ugly hack to handle multi-line definition in PDDL. In PDDL, type definition can be separated by newline.
        This kinds of breaks the parsing strategy that ignores all whitespaces. More specifically, consider the following two definitions:

        .. code-block:: pddl

            (:types
              a
              b - a
            )

        and

        .. code-block:: pddl

            (:types
              a b - a
            )
        """
        if isinstance(args[-1], Tree) and args[-1].data == "parent_type_name":
            parent_line, parent_name = args[-1].children[0]
            args = args[:-1]
        else:
            parent_line, parent_name = -1, 'object'

        for lineno, typedef in args:
            assert typedef is not AutoType, 'AutoType is not allowed in type definition.'

        for arg in args:
            arg_line, arg_name = arg
            if arg_line == parent_line:
                self.domain.define_type(arg_name, parent_name)
            else:
                self.domain.define_type(arg_name, parent_name)

    @inline_args
    def constants_definition(self, *args):
        for arg in args:
            self.domain.constants[arg.name] = arg

    @inline_args
    def predicate_definition(self, name, *args):
        name, kwargs = name

        return_type = kwargs.pop('return_type', BOOL)
        self._predicate_definition_inner(name, args, return_type, kwargs)

    @inline_args
    def predicate_definition2(self, name, *args):
        name, kwargs = name
        assert 'return_type' not in kwargs
        args, return_type = args[:-1], args[-1]

        return_type = self.domain.get_type(return_type)
        self._predicate_definition_inner(name, args, return_type, kwargs)

    def _predicate_definition_inner(self, name, args, return_type, kwargs):
        generators = kwargs.pop('generators', None)
        predicate_def = self.domain.define_predicate(name, args, return_type, **kwargs)
        self._define_inplace_generators(predicate_def, generators)

    def _define_inplace_generators(self, predicate: Predicate, generators) -> List[Generator]:
        defined_generators = list()
        if generators is not None:
            generators: List[str]
            for target_variable_name in generators:
                assert target_variable_name.startswith('?')
                parameters, certifies, context, generates = _canonize_inline_generator_def_predicate(self.domain, target_variable_name, predicate)
                generator_name = f'gen-{predicate.name}-{target_variable_name[1:]}' if len(generators) > 1 else f'gen-{predicate.name}'
                generator = self.domain.define_generator(generator_name, parameters, certifies, context, generates)
                defined_generators.append(generator)
        return defined_generators

    @inline_args
    def predicate_name(self, name, kwargs=None):
        if kwargs is None:
            kwargs = dict()
        return name.value, kwargs

    @inline_args
    def type_name(self, name):
        if name.value == 'auto':
            return name.line, AutoType
        # propagate the "lineno" of the type definition up.
        return name.line, name.value

    @inline_args
    def object_type_name(self, typedef):
        return typedef

    @inline_args
    def value_type_name(self, typedef):
        lineno, typedef = typedef
        if typedef is AutoType:
            return lineno, AutoType
        if isinstance(typedef, VectorValueType):
            return lineno, typedef
        assert isinstance(typedef, str)
        return lineno, self.domain.get_type(typedef)

    @inline_args
    def vector_type_name(self, dtype, dim, choices, kwargs=None):
        choices = choices.children[0] if len(choices.children) > 0 else 0
        if kwargs is None:
            kwargs = dict()
        lineno, dtype = dtype
        return lineno, VectorValueType(dtype, dim, choices, **kwargs)

    @inline_args
    def object_type_name_unwrapped(self, typedef):
        return typedef[1]

    @inline_args
    def value_type_name_unwrapped(self, typedef):
        return typedef[1]

    @inline_args
    def predicate_group_definition(self, *args):
        raise NotImplementedError()

    @inline_args
    def action_definition(self, name, *defs):
        name, kwargs = name

        parameters = tuple()
        precondition = None
        effect = None
        controller = None

        for def_ in defs:
            if isinstance(def_, _ParameterListWrapper):
                parameters = def_.parameters
            elif isinstance(def_, _PreconditionWrapper):
                precondition = def_.precondition
            elif isinstance(def_, _EffectWrapper):
                effect = def_.effect
            elif isinstance(def_, _ControllerWrapper):
                controller = def_.controller
            else:
                raise TypeError('Unknown definition type: {}.'.format(type(def_)))

        if precondition is not None:
            with ExpressionDefinitionContext(*parameters, domain=self.domain, scope=f"action::{name}").as_default():
                precondition = _canonize_precondition(precondition)
        else:
            precondition = list()

        if effect is not None:
            with ExpressionDefinitionContext(*parameters, domain=self.domain, scope=f"action::{name}", is_effect_definition=True).as_default():
                effect = _canonize_effect(effect)
        else:
            effect = list()

        if controller is not None:
            with ExpressionDefinitionContext(*parameters, domain=self.domain, scope=f"action::{name}").as_default():
                controller = _canonize_controller(controller)

        self.domain.define_operator(name, parameters, precondition, effect, controller, **kwargs)

    @inline_args
    def action_definition2(self, name, extends, *defs):
        name, kwargs = name

        assert 'extends' not in kwargs, 'Instantiation cannot be set using decorators. Use :extends instead.'
        kwargs['extends'] = extends

        template_op = self.domain.operators[extends]
        parameters = template_op.arguments

        precondition = None
        effect = None
        controller = None

        for def_ in defs:
            if isinstance(def_, _ParameterListWrapper):
                parameters += def_.parameters
            elif isinstance(def_, _PreconditionWrapper):
                precondition = def_.precondition
            elif isinstance(def_, _EffectWrapper):
                effect = def_.effect
            else:
                raise TypeError('Unknown definition type: {}.'.format(type(def_)))

        if precondition is not None:
            with ExpressionDefinitionContext(*parameters, domain=self.domain, scope=f"action::{name}").as_default():
                precondition = tuple(_canonize_precondition(precondition))
        else:
            precondition = tuple()
        if effect is not None:
            with ExpressionDefinitionContext(*parameters, domain=self.domain, scope=f"action::{name}", is_effect_definition=True).as_default():
                effect = tuple(_canonize_effect(effect))
        else:
            effect = tuple()
        if controller is not None:
            with ExpressionDefinitionContext(*parameters, domain=self.domain, scope=f"action::{name}").as_default():
                controller = _canonize_controller(controller)
        else:
            controller = template_op.controller

        self.domain.define_operator(name, parameters, template_op.preconditions + precondition, template_op.effects + effect, controller, **kwargs)

    def action_parameters(self, args):
        return _ParameterListWrapper(tuple(args))

    @inline_args
    def action_precondition(self, function_call):
        return _PreconditionWrapper(function_call)

    @inline_args
    def action_effect(self, function_call):
        return _EffectWrapper(function_call)

    @inline_args
    def action_controller(self, fcuntion_call):
        return _ControllerWrapper(fcuntion_call)

    @inline_args
    def action_name(self, name, kwargs=None):
        if kwargs is None:
            kwargs = dict()
        return name.value, kwargs

    @inline_args
    def action_instantiates(self, name):
        return name.value

    @inline_args
    def regression_definition(self, name, *defs):
        name, kwargs = name

        parameters = tuple()
        precondition = None
        preconstraint = None
        goal = None
        body = None

        for def_ in defs:
            if isinstance(def_, _ParameterListWrapper):
                parameters = def_.parameters
            elif isinstance(def_, _PreconditionWrapper):
                precondition = def_.precondition
            elif isinstance(def_, _PreconstraintWrapper):
                preconstraint = def_.precondition
            elif isinstance(def_, _EffectWrapper):
                goal = def_.effect
            else:
                assert body is None, 'Multiple bodies are not allowed.'
                body = def_

        assert goal is not None and body is not None, 'Both goal and body must be defined.'

        if precondition is not None:
            with ExpressionDefinitionContext(*parameters, domain=self.domain, scope=f"regression::{name}").as_default():
                precondition = tuple(_canonize_precondition(precondition))
        else:
            precondition = tuple()

        if preconstraint is not None:
            with ExpressionDefinitionContext(*parameters, domain=self.domain, scope=f"regression::{name}").as_default():
                preconstraint = tuple(_canonize_precondition(preconstraint))
        else:
            preconstraint = tuple()

        with ExpressionDefinitionContext(*parameters, domain=self.domain, scope=f"regression::{name}", is_effect_definition=False).as_default():
            goal_expression = goal.compose()

        if isinstance(goal, _FunctionApplicationImm):  # Simple goal
            with ExpressionDefinitionContext(*parameters, domain=self.domain, scope=f"regression::{name}", is_effect_definition=True).as_default():
                goal = tuple(_canonize_effect(goal))
        else:
            goal = tuple()

        with ExpressionDefinitionContext(*parameters, domain=self.domain, scope=f"regression::{name}").as_default():
            last_achieve = tuple()
            body_statements = list()
            for i, statement in enumerate(body):
                statement: _FunctionApplicationImm
                if statement.name == 'achieve':
                    if len(statement.arguments) == 1:
                        new_subgoal = statement.arguments[0].compose(BOOL)
                        this_subgoal = (new_subgoal, last_achieve)
                    elif len(statement.arguments) == 2:
                        new_subgoal = statement.arguments[0].compose(BOOL)
                        this_subgoal = (new_subgoal, statement.arguments[1].compose(BOOL))
                    else:
                        raise ValueError('Invalid achieve statement.')
                    body_statements.append(AchieveExpression(*this_subgoal, **statement.kwargs))
                    last_achieve = last_achieve + (new_subgoal, )
                else:
                    assert statement.name in self.domain.operators, f'Unknown operator: {statement.name}.'
                    operator = self.domain.operators[statement.name]
                    expression = _canonize_operator_expression_arguments(operator, statement.arguments)
                    body_statements.append(expression)
            body = tuple(body_statements)

        self.domain.define_regression_rule(name, parameters, precondition, preconstraint, goal, goal_expression, body, **kwargs)

    def regression_parameters(self, args):
        return _ParameterListWrapper(tuple(args))

    @inline_args
    def regression_precondition(self, function_call):
        return _PreconditionWrapper(function_call)

    @inline_args
    def regression_preconstraint(self, function_call):
        return _PreconstraintWrapper(function_call)

    @inline_args
    def regression_goal(self, function_call):
        return _EffectWrapper(function_call)

    @inline_args
    def regression_rule(self, *body_statements):
        return body_statements

    @inline_args
    def regression_name(self, name, kwargs=None):
        if kwargs is None:
            kwargs = dict()
        return name.value, kwargs

    @inline_args
    def axiom_definition(self, decorator, vars, context, implies):
        kwargs = dict() if len(decorator.children) == 0 else decorator.children[0]
        vars = tuple(vars.children)
        precondition = context.children[0]
        effect = implies.children[0]

        name = kwargs.pop('name', None)
        scope = None if name is None else f"axiom::{name}"

        with ExpressionDefinitionContext(*vars, domain=self.domain, scope=scope).as_default():
            precondition = _canonize_precondition(precondition)
        with ExpressionDefinitionContext(*vars, domain=self.domain, scope=scope, is_effect_definition=True).as_default():
            effect = _canonize_effect(effect)
        self.domain.define_axiom(name, vars, precondition, effect, **kwargs)

    @inline_args
    def macro_definition(self, signature, body):
        ((name, kwargs), parameters) = signature
        parameters_map = {p.name: p for p in parameters}

        sub_operators = list()
        for i, sub_call in enumerate(body):
            op_name = sub_call.name
            assert op_name in self.domain.operators, f'Unknown operator {op_name} in macro definition.'
            sub_op = self.domain.operators[op_name]
            sub_op = self._make_macro_sub_operator(name, sub_op, i)
            sub_parameters = [
                parameters_map[p.name] if p.name in parameters_map else p
                for p in sub_call.arguments
            ]
            sub_operators.append(sub_op(*sub_parameters))

        self.domain.define_macro(name, parameters, sub_operators, **kwargs)

    def _make_macro_sub_operator(self, macro_name: str, sub_op: Operator, index: int):
        sub_op = sub_op.rename(f'{macro_name}-{index}-{sub_op.name}', is_primitive=False)

        new_preconditions = list()
        for i, precondition in enumerate(sub_op.preconditions):
            expr = precondition.bool_expr
            if isinstance(expr, FunctionApplicationExpression):
                if expr.function.is_generator_placeholder:
                    new_function = expr.function.rename(f'{macro_name}-{index}-{expr.function.name}')
                    self.domain.define_predicate_inner(new_function.name, new_function)
                    for generator in self._define_inplace_generators(new_function, new_function.inplace_generators):
                        self.domain.declare_external_function_crossref(
                            f'generator::{generator.name}',
                            f'generator::{generator.name.replace(new_function.name, expr.function.name)}'
                        )
                    new_preconditions.append(Precondition(FunctionApplicationExpression(
                        new_function,
                        expr.arguments
                    )))
                else:
                    new_preconditions.append(precondition)
            else:
                new_preconditions.append(precondition)

        sub_op.preconditions = tuple(new_preconditions)
        self.domain.define_operator_inner(sub_op.name, sub_op)

        return sub_op

    @inline_args
    def macro_signature(self, name, *args):
        return name, args

    @inline_args
    def macro_name(self, name, kwargs=None):
        if kwargs is None:
            kwargs = dict()
        return name.value, kwargs

    @inline_args
    def macro_content(self, *body_statements):
        return body_statements

    @inline_args
    def derived_definition(self, signature, expr):
        name, args, kwargs = signature
        expr = expr

        return_type = kwargs.pop('return_type', BOOL)
        with ExpressionDefinitionContext(*args, domain=self.domain, scope=f"derived::{name}").as_default():
            if return_type is AutoType:
                expr = expr.compose()
                assert isinstance(expr, (E.VariableExpression, E.ValueOutputExpression))
                return_type = expr.return_type
            else:
                expr = expr.compose(return_type)
        self.domain.define_derived(name, args, return_type, expr=expr, **kwargs)

    @inline_args
    def derived_signature1(self, name, *args):
        name, kwargs = name
        return name, args, kwargs

    @inline_args
    def derived_signature2(self, name, *args):
        name, kwargs = name
        assert 'return_type' not in kwargs, 'Return type cannot be set using decorators.'
        kwargs['return_type'] = args[-1]
        return name, args[:-1], kwargs

    @inline_args
    def derived_name(self, name, kwargs=None):
        if kwargs is None:
            kwargs = dict()
        return name.value, kwargs

    @inline_args
    def generator_definition(self, name, parameters, certifies, context, generates):
        name, kwargs = name
        parameters = tuple(parameters.children)
        certifies = certifies.children[0]
        context = context.children[0]
        generates = generates.children[0]

        ctx = ExpressionDefinitionContext(*parameters, domain=self.domain, scope=f"generator::{name}")
        with ctx.as_default():
            certifies = certifies.compose(BOOL)
            assert context.name == 'and'
            context = [_compose(ctx, c) for c in context.arguments]
            assert generates.name == 'and'
            generates = [_compose(ctx, c) for c in generates.arguments]

        self.domain.define_generator(name, parameters, certifies, context, generates, **kwargs)

    @inline_args
    def fancy_generator_definition(self, name, parameters, certifies):
        name, kwargs = name
        certifies = certifies.children[0]

        ctx = ExpressionDefinitionContext(domain=self.domain, scope=f"generator::{name}")
        with ctx.as_default():
            certifies = certifies.compose(BOOL)

        self.domain.define_fancy_generator(name, certifies, **kwargs)

    @inline_args
    def generator_name(self, name, kwargs=None):
        if kwargs is None:
            kwargs = dict()
        return name.value, kwargs

    @inline_args
    def object_definition(self, constant):
        self.problem.add_object(constant.name, constant.typename)

    @inline_args
    def init_definition_item(self, function_call):
        if function_call.name not in self.domain.functions:
            if self.ignore_unknown_predicates:
                if function_call.name not in self.ignored_predicates:
                    logger.warning(f"Unknown predicate: {function_call.name}.")
                    self.ignored_predicates.add(function_call.name)
            else:
                raise ValueError(f"Unknown predicate: {function_call.name}.")
            return
        self.problem.add_proposition(function_call.compose())

    @inline_args
    def goal_definition(self, function_call):
        self.problem.set_goal(function_call.compose())

    @inline_args
    def variable(self, name) -> Variable:
        return Variable(name.value, AutoType)

    @inline_args
    def typedvariable(self, name, typename):
        # name is of type `Variable`.
        if typename is AutoType:
            return Variable(name.name, AutoType)
        return Variable(name.name, self.domain.get_type(typename))

    @inline_args
    def quantifiedvariable(self, quantifier, variable):
        variable.set_quantifier_flag(quantifier.value)
        return variable

    @inline_args
    def constant(self, name) -> ObjectConstant:
        assert self.allow_object_constants
        return ObjectConstant(name.value, AutoType)

    @inline_args
    def typedconstant(self, name, typename):
        return ObjectConstant(name.name, self.domain.get_type(typename))

    @inline_args
    def bool(self, v):
        return v.value == 'true'

    @inline_args
    def int(self, v):
        return int(v.value)

    @inline_args
    def float(self, v):
        return float(v.value)

    @inline_args
    def string(self, v):
        return v.value[1:-1]

    @inline_args
    def list(self, *args):
        return list(args)

    @inline_args
    def decorator_k(self, k):
        return k.value

    @inline_args
    def decorator_v(self, v):
        return v

    @inline_args
    def decorator_kwarg(self, k, v=True):
        return (k, v)

    def decorator_kwargs(self, args):
        return {k: v for k, v in args}

    @inline_args
    def slot(self, _, name, kwargs=None):
        return _Slot(name.children[0].value, kwargs)

    @inline_args
    def function_name(self, name, kwargs=None):
        return name, kwargs

    @inline_args
    def method_name(self, predicate_name, _, method_name):
        return _MethodName(predicate_name, method_name), None

    @inline_args
    def function_call(self, name, *args):
        name, kwargs = name
        if isinstance(name, (_MethodName, _Slot)):
            return _FunctionApplicationImm(name, args, kwargs=kwargs)
        else:
            return _FunctionApplicationImm(name.value, args, kwargs=kwargs)

    @inline_args
    def simple_function_call(self, name, *args):
        name, kwargs = name
        return _FunctionApplicationImm(name.value, args, kwargs=kwargs)

    @inline_args
    def pm_function_call(self, pm_sign, function_call):
        if pm_sign.value == '+':
            return function_call
        else:
            return _FunctionApplicationImm('not', [function_call])

    @inline_args
    def quantified_function_call(self, quantifier, variable, expr):
        return _QuantifierApplicationImm(quantifier, variable, expr)

    @inline_args
    def ellipsis(self):
        return Ellipsis


class _ParameterListWrapper(collections.namedtuple('_ParameterListWrapper', 'parameters')):
    pass


class _PreconditionWrapper(collections.namedtuple('_PreconditionWrapper', 'precondition')):
    pass


class _PreconstraintWrapper(collections.namedtuple('_PreconditionWrapper', 'precondition')):
    pass


class _EffectWrapper(collections.namedtuple('_EffectWrapper', 'effect')):
    pass


class _ControllerWrapper(collections.namedtuple('_ControllerWrapper', 'controller')):
    pass


class _FunctionApplicationImm(object):
    def __init__(self, name, arguments, kwargs=None):
        self.name = name
        self.arguments = arguments
        self.kwargs = kwargs

        if self.kwargs is None:
            self.kwargs = dict()

    def __str__(self):
        arguments_str = ', '.join([str(arg) for arg in self.arguments])
        return f'IMM::{self.name}({arguments_str})'

    __repr__ = jacinle.repr_from_str

    @_log_function
    def compose(self, expect_value_type: Optional[Union[ObjectType, TensorValueTypeBase]] = None):
        ctx = get_expression_definition_context()
        domain: Domain = ctx.domain

        if isinstance(self.name, _Slot):
            assert ctx.scope is not None, 'Cannot define slots inside anonymous actino/axioms.'

            name = ctx.scope + '::' + self.name.name
            arguments = self._compose_arguments(ctx, self.arguments)
            argument_types = [arg.return_type for arg in arguments]
            return_type = self.name.kwargs.pop('return_type', None)
            if return_type is None:
                assert expect_value_type is not None, f'Cannot infer return type for function {name}; please specify by [return_type=Type]'
                return_type = expect_value_type
            else:
                if expect_value_type is not None:
                    assert return_type == expect_value_type, f'Return type mismatch for function {name}: expect {expect_value_type}, got {return_type}.'
            function_def = domain.declare_external_function(name, argument_types, return_type, kwargs=self.name.kwargs)
            return E.FunctionApplicationExpression(function_def, arguments)
        elif isinstance(self.name, _MethodName):
            assert self.name.predicate_name in domain.functions, 'Unkwown predicate: {}.'.format(self.name.predicate_name)
            predicate_def = domain.functions[self.name.predicate_name]

            if self.name.method_name == 'equal':
                nr_index_arguments = len(self.arguments) - 1
            elif self.name.method_name == 'assign':
                nr_index_arguments = len(self.arguments) - 1
            elif self.name.method_name == 'cond-select':
                nr_index_arguments = len(self.arguments) - 1
            elif self.name.method_name == 'cond-assign':
                nr_index_arguments = len(self.arguments) - 2
            else:
                raise NameError('Unknown method name: {}.'.format(self.name.method_name))

            arguments = self._compose_arguments(ctx, self.arguments[:nr_index_arguments], predicate_def.arguments, is_variable_list=True)
            with ctx.mark_is_effect_definition(False):
                value = self._compose_arguments(ctx, [self.arguments[-1]], predicate_def.return_type.assignment_type())[0]

            feature = E.FunctionApplicationExpression(predicate_def, arguments)

            if self.name.method_name == 'equal':
                return E.PredicateEqualExpression(feature, value)
            elif self.name.method_name == 'assign':
                return E.AssignExpression(feature, value)
            elif self.name.method_name == 'cond-select':
                with ctx.mark_is_effect_definition(False):
                    condition = self._compose_arguments(ctx, [self.arguments[-1]], BOOL)[0]
                return E.ConditionalSelectExpression(feature, condition)
            elif self.name.method_name == 'cond-assign':
                with ctx.mark_is_effect_definition(False):
                    condition = self._compose_arguments(ctx, [self.arguments[-2]], BOOL)[0]
                return E.ConditionalAssignExpression(feature, value, condition)
            else:
                raise NameError('Unknown method name: {}.'.format(self.name.method_name))
        elif self.name == 'and':
            arguments = [arg.compose(expect_value_type) for arg in self.arguments]
            return E.AndExpression(*arguments)
        elif self.name == 'or':
            arguments = [arg.compose(expect_value_type) for arg in self.arguments]
            return E.OrExpression(*arguments)
        elif self.name == 'not':
            arguments = [arg.compose(expect_value_type) for arg in self.arguments]
            return E.NotExpression(*arguments)
        elif self.name == 'equal':
            assert len(self.arguments) == 2, 'PredicateEqualOp takes two arguments, got: {}.'.format(len(self.arguments))
            feature = self.arguments[0]
            feature = _compose(ctx, feature, None)
            value = self.arguments[1]
            value = _compose(ctx, value, feature.return_type.assignment_type())
            return E.PredicateEqualExpression(feature, value)
        elif self.name == 'assign':
            assert len(self.arguments) == 2, 'AssignOp takes two arguments, got: {}.'.format(len(self.arguments))
            assert isinstance(self.arguments[0], _FunctionApplicationImm)
            feature = self.arguments[0].compose(None)
            assert isinstance(feature, E.FunctionApplicationExpression)
            with ctx.mark_is_effect_definition(False):
                value = _compose(ctx, self.arguments[1], feature.return_type.assignment_type())
            return E.AssignExpression(feature, value)
        else:  # the name is a predicate name.
            if self.name in domain.functions or ctx.allow_auto_predicate_def:
                if self.name not in domain.functions:
                    arguments: List[Union[E.ValueOutputExpression, E.VariableExpression]] = self._compose_arguments(ctx, self.arguments, None)
                    argument_types = [arg.return_type for arg in arguments]
                    argument_defs = list()
                    for i, (arg, arg_type) in enumerate(zip(arguments, argument_types)):
                        if isinstance(arg, E.ValueOutputExpression):
                            argument_defs.append(Variable(f'?arg{i}', arg_type))
                        elif isinstance(arg, E.VariableExpression):
                            argument_defs.append(arg.variable)
                        else:
                            raise TypeError(f'Unexpected argument type: {type(arg)}. Can only be ValueOutputExpression or VariableExpression.')
                    self.kwargs.setdefault('state', False)
                    self.kwargs.setdefault('observation', False)
                    generators = self.kwargs.pop('generators', None)
                    self.kwargs['inplace_generators'] = generators
                    predicate_def = domain.define_predicate(self.name, argument_defs, BOOL, **self.kwargs)
                    rv = E.FunctionApplicationExpression(predicate_def, arguments)
                    logger.info(f'Auto-defined predicate {self.name} with arguments {argument_defs} and return type {BOOL}.')

                    # create generators inline
                    if generators is not None:
                        generators: List[str]
                        for i, target_variable_name in enumerate(generators):
                            assert target_variable_name.startswith('?')
                            parameters, context, generates = _canonize_inline_generator_def(ctx, target_variable_name, arguments)
                            generator_name = f'gen-{self.name}-{target_variable_name[1:]}' if len(generators) > 1 else f'gen-{self.name}'
                            domain.define_generator(generator_name, parameters=parameters, certifies=rv, context=context, generates=generates)
                    return rv
                else:
                    assert len(self.kwargs) == 0, 'Cannot specify decorators for non-auto predicate definition.'
                    predicate_def = domain.functions[self.name]
                    arguments = self._compose_arguments(ctx, self.arguments, predicate_def.arguments, is_variable_list=True)
                    return E.FunctionApplicationExpression(predicate_def, arguments)
            else:
                raise ValueError('Unknown function: {}.'.format(self.name))

    def _compose_arguments(self, ctx, arguments, expect_value_type=None, is_variable_list: bool = False) -> List[E.Expression]:
        if isinstance(expect_value_type, (tuple, list)):
            assert len(expect_value_type) == len(arguments), 'Mismatched number of arguments: expect {}, got {}. Expression: {}.'.format(len(expect_value_type), len(arguments), self)

            if is_variable_list:
                output_list = list()
                for arg, var in zip(arguments, expect_value_type):
                    rv = _compose(ctx, arg, var.dtype if var.dtype is not AutoType else None)
                    if var.dtype is AutoType:
                        var.dtype = rv.return_type
                    output_list.append(rv)
                return output_list

            return [_compose(ctx, arg, evt) for arg, evt in zip(arguments, expect_value_type)]
        return [_compose(ctx, arg, expect_value_type) for arg in arguments]


def _compose(ctx, arg, evt=None):
    if isinstance(arg, Variable):
        return ctx[arg.name]
    elif isinstance(arg, ObjectConstant):
        return E.ObjectConstantExpression(arg)
    else:
        return arg.compose(evt)


class _Slot(object):
    def __init__(self, name, kwargs=None):
        self.scope = None
        self.name = name
        self.kwargs = kwargs

        if self.kwargs is None:
            self.kwargs = dict()

    def set_scope(self, scope):
        self.scope = scope

    def __str__(self):
        kwargs = ', '.join([f'{k}={v}' for k, v in self.kwargs.items()])
        return f'??{self.name}[{kwargs}]'


class _MethodName(object):
    def __init__(self, predicate_name, method_name):
        self.predicate_name = predicate_name
        self.method_name = method_name

    def __str__(self):
        return f'{self.predicate_name}::{self.method_name}'


class _QuantifierApplicationImm(object):
    def __init__(self, quantifier, darg: Variable, expr: _FunctionApplicationImm):
        self.quantifier = quantifier
        self.darg = darg
        self.expr = expr

    def __str__(self):
        return f'QIMM::{self.quantifier}({self.darg}: {self.expr})'

    @_log_function
    def compose(self, expect_value_type: Optional[TensorValueTypeBase] = None):
        ctx = get_expression_definition_context()

        with ctx.new_variables(self.darg):
            if ctx.is_effect_definition:
                expr = _canonize_effect(self.expr)
            else:
                expr = self.expr.compose(expect_value_type)

        if self.quantifier in ('foreach', 'forall') and ctx.is_effect_definition:
            outputs = list()
            for e in expr:
                assert isinstance(e, Effect)
                outputs.append(E.DeicticAssignExpression(self.darg, e.assign_expr))
            return outputs
        if self.quantifier == 'foreach':
            assert E.is_value_output_expression(expr)
            return E.DeicticSelectExpression(self.darg, expr)

        assert E.is_value_output_expression(expr), f'Expect value output expression, got {expr}.'
        return E.QuantificationExpression(E.QuantificationOpType.from_string(self.quantifier), self.darg, expr)


@_log_function
def _canonize_precondition(precondition: Union[_FunctionApplicationImm, _QuantifierApplicationImm]):
    if isinstance(precondition, _FunctionApplicationImm) and precondition.name == 'and':
        return list(itertools.chain(*[_canonize_precondition(pre) for pre in precondition.arguments]))
    return [Precondition(precondition.compose(BOOL))]


@_log_function
def _canonize_effect(effect: Union[_FunctionApplicationImm, _QuantifierApplicationImm]):
    if isinstance(effect, _QuantifierApplicationImm):
        effect = effect.compose()
        if isinstance(effect, list):
            effect = [Effect(e) for e in effect]
    else:
        assert isinstance(effect, _FunctionApplicationImm)

        if effect.name == 'and':
            return list(itertools.chain(*[_canonize_effect(eff) for eff in effect.arguments]))

        if isinstance(effect.name, _MethodName):
            effect = effect.compose()
        elif effect.name == 'assign':
            effect = effect.compose()
        elif effect.name == 'not':
            assert len(effect.arguments) == 1, 'NotOp only takes 1 argument, got {}.'.format(len(effect.arguments))
            feat = effect.arguments[0].compose()
            assert feat.return_type == BOOL
            effect = E.AssignExpression(feat, E.ConstantExpression.FALSE)
        elif effect.name == 'when':
            assert len(effect.arguments) == 2, 'WhenOp takes two arguments, got: {}.'.format(len(effect.arguments))
            condition = effect.arguments[0].compose(BOOL)
            if effect.arguments[1].name == 'and':
                inner_effects = effect.arguments[1].arguments
            else:
                inner_effects = [effect.arguments[1]]
            inner_effects = list(itertools.chain(*[_canonize_effect(arg) for arg in inner_effects]))
            effect = list()
            for e in inner_effects:
                assert isinstance(e.assign_expr, E.AssignExpression)
                effect.append(Effect(E.ConditionalAssignExpression(e.assign_expr.predicate, e.assign_expr.value, condition)))
            return effect
        else:
            feat = effect.compose()
            assert isinstance(feat, E.FunctionApplicationExpression) and feat.return_type == BOOL
            effect = E.AssignExpression(feat, E.ConstantExpression.TRUE)

    if isinstance(effect, list):
        return effect

    assert isinstance(effect, E.VariableAssignmentExpression)
    return [Effect(effect)]


def _canonize_controller(controller: _FunctionApplicationImm):
    controller_name = controller.name
    assert isinstance(controller_name, str)

    ctx = get_expression_definition_context()
    arguments = [_compose(ctx, arg, None) for arg in controller.arguments]
    return Controller(controller_name, arguments)


def _canonize_operator_expression_arguments(operator: Operator, arguments: List[_FunctionApplicationImm]):
    ctx = get_expression_definition_context()
    if arguments[-1] == Ellipsis:
        arguments = [_compose(ctx, arg, operator.arguments[i]) for i, arg in enumerate(arguments[:-1])]
        arguments = arguments + [UnnamedPlaceholder(var.dtype) for var in operator.arguments[len(arguments):]]
    else:
        assert len(operator.arguments) == len(arguments), 'Operator {} takes {} arguments, got {}.'.format(operator.name, len(operator.arguments), len(arguments))
        arguments = [_compose(ctx, arg, operator.arguments[i]) for i, arg in enumerate(arguments)]
    return OperatorApplicationExpression(operator, arguments)


def _canonize_inline_generator_def(ctx: ExpressionDefinitionContext, variable_name: str, arguments: List[Union[E.VariableExpression, E.ValueOutputExpression]]):
    parameters, context, generates = list(), list(), list()
    used_parameters = set()

    for arg in arguments:
        if isinstance(arg, E.VariableExpression):
            if arg.name not in used_parameters:
                used_parameters.add(arg.name)
                parameters.append(arg.variable)
            if arg.name == variable_name:
                generates.append(arg)
            else:
                context.append(arg)
        else:
            context.append(arg)
            assert isinstance(arg, E.ValueOutputExpression)
            for sub_expr in E.iter_exprs(arg):
                if isinstance(sub_expr, E.VariableExpression):
                    if sub_expr.name not in used_parameters:
                        used_parameters.add(sub_expr.name)
                        parameters.append(sub_expr.variable)

    assert len(generates) == 1, f'Generator must generate exactly one variable, got {len(generates)}.'
    return parameters, context, generates


def _canonize_inline_generator_def_predicate(domain: Domain, variable_name: str, predicate_def: Predicate):
    parameters = predicate_def.arguments
    context, generates = list(), list()

    if predicate_def.return_type != BOOL:
        for arg in parameters:
            assert arg.name != '?rv', 'Arguments cannot be named ?rv.'
    parameters = tuple(parameters) + (Variable('?rv', predicate_def.return_type),)

    ctx = ExpressionDefinitionContext(*parameters, domain=domain)
    with ctx.as_default():
        for arg in predicate_def.arguments:
            assert isinstance(arg, Variable)
            if arg.name == variable_name:
                generates.append(ctx[arg.name])
            else:
                context.append(ctx[arg.name])
        certifies = E.FunctionApplicationExpression(predicate_def, [ctx[arg.name] for arg in predicate_def.arguments])
        if predicate_def.return_type != BOOL:
            context.append(ctx['?rv'])
            certifies = E.PredicateEqualExpression(certifies, ctx['?rv'])

    assert len(generates) == 1, f'Generator must generate exactly one variable, got {len(generates)}.'
    return parameters, certifies, context, generates


concepts/pdsketch/parsers/pdsketch_parser.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 11/17/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

# from .discrete_search import validate_plan, brute_force_search
# from .optimistic_search import construct_csp_from_optimistic_plan, solve_optimistic_plan, optimistic_search, optimistic_search_strips


concepts/pdsketch/planners/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : discrete_search.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 11/17/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Discrete-space search PDSketch planners."""

import jacinle
import jactorch
import heapq as hq
from typing import Optional, Any, Union, Tuple, Sequence, List, Dict, Callable
from dataclasses import dataclass

from concepts.dsl.dsl_types import PyObjValueType, NamedTensorValueType
from concepts.dsl.expression import ValueOutputExpression
from concepts.dsl.tensor_value import TensorValue
from concepts.algorithm.search.heuristic_search import QueueNode
from concepts.pdsketch.operator import OperatorApplier, gen_all_grounded_actions
from concepts.pdsketch.domain import State
from concepts.pdsketch.executor import PDSketchExecutor
from concepts.pdsketch.planners.search_utils import MostPromisingTrajectoryTracker
from concepts.pdsketch.csp_solvers.brute_force_sampling import ContinuousValueDict
from concepts.pdsketch.strips.strips_expression import SState
from concepts.pdsketch.strips.strips_grounding import GStripsTranslatorOptimistic
from concepts.pdsketch.strips.strips_heuristics import StripsHeuristic
from concepts.pdsketch.strips.strips_search import get_priority_func

__all__ = [
    'apply_action', 'goal_test', 'prepare_search',
    'validate_plan', 'brute_force_search'
]



def apply_action(executor: PDSketchExecutor, state: State, action: OperatorApplier, forward_derived=True) -> Tuple[bool, State]:
    """Apply an action to a state. If success, the function also forwards the axioms and optionally derived predicates.

    Args:
        executor: the executor.
        state: the state to be applied.
        action: the action to be applied.
        forward_derived: whether to forward the derived predicates.

    Returns:
        a tuple of (success, new_state).
    """
    succ, ns = executor.apply(action, state)
    # if succ:
    #     ns = executor.forward_predicates_and_axioms(ns, forward_state_variables=False, forward_axioms=True, forward_derived=forward_derived)
    return succ, ns


def goal_test(
    executor: PDSketchExecutor, state: State, goal_expr: ValueOutputExpression,
    trajectory: Sequence[OperatorApplier], mpt_tracker: Optional[MostPromisingTrajectoryTracker] = None, verbose=False
) -> bool:
    """Test whether a state satisfies a goal expression.

    Args:
        executor: the executor.
        state: the state to be tested.
        goal_expr: the goal expression.
        trajectory: the trajectory that leads to the current state.
        mpt_tracker: the tracker for the most promising trajectory.
        verbose: whether to print the verbose information.

    Returns:
        True if the state satisfies the goal expression.
    """
    score = executor.execute(goal_expr, state).item()

    threshold = 0.5
    if mpt_tracker is not None:
        if mpt_tracker.check(score):
            mpt_tracker.update(score, trajectory)
        threshold = mpt_tracker.threshold

    return score > threshold


def prepare_search(
    func_name,
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression], *,
    actions: Optional[Sequence[OperatorApplier]] = None, continuous_values: Optional[ContinuousValueDict] = None, action_filter: Callable[[OperatorApplier], bool] = None,
    forward_state_variables: bool = True, forward_derived: bool = True,
    verbose=False
) -> Tuple[State, ValueOutputExpression, List[OperatorApplier]]:
    """Prepare for discrete space search.

    Args:
        func_name: the name of the calling function.
        executor: the executor.
        state: the initial state.
        goal_expr: the goal expression.
        actions: the actions to be considered.
        continuous_values: the continuous values for action parameters. If None, all action parameters should be discrete.
        action_filter: the action filter. If None, use all possible actions. It should be a function that takes in an action and returns a boolean.
        forward_state_variables: whether to forward state variables.
        forward_derived: whether to forward derived variables.
        verbose: whether to print verbose information.

    Returns:
        the initial state, the goal expression, and the list of possible actions.
    """
    # NB(Jiayuan Mao @ 08/05): sanity check.
    # state = executor.forward_predicates_and_axioms(state, forward_state_variables, False, forward_derived)

    if actions is None:
        actions = gen_all_grounded_actions(executor, state, continuous_values, action_filter=action_filter)
    goal_expr = executor.parse(goal_expr)

    if verbose:
        # print(func_name+  '::initial_state', state)
        print(func_name + '::actions nr', len(actions))
        print(func_name + '::goal_expr', goal_expr)

    return state, goal_expr, actions


@jactorch.no_grad_func
def validate_plan(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression], actions: Sequence[OperatorApplier], *,
    forward_state_variables=True, forward_derived=True,
) -> Tuple[State, TensorValue]:
    """Validate a plan by executing it on the given state.

    Args:
        executor: the executor.
        state: the initial state.
        goal_expr: the goal expression.
        actions: the sequence of actions to execute.
        forward_state_variables: whether to forward state variables.
        forward_derived: whether to forward derived predicates.

    Returns:
        the final state and the goal value (the execution result of the goal expression).
    """
    # TODO(Jiayuan Mao @ 11/28): update!.
    # state = executor.forward_predicates_and_axioms(state, forward_state_variables, False, forward_derived)

    if isinstance(goal_expr, str):
        goal_expr = executor.domain.parse(goal_expr)
    else:
        assert isinstance(goal_expr, ValueOutputExpression)

    for action in actions:
        succ, state = apply_action(executor, state, action, forward_derived=forward_derived)
        assert succ, f'Action application failed: {action}.'

    score = executor.execute(goal_expr, state)
    return state, score


@jactorch.no_grad_func
def brute_force_search(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression], *,
    max_depth: int = 5,
    actions: Optional[Sequence[OperatorApplier]] = None, continuous_values: Optional[ContinuousValueDict] = None, action_filter: Callable[[OperatorApplier], bool] = None,
    use_tuple_desc: bool = True,
    forward_state_variables: bool = True, forward_derived: bool = True,
    verbose: bool = False, return_extra_info: bool = False
) -> Union[
    Optional[Sequence[OperatorApplier]],
    Tuple[Optional[Sequence[OperatorApplier]], Dict[str, Any]]
]:
    """Brute-force search for a plan that satisfies the goal expression.

    Args:
        executor: the executor.
        state: the initial state.
        goal_expr: the goal expression.
        max_depth: the maximum depth of the search.
        actions: the actions to use. If None, use all possible actions.
        continuous_values: the continuous values for action parameters. If None, all action parameters should be discrete.
        action_filter: the action filter. If None, use all possible actions. It should be a function that takes in an action and returns a boolean.
        use_tuple_desc: whether to use tuple description to prune the search space.
        forward_state_variables: whether to forward state variables before the search starts.
        forward_derived: whether to forward derived predicates after applying actions.
        verbose: whether to print verbose information.
        return_extra_info: whether to return extra information, such as the number of expanded nodes.

    Returns:
        the trajectory if succeeded, otherwise None.
        When `return_extra_info` is True, return a tuple of (trajectory, extra_info), where extra_info is a dictionary.
    """
    state, goal_expr, actions = prepare_search(
        'bfs', executor, state, goal_expr,
        actions=actions, continuous_values=continuous_values, action_filter=action_filter,
        verbose=verbose, forward_state_variables=forward_state_variables, forward_derived=forward_derived,
    )

    states = [(state, tuple())]
    visited = set()
    if use_tuple_desc:
        visited.add(state.generate_tuple_description(executor.domain))

    nr_expanded_states = 0
    nr_tested_actions = 0

    def wrap_extra_info(trajectory):
        if return_extra_info:
            return trajectory, {'nr_expansions': nr_expanded_states, 'nr_tested_actions': nr_tested_actions}
        return trajectory

    pbar = None
    for depth in range(max_depth):
        if verbose:
            pbar = jacinle.tqdm_pbar(desc=f'bfs::depth={depth}')
        with jacinle.cond_with(pbar, verbose):
            next_states = list()

            for s, traj in states:
                nr_expanded_states += 1
                # print(traj, domain.eval(goal_expr, s))
                if goal_test(executor, s, goal_expr, traj):
                    if verbose:
                        print('bfs::search succeeded.')
                        print('bfs::total_expansions:', nr_expanded_states)
                    return wrap_extra_info(traj)

                for a in actions:
                    nr_tested_actions += 1
                    succ, ns = apply_action(executor, s, a, forward_derived=forward_derived)
                    if verbose:
                        pbar.update()
                    nt = traj + (a, )
                    if succ:
                        if use_tuple_desc:
                            nst = ns.generate_tuple_description(executor.domain)
                            if nst not in visited:
                                next_states.append((ns, nt))
                                visited.add(nst)
                        else:  # unconditionally expand
                            next_states.append((ns, nt))

            states = next_states

            if verbose:
                pbar.set_description(f'bfs::depth={depth}, states={len(states)}')
                pbar.close()

    if verbose:
        print('bfs::search failed.')
        print('bfs::total_expansions:', nr_expanded_states)

    return None


@dataclass
class HeuristicSearchState(object):
    """The state for heuristic search."""

    state: State
    """The state."""

    strips_state: SState
    """The STRIPS state."""

    trajectory: Tuple[OperatorApplier, ...]
    """The trajectory."""

    g: float
    """The cost so far."""


@jactorch.no_grad_func
def heuristic_search_strips(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression],
    strips_heuristic: str = 'hff', *,
    max_expansions: int = 100000, max_depth: int = 100,  # search related parameters.
    heuristic_weight: float = float('inf'),  # heuristic related parameters.
    external_heuristic_function: Callable[[State, ValueOutputExpression], int] = None,  # external heuristic related parameters.
    actions: Optional[Sequence[OperatorApplier]] = None, continuous_values: Optional[ContinuousValueDict] = None, action_filter: Callable[[OperatorApplier], bool] = None,
    strips_forward_relevance_analysis: bool = False, strips_backward_relevance_analysis: bool = True,
    strips_use_sas: bool = False,  # whether to use SAS Strips compiler (AODiscretization)
    use_strips_op: bool = False,
    use_tuple_desc: bool = True, # pruning related parameters.
    forward_state_variables: bool = True, forward_derived: bool = False,  # initialization related parameters.
    track_most_promising_trajectory: bool = False, prob_goal_threshold: float = 0.5,  # non-optimal trajectory tracking related parameters.
    verbose: bool = False, return_extra_info: bool = False
) -> Union[
    Optional[Sequence[OperatorApplier]],
    Tuple[Optional[Sequence[OperatorApplier]], Dict[str, Any]]
]:
    """Perform heuristic search with STRIPS-based heuristics.

    Args:
        executor: the executor.
        state: the initial state.
        goal_expr: the goal expression.
        strips_heuristic: the heuristic to use. Should be a string. Use 'external' to use the external heuristic function.
        max_expansions: the maximum number of expanded nodes.
        max_depth: the maximum depth of the search.
        heuristic_weight: the weight of the heuristic. Use float('inf') to do greedy best-first search.
        external_heuristic_function: the external heuristic function. Should be a function that takes in a state and a goal expression, and returns an integer.
        actions: the actions to use. If None, use all possible actions.
        continuous_values: the continuous values for action parameters. If None, all action parameters should be discrete.
        action_filter: the action filter. If None, use all possible actions. It should be a function that takes in an action and returns a boolean.
        strips_forward_relevance_analysis: whether to perform forward relevance analysis when translating the problem into STRIPS.
        strips_backward_relevance_analysis: whether to perform backward relevance analysis when translating the problem into STRIPS.
        strips_use_sas: whether to use SAS Strips compiler (AODiscretization).
        use_strips_op: whether to use STRIPS operators when applying actions. Recommended to be False.
        use_tuple_desc: whether to use tuple description to prune the search space.
        forward_state_variables: whether to forward state variables before the search starts.
        forward_derived: whether to forward derived predicates after applying actions.
        track_most_promising_trajectory: whether to track the most promising trajectory.
        prob_goal_threshold: the probability threshold for the most promising trajectory. When there is a trajectory with probability greater than this threshold, the search will stop.
        verbose: whether to print verbose information.
        return_extra_info: whether to return extra information, such as the number of expanded nodes.

    Returns:
        the trajectory if succeeded, otherwise None.
        When `return_extra_info` is True, return a tuple of (trajectory, extra_info), where extra_info is a dictionary.
    """

    state, goal_expr, actions = prepare_search(
        'hsstrips', executor, state, goal_expr,
        actions=actions, action_filter=action_filter,
        continuous_values=continuous_values,
        forward_state_variables=forward_state_variables, forward_derived=forward_derived,
        verbose=verbose
    )

    if strips_use_sas:
        raise NotImplementedError('SAS strips is not implemented yet.')
    else:
        strips_translator = GStripsTranslatorOptimistic(executor, use_string_name=True, prob_goal_threshold=prob_goal_threshold)

    strips_task = strips_translator.compile_task(
        state, goal_expr, actions,
        is_relaxed=False,
        forward_relevance_analysis=strips_forward_relevance_analysis,
        backward_relevance_analysis=strips_backward_relevance_analysis,
    )

    if strips_heuristic == 'external' and external_heuristic_function is not None:
        pass
    else:
        heuristic = StripsHeuristic.from_type(
            strips_heuristic, strips_task, strips_translator,
            forward_relevance_analysis=strips_forward_relevance_analysis,
            backward_relevance_analysis=strips_backward_relevance_analysis,
        )

    # from IPython import embed; embed()
    # import ipdb; ipdb.set_trace()

    # print(strips_task.goal)
    # print(strips_task.operators)
    # print(heuristic.relaxed.goal)
    # print(heuristic.relaxed.operators)

    mpt_tracker = None
    if track_most_promising_trajectory:
        mpt_tracker = MostPromisingTrajectoryTracker(True, prob_goal_threshold)

    initial_state = HeuristicSearchState(state, strips_task.state, tuple(), 0)
    queue: List[QueueNode] = list()
    visited = set()

    if strips_heuristic == 'external' and external_heuristic_function is not None:
        def heuristic_fn(state: HeuristicSearchState, goal_expr=goal_expr) -> int:
            return external_heuristic_function(state.state, goal_expr)
    else:
        def heuristic_fn(state: HeuristicSearchState) -> int:
            return heuristic.compute(state.strips_state)
    priority_func = get_priority_func(heuristic_fn, heuristic_weight)

    def push_node(node: HeuristicSearchState):
        added = False
        if use_tuple_desc:
            nst = node.state.generate_tuple_description(executor.domain)
            if nst not in visited:
                added = True
                visited.add(nst)
        else:  # unconditionally expand
            added = True

        if added:
            hq.heappush(queue, QueueNode(priority_func(node, node.g), node))
            if heuristic_search_strips.DEBUG:
                print('  hsstrips::push_node:', *node.trajectory)
                print('   ', 'heuristic =', heuristic.compute(node.strips_state), 'g =', node.g)

    push_node(initial_state)

    nr_expanded_states = 0
    nr_tested_actions = 0

    def wrap_extra_info(trajectory):
        if return_extra_info:
            return trajectory, {'nr_expansions': nr_expanded_states, 'nr_tested_actions': nr_tested_actions}
        return trajectory

    pbar = None
    if verbose:
        pbar = jacinle.tqdm_pbar(desc='heuristic_search::expanding')

    while len(queue) > 0 and nr_expanded_states < max_expansions:
        priority, node = hq.heappop(queue)
        nr_expanded_states += 1

        """
        Name convention:
            - node: current node.
            - nnode: next node.
            - s: the state of the search tree.
            - ns: the state of the search tree after the action is applied.
            - a: the action applied.
            - ss: the strips state of the search tree.
            - nss: the strips state of the search tree after the action is applied.
            - traj: the path from the root to the node in the search tree.
            - nt: the path in the search tree after the action is applied.
            - g: the cost of the path from the root to the node in the search tree.
        """

        s, ss, traj = node.state, node.strips_state, node.trajectory
        if heuristic_search_strips.DEBUG:
            print('hsstrips::pop_node:')
            print('  trajectory:', *traj, sep='\n  ')
            print('  priority =', priority, 'g =', node.g)
            if heuristic_search_strips.DEBUG_INTERACTIVE:
                input('  Continue?')
        if verbose:
            pbar.set_description(f'heuristic_search::expanding: priority = {priority} g = {node.g}')
            pbar.update()

        if heuristic_search_strips.DEBUG:
            print('hsstrips::pop_node:', *traj)
            print('  priority =', priority, 'g =', node.g)

        goal_reached = goal_test(
            executor, s, goal_expr,
            trajectory=traj,
            verbose=verbose,
            mpt_tracker=mpt_tracker
        )
        if goal_reached:
            if verbose:
                print('hsstrips::search succeeded.')
                print('hsstrips::total_expansions:', nr_expanded_states)
            return wrap_extra_info(traj)

        if len(traj) >= max_depth:
            continue

        for sa in strips_task.operators:
            a = sa.raw_operator
            succ, ns = apply_action(executor, s, a)
            nt = traj + (a, )

            if succ:
                nss = sa.apply(ss) if use_strips_op else strips_translator.compile_state(ns.clone(), forward_derived=False)
                nnode = HeuristicSearchState(ns, nss, nt, node.g + 1)
                push_node(nnode)

    if verbose:
        print('hsstrips::search failed.')
        print('hsstrips::total_expansions:', nr_expanded_states)

    if mpt_tracker is not None:
        return wrap_extra_info(mpt_tracker.solution)

    return wrap_extra_info(None)


heuristic_search_strips.DEBUG = False
heuristic_search_strips.set_debug = lambda x = True: setattr(heuristic_search_strips, 'DEBUG', x)

heuristic_search_strips.DEBUG_INTERACTIVE = False
heuristic_search_strips.set_debug_interactive = lambda x = True: setattr(heuristic_search_strips, 'DEBUG_INTERACTIVE', x)


concepts/pdsketch/planners/discrete_search.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : optimistic_search.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 11/20/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""CSP-based PDSketch planner. At the high-level,
this planner defers the decision of all continuous parameters to a CSP solver. During forward search,
the planner only make decisions about discrete parameters (mostly object-typed parameters), and record
all preconditions and goal tests that involve continuous parameters to an accumulated CSP data structure.

At each step, the planner will internally call a CSP solver to solve the accumulated CSP. If the CSP solver
returns a solution, the planner will generate the corresponding fully-grounded trajectory based on the
discrete parameters and the CSP solution.
"""

import heapq as hq
from dataclasses import dataclass
from typing import Any, Optional, Union, Sequence, Tuple, List, Dict, Callable

import jactorch

from concepts.dsl.dsl_types import NamedTensorValueType, PyObjValueType, UnnamedPlaceholder
from concepts.dsl.tensor_value import TensorValue
from concepts.dsl.expression import ValueOutputExpression
from concepts.dsl.constraint import OptimisticValue, EqualityConstraint, ConstraintSatisfactionProblem, AssignmentDict, ground_assignment_value
from concepts.algorithm.search.heuristic_search import QueueNode
from concepts.pdsketch.operator import OperatorApplier, gen_all_partially_grounded_actions
from concepts.pdsketch.domain import State, Domain
from concepts.pdsketch.executor import PDSketchExecutor
from concepts.pdsketch.csp_solvers.dpll_sampling import csp_dpll_sampling_solve, csp_dpll_simplify, CSPNoGenerator
from concepts.pdsketch.planners.search_utils import MostPromisingTrajectoryTracker
from concepts.pdsketch.strips.strips_expression import SState
from concepts.pdsketch.strips.strips_heuristics import StripsHeuristic
from concepts.pdsketch.strips.strips_grounding import GStripsTranslatorOptimistic
from concepts.pdsketch.strips.strips_search import get_priority_func


__all__ = [
    'instantiate_action', 'ground_action', 'ground_actions', 'apply_action', 'goal_test',
    'prepare_optimistic_search', 'construct_csp_from_optimistic_plan', 'solve_optimistic_plan', 'optimistic_search', 'optimistic_search_strips'
]


def instantiate_action(csp: ConstraintSatisfactionProblem, action: OperatorApplier) -> OperatorApplier:
    """Instantiate an action by replacing all placeholder values with a new optimistic variable.

    Args:
        csp: the CSP to which the new optimistic variable will be added.
        action: the action to be instantiated.

    Returns:
        the instantiated action.
    """
    new_arguments = list()
    for arg in action.arguments:
        if isinstance(arg, UnnamedPlaceholder):
            assert isinstance(arg.dtype, (NamedTensorValueType, PyObjValueType))
            new_arguments.append(TensorValue.from_optimistic_value(csp.new_actionable_var(arg.dtype, wrap=True)))
        else:
            new_arguments.append(arg)
    return OperatorApplier(action.operator, new_arguments)


def ground_action(executor: PDSketchExecutor, action: OperatorApplier, assignments: AssignmentDict) -> OperatorApplier:
    """Ground a single action with a given assignment to the underlying CSP. Basically, this function looks up the
    assigned value of each optimistic variable that appear in action parameters.

    Args:
        executor: the executor.
        action: the action to be grounded.
        assignments: the solution to the underlying CSP.

    Returns:
        the grounded action.
    """
    new_arguments = list()
    for arg in action.arguments:
        if isinstance(arg, TensorValue):
            if arg.tensor_optimistic_values is not None:
                argv = arg.tensor_optimistic_values.item()
                assert isinstance(argv, int)
                argv = ground_assignment_value(assignments, argv)
                new_arguments.append(argv)
            else:
                new_arguments.append(arg)
        elif isinstance(arg, str):
            new_arguments.append(arg)
        else:
            raise TypeError(f'Unsupported argument type: {type(arg)}.')
    return OperatorApplier(action.operator, new_arguments)


def ground_actions(executor: PDSketchExecutor, actions: Sequence[OperatorApplier], assignments: AssignmentDict) -> List[OperatorApplier]:
    """Ground a list of actions with a given assignment to the underlying CSP. Basically, this function looks up the
    assigned value of each optimistic variable that appear in action parameters.

    Args:
        executor: the executor.
        actions: the list of actions to be grounded.
        assignments: the solution to the underlying CSP.

    Returns:
        the grounded actions.
    """

    return [ground_action(executor, action, assignments) for action in actions]


def apply_action(executor: PDSketchExecutor, state: State, action: OperatorApplier, csp: ConstraintSatisfactionProblem) -> Tuple[Tuple[bool, State], ConstraintSatisfactionProblem]:
    """Apply an action to a state.

    Args:
        executor: the executor.
        state: the state to be updated.
        action: the action to be applied.
        csp: the CSP to which the new constraints will be added.

    Returns:
        a tuple of (goal test result, the updated state), and the updated CSP.
    """

    csp = csp.clone()
    succ, state = executor.apply(action, state, csp=csp)
    return (succ, state), csp


def goal_test(
    executor: PDSketchExecutor, state: State, goal_expr: ValueOutputExpression, csp: ConstraintSatisfactionProblem, trajectory: Sequence[OperatorApplier],
    csp_max_generator_trials: int = 3, mpt_tracker: Optional[MostPromisingTrajectoryTracker] = None,
    verbose: bool = False
) -> Optional[List[OperatorApplier]]:
    """Test whether a state satisfies a goal expression with DPLL+Sampling CSP solver.

    Args:
        executor: the executor.
        state: the state to be tested.
        goal_expr: the goal expression.
        csp: the CSP that contains all constraints that has been accumulated so far.
        trajectory: the trajectory that leads to the current state.
        csp_max_generator_trials: the maximum number of trials for calling CSP generators.
        mpt_tracker: the tracker for the most promising trajectory.
        verbose: whether to print verbose information.

    Returns:
        the trajectory of grounded actions that leads to the goal state if the goal is satisfied, otherwise None.
    """

    csp = csp.clone()
    rv = executor.execute(goal_expr, state=state, csp=csp).item()
    # print("  opt::goal_test", *trajectory, sep='\n    ')
    # print("  opt::goal_test", rv)
    if isinstance(rv, OptimisticValue):
        csp.add_constraint(EqualityConstraint.from_bool(rv, True), note='goal_test')
        try:
            if verbose:
                pass
                # print("  opt::final_csp_solve", *trajectory, sep='\n    ')
                # print("  opt::final_csp_solve", jacinle.indent_text(csp).lstrip())
            assignments = csp_dpll_sampling_solve(executor, csp, max_generator_trials=csp_max_generator_trials)
        except CSPNoGenerator:
            return None
        if assignments is not None:
            # TODO(Jiayuan Mao @ 2022/12/16): implement mpt_tracker?
            return ground_actions(executor, trajectory, assignments)
        else:
            return None
    else:
        rv = float(rv)
        threshold = mpt_tracker.threshold if mpt_tracker is not None else 0.5
        if rv > threshold or mpt_tracker is not None:
            if rv <= threshold:
                if not mpt_tracker.check(rv):
                    return None

            try:
                assignments = csp_dpll_sampling_solve(executor, csp, max_generator_trials=csp_max_generator_trials)
            except CSPNoGenerator:
                raise None
            if assignments is not None:
                plan = ground_actions(executor, trajectory, assignments)
                if mpt_tracker is not None:
                    mpt_tracker.update(rv, plan)
                if rv > threshold:
                    return ground_actions(executor, trajectory, assignments)
            else:
                return None
        else:
            return None


def optimistic_search_domain_check(domain: Domain):
    """Check whether a domain is suitable for optimistic search.

    Args:
        domain: the domain to be checked.

    Raises:
        NotImplementedError: if the domain is not suitable for optimistic search.
    """
    for op in domain.operators.values():
        if op.is_axiom:
            raise NotImplementedError('Optimistic search does not support axioms.')


def prepare_optimistic_search(
    func_name: str,
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression], *,
    actions: Optional[Sequence[OperatorApplier]] = None, action_filter: Callable[[OperatorApplier], bool] = None,
    verbose: bool = False, forward_state_variables=True, forward_derived=True
) -> Tuple[State, ValueOutputExpression, Sequence[OperatorApplier]]:
    """Prepare for optimistic search.

    Args:
        func_name: the name of the calling function.
        executor: the executor.
        state: the initial state.
        goal_expr: the goal expression.
        actions: the actions to be considered.
        action_filter: the action filter. If None, use all possible actions. It should be a function that takes in an action and returns a boolean.
        verbose: whether to print verbose information.
        forward_state_variables: whether to forward state variables.
        forward_derived: whether to forward derived variables.

    Returns:
        the initial state, the goal expression, and the actions to be considered.
    """
    # state = session.forward_predicates_and_axioms(state, forward_state_variables, False, forward_derived)
    if actions is None:
        actions = gen_all_partially_grounded_actions(executor, state, action_filter=action_filter)

    goal_expr = executor.domain.parse(goal_expr)

    if verbose:
        print(func_name + '::initial_state', state)
        print(func_name + '::actions nr', len(actions))
        if len(actions) < 20:
            for action in actions:
                print(' ', action)
        print(func_name + '::goal_expr', goal_expr)

    return state, goal_expr, actions


@jactorch.no_grad_func
def construct_csp_from_optimistic_plan(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression], actions: Sequence[OperatorApplier], *,
    simplify: bool = False, verbose: bool = False
) -> Tuple[List[OperatorApplier], ConstraintSatisfactionProblem]:
    """Construct a CSP from an optimistic plan.

    Args:
        executor: the executor.
        state: the initial state.
        goal_expr: the goal expression.
        actions: a list of partially grounded actions.
        simplify: whether to simplify the CSP.
        verbose: whether to print verbose information.

    Returns:
        the list of grounded actions, and the constructed CSP.
    """

    optimistic_search_domain_check(executor.domain)

    csp = ConstraintSatisfactionProblem()
    if isinstance(goal_expr, str):
        goal_expr = executor.parse(goal_expr)
    else:
        assert isinstance(goal_expr, ValueOutputExpression)

    action_groundings = list()
    for a in actions:
        action_grounding = instantiate_action(csp, a)
        action_groundings.append(action_grounding)
        (succ, state), csp = apply_action(executor, state, action_grounding, csp)

        if succ:
            pass
        else:
            raise ValueError(f'Unable to perform action {action_grounding} at state {state}.')

    rv = executor.execute(goal_expr, state=state, csp=csp).item()
    if isinstance(rv, OptimisticValue):
        csp.add_constraint(EqualityConstraint.from_bool(rv, True), note='goal_test')

    if simplify:
        return action_groundings, csp_dpll_simplify(executor, csp)

    return action_groundings, csp


@jactorch.no_grad_func
def solve_optimistic_plan(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression], actions: Sequence[OperatorApplier], *,
    csp_max_generator_trials: int = 3, verbose: bool = False
) -> Tuple[State, ConstraintSatisfactionProblem, Optional[List[OperatorApplier]]]:
    """Solve an optimistic plan using the DPLL+Sampling algorithm.

    Args:
        executor: the executor.
        state: the initial state.
        goal_expr: the goal expression.
        actions: a sequence of partially grounded actions.
        csp_max_generator_trials: the maximum number of trials for calling CSP generators.
        verbose: whether to print verbose information.

    Returns:
        a list of actions that solves the plan, or None if no solution is found.
    """

    optimistic_search_domain_check(executor.domain)
    csp = ConstraintSatisfactionProblem()
    if isinstance(goal_expr, str):
        goal_expr = executor.parse(goal_expr)
    else:
        assert isinstance(goal_expr, ValueOutputExpression)

    action_groundings = list()
    for a in actions:
        action_grounding = instantiate_action(csp, a)
        action_groundings.append(action_grounding)
        (succ, state), csp = apply_action(executor, state, action_grounding, csp)
        if succ:
            pass
        else:
            raise ValueError(f'Unable to perform action {action_grounding} at state {state}.')

    plan = goal_test(executor, state, goal_expr, csp, trajectory=action_groundings, csp_max_generator_trials=csp_max_generator_trials, verbose=verbose)
    if plan is not None:
        return state, csp, plan
    return state, csp, None


@jactorch.no_grad_func
def optimistic_search(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression], *,
    max_depth: int = 5,
    csp_max_generator_trials: int = 3,
    actions: Optional[Sequence[OperatorApplier]] = None, action_filter: Callable[[OperatorApplier], bool] = None,
    use_tuple_desc: bool = False, use_csp_pruning: bool = True,
    forward_state_variables: bool = True, forward_derived: bool = False,
    verbose: bool = False, return_extra_info: bool = False
) -> Union[
    Optional[Sequence[OperatorApplier]],
    Tuple[Optional[Sequence[OperatorApplier]], Dict[str, Any]]
]:
    """Perform brute-force search with DPLL+Sampling algorithm for mixed discrete-continuous domains.

    Args:
        executor: the executor.
        state: the initial state.
        goal_expr: the goal expression.
        max_depth: the maximum depth of the search.
        csp_max_generator_trials: the maximum number of trials for calling CSP generators.
        actions: the actions to use. If None, use all possible actions. Partially grounded actions are allowed.
        action_filter: the action filter. If None, use all possible actions. It should be a function that takes in an action and returns a boolean.
        use_tuple_desc: whether to use tuple description to prune the search space.
        use_csp_pruning: whether to use partial CSP solver to prune the search space.
        forward_state_variables: whether to forward state variables before the search starts.
        forward_derived: whether to forward derived predicates after applying actions.
        verbose: whether to print verbose information.
        return_extra_info: whether to return extra information, such as the number of expanded nodes.

    Returns:
        the trajectory if succeeded, otherwise None.
        When `return_extra_info` is True, return a tuple of (trajectory, extra_info), where extra_info is a dictionary.
    """

    optimistic_search_domain_check(executor.domain)
    state, goal_expr, actions = prepare_optimistic_search(
        'opt', executor, state, goal_expr,
        actions=actions, action_filter=action_filter,
        verbose=verbose, forward_state_variables=forward_state_variables, forward_derived=forward_derived
    )

    assert not forward_derived, 'Not implemented.'

    states = [(state, tuple(), ConstraintSatisfactionProblem())]
    visited = set()
    if use_tuple_desc:
        visited.add(state.generate_tuple_description(executor.domain))

    nr_expanded_states = 0
    nr_tested_actions = 0

    def wrap_extra_info(trajectory):
        if return_extra_info:
            return trajectory, {'nr_expanded_states': nr_expanded_states, 'nr_tested_actions': nr_tested_actions}
        return trajectory

    for depth in range(max_depth):
        next_states = list()

        for s, traj, csp in states:
            nr_expanded_states += 1
            for a in actions:
                nr_tested_actions += 1
                ncsp = csp.clone()
                action_grounding = instantiate_action(ncsp, a)

                (succ, ns), ncsp = apply_action(executor, s, action_grounding, ncsp)
                nt = traj + (action_grounding, )

                if succ:
                    if use_csp_pruning:
                        try:
                            if not csp_dpll_sampling_solve(executor, ncsp, solvable_only=True, max_generator_trials=csp_max_generator_trials):
                                continue
                        except CSPNoGenerator:
                            pass

                    if use_tuple_desc:
                        nst = ns.generate_tuple_description(executor.domain)
                        if nst not in visited:
                            next_states.append((ns, nt, ncsp))
                            visited.add(nst)
                        else:
                            continue
                    else:  # unconditionally expand
                        next_states.append((ns, nt, ncsp))

                    plan = goal_test(executor, ns, goal_expr, ncsp, trajectory=nt, csp_max_generator_trials=csp_max_generator_trials, verbose=verbose)
                    if plan is not None:
                        if verbose:
                            print(f'opt::finished: depth={depth} nr_expanded_states={nr_expanded_states} nr_tested_actions={nr_tested_actions}.')
                        return wrap_extra_info(plan)

        states = next_states

        if verbose:
            print(f'opt::depth={depth}, this_layer_states={len(states)}')

    return wrap_extra_info(None)


@dataclass
class OptHeuristicSearchState(object):
    """The state for heuristic search."""

    state: State
    """The state."""

    strips_state: SState
    """The STRIPS state."""

    trajectory: Tuple[OperatorApplier, ...]
    """The trajectory."""

    csp: ConstraintSatisfactionProblem
    """The CSP that has been accumulated so far."""

    g: float
    """The cost so far."""


@jactorch.no_grad_func
def optimistic_search_strips(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression],
    strips_heuristic: str = 'hff', *,
    max_expansions=100000, max_depth=100,  # search related parameters.
    csp_max_generator_trials: int = 3,
    heuristic_weight: float = float('inf'),  # heuristic related parameters.
    actions: Optional[Sequence[OperatorApplier]] = None, action_filter: Callable[[OperatorApplier], bool] = None,
    use_strips_op: bool = False,
    use_tuple_desc: bool = False, use_csp_pruning: bool = True,  # pruning related parameters.
    forward_state_variables: bool = True, forward_derived: bool = False,  # initialization related parameters.
    track_most_promising_trajectory: bool = False, prob_goal_threshold: float = 0.5,  # non-optimal trajectory tracking related parameters.
    verbose: bool = False, return_extra_info: bool = False
):
    """Perform heuristic search with DPLL+Sampling algorithm and STRIPS-based heuristics, for mixed discrete-continuous domains.

    Args:
        executor: the executor.
        state: the initial state.
        goal_expr: the goal expression.
        strips_heuristic: the heuristic to use. Should be a string.
        max_expansions: the maximum number of expanded nodes.
        max_depth: the maximum depth of the search.
        csp_max_generator_trials: the maximum number of trials for calling CSP generators.
        heuristic_weight: the weight of the heuristic. Use float('inf') to do greedy best-first search.
        actions: the actions to use. If None, use all possible actions. Partially grounded actions are allowed.
        action_filter: the action filter. If None, use all possible actions. It should be a function that takes in an action and returns a boolean.
        use_strips_op: whether to use STRIPS operators when applying actions. Recommended to be False.
        use_tuple_desc: whether to use tuple description to prune the search space.
        forward_state_variables: whether to forward state variables before the search starts.
        forward_derived: whether to forward derived predicates after applying actions.
        track_most_promising_trajectory: whether to track the most promising trajectory.
        prob_goal_threshold: the probability threshold for the most promising trajectory. When there is a trajectory with probability greater than this threshold, the search will stop.
        verbose: whether to print verbose information.
        return_extra_info: whether to return extra information, such as the number of expanded nodes.

    Returns:
        the trajectory if succeeded, otherwise None.
        When `return_extra_info` is True, return a tuple of (trajectory, extra_info), where extra_info is a dictionary.
    """

    assert not forward_derived, 'Not implemented.'
    optimistic_search_domain_check(executor.domain)

    state, goal_expr, actions = prepare_optimistic_search(
        'optsstrips', executor, state, goal_expr,
        actions=actions, action_filter=action_filter,
        verbose=verbose, forward_state_variables=forward_state_variables, forward_derived=forward_derived
    )

    # TODO(Jiayuan Mao @ 2022/12/16): Relevance analysis for optimistic planning tasks.
    strips_translator = GStripsTranslatorOptimistic(executor, use_string_name=verbose)
    strips_task = strips_translator.compile_task(
        state, goal_expr, actions,
        is_relaxed=False,
        forward_relevance_analysis=False,
        backward_relevance_analysis=False
    )

    # TODO(Jiayuan Mao @ 2022/12/16): Support external heuristics.

    mpt_tracker = None
    if track_most_promising_trajectory:
        mpt_tracker = MostPromisingTrajectoryTracker(True, prob_goal_threshold)

    initial_state = OptHeuristicSearchState(state, strips_task.state, tuple(), ConstraintSatisfactionProblem(), 0)
    queue: List[QueueNode] = list()
    visited = set()

    heuristic = StripsHeuristic.from_type(strips_heuristic, strips_task, strips_translator)

    def heuristic_fn(state: OptHeuristicSearchState) -> int:
        return heuristic.compute(state.strips_state)
    priority_func = get_priority_func(heuristic_fn, heuristic_weight)

    def push_node(node: OptHeuristicSearchState):
        added = False
        if use_tuple_desc:
            nst = node.state.generate_tuple_description(executor.domain)
            if nst not in visited:
                added = True
                visited.add(nst)
        else:  # unconditionally expand
            added = True

        if added:
            hq.heappush(queue, QueueNode(priority_func(node, node.g), node))
            if optimistic_search_strips.DEBUG:
                print('  optsstrips::push_node:', *node.trajectory, sep='\n    ')
                print('   ', 'heuristic =', heuristic.compute(node.strips_state), 'g =', node.g)

    push_node(initial_state)

    nr_expanded_states = 0
    nr_tested_actions = 0

    def wrap_extra_info(trajectory):
        if return_extra_info:
            return trajectory, {'nr_expansions': nr_expanded_states, 'nr_tested_actions': nr_tested_actions}
        return trajectory

    is_initial_state = True
    while len(queue) > 0 and nr_expanded_states < max_expansions:
        priority, node = hq.heappop(queue)
        nr_expanded_states += 1

        s, ss, traj, csp = node.state, node.strips_state, node.trajectory, node.csp
        if optimistic_search_strips.DEBUG:
            print('optsstrips::pop_node:')
            print('  trajectory:', *traj, sep='\n  ')
            print('  priority =', priority, 'g =', node.g)
            if optimistic_search_strips.DEBUG_INTERACTIVE:
                input('  Continue?')

        if track_most_promising_trajectory and is_initial_state:
            is_initial_state = False
        else:
            plan = goal_test(
                executor, s, goal_expr, csp,
                trajectory=traj,
                csp_max_generator_trials=csp_max_generator_trials,
                verbose=verbose,
                mpt_tracker=mpt_tracker
            )
            if plan is not None:
                if verbose:
                    print('optsstrips::search succeeded.')
                    print('optsstrips::total_expansions:', nr_expanded_states)
                return wrap_extra_info(plan)

        if len(traj) >= max_depth:
            continue

        for sa in strips_task.operators:
            nr_tested_actions += 1
            a = sa.raw_operator
            ncsp = csp.clone()
            action_grounding = instantiate_action(ncsp, a)

            (succ, ns), ncsp = apply_action(executor, s, action_grounding, ncsp)
            nt = traj + (action_grounding, )

            if succ:
                if use_csp_pruning:
                    try:
                        if optimistic_search_strips.DEBUG:
                            print('  optsstrips::running CSP pruning...', *nt, sep='\n    ')
                        if not csp_dpll_sampling_solve(executor, ncsp, solvable_only=True, max_generator_trials=csp_max_generator_trials):
                            if optimistic_search_strips.DEBUG:
                                print('  optsstrips::pruned:', *nt, sep='\n    ')
                            continue
                    except CSPNoGenerator:
                        pass

                nss = sa.apply(ss) if use_strips_op else strips_translator.compile_state(ns.clone(), forward_derived=False)
                nnode = OptHeuristicSearchState(ns, nss, nt, ncsp, node.g + 1)
                push_node(nnode)

    if verbose:
        print('optsstrips::search failed.')
        print('optsstrips::total_expansions:', nr_expanded_states)

    if mpt_tracker is not None:
        return mpt_tracker.solution

    return None


optimistic_search_strips.DEBUG = False
optimistic_search_strips.set_debug = lambda x = True: setattr(optimistic_search_strips, 'DEBUG', x)

optimistic_search_strips.DEBUG_INTERACTIVE = False
optimistic_search_strips.set_debug_interactive = lambda x = True: setattr(optimistic_search_strips, 'DEBUG_INTERACTIVE', x)


concepts/pdsketch/planners/optimistic_search.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : optimistic_search_bilevel_legacy.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/10/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from typing import Union, Sequence, Iterable, Tuple, List, Set

import jacinle

from concepts.dsl.dsl_types import QINDEX
from concepts.dsl.expression import ValueOutputExpression, is_and_expr
from concepts.dsl.constraint import ConstraintSatisfactionProblem
from concepts.pdsketch.domain import State
from concepts.pdsketch.executor import PDSketchExecutor
from concepts.pdsketch.operator import OperatorApplier, AchieveExpression, OperatorApplicationExpression
from concepts.pdsketch.planners.optimistic_search_bilevel_utils import(
    ground_fol_expression, ground_operator_application_expression, surface_fol_downcast,
    extract_bounded_variables_from_nonzero, extract_bounded_variables_from_nonzero_dc,
    TotallyOrderedPlan, PartiallyOrderedPlan
)

__all__ = [
    'gen_applicable_regression_rules', 'enumerate_possible_symbolic_plans_regression_c',
    'gen_applicable_regression_rules_v2', 'enumerate_possible_symbolic_plans_regression_c_v2'
]


def enumerate_possible_symbolic_plans_regression_c(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression],
    max_actions: int = 10,
) -> Iterable[Tuple[Sequence[OperatorApplier], ConstraintSatisfactionProblem, State, State]]:
    """Enumerate all possible plans that can achieve the goal.

    Args:
        executor: the executor.
        state: the initial state.
        goal_expr: the goal expression.
        max_actions: the maximum number of actions in a plan.

    Returns:
        A list of plans. Each plan is a tuple of (actions, csp, initial_state, final_state).
    """
    if isinstance(goal_expr, str):
        goal_expr = executor.parse(goal_expr)

    @jacinle.log_function(verbose=False)
    def dfs(state, goal_set, maintains, csp, previous_actions, nr_high_level_actions=0):
        jacinle.log_function.print('Current goal_set', goal_set)
        rv = min([executor.execute(goal, state=state, bounded_variables=dict()).item() for goal in goal_set])
        if rv > 0.5:
            return [(state, None, previous_actions)]

        if nr_high_level_actions > max_actions:
            return []

        # jacinle.log_function.print('Current state', state)
        # import ipdb; ipdb.set_trace()

        all_possible_plans = list()
        candidate_regression_rules = gen_applicable_regression_rules(executor, state, goal_set, maintains)
        jacinle.log_function.print('Candidate regression rules', candidate_regression_rules)

        if len(candidate_regression_rules) == 0:
            print('Entering a debug mode: no candidate regression rules.')
            import ipdb; ipdb.set_trace()
            return []

        for goal_index, rule, bounded_variables in candidate_regression_rules:
            grounded_subgoals = list()
            for item in rule.body[:-1]:
                assert isinstance(item, AchieveExpression)
                grounded_subgoals.append(ground_fol_expression(item.goal, bounded_variables))
            assert isinstance(rule.body[-1], OperatorApplicationExpression)

            if len(goal_set) == 1:  # If there is only one single goal left, we should just directly try this regression rule.
                max_prefix_length = 0
            else:
                # TODO(Jiayuan Mao @ 2023/09/06): find the last [always] achievement statement.
                max_prefix_length = len(rule.body) - 1

            for prefix_length in range(0, max_prefix_length + 1):  # assuming the last item is OperatorApplicationExpression
                jacinle.log_function.print('Applying rule', rule, 'for', goal_set[goal_index], 'and prefix length', prefix_length)
                other_goals = goal_set[:goal_index] + goal_set[goal_index + 1:] + grounded_subgoals[:prefix_length]

                if len(other_goals) > 0:
                    # jacinle.log_function.print('Other goals', other_goals)
                    possible_branches = dfs(state, other_goals, maintains, csp, previous_actions, nr_high_level_actions=nr_high_level_actions + 1)
                else:
                    possible_branches = [(state, csp, previous_actions)]

                if len(possible_branches) == 0:
                    jacinle.log_function.print('Prefix planning failed!!! Stop.')
                    # If it's not possible to achieve the subset of goals, then it's not possible to achieve the whole goal set.
                    # Therefore, this is a break, not a continue.
                    break

                for i in range(prefix_length, len(rule.body)):
                    item = rule.body[i]
                    next_possible_branches = list()
                    jacinle.log_function.print(f'Now enter item #{i}: {item}')
                    for cur_state, cur_csp, actions in possible_branches:
                        if isinstance(item, AchieveExpression):
                            subgoal = grounded_subgoals[i]
                            # jacinle.log_function.print('Achieve subgoal', subgoal)
                            next_possible_branches.extend(dfs(cur_state, [subgoal], maintains, None, actions, nr_high_level_actions=nr_high_level_actions + 1))
                        elif isinstance(item, OperatorApplicationExpression):
                            cur_action = ground_operator_application_expression(item, bounded_variables)
                            succ, cur_state = executor.apply(cur_action, cur_state, csp=None, clone=True)
                            jacinle.log_function.print('Applying action', cur_action, 'with result', succ)
                            if succ:
                                next_possible_branches.append((cur_state, None, actions + [cur_action]))
                        else:
                            raise TypeError()
                    possible_branches = next_possible_branches

                # all_possible_plans.extend(possible_branches)

                found_plan = False
                for plan in possible_branches:
                    rv = min([executor.execute(goal, state=plan[0], bounded_variables=dict()).item() for goal in goal_set])
                    if rv > 0.5:
                        jacinle.log_function.print('Found plan', plan[2])
                        found_plan = True
                        all_possible_plans.append(plan)
                    else:
                        jacinle.log_function.print('Plan', plan[2], 'failed')
                if found_plan:
                    break

            # TODO(Jiayuan Mao @ 2023/09/06): think about what's the semantics for this early return.
            # The idea is that, as soon as we found a plan using one of the regression rules, we return it, and don't worry about the other rules.
            # Therefore, it is not guaranteed that the returned plan is the shortest (optimal) one.
            if len(all_possible_plans) > 0:
                break

        if len(all_possible_plans) == 0:
            jacinle.log_function.print('No possible plans for goal set', goal_set)
            return []

        shortest_plan = min(all_possible_plans, key=lambda x: len(x[2]))
        return [shortest_plan]

    if is_and_expr(goal_expr):
        goal_set = list(goal_expr.arguments)
    else:
        goal_set = [goal_expr]

    candidate_plans = dfs(state, goal_set, None, None, list())
    candidate_plans = [(actions, csp, state, final_state) for final_state, csp, actions in candidate_plans]
    return candidate_plans


def gen_applicable_regression_rules(executor: PDSketchExecutor, state: State, goal_set: List[ValueOutputExpression], maintains: Set[ValueOutputExpression]):
    """Generate applicable regression rules for the given goal set and maintains set."""

    # TODO(Jiayuan Mao @ 2023/09/04): implement CSP.
    # TODO(Jiayuan Mao @ 2023/09/04): implement "maintains".

    candidate_regression_rules = list()
    for i, goal in enumerate(goal_set):
        for regression_rule in executor.domain.regression_rules.values():
            goal_expr = regression_rule.goal_expression
            if (variable_binding := surface_fol_downcast(goal_expr, goal)) is None:
                continue
            bounded_variables = dict()
            for v in regression_rule.goal_arguments:
                bounded_variables[v] = variable_binding[v.name].name

            if len(regression_rule.binding_arguments) > 0:
                for v in regression_rule.binding_arguments:
                    bounded_variables[v] = QINDEX

            if len(regression_rule.preconditions_conjunction.arguments) > 0:
                rv = executor.execute(regression_rule.preconditions_conjunction, state=state, bounded_variables=bounded_variables)
            else:
                rv = None

            if rv is None:
                candidate_regression_rules.append((i, regression_rule, bounded_variables))
            else:
                all_forall, candidate_bounded_variables = extract_bounded_variables_from_nonzero(state, rv, regression_rule, default_bounded_variables=bounded_variables)
                if all_forall and regression_rule.always:
                    candidate_regression_rules.append((i, regression_rule, candidate_bounded_variables[0]))
                else:
                    candidate_regression_rules.extend([(i, regression_rule, cbv) for cbv in candidate_bounded_variables])
    return candidate_regression_rules


def enumerate_possible_symbolic_plans_regression_c_v2(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression],
    max_depth: int = 10, enable_reordering: bool = True,
    verbose: bool = False
) -> Iterable[Tuple[Sequence[OperatorApplier], ConstraintSatisfactionProblem, State, State]]:
    """Enumerate all possible plans that can achieve the goal.

    Args:
        executor: the executor.
        state: the initial state.
        goal_expr: the goal expression.
        max_depth: the maximum depth of goal regression search.
        enable_reordering: whether to enable reordering of the subgoals coming from different chains.
        verbose: whether to print verbose information.

    Returns:
        A list of plans. Each plan is a tuple of (actions, csp, initial_state, final_state).
    """
    if isinstance(goal_expr, str):
        goal_expr = executor.parse(goal_expr)

    search_cache = dict()
    search_stat = {'nr_expanded_nodes': 0}

    def return_with_cache(goal_set, previous_actions, rv):
        if len(previous_actions) == 0:
            goal_str = goal_set.gen_string()
            if goal_str not in search_cache:
                search_cache[goal_str] = rv
        return rv

    def try_retrive_cache(goal_set, previous_actions):
        if len(previous_actions) == 0:
            goal_str = goal_set.gen_string()
            if goal_str in search_cache:
                return search_cache[goal_str]
        return None

    @jacinle.log_function(verbose=False)
    def dfs(state: State, goal_set: PartiallyOrderedPlan, maintains, csp, previous_actions, nr_high_level_actions=0):
        if verbose:
            jacinle.log_function.print('Current goal_set', goal_set)
            jacinle.log_function.print('Previous actions', previous_actions)

        if (rv := try_retrive_cache(goal_set, previous_actions)) is not None:
            return rv

        if len(goal_set) == 0:
            return return_with_cache(goal_set, previous_actions, [(state, None, previous_actions)])
        rv = min([executor.execute(goal, state=state, bounded_variables=dict()).item() for goal in goal_set.iter_goals()])
        if rv > 0.5:
            return return_with_cache(goal_set, previous_actions, [(state, None, previous_actions)])

        if nr_high_level_actions > max_depth:
            return return_with_cache(goal_set, previous_actions, [])

        search_stat['nr_expanded_nodes'] += 1

        all_possible_plans = list()
        candidate_regression_rules = gen_applicable_regression_rules_v2(executor, state, goal_set, maintains)
        if verbose:
            jacinle.log_function.print('Candidate regression rules', candidate_regression_rules)

        if sum(len(r) for _, _, r in candidate_regression_rules) == 0:
            # print('Entering a debug mode: no candidate regression rules.')
            # import ipdb; ipdb.set_trace()
            return return_with_cache(goal_set, previous_actions, [])

        for chain_index, subgoal_index, this_candidate_regression_rules in candidate_regression_rules:
            other_goals = goal_set.exclude(chain_index, subgoal_index)
            cur_goal = goal_set.chains[chain_index].sequence[subgoal_index]
            if verbose:
                jacinle.log_function.print('Now trying to excluding goal', cur_goal)

            if len(other_goals) == 0:
                other_goals_plans = [(state, None, previous_actions)]
            else:
                other_goals_plans = dfs(state, other_goals, maintains, csp, previous_actions, nr_high_level_actions=nr_high_level_actions)
                for cur_state, cur_csp, cur_actions in other_goals_plans:
                    rv = executor.execute(cur_goal, state=state, bounded_variables=dict()).item()
                    if rv > 0.5:  # Found a plan that directly achieve the goal without any regression.
                        all_possible_plans.append((cur_state, cur_csp, cur_actions))

            if len(all_possible_plans) > 0:
                break

            if len(other_goals) == 0 or not enable_reordering:
                max_prefix_length = 0
            else:
                max_prefix_length = max(len(rule.body) - 1 for rule, _ in this_candidate_regression_rules)
            grounded_subgoals_cache = dict()

            for prefix_length in range(max_prefix_length + 1):
                for regression_rule_index, (rule, bounded_variables) in enumerate(this_candidate_regression_rules):
                    if prefix_length > len(rule.body) - 1:
                        continue
                    if verbose:
                        jacinle.log_function.print('Applying rule', rule, 'for', cur_goal, 'and prefix length', prefix_length, 'goal set is', goal_set)

                    # construct grounded_subgoals
                    if regression_rule_index in grounded_subgoals_cache:
                        grounded_subgoals = grounded_subgoals_cache[regression_rule_index]
                    else:
                        grounded_subgoals = list()
                        for item in rule.body:
                            if isinstance(item, AchieveExpression):
                                grounded_subgoals.append(ground_fol_expression(item.goal, bounded_variables))
                        grounded_subgoals_cache[regression_rule_index] = grounded_subgoals

                    if prefix_length == 0:
                        possible_branches = other_goals_plans
                    else:
                        cur_other_goals = other_goals.add_chain(grounded_subgoals[:prefix_length])
                        possible_branches = dfs(state, cur_other_goals, maintains, csp, previous_actions, nr_high_level_actions=nr_high_level_actions + 1)

                    if len(possible_branches) == 0:
                        if verbose:
                            jacinle.log_function.print('Prefix planning failed!!! Stop.')
                        # If it's not possible to achieve the subset of goals, then it's not possible to achieve the whole goal set.
                        # Therefore, this is a break, not a continue.
                        break

                    for i in range(prefix_length, len(rule.body)):
                        item = rule.body[i]
                        next_possible_branches = list()
                        if verbose:
                            jacinle.log_function.print(f'Now enter item #{i}: {item}')
                        for cur_state, cur_csp, actions in possible_branches:
                            if isinstance(item, AchieveExpression):
                                subgoal = grounded_subgoals[i]
                                # jacinle.log_function.print('Achieve subgoal', subgoal)
                                next_possible_branches.extend(dfs(cur_state, PartiallyOrderedPlan.from_single_goal(subgoal), maintains, None, actions, nr_high_level_actions=nr_high_level_actions + 1))
                            elif isinstance(item, OperatorApplicationExpression):
                                cur_action = ground_operator_application_expression(item, bounded_variables)
                                succ, cur_state = executor.apply(cur_action, cur_state, csp=None, clone=True)
                                # if verbose:
                                #     jacinle.log_function.print('Applying action', cur_action, 'with result', succ)
                                if succ:
                                    next_possible_branches.append((cur_state, None, actions + [cur_action]))
                            else:
                                raise TypeError()
                        possible_branches = next_possible_branches

                    # all_possible_plans.extend(possible_branches)

                    found_plan = False
                    for plan in possible_branches:
                        rv = min([executor.execute(goal, state=plan[0], bounded_variables=dict()).item() for goal in goal_set.iter_goals()])
                        if rv > 0.5:
                            if verbose:
                                jacinle.log_function.print('Found plan', plan[2], 'for goal', goal_set)
                            found_plan = True
                            all_possible_plans.append(plan)
                        else:
                            if verbose:
                                jacinle.log_function.print('Plan', plan[2], 'failed')
                    if found_plan:
                        # TODO(Jiayuan Mao @ 2023/09/06): since we have changed the order of prefix_length for-loop and the regression rule for-loop.
                        # We need to use an additional dictionary to store whether we have found a plan for a particular regression rule.
                        # Right now this doens't matter because we only use the first plan.
                        break

                # TODO(Jiayuan Mao @ 2023/09/06): think about what's the semantics for this early return.
                # The idea is that, as soon as we found a plan using one of the regression rules, we return it, and don't worry about the other rules.
                # Therefore, it is not guaranteed that the returned plan is the shortest (optimal) one.
                if len(all_possible_plans) > 0:
                    break

            if len(all_possible_plans) > 0:
                break

        if len(all_possible_plans) == 0:
            if verbose:
                jacinle.log_function.print('No possible plans for goal set', goal_set)
            return return_with_cache(goal_set, previous_actions, [])

        shortest_plan = min(all_possible_plans, key=lambda x: len(x[2]))
        return return_with_cache(goal_set, previous_actions, [shortest_plan])

    if is_and_expr(goal_expr):
        goal_set = list(goal_expr.arguments)
    else:
        goal_set = [goal_expr]
    goal_set = PartiallyOrderedPlan((TotallyOrderedPlan(goal_set, is_ordered=True),))

    candidate_plans = dfs(state, goal_set, None, None, list())
    candidate_plans = [(actions, csp, state, final_state) for final_state, csp, actions in candidate_plans]
    return candidate_plans, search_stat


def gen_applicable_regression_rules_v2(executor: PDSketchExecutor, state: State, goals: PartiallyOrderedPlan, maintains: Set[ValueOutputExpression]):
    # TODO(Jiayuan Mao @ 2023/09/10): implement maintains.
    candidate_regression_rules = list()
    for chain_index, chain in goals.iter_feasible_chains():
        if chain.is_ordered:
            subgoal_indices = [len(chain) - 1]
        else:
            subgoal_indices = list(range(len(chain)))

        for subgoal_index in subgoal_indices:
            subgoal = chain.sequence[subgoal_index]
            this_chain_candidate_regression_rules = list()
            for regression_rule in executor.domain.regression_rules.values():
                goal_expr = regression_rule.goal_expression
                if (variable_binding := surface_fol_downcast(goal_expr, subgoal)) is None:
                    continue

                bounded_variables = dict()
                for v in regression_rule.goal_arguments:
                    bounded_variables[v] = variable_binding[v.name].name

                if len(regression_rule.binding_arguments) > 0:
                    for v in regression_rule.binding_arguments:
                        bounded_variables[v] = QINDEX

                if len(regression_rule.preconditions_conjunction.arguments) > 0:
                    if len(regression_rule.binding_arguments) > 4 and len(regression_rule.preconditions_conjunction.arguments) >= 2:
                        rv = extract_bounded_variables_from_nonzero_dc(executor, state, regression_rule, bounded_variables, use_optimistic=False)
                    else:
                        rv = executor.execute(regression_rule.preconditions_conjunction, state=state, bounded_variables=bounded_variables)
                        rv = extract_bounded_variables_from_nonzero(state, rv, regression_rule, default_bounded_variables=bounded_variables, use_optimistic=False)
                else:
                    rv = None

                if rv is None:
                    this_chain_candidate_regression_rules = [(regression_rule, bounded_variables)]
                    if regression_rule.always:
                        break
                else:
                    all_forall, candidate_bounded_variables = rv
                    if all_forall and regression_rule.always:
                        this_chain_candidate_regression_rules = [(regression_rule, candidate_bounded_variables[0])]
                        break
                    else:
                        this_chain_candidate_regression_rules.extend([(regression_rule, cbv) for cbv in candidate_bounded_variables])
            candidate_regression_rules.append((chain_index, subgoal_index, this_chain_candidate_regression_rules))
    return candidate_regression_rules

concepts/pdsketch/planners/optimistic_search_bilevel_legacy.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : optimistic_search_bilevel_utils.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/17/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Utility functions to conduct bi-level search for optimistic planning."""

import collections
import itertools
from dataclasses import dataclass
from typing import Any, Optional, Union, Iterable, Iterator, Sequence, Tuple, List, Set, Dict

import torch
import jacinle

from concepts.dsl.dsl_types import BOOL, Variable, ObjectConstant, UnnamedPlaceholder, QINDEX
from concepts.dsl.expression import ValueOutputExpression, FunctionApplicationExpression, AssignExpression, ConstantExpression, VariableExpression, ObjectConstantExpression, BoolExpression, QuantificationExpression
from concepts.dsl.expression import AndExpression, is_and_expr
from concepts.dsl.constraint import ConstraintSatisfactionProblem, EqualityConstraint, is_optimistic_value
from concepts.dsl.tensor_value import TensorValue
from concepts.dsl.executors.tensor_value_executor import BoundedVariablesDictCompatible
from concepts.pdsketch.executor import PDSketchExecutor
from concepts.pdsketch.domain import State
from concepts.pdsketch.operator import OperatorApplier

from concepts.dsl.constraint import OptimisticValue
from concepts.pdsketch.operator import (
    RegressionRule, AchieveExpression, OperatorApplicationExpression, RegressionRuleApplier,
    gen_all_partially_grounded_actions, gen_all_grounded_regression_rules,
    SubgoalSerializability, SubgoalContSerializability
)
from concepts.pdsketch.planners.optimistic_search import instantiate_action

__all__ = [
    'enumerate_possible_symbolic_plans', 'enumerate_possible_symbolic_plans_regression_1', 'enumerate_possible_symbolic_plans_regression_c_v3',
    'surface_fol_downcast', 'ground_fol_expression', 'ground_operator_application_expression', 'evaluate_bool_scalar_expression',
    'TotallyOrderedPlan', 'PartiallyOrderedPlan',
    'gen_applicable_regression_rules_v3', 'extract_bounded_variables_from_nonzero', 'extract_bounded_variables_from_nonzero_dc'
]


@dataclass
class OptimisticSearchSymbolicPlan(object):
    actions: Sequence[OperatorApplier]
    csp: ConstraintSatisfactionProblem
    initial_state: State
    last_state: State


def enumerate_possible_symbolic_plans(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression],
    max_actions: int = 4, primitive_action_only: bool = False, verbose: bool = False
) -> Iterable[OptimisticSearchSymbolicPlan]:
    csp = ConstraintSatisfactionProblem()
    if isinstance(goal_expr, str):
        goal_expr = executor.parse(goal_expr)
    assert isinstance(goal_expr, ValueOutputExpression)

    initial_state, state = state, state.clone()

    all_actions = gen_all_partially_grounded_actions(executor, state, action_filter=lambda x: x.operator.is_primitive)

    def dfs(state: State, actions: List[OperatorApplier], csp: ConstraintSatisfactionProblem, nr_high_level_actions: int = 0) -> Iterator[OptimisticSearchSymbolicPlan]:
        if nr_high_level_actions >= max_actions:
            return
        for a in all_actions:
            ncsp = csp.clone()
            a = instantiate_action(ncsp, a)
            succ, nstate = executor.apply(a, state, csp=ncsp, clone=True, action_index=len(actions))
            # print(f'Actions {actions} || {a} -> {succ}')
            if succ:
                ncsp_with_goal = ncsp.clone()
                nactions = actions + [a]
                rv = executor.execute(goal_expr, state=nstate, csp=ncsp_with_goal).item()
                if isinstance(rv, OptimisticValue):
                    ncsp_with_goal.add_constraint(EqualityConstraint.from_bool(rv, True), note='goal_test')
                    yield OptimisticSearchSymbolicPlan(nactions, ncsp_with_goal, initial_state, nstate)
                elif float(rv) >= 0.5:
                    yield OptimisticSearchSymbolicPlan(nactions, ncsp_with_goal, initial_state, nstate)

                yield from dfs(nstate, nactions, ncsp, nr_high_level_actions=nr_high_level_actions + 1)

    plans = list(dfs(state, [], csp))
    plans.sort(key=lambda x: len(x.actions))
    return plans


def enumerate_possible_symbolic_plans_regression_1(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression],
    max_actions: int = 4,
) -> Iterable[OptimisticSearchSymbolicPlan]:
    domain = executor.domain
    if isinstance(goal_expr, str):
        goal_expr = executor.parse(goal_expr)
    assert isinstance(goal_expr, ValueOutputExpression)

    assert isinstance(goal_expr, FunctionApplicationExpression), 'Only support a single function application as goal expression.'
    for regression_rule in domain.regression_rules.values():
        regression_rule: RegressionRule
        for action_item in regression_rule.body:
            if isinstance(action_item, AchieveExpression) and len(action_item.maintains) > 0:
                raise NotImplementedError('enumerate_possible_symbolic_plans_regression_1 does not support maintain expressions.')
        if len(regression_rule.preconstraints) > 0:
            raise NotImplementedError('enumerate_possible_symbolic_plans_regression_1 does not support preconstraints.')
        if len(regression_rule.goal) > 1:
            raise NotImplementedError('enumerate_possible_symbolic_plans_regression_1 does not support multiple goal expressions.')
        goal = regression_rule.goal[0]
        if not isinstance(goal.assign_expr, AssignExpression):
            raise NotImplementedError('enumerate_possible_symbolic_plans_regression_1 does not support non-assign goal expressions.')
        if not isinstance(goal.assign_expr.predicate, FunctionApplicationExpression) or goal.assign_expr.predicate.return_type != BOOL:
            raise NotImplementedError('enumerate_possible_symbolic_plans_regression_1 does not support non-boolean goal expressions.')
        if not isinstance(goal.assign_expr.value, ConstantExpression) or bool(goal.assign_expr.value.constant.item()) is not True:
            raise NotImplementedError('enumerate_possible_symbolic_plans_regression_1 does not support negations or non-constant goal expressions.')

    all_actions = gen_all_partially_grounded_actions(executor, state)
    all_regression_rules = gen_all_grounded_regression_rules(executor, state)

    def ground_object_function_application(fapp: FunctionApplicationExpression, variable_mapping: Dict[str, str]):
        arguments = list()
        for arg in fapp.arguments:
            if isinstance(arg, VariableExpression):
                arguments.append(ObjectConstantExpression(ObjectConstant(variable_mapping[arg.variable.name], arg.return_type)))
            elif isinstance(arg, ObjectConstantExpression):
                arguments.append(arg)
            else:
                raise TypeError(f'Unknown argument type: {type(arg)}')
        return FunctionApplicationExpression(fapp.function, arguments)

    def ground_object_operator_application(fapp: OperatorApplicationExpression, variable_mapping: Dict[str, str], csp: ConstraintSatisfactionProblem) -> OperatorApplier:
        arguments = list()
        for arg in fapp.arguments:
            if isinstance(arg, VariableExpression):
                arguments.append(variable_mapping[arg.variable.name])
            elif isinstance(arg, UnnamedPlaceholder):
                arguments.append(TensorValue.from_optimistic_value(csp.new_var(arg.dtype, wrap=True)))
            else:
                raise TypeError(f'Unknown argument type: {type(arg)}')
        return OperatorApplier(fapp.operator, arguments)

    def is_grounded_object_function_application_matching(expr1: FunctionApplicationExpression, expr2: FunctionApplicationExpression):
        if expr1.function != expr2.function:
            return False
        for arg1, arg2 in zip(expr1.arguments, expr2.arguments):
            if isinstance(arg1, ObjectConstantExpression) and isinstance(arg2, ObjectConstantExpression):
                if arg1.constant.name != arg2.constant.name:
                    return False
            else:
                return False
        return True

    def is_rule_applicable(state, regression_rule: RegressionRuleApplier, goal: FunctionApplicationExpression, csp: ConstraintSatisfactionProblem):
        goal_function = regression_rule.regression_rule.goal[0].assign_expr.predicate
        # TODO(Jiayuan Mao @ 2023/08/17): use "grounding" or "replace feature".
        variable_mapping = {variable.name: value for variable, value in zip(regression_rule.regression_rule.arguments, regression_rule.arguments)}
        grounded_goal_function = ground_object_function_application(goal_function, variable_mapping)
        if is_grounded_object_function_application_matching(goal, grounded_goal_function):
            for precondition in regression_rule.regression_rule.preconditions:
                rv = executor.execute(precondition.bool_expr, state=state, bounded_variables=variable_mapping, csp=csp).item()
                if isinstance(rv, OptimisticValue):
                    csp.add_constraint(EqualityConstraint.from_bool(rv, True), note='precondition')
                elif float(rv) < 0.5:
                    return False
            return True
        return False

    def dfs(state: State, goal: FunctionApplicationExpression, maintains: Set[FunctionApplicationExpression], csp: ConstraintSatisfactionProblem, previous_actions: List[OperatorApplier], nr_high_level_actions: int = 0):
        if nr_high_level_actions >= max_actions:
            return

        all_possible_plans = list()
        for rule in all_regression_rules:
            cur_csp = csp.clone()
            if not is_rule_applicable(state, rule, goal, cur_csp):
                continue

            cur_state = state.clone()
            possible_branches = [(cur_state, cur_csp, previous_actions)]
            variable_mapping = {variable.name: value for variable, value in zip(rule.regression_rule.arguments, rule.arguments)}
            for item in rule.regression_rule.body:
                next_possible_branches = list()
                for cur_state, cur_csp, actions in possible_branches:
                    if isinstance(item, AchieveExpression):
                        cur_csp_2 = cur_csp.clone()
                        rv = executor.execute(item.goal, state=cur_state, bounded_variables=variable_mapping, csp=cur_csp_2).item()
                        subgoal = ground_object_function_application(item.goal, variable_mapping)
                        if isinstance(rv, OptimisticValue):
                            cur_csp_2.add_constraint(EqualityConstraint.from_bool(rv, True), note='subgoal_test')
                            next_possible_branches.append((cur_state, cur_csp_2, actions))

                            next_possible_branches.extend(dfs(cur_state, subgoal, maintains, cur_csp, actions, nr_high_level_actions=nr_high_level_actions))
                        elif float(rv) > 0.5:
                            next_possible_branches.append((cur_state, cur_csp_2, actions))
                        else:
                            next_possible_branches.extend(dfs(cur_state, subgoal, maintains, cur_csp, actions, nr_high_level_actions=nr_high_level_actions + 1))
                    elif isinstance(item, OperatorApplicationExpression):
                        cur_csp_2 = cur_csp.clone()
                        cur_action = ground_object_operator_application(item, variable_mapping, cur_csp_2)
                        succ, cur_state = executor.apply(cur_action, cur_state, csp=cur_csp_2, clone=True, action_index=len(actions))
                        if succ:
                            next_possible_branches.append((cur_state, cur_csp_2, actions + [cur_action]))
                possible_branches = next_possible_branches

            all_possible_plans.extend(possible_branches)

        # TODO(Jiayuan Mao @ 2023/08/18): add a flag to control this behavior.
        if len(all_possible_plans) == 0:
            for action in all_actions:
                cur_csp = csp.clone()
                cur_action = instantiate_action(cur_csp, action)
                succ, cur_state = executor.apply(cur_action, state, csp=cur_csp, clone=True, action_index=len(previous_actions))
                if succ:
                    cur_csp_2 = cur_csp.clone()
                    rv = executor.execute(goal, state=cur_state, csp=cur_csp_2).item()
                    if isinstance(rv, OptimisticValue):
                        cur_csp_2.add_constraint(EqualityConstraint.from_bool(rv, True), note='goal_test')
                        all_possible_plans.append((cur_state, cur_csp_2, previous_actions + [cur_action]))
                        all_possible_plans.extend(dfs(cur_state, goal, maintains, cur_csp, previous_actions + [cur_action], nr_high_level_actions=nr_high_level_actions + 1))
                    elif float(rv) > 0.5:
                        all_possible_plans.append((cur_state, cur_csp_2, previous_actions + [cur_action]))
                    else:
                        all_possible_plans.extend(dfs(cur_state, goal, maintains, cur_csp, previous_actions + [cur_action], nr_high_level_actions=nr_high_level_actions + 1))

        if len(all_possible_plans) == 0:
            return []

        # TODO(Jiayuan Mao @ 2023/08/18): add another flag to control this argmin behavior.
        shortest_plan = min(all_possible_plans, key=lambda x: len(x[2]))
        return [shortest_plan]

    initial_state = state
    csp = ConstraintSatisfactionProblem()
    candidate_plans = dfs(state, goal_expr, set(), csp, list())
    candidate_plans = [OptimisticSearchSymbolicPlan(actions, csp, initial_state, state) for state, csp, actions in candidate_plans]
    candidate_plans.sort(key=lambda x: len(x.actions))
    return candidate_plans


def enumerate_possible_symbolic_plans_regression_c_v3(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression], *,
    is_goal_serialized: bool = True,
    max_actions: int = 10, enable_csp: bool = False, verbose: bool = False
) -> Tuple[Iterable[OptimisticSearchSymbolicPlan], Dict[str, Any]]:
    """Enumerate all possible plans that can achieve the goal.

    Args:
        executor: the executor.
        state: the initial state.
        goal_expr: the goal expression.
        max_actions: the maximum number of actions in a plan.
        verbose: whether to print verbose information.

    Returns:
        A list of plans. Each plan is a tuple of (actions, csp, initial_state, final_state).
    """
    if isinstance(goal_expr, str):
        goal_expr = executor.parse(goal_expr)

    search_cache = dict()
    search_stat = {'nr_expanded_nodes': 0}

    # NB(Jiayuan Mao @ 2023/09/09): the cache only works for previous_actions == [].
    # That is, we only cache the search results that starts from the initial state.
    def return_with_cache(goal_set, previous_actions, rv):
        if len(previous_actions) == 0:
            goal_str = goal_set.gen_string()
            if goal_str not in search_cache:
                search_cache[goal_str] = rv
        return rv

    def try_retrive_cache(goal_set, previous_actions):
        if len(previous_actions) == 0:
            goal_str = goal_set.gen_string()
            if goal_str in search_cache:
                return search_cache[goal_str]
        return None

    @jacinle.log_function(verbose=False)
    def dfs(
        state: State, goal_set: PartiallyOrderedPlan, maintains: List[ValueOutputExpression],
        csp: Optional[ConstraintSatisfactionProblem], previous_actions: List[OperatorApplier],
        return_all: bool = False,
        nr_high_level_actions: int = 0
    ) -> Iterator[Tuple[State, ConstraintSatisfactionProblem, List[OperatorApplier]]]:
        """Depth-first search for all possible plans.

        Args:
            state: the current state.
            goal_set: the goal set.
            maintains: the list of maintains.
            csp: the current constraint satisfaction problem.
            previous_actions: the previous actions.
            return_all: whether to return all possible plans. If False, only return first plan found.
            nr_high_level_actions: the number of high-level actions.

        Returns:
            a list of plans. Each plan is a tuple of (final_state, csp, actions).
        """

        if verbose:
            jacinle.log_function.print('Current goal_set', goal_set, f'return_all={return_all}')
        #     jacinle.log_function.print('Previous actions', previous_actions)

        if (rv := try_retrive_cache(goal_set, previous_actions)) is not None:
            return rv

        all_possible_plans = list()

        flatten_goals = list(goal_set.iter_goals())
        rv, is_optimistic, new_csp = evaluate_bool_scalar_expression(executor, flatten_goals, state, dict(), csp, csp_note='goal_test')
        if rv:
            all_possible_plans.append((state, new_csp, previous_actions))
            if not is_optimistic:  # If there is no optimistic value, we can stop the search from here.
                # NB(Jiayuan Mao @ 2023/09/11): note that even if `return_all` is True, we still return here.
                # This corresponds to an early stopping behavior that defines the space of all possible plans.
                return return_with_cache(goal_set, previous_actions, all_possible_plans)

        if nr_high_level_actions > max_actions:
            return return_with_cache(goal_set, previous_actions, all_possible_plans)

        search_stat['nr_expanded_nodes'] += 1

        candidate_regression_rules = gen_applicable_regression_rules_v3(executor, state, goal_set, maintains, csp=csp)
        # if verbose:
        #     jacinle.log_function.print('Candidate regression rules', candidate_regression_rules)

        if sum(len(r) for _, _, r in candidate_regression_rules) == 0:
            return return_with_cache(goal_set, previous_actions, all_possible_plans)

        for chain_index, subgoal_index, this_candidate_regression_rules in candidate_regression_rules:
            other_goals = goal_set.exclude(chain_index, subgoal_index)
            cur_goal = goal_set.chains[chain_index].sequence[subgoal_index]
            if verbose:
                jacinle.log_function.print('Now trying to excluding goal', cur_goal)

            if len(other_goals) == 0:
                other_goals_plans = [(state, csp, previous_actions)]
            else:
                # TODO(Jiayuan Mao @ 2023/09/09): change this list to an actual generator call.
                need_return_all = any(rule.max_rule_prefix_length > 0 for rule, _ in this_candidate_regression_rules)
                other_goals_plans = list(dfs(state, other_goals, maintains, csp, previous_actions, nr_high_level_actions=nr_high_level_actions, return_all=need_return_all))
                for cur_state, cur_csp, cur_actions in other_goals_plans:
                    rv, is_optimistic, new_csp = evaluate_bool_scalar_expression(executor, [cur_goal], cur_state, dict(), cur_csp, csp_note='goal_test')
                    if rv:
                        all_possible_plans.append((cur_state, new_csp, cur_actions))
                        if not is_optimistic:
                            # NB(Jiayuan Mao @ 2023/09/11): another place where we stop the search and ignores the `return_all` flag.
                            continue

            if len(this_candidate_regression_rules) == 0:
                continue

            if len(other_goals_plans) == 0:
                continue

            if len(other_goals) == 0:
                max_prefix_length = 0
            else:
                max_prefix_length = max(rule.max_reorder_prefix_length for rule, _ in this_candidate_regression_rules)

            grounded_subgoals_cache = dict()
            prefix_stop_mark = dict()

            for prefix_length in range(max_prefix_length + 1):
                for regression_rule_index, (rule, bounded_variables) in enumerate(this_candidate_regression_rules):
                    if csp is not None:
                        regression_rule_rv, bounded_variables = bounded_variables
                    else:
                        regression_rule_rv = 1
                    if prefix_length > rule.max_reorder_prefix_length:
                        continue
                    if regression_rule_index in prefix_stop_mark and prefix_stop_mark[regression_rule_index]:
                        continue

                    if verbose:
                        jacinle.log_function.print('Applying rule', rule, 'for', cur_goal, 'and prefix length', prefix_length, 'goal set is', goal_set)

                    # construct grounded_subgoals
                    if regression_rule_index in grounded_subgoals_cache:
                        grounded_subgoals = grounded_subgoals_cache[regression_rule_index]
                    else:
                        grounded_subgoals = list()
                        for item in rule.body:
                            if isinstance(item, AchieveExpression):
                                grounded_subgoals.append(ground_fol_expression(item.goal, bounded_variables))
                        grounded_subgoals_cache[regression_rule_index] = grounded_subgoals

                    if prefix_length == 0:
                        if rule.max_rule_prefix_length > 0:
                            possible_branches = other_goals_plans
                        else:
                            possible_branches = [other_goals_plans[0]]
                    else:
                        cur_other_goals = other_goals.add_chain(grounded_subgoals[:prefix_length])
                        possible_branches = list(dfs(state, cur_other_goals, maintains, csp, previous_actions, nr_high_level_actions=nr_high_level_actions + 1, return_all=rule.max_rule_prefix_length > 0))

                    if len(possible_branches) == 0:
                        if verbose:
                            jacinle.log_function.print('Prefix planning failed!!! Stop.')
                        # If it's not possible to achieve the subset of goals, then it's not possible to achieve the whole goal set.
                        # Therefore, this is a break, not a continue.
                        prefix_stop_mark[regression_rule_index] = True
                        continue

                    for i in range(prefix_length, len(rule.body)):
                        item = rule.body[i]
                        next_possible_branches = list()

                        if isinstance(item, AchieveExpression):
                            if item.serializability is SubgoalSerializability.STRONG and len(possible_branches) > 1:
                                possible_branches = [min(possible_branches, key=lambda x: len(x[2]))]
                        need_return_all = i < rule.max_rule_prefix_length

                        for cur_state, cur_csp, actions in possible_branches:
                            if isinstance(item, AchieveExpression):
                                subgoal = grounded_subgoals[i]
                                next_possible_branches.extend(dfs(cur_state, PartiallyOrderedPlan.from_single_goal(subgoal), maintains, cur_csp, actions, nr_high_level_actions=nr_high_level_actions + 1))
                            elif isinstance(item, OperatorApplicationExpression):
                                # TODO(Jiayuan Mao @ 2023/09/11): vectorize this operation, probably only useful when `return_all` is True.
                                new_csp = cur_csp.clone() if cur_csp is not None else None
                                cur_action = ground_operator_application_expression(item, bounded_variables, csp=new_csp)
                                succ, new_state = executor.apply(cur_action, cur_state, csp=new_csp, clone=True, action_index=len(actions))
                                if succ:
                                    next_possible_branches.append((new_state, new_csp, actions + [cur_action]))
                            else:
                                raise TypeError()
                        possible_branches = next_possible_branches

                    # all_possible_plans.extend(possible_branches)

                    found_plan = False
                    # TODO(Jiayuan Mao @ 2023/09/11): implement this via maintains checking.
                    for cur_state, cur_csp, actions in possible_branches:
                        rv, is_optimistic, new_csp = evaluate_bool_scalar_expression(executor, flatten_goals, cur_state, dict(), csp=cur_csp, csp_note='goal_test')
                        if rv:
                            if verbose:
                                jacinle.log_function.print('Found a plan', [str(x) for x in actions], 'for goal set', goal_set)
                            all_possible_plans.append((cur_state, new_csp, actions))
                            found_plan = True

                    if found_plan:
                        prefix_stop_mark[regression_rule_index] = True
                        # TODO(Jiayuan Mao @ 2023/09/06): since we have changed the order of prefix_length for-loop and the regression rule for-loop.
                        # We need to use an additional dictionary to store whether we have found a plan for a particular regression rule.
                        # Right now this doens't matter because we only use the first plan.

                    if not return_all and len(all_possible_plans) > 0:
                        break

                if not return_all and len(all_possible_plans) > 0:
                    break

            if not return_all and len(all_possible_plans) > 0:
                break

        if len(all_possible_plans) == 0:
            if verbose:
                jacinle.log_function.print('No possible plans for goal set', goal_set)
            return return_with_cache(goal_set, previous_actions, [])

        unique_all_possible_plans = _unique_plans(all_possible_plans)
        if len(unique_all_possible_plans) != len(all_possible_plans):
            if verbose:
                jacinle.log_function.print('Warning: there are duplicate plans for goal set', goal_set, f'({len(unique_all_possible_plans)} unique plans vs {len(all_possible_plans)} total plans)')
                # import ipdb; ipdb.set_trace()

        unique_all_possible_plans = sorted(unique_all_possible_plans, key=lambda x: len(x[2]))
        return return_with_cache(goal_set, previous_actions, unique_all_possible_plans)

    if is_and_expr(goal_expr):
        goal_set = list(goal_expr.arguments)
    else:
        goal_set = [goal_expr]
    goal_set = PartiallyOrderedPlan((TotallyOrderedPlan(goal_set, is_ordered=is_goal_serialized),))

    if enable_csp:
        csp = ConstraintSatisfactionProblem()
    else:
        csp = None

    candidate_plans = dfs(state, goal_set, maintains=None, csp=csp, previous_actions=list(), return_all=enable_csp)
    candidate_plans = [OptimisticSearchSymbolicPlan(actions, csp, state, final_state) for final_state, csp, actions in candidate_plans]
    return candidate_plans, search_stat


def surface_fol_downcast(expression_1: ValueOutputExpression, expression_2: ValueOutputExpression) -> Optional[Dict[Variable, Union[Variable, ObjectConstant]]]:
    """Trying to downcast the `expression_1` to the same form as `expression_2`. Downcasting means that
    we try to replace variables in `expression_1` with constants in `expression_2` to make them the same.

    Args:
        expression_1: the first expression.
        expression_2: the second expression.

    Returns:
        the downcasted mapping if the downcasting is successful, otherwise None.
    """
    current_mapping = dict()

    def dfs(expr1, expr2):
        nonlocal current_mapping

        if isinstance(expr1, VariableExpression):
            if expr1.name in current_mapping:
                expr2_name = _get_variable_or_constant_name(expr2)
                return expr2_name == current_mapping[expr1.name].name
            else:
                if isinstance(expr2, VariableExpression):
                    current_mapping[expr1.name] = expr2.variable
                    return True
                elif isinstance(expr2, ObjectConstantExpression):
                    current_mapping[expr1.name] = expr2.constant
                    return True
                else:
                    return False
        elif isinstance(expr1, ObjectConstantExpression):
            if isinstance(expr2, ObjectConstantExpression):
                return expr1.name == expr2.name
            else:
                return False
        elif isinstance(expr1, FunctionApplicationExpression):
            if not isinstance(expr2, FunctionApplicationExpression):
                return False
            if expr1.function.name != expr2.function.name:
                return False
            if len(expr1.arguments) != len(expr2.arguments):
                return False
            for arg1, arg2 in zip(expr1.arguments, expr2.arguments):
                if not dfs(arg1, arg2):
                    return False
            return True
        elif isinstance(expr1, BoolExpression):
            if not isinstance(expr2, BoolExpression):
                return False
            if expr1.bool_op != expr2.bool_op:
                return False
            if len(expr1.arguments) != len(expr2.arguments):
                return False
            for arg1, arg2 in zip(expr1.arguments, expr2.arguments):
                if not dfs(arg1, arg2):
                    return False
            return True
        elif isinstance(expr1, QuantificationExpression):
            if not isinstance(expr2, QuantificationExpression):
                return False
            if expr1.quantification_op != expr2.quantification_op:
                return False
            assert expr1.variable.name not in current_mapping
            current_mapping[expr1.variable.name] = expr2.variable
            try:
                return dfs(expr1.expression, expr2.expression)
            finally:
                del current_mapping[expr1.variable.name]
        else:
            raise TypeError(f'Unsupported expression type: {type(expr1)}')

    rv = dfs(expression_1, expression_2)
    if rv:
        return current_mapping
    return None


def ground_fol_expression(expression: ValueOutputExpression, variable_mapping: Dict[Variable, str]) -> ValueOutputExpression:
    """Ground the given FOL expression with the given variable mapping.

    Args:
        expression: the expression to ground.
        variable_mapping: the variable mapping, which is a mapping from the Variable object to the constant name.

    Returns:
        the grounded expression.
    """

    name2symbol = {var.name: ObjectConstantExpression(ObjectConstant(name, var.dtype)) for var, name in variable_mapping.items()}
    bounded_variables = set()

    def dfs(e):
        if isinstance(e, VariableExpression):
            if e.name not in bounded_variables:
                return name2symbol[e.name]
            return e
        elif isinstance(e, ObjectConstantExpression):
            return e
        elif isinstance(e, FunctionApplicationExpression):
            return FunctionApplicationExpression(e.function, [dfs(arg) for arg in e.arguments])
        elif isinstance(e, BoolExpression):
            return BoolExpression(e.bool_op, [dfs(arg) for arg in e.arguments])
        elif isinstance(e, QuantificationExpression):
            bounded_variables.add(e.variable.name)
            rv = QuantificationExpression(e.quantification_op, e.variable, dfs(e.expression))
            bounded_variables.remove(e.variable.name)
            return rv
        else:
            raise TypeError(f'Unsupported expression type: {type(e)}')

    return dfs(expression)


def _get_variable_or_constant_name(expr: Union[VariableExpression, ObjectConstantExpression]) -> str:
    """Get the name of the given variable or constant expression."""
    if isinstance(expr, VariableExpression):
        return expr.name
    elif isinstance(expr, ObjectConstantExpression):
        return expr.name
    else:
        raise TypeError(f'Unsupported type: {type(expr)} for _get_variable_or_constant_name.')


def ground_operator_application_expression(expression: OperatorApplicationExpression, variable_mapping: Dict[Variable, str], csp: Optional[ConstraintSatisfactionProblem] = None) -> OperatorApplicationExpression:
    """Ground the given operator application expression with the given variable mapping.

    Args:
        expression: the expression to ground.
        variable_mapping: the variable mapping, which is a mapping from the Variable object to the constant name.
        csp: the constraint satisfaction problem to add the constraints to.

    Returns:
        the grounded expression.
    """
    name2symbol = {var.name: name for var, name in variable_mapping.items()}
    arguments = list()
    for arg in expression.arguments:
        if isinstance(arg, VariableExpression):
            arguments.append(name2symbol[arg.variable.name])
        elif isinstance(arg, UnnamedPlaceholder):
            assert csp is not None
            arguments.append(TensorValue.from_optimistic_value(csp.new_var(arg.dtype, wrap=True)))
        elif isinstance(arg, str):
            arguments.append(arg)
        else:
            raise TypeError(f'Unknown argument type: {type(arg)}')
    return OperatorApplier(expression.operator, arguments)


def evaluate_bool_scalar_expression(
    executor: PDSketchExecutor, expr: Union[ValueOutputExpression, Sequence[ValueOutputExpression]], state: State,
    bounded_variables: BoundedVariablesDictCompatible, csp: ConstraintSatisfactionProblem, csp_note: str = ''
) -> Tuple[bool, bool, Optional[ConstraintSatisfactionProblem]]:
    if csp is not None:
        csp = csp.clone()
    is_optimistic = False
    if not isinstance(expr, (list, tuple)):
        expr = [expr]
    for e in expr:
        rv = executor.execute(e, state=state, bounded_variables=bounded_variables, csp=csp).item()
        if isinstance(rv, OptimisticValue):
            csp.add_constraint(EqualityConstraint.from_bool(rv, True), note=csp_note)
            is_optimistic = True
        elif float(rv) < 0.5:
            return False, is_optimistic, csp
    return True, is_optimistic, csp


class TotallyOrderedPlan(object):
    """A totally ordered plan sequence."""

    def __init__(self, sequence: Sequence[ValueOutputExpression], is_ordered: bool = True):
        """Initialize the totally ordered plan sequence.

        Args:
            sequence: the sequence of the plan.
            is_ordered: whether the sequence is ordered.
        """
        self.sequence = tuple(sequence)
        self.is_ordered = is_ordered

    sequence: Tuple[ValueOutputExpression, ...]
    """The sequence of the plan."""

    is_ordered: bool
    """Whether the sequence is ordered."""

    def exclude(self, index: int):
        """Exclude the given index from the sequence; return a new sequence."""
        return TotallyOrderedPlan(self.sequence[:index] + self.sequence[index + 1:], self.is_ordered)

    def gen_string(self):
        """Generate the string representation of the plan."""
        if self.is_ordered:
            return '(then ' + ' '.join([str(e) for e in self.sequence]) + ')'
        else:
            return '(and ' + ' '.join([str(e) for e in self.sequence]) + ')'

    def __len__(self):
        return len(self.sequence)

    def __str__(self):
        return self.gen_string()

    __repr__ = jacinle.repr_from_str


class PartiallyOrderedPlan(object):
    """A partially ordered plan is a set of totally ordered plan sequences."""

    def __init__(self, chains: Sequence[TotallyOrderedPlan]):
        """Initialize the partially ordered plan.

        Args:
            chains: a collection of the totally ordered plan sequences.
        """
        self.chains = tuple(chains)
        self.infeasible_chain_inedx = None

    chains: Tuple[TotallyOrderedPlan, ...]
    """The totally ordered plan sequences."""

    infeasible_index: Optional[int]
    """One can optionally mark one of the chain as infeasible to be the last chain. Currently this flag is not used."""

    def exclude(self, chain_index: int, item_index: int) -> 'PartiallyOrderedPlan':
        """Exclude the given item from the given chain; return a new plan."""
        if len(self.chains[chain_index]) == 1:
            assert item_index == 0
            return PartiallyOrderedPlan(self.chains[:chain_index] + self.chains[chain_index + 1:])
        return PartiallyOrderedPlan(self.chains[:chain_index] + (self.chains[chain_index].exclude(item_index),) + self.chains[chain_index + 1:])

    def add_chain(self, chain: Sequence[ValueOutputExpression]) -> 'PartiallyOrderedPlan':
        """Add a new chain to the plan; return a new plan."""
        plan = PartiallyOrderedPlan(self.chains + (TotallyOrderedPlan(chain),))
        return plan

    @classmethod
    def from_single_goal(cls, goal: ValueOutputExpression):
        """Create a plan from a single goal."""
        return cls((TotallyOrderedPlan((goal,), True),))

    @property
    def nr_chains(self):
        return len(self.chains)

    @property
    def total_length(self):
        return sum(len(chain) for chain in self.chains)

    def set_infeasible_index(self, chain_index: int):
        self.infeasible_chain_inedx = chain_index

    def iter_feasible_chains(self) -> Iterator[Tuple[int, TotallyOrderedPlan]]:
        for i, seq in enumerate(self.chains):
            if i == self.infeasible_chain_inedx:
                continue
            yield i, seq

    def iter_goals(self) -> Iterator[ValueOutputExpression]:
        for chain in self.chains:
            for goal in chain.sequence:
                yield goal

    def gen_string(self):
        return '(and ' + ' '.join([chain.gen_string() for chain in self.chains]) + ')'

    def __len__(self):
        return self.total_length

    def __str__(self):
        return self.gen_string()

    __repr__ = jacinle.repr_from_str


def gen_applicable_regression_rules_v3(
    executor: PDSketchExecutor, state: State, goals: PartiallyOrderedPlan,
    maintains: Set[ValueOutputExpression], csp: Optional[ConstraintSatisfactionProblem] = None
) -> List[Tuple[int, int, List[Tuple[RegressionRule, BoundedVariablesDictCompatible]]]]:
    # TODO(Jiayuan Mao @ 2023/09/10): implement maintains.
    candidate_regression_rules = list()
    for chain_index, chain in goals.iter_feasible_chains():
        if chain.is_ordered:
            subgoal_indices = [len(chain) - 1]
        else:
            subgoal_indices = list(range(len(chain)))

        for subgoal_index in subgoal_indices:
            subgoal = chain.sequence[subgoal_index]
            this_chain_candidate_regression_rules = list()

            # For each subgoal in the goal_set, try to find a list of applicable regression rules.
            # If one of the regression rules is always applicable, then we can stop searching.
            for regression_rule in executor.domain.regression_rules.values():
                goal_expr = regression_rule.goal_expression
                if (variable_binding := surface_fol_downcast(goal_expr, subgoal)) is None:
                    continue

                bounded_variables = dict()
                for v in regression_rule.goal_arguments:
                    bounded_variables[v] = variable_binding[v.name].name

                if len(regression_rule.binding_arguments) > 0:
                    for v in regression_rule.binding_arguments:
                        bounded_variables[v] = QINDEX

                if len(regression_rule.preconditions_conjunction.arguments) > 0:
                    if len(regression_rule.binding_arguments) > 4 and len(regression_rule.preconditions_conjunction.arguments) >= 2:
                        rv = extract_bounded_variables_from_nonzero_dc(executor, state, regression_rule, bounded_variables, use_optimistic=csp is not None)
                    else:
                        if csp is not None:
                            new_csp = csp.clone()
                            rv = executor.execute(regression_rule.preconditions_conjunction, state=state, bounded_variables=bounded_variables, csp=new_csp, optimistic_execution=True)
                            rv = extract_bounded_variables_from_nonzero(state, rv, regression_rule, default_bounded_variables=bounded_variables, use_optimistic=True)
                        else:
                            rv = executor.execute(regression_rule.preconditions_conjunction, state=state, bounded_variables=bounded_variables)
                            rv = extract_bounded_variables_from_nonzero(state, rv, regression_rule, default_bounded_variables=bounded_variables, use_optimistic=False)
                else:
                    rv = None

                if rv is None:
                    type_binding_arguments = [state.object_type2name[v.dtype.typename] for v in regression_rule.binding_arguments]
                    all_forall = True
                    for i, arg in enumerate(regression_rule.binding_arguments):
                        if arg.quantifier_flag == 'forall':
                            type_binding_arguments[i] = type_binding_arguments[i][:1] if len(type_binding_arguments[i]) > 0 else []
                        else:
                            all_forall = False
                    for binding_arguments in itertools.product(*type_binding_arguments):
                        cbv = bounded_variables.copy()
                        for variable, value in zip(regression_rule.binding_arguments, binding_arguments):
                            cbv[variable] = value
                        if all_forall and regression_rule.always:
                            this_chain_candidate_regression_rules = [(regression_rule, cbv)] if csp is None else [(regression_rule, (1, cbv))]
                            break
                        this_chain_candidate_regression_rules.append((regression_rule, cbv) if csp is None else (regression_rule, (1, cbv)))
                else:
                    all_forall, candidate_bounded_variables = rv
                    if all_forall and regression_rule.always:
                        this_chain_candidate_regression_rules = [(regression_rule, candidate_bounded_variables[0])]
                        break
                    else:
                        this_chain_candidate_regression_rules.extend([(regression_rule, cbv) for cbv in candidate_bounded_variables])
            candidate_regression_rules.append((chain_index, subgoal_index, this_chain_candidate_regression_rules))
    return candidate_regression_rules


def extract_bounded_variables_from_nonzero(
    state: State, value: TensorValue, regression_rule: RegressionRule,
    default_bounded_variables: Dict[Variable, Union[str, slice]],
    use_optimistic: bool = False
) -> Union[Tuple[bool, List[Dict[Variable, str]]], Tuple[bool, List[Tuple[int, Dict[Variable, str]]]]]:
    """Extract the indices of the given tensor value. The return value is a list of bounded variable dictionaries."""

    if value.tensor_optimistic_values is not None:
        assert use_optimistic, "Optimistic values are not allowed."
        possibly_true_values = torch.where(is_optimistic_value(value.tensor_optimistic_values), value.tensor_optimistic_values, (value.tensor > 0.5))
    else:
        possibly_true_values = (value.tensor > 0.5)

    indices = torch.nonzero(possibly_true_values, as_tuple=True)
    if len(indices[0]) == 0:
        return False, []

    variable_name2var = {v.name: v for v in regression_rule.arguments}
    tensor_variables = value.batch_variables
    exists_indices = [tensor_variables.index(v.name) for v in regression_rule.binding_arguments if v.quantifier_flag == 'exists']
    all_forall = len(exists_indices) == 0

    if not use_optimistic:
        indices = list(zip(*[i.tolist() for i in indices]))
        if possibly_true_values.ndim == 0:
            value = possibly_true_values.item()
            unique_mapping = {tuple(): tuple()}
        else:
            if not all_forall:
                unique_mapping = dict()
                for index in indices:
                    exists_indices = [index[i] for i in exists_indices]
                    unique_mapping[tuple(exists_indices)] = index
            else:
                unique_mapping = {tuple(): indices[0]}

        candidate_bounded_variables = list()
        for index in unique_mapping.values():
            bounded_variables = default_bounded_variables.copy()
            for i, v in enumerate(index):
                bounded_variables[variable_name2var[tensor_variables[i]]] = state.object_type2name[variable_name2var[tensor_variables[i]].dtype.typename][v]
            if all_forall:
                return True, [bounded_variables]
            candidate_bounded_variables.append(bounded_variables)
        return False, candidate_bounded_variables
    else:
        if possibly_true_values.ndim == 0:
            value = possibly_true_values.item()
            unique_mapping = {tuple(): (tuple(), value)}
        else:
            values = possibly_true_values[indices]
            indices = list(zip(*[i.tolist() for i in indices]))

            unique_mapping = collections.defaultdict(list)
            found_determined_solution = dict()
            for index, value in zip(indices, values.tolist()):
                exists_indices = tuple([index[i] for i in exists_indices])
                if value == 1:
                    if not found_determined_solution.get(exists_indices, False):
                        found_determined_solution[exists_indices] = True
                        unique_mapping[exists_indices] = [(index, value)]
                else:
                    unique_mapping[exists_indices].append((index, value))

        candidate_bounded_variables = list()
        for exists_indices, indices_and_values in unique_mapping.items():
            for index, value in indices_and_values:
                bounded_variables = default_bounded_variables.copy()
                for i, v in enumerate(index):
                    bounded_variables[variable_name2var[tensor_variables[i]]] = state.object_type2name[variable_name2var[tensor_variables[i]].dtype.typename][v]
                if all_forall and value == 1:
                    return True, [(value, bounded_variables)]
                candidate_bounded_variables.append((value, bounded_variables))
        return False, candidate_bounded_variables


def extract_bounded_variables_from_nonzero_dc(
    executor: PDSketchExecutor, state: State, regression_rule: RegressionRule,
        default_bounded_variables: Dict[Variable, Union[str, slice]], use_optimistic: bool = False
) -> Union[Tuple[bool, List[Dict[Variable, str]]], Tuple[bool, List[Tuple[bool, Dict[Variable, str]]], ConstraintSatisfactionProblem]]:
    """Extract bounded variables from nonzero indices in the preconditions of a regression rule.
    The function uses divide-and-conquer to find the bounded variables so it works for regression rules with
    a large number of variabels."""

    if use_optimistic:
        raise NotImplementedError('Optimistic mode is not implemented for extract_bounded_variables_from_nonzero_dc.')

    # NB(Jiayuan Mao @ 2023/09/10): the commented-out code is part of an implementation that performs divide-and-conquer
    # on variables. The current implementation does divide-and-conquer on the precondition list instead.
    # The new version is a bit easier to implement and easier to parallelize.

    # total_binding_variables = len(regression_rule.binding_arguments)

    # group1 = regression_rule.binding_arguments[:total_binding_variables // 2]
    # group2 = regression_rule.binding_arguments[total_binding_variables // 2:]
    # group1_names = {v.name for v in group1}
    # group2_names = {v.name for v in group2}

    # expression_variables = [find_free_variables(e) for e in regression_rule.preconditions_conjunction.arguments]
    # expression_variable_names = [set(v.name for v in vs) for vs in expression_variables]

    # group1_expressions, group2_expressions, cross_expressions = list(), list(), list()
    # for e, variables in zip(regression_rule.preconditions_conjunction.arguments, expression_variable_names):
    #     k1 = len(group1_names.intersection(variables)) > 0
    #     k2 = len(group2_names.intersection(variables)) > 0
    #     if k1 and k2:
    #         cross_expressions.append(e)
    #     elif k1:
    #         group1_expressions.append(e)
    #     elif k2:
    #         group2_expressions.append(e)
    #     else:
    #         group1_expressions.append(e)
    #         group2_expressions.append(e)

    # group1_expression = AndExpression(*group1_expressions) if len(group1_expressions) > 0 else None
    # group2_expression = AndExpression(*group2_expressions) if len(group2_expressions) > 0 else None
    # cross_expression = AndExpression(*cross_expressions) if len(cross_expressions) > 0 else None

    # rv1 = executor.execute(group1_expression, state=state, bounded_variables=bounded_variables) if group1_expression is not None else None
    # rv2 = executor.execute(group2_expression, state=state, bounded_variables=bounded_variables) if group2_expression is not None else None

    # if rv1 is None or rv2 is None:
    #     raise NotImplementedError('not implemented yet.')

    total_terms = len(regression_rule.preconditions_conjunction.arguments)
    group1_expressions = regression_rule.preconditions_conjunction.arguments[:total_terms // 2]
    group2_expressions = regression_rule.preconditions_conjunction.arguments[total_terms // 2:]

    group1_expression = AndExpression(*group1_expressions)
    group2_expression = AndExpression(*group2_expressions)
    rv1 = executor.execute(group1_expression, state=state, bounded_variables=default_bounded_variables)
    rv2 = executor.execute(group2_expression, state=state, bounded_variables=default_bounded_variables)

    shared_variables = list(set(rv1.batch_variables).intersection(rv2.batch_variables))
    shared_var_indices_in_group1 = [rv1.batch_variables.index(v) for v in shared_variables]
    shared_var_indices_in_group1_neg = [i for i in range(len(rv1.batch_variables)) if i not in shared_var_indices_in_group1]
    shared_var_indices_in_group2 = [rv2.batch_variables.index(v) for v in shared_variables]
    shared_var_indices_in_group2_neg = [i for i in range(len(rv2.batch_variables)) if i not in shared_var_indices_in_group2]

    rv1_indices = collections.defaultdict(list)
    for x in rv1.tensor.nonzero().tolist():
        rv1_indices[tuple(x[i] for i in shared_var_indices_in_group1)].append(tuple(x[i] for i in shared_var_indices_in_group1_neg))
    rv2_indices = collections.defaultdict(list)
    for x in rv2.tensor.nonzero().tolist():
        rv2_indices[tuple(x[i] for i in shared_var_indices_in_group2)].append(tuple(x[i] for i in shared_var_indices_in_group2_neg))

    output_variables = shared_variables + [v for v in rv1.batch_variables if v not in shared_variables] + [v for v in rv2.batch_variables if v not in shared_variables]
    variable_name2var = {v.name: v for v in regression_rule.arguments}
    exists_indices = [output_variables.index(v.name) for v in regression_rule.binding_arguments if v.quantifier_flag == 'exists']
    all_forall = len(exists_indices) == 0

    if len(rv1_indices) == 0 or len(rv2_indices) == 0:
        return False, []

    indices = list()
    for k in set(rv1_indices.keys()).intersection(rv2_indices.keys()):
        indices.extend([k + x + y for x in rv1_indices[k] for y in rv2_indices[k]])

    if len(indices) == 0:
        return False, []

    if not all_forall:
        unique_mapping = dict()
        for index in indices:
            exists_indices = [index[i] for i in exists_indices]
            unique_mapping[tuple(exists_indices)] = index
    else:
        unique_mapping = {tuple(): indices[0]}

    candidate_bounded_variables = list()
    for index in unique_mapping.values():
        bounded_variables = default_bounded_variables.copy()
        for i, v in enumerate(index):
            bounded_variables[variable_name2var[output_variables[i]]] = state.object_type2name[variable_name2var[output_variables[i]].dtype.typename][v]
        if all_forall:
            return True, [bounded_variables]
        candidate_bounded_variables.append(bounded_variables)
    return False, candidate_bounded_variables


def _unique_plans(plans):
    plan_string_set = set()
    unique_plans = list()
    for state, csp, actions in plans:
        plan_string = ' '.join(str(x) for x in actions)
        if plan_string not in plan_string_set:
            plan_string_set.add(plan_string)
            unique_plans.append((state, csp, actions))
    return unique_plans


concepts/pdsketch/planners/optimistic_search_bilevel_utils.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : optimistic_search_with_simulation.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/26/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Planners using CSP solvers and simulators. The main difference between this file and the :mod:`concepts.pdsketch.csp_solvers.dpll_sampling` is that
this file assumes that the CSP is a grounding of a planning domain problem. Therefore, it does not rely on a PDSketch definition of the
underlying physical system. Instead, it can directly simulate the physical system to check the validity of a solution."""

import collections
from typing import Any, Optional, Union, Sequence, Tuple, List, Dict

import jacinle

from concepts.dsl.dsl_types import BOOL
from concepts.dsl.expression import ValueOutputExpression
from concepts.dsl.constraint import ConstraintSatisfactionProblem, EqualityConstraint
from concepts.pdsketch.executor import PDSketchExecutor, GeneratorManager
from concepts.pdsketch.domain import State
from concepts.pdsketch.operator import OperatorApplier

from concepts.dsl.constraint import AssignmentDict, Assignment, AssignmentType, OptimisticValue, SimulationFluentConstraintFunction
from concepts.dsl.tensor_value import TensorValue
from concepts.pdsketch.predicate import Predicate
from concepts.pdsketch.csp_solvers.dpll_sampling import CSPNotSolvable, CSPNoGenerator
from concepts.pdsketch.csp_solvers.dpll_sampling import _filter_deterministic_equal, _filter_unused_rhs, _filter_deterministic_clauses
from concepts.pdsketch.csp_solvers.dpll_sampling import _find_bool_variable, _apply_assignments as _apply_assignments_old, _find_grounded_function_application, _find_gen_variable_combined, _find_typegen_variable
from concepts.pdsketch.csp_solvers.dpll_sampling import ConstraintList

from concepts.pdsketch.planners.optimistic_search import optimistic_search_domain_check, instantiate_action, ground_action, ground_actions
from concepts.pdsketch.planners.optimistic_search_bilevel_utils import (
    enumerate_possible_symbolic_plans,
    enumerate_possible_symbolic_plans_regression_1,
    enumerate_possible_symbolic_plans_regression_c_v3
)
from concepts.pdsketch.simulator_interface import PDSketchSimulatorInterface

__all__ = ['construct_csp_from_optimistic_plan', 'solve_optimistic_plan_with_simulation', 'optimistic_search_with_simulation', 'csp_dpll_sampling_solve_with_simulation']


def construct_csp_from_optimistic_plan(
    executor: PDSketchExecutor, state: State, goal_expr: Union[str, ValueOutputExpression], actions: Sequence[OperatorApplier]
) -> Tuple[Sequence[OperatorApplier], ConstraintSatisfactionProblem, State, State]:
    csp = ConstraintSatisfactionProblem()
    if isinstance(goal_expr, str):
        goal_expr = executor.parse(goal_expr)
    assert isinstance(goal_expr, ValueOutputExpression)

    action_index = 0
    action_groundings = list()
    initial_state, state = state, state.clone()
    for a in actions:
        action_grounding = instantiate_action(csp, a)

        for sub_action in action_grounding.iter_sub_operator_appliers():
            action_groundings.append(sub_action)
            succ, state = executor.apply(sub_action, state, csp=csp, clone=False, action_index=action_index)
            action_index += 1
            if succ:
                pass
            else:
                executor.apply_precondition_debug(sub_action, state, csp=csp)
                raise ValueError(f'Unable to perform action {sub_action} at state {state}.')

    rv = executor.execute(goal_expr, state=state, csp=csp).item()
    if isinstance(rv, OptimisticValue):
        csp.add_constraint(EqualityConstraint.from_bool(rv, True), note='goal_test')
    else:
        rv = float(rv)
        if rv < 0.5:
            raise ValueError(f'Unable to satisfy the goal {goal_expr} with the given action skeleton.')

    return action_groundings, csp, initial_state, state


def solve_optimistic_plan_with_simulation(
    executor: PDSketchExecutor, simulator: 'PDSketchSimulatorInterface', state: State, goal_expr: Union[str, ValueOutputExpression], actions: Sequence[OperatorApplier], *,
    generator_manager: Optional[GeneratorManager] = None,
    csp_max_trials: int = 1, csp_max_generator_trials: int = 3, verbose: bool = True
) -> Tuple[State, ConstraintSatisfactionProblem, Optional[List[OperatorApplier]]]:
    """Solve an optimistic plan using the DPLL+Sampling algorithm and an actual physics simulator."""

    optimistic_search_domain_check(executor.domain)
    simulator.set_init_state(state)
    action_groundings, csp, initial_state, last_state = construct_csp_from_optimistic_plan(executor, state, goal_expr, actions)

    if verbose:
        print('CSP', csp)

    for _ in range(csp_max_trials):
        if verbose:
            print('CSP trial', _)

        with executor.checkpoint_storage():
            assignments = csp_dpll_sampling_solve_with_simulation(
                executor, simulator, csp, initial_state, action_groundings,
                generator_manager=generator_manager,
                max_generator_trials=csp_max_generator_trials,
                verbose=verbose
            )

            if assignments is not None:
                return last_state, csp, ground_actions(executor, action_groundings, assignments)

    return last_state, csp, None


class OptimisticSearchSymbolicPlanner(jacinle.JacEnum):
    BRUTE_FORCE = 'brute_force'
    REGRESSION_1 = 'regression_1'
    REGRESSION_C = 'regression_c'


def optimistic_search_with_simulation(
    executor: PDSketchExecutor, simulator: 'PDSketchSimulatorInterface', state: State, goal_expr: Union[str, ValueOutputExpression], *,
    generator_manager: Optional[GeneratorManager] = None,
    symbolic_planner: Union[str, OptimisticSearchSymbolicPlanner] = OptimisticSearchSymbolicPlanner.BRUTE_FORCE,
    max_actions: int = 4, symbolic_planner_kwargs: Optional[Dict[str, Any]] = None,
    csp_max_trials: int = 1, csp_max_generator_trials: int = 3, verbose: bool = True, solver_verbose: bool = False, use_tqdm: bool = True
) -> Tuple[Optional[State], Optional[ConstraintSatisfactionProblem], Optional[List[OperatorApplier]]]:
    """Solve an optimistic plan using the DPLL+Sampling algorithm and an actual physics simulator."""

    symbolic_planner = OptimisticSearchSymbolicPlanner.from_string(symbolic_planner)
    optimistic_search_domain_check(executor.domain)

    if symbolic_planner is OptimisticSearchSymbolicPlanner.BRUTE_FORCE:
        symbolic_plan_generator = enumerate_possible_symbolic_plans(executor, state, goal_expr, max_actions=max_actions, verbose=verbose, **(symbolic_planner_kwargs or {}))
    elif symbolic_planner is OptimisticSearchSymbolicPlanner.REGRESSION_1:
        symbolic_plan_generator = enumerate_possible_symbolic_plans_regression_1(executor, state, goal_expr, max_actions=max_actions, verbose=verbose, **(symbolic_planner_kwargs or {}))
    elif symbolic_planner is OptimisticSearchSymbolicPlanner.REGRESSION_C:
        symbolic_planner_kwargs = symbolic_planner_kwargs or {}
        symbolic_planner_kwargs['enable_csp'] = True
        symbolic_plan_generator, _ = enumerate_possible_symbolic_plans_regression_c_v3(executor, state, goal_expr, max_actions=max_actions, verbose=verbose, **symbolic_planner_kwargs)
    else:
        raise ValueError(f'Unknown symbolic planner {symbolic_planner}.')

    simulator.set_init_state(state)
    # for actions, csp, initial_state, last_state in symbolic_plan_generator:
    for symbolic_plan in symbolic_plan_generator:
        actions, csp, initial_state, last_state = symbolic_plan.actions, symbolic_plan.csp, symbolic_plan.initial_state, symbolic_plan.last_state
        if verbose:
            print('Solving for: ', actions)

        if use_tqdm:
            trials = jacinle.tqdm(range(csp_max_trials))
        else:
            trials = range(csp_max_trials)
        for csp_trial_index in trials:
            if verbose and not use_tqdm:
                print('CSP trial', csp_trial_index)

            with executor.checkpoint_storage(), simulator.restore_context():
                assignments = csp_dpll_sampling_solve_with_simulation(
                    executor, simulator, csp, initial_state, actions,
                    generator_manager=generator_manager,
                    max_generator_trials=csp_max_generator_trials,
                    verbose=solver_verbose
                )

                if assignments is not None:
                    return last_state, csp, ground_actions(executor, actions, assignments)

    return None, None, None


def _filter_unused_simulation_rhs(executor: PDSketchExecutor, constraints: ConstraintList, assignments: Dict[int, Assignment]) -> ConstraintList:
    """Filter out simulation constraints that only appear once in the RHS of the constraints. In this case, the variable can be ignored and the related constraints can be removed.

    Args:
        executor: the executor.
        constraints: the list of constraints.
        assignments: the dictionary of assignments.

    Returns:
        the list of constraints that have not been satisfied, after removing all unused variables.
    """
    used = collections.defaultdict(int)
    for c in constraints:
        for x in c.arguments:
            if isinstance(x, OptimisticValue):
                used[x.identifier] += 100  # as long as a variable has appeared in the lhs of a constraint, it is used.
        if isinstance(c.rv, OptimisticValue) and isinstance(c.function, SimulationFluentConstraintFunction):
            used[c.rv.identifier] += 1  # if the variable has only appeared in the rhs of a constraint for once, it is not used.
    for k, v in used.items():
        if v == 1:
            assignments[k] = Assignment(AssignmentType.IGNORE, None)
    return _apply_assignments_old(executor, constraints, assignments)


def _apply_assignments_with_simulation(
    executor: PDSketchExecutor,
    constraints: ConstraintList, assignments: Dict[int, Assignment],
    simulator: PDSketchSimulatorInterface, actions: Sequence[OperatorApplier],
    verbose: bool = False
) -> Tuple[ConstraintList, AssignmentDict]:
    new_constraints = _apply_assignments_old(executor, constraints, assignments)
    new_assignments = None

    while True:
        next_action_index = simulator.last_action_index + 1
        if next_action_index < len(actions):
            action = actions[next_action_index]

            try:
                grounded_action = ground_action(executor, action, assignments)
            except AssertionError:
                # The action is not applicable.
                break

            if verbose:
                jacinle.log_function.print(f'Executing grounded action: #{next_action_index}')

            action_name = grounded_action.operator.controller.name
            action_args = executor.get_controller_args(grounded_action, simulator.get_latest_pd_state())
            succ, state = simulator.run(
                next_action_index,
                action_name,
                action_args
            )
            if not succ:
                # jacinle.log_function.print(f'Action {grounded_action} failed.')
                raise CSPNotSolvable(f'Unable to perform action {action}.')
            # jacinle.log_function.print(f'Action {grounded_action} succeeded.')

            if new_assignments is None:
                new_assignments = assignments.copy()

            # Update the assignments.
            for i, c in enumerate(new_constraints):
                if isinstance(c.function, SimulationFluentConstraintFunction):
                    function: SimulationFluentConstraintFunction = c.function
                    if function.action_index == next_action_index:
                        if isinstance(c.rv, TensorValue):
                            if c.rv.dtype != BOOL:
                                raise NotImplementedError('Only bool is supported for simulation constraints.')

                            # print('Simulation constraint', c)
                            # print(state.features[function.predicate.name])
                            # print('Desired value =', c.rv.value, 'Actual value =', state.features[function.predicate.name][function.arguments].item())
                            # import pybullet
                            # pybullet.stepSimulation()
                            # import ipdb; ipdb.set_trace()

                            if state.features[function.predicate.name][function.arguments].item() != c.rv.item():
                                raise CSPNotSolvable()
                            else:
                                new_constraints[i] = None
                        else:
                            new_assignments[c.rv.identifier] = Assignment(
                                AssignmentType.VALUE,
                                state.features[function.predicate.name][function.arguments]
                            )
                            new_constraints[i] = None

            new_constraints = _apply_assignments_old(executor, new_constraints, new_assignments)
        else:
            break

    return new_constraints, new_assignments if new_assignments is not None else assignments


def csp_dpll_sampling_solve_with_simulation(
    executor: PDSketchExecutor, simulator: PDSketchSimulatorInterface, csp: ConstraintSatisfactionProblem,
    state: State, actions: Sequence[OperatorApplier], *,
    generator_manager: Optional[GeneratorManager] = None,
    max_generator_trials: int = 3,
    enable_ignore: bool = False, solvable_only: bool = False,
    verbose: bool = False
) -> Optional[Union[bool, AssignmentDict]]:
    if generator_manager is None:
        generator_manager = GeneratorManager(executor, store_history=False)

    constraints = csp.constraints.copy()

    @jacinle.log_function(verbose=False)
    def dfs(constraints, assignments):
        if len(constraints) == 0:
            return assignments

        progress = True
        while progress:
            progress, constraints = _filter_deterministic_equal(executor, constraints, assignments)
        if enable_ignore:
            constraints = _filter_unused_rhs(executor, constraints, assignments, csp.index2actionable)
        else:
            # NB(Jiayuan Mao @ 2023/03/11): for simulation constraints, we can remove them if they are not used in any other constraints.
            # TODO(Jiayuan Mao @ 2023/03/15): actually, I just noticed that this is a "bug" in the implementation of `_find_bool_variable`.
            # Basically, if the variable is a simulation variable and it only appears in the RHS of a constraint, then it will be ignored.
            # The current handling will work, but probably I need to think about a better way to handle this.
            constraints = _filter_unused_simulation_rhs(executor, constraints, assignments)

        progress = True
        while progress:
            progress, constraints = _filter_deterministic_clauses(executor, constraints, assignments)

        if verbose:
            jacinle.log_function.print('Remaining constraints:', len(constraints))
            jacinle.log_function.print(*constraints, sep='\n')

        if len(constraints) == 0:
            return assignments

        if (next_bool_var := _find_bool_variable(executor, constraints, assignments)) is not None:
            assignments_true = assignments.copy()
            assignments_true[next_bool_var] = Assignment(AssignmentType.VALUE, True)

            with simulator.restore_context(verbose):
                try:
                    constraints_true, assignments_true = _apply_assignments_with_simulation(executor, constraints, assignments_true, simulator, actions, verbose=verbose)
                    return dfs(constraints_true, assignments_true)
                except CSPNotSolvable:
                    pass

            assignments_false = assignments.copy()
            assignments_false[next_bool_var] = Assignment(AssignmentType.VALUE, False)
            with simulator.restore_context(verbose):
                try:
                    constraints_false, assignments_false = _apply_assignments_with_simulation(executor, constraints, assignments_false, simulator, actions, verbose=verbose)
                    return dfs(constraints_false, assignments_false)
                except CSPNotSolvable:
                    pass

            raise CSPNotSolvable()
        elif (next_fapp := _find_grounded_function_application(executor, constraints)) is not None:
            function: Predicate = next_fapp.function
            arguments = next_fapp.arguments

            external_function = executor.get_function_implementation(f'predicate::{function.name}')
            output = external_function(*arguments)

            target = next_fapp.rv
            new_assignments = assignments.copy()
            new_assignments[target.identifier] = Assignment(AssignmentType.VALUE, output)
            with simulator.restore_context(verbose):
                try:
                    new_constraints = constraints.copy()
                    new_constraints[new_constraints.index(next_fapp)] = None
                    new_constraints, new_assignments = _apply_assignments_with_simulation(executor, new_constraints, new_assignments, simulator, actions, verbose=verbose)
                    return dfs(new_constraints, new_assignments)
                except CSPNotSolvable:
                    pass

            raise CSPNotSolvable()
        elif (next_gen_vars := _find_gen_variable_combined(executor, csp, constraints, assignments)) is not None:
            if len(next_gen_vars) > 1:
                # jacinle.log_function.print('Generator orders', *[str(vv[1]).split('\n')[0] for vv in next_gen_vars], sep='\n  ')
                pass

            for vv in next_gen_vars:
                c, g, args, outputs_target = vv
                if g.unsolvable:
                    raise CSPNotSolvable('Hit unsolvable generator.')

                if verbose:
                    jacinle.log_function.print(f'Generator: {g}\nArgs: {args}')
                generator = generator_manager.call(g, max_generator_trials, args, c)

                # NB(Jiayuan Mao @ 2023/03/03): I didn't write for x in generator in order to make the verbose output more readable.
                generator = iter(generator)
                for j in range(max_generator_trials):
                    if verbose:
                        jacinle.log_function.print('Running generator', g.name, f'count = {j + 1} / {max_generator_trials}')

                    try:
                        output_index, outputs = next(generator)
                    except StopIteration:
                        if verbose:
                            jacinle.log_function.print('Generator', g.name, 'exhausted.')
                        break

                    new_assignments = assignments.copy()
                    for output, target in zip(outputs, outputs_target):
                        new_assignments[target.identifier] = Assignment(AssignmentType.VALUE, output, generator_index=output_index)
                        # jacinle.log_function.print('Assigned', target, output)

                    with simulator.restore_context(verbose):
                        try:
                            new_constraints = constraints.copy()
                            if isinstance(c, list):
                                for cc in c:
                                    new_constraints[new_constraints.index(cc)] = None
                            else:
                                new_constraints[new_constraints.index(c)] = None

                            new_constraints, new_assignments = _apply_assignments_with_simulation(executor, new_constraints, new_assignments, simulator, actions, verbose=verbose)
                            return dfs(new_constraints, new_assignments)
                        except CSPNotSolvable as e:
                            if verbose:
                                jacinle.log_function.print('Failed to apply assignments. Reason:', e)
                            pass

            raise CSPNotSolvable()
        else:
            # jacinle.log_function.print('Can not find a generator. Constraints:\n  ' + '\n  '.join([str(x) for x in constraints]))
            raise CSPNoGenerator('Can not find a generator. Constraints:\n  ' + '\n  '.join([str(x) for x in constraints]))

    try:
        assignments = dfs(constraints, {})
        if not solvable_only:
            for name, dtype in csp.index2type.items():
                if name not in assignments:
                    g = _find_typegen_variable(executor, dtype)
                    if g is None:
                        raise NotImplementedError('Can not find a generator for unbounded variable {}, type {}.'.format(name, dtype))
                    else:
                        try:
                            output_index, (output, ) = next(iter(generator_manager.call(g, 1, [], None)))
                            assignments[name] = Assignment(AssignmentType.VALUE, output, generator_index=output_index)
                        except StopIteration:
                            raise CSPNotSolvable

        if generator_manager.store_history:
            generator_manager.mark_success(assignments)
        if solvable_only:
            return True
        return assignments
    except CSPNotSolvable:
        return None
    except CSPNoGenerator:
        raise


concepts/pdsketch/planners/optimistic_search_with_simulation.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/09/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/pdsketch/planners/regression/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : search_utils.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 11/20/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Utilities for PDSketch planners."""

from typing import Any

__all__ = ['MostPromisingTrajectoryTracker']


class MostPromisingTrajectoryTracker(object):
    """This is a tracker for tracking the most promising next action, used in joint learning settings
    where we don't have a knowledge about the actual transition function or the goal."""

    def __init__(self, threshold: float):
        """Initialize the tracker.

        Args:
            threshold: the threshold for the score.
        """

        self.threshold = threshold
        self.best_score = float('-inf')
        self.solution = None

    threshold: float
    """A score threshold."""

    best_score: float
    """The best score achived so far."""

    solution: Any
    """The solution associated with the `best_score`."""

    def check(self, new_score: float) -> bool:
        """If the new score is better than the current best score, return True.

        Args:
            new_score: the new score.

        Returns:
            True if the new score is better than the current best score.
        """
        return new_score > self.best_score

    def update(self, new_score: float, solution: Any):
        """Update the best score and the solution.

        Args:
            new_score: the new score.
            solution: the new solution.
        """
        assert new_score > self.best_score
        self.best_score = new_score
        self.solution = solution


concepts/pdsketch/planners/search_utils.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : predicate.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/26/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import contextlib
from typing import Optional, Union, Sequence, Tuple, Set, Dict, Callable

from jacinle.utils.meta import repr_from_str

from concepts.dsl.dsl_types import ObjectType, Variable
from concepts.dsl.dsl_functions import FunctionType, Function
from concepts.dsl.expression import Expression, ExpressionDefinitionContext, iter_exprs
from concepts.dsl.expression import FunctionApplicationExpression, VariableExpression, ObjectConstantExpression, ValueOutputExpression, VariableAssignmentExpression
from concepts.dsl.expression import BoolExpression, QuantificationExpression, BoolOpType
from concepts.dsl.expression import PredicateEqualExpression, AssignExpression, ConditionalSelectExpression, DeicticSelectExpression, ConditionalAssignExpression, DeicticAssignExpression
from concepts.dsl.expression_visitor import ExpressionVisitor

__all__ = ['Predicate', 'flatten_expression', 'get_used_state_variables', 'is_simple_bool', 'split_simple_bool', 'get_simple_bool_predicate']


class Predicate(Function):
    """A predicate is a special function that can part of the observation or state space.
    Predicate does not support overloaded function types. Currently, predicates are only used in planning domains."""

    def __init__(
        self,
        name: str,
        ftype: FunctionType,
        derived_expression: Optional[str] = None,
        overridden_call: Optional[Callable] = None,
        observation: Optional[bool] = None,
        state: Optional[bool] = None,
        generator_placeholder: Optional[bool] = None,
        inplace_generators: Optional[Sequence[str]] = None,
    ):
        assert isinstance(name, str)
        super().__init__(name, ftype, derived_expression=derived_expression, overridden_call=overridden_call)

        # self.is_derived has been coomputed in super().__init__
        self.is_static = False
        self.is_cacheable = self._guess_is_cacheable()

        self.is_observation_variable = observation if observation is not None else self.is_cacheable
        self.is_state_variable = state if state is not None else self.is_cacheable
        self.is_generator_placeholder = generator_placeholder if generator_placeholder is not None else False
        self.inplace_generators = tuple(inplace_generators) if inplace_generators is not None else tuple()

        # from concepts.pdsketch.strips.ao_discretization import AOFeatureCodebook
        # self.ao_discretization: Optional[AOFeatureCodebook] = None  # for AODiscretization
        self.ao_discretization = None

        self._check_flags_sanity()

    is_static: bool
    """Whether the predicate is a static predicate. I.e., the predicate will never be changed."""

    is_cacheable: bool
    """Whether the predicate can be cached. Specifically, if it contains only "ObjectTypes" as arguments."""

    is_observation_variable: bool
    """Whether the predicate is an observation variable."""

    is_state_variable: bool
    """Whether the predicate is a state variable."""

    is_generator_placeholder: bool
    """Whether the predicate is a generator placeholder."""

    inplace_generators: Tuple[str, ...]
    """The names of the inplace defined generators associated with this predicate."""

    # ao_discretization: Optional['AOFeatureCodebook']
    # """The discretization codebook for this predicate."""

    def _guess_is_cacheable(self) -> bool:
        """Return whether the function can be cached. Specifically, if it contains only "ObjectTypes" as arguments, it can be statically evaluated."""
        for arg_def in self.arguments:
            if not isinstance(arg_def.dtype, ObjectType):
                return False
        return True

    def mark_static(self, flag: bool = True):
        """Mark a predicate as static (i.e., its grounded value will never change).

        Args:
            flag: Whether to mark the predicate as static.
        """
        self.is_static = flag

    def _check_flags_sanity(self):
        if self.is_observation_variable:
            assert self.is_cacheable and not self.is_derived
        if self.is_state_variable:
            assert self.is_cacheable
            if self.is_derived:
                for predicate_def in get_used_state_variables(self.derived_expression):
                    assert predicate_def.is_observation_variable and not predicate_def.is_state_variable
            else:
                assert self.is_cacheable
        if self.is_generator_placeholder:
            assert not self.is_derived and not self.is_state_variable and not self.is_observation_variable

    @property
    def flags(self) -> Dict[str, bool]:
        """Return the flags of the predicate, which is a dictionary of {flag_name: flag_value}."""
        return {
            'is_derived': self.is_derived,
            'is_static': self.is_static,
            'is_cacheable': self.is_cacheable,
            'is_observation_variable': self.is_observation_variable,
            'is_state_variable': self.is_state_variable,
            'is_generator_placeholder': self.is_generator_placeholder,
        }

    def rename(self, new_name: str) -> 'Predicate':
        """Rename the predicate."""
        return Predicate(
            new_name, self.ftype,
            derived_expression=self.derived_expression,
            overridden_call=self.overridden_call,
            observation=self.is_observation_variable,
            state=self.is_state_variable,
            generator_placeholder=self.is_generator_placeholder,
            inplace_generators=self.inplace_generators
        )

    def __str__(self) -> str:
        flags = list()

        if self.is_observation_variable:
            flags.append('observation')
        if self.is_state_variable:
            flags.append('state')
        if self.is_generator_placeholder:
            flags.append('gen')
        if self.is_cacheable:
            flags.append('cacheable')
        if self.is_static:
            flags.append('static')
        flags_string = '[' + ', '.join(flags) + ']' if len(flags) > 0 else ''
        arguments = ', '.join([str(arg) for arg in self.arguments])
        fmt = f'{self.name}{flags_string}({arguments}) -> {self.return_type}'
        if self.derived_expression is not None:
            fmt += ' {\n'
            fmt += '  ' + str(self.derived_expression)
            fmt += '\n}'
        return fmt

    __repr__ = repr_from_str


def flatten_expression(
    expr: Expression,
    mappings: Optional[Dict[Union[FunctionApplicationExpression, VariableExpression], Union[Variable, ValueOutputExpression]]] = None,
    ctx: Optional[ExpressionDefinitionContext] = None,
    flatten_cacheable_expression: bool = True,
) -> Union[VariableExpression, ObjectConstantExpression, ValueOutputExpression, VariableAssignmentExpression]:
    """Flatten an expression by replacing certain variables or function applications with sub-expressions.
    The input mapping is a dictionary of {expression: sub-expression}. There are two cases:

    - The expression is a :class:`~concepts.dsl.expression.VariableExpression`, and the sub-expression is a
        :class:`~concepts.dsl.dsl_types.Variable` or a :class:`~concepts.dsl.expression.ValueOutputExpression`. In this case,
        the variable expression will is the sub-expression used for replacing the variable.
    - The expression is a :class:`~concepts.dsl.expression.FunctionApplicationExpression`, and the sub-expression is a
        :class:`~concepts.dsl.dsl_types.Variable`. Here, the function
        application expression must be a "simple" function application expression, i.e., it contains only variables
        as arguments. The entire function application expression will be replaced by the Variable

    Args:
        expr: the expression to flatten.
        mappings: a dictionary of {expression: sub-expression} to replace the expression with the sub-expression.
        ctx: a :class:`~concepts.dsl.expression.ExpressionDefinitionContext`.
        flatten_cacheable_expression: whether to flatten cacheable expressions. If False, cacheable function applications will be kept as-is.

    Returns:
        the flattened expression.
    """

    if mappings is None:
        mappings = dict()

    if ctx is None:
        ctx = ExpressionDefinitionContext()

    with ctx.as_default():
        return _FlattenExpressionVisitor(ctx, mappings, flatten_cacheable_expression).visit(expr)


class _FlattenExpressionVisitor(ExpressionVisitor):
    def __init__(
        self,
        ctx: ExpressionDefinitionContext,
        mappings: Dict[Union[FunctionApplicationExpression, VariableExpression], Union[Variable, ValueOutputExpression]],
        flatten_cacheable_expression: bool = True,
    ):

        self.ctx = ctx
        self.mappings = mappings
        self.flatten_cacheable_expression = flatten_cacheable_expression

    def visit_variable_expression(self, expr: VariableExpression) -> Union[VariableExpression, ValueOutputExpression]:
        rv = expr
        for k, v in self.mappings.items():
            if isinstance(k, VariableExpression):
                if k.name == expr.name:
                    if isinstance(v, Variable):
                        rv = VariableExpression(v)
                    else:
                        rv = v
                    break
        return rv

    def visit_function_application_expression(self, expr: FunctionApplicationExpression) -> Union[VariableExpression, ValueOutputExpression]:
        # Case 1: the function application will be replaced by something in the mappings.
        for k, v in self.mappings.items():
            if isinstance(k, FunctionApplicationExpression):
                if expr.function.name == k.function.name and all(
                    isinstance(a1, VariableExpression) and isinstance(a2, VariableExpression) and a1.name == a2.name for a1, a2 in zip(expr.arguments, k.arguments)
                ):
                    assert isinstance(v, Variable)
                    return VariableExpression(v)

        # Case 2 contains three sub-cases:
        #   (1) the function is not a derived function
        #   (2) the function corresponds to a state variable
        #   (3) the function is a cacheable function and we do not want to flatten it.
        if not expr.function.is_derived or (isinstance(expr.function, Predicate) and expr.function.is_state_variable) or (not self.flatten_cacheable_expression and expr.function.ftype.is_cacheable):
            return type(expr)(expr.function, [self.visit(e) for e in expr.arguments])

        # Case 3: the function is a derived function and we want to flatten it.
        for arg in expr.function.arguments:
            if not isinstance(arg, Variable):
                raise TypeError(f'Cannot flatten function application {expr} because it contains non-variable arguments.')

        # (1) First resolve the arguments.
        argvs = [self.visit(e) for e in expr.arguments]

        # (2) Make a backup of the current context, and then create a new context using the arguments.
        old_ctx = self.ctx
        old_mappings = self.mappings
        self.ctx = ExpressionDefinitionContext(*expr.function.arguments)
        self.mappings = dict()
        for arg, argv in zip(expr.function.arguments, argvs):
            if isinstance(arg, Variable):
                self.mappings[VariableExpression(arg)] = argv

        # (3) Flatten the derived expression.
        with self.ctx.as_default():
            rv = self.visit(expr.function.derived_expression)

        # (4) Restore the old context.
        self.ctx = old_ctx
        self.mappings = old_mappings

        # (5) Flatten the result again, using the old context + mappings.
        return type(rv)(rv.function, [self.visit(e) for e in rv.arguments])

    def visit_bool_expression(self, expr: BoolExpression) -> BoolExpression:
        return BoolExpression(expr.bool_op, [self.visit(child) for child in expr.arguments])

    def visit_quantification_expression(self, expr: QuantificationExpression) -> QuantificationExpression:
        with self.make_dummy_variable(expr.variable) as dummy_variable:
            return QuantificationExpression(expr.quantification_op, dummy_variable, self.visit(expr.expr))

    def visit_predicate_equal_expression(self, expr: PredicateEqualExpression) -> PredicateEqualExpression:
        return type(expr)(self.visit(expr.predicate), self.visit(expr.value))

    def visit_assign_expression(self, expr: AssignExpression) -> AssignExpression:
        return type(expr)(self.visit(expr.predicate), self.visit(expr.value))

    def visit_conditional_select_expression(self, expr: ConditionalSelectExpression) -> ConditionalSelectExpression:
        return type(expr)(self.visit(expr.predicate), self.visit(expr.condition))

    def visit_deictic_select_expression(self, expr: DeicticSelectExpression) -> DeicticSelectExpression:
        with self.make_dummy_variable(expr.variable) as dummy_variable:
            return type(expr)(dummy_variable, self.visit(expr.expression))

    def visit_conditional_assign_expression(self, expr: ConditionalAssignExpression) -> ConditionalAssignExpression:
        return type(expr)(self.visit(expr.predicate), self.visit(expr.value), self.visit(expr.condition))

    def visit_deictic_assign_expression(self, expr: DeicticAssignExpression) -> DeicticAssignExpression:
        with self.make_dummy_variable(expr.variable) as dummy_variable:
            return type(expr)(dummy_variable, self.visit(expr.expr))

    def visit_constant_expression(self, expr: Expression) -> Expression:
        return expr

    def visit_object_constant_expression(self, expr: Expression) -> Expression:
        return expr

    @contextlib.contextmanager
    def make_dummy_variable(self, variable: Variable):
        dummy_variable = self.ctx.gen_random_named_variable(variable.dtype)
        dummy_variable_expr = VariableExpression(variable)

        old_mapping = self.mappings.get(dummy_variable_expr, None)
        self.mappings[dummy_variable_expr] = dummy_variable

        yield dummy_variable

        self.mappings[dummy_variable_expr] = old_mapping


def get_used_state_variables(expr: ValueOutputExpression) -> Set[Predicate]:
    """Return the set of state variables used in the given expression.

    Args:
        expr: the expression to be analyzed.

    Returns:
        the set of state variables (the :class:`Predicate` objects) used in the given expression.
    """

    assert isinstance(expr, ValueOutputExpression), (
        'Only value output expression has well-defined used-state-variables.\n'
        'For value assignment expressions, please separately process the targets, conditions, and values.'
    )

    used_svs = set()

    def dfs(this):
        nonlocal used_svs
        for e in iter_exprs(this):
            if isinstance(e, FunctionApplicationExpression):
                if isinstance(e.function, Predicate) and e.function.is_state_variable:
                    used_svs.add(e.function)
                elif e.function.derived_expression is not None:
                    dfs(e.function.derived_expression)

    dfs(expr)
    return used_svs


def is_simple_bool(expr: Expression) -> bool:
    """Check if the expression is a simple Boolean expression. That is, it is either a Boolean state variable,
    or the negation of a Boolean state variable.

    Args:
        expr: the expression to check.

    Returns:
        True if the expression is a simple boolean expression, False otherwise.
    """

    if isinstance(expr, FunctionApplicationExpression) and isinstance(expr.function, Predicate) and expr.function.is_state_variable:
        return True
    if isinstance(expr, BoolExpression) and expr.bool_op is BoolOpType.NOT:
        return is_simple_bool(expr.arguments[0])
    return False


def split_simple_bool(expr: Expression, initial_negated: bool = False) -> Tuple[Optional[FunctionApplicationExpression], bool]:
    """
    If the expression is a simple Boolean expression (see :func:`is_simple_bool`),
    it returns the feature definition and a boolean indicating whether the expression is negated.

    Args:
        expr (Expression): the expression to be checked.
        initial_negated (bool, optional): whether outer context of the feature expression is a negated function.

    Returns:
        a tuple of the feature application and a boolean indicating whether the expression is negated.
        The first element is None if the feature is not a simple Boolean feature application.
    """
    if isinstance(expr, FunctionApplicationExpression) and isinstance(expr.function, Predicate) and expr.function.is_state_variable:
        return expr, initial_negated
    if isinstance(expr, BoolExpression) and expr.bool_op is BoolOpType.NOT:
        return split_simple_bool(expr.arguments[0], not initial_negated)
    return None, initial_negated


def get_simple_bool_predicate(expr: Expression) -> Predicate:
    """If the expression is a simple bool (see :func:`is_simple_bool`), it returns the underlying predicate.

    Args:
        expr (Expression): the expression, assumed to be a simple Boolean expression.

    Returns:
        the underlying predicate.
    """
    if isinstance(expr, FunctionApplicationExpression) and isinstance(expr.function, Predicate) and expr.function.is_state_variable:
        return expr.function
    assert isinstance(expr, BoolExpression) and expr.bool_op is BoolOpType.NOT
    return get_simple_bool_predicate(expr.arguments[0])

concepts/pdsketch/predicate.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : simulator_interface.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/16/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import contextlib
from typing import Any, Iterable, Tuple, Dict, Callable

import jacinle

from concepts.pdsketch.operator import OperatorApplier
from concepts.pdsketch.executor import PythonFunctionRef, PDSketchExecutor
from concepts.pdsketch.domain import State

__all__ = ['PDSketchSimulatorInterface']


class PDSketchSimulatorInterface(object):
    """The base class for interaction with a (physical simulator). This class serves two purposes:

    1. perform simulation for planners. Therefore, it requires actions to support `restore` operations.
    2. perform simulation when executing a given plan. In this case, restore operations are not required.
    """

    def __init__(self, executor: PDSketchExecutor):
        """Initialize the simulator interface.

        Args:
            executor: the :class:`~concepts.pdsketch.executor.PDSketchExecutor` instance for grounding action parameters.
        """
        self._executor = executor
        self._last_action_index = -1

        self.controllers = dict()
        self.pd_states = dict()
        self.restore_functions = dict()

    controllers: Dict[str, Callable]
    """Registered controllers."""

    pd_states: Dict[int, State]
    """Mappings from action indices to PDSketch states. -1 is the initial state."""

    restore_functions: Dict[int, Any]
    """Mappings from action indices to controller states. -1 is the initial state."""

    def register_action_controller(self, name: str, controller_function: Callable):
        """Register a controller function for a controller name.

        Args:
            name: the name of the controller.
            controller_function: the controller function. It should take a PDSketch state as the first argument, together with other action arguments,
                and return a tuple of (success, new PDSketch state, restore function).
        """
        if not isinstance(controller_function, PythonFunctionRef):
            controller_function = PythonFunctionRef(controller_function, executor=self._executor)

        self.controllers[name] = controller_function.set_executor(self._executor)

    def set_init_state(self, state: State):
        """Set the initial PDSketch state of the simulator."""
        self.pd_states[-1] = state

    def get_pd_state(self, action_index: int) -> State:
        """Get the PDSketch state at a given action index. -1 is the initial state.

        Args:
            action_index: the action index.
        """
        return self.pd_states[action_index]

    def get_latest_pd_state(self) -> State:
        """Get the latest PDSketch state."""
        return self.pd_states[self._last_action_index]

    def get_restore_function(self, action_index: int) -> Any:
        """Get the restore function for a given action index.

        Args:
            action_index: the action index.
        """
        return self.restore_functions[action_index]

    def run_operator_applier(self, action_index: int, action: OperatorApplier) -> Tuple[bool, State]:
        """Execute an action, and return the success flag and the new PDSketch state. This function will not store the restore function corresponding to the action.

        Args:
            action_index: the action index.
            action: the action to execute.

        Returns:
            a tuple of (success, new PDSketch state).
        """

        action_name = action.operator.controller.name
        action_args = self._executor.get_controller_args(action, self.get_latest_pd_state())
        return self.run(action_index, action_name, action_args)

    def run(self, action_index: int, action_name: str, action_args: Tuple[Any, ...], verbose: bool = False) -> Tuple[bool, State]:
        """Run an action, and return the success flag and the new PDSketch state. This function will also store the restore function corresponding to the action.

        Args:
            action_index: the action index.
            action_name: the name of the action (the controller).
            action_args: the arguments to the controller.
            verbose: whether to print verbose information.

        Returns:
            a tuple of (success, new PDSketch state).
        """

        assert action_index == self._last_action_index + 1, f'Action index {action_index} is not the next action index.'
        if action_name not in self.controllers:
            raise ValueError(f'No controller registered for {action_name}.')

        if verbose:
            jacinle.log_function.print('Running action', action_index, action_name)

        state = self.pd_states[action_index - 1]

        # TODO (Jiayuan@2023/08/16): support wrap_rv, or handle it at the definition of the controller.
        succ, pd_state, restore_function = self.controllers[action_name](state, *action_args, wrap_rv=False)

        if not succ:
            self.restore(action_index - 1, verbose=verbose)
            return False, None

        self.pd_states[action_index] = pd_state
        self.restore_functions[action_index] = restore_function
        self._last_action_index = action_index
        return succ, pd_state

    def restore(self, target_action_index: int, verbose: bool = False) -> bool:
        """Restore the simulator to a given action index.

        Args:
            target_action_index: the target action index.
            verbose: whether to print verbose information.

        Returns:
            whether the restore operation is successful.
        """
        if target_action_index == self._last_action_index:
            return True

        for action_index in range(self._last_action_index, target_action_index, -1):
            if verbose:
                jacinle.log_function.print(f'Restoring action {action_index}')
            if action_index not in self.restore_functions:
                raise ValueError(f'No restore function for action {action_index}.')
            if self.restore_functions[action_index] is None:
                raise ValueError(f'Empty restore function for action {action_index}.')
            self.restore_functions[action_index]()
            del self.restore_functions[action_index]  # safe to delete the restore function after it is used.

        self._last_action_index = target_action_index
        return True

    @contextlib.contextmanager
    def restore_context(self, verbose: bool = False) -> Iterable[int]:
        """A context manager for restoring the simulator to the current action index."""
        action_index = self._last_action_index
        yield self._last_action_index
        self.restore(action_index, verbose=verbose)

    @property
    def last_action_index(self):
        """The last action index."""
        return self._last_action_index


concepts/pdsketch/simulator_interface.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 03/19/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from .atomic_strips_domain import AtomicStripsOperator, AtomicStripsDomain, AtomicStripsOperatorApplier, AtomicStripsProblem, load_astrips_domain_file, load_astrips_domain_string, load_astrips_problem_file, load_astrips_problem_string
from .strips_expression import *
from .strips_grounded_expression import *
from .strips_grounding import *
from .strips_heuristics import *
from .strips_search import *


concepts/pdsketch/strips/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : atomic_strips_domain.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 04/27/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from typing import Optional, Union, Sequence, Tuple, Mapping, Dict

import jacinle
from jacinle.utils.cache import cached_property
from concepts.dsl.dsl_types import ObjectType, Variable, ObjectConstant
from concepts.dsl.expression import VariableExpression, FunctionApplicationExpression, NotExpression, AssignExpression, ConstantExpression, is_and_expr
from concepts.pdsketch.predicate import Predicate
from concepts.pdsketch.operator import Operator, RegressionRule, OperatorApplicationExpression, AchieveExpression
from concepts.pdsketch.domain import Domain, Problem
from concepts.pdsketch.parsers.pdsketch_parser import load_domain_file, load_domain_string, load_problem_file, load_problem_string
from concepts.pdsketch.strips.strips_expression import SBoolPredicateApplicationExpression, SProposition, SState, SStateCompatible, make_sproposition_from_function_application


# TODO(Jiayuan Mao @ 2023/04/27): implement ADL operator.

class AtomicStripsOperator(object):
    def __init__(
        self,
        arguments: Sequence[Variable],
        preconditions: Sequence[SBoolPredicateApplicationExpression],
        add_effects: Sequence[SBoolPredicateApplicationExpression],
        del_effects: Sequence[SBoolPredicateApplicationExpression],
        raw_operator: Operator
    ):
        self.name = raw_operator.name
        self.arguments = tuple(arguments)
        self.preconditions = tuple(preconditions)
        self.add_effects = tuple(add_effects)
        self.del_effects = tuple(del_effects)
        self.raw_operator = raw_operator

    arguments: Tuple[Variable]
    """The arguments of the operator, as a tuple of :class:`~concepts.dsl.dsl_types.Variable`."""

    preconditions: Tuple[SBoolPredicateApplicationExpression, ...]
    """The precondition of the operator, as a tuple of :class:`~concepts.pdsketch.strips.strips_expression.SBoolPredicateApplicationExpression`."""

    add_effects: Tuple[SBoolPredicateApplicationExpression, ...]
    """The add effects of the operator, as a tuple of :class:`~concepts.pdsketch.strips.strips_expression.SBoolPredicateApplicationExpression`."""

    del_effects: Tuple[SBoolPredicateApplicationExpression, ...]
    """The delete effects of the operator, as a tuple of :class:`~concepts.pdsketch.strips.strips_expression.SBoolPredicateApplicationExpression`."""

    raw_operator: Operator
    """The raw operator that this strips operator is derived from."""

    def relax(self) -> 'AtomicStripsOperator':
        return type(self)(
            arguments=self.arguments,
            preconditions=self.preconditions,
            add_effects=self.add_effects,
            del_effects=[],
            raw_operator=self.raw_operator
        )

    @classmethod
    def from_operator(cls, operator: Operator) -> 'AtomicStripsOperator':
        """Create a strips operator from a PDSketch operator.

        Args:
            operator: the PDSketch operator.

        Returns:
            the strips operator.
        """
        arguments = operator.arguments

        preconditions = list()
        for precondition in operator.preconditions:
            precondition = precondition.bool_expr
            preconditions.append(SBoolPredicateApplicationExpression.from_function_application_expression(precondition))

        add_effects = list()
        del_effects = list()
        for effect in operator.effects:
            effect = effect.assign_expr
            assert isinstance(effect, AssignExpression), f'Not supported assign expression: {effect}.'

            sassign = SBoolPredicateApplicationExpression.from_function_application_expression(effect.predicate)
            assert isinstance(effect.value, ConstantExpression), f'Not supported assign expression: {effect}.'
            value = effect.value.constant.item()
            assert isinstance(value, int)
            if value > 0.5:
                add_effects.append(sassign)
            else:
                del_effects.append(sassign)

        return cls(arguments, preconditions, add_effects, del_effects, operator)

    def __str__(self) -> str:
        return f'{self.name}({", ".join(arg.name for arg in self.arguments)})'

    def __repr__(self) -> str:
        return f'StripsOperator({self.name}, arguments={self.arguments}, preconditions={self.preconditions}, {self.add_effects}, {self.del_effects})'

    def to_applier_str(self, bound_arguments: Dict[str, str]):
        return self.name + '(' + ', '.join(f'{arg.name}={bound_arguments[arg.name]}' for arg in self.arguments) + ')'

    def to_applier_pddl_str(self, bound_arguments: Dict[str, str]):
        return '(' + self.name + ' ' + ' '.join(bound_arguments[arg.name] for arg in self.arguments) + ')'

    def ground(self, variable_dict: Dict[str, str], state: Optional[SStateCompatible] = None):
        # TODO(Jiayuan Mao @ 2023/04/27): include the state in the grounding.
        return AtomicStripsOperatorApplier(self, variable_dict)

    def __call__(self, *args: str):
        assert len(args) == len(self.arguments), f'Expected {len(self.arguments)} arguments, got {len(args)}.'
        bound_arguments = {arg.name: argv for arg, argv in zip(self.arguments, args)}
        return self.ground(bound_arguments)


class AtomicStripsOperatorApplier(object):
    def __init__(self, operator: AtomicStripsOperator, bound_arguments: Dict[str, str]):
        self.operator = operator
        self.arguments = tuple(bound_arguments[arg.name] for arg in operator.arguments)
        self.bound_arguments = bound_arguments

        self.preconditions = tuple(precondition.ground(bound_arguments, return_proposition=True) for precondition in operator.preconditions)
        self.add_effects = tuple(add_effect.ground(bound_arguments, return_proposition=True) for add_effect in operator.add_effects)
        self.del_effects = tuple(del_effect.ground(bound_arguments, return_proposition=True) for del_effect in operator.del_effects)

    operator: AtomicStripsOperator
    """The operator that this applier is derived from."""

    arguments: Tuple[str, ...]
    """The arguments of the operator, as a tuple of strings."""

    bound_arguments: Dict[str, str]
    """The bound arguments of the operator."""

    preconditions: Tuple[SProposition, ...]
    """The preconditions of the operator, as a tuple of :class:`~concepts.pdsketch.strips.strips_expression.SProposition`."""

    add_effects: Tuple[SProposition, ...]
    """The add effects of the operator, as a tuple of :class:`~concepts.pdsketch.strips.strips_expression.SProposition`."""

    del_effects: Tuple[SProposition, ...]
    """The delete effects of the operator, as a tuple of :class:`~concepts.pdsketch.strips.strips_expression.SProposition`."""

    def __str__(self) -> str:
        return f'{self.operator.name}({", ".join(f"{arg.name}={self.bound_arguments[arg.name]}" for arg in self.operator.arguments)})'

    __repr__ = jacinle.repr_from_str


class AtomicStripsOperatorApplicationExpression(object):
    def __init__(self, operator: AtomicStripsOperator, arguments: Sequence[Variable]):
        self.operator = operator
        self.arguments = tuple(arguments)

    operator: AtomicStripsOperator
    """The operator to be applied."""

    arguments: Tuple[Variable, ...]
    """The arguments of the operator."""

    def __str__(self) -> str:
        return f'{self.operator.name}({", ".join(arg.name if isinstance(arg, Variable) else arg for arg in self.arguments)})'

    __repr__ = jacinle.repr_from_str

    @classmethod
    def from_operator_application_expression(cls, expression: OperatorApplicationExpression, atomic_strips_operators: Mapping[str, AtomicStripsOperator]) -> 'AtomicStripsOperatorApplicationExpression':
        """Create a strips operator application expression from a PDSketch operator application expression.

        Args:
            expression: the PDSketch operator application expression.

        Returns:
            the strips operator application expression.
        """
        operator = atomic_strips_operators[expression.operator.name]
        arguments = list()
        for arg in expression.arguments:
            assert isinstance(arg, VariableExpression), f'Not supported argument expression: {arg}.'
            arguments.append(arg.variable)
        return cls(operator, arguments)

    def ground(self, variable_dict: Dict[str, str], state: Optional[SStateCompatible] = None):
        return AtomicStripsOperatorApplier(self.operator, {arg.name: variable_dict[arg.name] for arg in self.arguments})

    def __call__(self, *args):
        assert len(args) == len(self.arguments), f'Expected {len(self.arguments)} arguments, got {len(args)}.'
        bound_arguments = {arg.name: argv for arg, argv in zip(self.arguments, args)}
        return self.ground(bound_arguments)


class AtomicStripsAchieveExpression(object):
    def __init__(self, goal: SBoolPredicateApplicationExpression, maintains: Sequence[SBoolPredicateApplicationExpression]):
        self.goal = goal
        self.maintains = tuple(maintains)

    goal: SBoolPredicateApplicationExpression
    """The goal of the achieve expression."""

    maintains: Tuple[SBoolPredicateApplicationExpression, ...]
    """A list of expression to be maintained."""

    def __str__(self) -> str:
        return f'achieve({self.goal}, maintains={{{self.maintains}}})'

    __repr__ = jacinle.repr_from_str

    @classmethod
    def from_achieve_expression(cls, expression: AchieveExpression) -> 'AtomicStripsAchieveExpression':
        """Create a strips achieve expression from a PDSketch achieve expression.

        Args:
            expression: the PDSketch achieve expression.

        Returns:
            the strips achieve expression.
        """
        goal = SBoolPredicateApplicationExpression.from_function_application_expression(expression.goal)
        maintains = tuple(SBoolPredicateApplicationExpression.from_function_application_expression(maintain) for maintain in expression.maintains)
        return cls(goal, maintains)

    def ground(self, variable_dict: Dict[str, str], state: Optional[SStateCompatible] = None):
        return AtomicStripsGroundedAchieveExpression(
            self.goal.ground(variable_dict, return_proposition=True),
            tuple(maintain.ground(variable_dict, return_proposition=True) for maintain in self.maintains)
        )

    def __call__(self, *args):
        assert len(args) == len(self.arguments), f'Expected {len(self.arguments)} arguments, got {len(args)}.'
        bound_arguments = {arg.name: argv for arg, argv in zip(self.arguments, args)}
        return self.ground(bound_arguments)


class AtomicStripsGroundedAchieveExpression(object):
    def __init__(self, goal: SProposition, maintains: Sequence[SProposition]):
        self.goal = goal
        self.maintains = tuple(maintains)

    goal: SProposition
    """The goal of the achieve expression."""

    maintains: Tuple[SProposition, ...]
    """A list of expression to be maintained."""

    def __str__(self) -> str:
        return f'achieve({self.goal}, maintains={{{", ".join(map(str, self.maintains))}}})'

    __repr__ = jacinle.repr_from_str


class AtomicStripsRegressionRule(object):
    def __init__(
        self,
        arguments: Sequence[Variable],
        preconditions: Sequence[SBoolPredicateApplicationExpression],
        preconstraints: Sequence[SBoolPredicateApplicationExpression],
        goal: SBoolPredicateApplicationExpression,
        body: Sequence[Union[AtomicStripsOperatorApplicationExpression, AtomicStripsAchieveExpression]],
        raw_regression_rule: RegressionRule,
    ):
        self.name = raw_regression_rule.name
        self.arguments = tuple(arguments)
        self.preconditions = tuple(preconditions)
        self.preconstraints = tuple(preconstraints)
        self.goal = goal
        self.body = tuple(body)
        self.raw_regression_rule = raw_regression_rule

    arguments: Tuple[Variable, ...]
    """The arguments of the regression rule."""

    preconditions: Tuple[SBoolPredicateApplicationExpression, ...]
    """The preconditions of the regression rule."""

    preconstraints: Tuple[SBoolPredicateApplicationExpression, ...]
    """The preconstraints of the regression rule."""

    goal: SBoolPredicateApplicationExpression
    """The goal of the regression rule."""

    body: Tuple[Union[AtomicStripsOperatorApplicationExpression, AtomicStripsAchieveExpression], ...]
    """The body of the regression rule."""

    raw_regression_rule: RegressionRule
    """The raw regression rule."""

    def __str__(self) -> str:
        return f'{self.goal} <- {", ".join(str(item) for item in self.body)}'

    __repr__ = jacinle.repr_from_str

    @classmethod
    def from_regression_rule(cls, regression_rule: RegressionRule, atomic_strips_operators: Mapping[str, AtomicStripsOperator]):
        arguments = regression_rule.arguments
        preconditions = tuple(SBoolPredicateApplicationExpression.from_function_application_expression(precondition.bool_expr) for precondition in regression_rule.preconditions)
        preconstraints = tuple(SBoolPredicateApplicationExpression.from_function_application_expression(preconstraint.bool_expr) for preconstraint in regression_rule.preconstraints)

        assert len(regression_rule.goal) == 1, f'Expected exactly one goal, got {len(regression_rule.goal)}.'
        goal_expression = regression_rule.goal[0].assign_expr
        assert isinstance(goal_expression, AssignExpression), f'Expected an assign expression, got {type(goal_expression)}.'
        assert isinstance(goal_expression.value, ConstantExpression), f'Expected a constant expression for the goal expression assignment, got {type(goal_expression.value)}.'
        value = goal_expression.value.constant.item()
        assert isinstance(value, int), f'Expected an integer for the goal expression assignment, got {type(value)}.'
        sassign = SBoolPredicateApplicationExpression.from_function_application_expression(goal_expression.predicate, negated=value < 0.5)

        body = []
        for body_expression in regression_rule.body:
            if isinstance(body_expression, OperatorApplicationExpression):
                body.append(AtomicStripsOperatorApplicationExpression.from_operator_application_expression(body_expression, atomic_strips_operators))
            elif isinstance(body_expression, AchieveExpression):
                body.append(AtomicStripsAchieveExpression.from_achieve_expression(body_expression))
            else:
                raise NotImplementedError(f'Not supported body expression: {body_expression}.')

        return cls(arguments, preconditions, preconstraints, sassign, body, regression_rule)

    def ground(self, variable_dict: Dict[str, str], state: Optional[SStateCompatible] = None):
        # TODO(Jiayuan Mao @ 2023/07/12): include the state in the grounding.
        return AtomicStripsRegressionRuleApplier(self, variable_dict)

    def __call__(self, *args: str):
        assert len(args) == len(self.arguments), f'Expected {len(self.arguments)} arguments, got {len(args)}.'
        bound_arguments = {arg.name: argv for arg, argv in zip(self.arguments, args)}
        return self.ground(bound_arguments)


class AtomicStripsRegressionRuleApplier(object):
    def __init__(self, regression_rule: AtomicStripsRegressionRule, bound_arguments: Dict[str, str]):
        self.regression_rule = regression_rule
        self.arguments = tuple(bound_arguments[arg.name] for arg in regression_rule.arguments)
        self.bound_arguments = bound_arguments

        self.preconditions = tuple(precondition.ground(bound_arguments, return_proposition=True) for precondition in regression_rule.preconditions)
        self.preconstraints = tuple(preconstraint.ground(bound_arguments, return_proposition=True) for preconstraint in regression_rule.preconstraints)
        self.goal = regression_rule.goal.ground(bound_arguments, return_proposition=True)
        self.body = tuple(item.ground(bound_arguments) for item in regression_rule.body)

    regression_rule: AtomicStripsRegressionRule
    """The regression rule."""

    arguments: Tuple[str, ...]
    """The arguments of the operator, as a tuple of strings."""

    bound_arguments: Dict[str, str]
    """The bound arguments of the operator."""

    def __str__(self) -> str:
        return f'{self.regression_rule.name}({", ".join(f"{arg}={argv}" for arg, argv in self.bound_arguments.items())})'

    __repr__ = jacinle.repr_from_str


class AtomicStripsDomain(object):
    """The domain of the atomic STRIPS planning problem."""

    def __init__(self, types: Dict[str, ObjectType], predicates: Dict[str, Predicate], operators: Dict[str, AtomicStripsOperator], regression_rules: Dict[str, AtomicStripsRegressionRule], constants: Dict[str, ObjectConstant]):
        self.types = types
        self.predicates = predicates
        self.operators = operators
        self.regression_rules = regression_rules
        self.constants = constants

    types: Dict[str, ObjectType]
    """The types of the domain, as a dictionary from type name to :class:`~concepts.pdsketch.strips.strips_expression.ObjectType`."""

    predicates: Dict[str, Predicate]
    """The predicates of the domain, as a dictionary from predicate name to :class:`~concepts.pdsketch.predicate.Predicate`."""

    operators: Dict[str, AtomicStripsOperator]
    """The operators of the domain, as a dictionary from operator name to :class:`~concepts.pdsketch.strips.atomic_strips_domain.AtomicStripsOperator."""

    regression_rules: Dict[str, AtomicStripsRegressionRule]
    """The regression rules of the domain, as a dictionary from regression rule name to :class:`~concepts.pdsketch.strips.atomic_strips_domain.AtomicStripsRegressionRule`."""

    constants: Dict[str, ObjectConstant]
    """The constants of the domain, as a dictionary from constant name to :class:`~concepts.dsl.dsl_types.ObjectConstant`."""

    @classmethod
    def from_domain(cls, domain: Domain) -> 'AtomicStripsDomain':
        operators = dict()
        for operator in domain.operators.values():
            if isinstance(operator, Operator):
                if operator.is_macro:
                    continue
                # TODO(Jiayuan Mao @ 2023/03/19): support macro operator here.
                strips_operator = AtomicStripsOperator.from_operator(operator)
                operators[operator.name] = strips_operator
        regression_rules = dict()
        for regression_rule in domain.regression_rules.values():
            strips_regression_rule = AtomicStripsRegressionRule.from_regression_rule(regression_rule, operators)
            regression_rules[regression_rule.name] = strips_regression_rule

        return cls(
            types=domain.types.copy(),
            predicates=domain.functions.copy(),
            operators=operators,
            regression_rules=regression_rules,
            constants=domain.constants.copy(),
        )


class AtomicStripsProblem(object):
    def __init__(self, domain: AtomicStripsDomain, objects: Dict[str, str], initial_state: SStateCompatible, conjunctive_goal: Sequence[SProposition]):
        self.domain = domain
        self.objects = objects
        self.initial_state = SState(initial_state)
        self.conjunctive_goal = tuple(conjunctive_goal)

    domain: AtomicStripsDomain
    """The domain of the problem, as a :class:`~concepts.pdsketch.strips.strips_domain.AtomicStripsDomain`."""

    objects: Dict[str, str]
    """The objects of the problem, as a dictionary from object name to object type."""

    initial_state: SState
    """The initial state of the problem, as a :class:`~concepts.pdsketch.strips.strips_expression.SState`."""

    conjunctive_goal: Tuple[SProposition, ...]
    """The conjunctive goal of the problem, as a tuple of :class:`~concepts.pdsketch.strips.strips_expression.SProposition`."""

    @cached_property
    def objects_type2names(self) -> Dict[str, Tuple[str, ...]]:
        objects_type2names = dict()
        for name, type_ in self.objects.items():
            objects_type2names.setdefault(type_, []).append(name)
        return {type_: tuple(names) for type_, names in objects_type2names.items()}

    @classmethod
    def from_domain_and_problem(cls, domain: AtomicStripsDomain, problem: Problem) -> 'AtomicStripsProblem':
        objects = problem.objects.copy()
        for _, typename in objects.items():
            if typename not in domain.types:
                raise ValueError(f'Unknown type {typename}.')

        initial_state = list()
        for proposition in problem.predicates:
            if proposition.function.name not in domain.predicates:
                raise ValueError(f'Unknown predicate {proposition.function.name}.')
            initial_state.append(make_sproposition_from_function_application(proposition, objects))

        conjunctive_goal = list()
        if not is_and_expr(problem.goal):
            raise ValueError(f'Expected conjunctive goal, got {problem.goal}.')
        for proposition in problem.goal.arguments:
            if not isinstance(proposition, FunctionApplicationExpression):
                raise ValueError(f'Expected proposition, got {proposition}.')
            if proposition.function.name not in domain.predicates:
                raise ValueError(f'Unknown predicate {proposition.function.name}.')
            conjunctive_goal.append(make_sproposition_from_function_application(proposition, objects))

        return cls(
            domain,
            objects=objects,
            initial_state=initial_state,
            conjunctive_goal=conjunctive_goal,
        )


def load_astrips_domain_file(filename: str, return_raw_domain: bool = False) -> Union[AtomicStripsDomain, Tuple[Domain, AtomicStripsDomain]]:
    """Load a strips domain from a domain file.

    Args:
        filename: the domain file.
        return_raw_domain: whether to return the raw domain.

    Returns:
        the strips domain, or a tuple of the raw domain and the strips domain.
    """
    domain = load_domain_file(filename)
    strips_domain = AtomicStripsDomain.from_domain(domain)
    if return_raw_domain:
        return domain, strips_domain
    return strips_domain


def load_astrips_domain_string(domain_string: str, return_raw_domain: bool = False) -> Union[AtomicStripsDomain, Tuple[Domain, AtomicStripsDomain]]:
    """Load a strips domain from a domain string.

    Args:
        domain_string: the domain string.
        return_raw_domain: whether to return the raw domain.

    Returns:
        the strips domain, or a tuple of the raw domain and the strips domain.
    """
    domain = load_domain_string(domain_string)
    strips_domain = AtomicStripsDomain.from_domain(domain)
    if return_raw_domain:
        return domain, strips_domain
    return strips_domain


def load_astrips_problem_file(domain: Domain, astrips_domain: AtomicStripsDomain, filename: str) -> AtomicStripsProblem:
    """Load a strips problem from a problem file.

    Args:
        domain: the domain of the problem.
        filename: the problem file.

    Returns:
        the strips problem.
    """
    problem = load_problem_file(filename, domain=domain, return_tensor_state=False)
    return AtomicStripsProblem.from_domain_and_problem(astrips_domain, problem)


def load_astrips_problem_string(domain: Domain, astrips_domain: AtomicStripsDomain, problem_string: str) -> AtomicStripsProblem:
    """Load a strips problem from a problem string.

    Args:
        domain: the domain of the problem.
        problem_string: the problem string.

    Returns:
        the strips problem.
    """
    problem = load_problem_string(problem_string, domain=domain, return_tensor_state=False)
    return AtomicStripsProblem.from_domain_and_problem(astrips_domain, problem)


concepts/pdsketch/strips/atomic_strips_domain.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : atomic_strips_onthefly_search.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 07/14/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import time
from typing import Optional, Union, Iterator, Tuple, List, Dict
from collections import defaultdict, deque

from concepts.dsl.dsl_types import Variable
from concepts.pdsketch.strips.strips_expression import SStateDict, SBoolPredicateApplicationExpression
from concepts.pdsketch.strips.atomic_strips_domain import AtomicStripsDomain, AtomicStripsProblem, AtomicStripsOperator, AtomicStripsOperatorApplier


def _bind_arguments(predicate: SBoolPredicateApplicationExpression, bound_arguments: Dict[str, Union[int, str]]):
    return predicate.name, tuple(bound_arguments[arg.name] if isinstance(arg, Variable) else arg for arg in predicate.arguments)


def _gen_applicable_actions(domain: AtomicStripsDomain, objects: Dict[str, List[str]], state: SStateDict, check_negation: bool = False) -> Iterator[Tuple[AtomicStripsOperator, Dict[str, str]]]:
    TOO_MANY, FAILED, PASS = object(), object(), object()

    def compute_possible_grounding(predicate: SBoolPredicateApplicationExpression, bound_arguments: Dict[str, Union[int, str]]):
        unbound_arguments = [arg for arg in predicate.arguments if isinstance(arg, Variable) and arg.name not in bound_arguments]
        if len(unbound_arguments) == 0:
            name, arguments = _bind_arguments(predicate, bound_arguments)
            rv = state.contains(name, arguments, predicate.negated, check_negation=check_negation)
            if not rv:
                return '', FAILED
            return '', PASS
        elif len(unbound_arguments) == 1:
            arg = unbound_arguments[0]
            valid_arguments = list()

            options = objects[arg.typename]
            for o in options:
                bound_arguments[arg.name] = o
                name, arguments = _bind_arguments(predicate, bound_arguments)
                rv = state.contains(name, arguments, predicate.negated, check_negation=check_negation)
                if rv:
                    valid_arguments.append(o)
                del bound_arguments[arg.name]
            return arg.name, valid_arguments
        else:
            return '', TOO_MANY

    # @jacinle.log_function(verbose=False)
    def dfs(preconditions: Tuple[SBoolPredicateApplicationExpression, ...], bound_arguments: Dict[str, int]):
        """Inner DFS function.

        Args:
            preconditions: the preconditions to be satisfied.
            bound_arguments: a mapping from variable name to object.
        """

        # jacinle.log_function.print('dfs', bound_arguments, 'remaining preconditions:', len(preconditions))
        # import ipdb; ipdb.set_trace()

        for i, precondition in enumerate(preconditions):
            name, valid_arguments = compute_possible_grounding(precondition, bound_arguments)
            if valid_arguments == FAILED:
                # jacinle.log_function.print('Failed.')
                return list()
            elif valid_arguments == PASS:
                # jacinle.log_function.print('Pass.')
                return dfs(preconditions[:i] + preconditions[i + 1:], bound_arguments)
            elif valid_arguments == TOO_MANY:
                pass
            else:
                outputs = list()
                for arg in valid_arguments:
                    bound_arguments[name] = arg
                    outputs.extend(dfs(preconditions[:i] + preconditions[i + 1:], bound_arguments))
                    del bound_arguments[name]
                return outputs

        unbound_arguments = [arg for arg in operator.arguments if isinstance(arg, Variable) and arg.name not in bound_arguments]
        # print('unbound_arguments', unbound_arguments, bound_arguments)
        if len(unbound_arguments) == 0:
            # jacinle.log_function.print('Found a grounding:', bound_arguments)
            return [bound_arguments.copy()]

        unbound_arguments_possible_values = {arg.name: objects[arg.typename] for arg in unbound_arguments}

        name, valid_arguments = min(unbound_arguments_possible_values.items(), key=lambda x: len(x[1]))
        outputs = list()
        for arg in valid_arguments:
            bound_arguments[name] = arg
            # jacinle.log_function.print('{} = {}'.format(name, arg))
            outputs.extend(dfs(preconditions, bound_arguments))
        del bound_arguments[name]
        return outputs

    for operator in domain.operators.values():
        # jacinle.log_function.print(f'operator: {operator.name}')
        for bound_arguments in dfs(operator.preconditions, dict()):
            # jacinle.log_function.print('yield bound_arguments:', bound_arguments)
            yield operator, bound_arguments


def _check_precondition(state: SStateDict, operator: AtomicStripsOperator, bound_arguments: Dict[str, Union[int, str]]):
    for precondition in operator.preconditions:
        name, arguments = _bind_arguments(precondition, bound_arguments)
        if not state.contains(name, arguments, precondition.negated):
            return False
    return True


def _apply_operator(state: SStateDict, operator: AtomicStripsOperator, bound_arguments: Dict[str, Union[int, str]]):
    new_state = state.clone()
    for predicate in operator.del_effects:
        name, arguments = _bind_arguments(predicate, bound_arguments)
        new_state.remove(name, arguments)
    for predicate in operator.add_effects:
        name, arguments = _bind_arguments(predicate, bound_arguments)
        new_state.add(name, arguments)
    return new_state


def _ground_actions(actions: Tuple[Tuple[AtomicStripsOperator, Dict[str, str]], ...]) -> Tuple[AtomicStripsOperatorApplier, ...]:
    ground_operators = list()
    for operator, bound_arguments in actions:
        ground_operators.append(operator.ground(bound_arguments))
    return tuple(ground_operators)


def astrips_onthefly_search(domain: AtomicStripsDomain, problem: AtomicStripsProblem, verbose: bool = False, timeout: float = 300, max_expanded_nodes: int = 1000000) -> Optional[Tuple[AtomicStripsOperatorApplier, ...]]:
    objects = defaultdict(list)
    object2index = dict()

    for name, constant in domain.constants.items():
        objects[constant.dtype.typename].append(name)
        object2index[name] = len(object2index) - 1
    for name, typename in problem.objects.items():
        objects[typename].append(name)
        object2index[name] = len(object2index) - 1

    initial_state = SStateDict()
    for predicate in problem.initial_state:
        name, *args = predicate.split()
        initial_state.add(name, args)

    goal_conditions = problem.conjunctive_goal

    frontier = deque()
    frontier.append((initial_state, tuple()))
    explored = set()

    start_time = time.time()
    nr_expanded_nodes = 0
    while len(frontier) > 0:
        nr_expanded_nodes += 1
        if nr_expanded_nodes > max_expanded_nodes:
            break

        if nr_expanded_nodes % 100 == 0:
            if time.time() - start_time > timeout:
                print('astrips_onthefly_search::Timeout.')
                break

        state, plan = frontier.popleft()

        # action_strings = [f"{operator.name}({', '.join(bound_arguments.values())})" for operator, bound_arguments in plan]
        # print('State', state, 'Plan', action_strings)
        # print('Plan', action_strings)

        for operator, bound_arguments in _gen_applicable_actions(domain, objects, state):
            new_state = _apply_operator(state, operator, bound_arguments)
            new_state_set = new_state.as_state()
            if new_state_set not in explored:
                if new_state_set.issuperset(goal_conditions):
                    return _ground_actions(plan + ((operator, bound_arguments), ))

                frontier.append((new_state, plan + ((operator, bound_arguments), )))
                explored.add(new_state_set)

    return None


concepts/pdsketch/strips/atomic_strips_onthefly_search.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : atomic_strips_regression_search.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 07/13/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import itertools
from typing import List, Set

import jacinle

from concepts.pdsketch.strips.strips_expression import SProposition
from concepts.pdsketch.strips.atomic_strips_domain import (
    AtomicStripsDomain, AtomicStripsProblem,
    AtomicStripsOperator, AtomicStripsRegressionRule,
    AtomicStripsOperatorApplier, AtomicStripsRegressionRuleApplier,
    AtomicStripsGroundedAchieveExpression
)


def gen_all_grounded_actions_and_rules(domain: AtomicStripsDomain, problem: AtomicStripsProblem, mode: str) -> List[AtomicStripsOperatorApplier]:
    assert mode in ('action', 'rule')
    all_actions = list()

    if mode == 'action':
        operator_list = domain.operators
    elif mode == 'rule':
        operator_list = domain.regression_rules
    else:
        raise ValueError('Unknown mode: {}.'.format(mode))

    for operator in operator_list.values():
        candidate_arguments = [
            problem.objects_type2names[operator.arguments[i].dtype.typename]
            for i in range(len(operator.arguments))
        ]
        for arg_list in itertools.product(*candidate_arguments):
            action = operator(*arg_list)
            static_check = True
            for i, prop in enumerate(action.preconditions):
                if domain.predicates[operator.preconditions[i].name].is_static:
                    if prop not in problem.initial_state:
                        static_check = False
                        break
            if not static_check:
                continue
            all_actions.append(action)

    return all_actions


def astrips_regression_search_1(domain: AtomicStripsDomain, problem: AtomicStripsProblem, verbose: bool = False) -> List[AtomicStripsOperatorApplier]:
    """Search for a plan using regression search."""

    assert problem.conjunctive_goal is not None, "Only conjunctive goals are supported."
    assert len(problem.conjunctive_goal) == 1, "Only single-goal problems are supported."

    for operator in domain.operators.values():
        for precondition in operator.preconditions:
            if precondition.negated:
                raise NotImplementedError('astrips_regression_search does not support negated preconditions.')
    for regression_rule in domain.regression_rules.values():
        for precondition in regression_rule.preconditions:
            if precondition.negated:
                raise NotImplementedError('astrips_regression_search does not support negated preconditions.')
        if len(regression_rule.preconstraints) > 0:
            raise NotImplementedError('astrips_regression_search does not support preconstraints.')

    all_rules = gen_all_grounded_actions_and_rules(domain, problem, 'rule')

    def find_applicable_rules(state: Set[SProposition], goal: SProposition, maintains: Set[SProposition]):
        applicable_rules = list()
        for rule in all_rules:
            if rule.goal == goal and state.issuperset(rule.preconditions):
                applicable_rules.append(rule)
        assert len(applicable_rules) == 1, "Only one applicable rule is allowed."
        return applicable_rules[0]

    @jacinle.log_function(verbose=False)
    def dfs(state: Set[SProposition], goal: SProposition, maintains: Set[SProposition]):
        rule = find_applicable_rules(state, goal, maintains)
        if verbose:
            jacinle.log_function.print(f'Current state: {state}, goal: {goal}, maintains: {maintains} => rule: {rule}')
        actions = list()
        for item in rule.body:
            if isinstance(item, AtomicStripsGroundedAchieveExpression):
                if item.goal not in state:
                    if verbose:
                        jacinle.log_function.print(f'{str(item)}')
                    state, sub_actions = dfs(state, item.goal, maintains.union(item.maintains))
                    actions.extend(sub_actions)
                else:
                    if verbose:
                        jacinle.log_function.print(f'{str(item)} (skipped)')
            elif isinstance(item, AtomicStripsOperatorApplier):
                if verbose:
                    jacinle.log_function.print(f'do({str(item)})')
                actions.append(item)
                assert state.issuperset(item.preconditions), "The preconditions of an action should be satisfied."
                state = (state - frozenset(item.del_effects)).union(frozenset(item.add_effects))
            else:
                raise ValueError('Unknown item: {}.'.format(item))
        return state, actions

    end_state, actions = dfs(problem.initial_state, problem.conjunctive_goal[0], set())
    return actions


concepts/pdsketch/strips/atomic_strips_regression_search.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : strips_expression.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 04/26/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""
This file defines a set of classes for representing STRIPS expressions, including:

- Boolean constant expressions.
- Boolean predicate expressions.
- SAS predicate expressions.
- And/Or expressions.
- Not expressions.
- Forall/Exists expressions.
- SAS expressions.
- Single predicate assignment expressions.
- Conditional assignment expressions.
- Deictic assignment expressions.

At the highest level, the STRIPS expressions are categorized into two types:

- StripsValueOutputExpression, which outputs a value. There is a special instantiation of this type, StripsBooleanOutputExpression, which outputs a boolean value.
- StripsVariableAssignmentExpression, which assigns a value to a state variable.
"""

import jacinle
from copy import deepcopy
from abc import abstractmethod, ABC
from typing import Optional, Union, Iterable, Sequence, Tuple, Set, FrozenSet, Dict

from concepts.dsl.dsl_types import Variable, ObjectConstant, BOOL
from concepts.dsl.expression import FunctionApplicationExpression, VariableExpression, ObjectConstantExpression, BoolExpression, BoolOpType

__all__ = [
    'SPredicateName', 'SProposition', 'make_sproposition',
    'SState', 'SStateCompatible', 'SStateDict',
    'SExpression', 'SValueOutputExpression', 'SVariableAssignmentExpression', 'SBoolOutputExpression',
    'SBoolConstant', 'SBoolPredicateApplicationExpression', 'SSASPredicateApplicationExpression',
    'SSimpleBoolExpression', 'SBoolNot', 'SQuantificationExpression', 'SSASExpression',
    'SAssignExpression', 'SConditionalAssignExpression', 'SDeicticAssignExpression'
]

"""The name of predicates, represented as strings."""
SPredicateName = str

"""The name of propositions. A proposition is a predicate grounded on a set of arguments."""
SProposition = str

# """The representation of a SAS proposition. It is a tuple of (predicate name, SAS value)."""
# StripsSASProposition = Tuple[str, int]


def _variable_or_constant_to_str(x: Union[str, Variable, ObjectConstant]) -> str:
    if isinstance(x, str):
        return x
    elif isinstance(x, int):
        return str(x)
    elif isinstance(x, Variable):
        return x.name
    elif isinstance(x, ObjectConstant):
        return x.name
    else:
        raise TypeError(f'Unknown type: {type(x)}.')


def make_sproposition(name: SPredicateName, *args: Union[str, Variable, ObjectConstantExpression]) -> SProposition:
    """
    Compose a proposition from a predicate name and a list of arguments.
    """
    if len(args) == 0:
        return name
    return '{} {}'.format(name, ' '.join(_variable_or_constant_to_str(x) for x in args))


def make_sproposition_from_function_application(expr: FunctionApplicationExpression, objects: Optional[Dict[str, str]] = None) -> SProposition:
    """
    Compose a proposition from a function application expression.
    """
    name = expr.function.name

    arguments = list()
    for arg in expr.arguments:
        if not isinstance(arg, ObjectConstantExpression):
            raise ValueError(f'Expected object constant, got {arg}.')
        if objects is not None and arg.name not in objects:
            raise ValueError(f'Unknown object {arg.name}.')
        assert arg.name in objects, f'Unknown object {arg.name}.'
        arguments.append(arg.name)

    return make_sproposition(name, *arguments)


class SState(frozenset, FrozenSet[SProposition]):
    """The representation of a STRIPS state, which is a set of propositions."""
    pass


class SStateDict(dict, Dict[str, Set[Tuple[Union[int, str], ...]]]):
    def add(self, predicate_name: SPredicateName, arguments: Sequence[Union[int, str]]):
        if predicate_name not in self:
            self[predicate_name] = set()
        self[predicate_name].add(tuple(arguments))

    def remove(self, predicate_name: SPredicateName, arguments: Sequence[Union[int, str]]):
        if predicate_name in self:
            self[predicate_name].discard(tuple(arguments))

    def contains(self, predicate_name: SPredicateName, arguments: Sequence[Union[int, str]], negated: bool = False, check_negation: bool = False) -> bool:
        """Check whether the state contains the given proposition.

        Args:
            predicate_name: the name of the predicate.
            arguments: the arguments of the predicate, as a tuple of integers or strings.
            negated: whether the proposition is negated. If True, the function will check whether the state does not contain the proposition.
            check_negation: whether the function should also check "{predicate_name}_not" in the state. This will only be used when `negated` is True.
                This is useful for delete-relaxed planning.

        Returns:
            True if the state contains the proposition, False otherwise.
        """
        if not check_negation:
            if predicate_name in self:
                return (arguments in self[predicate_name]) ^ negated
            return negated  # if the predicate is not in the state and it is not negated, we return False.
        else:
            if not negated:
                if predicate_name in self:
                    return arguments in self[predicate_name]
                return False
            else:
                true_set = self.get(predicate_name, None)
                false_set = self.get(f'{predicate_name}_not', None)
                return (false_set is not None and arguments in false_set) or (true_set is None) or (true_set is not None and arguments not in true_set)

    def clone(self):
        return deepcopy(self)

    def as_state(self) -> SState:
        return SState([f'{predicate_name} {" ".join(map(str, arguments))}' for predicate_name, list_of_arguments in self.items() for arguments in list_of_arguments])


SStateCompatible = Union[SState, Set[SProposition]]


# class StripsSASState(dict, Dict[Tuple[StripsPredicateName, str], int]):
#     """The representation of a SAS state, which is a mapping from (predicate name, SAS value) to the number of occurrences."""
#     """StripsSASState is a mapping from (predicate name, arguments_str) to value."""
#     pass


class SExpression(ABC):
    """The base class for STRIPS expressions."""
    __repr__ = jacinle.repr_from_str

    def ground(self, variable_dict: Dict[str, str], state: Optional[SStateCompatible] = None):
        """Return a new expression with all variables grounded according to the given variable dictionary."""
        raise NotImplementedError()

    @abstractmethod
    def iter_precondition_predicates(self) -> Iterable[SPredicateName]:
        """Iterate over the precondition predicate names in the expression."""
        raise NotImplementedError()

    @abstractmethod
    def iter_effect_predicates(self) -> Iterable[SPredicateName]:
        """Iterate over the effect predicate names in the expression."""
        raise NotImplementedError()


class SValueOutputExpression(SExpression, ABC):
    """The base class for STRIPS expressions that output a value."""

    @abstractmethod
    def iter_precondition_predicates(self) -> Iterable[SPredicateName]:
        """Iterate over the precondition predicate names in the expression."""
        raise NotImplementedError()

    def iter_effect_predicates(self) -> Iterable[SPredicateName]:
        """Iterate over the effect predicate names in the expression."""
        return set()


class SVariableAssignmentExpression(SExpression, ABC):
    """The base class for STRIPS expressions that assign a value to a state variable."""

    def iter_precondition_predicates(self) -> Iterable[SPredicateName]:
        """Iterate over the precondition predicate names in the expression."""
        return set()

    @abstractmethod
    def iter_effect_predicates(self) -> Iterable[SPredicateName]:
        """Iterate over the effect predicate names in the expression."""
        raise NotImplementedError()


class SBoolOutputExpression(SValueOutputExpression, ABC):
    """The base class for STRIPS expressions that output a boolean value."""
    pass


class SBoolConstant(SBoolOutputExpression):
    """The representation of a boolean constant."""

    def __init__(self, constant: bool):
        """Initialize a boolean constant.

        Args:
            constant: the value of the constant.
        """

        self.constant = constant

    def ground(self, variable_dict: Dict[str, str], state: Optional[SStateCompatible] = None):
        from concepts.pdsketch.strips.strips_grounded_expression import GSBoolConstantExpression
        return GSBoolConstantExpression(self.constant)

    def iter_precondition_predicates(self) -> Iterable[SPredicateName]:
        return set()

    def __str__(self) -> str:
        return 'true' if self.constant else 'false'


class SBoolPredicateApplicationExpression(SBoolOutputExpression):
    """The base class for STRIPS expressions that output a boolean value based on a predicate."""

    def __init__(self, name: SPredicateName, negated: bool, arguments: Sequence[Union[Variable, str]]):
        """Initialize a boolean predicate expression.

        Args:
            name: the name of the predicate.
            negated: whether the predicate is negated.
            arguments: the arguments of the predicate. Either variables or str (constants).
        """

        self.name = name
        self.arguments = tuple(arguments)
        self.negated = negated

    name: SPredicateName
    """The name of the predicate."""

    arguments: Tuple[Union[Variable, str]]
    """The arguments of the predicate."""

    negated: bool
    """Whether the predicate is negated."""

    def ground(self, variable_dict: Dict[str, str], state: Optional[SStateCompatible] = None, negated: bool = False, return_proposition: bool = False):
        """Ground the expression according to the given variable dictionary.

        Args:
            variable_dict: the variable dictionary.
            state: the state to ground the expression on. If None, the expression will be grounded without considering the state.
            negated: whether the predicate is negated.
            return_proposition: whether to return a SProposition instead of a GSSimpleBoolExpression

        Returns:
            the grounded expression. Will be a GSSimpleBoolExpression if `return_proposition` is False, otherwise a SProposition.
        """
        from concepts.pdsketch.strips.strips_grounded_expression import GSSimpleBoolExpression
        identifier = self.name + '_not' if (self.negated ^ negated) else self.name
        proposition = make_sproposition(identifier, *tuple(variable_dict[argument.name] if isinstance(argument, Variable) else argument for argument in self.arguments))

        if return_proposition:
            return proposition
        return GSSimpleBoolExpression({proposition})

    def iter_precondition_predicates(self) -> Iterable[SPredicateName]:
        return {self.name}

    def __str__(self) -> str:
        if len(self.arguments) == 0:
            fmt = f'({self.name})'
        else:
            argument_str = ' '.join(x.name if isinstance(x, Variable) else x for x in self.arguments)
            fmt = f'({self.name} {argument_str})'
        if self.negated:
            return f'(not {fmt})'
        return fmt

    @classmethod
    def from_function_application_expression(cls, expression: Union[FunctionApplicationExpression, BoolExpression], negated: bool = False):
        assert isinstance(expression, (FunctionApplicationExpression, BoolExpression)), f'Invalid expression type: {type(expression)}.'

        if isinstance(expression, BoolExpression):
            assert expression.bool_op is BoolOpType.NOT
            assert len(expression.arguments) == 1
            expression = expression.arguments[0]
            assert isinstance(expression, FunctionApplicationExpression)
            negated = not negated
            return cls.from_function_application_expression(expression, negated)

        assert expression.function.return_type == BOOL
        new_arguments = list()
        for arg in expression.arguments:
            if isinstance(arg, VariableExpression):
                new_arguments.append(arg.variable)
            elif isinstance(arg, ObjectConstantExpression):
                new_arguments.append(arg.constant.name)
            else:
                raise TypeError(f'Invalid argument type: {type(arg)}.')
        return cls(expression.function.name, negated, new_arguments)


class SSASPredicateApplicationExpression(SBoolPredicateApplicationExpression):
    """The representation for an SAS predicate expression. It is composed of a predicate name and an SAS index."""

    def __init__(self, sas_name: SPredicateName, sas_index: Optional[int], negated: bool, arguments: Sequence[Variable]):
        """Initialize an SAS predicate expression.

        Args:
            sas_name: the name of the SAS predicate.
            sas_index: the index of the SAS predicate.
            negated: whether the predicate is negated.
            arguments: the arguments of the predicate.
        """
        if sas_index is None:
            super().__init__(sas_name, negated, arguments)
        else:
            super().__init__(sas_name + '@' + str(sas_index), negated, arguments)

        self.sas_name = sas_name
        self.sas_index = sas_index

    name: SPredicateName
    arguments: Tuple[Variable]
    negated: bool

    sas_name: SPredicateName
    """The name of the SAS predicate."""

    sas_index: Optional[int]
    """The index of the SAS predicate. If None, it is a normal predicate."""


class SSimpleBoolExpression(SBoolOutputExpression):
    """The representation of a boolean expression. Note that since the negation is recorded in the raw :class:`StripsBoolPredicateApplicationExpression`,
    we do not need to record it here. Therefore, in this class, we only need to record whether the expression is an AND or an OR."""

    def __init__(self, arguments: Sequence[SBoolOutputExpression], is_disjunction: bool):
        """Initialize a boolean expression.

        Args:
            arguments: the arguments of the expression.
            is_disjunction: whether the expression is a disjunction.
        """
        self.arguments = arguments
        self.is_disjunction = is_disjunction

    arguments: Sequence[SBoolOutputExpression]
    """The arguments of the expression."""

    is_disjunction: bool
    """Whether the expression is a disjunction."""

    def ground(self, variable_dict: Dict[str, str], state: Optional[SStateCompatible] = None):
        from concepts.pdsketch.strips.strips_grounded_expression import gs_compose_bool_expressions
        return gs_compose_bool_expressions(
            [argument.ground(variable_dict) for argument in self.arguments],
            is_disjunction=self.is_disjunction,
        )

    @property
    def is_conjunction(self) -> bool:
        """Whether the expression is a conjunction."""
        return not self.is_disjunction

    def iter_precondition_predicates(self) -> Iterable[SPredicateName]:
        return set.union(*(x.iter_precondition_predicates() for x in self.arguments))

    def __str__(self) -> str:
        arguments_str = [str(arg) for arg in self.arguments]
        if sum(len(s) for s in arguments_str) > 120:
            arguments_str = [jacinle.indent_text(s) for s in arguments_str]
            fmt = '\n'.join(arguments_str)
            return f'(or\n{fmt}\n)' if self.is_disjunction else f'(and\n{fmt}\n)'
        return '({} {})'.format('or' if self.is_disjunction else 'and', ' '.join(arguments_str))


class SBoolNot(SBoolOutputExpression):
    """The representation of a boolean NOT expression. Note that this class is usually only used as a temporary expression during parsing.
    At the end, the negation is recorded in the raw :class:`StripsBoolPredicateApplicationExpression`."""

    def __init__(self, expr: SBoolOutputExpression):
        """Initialize a boolean NOT expression.

        Args:
            expr: the expression to be negated.
        """
        self.expr = expr

    expr: SBoolOutputExpression
    """The expression to be negated."""

    def ground(self, variable_dict: Dict[str, str], state: Optional[SStateCompatible] = None):
        if isinstance(self.expr, SBoolPredicateApplicationExpression):
            return self.expr.ground(variable_dict, negated=True)
        raise NotImplementedError()

    def iter_precondition_predicates(self) -> Iterable[SPredicateName]:
        return self.expr.iter_precondition_predicates()

    def __str__(self) -> str:
        return '(not {})'.format(str(self.expr))


class SQuantificationExpression(SBoolOutputExpression):
    """The representation of a quantification expression."""

    def __init__(self, variable: Variable, expr: SBoolOutputExpression, is_disjunction: bool):
        """Initialize a quantification expression.

        Args:
            variable: the variable to be quantified.
            expr: the expression to be quantified.
            is_disjunction: whether the expression is a disjunction (EXISTS quantification).
        """

        self.variable = variable
        self.expr = expr
        self.is_disjunction = is_disjunction

    variable: Variable
    """The variable to be quantified."""

    expr: SBoolOutputExpression
    """The expression to be quantified."""

    is_disjunction: bool
    """Whether the expression is a disjunction (EXISTS quantification)."""

    @property
    def is_conjunction(self) -> bool:
        """Whether the expression is a conjunction (FORALL quantification)."""
        return not self.is_disjunction

    @property
    def is_forall(self) -> bool:
        """Whether the expression is a conjunction (FORALL quantification)."""
        return not self.is_disjunction

    @property
    def is_exists(self) -> bool:
        """Whether the expression is a disjunction (EXISTS quantification)."""
        return self.is_disjunction

    def ground(self, variable_dict: Dict[str, str], state: Optional[SStateCompatible] = None):
        raise NotImplementedError()

    def iter_precondition_predicates(self) -> Iterable[SPredicateName]:
        return self.expr.iter_precondition_predicates()

    def __str__(self) -> str:
        return '({} ({}) {})'.format('exists' if self.is_disjunction else 'forall', str(self.variable), str(self.expr))


class SSASExpression(SValueOutputExpression):  # For all external functions.
    """The representation of an SAS expression. The return value of the expression is an SAS index, therefore it can be represented as a dictionary,
    mapping from Boolean expressions to SAS indices. The execution procedure is to first evaluate all Boolean expressions, and then set the SAS index.
    Suggested implementation is:

    .. code-block:: python

        for sas_index, expr in self.mappings.items():
            if evaluate(expr, state):
                return sas_index
    """

    def __init__(self, mappings: Dict[int, SBoolOutputExpression]):
        """Initialize an SAS expression.

        Args:
            mappings: the mappings from SAS indices to Boolean expressions.
        """
        self.mappings: Dict[int, SBoolOutputExpression] = mappings

    mappings: Dict[int, SBoolOutputExpression]
    """The mappings from SAS indices to Boolean expressions."""

    def ground(self, variable_dict: Dict[str, str], state: Optional[SStateCompatible] = None):
        raise NotImplementedError()

    def __str__(self) -> str:
        return '(SAS\n{}\n)'.format('\n'.join('  ' + str(i) + ' <- ' + str(self.mappings[i]) for i in self.mappings))


class SAssignExpression(SVariableAssignmentExpression):
    """The representation of an assignment expression."""

    def __init__(self, predicate: Union[SBoolPredicateApplicationExpression, SSASPredicateApplicationExpression], value: Union[SBoolOutputExpression, SSASExpression]):
        """Initialize an assignment expression.

        Args:
            predicate: the predicate in the state representation to be assigned.
            value: the value to be assigned.
        """
        self.predicate = predicate
        self.value = value

    predicate: Union[SBoolPredicateApplicationExpression, SSASPredicateApplicationExpression]
    """The predicate in the state representation to be assigned."""

    value: Union[SBoolOutputExpression, SSASExpression]
    """The value to be assigned."""

    def iter_effect_predicates(self) -> Iterable[SPredicateName]:
        return self.predicate.iter_precondition_predicates()

    def __str__(self) -> str:
        return '({} <- {})'.format(str(self.predicate), str(self.value))


class SConditionalAssignExpression(SVariableAssignmentExpression):
    """The representation of a conditional assignment expression. Note that the inner assignment expression is always a :class:`StripsAssignment`."""

    def __init__(self, assign_op: SAssignExpression, condition: SBoolOutputExpression):
        """Initialize a conditional assignment expression.

        Args:
            assign_op: the assignment expression.
            condition: the condition expression.
        """
        self.assign_op = assign_op
        self.condition = condition

    assign_op: SAssignExpression
    """The assignment expression."""

    condition: SBoolOutputExpression
    """The condition expression."""

    @property
    def predicate(self):
        """The predicate in the state representation to be assigned."""
        return self.assign_op.predicate

    @property
    def value(self):
        """The value to be assigned, if the condition is satisfied."""
        return self.assign_op.value

    def iter_precondition_predicates(self) -> Iterable[SPredicateName]:
        return self.condition.iter_precondition_predicates()

    def iter_effect_predicates(self) -> Iterable[SPredicateName]:
        return self.assign_op.iter_effect_predicates()

    def __str__(self) -> str:
        return '({} if {})'.format(str(self.assign_op), str(self.condition))


class SDeicticAssignExpression(SVariableAssignmentExpression):
    """The representation of a deictic assignment expression."""

    def __init__(self, variable: Variable, expression: SVariableAssignmentExpression):
        """Initialize a deictic assignment expression.

        Args:
            variable: the deictic variable.
            expression: the inner assignment expression.
        """
        self.variable = variable
        self.expression = expression

    variable: Variable
    """The deictic expression."""

    expression: SVariableAssignmentExpression
    """The inner assignment expression."""

    def iter_precondition_predicates(self) -> Iterable[SPredicateName]:
        return self.expression.iter_precondition_predicates()

    def iter_effect_predicates(self) -> Iterable[SPredicateName]:
        return self.expression.iter_effect_predicates()

    def __str__(self) -> str:
        return '(foreach ({}) {})'.format(self.variable, str(self.expression))


concepts/pdsketch/strips/strips_expression.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : strips_grounded_expression.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 04/27/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""This file contains the implementation of grounded STRIPS expressions.
Here, the word "grounded" means that the expressions have been instantiated on a specific state. Therefore, the expressions
will not contain any variables (e.g., quantification or deictic expressions). This will make the representations significantly simpler.
For example, the :class:`~concepts.pdsketch.strips.strips_expression.SQuantificationExpression` will be translated into a simple
:class:`GSSimpleBoolExpression` (explicit conjunction or disjunction).

To maximize the execution efficiency of grounded expressions (because they will typically be used in the inner loop of the planner for heuristic evaluation),
we use a compilation strategy to convert each grounded expression into a function that can be called directly. Therefore, each grounded expression class
will have a ``compile`` method that returns a function that takes a state as input and returns a Boolean value.

Note that since we want to support delete-relaxation heuristics, the compiled function will support a feature called "forward difference." This means,
when we compute a Boolean expression on a state, we will also keep track of the branches that makes the expression true. Specifically, we will keep track
of all state variable-level propositions that have contributed to the "true" value of the expression. During the backtracking stage of the heuristic
computation, we will use this information to mark corresponding state variables. See the documentation of :meth:`GSBoolExpression.enable_forward_diff_ctx` for more details.
"""

import warnings
import contextlib

from abc import ABC, abstractmethod
from typing import Any, Optional, Union, Callable, Iterable, Sequence, Tuple, List, Set, FrozenSet, Dict
from dataclasses import dataclass

import jacinle
from concepts.pdsketch.strips.strips_expression import SPredicateName, SProposition, SState, SStateCompatible

__all__ = [
    'GSOptimisticStaticObjectType', 'GS_OPTIMISTIC_STATIC_OBJECT', 'GSBoolForwardDiffReturn',
    'GSExpression', 'GSBoolOutputExpression', 'GSBoolConstantExpression', 'GSSimpleBoolExpression', 'GSComplexBoolExpression',
    'GSVariableAssignmentExpression', 'GSSimpleBoolAssignExpression', 'GSConditionalAssignExpression', 'GSSASAssignExpression', 'GStripsDerivedPredicate',
    'gs_compose_bool_expressions',
    'gs_is_constant_true', 'gs_is_constant_false', 'gs_is_empty_bool_expression', 'gs_is_simple_conjunctive_classifier',
]


class GSOptimisticStaticObjectType(object):
    """The underlying type for :class:`GS_OPTIMISTIC_STATIC_OBJECT`."""
    pass


GS_OPTIMISTIC_STATIC_OBJECT = GSOptimisticStaticObjectType()
"""OptimisticObject only occurs when the arguments to an operator is a complex-typed (Tensor or PyObject) value."""


@dataclass
class GSBoolForwardDiffReturn(object):
    """The return type of the forward difference function of a Boolean expression."""

    rv: bool
    """The return value of the expression."""

    propositions: Union[FrozenSet[SProposition], Set[SProposition]]
    """The set of propositions that have contributed to the return value."""


class GSExpression(ABC):
    """The base class for all grounded STRIPS expressions."""

    @abstractmethod
    def iter_propositions(self) -> Iterable[SProposition]:
        """Iterate over all propositions that are used in this expression."""
        raise NotImplementedError()

    @abstractmethod
    def filter_propositions(self, propositions: SStateCompatible, state: Optional[SState] = None) -> 'GSExpression':
        """Filter the propositions in this expression. Only the propositions in the given set will be kept. Note that this function also takes a state as input,
        essentially, the state is the initial state of the environment, and the `propositions` contains all propositions that will be possibly changed
        by actions. Therefore, for propositions outside the `propositions` set, their value will stay as the same value as in the initial state. See, for example,
        the implementation for :meth:`GSSimpleBoolExpression.filter_propositions` for more details.

        Args:
            propositions: the propositions that should be kept.
            state: the initial state, default to None.

        Returns:
            the filtered expression.
        """
        raise NotImplementedError()

    @abstractmethod
    def compile(self) -> Callable[[SState], Any]:
        """Compile the expression into a function that takes a state as input and returns a boolean value."""
        raise NotImplementedError()

    def __str__(self) -> str:
        raise NotImplementedError()

    __repr__ = jacinle.repr_from_str


class GSBoolOutputExpression(GSExpression, ABC):
    """The base class for all Boolean expressions."""

    FORWARD_DIFF = False
    """A static variable that indicates whether we want to enable forward difference for expressions."""

    @staticmethod
    def set_forward_diff(value: bool = True):
        """Set the forward difference flag.
        If the forward difference flag is set to True, the compiled function will return a tuple of (bool, Set[SProposition]) instead of a bool.

        Args:
            value: the value to set, default to True.
        """
        GSBoolOutputExpression.FORWARD_DIFF = value

    @staticmethod
    @contextlib.contextmanager
    def enable_forward_diff_ctx():
        """A context manager that enables the forward difference flag. Example:

        .. code-block:: python

            state = SState({'a', 'b'})
            expr = GSSimpleBoolExpression({'a', 'b'})
            compiled_function = expr.compile()
            with GSBoolExpression.enable_forward_diff_ctx():
                rv = compiled_function(state)
                assert isinstance(rv, GSBoolForwardDiffReturn)
                assert rv.rv is True
                assert rv.propositions == {'a', 'b'}
        """
        old_value = GSBoolOutputExpression.FORWARD_DIFF
        GSBoolOutputExpression.FORWARD_DIFF = True
        yield
        GSBoolOutputExpression.FORWARD_DIFF = old_value

    def compile(self) -> Callable[[SState], Union[bool, GSBoolForwardDiffReturn]]:
        """Compile the expression into a function that takes a state as input."""
        raise NotImplementedError()

    def forward(self, state: SState) -> Union[bool, GSBoolForwardDiffReturn]:
        """Compute the expression on the given state.

        Args:
            state: the state to compute the expression on.

        Returns:
            the return value of the expression.
        """

        if hasattr(self, '_compiled_function'):
            compiled_function = self._compiled_function
        else:
            compiled_function = self._compiled_function = self.compile()

        return compiled_function(state)

    def filter_propositions(self, propositions: SStateCompatible, state: Optional[SState] = None) -> 'GSBoolOutputExpression':
        raise NotImplementedError()


class GSVariableAssignmentExpression(GSExpression, ABC):
    """The base class for assignment expressions."""

    @abstractmethod
    def compile(self) -> Callable[[SState], SState]:
        """Compile the assignment expression into a function. By default, the function will return the state with the add and delete effects applied."""
        raise NotImplementedError()

    @abstractmethod
    def relax(self) -> Union['GSVariableAssignmentExpression', List['GSVariableAssignmentExpression']]:
        """Relax the expression (a.k.a. delete relaxation)."""
        raise NotImplementedError()

    def filter_propositions(self, propositions: SStateCompatible, state: Optional[SState] = None) -> 'GSVariableAssignmentExpression':
        raise NotImplementedError()


class GSBoolConstantExpression(GSBoolOutputExpression):
    """A constant Boolean expression."""

    def __init__(self, constant: bool):
        """Initialize the constant expression.

        Args:
            constant: the constant value.
        """
        super().__init__()
        self.constant = constant

    constant: bool
    """The constant value."""

    def compile(self) -> Callable[[SState], Union[bool, GSBoolForwardDiffReturn]]:
        def function(state: SState, constant=self.constant):
            if GSBoolOutputExpression.FORWARD_DIFF:
                return GSBoolForwardDiffReturn(constant, set())
            return constant
        return function

    def iter_propositions(self) -> Iterable[SProposition]:
        return tuple()

    def filter_propositions(self, propositions: Set[SProposition], state: Optional[SState] = None) -> 'GSBoolConstantExpression':
        """Filter the propositions in this expression. Since this is a constant expression, the result will always be the same expression.

        Args:
            propositions: the propositions that should be kept.
            state: the initial state, default to None.

        Returns:
            the Boolean expression after filtering.
        """
        return self

    def __str__(self):
        return '{}'.format(self.constant).upper()


class GSSimpleBoolExpression(GSBoolOutputExpression):
    """A simple Boolean expression. Here, a simple Boolean expression is a conjunction or disjunction of propositions. Therefore, internally,
    we use a frozen set to store the propositions. This also accelerates the testing process."""

    def __init__(self, propositions: Union[Sequence[SProposition], FrozenSet[str]], is_disjunction: bool = False):
        """Initialize the simple Boolean expression.

        Args:
            propositions: the propositions in the expression.
            is_disjunction: whether the expression is a disjunction, default to False (a.k.a. conjunction).
        """
        super().__init__()

        self.propositions = frozenset(propositions)
        if len(self.propositions) > 1:
            self.is_disjunction = is_disjunction
        else:
            self.is_disjunction = False  # prefer to represent a single prop. as a conjunction.

    propositions: FrozenSet[SProposition]
    """The propositions in the expression."""

    is_disjunction: bool
    """Whether the expression is a disjunction."""

    @property
    def is_conjunction(self) -> bool:
        """Whether the expression is a conjunction."""
        return not self.is_disjunction

    def compile(self) -> Callable[[SState], Union[bool, GSBoolForwardDiffReturn]]:
        """Compile the expression into a function. When forward diff is enabled, the function will use the following strategy:

        1. If the expression is a conjunction, the function will return True if all propositions are true, and return all propositions as the forward diff.
        2. If the expression is a disjunction, the function will return True if any proposition is true, and return the first true proposition as the forward diff.
        """
        if self.is_disjunction:
            def function(state: SState, classifier=self.propositions) -> Optional[Union[bool, GSBoolForwardDiffReturn]]:
                if GSBoolOutputExpression.FORWARD_DIFF:
                    intersection = classifier.intersection(state)
                    if intersection:
                        return GSBoolForwardDiffReturn(True, {next(iter(intersection))})
                    else:
                        return GSBoolForwardDiffReturn(False, None)
                return len(classifier.intersection(state)) > 0
            return function
        else:
            def function(state: SState, classifier=self.propositions) -> Optional[Union[bool, GSBoolForwardDiffReturn]]:
                if GSBoolOutputExpression.FORWARD_DIFF:
                    return GSBoolForwardDiffReturn(classifier <= state, classifier)
                return classifier <= state
            return function

    def iter_propositions(self) -> Iterable[SProposition]:
        yield from iter(self.propositions)

    def filter_propositions(self, propositions: Set[SProposition], state: Optional[SState] = None) -> 'GSBoolOutputExpression':
        """Filter the propositions in this expression. Only the propositions in the given set will be kept. Note that this function will automatically
        handle cases where the filtered expression is a constant expression. The rules are:

        1. If the expression is a conjunction, denote the set of propositions that will be filtered out as `A`, and the set of propositions that
        are True in the initial state as `B`:

            1. If `A` is a not subset of `B`, then the filtered expression is a constant False expression (i.e., there are propositions that are
            not True in the initial state but will never be changed by actions).

            2. Otherwise, if the remaining set of propositions is empty, then the filtered expression is a constant True expression. Otherwise,
            the filtered expression is a conjunction of the remaining propositions.

        2. If the expression is a disjunction, denote the set of propositions that will be filtered out as `A`, and the set of propositions that
        are True in the initial state as `B`:

            1. If `A` and `B` has non-empty intersection, then the filtered expression is a constant True expression (i.e., there are propositions that are
            True in the initial state and will never be changed by actions).

            2. Otherwise, if the remaining set of propositions is empty, then the filtered expression is a constant False expression (i.e., all propositions
            are False in the initial state and will never be changed by actions). Otherwise, the filtered expression is a disjunction of the remaining propositions.

        Args:
            propositions: the propositions that should be kept.
            state: the initial state, default to None.

        Returns:
            the Boolean expression after filtering.
        """
        if state is None:
            state = set()

        diff = self.propositions - propositions
        new_classifiers = frozenset(self.propositions & propositions)

        if len(diff) == 0:
            return self
        else:
            if self.is_disjunction:
                if len(diff & state) > 0:
                    return GSBoolConstantExpression(True)
                else:
                    if len(new_classifiers) == 0:
                        return GSBoolConstantExpression(False)
                    else:
                        return GSSimpleBoolExpression(new_classifiers, self.is_disjunction)
            else:
                if not diff <= state:
                    return GSBoolConstantExpression(False)
                else:
                    if len(new_classifiers) == 0:
                        return GSBoolConstantExpression(True)
                    else:
                        return GSSimpleBoolExpression(new_classifiers, self.is_disjunction)

    def __str__(self) -> str:
        name = 'CONJ' if not self.is_disjunction else 'DISJ'
        classifier_str = ', '.join([str(x) for x in self.propositions])
        return f'{name}({classifier_str})'


class GSComplexBoolExpression(GSBoolOutputExpression):
    """A complex Boolean expression. Here, a complex Boolean expression is a conjunction or disjunction of Boolean sub-expressions.
    In most of the scenarios, you should directly call the constructor for this class. Instead, if you want to compose multiple Boolean
    expressions, you should use the function :func:`gs_compose_bool_expressions` instead."""

    def __init__(self, expressions: Sequence[GSBoolOutputExpression], is_disjunction: Optional[bool] = False):
        """Initialize the complex Boolean expression.

        Args:
            expressions: the sub-expressions.
            is_disjunction: whether the expression is a disjunction, default to False (a.k.a. conjunction).
        """
        super().__init__()
        self.expressions = tuple(expressions)
        self.is_disjunction = is_disjunction

    expressions: Tuple[GSBoolOutputExpression]
    """The sub-expressions."""

    is_disjunction: bool
    """Whether the expression is a disjunction."""

    @property
    def is_conjunction(self) -> bool:
        """Whether the expression is a conjunction."""
        return not self.is_disjunction

    def compile(self) -> Callable[[SState], Union[bool, GSBoolForwardDiffReturn]]:
        """Compile the expression into a function. When forward diff is enabled, the function will use the following strategy:

        1. If the expression is a conjunction, the function will return True if all sub-expressions are true, and return the union of propositions as the forward diff.
        2. If the expression is a disjunction, the function will return True if any sub-expression is true, and return the first true sub-expression as the forward diff.
        """
        compiled_functions = tuple(e.compile() for e in self.expressions)
        if self.is_disjunction:
            def function(state: SState, functions=compiled_functions) -> Union[bool, GSBoolForwardDiffReturn]:
                if GSBoolOutputExpression.FORWARD_DIFF:
                    for f in functions:
                        rv = f(state)
                        result, propositions = rv.rv, rv.propositions
                        if result:
                            return GSBoolForwardDiffReturn(True, propositions)
                    return GSBoolForwardDiffReturn(False, set())

                return any(f(state) for f in functions)
        else:
            def function(state: SState, functions=compiled_functions) -> Union[bool, GSBoolForwardDiffReturn]:
                if GSBoolOutputExpression.FORWARD_DIFF:
                    all_propositions = list()
                    for f in functions:
                        rv = f(state)
                        result, propositions = rv.rv, rv.propositions
                        if not result:
                            return GSBoolForwardDiffReturn(False, set())
                        all_propositions.append(propositions)
                    propositions = frozenset.union(*all_propositions)
                    return GSBoolForwardDiffReturn(True, propositions)

                return all(f(state) for f in functions)
        return function

    def iter_propositions(self) -> Iterable[SProposition]:
        for e in self.expressions:
            yield from e.iter_propositions()

    def filter_propositions(self, propositions: Set[SProposition], state: Optional[SState] = None) -> 'GSBoolOutputExpression':
        """Filter the given propositions from the expression. See the documentation of :func:`gstrips_compose_classifiers` for more details.

        Args:
            propositions: the propositions that should be kept.
            state: the initial state, default to None.

        Returns:
            the Boolean expression after filtering.
        """
        expressions = [e.filter_propositions(propositions, state=state) for e in self.expressions]
        return gs_compose_bool_expressions(expressions, self.is_disjunction, propagate_implicit_propositions=False)

    def __str__(self):
        name = 'ComplexAND' if not self.is_disjunction else 'ComplexOR'
        expressions = [str(x) for x in self.expressions]
        if sum([len(x) for x in expressions]) > 120:
            return f'{name}(\n' + ',\n'.join([jacinle.indent_text(jacinle.stformat(x)).rstrip() for x in expressions]) + '\n)'
        return f'{name}(' + ', '.join([str(x) for x in self.expressions]) + ')'


class GSSimpleBoolAssignExpression(GSVariableAssignmentExpression):
    """A simple assignment expression. Here, a simple assignment expression is represented by the set of add effects and the set of delete effects, both
    of which are a set of propositions."""

    def __init__(self, add_effects: Iterable[SProposition], del_effects: Iterable[SProposition]):
        """Initialize the simple assignment expression.

        Args:
            add_effects: the set of add effects.
            del_effects: the set of delete effects.
        """
        super().__init__()
        self.add_effects = frozenset(add_effects)
        self.del_effects = frozenset(del_effects)

    add_effects: FrozenSet[SProposition]
    """The set of add effects."""

    del_effects: FrozenSet[SProposition]
    """The set of delete effects."""

    def compile(self) -> Callable[[SState], SState]:
        def function(state: SState, del_effects=self.del_effects, add_effects=self.add_effects) -> SState:
            new_state = (state - del_effects) | add_effects
            return SState(new_state)
        return function

    def iter_propositions(self) -> Iterable[SProposition]:
        yield from iter(self.add_effects)
        yield from iter(self.del_effects)

    def filter_propositions(self, propositions: Set[SProposition], state: Optional[SState] = None) -> 'GSSimpleBoolAssignExpression':
        """Filter the given propositions from the expression. For simple assignment expressions, the filtering is done simply by removing irrelevant propositions
        from the add and delete effects."""
        return GSSimpleBoolAssignExpression(propositions & self.add_effects, propositions & self.del_effects)

    def relax(self) -> 'GSSimpleBoolAssignExpression':
        """Delete relaxation of a simple assignment. Essentially, it removes all delete effects."""
        return GSSimpleBoolAssignExpression(self.add_effects, set())

    def __str__(self) -> str:
        return f'EFF[add={self.add_effects}, del={self.del_effects}]'


class GSConditionalAssignExpression(GSVariableAssignmentExpression):
    """A conditional assignment expression. Note that the inner expression must be a simple assignment expression.
    Therefore , a conditional assignment expression is represented by a condition, a set of add effects, and a set of delete effects."""

    def __init__(self, condition: GSBoolOutputExpression, assignment: GSSimpleBoolAssignExpression):
        """Initialize the conditional assignment expression.

        Args:
            condition: the condition.
            assignment: the simple assignment expression.
        """

        super().__init__()
        self.condition = condition
        self.assignment = assignment
        self.condition_func = None
        self.assignment_func = None

    condition: GSBoolOutputExpression
    """The condition of the conditional assignment expression."""

    assignment: GSSimpleBoolAssignExpression
    """The inner (simple) assignment expression of the conditional assignment expression."""

    condition_func: Optional[Callable[[SState], Union[bool, GSBoolForwardDiffReturn]]]
    """The compiled condition function."""

    assignment_func: Optional[Callable[[SState], SState]]
    """The compiled assignment function."""

    @property
    def add_effects(self) -> FrozenSet[SProposition]:
        """The add effects of the conditional assignment expression."""
        assert isinstance(self.assignment, GSSimpleBoolAssignExpression)
        return self.assignment.add_effects

    @property
    def del_effects(self) -> FrozenSet[SProposition]:
        """The delete effects of the conditional assignment expression."""
        assert isinstance(self.assignment, GSSimpleBoolAssignExpression)
        return self.assignment.del_effects

    def compile(self) -> Callable[[SState], SState]:
        condition_func = self.condition.compile()
        assignment_func = self.assignment.compile()

        def function(state: SState, condition_func=condition_func, assignment=assignment_func) -> SState:
            if condition_func(state):
                return assignment(state)
            return state
        self.condition_func = condition_func
        self.assignment_func = assignment_func
        return function

    def applicable(self, state: SState) -> Union[bool, GSBoolForwardDiffReturn]:
        """Check if the conditional assignment expression is applicable in the given state."""
        return self.condition_func(state)

    def apply(self, state: SState) -> SState:
        """Apply the conditional assignment expression in the given state."""
        return self.assignment_func(state)

    def iter_propositions(self) -> Iterable[SProposition]:
        yield from self.assignment.iter_propositions()

    def filter_propositions(self, propositions: Set[SProposition], state: Optional[SState] = None) -> 'GSConditionalAssignExpression':
        """Filter the given propositions from the expression. For conditional assignment expressions, the filtering is done by filtering the inner simple
        assignment expression and the condition.

        Args:
            propositions: the propositions that should be kept.
            state: the initial state, default to None.

        Returns:
            the conditional assignment expression after filtering.
        """
        return GSConditionalAssignExpression(
            self.condition.filter_propositions(propositions, state=state),
            self.assignment.filter_propositions(propositions, state=state)
        )

    def relax(self) -> 'GSConditionalAssignExpression':
        """Delete relaxation of a conditional assignment. Essentially, it removes all delete effects for the inner simple assignment expression.

        Returns:
            the delete relaxed conditional assignment expression.
        """
        return GSConditionalAssignExpression(self.condition, self.assignment.relax())

    def __str__(self) -> str:
        if isinstance(self.assignment, GSSimpleBoolAssignExpression):
            return f'CONDEFF[{self.condition} => add={self.assignment.add_effects}, del={self.assignment.del_effects}]'
        else:
            return f'CONDEFF[{self.condition} => {self.assignment}]'


class GSSASAssignExpression(GSVariableAssignmentExpression):
    """A SAS assignment expression. It is represented as the name for the SAS predicate, the size of the SAS predicate, a list of assignments,
    and a mapping from Boolean expressions to integers (the value of the SAS assignment)."""

    def __init__(
        self,
        sas_name: SPredicateName, sas_size: int, arguments: Sequence[str],
        expression: Dict[int, GSBoolOutputExpression]
    ):
        """Initialize the SAS assignment expression.

        Args:
            sas_name: the name of the SAS predicate.
            sas_size: the size of the SAS predicate.
            arguments: the arguments of the SAS predicate.
            expression: the expressions for SAS assignments.
        """

        super().__init__()
        self.sas_name = sas_name
        self.sas_size = sas_size
        self.arguments = arguments
        self.arguments_str = ' '.join(str(arg) for arg in arguments)
        self.all_bool_predicates = frozenset({f'{sas_name}@{i} {self.arguments_str}' for i in range(sas_size)})
        self.all_neg_bool_predicates = frozenset({f'{sas_name}@{i}_not {self.arguments_str}' for i in range(sas_size)})
        self.expression = expression

    sas_name: SPredicateName
    """The name of the SAS predicate."""

    sas_size: int
    """The size of the SAS predicate."""

    arguments: Sequence[str]
    """The arguments of the SAS predicate."""

    arguments_str: str
    """A string format for the arguments of the SAS predicate."""

    all_bool_predicates: FrozenSet[SProposition]
    """All the Boolean propositions of the SAS predicate."""

    all_neg_bool_predicates: FrozenSet[SProposition]
    """All the negated Boolean propositions of the SAS predicate."""

    expression: Dict[int, GSBoolOutputExpression]
    """The expressions for SAS assignments, which is a mapping from integers to Boolean expressions."""

    def compile(self) -> Callable[[SState], SState]:
        """Compile the SAS assignment expression to a function.

        .. warning::

            This method is deprecated and will be removed in the future. Use the :meth:`to_conditional_assignments` or :meth:`relax` methods instead.

        Returns:
            the compiled function.
        """
        warnings.warn('SAS assignments are not supported. Run to_conditional_assignments() before compilation.', RuntimeWarning)
        compiled_expressions = tuple((i, e.compile()) for i, e in self.expression.items())

        def function(
            state: SState, *,
            compiled_expressions=compiled_expressions,
            sas_name=self.sas_name, arguments_str=self.arguments_str, all_bool_predicates=self.all_bool_predicates,
        ) -> SState:
            new_value = None
            for k, v in compiled_expressions:
                if v(state):
                    new_value = k
                    break
            if new_value is not None:
                current_value = state.intersection(all_bool_predicates)
                diff: FrozenSet[SProposition] = state - all_bool_predicates
                current_value_not: Set[SProposition] = set()
                for v in current_value:
                    a = v.split()
                    a[0] = f'{a[0]}_not'
                    current_value_not.add(' '.join(a))
                new_value_set: Set[SProposition] = {f'{sas_name}@{new_value} {arguments_str}'}

                state: SState = SState(diff | current_value_not | new_value_set)
                return state
            return state
        return function

    def iter_propositions(self) -> Iterable[SProposition]:
        yield from self.all_bool_predicates
        for v in self.expression.values():
            yield from v.iter_propositions()

    def filter_propositions(self, propositions: Set[SProposition], state: Optional[SState] = None) -> 'GSSASAssignExpression':
        return GSSASAssignExpression(
            self.sas_name, self.sas_size, self.arguments,
            expression={k: v.filter_propositions(propositions, state=state) for k, v in self.expression.items()}
        )

    def to_conditional_assignments(self) -> List[GSVariableAssignmentExpression]:
        """Convert the SAS assignment to a list of conditional assignments.

        Returns:
            a list of conditional assignments.
        """
        rvs = list()
        for k, v in self.expression.items():
            this_add = {f'{self.sas_name}@{k} {self.arguments_str}'} | self.all_neg_bool_predicates - {f'{self.sas_name}@{k}_not {self.arguments_str}'}
            this_del = self.all_bool_predicates | {f'{self.sas_name}@{k}_not {self.arguments_str}'} - {f'{self.sas_name}@{k} {self.arguments_str}'}
            rvs.append(GSConditionalAssignExpression(v, GSSimpleBoolAssignExpression(this_add, this_del)))
        return rvs

    def relax(self) -> List['GSConditionalAssignExpression']:
        """Delete relaxation of a SAS assignment. Essentially, it removes all delete effects for the inner simple assignment expression.
        Note that this method returns a list of conditional assignments, instead of a single SASAssignment.

        Returns:
            a list of delete-relaxed SAS assignment expression.
        """
        rvs = list()
        for k, v in self.expression.items():
            add_effects = {f'{self.sas_name}@{k} {self.arguments_str}'} | {f'{self.sas_name}@{i}_not {self.arguments_str}' for i in range(self.sas_size) if i != k}
            rvs.append(GSConditionalAssignExpression(v, GSSimpleBoolAssignExpression(add_effects, set())))
        return rvs

    def __str__(self) -> str:
        expression_str = jacinle.stformat(self.expression).rstrip()
        return f'SAS[target={self.sas_name} {self.arguments_str}, value={expression_str}]'


class GStripsDerivedPredicate(GSExpression):
    """Grounded STRIPS version of derived predicates."""

    def __init__(
        self, name: str, arguments: Sequence[str],
        expression_true: Optional[GSBoolOutputExpression] = None, expression_false: Optional[GSBoolOutputExpression] = None, is_relaxed: bool = False,
        effects: Optional[Iterable[GSConditionalAssignExpression]] = None
    ):
        """Initialize the grounded derived predicate.

        Args:
            name: the name of the derived predicate.
            arguments: the arguments of the derived predicate.
            expression_true: the expression for the true case. Optional if ``effects`` is provided.
            expression_false: the expression for the false case. Optional if ``effects`` is provided.
            is_relaxed: whether the derived predicate has been delete-relaxed.
            effects: the effects of the derived predicate. Optional if ``expression_true`` and ``expression_false`` are provided.
        """

        super().__init__()
        self.name = name
        self.arguments = tuple(arguments)

        self.pos_name = self.name + ' ' + ' '.join(str(x) for x in self.arguments)
        self.neg_name = self.name + '_not ' + ' '.join(str(x) for x in self.arguments)

        if effects is None:
            assert expression_true is not None and expression_false is not None
            self.pos_effect = GSConditionalAssignExpression(expression_true, GSSimpleBoolAssignExpression({self.pos_name}, {self.neg_name} if not is_relaxed else set()))
            self.neg_effect = GSConditionalAssignExpression(expression_false, GSSimpleBoolAssignExpression({self.neg_name}, {self.pos_name} if not is_relaxed else set()))
            self.effects = (self.pos_effect, self.neg_effect)
        else:
            self.effects = tuple(effects)
            assert len(self.effects) == 2

    name: str
    """The name of the derived predicate."""

    arguments: Tuple[str, ...]
    """The arguments of the derived predicate."""

    pos_name: str
    """The name of the positive proposition corresponding to this derived predicate."""

    neg_name: str
    """The name of the negative proposition corresponding to this derived predicate."""

    pos_effect: GSConditionalAssignExpression
    """The positive effect of the derived predicate."""

    neg_effect: GSConditionalAssignExpression
    """The negative effect of the derived predicate."""

    effects: Tuple[GSConditionalAssignExpression, ...]
    """A tuple of the positive and negative effects of the derived predicate."""

    def compile(self) -> Callable[[SState], SState]:
        effects_func = tuple(effect.compile() for effect in self.effects)

        def function(state: SState, effects_func=effects_func) -> SState:
            for func in effects_func:
                state = func(state)
            return state
        return function

    def iter_propositions(self) -> Iterable[SProposition]:
        for eff in self.effects:
            yield from eff.iter_propositions()

    def filter_propositions(self, propositions: Set[SProposition], state: Optional[SState] = None) -> 'GStripsDerivedPredicate':
        return GStripsDerivedPredicate(
            self.name,
            self.arguments,
            effects=[eff.filter_propositions(propositions, state=state) for eff in self.effects]
        )

    def relax(self) -> 'GStripsDerivedPredicate':
        """Delete relaxation of a derived predicate. Essentially, it removes all delete effects for the inner simple assignment expression."""
        return GStripsDerivedPredicate(self.name, self.arguments, effects=[eff.relax() for eff in self.effects])

    def __str__(self) -> str:
        effects_str = '\n'.join(jacinle.indent_text(str(eff)) for eff in self.effects)
        return f'DERIVED[\n{effects_str}\n]'


# @jacinle.log_function
def _compose_strips_classifiers_inner(classifiers: Sequence[GSBoolOutputExpression], is_disjunction: Optional[bool] = False) -> GSBoolOutputExpression:
    new_classifiers = list()
    visited = [False for _ in classifiers]

    for i in range(len(classifiers)):
        if not visited[i]:
            c = classifiers[i]
            if gs_is_constant_true(c):
                visited[i] = True
                if is_disjunction:
                    return GSBoolConstantExpression(True)
            elif gs_is_constant_false(c):
                visited[i] = True
                if not is_disjunction:
                    return GSBoolConstantExpression(False)
            elif c == GS_OPTIMISTIC_STATIC_OBJECT:
                visited[i] = True

    new_set = set()

    def add_simple_classifier(c: GSSimpleBoolExpression):
        if c.is_disjunction == is_disjunction or len(c.propositions) == 1:
            new_set.update(c.propositions)
            return True
        return False

    for i in range(len(classifiers)):
        if not visited[i]:
            c = classifiers[i]
            if isinstance(c, GSSimpleBoolExpression):
                visited[i] = add_simple_classifier(c)

    complex = list()
    for i in range(len(classifiers)):
        if not visited[i]:
            c = classifiers[i]
            assert isinstance(c, (GSComplexBoolExpression, GSSimpleBoolExpression))
            if c.is_disjunction == is_disjunction:
                assert isinstance(c, GSComplexBoolExpression)
                for e in c.expressions:
                    if isinstance(e, GSSimpleBoolExpression):
                        if add_simple_classifier(e):
                            continue
                        else:
                            complex.append(e)
                    else:
                        complex.append(e)
            else:
                complex.append(c)

    if len(new_set) > 0:
        new_classifiers.append(GSSimpleBoolExpression(new_set, is_disjunction))
    new_classifiers.extend(complex)

    if len(new_classifiers) == 0:
        return GSBoolConstantExpression(True if not is_disjunction else False)
    elif len(new_classifiers) == 1:
        return new_classifiers[0]
    else:
        return GSComplexBoolExpression(new_classifiers, is_disjunction)


def gs_compose_bool_expressions(
    expressions: Union[Sequence[Tuple[GSBoolOutputExpression, Set[SProposition]]], Sequence[GSBoolOutputExpression]],
    is_disjunction: Optional[bool] = False,
    propagate_implicit_propositions: bool = True
) -> Union[Tuple[GSBoolOutputExpression, Set[SProposition]], GSBoolOutputExpression]:
    """Compose a list of Boolean expressions into a single expression, by taking the conjunction or disjunction of them.
    This function will automatically handles possible merging of simple expressions to reduce the complexity of the final expression.
    For example, if we are taking the conjunction of two simple conjunctive expressions, then the two expressions will be merged into one single conjunction.

    This function takes an additional argument `propagate_implicit_propositions`, which controls whether the implicit propositions should be propagated to the final expression.
    Specifically, when `propagate_implicit_propositions` is set to `True`, the argument `expression` should be a list of tuples,
    where the first element is the Boolean expression and the second element is the set of implicit propositions.
    In this case, the return will also be a tuple, where the first element is the final expression and the second element is the set of implicit propositions.

    When `propagate_implicit_propositions` is set to `False`, the argument `expression` should be a list of Boolean expressions,
    and the return is also a single Boolean expression.

    Args:
        expressions: a list of Boolean expressions.
        is_disjunction: whether to take disjunction of the expressions, default to False (conjunction).
        propagate_implicit_propositions: whether to keep the implicit propositions in the final expression, default to True.
    """
    if propagate_implicit_propositions:
        rv = _compose_strips_classifiers_inner([c[0] for c in expressions], is_disjunction)
        return rv, set.union(*[c[1] for c in expressions]) if len(expressions) > 0 else set()
    else:
        rv = _compose_strips_classifiers_inner(expressions, is_disjunction)
        return rv


def gs_is_constant_true(classifier: GSBoolOutputExpression) -> bool:
    """Check if the given classifier is a constant true expression."""
    return isinstance(classifier, GSBoolConstantExpression) and classifier.constant


def gs_is_constant_false(classifier: GSBoolOutputExpression) -> bool:
    """Check if the given classifier is a constant false expression."""
    return isinstance(classifier, GSBoolConstantExpression) and not classifier.constant


def gs_is_empty_bool_expression(expression: GSBoolOutputExpression) -> bool:
    """Check if the given classifier is an empty Boolean expression."""
    return isinstance(expression, GSSimpleBoolExpression) and len(expression.propositions) == 0


def gs_is_simple_conjunctive_classifier(classifier: GSBoolOutputExpression) -> bool:
    """Check if the given classifier is a simple conjunctive classifier."""
    return isinstance(classifier, GSSimpleBoolExpression) and not classifier.is_disjunction


concepts/pdsketch/strips/strips_grounded_expression.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : strips_grounding.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 03/19/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import itertools
from typing import Optional, Union, Callable, Iterable, Sequence, Iterator, Tuple, List, Set, Dict

import jacinle
from jacinle.utils.printing import indent_text

import concepts.dsl.expression as E
from concepts.dsl.dsl_types import BOOL
from concepts.dsl.tensor_value import TensorValue
from concepts.dsl.constraint import OptimisticValue
from concepts.dsl.expression import Expression, ExpressionDefinitionContext
from concepts.dsl.executors.tensor_value_executor import compose_bvdict_args

from concepts.pdsketch.predicate import Predicate, flatten_expression, split_simple_bool
from concepts.pdsketch.domain import Domain, State
from concepts.pdsketch.operator import OperatorApplier, gen_all_partially_grounded_actions, RegressionRuleApplier, gen_all_grounded_regression_rules
from concepts.pdsketch.executor import PDSketchExecutor

from concepts.pdsketch.strips.strips_expression import SPredicateName, SProposition, SState, SStateCompatible, make_sproposition

from concepts.pdsketch.strips.strips_grounded_expression import GSBoolOutputExpression, GSVariableAssignmentExpression, GSBoolForwardDiffReturn
from concepts.pdsketch.strips.strips_grounded_expression import GS_OPTIMISTIC_STATIC_OBJECT, GSOptimisticStaticObjectType
from concepts.pdsketch.strips.strips_grounded_expression import GSBoolConstantExpression, GSSimpleBoolExpression, gs_compose_bool_expressions, gs_is_empty_bool_expression
from concepts.pdsketch.strips.strips_grounded_expression import GSSimpleBoolAssignExpression, GSConditionalAssignExpression, GStripsDerivedPredicate


__all__ = [
    'GStripsOperator', 'GStripsProblem',
    'GStripsTranslatorBase', 'GStripsTranslatorOptimistic',
    'relevance_analysis_v1', 'relevance_analysis_v2'
]


class GStripsOperator(object):
    """A grounded STRIPS operator, the STRIPS version of the OperatorApplier.
    By grounded we mean that all the preconditions / effects are propositions."""

    def __init__(
        self,
        precondition: GSBoolOutputExpression,
        effects: Sequence[GSVariableAssignmentExpression],
        raw_operator: OperatorApplier,
        implicit_propositions: Optional[Set[SProposition]] = None
    ):
        self.precondition = precondition
        self.implicit_propositions = implicit_propositions if implicit_propositions is not None else set()
        self.effects = tuple(effects)
        self.raw_operator = raw_operator
        self.precondition_func: Optional[Callable[[SState], bool]] = None
        self.effects_func: Optional[Tuple[Callable[[SState], SState], ...]] = None

    precondition: GSBoolOutputExpression
    """The precondition for the operator."""

    effects: Tuple[GSVariableAssignmentExpression, ...]
    """The effects for the operator."""

    implicit_propositions: Set[SProposition]
    """Implicit preconditions."""

    raw_operator: OperatorApplier
    """The raw operator.
    Note that since the strips operator is grounded, the raw operator is the :class:`concepts.pdsketch.operator.OperatorApplier` instance."""

    precondition_func: Optional[Callable[[SState], bool]]
    """An compiled version of the precondition tester."""

    effects_func: Optional[Tuple[Callable[[SState], SState], ...]]
    """A list of compiled version of the effect appliers."""

    def compile(self):
        """Generate the compiled version of the operator (preconditions and effects)."""
        self.precondition_func = self.precondition.compile()
        self.effects_func = tuple(e.compile() for e in self.effects)

    def applicable(self, state: Union[SState, Set[SProposition]]) -> Union[bool, GSBoolForwardDiffReturn]:
        """Return whether the operator is applicable to the state.

        Args:
            state: the current state.

        Returns:
            True if the operator is applicable, False otherwise. When the `FORWARD_DIFF` mode is enabled for :class:`~concepts.pdsketch.strips.strips_grounded_expression.GSBoolExpression`,
            this function will return a :class:`~concepts.pdsketch.strips.strips_grounded_expression.GSBoolForwardDiffReturn` object.
        """
        return self.precondition_func(state)

    def apply(self, state: SState) -> SState:
        """Apply the operator to the state. Note that this function does not check the precondition. To do that, use :meth:`applicable` before calling this function.

        Args:
             state: the current state.

        Returns:
             the new state after applying the operator.
        """
        for effect_func in self.effects_func:
            state = effect_func(state)
        return state

    def iter_propositions(self) -> Iterable[SProposition]:
        """Iterate over all the propositions used in the operator, including both preconditions and effects."""
        yield from self.precondition.iter_propositions()
        for e in self.effects:
            yield from e.iter_propositions()

    def filter_propositions(self, propositions: Set[SProposition], state: Optional[SState] = None) -> 'GStripsOperator':
        """Filter the propositions in this expression. Only the propositions in the given set will be kept. Note that this function also takes a state as input,
        essentially, the state is the initial state of the environment, and the `propositions` contains all propositions that will be possibly changed
        by actions. See :meth:`~concepts.pdsketch.strips.strips_grounded_expression.GSExpression.filter_propositions` for more details.

        Args:
            propositions: the propositions to keep.
            state: the initial state for filtering the propositions.

        Returns:
            a new operator with filtered propositions.
        """

        return GStripsOperator(
            self.precondition.filter_propositions(propositions, state=state),
            [e.filter_propositions(propositions, state=state) for e in self.effects],
            self.raw_operator
        )

    def __str__(self) -> str:
        return (
            type(self).__name__
            + '{' + f'{self.raw_operator}' + '}'
            + '{\n' + indent_text(self.precondition)
            + '\n' + '\n'.join(indent_text(x) for x in self.effects) if len(self.effects) > 0 else '<Empty Effect>'
            + '\n}'
        )

    def __repr__(self) -> str:
        return self.__str__()


class GStripsAchieveExpression(object):
    def __init__(self, goal: GSBoolOutputExpression, maintains: GSBoolOutputExpression):
        self.goal = goal
        self.maintains = maintains

    goal: GSBoolOutputExpression
    maintains: GSBoolOutputExpression


class GStripsRegressionRule(object):
    def __init__(self, precondition: GSBoolOutputExpression, preconstraints: GSBoolOutputExpression, goal: GSBoolOutputExpression, body: Sequence[Union[GStripsOperator, GStripsAchieveExpression]]):
        self.precondition = precondition
        self.preconstraints = preconstraints
        self.goal = goal
        self.body = tuple(body)

    precondition: GSBoolOutputExpression
    preconstraints: GSBoolOutputExpression
    goal: GSBoolOutputExpression
    body: Tuple[Union[GStripsOperator, GStripsAchieveExpression], ...]


class GStripsProblem(object):
    """A grounded STRIPS task. It is composed of three parts:

    - The current state, which is a set of propositions.
    - The goal state, which is a grounded STRIPS classifier.
    - The operators, which is a list of applicable grounded STRIPS operators.
    """

    def __init__(
        self,
        state: SStateCompatible,
        goal: GSBoolOutputExpression,
        operators: Sequence[GStripsOperator],
        is_relaxed: bool = False,
        goal_implicit_propositions: Optional[Set[SProposition]] = None,
        regression_rules: Optional[Sequence[GStripsRegressionRule]] = None,
        derived_predicates: Optional[Sequence[GStripsDerivedPredicate]] = None,
        facts: Optional[Set[SProposition]] = None,
    ):
        self.state = SState(state)
        self.goal = goal
        self.operators = tuple(operators)
        self.regression_rules = tuple(regression_rules) if regression_rules is not None else tuple()
        self.derived_predicates = tuple(derived_predicates) if derived_predicates is not None else tuple()
        self.is_relaxed = is_relaxed
        self.goal_implicit_propositions = goal_implicit_propositions if goal_implicit_propositions is not None else set()
        self._facts = facts

    state: SState
    """The initial state."""

    goal: GSBoolOutputExpression
    """The goal expression."""

    operators: Tuple[GStripsOperator]
    """The list of operators."""

    regression_rules: Tuple[GStripsRegressionRule]
    """The list of regression rules."""

    derived_predicates: Tuple[GStripsDerivedPredicate]
    """The list of derived predicates."""

    is_relaxed: bool
    """Whether the task is a delete-relaxed task."""

    goal_implicit_propositions: Set[SProposition]
    """The implicit propositions in the goal."""

    def compile(self) -> 'GStripsProblem':
        """Compile all operators and derived predicates.

        Returns:
            self
        """
        for op in self.operators:
            op.compile()
        for op in self.derived_predicates:
            op.compile()
        return self

    @property
    def facts(self):
        """The relevant facts in the task."""
        return self._facts

    def __str__(self):
        operator_str = '\n'.join(str(op) for op in self.operators)
        derived_predicate_str = '\n'.join(str(dp) for dp in self.derived_predicates)
        return f"""{type(self).__name__}{{
  state: {self.state}
  goal: {self.goal}
  operators:
    {indent_text(operator_str, 2).strip()}
  derived_predicates:
    {indent_text(derived_predicate_str, 2).strip()}
  facts: {self._facts}
}}"""

    def __repr__(self) -> str:
        return self.__str__()


class GStripsTranslatorBase(object):
    def __init__(self, executor: PDSketchExecutor, use_string_name: bool = True, prob_goal_threshold: float = 0.5, use_derived_predicates: bool = False, use_regression_rules: bool = True):
        assert isinstance(executor, PDSketchExecutor)
        self._executor = executor
        self._use_string_name = use_string_name
        self._prob_goal_threshold = prob_goal_threshold
        self._use_derived_predicates = use_derived_predicates
        self._use_regression_rules = use_regression_rules
        self._predicate2index = dict()
        self._init_indices()

    @property
    def executor(self) -> PDSketchExecutor:
        """The executor for the domain."""
        return self._executor

    @property
    def domain(self) -> Domain:
        """The domain."""
        return self._executor.domain

    @property
    def predicate2index(self) -> Dict[Tuple[SPredicateName, str], int]:
        """The mapping from predicate names to index."""
        return self._predicate2index

    def _init_indices(self):
        raise NotImplementedError()

    def compile_expr(self, expr: Union[str, Expression], state: State) -> Tuple[GSBoolOutputExpression, Set[SProposition]]:
        raise NotImplementedError()

    def compile_operator(self, op: OperatorApplier, state: State, is_relaxed=False) -> GStripsOperator:
        raise NotImplementedError()

    def compile_regression_rule(self, op: OperatorApplier, state: State) -> GStripsRegressionRule:
        raise NotImplementedError()

    def compile_derived_predicate(self, dp: Predicate, state: State, is_relaxed=False) -> List[GStripsDerivedPredicate]:
        raise NotImplementedError()

    def compile_state(self, state: State, forward_derived: bool = False) -> SState:
        raise NotImplementedError()

    def relevance_analysis(self, task: GStripsProblem, relaxed_relevance: bool = False, forward: bool = True, backward: bool = True) -> GStripsProblem:
        raise NotImplementedError()

    def define_grounded_predicate(self, name: str, modifier: Optional[str] = None) -> int:
        """Allocate a new identifier for the predicate (with modifier).

        Args:
            name: the name of the predicate.
            modifier: an optional modifier (e.g., not)

        Returns:
            the index of the predicate.
        """
        if not self._use_string_name:
            identifier = len(self.predicate2index)
        else:
            identifier = name + (f'_{modifier}' if modifier is not None else '')
        self.predicate2index[(name, modifier)] = identifier
        return identifier

    def get_grounded_predicate_indentifier(self, name: str, modifier: Optional[str] = None):
        return self.predicate2index[(name, modifier)]

    def compile_task(
        self,
        state: State,
        goal_expr: Union[str, Expression],
        actions: Optional[Sequence[OperatorApplier]] = None,
        regression_rules: Optional[Sequence[RegressionRuleApplier]] = None,
        is_relaxed: bool = False,
        forward_relevance_analysis: bool = True,
        backward_relevance_analysis: bool = True,
        verbose: bool = False
    ) -> GStripsProblem:
        """Compile a grounded STRIPS task.

        Args:
            state: the initial state.
            goal_expr: the goal expression.
            actions: the list of actions. If not specified, all actions in the domain will be used.
            is_relaxed: whether the task is a delete-relaxed task.
            forward_relevance_analysis: whether to perform forward relevance analysis.
            backward_relevance_analysis: whether to perform backward relevance analysis.
            verbose: whether to print verbose information.

        Returns:
            the compiled task.
        """
        with jacinle.cond_with(jacinle.time('compile_task::actions'), verbose):
            if actions is None:
                actions = gen_all_partially_grounded_actions(self._executor, state, filter_static=True)
        with jacinle.cond_with(jacinle.time('compile_task::regression_rules'), verbose):
            if regression_rules is None:
                regression_rules = gen_all_grounded_regression_rules(self._executor, state, filter_static=True)

        with jacinle.cond_with(jacinle.time('compile_task::state'), verbose):
            strips_state = self.compile_state(state)
        with jacinle.cond_with(jacinle.time('compile_task::operators'), verbose):
            strips_operators = [self.compile_operator(op, state, is_relaxed=is_relaxed) for op in actions]
        derived_predicates = list()
        if self._use_derived_predicates:
            with jacinle.cond_with(jacinle.time('compile_task::derived_predicates'), verbose):
                for pred in self.domain.functions.values():
                    if not pred.is_state_variable and pred.is_cacheable and pred.return_type == BOOL and not pred.is_static:
                        derived_predicates.extend(self.compile_derived_predicate(pred, state, is_relaxed=is_relaxed))
        regression_rules = list()
        with jacinle.cond_with(jacinle.time('compile_task::regression_rules'), verbose):
            strips_regression_rules = [self.compile_regression_rule(op, state) for op in regression_rules]
        with jacinle.cond_with(jacinle.time('compile_task::goal'), verbose):
            strips_goal, strips_goal_ip = self.compile_expr(goal_expr, state)
        task = GStripsProblem(strips_state, strips_goal, strips_operators, regression_rules=strips_regression_rules, is_relaxed=is_relaxed, goal_implicit_propositions=strips_goal_ip, derived_predicates=derived_predicates)
        with jacinle.cond_with(jacinle.time('compile_task::relevance_analysis'), verbose):
            if forward_relevance_analysis or backward_relevance_analysis:
                task = self.relevance_analysis(task, forward=forward_relevance_analysis, backward=backward_relevance_analysis)
        return task.compile()

    def recompile_relaxed_task(self, task: GStripsProblem, forward_relevance_analysis: bool = True, backward_relevance_analysis: bool = True) -> GStripsProblem:
        """Recompile a task to a delete-relaxed task.

        Args:
            task: the task to be recompiled.
            forward_relevance_analysis: whether to perform forward relevance analysis.
            backward_relevance_analysis: whether to perform backward relevance analysis.

        Returns:
            the recompiled task.
        """
        new_operators = list(task.operators)
        for i, op in enumerate(new_operators):
            new_effects = []
            for e in op.effects:
                new_e = e.relax()
                if isinstance(new_e, (tuple, list)):  # Note that GSSASAssignment will be relaxed into a tuple.
                    new_effects.extend(new_e)
                else:
                    new_effects.append(new_e)
            new_operators[i] = GStripsOperator(op.precondition, new_effects, op.raw_operator)
        task = GStripsProblem(task.state, task.goal, new_operators, is_relaxed=True, derived_predicates=[dp.relax() for dp in task.derived_predicates], facts=task.facts)
        if forward_relevance_analysis or backward_relevance_analysis:
            task = self.relevance_analysis(task, forward=forward_relevance_analysis, backward=backward_relevance_analysis)
        return task.compile()

    def recompile_task_new_state(
        self,
        task: GStripsProblem, new_state: Union[State, SStateCompatible],
        forward_relevance_analysis: bool = True, backward_relevance_analysis: bool = True,
        forward_derived: bool = False
    ) -> GStripsProblem:
        """
        Compile a new GStripsTask from a new state.

        Args:
            task: the original task.
            new_state: the new state.
            forward_relevance_analysis: whether to perform forward relevance analysis. Defaults to True.
            backward_relevance_analysis: whether to perform backward relevance analysis. Defaults to True.
            forward_derived: whether to forward derived predicates. Defaults to False.

        Returns:
            the new task.
        """
        if isinstance(new_state, State):
            new_state = self.compile_state(new_state.clone(), forward_derived=forward_derived)
        if task.facts is not None:
            new_state = new_state & task.facts
        task = GStripsProblem(new_state, task.goal, task.operators, is_relaxed=task.is_relaxed, derived_predicates=task.derived_predicates, facts=task.facts)
        if forward_relevance_analysis or backward_relevance_analysis:
            task = self.relevance_analysis(task, forward=forward_relevance_analysis, backward=backward_relevance_analysis)
            return task.compile()
        return task


class GStripsTranslatorOptimistic(GStripsTranslatorBase):
    def _init_indices(self):
        for pred in _find_cached_predicates(self.domain):
            if pred.return_type == BOOL:
                self.define_grounded_predicate(pred.name)
                self.define_grounded_predicate(pred.name, 'not')
            else:
                self.define_grounded_predicate(pred.name, 'initial')
                self.define_grounded_predicate(pred.name, 'optimistic')

    def compose_grounded_predicate(
        self, predicate_app: E.FunctionApplicationExpression,
        negated: bool = False, optimistic: Optional[bool] = None, allow_set: bool = False, return_argument_indices: bool = False
    ) -> Union[SProposition, Tuple[SProposition, List[int]]]:
        state = self.executor.state
        arguments = list()
        for arg_index, arg in enumerate(predicate_app.arguments):
            assert isinstance(arg, (E.ObjectConstantExpression, E.VariableExpression))
            if isinstance(arg, E.ObjectConstantExpression):
                arg = state.get_typed_index(arg.name)
                if allow_set:
                    arg = [arg]
            else:
                if arg.variable.name == '??':
                    assert allow_set
                    arg = list(range(state.get_nr_objects_by_type(predicate_app.function.arguments[arg_index].typename)))
                else:
                    arg = self.executor.get_bounded_variable(arg.variable).index
                    if allow_set:
                        arg = [arg]
            assert isinstance(arg, list) if allow_set else isinstance(arg, int)
            arguments.append(arg)

        if predicate_app.return_type == BOOL:
            assert optimistic is None
            if allow_set:
                rv = set(
                    _format_proposition((self.get_grounded_predicate_indentifier(predicate_app.function.name, 'not' if negated else None),) + tuple(args))
                    for args in itertools.product(*arguments)
                )
            else:
                rv = _format_proposition((self.get_grounded_predicate_indentifier(predicate_app.function.name, 'not' if negated else None),) + tuple(arguments))
        else:
            assert not negated and optimistic is not None
            modifier = 'optimistic' if optimistic else 'initial'
            if allow_set:
                rv = set(
                    _format_proposition((self.get_grounded_predicate_indentifier(predicate_app.function.name, modifier),) + tuple(args))
                    for args in itertools.product(*arguments)
                )
            else:
                rv = _format_proposition((self.get_grounded_predicate_indentifier(predicate_app.function.name, modifier),) + tuple(arguments))

        if return_argument_indices:
            return rv, arguments
        return rv

    # @jacinle.log_function
    def compose_bool_expression(self, expr: E.ValueOutputExpression, negated: bool = False) -> Tuple[Union[GSBoolOutputExpression, GSOptimisticStaticObjectType], Set[SProposition]]:
        state = self.executor.state
        predicate_app, this_negated = split_simple_bool(expr, initial_negated=negated)
        if predicate_app is not None:
            # jacinle.log_function.print(predicate_app, this_negated)
            if predicate_app.function.is_static:
                if predicate_app.return_type == BOOL:
                    if predicate_app.function.name in state.features:
                        _, arguments = self.compose_grounded_predicate(predicate_app, this_negated, return_argument_indices=True)
                        init_value = state.features[predicate_app.function.name][tuple(arguments)]
                    else:
                        init_value = self.executor.execute(expr)
                    return GSBoolConstantExpression(bool(init_value) ^ this_negated), set()
                else:
                    return GS_OPTIMISTIC_STATIC_OBJECT, set()

            if predicate_app.return_type == BOOL:
                return GSSimpleBoolExpression(self.compose_grounded_predicate(predicate_app, this_negated, allow_set=True), is_disjunction=True), set()
            else:
                return GSSimpleBoolExpression(self.compose_grounded_predicate(predicate_app, this_negated, optimistic=True, allow_set=True), is_disjunction=True), set()
        elif E.is_not_expr(expr):
            return self.compose_bool_expression(expr.arguments[0], negated=not negated)
        elif E.is_and_expr(expr) and not negated or E.is_or_expr(expr) and negated:
            classifiers = [self.compose_bool_expression(e, negated=negated) for e in expr.arguments]
            rv = gs_compose_bool_expressions(classifiers, is_disjunction=False)
            return rv
        elif E.is_and_expr(expr) and negated or E.is_or_expr(expr) and not negated:
            classifiers = [self.compose_bool_expression(e, negated=negated) for e in expr.arguments]
            return gs_compose_bool_expressions(classifiers, is_disjunction=True)
        elif E.is_forall_expr(expr) and not negated or E.is_exists_expr(expr) and negated:
            classifiers = list()
            for index in range(state.get_nr_objects_by_type(expr.variable.typename)):
                with self.executor.new_bounded_variables({expr.variable: index}):
                    classifiers.append(self.compose_bool_expression(expr.expr, negated=negated))
            return gs_compose_bool_expressions(classifiers, is_disjunction=False)
        elif E.is_forall_expr(expr) and negated or E.is_exists_expr(expr) and not negated:
            classifiers = list()
            for index in range(state.get_nr_objects_by_type(expr.variable.typename)):
                with self.executor.new_bounded_variables({expr.variable: index}):
                    classifiers.append(self.compose_bool_expression(expr.expression, negated=negated))
            return gs_compose_bool_expressions(classifiers, is_disjunction=True)
        elif isinstance(expr, E.DeicticSelectExpression):
            classifiers = list()
            for index in range(state.get_nr_objects_by_type(expr.variable.typename)):
                with self.executor.new_bounded_variables({expr.variable: index}):
                    classifiers.append(self.compose_bool_expression(expr.expr, negated=negated))
            return gs_compose_bool_expressions(classifiers, is_disjunction=True)
        elif isinstance(expr, E.ConditionalSelectExpression):
            classifiers = [
                self.compose_bool_expression(expr.condition),
                self.compose_bool_expression(expr.predicate)
            ]
            return gs_compose_bool_expressions(classifiers, is_disjunction=False)
        elif isinstance(expr, E.PredicateEqualExpression):
            argument_values = [self.compose_bool_expression(arg) for arg in [expr.predicate, expr.value]]
            has_optimistic_object = any(c[0] == GS_OPTIMISTIC_STATIC_OBJECT for c in argument_values)
            if has_optimistic_object:
                if expr.return_type == BOOL:
                    return GSBoolConstantExpression(True), _extract_all_propositions(argument_values)
                else:
                    return GS_OPTIMISTIC_STATIC_OBJECT, _extract_all_propositions(argument_values)

            argument_values = [argv for argv in argument_values if not gs_is_empty_bool_expression(argv)]
            init_value = self.executor.execute(expr)
            if init_value.item() > self._prob_goal_threshold and not negated or init_value.item() < 1 - self._prob_goal_threshold and negated:
                return GSBoolConstantExpression(True), _extract_all_propositions(argument_values)
            else:
                return gs_compose_bool_expressions(argument_values, is_disjunction=True)
        elif isinstance(expr, E.FunctionApplicationExpression):
            argument_values = [self.compose_bool_expression(arg) for arg in expr.arguments]
            # jacinle.log_function.print(argument_values)
            if GS_OPTIMISTIC_STATIC_OBJECT in argument_values:
                if expr.return_type == BOOL:
                    return GSBoolConstantExpression(True), _extract_all_propositions(argument_values)
                else:
                    return GS_OPTIMISTIC_STATIC_OBJECT, _extract_all_propositions(argument_values)

            argument_values = [argv for argv in argument_values if not gs_is_empty_bool_expression(argv)]
            # jacinle.log_function.print('computing initial value.')
            if expr.return_type == BOOL:
                # Theoretically, we can compute these values bottom-up together with the transformation.
                # In practice, this requires much more code to do...
                init_value = self.executor.execute(expr)
                # jacinle.log_function.print('computed initial value:', init_value)
                if init_value.item() > self._prob_goal_threshold and not negated or init_value.item() < 1 - self._prob_goal_threshold and negated:
                    return GSBoolConstantExpression(True), _extract_all_propositions(argument_values)
                else:
                    return gs_compose_bool_expressions(argument_values, is_disjunction=True)
            else:
                return gs_compose_bool_expressions(argument_values, is_disjunction=True)
        elif isinstance(expr, E.VariableExpression):
            assert expr.return_type != BOOL
            assert isinstance(self.executor.get_bounded_variable(expr.variable), TensorValue), 'Most likely you are accessing a non-optimistic object.'
            assert isinstance(self.executor.get_bounded_variable(expr.variable).single_elem(), OptimisticValue)
            return GS_OPTIMISTIC_STATIC_OBJECT, set()
        else:
            raise TypeError('Unsupported expression grounding: {}.'.format(expr))

    def compose_grounded_assignment(self, assignments: Sequence[E.VariableAssignmentExpression], is_relaxed: bool = False) -> Tuple[List[GSVariableAssignmentExpression], Set[SProposition]]:
        state = self.executor.state
        add_effects = set()
        del_effects = set()
        implicit_propositions = set()
        outputs = list()
        for assign_expr in assignments:
            if isinstance(assign_expr, E.AssignExpression):
                feat = assign_expr.predicate
                if feat.return_type == BOOL:
                    assert E.is_constant_bool_expr(assign_expr.value)
                    if assign_expr.value.constant.item():
                        add_effects.add(self.compose_grounded_predicate(feat, negated=False))
                        if not is_relaxed:
                            del_effects.add(self.compose_grounded_predicate(feat, negated=True))
                    else:
                        add_effects.add(self.compose_grounded_predicate(feat, negated=True))
                        if not is_relaxed:
                            del_effects.add(self.compose_grounded_predicate(feat, negated=False))
                else:
                    # For customized feature types, the "feat(...)" means that "this state variable has been set to an optimistic value."
                    add_effects.add(self.compose_grounded_predicate(feat, optimistic=True))
                    if not is_relaxed:
                        del_effects.add(self.compose_grounded_predicate(feat, optimistic=False))
                    value, ip = self.compose_bool_expression(assign_expr.value)
                    implicit_propositions = ip
                    if isinstance(value, GSBoolOutputExpression):
                        implicit_propositions.update(set(value.iter_propositions()))
            elif isinstance(assign_expr, E.ConditionalAssignExpression):
                assignment, ass_ip = self.compose_grounded_assignment([E.AssignExpression(assign_expr.predicate, assign_expr.value)], is_relaxed=is_relaxed)
                condition_classifier, cond_ip = self.compose_bool_expression(assign_expr.condition)
                outputs.append(GSConditionalAssignExpression(condition_classifier, assignment[0]))
                implicit_propositions = cond_ip | ass_ip
            elif isinstance(assign_expr, E.DeicticAssignExpression):
                for index in range(state.get_nr_objects_by_type(assign_expr.variable.typename)):
                    with self.executor.new_bounded_variables({assign_expr.variable: index}):
                        assignments, ass_ip = self.compose_grounded_assignment([assign_expr.expr], is_relaxed=is_relaxed)
                        for assignment in assignments:
                            if isinstance(assignment, GSSimpleBoolAssignExpression):
                                add_effects.update(assignment.add_effects)
                                del_effects.update(assignment.del_effects)
                            else:
                                outputs.append(assignment)
                        implicit_propositions.update(ass_ip)
        if len(add_effects) > 0 or len(del_effects) > 0:
            outputs.append(GSSimpleBoolAssignExpression(add_effects, del_effects))
        return outputs, implicit_propositions

    def compile_expr(self, expr: Union[str, Expression], state: State) -> Tuple[GSBoolOutputExpression, Set[SProposition]]:
        expr = self.domain.parse(expr)
        expr = flatten_expression(expr)
        with self.executor.with_state(state):
            return self.compose_bool_expression(expr)

    def compile_operator(self, op: OperatorApplier, state: State, is_relaxed=False) -> GStripsOperator:
        # print('compile_operator:: {}'.format(op))
        if getattr(op.operator, 'flatten_precondition', None) is None:
            ctx = ExpressionDefinitionContext(*op.operator.arguments, domain=self.domain)
            precondition = E.AndExpression(*[flatten_expression(e.bool_expr, ctx=ctx) for e in op.operator.preconditions])
            op.operator.flatten_precondition = precondition
        else:
            precondition = op.operator.flatten_precondition

        if getattr(op.operator, 'flatten_effects', None) is None:
            ctx = ExpressionDefinitionContext(*op.operator.arguments, domain=self.domain)
            effects = [flatten_expression(e.assign_expr, ctx=ctx) for e in op.operator.effects]
            op.operator.flatten_effects = effects
        else:
            effects = op.operator.flatten_effects

        # print('  precondition: {}'.format(precondition))
        with self.executor.with_state(state), self.executor.with_bounded_variables(compose_bvdict_args(op.operator.arguments, op.arguments, state=state)):
            precondition, pre_ip = self.compose_bool_expression(precondition)
            effects, eff_ip = self.compose_grounded_assignment(effects, is_relaxed=is_relaxed)
        # print('  compiled precondition: {}'.format(precondition))

        return GStripsOperator(precondition, effects, op, implicit_propositions=pre_ip | eff_ip)

    def compile_state(self, state: State, forward_derived: bool = False) -> SState:
        predicates = set()
        for name, feature in state.features.items():
            if self.domain.functions[name].is_static:
                continue
            if not self.domain.functions[name].is_state_variable:
                continue
            if feature.dtype == BOOL:
                for args, v in _iter_value(feature):
                    if v > 0.5:
                        predicates.add(_format_proposition((self.get_grounded_predicate_indentifier(name),) + args))
                    else:
                        predicates.add(_format_proposition((self.get_grounded_predicate_indentifier(name, 'not'),) + args))
            else:
                for args, _ in _iter_value(feature):
                    predicates.add(_format_proposition((self.get_grounded_predicate_indentifier(name, 'initial'),) + args))
        return SState(predicates)

    def compile_derived_predicate(self, dp: Predicate, state: State, is_relaxed=False) -> List[GStripsDerivedPredicate]:
        raise NotImplementedError('Derived predicates are not supported in Optimistic GStrips translation.')

    def relevance_analysis(self, task: GStripsProblem, relaxed_relevance: bool = False, forward: bool = True, backward: bool = True) -> GStripsProblem:
        return relevance_analysis_v1(task, relaxed_relevance=relaxed_relevance, forward=forward, backward=backward)


"""
class GStripsTranslatorSAS(GStripsTranslatorBase):
    def __init__(
        self,
        session: PDSketchExecutor,
        use_string_name: Optional[bool] = True,
        prob_goal_threshold: float = 0.5,
        cache_bool_predicates: bool = False
    ):
        self.cache_bool_predicates = cache_bool_predicates
        super().__init__(session, use_string_name, prob_goal_threshold, use_derived_predicates=cache_bool_predicates)

    def _init_indices(self):
        for pred in _find_cached_predicates(self.domain, allow_cacheable_bool=self.cache_bool_predicates):
            if pred.return_type == BOOL:
                self.define_grounded_predicate(pred.name)
                self.define_grounded_predicate(pred.name, 'not')
            else:
                for i in range(pred.ao_discretization.size):
                    self.define_grounded_predicate(f'{pred.name}@{i}')
                    self.define_grounded_predicate(f'{pred.name}@{i}', 'not')

    def compose_grounded_predicate_strips(
        self, ctx: ExpressionExecutionContext, feature_app: SE.StripsBoolPredicate,
        negated: bool = False
    ) -> Union[GSSimpleBoolExpression, GSBoolConstantExpression]:
        argument_indices = list()
        for arg_index, arg in enumerate(feature_app.arguments):
            argument_indices.append(ctx.get_bounded_variable(arg))

        feature_name = feature_app.sas_name if isinstance(feature_app, SE.SSASPredicateApplicationExpression) else feature_app.name
        predicate_def = self.domain.functions[feature_name]
        if predicate_def.is_static:
            if isinstance(feature_app, SE.SSASPredicateApplicationExpression):
                value = ctx.state.features[feature_name].tensor_indices[tuple(argument_indices)]
                return GSBoolConstantExpression(value.item() == feature_app.sas_index ^ negated ^ feature_app.negated)
            else:
                value = ctx.state.features[feature_name][tuple(argument_indices)]
                return GSBoolConstantExpression((value.item() > 0.5) ^ negated ^ feature_app.negated)

        predicate_name = self.get_grounded_predicate_indentifier(feature_app.name, 'not' if negated ^ feature_app.negated else None)
        return GSSimpleBoolExpression({_format_proposition((predicate_name,) + tuple(argument_indices))})

    def compose_grounded_predicate(self, ctx: ExpressionExecutionContext, feature_app: E.FunctionApplicationExpression, negated: bool = False) -> Union[GSSimpleBoolExpression, GSBoolConstantExpression]:
        argument_indices = list()
        for arg_index, arg in enumerate(feature_app.arguments):
            if isinstance(arg, E.ObjectConstantExpression):
                arg = ctx.state.get_typed_index(arg.name)
            else:
                assert isinstance(arg, E.VariableExpression)
                arg = ctx.get_bounded_variable(arg.variable)
            argument_indices.append(arg)
        predicate_def = feature_app.function
        feature_name = predicate_def.name

        if predicate_def.is_static:
            value = ctx.state.features[feature_name][tuple(argument_indices)]
            assert value.dtype == BOOL
            return GSBoolConstantExpression((value.item() > 0.5) ^ negated)

        predicate_name = self.get_grounded_predicate_indentifier(predicate_def.name, 'not' if negated else None)
        return GSSimpleBoolExpression({_format_proposition((predicate_name,) + tuple(argument_indices))})

    def _compose_grounded_classifier_strips(self, ctx: ExpressionExecutionContext, expr: SExpression, negated: bool = False) -> Union[GSBoolExpression, SProposition, GSOptimisticStaticObjectType]:
        if isinstance(expr, SE.SBoolConstant):
            return GSBoolConstantExpression(expr.constant ^ negated)
        elif isinstance(expr, SE.SBoolNot):
            return self._compose_grounded_classifier_strips(ctx, expr.expr, not negated)
        elif isinstance(expr, SE.SBoolExpression):
            classifiers = [self._compose_grounded_classifier_strips(ctx, e, negated) for e in expr.arguments]
            if expr.is_conjunction and not negated or expr.is_disjunction and negated:
                return gs_compose_bool_expressions(classifiers, is_disjunction=False, propagate_implicit_propositions=False)
            else:
                return gs_compose_bool_expressions(classifiers, is_disjunction=True, propagate_implicit_propositions=False)
        elif isinstance(expr, SE.SQuantificationExpression):
            classifiers = list()
            for index in range(ctx.state.get_nr_objects_by_type(expr.variable.typename)):
                with ctx.new_bounded_variables({expr.variable: index}):
                    classifiers.append(self._compose_grounded_classifier_strips(ctx, expr.expr, negated))
            if expr.is_forall and not negated or expr.is_exists and negated:
                return gs_compose_bool_expressions(classifiers, is_disjunction=False, propagate_implicit_propositions=False)
            else:
                return gs_compose_bool_expressions(classifiers, is_disjunction=True, propagate_implicit_propositions=False)
        elif isinstance(expr, SE.StripsBoolPredicate):
            return self.compose_grounded_predicate_strips(ctx, expr, negated)
        else:
            raise TypeError('Unknown expression type: {}.'.format(expr))

    def compose_grounded_classifier(
        self,
        ctx: ExpressionExecutionContext,
        expr: E.ValueOutputExpression,
        negated: bool = False
    ) -> Union[GSBoolExpression, GSOptimisticStaticObjectType]:
        if isinstance(expr, E.FunctionApplicationExpression):
            predicate_def = expr.function
            assert predicate_def.is_cacheable and predicate_def.return_type == BOOL
            if predicate_def.expr is None or self.cache_bool_predicates:  # a basic predicate.
                return self.compose_grounded_predicate(ctx, expr, negated)
            else:
                return self._compose_grounded_classifier_strips(ctx, predicate_def.ao_discretization, negated)
        elif E.is_not_expr(expr):
            return self.compose_grounded_classifier(ctx, expr.arguments[0], negated=not negated)
        elif E.is_and_expr(expr) and not negated or E.is_or_expr(expr) and negated:
            classifiers = [self.compose_grounded_classifier(ctx, e, negated=negated) for e in expr.arguments]
            rv = gs_compose_bool_expressions(classifiers, is_disjunction=False, propagate_implicit_propositions=False)
            return rv
        elif E.is_and_expr(expr) and negated or E.is_or_expr(expr) and not negated:
            classifiers = [self.compose_grounded_classifier(ctx, e, negated=negated) for e in expr.arguments]
            return gs_compose_bool_expressions(classifiers, is_disjunction=True, propagate_implicit_propositions=False)
        elif E.is_forall_expr(expr) and not negated or E.is_exists_expr(expr) and negated:
            classifiers = list()
            for index in range(ctx.state.get_nr_objects_by_type(expr.variable.typename)):
                with ctx.new_bounded_variables({expr.variable: index}):
                    classifiers.append(self.compose_grounded_classifier(ctx, expr.expr, negated=negated))
            return gs_compose_bool_expressions(classifiers, is_disjunction=False, propagate_implicit_propositions=False)
        elif E.is_forall_expr(expr) and negated or E.is_exists_expr(expr) and not negated:
            classifiers = list()
            for index in range(ctx.state.get_nr_objects_by_type(expr.variable.typename)):
                with ctx.new_bounded_variables({expr.variable: index}):
                    classifiers.append(self.compose_grounded_classifier(ctx, expr.expr, negated=negated))
            return gs_compose_bool_expressions(classifiers, is_disjunction=True, propagate_implicit_propositions=False)
        else:
            raise TypeError('Unsupported expression grounding: {}.'.format(expr))

    def _compose_grounded_assignment_strips(
        self,
        ctx: ExpressionExecutionContext,
        assignments: Sequence[SE.SVariableAssignmentExpression],
        is_relaxed: bool = False
    ) -> List[GSAssignmentExpression]:
        add_effects = set()
        del_effects = set()
        outputs = list()

        for expr in assignments:
            if isinstance(expr, SE.SDeicticAssignment):
                for index in range(ctx.state.get_nr_objects_by_type(expr.variable.typename)):
                    with ctx.new_bounded_variables({expr.variable: index}):
                        this_outputs = self._compose_grounded_assignment_strips(ctx, [expr.expression], is_relaxed)
                        for ass in this_outputs:
                            if isinstance(ass, GSSimpleAssignment):
                                add_effects.update(ass.add_effects)
                                del_effects.update(ass.del_effects)
                            else:
                                outputs.append(ass)
            elif isinstance(expr, SE.SConditionalAssignment):
                assignments = self._compose_grounded_assignment_strips(ctx, [expr.assign_op], is_relaxed)
                condition = self._compose_grounded_classifier_strips(ctx, expr.condition)
                for ass in assignments:
                    if isinstance(ass, GSSimpleAssignment):
                        outputs.append(GStripsConditionalAssignment(condition, ass))
                    elif isinstance(ass, GStripsConditionalAssignment):
                        outputs.append(GStripsConditionalAssignment(gs_compose_bool_expressions([condition, ass.condition], propagate_implicit_propositions=False), ass.assignment))
                    else:
                        raise TypeError('Invalid assignment type: {}.'.format(ass))
            elif isinstance(expr, SE.SAssignment):
                if isinstance(expr.predicate, SE.SSASPredicateApplicationExpression):
                    if is_relaxed:
                        raise NotImplementedError('Relaxed assignment to SAS predicate not supported during compilation. First compile it without is_relaxed, and re-run recompile_relaxed_operators.')
                    feature = expr.predicate
                    feature_name = feature.sas_name
                    predicate_def = self.domain.functions[feature_name]
                    feature_sas_size = predicate_def.ao_discretization.size
                    assert isinstance(expr.value, SE.SSASExpression)
                    argument_indices = list()
                    for arg_index, arg in enumerate(feature.arguments):
                        argument_indices.append(ctx.get_bounded_variable(arg))
                    expression = {k: self._compose_grounded_classifier_strips(ctx, v) for k, v in expr.value.mappings.items()}
                    sas_assignment = GSSASAssignment(feature_name, feature_sas_size, argument_indices, expression)
                    outputs.extend(sas_assignment.to_conditional_assignments())
                else:
                    feature = expr.predicate
                    value = bool(expr.value)
                    if value:
                        add_effects.add(self.compose_grounded_predicate_strips(ctx, feature))
                        if not is_relaxed:
                            add_effects.add(self.compose_grounded_predicate_strips(ctx, feature, negated=True))
                    else:
                        add_effects.add(self.compose_grounded_predicate_strips(ctx, feature, negated=True))
                        if not is_relaxed:
                            add_effects.add(self.compose_grounded_predicate_strips(ctx, feature))

        if len(add_effects) > 0 or len(del_effects) > 0:
            outputs.append(GSSimpleAssignment(add_effects, del_effects))
        return outputs

    def compile_expr(self, expr: Union[str, Expression], state: State) -> Tuple[GSBoolExpression, Set[SProposition]]:
        expr = self.domain.parse(expr)
        expr = flatten_expression(expr, flatten_cacheable_bool=not self.cache_bool_predicates)
        ctx = ExpressionExecutionContext(self._executor, state, {})
        return self.compose_grounded_classifier(ctx, expr), set()

    def compile_operator(self, op: OperatorApplier, state: State, is_relaxed=False) -> GStripsOperator:
        ctx = ExpressionExecutionContext(self._executor, state, compose_bvdict_args(op.operator.arguments, op.arguments, state=state))
        preconditions = list()
        for pred in op.operator.preconditions:
            preconditions.append(self._compose_grounded_classifier_strips(ctx, pred.ao_discretization))
        precondition = gs_compose_bool_expressions(preconditions, is_disjunction=False, propagate_implicit_propositions=False)
        effects = self._compose_grounded_assignment_strips(ctx, [eff.ao_discretization for eff in op.operator.effects], is_relaxed)
        return GStripsOperator(precondition, effects, op, implicit_propositions=set())

    def compile_derived_predicate(self, dp: Predicate, state: State, is_relaxed=False) -> List[GSDerivedPredicate]:
        arguments = list()
        for arg in dp.arguments:
            arguments.append(range(state.get_nr_objects_by_type(arg.dtype.typename)))

        rvs = list()
        for arg_indices in itertools.product(*arguments):
            bounded_variables = dict()
            for arg, arg_index in zip(dp.arguments, arg_indices):
                bounded_variables.setdefault(arg.typename, dict())[arg.name] = arg_index
            rvs.append(GSDerivedPredicate(
                dp.name, arg_indices,
                self._compose_grounded_classifier_strips(ctx, dp.ao_discretization),
                self._compose_grounded_classifier_strips(ctx, dp.ao_discretization, negated=True),
                is_relaxed=is_relaxed
            ))
        return rvs

    def compile_state(self, state: State, forward_derived: bool = False) -> SState:
        # Note: this function will change the original values of the state.
        # So be sure to make a copy of the state before calling this function.
        # This copying behavior is implemented in the compile_task function. If you are calling this function
        # directly, make sure to copy the state before calling this function.

        if forward_derived and self.cache_bool_predicates:
            self._executor.forward_predicates_and_axioms(state, forward_state_variables=False, forward_derived=True, forward_axioms=False)

        for name, feature in state.features.items():
            predicate_def = self.domain.functions[name]
            if predicate_def.is_state_variable and not (predicate_def.return_type == BOOL):
                state.features[name].tensor_indices = predicate_def.ao_discretization.quantize(feature).tensor

        predicates = set()
        for name, feature in state.features.items():
            predicate_def = self.domain.functions[name]
            if predicate_def.is_state_variable or (self.cache_bool_predicates and predicate_def.return_type == BOOL):
                if feature.dtype == BOOL:
                    for args, v in _iter_value(feature):
                        if v > 0.5:
                            predicates.add(_format_proposition((self.get_grounded_predicate_indentifier(name),) + args))
                        else:
                            predicates.add(_format_proposition((self.get_grounded_predicate_indentifier(name, 'not'),) + args))
                else:
                    codebook = predicate_def.ao_discretization
                    quantized_feature = codebook.quantize(feature)
                    for args, v in _iter_value(quantized_feature):
                        v = int(v)
                        for i in range(codebook.size):
                            if i == v:
                                predicates.add(_format_proposition((self.get_grounded_predicate_indentifier(f'{name}@{i}'),) + args))
                            else:
                                predicates.add(_format_proposition((self.get_grounded_predicate_indentifier(f'{name}@{i}', 'not'),) + args))

        return SState(predicates)

    def compile_task(
        self,
        state: State,
        goal_expr: Union[str, Expression],
        actions: Optional[Sequence[OperatorApplier]] = None,
        is_relaxed = False,
        forward_relevance_analysis: bool = True,
        backward_relevance_analysis: bool = True,
        verbose: bool = False
    ) -> GStripsProblem:
        state = state.clone()
        if self.cache_bool_predicates:
            self._executor.forward_predicates_and_axioms(state, forward_state_variables=True, forward_derived=True, forward_axioms=False)
        return super().compile_task(
            state, goal_expr, actions, is_relaxed,
            forward_relevance_analysis=forward_relevance_analysis, backward_relevance_analysis=backward_relevance_analysis,
            verbose=verbose
        )

    def relevance_analysis(self, task: GStripsProblem, relaxed_relevance: bool = False, forward: bool = True, backward: bool = True) -> GStripsProblem:
        return relevance_analysis_v2(task, relaxed_relevance=relaxed_relevance, forward=forward, backward=backward)
"""


class GStripsTranslatorFromAtomicStrips(GStripsTranslatorBase):
    def compile_expr(self, expr: Union[str, Expression], state: State) -> Tuple[GSBoolOutputExpression, Set[SProposition]]:
        raise NotImplementedError()

    def compile_operator(self, op: OperatorApplier, state: State, is_relaxed=False) -> GStripsOperator:
        raise NotImplementedError()

    def compile_regression_rule(self, op: OperatorApplier, state: State) -> GStripsRegressionRule:
        raise NotImplementedError()

    def compile_state(self, state: State, forward_derived: bool = False) -> SState:
        raise NotImplementedError()

    def compile_derived_predicate(self, dp: Predicate, state: State, is_relaxed=False) -> List[GStripsDerivedPredicate]:
        raise NotImplementedError('Derived predicates are not supported in Optimistic GStrips translation.')

    def relevance_analysis(self, task: GStripsProblem, relaxed_relevance: bool = False, forward: bool = True, backward: bool = True) -> GStripsProblem:
        return relevance_analysis_v1(task, relaxed_relevance=relaxed_relevance, forward=forward, backward=backward)


def _find_cached_predicates(domain: Domain, allow_cacheable_bool: bool = False) -> Iterable[Predicate]:
    """
    Return the set of predicates that are either in the `basic` or the `augmented` group.
    When the flag allow_cacheable_bool is set to True, also return the set of boolean predicates that are cacheable.

    Args:
        domain: the domain to search for predicates
        allow_cacheable_bool: whether to return the set of boolean predicates that are cacheable

    Returns:
        the set of predicates that are either in the `basic` or the `augmented` group and optionally cacheable boolean predicates.
    """
    for f in domain.functions.values():
        if f.is_state_variable:
            yield f
        elif allow_cacheable_bool and f.is_cacheable and f.return_type == BOOL:
            yield f


def _iter_value(value: TensorValue) -> Iterator[Tuple[Tuple[int, ...], Union[bool, int, float]]]:
    indices = [list(range(value.tensor.size(i))) for i in range(value.total_batch_dims)]
    for args, x in zip(itertools.product(*indices), value.tensor.flatten().tolist()):
        yield args, x


def _extract_all_propositions(classifiers: Sequence[Tuple[GSBoolOutputExpression, Set[SProposition]]]) -> Set[SProposition]:
    return set.union(*[c[1] for c in classifiers], *[c[0].iter_propositions() for c in classifiers if isinstance(c[0], GSBoolOutputExpression)])


def _format_proposition(pred_application: tuple[Union[SPredicateName, int], ...]) -> SProposition:
    return make_sproposition(*pred_application)


def relevance_analysis_v1(task: GStripsProblem, relaxed_relevance: bool = False, forward: bool = True, backward: bool = True) -> GStripsProblem:
    """
    Relevance analysis for a task.

    Args:
        task: the StripsTask to be analyzed.
        relaxed_relevance: whether to use relaxed relevance analysis. If this is set to False, we will not drop functions that are "implicitly" used by
            the operators. One example is that if one of the effect of the operator is a function (instead of simply Y/N).
        forward: whether to perform forward relevance analysis.
        backward: whether to perform backward relevance analysis.

    Returns:
        the updated task, removing all irrelevant propositions and operators.
    """

    # forward_analysis. a.k.a. static analysis.
    # useful when most objects are "non-movable". Then we don't need to track
    # their state variables/pose variables.

    # print('relevance_analysis::before')
    # print(task)
    if len(task.derived_predicates) > 0:
        raise RuntimeError('relevance_analysis::task.derived_predicates is not supported in relevance_analysis_v1.')

    if forward:
        task.compile()
        achievable_facts = set(task.state)
        changed = True
        visited = [False for _ in range(len(task.operators))]
        while changed:
            old_lengths = len(achievable_facts)
            for i, op in enumerate(task.operators):
                if not visited[i] and op.applicable(achievable_facts):
                    for e in op.effects:
                        achievable_facts.update(e.add_effects)
                    visited[i] = True
            changed = len(achievable_facts) != old_lengths
        new_operators = [op for i, op in enumerate(task.operators) if visited[i]]

        relevant_facts = set()
        for op in new_operators:
            for e in op.effects:
                relevant_facts.update(e.iter_propositions())

        # Note:: it seems that even when the task is not relaxed, we can also
        # do this drop...
        # Basically, if goal - relevant_facts isn't a subset of the initial
        # state, the problem is just unsolvable.
        # But when there are disjunctions, it's a bit hard to check.
        # relevant_facts.update(task.goal.iter_propositions())

        new_state = task.state.intersection(relevant_facts)
        new_operators = [op.filter_propositions(relevant_facts, state=task.state) for op in new_operators]
        new_goal = task.goal.filter_propositions(relevant_facts, state=task.state)

        task = GStripsProblem(new_state, new_goal, new_operators, is_relaxed=task.is_relaxed, facts=relevant_facts)
        task.compile()

        # print('relevance_analysis::forward')
        # print(task)

    # backward analysis.
    if backward:
        relevant_facts = set()
        relevant_facts.update(task.goal.iter_propositions())
        relevant_facts.update(task.goal_implicit_propositions)

        op_eff_facts = list()
        for op in task.operators:
            effects = set()
            for e in op.effects:
                effects.update(e.iter_propositions())
            op_eff_facts.append(effects)

        changed = True
        while changed:
            old_lengths = len(relevant_facts)
            for op, eff_facts in zip(task.operators, op_eff_facts):
                if set.intersection(eff_facts, relevant_facts):
                    relevant_facts |= set(op.precondition.iter_propositions())
                    if not relaxed_relevance:
                        relevant_facts |= set(op.implicit_propositions)
                    for e in op.effects:
                        if isinstance(e, GSConditionalAssignExpression):
                            relevant_facts |= set(e.condition.iter_propositions())
            changed = len(relevant_facts) != old_lengths

        new_operators = list()
        for op in task.operators:
            new_op = op.filter_propositions(relevant_facts, state=task.state)
            empty = True
            for e in new_op.effects:
                assert isinstance(e, (GSSimpleBoolAssignExpression, GSConditionalAssignExpression))
                if len(e.add_effects.symmetric_difference(e.del_effects)) > 0:
                    empty = False
                    break
                if isinstance(e, GSConditionalAssignExpression) and len(list(e.condition.iter_propositions())) > 0:
                    empty = False
                    break
            if not empty:
                new_operators.append(new_op)

        new_state = task.state.intersection(relevant_facts)
        task = GStripsProblem(new_state, task.goal, new_operators, is_relaxed=task.is_relaxed, facts=relevant_facts)

        # print('relevance_analysis::backward')
        # print(task)

    return task


def relevance_analysis_v2(task: GStripsProblem, relaxed_relevance: bool = False, forward: bool = True, backward: bool = True) -> GStripsProblem:
    """
    Relevance analysis for a task.

    Args:
        task: the StripsTask to be analyzed.
        relaxed_relevance: whether to use relaxed relevance analysis. If this is set to False, we will not drop functions that are "implicitly" used by
        the operators. One example is that if one of the effect of the operator is a function (instead of simply Y/N).
        forward: whether to run the forward pruning (forward reachability checking).
        backward: whether to run the backward pruning (goal regression).

    Returns:
        the updated task, removing all irrelevant propositions and operators.
    """

    # forward_analysis. a.k.a. static analysis.
    # useful when most objects are "non-movable". Then we don't need to track
    # their state variables/pose variables.

    # import ipdb; ipdb.set_trace()
    # print('relevance_analysis::before')
    # print(task)

    if forward:
        # collect all operators and derived predicates applicable.
        used_ops = set()
        used_dps = set()
        task.compile()
        achievable_facts = set(task.state)
        for i, dp in enumerate(task.derived_predicates):
            for j, eff in enumerate(dp.effects):
                if eff.applicable(achievable_facts):
                    achievable_facts.update(eff.add_effects)
                    used_dps.add((i, j))

        changed = True
        while changed:
            old_lengths = len(achievable_facts)
            for i, op in enumerate(task.operators):
                applicable = op.applicable(achievable_facts)
                if applicable:
                    for j, eff in enumerate(op.effects):
                        if (i, j) not in used_ops:
                            if isinstance(eff, GSSimpleBoolAssignExpression) or isinstance(eff, GSConditionalAssignExpression) and eff.applicable(achievable_facts):
                                achievable_facts.update(eff.add_effects)
                                used_ops.add((i, j))
            for i, dp in enumerate(task.derived_predicates):
                for j, eff in enumerate(dp.effects):
                    if (i, j) not in used_dps and eff.applicable(achievable_facts):
                        achievable_facts.update(eff.add_effects)
                        used_dps.add((i, j))
            changed = len(achievable_facts) != old_lengths

        new_operators = list()
        for i, op in enumerate(task.operators):
            used_effects = list()
            for j, eff in enumerate(op.effects):
                if (i, j) in used_ops:
                    used_effects.append(eff)
            if len(used_effects) > 0:
                new_operators.append(GStripsOperator(op.precondition, used_effects, op.raw_operator, op.implicit_propositions))
        new_derived_predicates = list()
        for i, dp in enumerate(task.derived_predicates):
            used_effects = list()
            for j, eff in enumerate(dp.effects):
                if (i, j) in used_dps:
                    used_effects.append(eff)
            if len(used_effects) > 0:
                new_derived_predicates.append(GStripsDerivedPredicate(dp.name, dp.arguments, effects=used_effects))

        relevant_facts = set()
        for op in new_operators:
            for e in op.effects:
                relevant_facts.update(e.iter_propositions())
        for i, dp in enumerate(task.derived_predicates):
            for e in dp.effects:
                relevant_facts.update(e.assignment.iter_propositions())

        # Note:: it seems that even when the task is not relaxed, we can also
        # do this drop...
        # Basically, if goal - relevant_facts isn't a subset of the initial
        # state, the problem is just unsolvable.
        # But when there are disjunctions, it's a bit hard to check.
        # relevant_facts.update(task.goal.iter_propositions())

        new_state = task.state.intersection(relevant_facts)
        new_operators = [op.filter_propositions(relevant_facts, state=task.state) for op in new_operators]
        new_derived_predicates = [dp.filter_propositions(relevant_facts, state=task.state) for dp in task.derived_predicates]
        new_goal = task.goal.filter_propositions(relevant_facts, state=task.state)

        task = GStripsProblem(new_state, new_goal, new_operators, is_relaxed=task.is_relaxed, derived_predicates=new_derived_predicates, facts=relevant_facts)
        task.compile()

        # import ipdb; ipdb.set_trace()
        # print('relevance_analysis::forward')
        # print(task)

    if backward:
        # backward analysis.
        relevant_facts = set()
        relevant_facts.update(task.goal.iter_propositions())
        relevant_facts.update(task.goal_implicit_propositions)

        changed = True
        while changed:
            old_lengths = len(relevant_facts)
            for i, dp in enumerate(task.derived_predicates):
                for j, eff in enumerate(dp.effects):
                    if set.intersection(set(eff.iter_propositions()), relevant_facts):
                        relevant_facts.update(eff.condition.iter_propositions())

            for i, op in enumerate(task.operators):
                for j, eff in enumerate(op.effects):
                    if set.intersection(set(eff.iter_propositions()), relevant_facts):
                        relevant_facts |= set(op.precondition.iter_propositions())
                        if not relaxed_relevance:
                            relevant_facts |= set(op.implicit_propositions)
                        if isinstance(eff, GSConditionalAssignExpression):
                            relevant_facts |= set(eff.condition.iter_propositions())
            changed = len(relevant_facts) != old_lengths

        new_operators = list()
        for op in task.operators:
            new_op = op.filter_propositions(relevant_facts, state=task.state)
            new_effects = list()
            for j, eff in enumerate(new_op.effects):
                assert isinstance(eff, (GSSimpleBoolAssignExpression, GSConditionalAssignExpression))
                if len(eff.add_effects.symmetric_difference(eff.del_effects)) > 0:
                    new_effects.append(eff)
            if len(new_effects) > 0:
                new_op = GStripsOperator(new_op.precondition, new_effects, new_op.raw_operator, new_op.implicit_propositions)
                new_operators.append(new_op)

        new_derived_predicates = list()
        for dp in task.derived_predicates:
            new_dp = dp.filter_propositions(relevant_facts, state=task.state)
            new_effects = list()
            for j, eff in enumerate(new_dp.effects):
                if len(eff.add_effects.symmetric_difference(eff.del_effects)) > 0:
                    new_effects.append(eff)
            if len(new_effects) > 0:
                new_dp = GStripsDerivedPredicate(new_dp.name, new_dp.arguments, effects=new_effects)
                new_derived_predicates.append(new_dp)

        new_state = task.state.intersection(relevant_facts)
        task = GStripsProblem(new_state, task.goal, new_operators, is_relaxed=task.is_relaxed, derived_predicates=new_derived_predicates, facts=relevant_facts)

        # import ipdb; ipdb.set_trace()
        # print('relevance_analysis::backward')
        # print(task)

    return task

concepts/pdsketch/strips/strips_grounding.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : strips_grounding_onthefly.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/18/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Strips grounding and seawrch using on-the-fly grounding of objects.

TODO
----

- [ ] Rewrite this file using the AtomicStrips interface.
"""

import time
import itertools
import warnings
import functools
import heapq as hq
import jacinle  # noqa
from collections import deque
from typing import Optional, Union, Sequence, Tuple, List, Dict

from concepts.dsl.dsl_types import Variable, ObjectConstant
from concepts.dsl.expression import ObjectConstantExpression, AndExpression
from concepts.pdsketch.predicate import Predicate
from concepts.pdsketch.operator import Operator
from concepts.pdsketch.domain import Domain, Problem
from concepts.pdsketch.strips.strips_expression import SProposition, SState, SStateDict, SBoolPredicateApplicationExpression
from concepts.pdsketch.strips.strips_grounded_expression import GSBoolOutputExpression
from concepts.pdsketch.strips.atomic_strips_domain import AtomicStripsOperator


class GoalNotAConjunctionError(Exception):
    pass


class OnTheFlyGStripsProblem(object):
    def __init__(
        self,
        objects: Dict[str, Sequence[str]],
        object2index: Dict[str, int],
        initial_state: SStateDict,
        predicates: Dict[str, Predicate],
        operators: Dict[str, AtomicStripsOperator],
        constants: Dict[str, ObjectConstant],
        conjunctive_goal: Optional[Sequence[SProposition]] = None,
        complex_goal: Optional[GSBoolOutputExpression] = None,
        use_integer_constants: bool = False
    ):
        self.objects = {k: tuple(v) for k, v in objects.items()}
        self.object2index = object2index
        self.initial_state = initial_state
        self.predicates = predicates
        self.operators = operators
        self.constants = constants.copy()

        assert (conjunctive_goal is None) ^ (complex_goal is None), 'Only one of conjunctive_goal and complex_goal can be specified.'
        self.conjunctive_goal = frozenset(conjunctive_goal) if conjunctive_goal is not None else None
        self.complex_goal = complex_goal
        self.complex_goal_func = complex_goal.compile() if complex_goal is not None else None

        self.use_integer_constants = use_integer_constants

    def goal_test(self, state: SState):
        if self.conjunctive_goal is not None:
            return self.conjunctive_goal.issubset(state)
        else:
            return self.complex_goal_func(state)

    @property
    def has_complex_goal(self):
        return self.complex_goal is not None

    @classmethod
    def from_domain_and_problem(cls, domain: Domain, problem: Problem, use_integer_constants: bool = False) -> 'OnTheFlyGStripsProblem':
        operators = dict()
        for operator in domain.operators.values():
            if isinstance(operator, Operator):
                if operator.is_macro:
                    continue
                # TODO(Jiayuan Mao @ 2023/03/19): support macro operator here.
                strips_operator = AtomicStripsOperator.from_operator(operator)
                operators[operator.name] = strips_operator
        constants = domain.constants

        objects = dict()
        object2index = dict()

        for name, constant in constants.items():
            typename = constant.dtype.typename
            if name not in object2index:
                if typename not in objects:
                    objects[typename] = list()
                objects[typename].append(name)
                object2index[name] = len(objects[typename]) - 1

        for k, typename in problem.objects.items():
            if k not in object2index:
                if typename not in objects:
                    objects[typename] = list()
                objects[typename].append(k)
                object2index[k] = len(objects[typename]) - 1

        initial_state = SStateDict()
        for predicate in problem.predicates:
            name = predicate.function.name
            if use_integer_constants:
                args = [object2index[arg.constant.name] for arg in predicate.arguments]
            else:
                args = [arg.constant.name for arg in predicate.arguments]
            initial_state.add(name, args)

        conjunctive_goal = None
        complex_goal = None
        try:
            conjunctive_goal = list()
            if not isinstance(problem.goal, AndExpression):
                raise GoalNotAConjunctionError()
            for arg in problem.goal.arguments:
                predicate_name = arg.function.name
                args = list()
                for arg in arg.arguments:
                    if not isinstance(arg, ObjectConstantExpression):
                        raise GoalNotAConjunctionError()
                    arg = arg.constant.name
                    args.append(arg)
                conjunctive_goal.append(f'{predicate_name} {" ".join(args)}')
        except GoalNotAConjunctionError:
            from concepts.pdsketch.executor import PDSketchExecutor
            from concepts.pdsketch.strips.strips_grounding import GStripsTranslatorOptimistic
            executor = PDSketchExecutor(domain)
            tensor_state = problem.to_state(executor)
            translator = GStripsTranslatorOptimistic(executor)
            with executor.with_state(tensor_state):
                gstrips_goal, _ = translator.compose_bool_expression(problem.goal)
            conjunctive_goal = None
            complex_goal = gstrips_goal

        return cls(
            objects, object2index, initial_state,
            predicates=domain.functions.copy(),
            operators=operators,
            constants=constants,
            conjunctive_goal=conjunctive_goal,
            complex_goal=complex_goal,
            use_integer_constants=use_integer_constants,
        )

    def decode_plan(self, plan: Sequence[Tuple[AtomicStripsOperator, Dict[str, Union[int, str]]]]) -> List[Tuple[AtomicStripsOperator, Dict[str, str]]]:
        if not self.use_integer_constants:
            return plan

        decoded_plan = list()
        for op, bound_arguments in plan:
            decoded_bound_arguments = dict()
            for arg in op.arguments:
                assert arg.name in bound_arguments, f'Argument {arg.name} is not bound in operator {op.name}.'
                decoded_bound_arguments[arg.name] = self.objects[arg.dtype.typename][bound_arguments[arg.name]]
            decoded_plan.append((op, decoded_bound_arguments))
        return decoded_plan

    def __str__(self) -> str:
        return f"""OnTheFlyGStripsProblem(
  objects={self.objects},
  initial_state={self.initial_state},
  goal={self.conjunctive_goal if self.conjunctive_goal is not None else self.complex_goal},
  operators={self.operators},
  constants={self.constants}
)"""

    def __repr__(self) -> str:
        return self.__str__()


def ogstrips_expand_state_with_negation(problem: OnTheFlyGStripsProblem, state: SState) -> SState:
    warnings.warn('ogstrips_expand_state_with_negation is deprecated. You should avoid using it because it is slow.', DeprecationWarning)
    expanded_state = SStateDict()

    for predicate in problem.predicates.values():
        if problem.use_integer_constants:
            options = itertools.product(*[range(len(problem.objects[arg.typename])) for arg in predicate.arguments])
        else:
            options = itertools.product(*[problem.objects[arg.typename] for arg in predicate.arguments])
        for args in options:
            args = tuple(args)
            if state.contains(predicate.name, args):
                expanded_state.add(predicate.name, args)
            else:
                expanded_state.add(f'{predicate.name}_not', args)

    return expanded_state


def ogstrips_bind_arguments(predicate: SBoolPredicateApplicationExpression, bound_arguments: Dict[str, Union[int, str]], object2index: Dict[str, int]):
    return predicate.name, tuple(bound_arguments[arg.name] if isinstance(arg, Variable) else arg for arg in predicate.arguments)

    # TODO(Jiayuan Mao @ 2023/03/27): bring back.
    boudned_arguments = list()
    for arg in predicate.arguments:
        if isinstance(arg, Variable):
            boudned_arguments.append(bound_arguments[arg.name])
        elif isinstance(arg, str):
            boudned_arguments.append(object2index.get(arg, 100000))
        else:
            raise ValueError(f'Unexpected argument type: {type(arg)}')
    return predicate.name, tuple(boudned_arguments)


def ogstrips_generate_applicable_actions(problem: OnTheFlyGStripsProblem, state: SStateDict, check_negation: bool = False) -> List[AtomicStripsOperator]:
    TOO_MANY, FAILED, PASS = object(), object(), object()

    def compute_possible_grounding(predicate: SBoolPredicateApplicationExpression, bound_arguments: Dict[str, Union[int, str]]):
        unbound_arguments = [arg for arg in predicate.arguments if isinstance(arg, Variable) and arg.name not in bound_arguments]
        if len(unbound_arguments) == 0:
            name, arguments = ogstrips_bind_arguments(predicate, bound_arguments, problem.object2index)
            rv = state.contains(name, arguments, predicate.negated, check_negation=check_negation)
            if not rv:
                return '', FAILED
            return '', PASS
        elif len(unbound_arguments) == 1:
            arg = unbound_arguments[0]
            valid_arguments = list()

            if problem.use_integer_constants:
                options = range(len(problem.objects[arg.typename]))
            else:
                options = problem.objects[arg.typename]

            for o in options:
                bound_arguments[arg.name] = o
                name, arguments = ogstrips_bind_arguments(predicate, bound_arguments, problem.object2index)
                rv = state.contains(name, arguments, predicate.negated, check_negation=check_negation)
                if rv:
                    valid_arguments.append(o)
                del bound_arguments[arg.name]
            return arg.name, valid_arguments
        else:
            return '', TOO_MANY

    # @jacinle.log_function(verbose=False)
    def dfs(preconditions: Sequence[SBoolPredicateApplicationExpression], bound_arguments: Dict[str, int]):
        """Inner DFS function.

        Args:
            preconditions: the preconditions to be satisfied.
            bound_arguments: a mapping from variable name to object.
        """

        # jacinle.log_function.print('dfs', bound_arguments, 'remaining preconditions:', len(preconditions))
        # import ipdb; ipdb.set_trace()

        for i, precondition in enumerate(preconditions):
            name, valid_arguments = compute_possible_grounding(precondition, bound_arguments)
            if valid_arguments == FAILED:
                # jacinle.log_function.print('Failed.')
                return list()
            elif valid_arguments == PASS:
                # jacinle.log_function.print('Pass.')
                return dfs(preconditions[:i] + preconditions[i + 1:], bound_arguments)
            elif valid_arguments == TOO_MANY:
                pass
            else:
                outputs = list()
                for arg in valid_arguments:
                    bound_arguments[name] = arg
                    outputs.extend(dfs(preconditions[:i] + preconditions[i + 1:], bound_arguments))
                    del bound_arguments[name]
                return outputs

        unbound_arguments = [arg for arg in operator.arguments if isinstance(arg, Variable) and arg.name not in bound_arguments]
        # print('unbound_arguments', unbound_arguments, bound_arguments)
        if len(unbound_arguments) == 0:
            # jacinle.log_function.print('Found a grounding:', bound_arguments)
            return [bound_arguments.copy()]

        if problem.use_integer_constants:
            unbound_arguments_possible_values = {arg.name: range(len(problem.objects[arg.typename])) for arg in unbound_arguments}
        else:
            unbound_arguments_possible_values = {arg.name: problem.objects[arg.typename] for arg in unbound_arguments}

        name, valid_arguments = min(unbound_arguments_possible_values.items(), key=lambda x: len(x[1]))
        outputs = list()
        for arg in valid_arguments:
            bound_arguments[name] = arg
            # jacinle.log_function.print('{} = {}'.format(name, arg))
            outputs.extend(dfs(preconditions, bound_arguments))
        del bound_arguments[name]
        return outputs

    for operator in problem.operators.values():
        # jacinle.log_function.print(f'operator: {operator.name}')
        for bound_arguments in dfs(operator.preconditions, dict()):
            # jacinle.log_function.print('yield bound_arguments:', bound_arguments)
            yield operator, bound_arguments


def ogstrips_check_precondition(state: SStateDict, operator: AtomicStripsOperator, bound_arguments: Dict[str, Union[int, str]], object2index: Dict[str, int]):
    for precondition in operator.preconditions:
        name, arguments = ogstrips_bind_arguments(precondition, bound_arguments, object2index)
        if not state.contains(name, arguments, precondition.negated):
            return False
    return True


def ogstrips_apply_operator(state: SStateDict, operator: AtomicStripsOperator, bound_arguments: Dict[str, Union[int, str]], object2index: Dict[str, int]):
    new_state = state.clone()
    for predicate in operator.del_effects:
        name, arguments = ogstrips_bind_arguments(predicate, bound_arguments, object2index)
        new_state.remove(name, arguments)
    for predicate in operator.add_effects:
        name, arguments = ogstrips_bind_arguments(predicate, bound_arguments, object2index)
        new_state.add(name, arguments)
    return new_state


def ogstrips_search(problem: OnTheFlyGStripsProblem, initial_actions: Sequence[Tuple[AtomicStripsOperator, Dict[str, Union[int, str]]]] = tuple(), max_expanded_nodes: int = 1000000, timeout: float = 10.0):
    frontier = deque()

    initial_state = problem.initial_state
    for operator, arguments in initial_actions:
        ogstrips_check_precondition(initial_state, operator, arguments, problem.object2index)
        initial_state = ogstrips_apply_operator(initial_state, operator, arguments, problem.object2index)

    frontier.append((initial_state, list(initial_actions)))
    explored = set()

    start_time = time.time()
    nr_expanded_nodes = 0
    while len(frontier) > 0:
        nr_expanded_nodes += 1
        if nr_expanded_nodes > max_expanded_nodes:
            break

        if nr_expanded_nodes % 100 == 0:
            if time.time() - start_time > timeout:
                print('ogstrips_search::Timeout.')
                break

        state, plan = frontier.popleft()

        # action_strings = [f"{operator.name}({', '.join(bound_arguments.values())})" for operator, bound_arguments in plan]
        # print('State', state, 'Plan', action_strings)
        # print('Plan', action_strings)

        for operator, bound_arguments in ogstrips_generate_applicable_actions(problem, state):
            new_state = ogstrips_apply_operator(state, operator, bound_arguments, problem.object2index)
            new_state_set = new_state.as_state()
            if new_state_set not in explored:
                if problem.goal_test(new_state_set):
                    return plan + [(operator, bound_arguments)]

                frontier.append((new_state, plan + [(operator, bound_arguments)]))
                explored.add(new_state_set)

    return None


def ogstrips_search_with_heuristics(
    problem: OnTheFlyGStripsProblem,
    initial_actions: Sequence[Tuple[AtomicStripsOperator, Dict[str, Union[int, str]]]] = tuple(),
    hfunc_name: str = 'hmax', h_weight: float = 1, g_weight: float = 1,
    max_expanded_nodes: int = 1000000, timeout: float = 10.0,
    verbose: bool = False, hfunc_verbose: bool = False
):

    if hfunc_name == 'hmax':
        hfunc = ogstrips_hmax
    elif hfunc_name == 'hadd':
        hfunc = ogstrips_hadd
    elif hfunc_name == 'hff':
        hfunc = ogstrips_hff
    else:
        raise ValueError(f'Unknown heuristic function name: {hfunc_name}.')

    queue = list()

    initial_state = problem.initial_state
    for operator, arguments in initial_actions:
        ogstrips_check_precondition(initial_state, operator, arguments, problem.object2index)
        initial_state = ogstrips_apply_operator(initial_state, operator, arguments, problem.object2index)

    counter = 0
    # NB(Jiayuan Mao @ 2023/04/04): only verbose printing the hfunc computation for the initial state.
    h = hfunc(problem, initial_state, verbose=hfunc_verbose)
    queue.append((h * h_weight, counter, h, initial_state, list()))
    counter += 1
    explored = set()

    if verbose:
        print('Initial heuristic:', queue[0][2])
        import ipdb; ipdb.set_trace()

    start_time = time.time()
    nr_expanded_nodes = 0
    while len(queue) > 0:
        nr_expanded_nodes += 1
        if nr_expanded_nodes > max_expanded_nodes:
            break

        if nr_expanded_nodes % 10 == 0:
            if time.time() - start_time > timeout:
                print('ogstrips_search_with_heuristics::Timeout.')
                break

        prio, _, h, state, plan = hq.heappop(queue)
        if verbose:
            print('  h =', h, 'prio =', prio, 'plan =', [f"{operator.name}({', '.join(bound_arguments.values())})" for operator, bound_arguments in plan])

        # action_strings = [f"{operator.name}({', '.join(bound_arguments.values())})" for operator, bound_arguments in plan]
        # print('State', state, 'Plan', action_strings)
        # print('Plan', action_strings)

        for operator, bound_arguments in ogstrips_generate_applicable_actions(problem, state):
            new_state = ogstrips_apply_operator(state, operator, bound_arguments, problem.object2index)
            new_state_set = new_state.as_state()
            if new_state_set not in explored:
                new_plan = plan + [(operator, bound_arguments)]
                h = hfunc(problem, new_state)

                if verbose:
                    print('    add new plan, h =', h, f'new action={operator.name}({", ".join(bound_arguments.values())})')

                # if problem.goal_test(new_state_set):
                if h == 0:
                    return problem.decode_plan(new_plan)

                hq.heappush(queue, (h * h_weight + len(new_plan) * g_weight, counter, h, new_state, new_plan))  # GBF!
                counter += 1
                explored.add(new_state_set)


def ogstrips_verify(problem: OnTheFlyGStripsProblem, plan: Sequence[str], from_fast_downward: bool = False):
    state = problem.initial_state

    # For some reason, all the names from Fast Downward are lower-cased.
    op_name_mapping = {op.lower(): op for op in problem.operators.keys()}
    object_name_mapping = {obj.lower(): obj for obj in problem.object2index.keys()}

    for plan_string in plan:
        assert plan_string.startswith('(') and plan_string.endswith(')')
        plan_string = plan_string[1:-1]
        operator_name, arguments = plan_string.split(' ', 1)
        arguments = arguments.split(' ')

        if from_fast_downward:
            operator_name = op_name_mapping[operator_name]
            arguments = [object_name_mapping[arg] for arg in arguments]

        operator = problem.operators[operator_name]
        bound_arguments = {arg.name: problem.object2index[argv] for arg, argv in zip(operator.arguments, arguments)}
        assert ogstrips_check_precondition(state, operator, bound_arguments, problem.object2index)
        state = ogstrips_apply_operator(state, operator, bound_arguments, problem.object2index)

    assert problem.goal_test(state.as_state())


def _ogstrips_backward_relavance(problem: OnTheFlyGStripsProblem, state_dict: SStateDict) -> Tuple[OnTheFlyGStripsProblem, SStateDict]:
    warnings.warn('This function is not implemented yet.', NotImplementedError)
    assert problem.conjunctive_goal is not None

    used_predicates = set()
    used_operators = set()

    for proposition in problem.conjunctive_goal:
        used_predicates.add(proposition.split()[0])

    while True:
        new_operator_added = False
        for operator in problem.operators.values():
            if operator.name in used_operators:
                continue

            op_useful = False
            for effect in itertools.chain(operator.add_effects, operator.del_effects):
                if effect.name in used_predicates:
                    op_useful = True
                    break

            if not op_useful:
                continue

            used_operators.add(operator.name)
            new_operator_added = True

            for precondition in operator.preconditions:
                used_predicates.add(precondition.name)

        if not new_operator_added:
            break

    return problem, state_dict


def ogstrips_delete_relaxation_heuristic(problem: OnTheFlyGStripsProblem, state_dict: SStateDict, htype: str, verbose: bool = False) -> float:
    """Heuristic function for the HFF planner.

    Args:
        state_dict: the current state.
        problem: the problem.
        htype: the type of heuristic function to use: 'hmax' or 'hadd' or 'hff'.

    Returns:
        the heuristic value.
    """
    assert htype in ('hadd', 'hmax', 'hff')

    state = state_dict.as_state()
    F_sets = [set(state)]
    A_sets = []
    F_to_A = dict()

    used_operators = set()

    goal_rv = problem.goal_test(state)
    while not goal_rv:
        if verbose:
            print(f'heuristic::current_level = {len(F_sets)}')
        new_state_dict = state_dict.clone()
        new_ops = list()
        new_facts = list()
        for operator, bound_arguments in ogstrips_generate_applicable_actions(problem, state_dict, check_negation=True):
            if verbose:
                print('  operator:', operator.name, bound_arguments)

            grounded_operator_identifier = (operator.name, ) + tuple(bound_arguments[x.name] for x in operator.arguments)
            if grounded_operator_identifier not in used_operators:
                new_ops.append(grounded_operator_identifier)
                used_operators.add(grounded_operator_identifier)

                for predicate in operator.add_effects:
                    name, arguments = ogstrips_bind_arguments(predicate, bound_arguments, problem.object2index)
                    if not new_state_dict.contains(name, arguments):
                        new_state_dict.add(name, arguments)
                        fact_name = f'{name} {" ".join(map(str, arguments))}'
                        new_facts.append(fact_name)
                        F_to_A[fact_name] = (operator, bound_arguments)
                for predicate in operator.del_effects:
                    name, arguments = ogstrips_bind_arguments(predicate, bound_arguments, problem.object2index)
                    name = name + '_not'
                    if not new_state_dict.contains(name, arguments):
                        new_state_dict.add(name, arguments)
                        fact_name = f'{name} {" ".join(map(str, arguments))}'
                        new_facts.append(fact_name)
                        F_to_A[fact_name] = (operator, bound_arguments)

        if len(new_facts) == 0:
            break

        F_sets.append(set(new_facts))
        A_sets.append(set(new_ops))

        state_dict = new_state_dict
        state = state_dict.as_state()
        goal_rv = problem.goal_test(state)

    if not goal_rv:
        return int(1e9)

    F_levels = dict()
    for i, F_set in enumerate(F_sets):
        for fact in F_set:
            F_levels[fact] = i

    if verbose:
        print('F_levels:')
        for fact, level in F_levels.items():
            print(f'  {fact} -> {level}')

    goal_propositions = list()
    if problem.has_complex_goal:
        with GSBoolOutputExpression.enable_forward_diff_ctx():
            # We need to compute the forward-diff of the goal expression.
            goal_rv = problem.goal_test(state)
        goal_propositions = list(goal_rv.propositions)
    else:
        goal_propositions = problem.conjunctive_goal

    if htype == 'hadd':
        h = 0
        for proposition in goal_propositions:
            h += F_levels[proposition]
        return h
    elif htype == 'hmax':
        h = 0
        for proposition in goal_propositions:
            h = max(h, F_levels[proposition])
        return h
    elif htype == 'hff':
        used_actions = set()

        def backtrace(fact):
            if fact in F_to_A:
                operator, bound_arguments = F_to_A[fact]
                grounded_operator_identifier = (operator.name, ) + tuple(bound_arguments[x.name] for x in operator.arguments)
                if grounded_operator_identifier not in used_actions:
                    used_actions.add(grounded_operator_identifier)
                    if verbose:
                        print('  hff::used_action:', operator.name, bound_arguments, '->', fact)
                    for precondition in operator.preconditions:
                        name, arguments = ogstrips_bind_arguments(precondition, bound_arguments, problem.object2index)
                        new_fact = f'{name} {" ".join(map(str, arguments))}'
                        if new_fact in F_levels and (F_levels[new_fact] < F_levels[fact]):
                            backtrace(new_fact)

        for proposition in goal_propositions:
            backtrace(proposition)

        # print('  hff::used_actions:', len(used_actions))
        return len(used_actions)


ogstrips_hadd = functools.partial(ogstrips_delete_relaxation_heuristic, htype='hadd')
ogstrips_hmax = functools.partial(ogstrips_delete_relaxation_heuristic, htype='hmax')
ogstrips_hff = functools.partial(ogstrips_delete_relaxation_heuristic, htype='hff')


concepts/pdsketch/strips/strips_grounding_onthefly.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : strips_heuristics.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 03/20/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from typing import Optional, Union, Tuple, List, Set, Dict, Callable

from concepts.pdsketch.strips.strips_expression import SProposition, SState
from concepts.pdsketch.strips.strips_grounded_expression import GSBoolForwardDiffReturn, GSBoolOutputExpression, GSSimpleBoolAssignExpression, GSConditionalAssignExpression
from concepts.pdsketch.strips.strips_grounding import GStripsProblem, GStripsTranslatorBase

__all__ = ['StripsHeuristic', 'StripsBlindHeuristic', 'StripsRPGHeuristic', 'StripsHFFHeuristic']


class StripsHeuristic(object):
    """The base class for STRIPS heuristics."""

    def __init__(self, task: GStripsProblem, translator: Optional[GStripsTranslatorBase] = None):
        """Initialize the heuristic.

        Args:
            task: the task to compute the heuristic value (including the goal and the operators)
            translator: the translator to translate the task representation into the `GS*` (grounded STRIPS) representation.
        """
        self._task = task
        self._goal_func = task.goal.compile()
        self._translator = translator

    @property
    def task(self) -> GStripsProblem:
        """The grounded STRIPS planning task."""
        return self._task

    @property
    def goal_func(self) -> Callable[[SState], Union[bool, GSBoolForwardDiffReturn]]:
        """The (compiled version) of goal function of the task."""
        return self._goal_func

    @property
    def translator(self) -> Optional[GStripsTranslatorBase]:
        """The translator used to translate the task representation into the `GS*` (grounded STRIPS) representation."""
        return self._translator

    @classmethod
    def from_type(cls, type_identifier: str, task: GStripsProblem, translator: Optional[GStripsTranslatorBase] = None, **kwargs) -> 'StripsHeuristic':
        """Create a heuristic from the given type identifier.

        Args:
            type_identifier: the type identifier of the heuristic.
            task: the grounded STRIPS planning task.
            translator: the translator used to translate the planning task to the `GS*` (grounded STRIPS) representation.
            **kwargs: additional arguments to pass to the constructor.

        Returns:
            the created heuristic.
        """
        if type_identifier == 'hff':
            return StripsHFFHeuristic(task, translator, **kwargs)
        elif type_identifier == 'blind':
            return StripsBlindHeuristic(task, translator)
        else:
            raise ValueError(f'Unknown heuristic type {type_identifier}.')

    def compute(self, state: SState) -> int:
        """Compute the heuristic value of the given state.

        Args:
            state: the state to compute the heuristic value.

        Returns:
            the heuristic value of the given state.
        """
        raise NotImplementedError()


class StripsBlindHeuristic(StripsHeuristic):
    """A blind heuristic that always returns 1 if the goal is not satisfied, and 0 otherwise."""

    def compute(self, state: SState) -> int:
        goal_rv = self._goal_func(state)
        return 0 if goal_rv else 1


class StripsRPGHeuristic(StripsHeuristic):
    """RPG heuristic (relaxed planning graph)."""

    def __init__(self, task: GStripsProblem, translator: Optional[GStripsTranslatorBase] = None, forward_relevance_analysis: bool = True, backward_relevance_analysis: bool = True):
        """Initialize an RPG heuristic.

        Args:
            task: the grounded STRIPS planning task.
            translator: the translator used to translate the planning task to the `GS*` (grounded STRIPS) representation.
            forward_relevance_analysis: whether to perform forward relevance analysis.
            backward_relevance_analysis: whether to perform backward relevance analysis.
        """
        super().__init__(task, translator)

        self._forward_relevance_analysis = forward_relevance_analysis
        self._backward_relevance_analysis = backward_relevance_analysis

        if task.is_relaxed:
            self._relaxed = task
        else:
            assert translator is not None
            self._relaxed = translator.recompile_relaxed_task(task, forward_relevance_analysis=forward_relevance_analysis, backward_relevance_analysis=backward_relevance_analysis)

    @property
    def relaxed(self) -> GStripsProblem:
        """The relaxed version of the task."""
        return self._relaxed

    def compute_rpg_forward_diff(self, state: SState, relaxed_task: Optional[GStripsProblem] = None) -> Tuple[
        List[Set[SState]],
        Dict[SProposition, int],
        List[List[Tuple[int, int, Set[SProposition]]]],
        List[List[Tuple[int, int, Set[SProposition]]]],
        GSBoolForwardDiffReturn
    ]:
        """Compute the relaxed planning graph using forward differentiation.

        Args:
            state: the state to compute the relaxed planning graph.
            relaxed_task: the relaxed version of the task. If not provided, the task will be relaxed using the translator.

        Returns:
            the relaxed planning graph, which is represented as a tuple of the following elements:
            - a list of sets of states, where each set contains all the states that can be reached from the initial state in the given number of steps.
            - a dictionary that maps propositions to their level.
            - a list of lists of tuples, where each tuple contains the following elements:
                - the index of the action in the relaxed task operator list.
                - the index of the effect in the operator.
                - the propositions that contributes to the effect (computed using Forward Diff).
            - a list of lists of tuples, where each tuple contains the following elements:
                - the index of the action that achieves the proposition.
                - the level of the proposition.
                - the set of propositions that are achieved by the action.
            - the result of forward differentiation of the goal function.
        """

        if relaxed_task is None:
            relaxed_task = self._relaxed

        with GSBoolOutputExpression.enable_forward_diff_ctx():
            F_sets = [set(state)]
            A_sets = []
            D_sets = []

            used_operators = set()
            used_derived_predicates = set()
            # print('rpginit', F_sets[-1])

            goal_rv = self._goal_func(F_sets[-1])
            while not goal_rv.rv:
                # for op in relaxed_task.operators:
                #     print(' ', op.raw_operator, op.precondition, op.applicable(F_sets[-1]))

                new_ops: List[Tuple[int, int, Set[SProposition]]] = list()  # op_index, eff_index, op_precondition
                # print('Starting new step')
                for i, op in enumerate(relaxed_task.operators):
                    op_pred_rv = None
                    for j, e in enumerate(op.effects):
                        if (i, j) not in used_operators:
                            if op_pred_rv is None:
                                op_pred_rv = op.applicable(F_sets[-1])
                                # print('Evaluating op', op.raw_operator, op_pred_rv)
                            if op_pred_rv.rv:
                                if isinstance(e, GSSimpleBoolAssignExpression):
                                    # TODO:: check if it should be (i, j) or (i, 0, op_pred_rv.propositions)
                                    new_ops.append((i, j, op_pred_rv.propositions))
                                    used_operators.add((i, j))
                                elif isinstance(e, GSConditionalAssignExpression):
                                    eff_pred_rv = e.applicable(F_sets[-1])
                                    if eff_pred_rv.rv:
                                        new_ops.append((i, j, op_pred_rv.propositions | eff_pred_rv.propositions))
                                        # print('  Use operator', i, j)
                                        used_operators.add((i, j))
                            else:
                                break

                new_F = F_sets[-1].copy()
                for op_index, effect_index, _ in new_ops:
                    op = relaxed_task.operators[op_index]
                    eff = op.effects[effect_index]
                    new_F.update(eff.add_effects)

                new_dps: List[Tuple[int, int, Set[SProposition]]] = list()  # dp_index, eff_index, dp_precondition
                for i, dp in enumerate(relaxed_task.derived_predicates):
                    for j, e in enumerate(dp.effects):
                        if (i, j) not in used_derived_predicates:
                            dp_pred_rv = e.applicable(new_F)
                            # print((i, j), dp_pred_rv)
                            if dp_pred_rv.rv:
                                used_derived_predicates.add((i, j))
                                new_dps.append((i, j, dp_pred_rv.propositions))

                for dp_index, effect_index, _ in new_dps:
                    dp = relaxed_task.derived_predicates[dp_index]
                    eff = dp.effects[effect_index]
                    new_F.update(eff.add_effects)

                # print('depth', len(F_sets), new_F)
                if len(new_F) == len(F_sets[-1]):
                    break

                A_sets.append(new_ops)
                D_sets.append(new_dps)
                F_sets.append(new_F)
                goal_rv = self._goal_func(F_sets[-1])

            F_levels = {}
            for i in range(0, len(F_sets)):
                for f in F_sets[i]:
                    if f not in F_levels:
                        F_levels[f] = i

            return F_sets, F_levels, A_sets, D_sets, goal_rv

    def compute(self, state: SState) -> int:
        raise NotImplementedError()


class StripsHFFHeuristic(StripsRPGHeuristic):
    def compute(self, state: SState) -> int:
        """Compute the hFF heuristic value for the given state.

        Args:
            state: the state to compute the heuristic value for.

        Returns:
            the heuristic value.
        """

        # NB(Jiayuan Mao @ 12/04): We don't need to do relevance analysis because the actual computation for hFF will do similar things anyway...
        relaxed = self._translator.recompile_task_new_state(self._relaxed, state, forward_relevance_analysis=False, backward_relevance_analysis=False)
        state = relaxed.state

        F_sets, F_levels, A_sets, D_sets, goal_rv = self.compute_rpg_forward_diff(state, relaxed_task=relaxed)

        if not goal_rv.rv:
            return int(1e9)

        selected: Set[Tuple[int, int]] = set()
        selected_facts = [set() for _ in F_sets]

        goal_propositions = goal_rv.propositions
        for pred in goal_propositions:
            if pred not in F_levels:
                return int(1e9)
            selected_facts[F_levels[pred]].add(pred)

        for t in reversed(list(range(len(A_sets)))):
            for dp_index, effect_index, propositions in D_sets[t]:
                dp = self._relaxed.derived_predicates[dp_index]
                eff = dp.effects[effect_index]
                if eff.add_effects.intersection(selected_facts[t+1]):
                    # print('Selecting Derived Predicate', dp_index, effect_index, propositions, eff.add_effects)
                    selected_facts[t + 1] -= eff.add_effects
                    for pred in propositions:
                        # print('  Selecting predicate', pred)
                        selected_facts[F_levels[pred]].add(pred)
            for op_index, effect_index, propositions in A_sets[t]:
                op = self._relaxed.operators[op_index]
                eff = op.effects[effect_index]
                if eff.add_effects.intersection(selected_facts[t+1]):
                    # print('Selecting Action', op_index, effect_index, op.raw_operator, propositions, eff.add_effects)
                    selected.add((op_index, effect_index))  # only add the operator or both? Maybe at each level, just add the operator? or add a flag?
                    selected_facts[t + 1] -= eff.add_effects
                    for pred in propositions:
                        # print('  Selecting predicate', pred)
                        selected_facts[F_levels[pred]].add(pred)
        return len(selected)

concepts/pdsketch/strips/strips_heuristics.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : strips_search.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 03/20/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import jacinle
import jactorch
from typing import Callable

from concepts.pdsketch.strips.strips_expression import SState
from concepts.pdsketch.strips.strips_grounding import GStripsProblem
from concepts.pdsketch.strips.strips_heuristics import StripsHeuristic
from concepts.algorithm.search.heuristic_search import run_heuristic_search

__all__ = ['strips_brute_force_search', 'strips_heuristic_search', 'get_priority_func']


def print_task(task: GStripsProblem, func):
    func_name = func.__name__

    print(f'{func_name}::task.goal={task.goal}')
    print(f'{func_name}::task.facts={len(task.facts) if task.facts is not None else "N/A"}')
    print(f'{func_name}::task.operators={len(task.operators)}')


@jactorch.no_grad_func
def strips_brute_force_search(
    task: GStripsProblem, max_depth=5, verbose=True
):
    if verbose:
        print_task(task, strips_brute_force_search)

    goal_func = task.goal.compile()
    states = [(task.state, tuple())]
    visited = set()
    visited.add(task.state)

    pbar = None
    if verbose:
        pbar = jacinle.tqdm_pbar(desc='strips_brute_force_search::depth=0')
    with jacinle.cond_with(pbar, verbose):
        for depth in range(max_depth):
            next_states = list()
            for s, traj in states:
                for a in task.operators:
                    if verbose:
                        pbar.update()
                    if a.applicable(s):
                        ns = a.apply(s)

                        if ns not in visited:
                            visited.add(ns)
                        nt = traj + (a, )
                        next_states.append((ns, nt))
                        if goal_func(ns):
                            return nt
            states = next_states
            if verbose:
                pbar.set_description(f'strips_brute_force_search::depth={depth}, states={len(states)}')
    return None


def get_priority_func(heuristic: Callable[[SState], float], weight: float) -> Callable[[SState, int], float]:
    if weight == 1:
        def priority_fun(state, g, heuristic=heuristic):
            return heuristic(state) + g
    elif weight == float('inf'):
        def priority_fun(state, g, heuristic=heuristic):
            return heuristic(state)
    else:
        def priority_fun(state, g, heuristic=heuristic, weight=weight):
            return heuristic(state) + g * weight
    return priority_fun


@jactorch.no_grad_func
def strips_heuristic_search(
    task: GStripsProblem, heuristic: StripsHeuristic, *,
    max_expansions=int(1e9), verbose=False,
    heuristic_weight=float('inf')
):
    if verbose:
        print_task(task, strips_heuristic_search)
        print('strips_heuristic_search::init_heuristic={}'.format(heuristic.compute(task.state)))

    goal_func = task.goal.compile()

    pbar = None
    if verbose:
        pbar = jacinle.tqdm_pbar(desc='strips_heuristic_search::expanding')

    def check_goal(state: SState, gf=goal_func):
        return gf(state)

    def get_successors(state: GStripsProblem, actions=task.operators):
        for action in actions:
            if verbose:
                pbar.update()
            if action.applicable(state):
                yield action, action.apply(state), 1

    with jacinle.cond_with(pbar, verbose):
        return run_heuristic_search(
            task.state,
            check_goal,
            get_priority_func(heuristic.compute, heuristic_weight),
            get_successors,
            max_expansions=max_expansions
        )[1]


concepts/pdsketch/strips/strips_search.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 06/12/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/sandbox/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/14/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/simulator/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 07/10/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/simulator/pybullet/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : camera.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import math
import collections
import numpy as np
import pybullet as p

from typing import Optional, Tuple
from .rotation_utils import rpy

__all__ = ['CameraConfig', 'RealSenseD415', 'TopDownOracle', 'RS200Gazebo', 'KinectFranka', 'Cliport', 'get_point_cloud', 'lookat_rpy']


class CameraConfig(collections.namedtuple('_CameraConfig', ['image_size', 'intrinsics', 'position', 'rotation', 'zrange', 'shadow'])):
    """
    Mostly based on https://github.com/cliport/cliport/blob/e9cde74754606448d8a0495c1efea36c29b201f1/cliport/tasks/cameras.py
    """

    def get_view_and_projection_matricies(self, image_size: Optional[Tuple[int, int]] = None) -> Tuple[np.ndarray, np.ndarray]:
        if image_size is None:
            image_size = self.image_size

        view_matrix = self.get_view_matrix()
        projection_matrix = self.get_projection_matrix(image_size)
        return view_matrix, projection_matrix

    def get_view_matrix(self) -> np.ndarray:
        lookdir = np.float32([0, 0, 1]).reshape(3, 1)
        updir = np.float32([0, -1, 0]).reshape(3, 1)
        rotation = p.getMatrixFromQuaternion(rpy(*self.rotation))
        rotm = np.float32(rotation).reshape(3, 3)
        lookdir = (rotm @ lookdir).reshape(-1)
        updir = (rotm @ updir).reshape(-1)
        lookat = self.position + lookdir
        return p.computeViewMatrix(self.position, lookat, updir)

    def get_projection_matrix(self, image_size: Tuple[int, int]) -> np.ndarray:
        focal_len = self.intrinsics[0]
        znear, zfar = self.zrange
        fovh = (image_size[0] / 2) / focal_len
        fovh = 180 * np.arctan(fovh) * 2 / np.pi        # Notes: 1) FOV is vertical FOV 2) aspect must be float
        aspect_ratio = float(image_size[1]) / image_size[0]
        return p.computeProjectionMatrixFOV(fovh, aspect_ratio, znear, zfar)


class RealSenseD415(object):
    """Default configuration with 3 RealSense RGB-D cameras."""

    # Mimic RealSense D415 RGB-D camera parameters.
    image_size = (480, 640)
    intrinsics = (450., 0, 320., 0, 450., 240., 0, 0, 1)

    # Set default camera poses.
    front_position = (1., 0, 0.75)
    front_rotation = np.rad2deg((np.pi / 4, np.pi, -np.pi / 2))
    left_position = (0, 0.5, 0.75)
    left_rotation = np.rad2deg((np.pi / 4.5, np.pi, np.pi / 4))
    right_position = (0, -0.5, 0.75)
    right_rotation = np.rad2deg((np.pi / 4.5, np.pi, 3 * np.pi / 4))

    @classmethod
    def get_configs(cls):
        return [
            CameraConfig(cls.image_size, cls.intrinsics, cls.front_position, cls.front_rotation, (0.1, 10.), True),
            CameraConfig(cls.image_size, cls.intrinsics, cls.left_position, cls.left_rotation, (0.1, 10.), True),
            CameraConfig(cls.image_size, cls.intrinsics, cls.right_position, cls.right_rotation, (0.1, 10.), True)
        ]


class TopDownOracle(object):
    """Top-down noiseless image used only by the oracle demonstrator."""

    # Near-orthographic projection.
    image_size = (480, 640)
    intrinsics = (63e4, 0, 320., 0, 63e4, 240., 0, 0, 1)
    position = (0.5, 0, 1000.)
    rotation = np.rad2deg((0, np.pi, -np.pi / 2))

    @classmethod
    def get_configs(cls):
        return [CameraConfig(cls.image_size, cls.intrinsics, cls.position, cls.rotation, (999.7, 1001.), False)]


class RS200Gazebo(object):
    """Gazebo Camera"""

    # Near-orthographic projection.
    image_size = (480, 640)
    intrinsics = (554.3826904296875, 0.0, 320.0, 0.0, 554.3826904296875, 240.0, 0.0, 0.0, 1.0)
    position = (0.5, 0, 1.0)
    rotation = np.rad2deg((0, np.pi, np.pi / 2))

    @classmethod
    def get_configs(cls):
        return [CameraConfig(cls.image_size, cls.intrinsics, cls.position, cls.rotation, (0.1, 10.), False)]


class KinectFranka(object):
    """Kinect Franka Camera"""

    # Near-orthographic projection.
    image_size = (424, 512)
    intrinsics = (365.57489013671875, 0.0, 257.5205078125, 0.0, 365.57489013671875, 205.26710510253906, 0.0, 0.0, 1.0)
    position = (1.082, -0.041, 1.027)
    rotation = np.rad2deg((-2.611, 0.010, 1.553))

    @classmethod
    def get_configs(cls):
        return [CameraConfig(cls.image_size, cls.intrinsics, cls.position, cls.rotation, (0.1, 10.), True)]


class Cliport(object):
    """Cliport camera configuration."""

    image_size = (480, 640)
    intrinsics = (450., 0, 320., 0, 450., 240., 0, 0, 1)
    position1 = (1.5, 0, 0.8)
    rotation1 = (0, -115, 90)
    position2 = (0, -1, 0.8)
    rotation2 = (-120, 20, 10)
    position3 = (0, 1, 0.8)
    rotation3 = (120, 20, 170)

    @classmethod
    def get_configs(cls):
        return [
            CameraConfig(cls.image_size, cls.intrinsics, cls.position1, cls.rotation1, (0.1, 10.), False),
            CameraConfig(cls.image_size, cls.intrinsics, cls.position2, cls.rotation2, (0.1, 10.), False),
            CameraConfig(cls.image_size, cls.intrinsics, cls.position3, cls.rotation3, (0.1, 10.), False)
        ]


def get_point_cloud(config, color, depth, segmentation=None):
    """Reconstruct point clouds from the color and depth images."""
    # based on https://stackoverflow.com/questions/59128880/getting-world-coordinates-from-opengl-depth-buffer
    height, width = depth.shape[0:2]
    view_matrix, proj_matrix = config.get_view_and_projection_matricies()

    # create a 4x4 transform matrix that goes from pixel coordinates (and depth values) to world coordinates
    proj_matrix = np.asarray(proj_matrix).reshape([4, 4], order="F")
    view_matrix = np.asarray(view_matrix).reshape([4, 4], order="F")
    tran_pix_world = np.linalg.inv(np.matmul(proj_matrix, view_matrix))

    # create a grid with pixel coordinates and depth values
    y, x = np.mgrid[-1:1:2 / height, -1:1:2 / width]
    y *= -1.
    x, y, z = x.reshape(-1), y.reshape(-1), depth.reshape(-1)
    h = np.ones_like(z)

    pixels = np.stack([x, y, z, h], axis=1)
    mask = z < 0.99
    # mask = ...
    # filter out "infinite" depths
    pixels = pixels[mask]
    pixels[:, 2] = 2 * pixels[:, 2] - 1

    # turn pixels to world coordinates
    points = np.matmul(tran_pix_world, pixels.T).T
    points /= points[:, 3: 4]
    points = points[:, :3]

    if segmentation is not None:
        return points, color.reshape(-1, color.shape[-1])[mask], segmentation.reshape(-1)[mask]
    return points, color.reshape(-1, color.shape[-1])[mask]


def lookat_rpy(camera_pos: np.ndarray, target_pos: np.ndarray, roll: float = 0) -> np.ndarray:
    """Construct the roll, pitch, yaw angles of a camera looking at a target.
    This function assumes that the camera is pointing to the z-axis ([0, 0, 1]),
    in the camera frame.
    Args:
        camera_pos: the position of the camera.
        target_pos: the target position.
        roll: the roll angle of the camera.

    Returns:
        a numpy array of the roll, pitch, yaw angles in degrees.
    """
    camera_pos = np.asarray(camera_pos)
    target_pos = np.asarray(target_pos)

    delta = target_pos - camera_pos
    pitch = math.atan2(-np.linalg.norm(delta[:2]), delta[2])
    yaw = math.atan2(-delta[1], -delta[0])
    rv = np.array([np.deg2rad(roll), pitch, yaw], dtype="float32")

    return np.rad2deg(rv)


concepts/simulator/pybullet/camera.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : client.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/17/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import time
import six
import os.path as osp
import tempfile
import threading
import collections
import functools
import contextlib
from typing import Any, Optional, Union, Tuple, List, Dict

import pybullet as p
import pybullet_data
import jacinle
import jacinle.io as io

from concepts.simulator.pybullet.world import BulletWorld

__all__ = ['BulletClient']


class BulletP(object):
    def __init__(self, client_id=None):
        self.client_id = client_id

    def set_client_id(self, client_id):
        self.client_id = client_id

    def __getattr__(self, item):
        assert self.client_id is not None
        func = getattr(p, item)
        if callable(func):
            return functools.partial(func, physicsClientId=self.client_id)
        return func


class MouseEvent(collections.namedtuple('_MouseEvent', ['eventType', 'mousePosX', 'mousePosY', 'buttonIndex', 'buttonState'])):
    pass


class BulletClient(object):
    DEFAULT_ENGINE_PARAMETERS = {'numSolverIterations': 10}
    DEFAULT_FPS = 120
    DEFAULT_GRAVITY = (0, 0, -9.8)
    DEFAULT_ASSETS_ROOT = osp.join(osp.dirname(osp.dirname(osp.dirname(osp.abspath(__file__)))), 'assets')

    def __init__(
        self,
        assets_root: Optional[str] = None,
        is_gui: bool = False,
        realtime_gui: bool = False,
        fps: Optional[int] = None,
        render_fps: Optional[int] = None,
        gravity: Optional[Union[Tuple[float], float]] = None,
        engine_parameters: Optional[Dict[str, Any]] = None,
        connect: bool = True,
        client_id: int = -1,
        save_video: Optional[str] = None,
        width: Optional[int] = 960,
        height: Optional[int] = 960,
    ):
        if not is_gui:
            render_fps = 0

        self.is_gui = is_gui
        self.realtime_gui = realtime_gui
        self.fps = fps if fps is not None else type(self).DEFAULT_FPS
        self.render_fps = render_fps if render_fps is not None else self.fps
        self.gravity = canonize_gravity(gravity if gravity is not None else type(self).DEFAULT_GRAVITY)
        self.engine_parameters = engine_parameters
        self.client_id = None
        self.assets_root = assets_root if assets_root is not None else type(self).DEFAULT_ASSETS_ROOT
        self.save_video = save_video

        self.w = BulletWorld()
        self.p = BulletP()
        self.width = width
        self.height = height

        if client_id == -1:
            if connect:
                self.connect()
        else:
            self.client_id = client_id
            self.w.set_client_id(self.client_id)
            self.p.set_client_id(self.client_id)

    @property
    def world(self):
        """Alias for `self.w`."""
        return self.w

    @contextlib.contextmanager
    def with_fps(self, fps: Optional[int] = None, render_fps: Optional[int] = None):
        current_fps, current_render_fps = self.fps, self.render_fps
        if fps is not None:
            self.fps = fps
        if render_fps is not None:
            self.render_fps = render_fps
        elif fps is not None:
            self.render_fps = fps
        yield
        self.fps, self.render_fps = current_fps, current_render_fps

    def connect(self, suppress_warnings: bool = True):
        if suppress_warnings:
            with jacinle.suppress_stdout():
                self._connect()
        else:
            self._connect()

    def _connect(self):
        options = ''
        if self.save_video:
            options += f'--mp4="{self.save_video}" --mp4fps=60'
        if self.width is not None:
            options += ' --width={}'.format(self.width)
        if self.height is not None:
            options += ' --height={}'.format(self.height)
        self.client_id = p.connect(p.GUI if self.is_gui else p.DIRECT, options=options)

        if self.save_video or (self.is_gui and self.realtime_gui):
            p.configureDebugVisualizer(p.COV_ENABLE_SINGLE_STEP_RENDERING, 1, physicsClientId=self.client_id)

        p.setAdditionalSearchPath(pybullet_data.getDataPath(), physicsClientId=self.client_id)
        if self.engine_parameters is not None:
            p.setPhysicsEngineParameter(physicsClientId=self.client_id, **self.engine_parameters)
        else:
            p.setPhysicsEngineParameter(physicsClientId=self.client_id, **type(self).DEFAULT_ENGINE_PARAMETERS)

        p.configureDebugVisualizer(p.COV_ENABLE_MOUSE_PICKING, 0, physicsClientId=self.client_id)
        p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 1, physicsClientId=self.client_id)
        p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 1, physicsClientId=self.client_id)
        p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 1, physicsClientId=self.client_id)

        # Disable the GUI (e.g., synthetic camera views and parameters).
        p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.client_id)
        # p.configureDebugVisualizer(p.COV_ENABLE_GUI, 1, physicsClientId=self.client_id)

        p.setPhysicsEngineParameter(enableFileCaching=0, physicsClientId=self.client_id)

        if self.assets_root is not None:
            file_io = p.loadPlugin('fileIOPlugin', physicsClientId=self.client_id)
            if file_io >= 0:
                p.executePluginCommand(file_io, textArgument=self.assets_root, intArgs=[p.AddFileIOAction], physicsClientId=self.client_id)
            else:
                raise RuntimeError('pybullet: cannot load FileIO!')
            p.setAdditionalSearchPath(self.assets_root, physicsClientId=self.client_id)

        # NB(Jiayuan Mao @ 10/04): also add the temp dir to the asset path so that we can load JIT URDF files.
        p.setAdditionalSearchPath(tempfile.gettempdir(), physicsClientId=self.client_id)

        p.setGravity(*self.gravity, physicsClientId=self.client_id)
        p.setTimeStep(1.0 / self.fps, physicsClientId=self.client_id)
        self.w.set_client_id(self.client_id)
        self.p.set_client_id(self.client_id)

    def is_connected(self):
        return p.isConnected(physicsClientId=self.client_id)

    def has_gui(self):
        return p.getConnectionInfo(physicsClientId=self.client_id)['connectionMethod'] == p.GUI

    def reset_world(self):
        p.resetSimulation(physicsClientId=self.client_id)
        p.setGravity(*self.gravity, physicsClientId=self.client_id)
        p.setTimeStep(1.0 / self.fps, physicsClientId=self.client_id)

        # Should also remember to reset the world record.
        self.w = BulletWorld()
        self.w.set_client_id(self.client_id)

    def disconnect(self):
        p.disconnect(physicsClientId=self.client_id)

    @contextlib.contextmanager
    def disable_rendering(self, disable_rendering: bool = True, reset: bool = False, suppress_stdout: bool = False):
        if reset:
            self.reset_world()

        with jacinle.cond_with(
            jacinle.suppress_stdout(),
            suppress_stdout
        ):
            if disable_rendering:
                p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 0, physicsClientId=self.client_id)
            yield
            if disable_rendering and self.is_connected():
                p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 1, physicsClientId=self.client_id)

    @contextlib.contextmanager
    def disable_world_update(self):
        """Temproary disable the world update. Specifically, when loading a new model, the world object `self.w` will not be updated.
        This function also disables rendering of the pybullet debug renderer.
        Thus, this functionality is useful when loading a large number of models."""
        with self.disable_rendering(suppress_stdout=True):
            yield

    def step(self, steps=1):
        clock = None
        if self.render_fps > 0:
            clock = jacinle.Clock(1 / self.render_fps)
        for i in range(steps):
            p.stepSimulation(physicsClientId=self.client_id)
            if self.render_fps > 0:
                clock.tick()

    def perform_collision_detection(self):
        p.performCollisionDetection(physicsClientId=self.client_id)

    def load_urdf(self, xml_path, pos=(0, 0, 0), quat=(0, 0, 0, 1), body_name: Optional[str] = None, group: Optional[str] = '__UNSET__', static=False, scale: float = 1.0, rgba=None, notify_world_update=True) -> int:
        xml_path = self._canonize_asset_path(xml_path)
        pos, quat = canonize_default_pos_and_quat(pos, quat)
        try:
            ret = p.loadURDF(xml_path, pos, quat, useFixedBase=static, globalScaling=scale, physicsClientId=self.client_id, flags=p.URDF_USE_SELF_COLLISION)
        except p.error as e:
            raise RuntimeError('pybullet: cannot load URDF file: {}'.format(xml_path)) from e
        if notify_world_update:
            if group == '__UNSET__':
                group = 'fixed' if static else 'rigid'
            self.w.notify_update(ret, body_name=body_name, group=group)
        if rgba is not None:
            self.w.change_visual_color(ret, rgba=rgba)
        return ret

    def load_urdf_template(self, xml_path: str, replaces: Dict[str, Any], pos=None, quat=None, **kwargs) -> int:
        xml_path = self._canonize_asset_path(xml_path)
        with open(xml_path) as f:
            xml_content = f.read()
        for k, v in sorted(replaces.items(), key=lambda x: len(x[0]), reverse=True):
            if isinstance(v, (tuple, list)):
                for i in range(len(v)):
                    xml_content = xml_content.replace(k + str(i), str(v[i]))
            else:
                xml_content = xml_content.replace(k, str(v))

        with io.tempfile('w', '.xml') as f:
            f.write(xml_content)
            f.flush()
            return self.load_urdf(f.name, pos=pos, quat=quat, **kwargs)

    def load_mjcf(self, xml_path, pos=(0, 0, 0), quat=(0, 0, 0, 1), body_name=None, group='__UNSET__', static=False, notify_world_update=True) -> int:
        xml_path = self._canonize_asset_path(xml_path)
        pos, quat = canonize_default_pos_and_quat(pos, quat)
        ret = p.loadMJCF(xml_path, pos, quat, useFixedBase=static, physicsClientId=self.client_id, flags=p.MJCF_COLORS_FROM_FILE)
        if notify_world_update:
            if group == '__UNSET__':
                group = 'fixed' if static else 'rigid'
            self.w.notify_update(ret, body_name=body_name, group=group)
        return ret

    def loads_mjcf(self, xml_content, pos=None, quat=None, save_to=None, **kwargs) -> int:
        if not isinstance(xml_content, six.string_types):
            xml_content = io.dumps_xml(xml_content)

        if save_to is not None:
            with open(save_to, 'w') as f:
                f.write(xml_content)

        with io.tempfile('w', '.xml') as f:
            f.write(xml_content)
            f.flush()
            return self.load_mjcf(f.name, pos=pos, quat=quat, **kwargs)

    def _canonize_asset_path(self, path):
        return path.replace('assets://', self.assets_root + '/')

    def remove_body(self, body_id):
        return p.removeBody(body_id, physicsClientId=self.client_id)

    def get_mouse_events(self) -> List[MouseEvent]:
        return list(MouseEvent(*event) for event in self.p.getMouseEvents())

    def update_viewer(self):
        self.get_mouse_events()

    def wait_for_duration(self, duration):
        t0 = time.time()
        while time.time() - t0 <= duration:
            self.update_viewer()

    def wait_forever(self):
        while True:
            self.update_viewer()

    def wait_for_user(self, message='Press enter to continue...'):
        import platform
        if self.has_gui() and platform.system() == 'Darwin':
            # OS X doesn't multi-thread the OpenGL visualizer
            return self._threaded_input(message)
        return input(message)

    def timeout(self, duration: float):
        return jacinle.timeout(duration, fps=self.fps)

    def _threaded_input(self, *args, **kwargs):
        # OS X doesn't multi-thread the OpenGL visualizer
        data = []
        thread = threading.Thread(target=lambda: data.append(input(*args, **kwargs)), args=[])
        thread.start()
        try:
            while thread.is_alive():
                self.update_viewer()
        finally:
            thread.join()
        return data[-1]


def canonize_gravity(gravity):
    if isinstance(gravity, (int, float)):
        return (0, 0, gravity)
    else:
        gravity = tuple(gravity)
        assert len(gravity) == 3
        return gravity


def canonize_default_pos_and_quat(pos, quat):
    if pos is None:
        pos = (0, 0, 0)
    if quat is None:
        quat = (0, 0, 0, 1)
    return pos, quat


concepts/simulator/pybullet/client.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 07/10/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/simulator/pybullet/components/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : component_base.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/24/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from concepts.simulator.pybullet.client import BulletClient


class BulletComponent(object):
    def __init__(self, client: BulletClient):
        self.client = client

    @property
    def p(self):
        return self.client.p

    @property
    def world(self):
        return self.client.w

    @property
    def w(self):
        return self.client.w


concepts/simulator/pybullet/components/component_base.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/23/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.



concepts/simulator/pybullet/components/panda/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : panda_robot.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/23/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import itertools
import contextlib
import numpy as np
import pybullet as p
from typing import Any, Optional, Union, Iterable, Tuple, Dict

import jacinle
from jacinle.logging import get_logger

from concepts.algorithm.configuration_space import BoxConfigurationSpace, CollisionFreeProblemSpace
from concepts.algorithm.rrt.rrt import birrt
from concepts.utils.range import Range
from concepts.utils.interpolation_utils import gen_cubic_spline, get_next_target_cubic_spline
from concepts.utils.interpolation_utils import gen_linear_spline, get_next_target_linear_spline

from concepts.simulator.pybullet.client import BulletClient
from concepts.simulator.pybullet.world import BodyFullStateSaver, ConstraintInfo
from concepts.simulator.pybullet.components.robot_base import Robot, GripperObjectIndices, RobotActionPrimitive
from concepts.simulator.pybullet.rotation_utils import rotate_vector, quat_mul, quat_conjugate

logger = get_logger(__file__)

np.set_printoptions(precision=5, suppress=True)

__all__ = ['PandaRobot']


class PandaRobot(Robot):
    # PANDA_FILE = 'assets://franka_panda/panda.urdf'
    # PANDA_JOINT_HOMES = np.array([0.98, 0.458, 0.31, -2.24, -0.30, 2.66, 2.32])
    PANDA_FILE = 'assets://robots/franka_description/robots/panda_arm_hand.urdf'
    PANDA_JOINT_HOMES = np.array([-0.45105, -0.38886, 0.45533, -2.19163, 0.13169, 1.81720, 0.51563])

    PANDA_GRIPPER_HOME = 0.02
    PANDA_GRIPPER_OPEN = 0.04
    PANDA_GRIPPER_CLOSE = 0.00

    def __init__(
        self,
        client: BulletClient,
        body_name='panda',
        gripper_objects: Optional[GripperObjectIndices] = None,
        use_magic_gripper: bool = True,
        pos: Optional[np.ndarray] = None,
        # pos: Optional[np.ndarray] = (0., 0.172, 0.),
        quat: Optional[np.ndarray] = None,
        # quat: Optional[np.ndarray] = (0, 0, -0.707, 0.707),
        # quat: Optional[np.ndarray] = (0, 0, -0.707, 0.707),
        # quat: Optional[np.ndarray] = (0, 0, 0.707, 0.707),
        # quat: Optional[np.ndarray] = (0, 0, 1, 0),
        ik_fast: bool = False,
    ):
        super().__init__(client, body_name, gripper_objects)

        self.panda = client.load_urdf(type(self).PANDA_FILE, pos=pos, quat=quat, static=True, body_name=body_name, group=None)

        all_joints = client.w.get_joint_info_by_body(self.panda)
        self.panda_joints = [joint.joint_index for joint in all_joints if joint.joint_type == p.JOINT_REVOLUTE]
        self.gripper_joints = [joint.joint_index for joint in all_joints if joint.joint_type == p.JOINT_PRISMATIC]
        assert set(self.gripper_joints) == {9, 10}
        self.panda_ee_tip = 11
        self.ik_fast = ik_fast
        self.ik_fast_wrapper = None

        # for joint_index in self.panda_joints:
        #     self.client.p.changeDynamics(self.panda, joint_index, linearDamping=0.0, angularDamping=0.0)

        # Create a constraint to keep the fingers centered
        # https://github.com/bulletphysics/bullet3/blob/daadfacfff365852ffc96f373c834216a25b11e5/examples/pybullet/gym/pybullet_robots/panda/panda_sim_grasp.py#L46-L54

        c = self.client.p.createConstraint(
            self.panda, 9, self.panda, 10,
            jointType=p.JOINT_GEAR, jointAxis=[1, 0, 0],
            parentFramePosition=[0, 0, 0], childFramePosition=[0, 0, 0]
        )
        self.client.p.changeConstraint(c, gearRatio=-1, erp=0.1, maxForce=50)

        joint_info = [self.client.w.get_joint_info_by_id(self.panda, i) for i in self.panda_joints]
        self.panda_joints_lower = np.array([joint.joint_lower_limit for joint in joint_info])
        self.panda_joints_upper = np.array([joint.joint_upper_limit for joint in joint_info])

        # State variable: gripper_activated. True if the gripper is activated.
        # For all control actions, make sure to call self._set_gripper_control before steping the simulation.
        self.griper_activated = None

        self.use_magic_gripper = use_magic_gripper
        self.gripper_constraint = None

        self.reset_home_qpos()
        self.ee_home_pos, self.ee_home_quat = self.get_ee_pose()

        self.reach_two_stage = PandaReachTwoStage(self)
        self.grasp = PandaGrasp(self)
        self.planar_push = PandaPlanarPush(self)
        self.push_two_stage = PandaPushTwoStage(self)
        self.pick_and_place = PandaPickAndPlace(self)

        self._cspace = None
        self._cfree_pspace = None

        self.register_action('move_pose', self.move_pose)
        self.register_action('move_home', self.move_home)
        self.register_action('move_home_cfree', self.move_home_cfree)

        self.register_action('grasp', self.grasp)
        self.register_action('open_gripper_free', self.open_gripper_free)
        self.register_action('close_gripper_free', self.close_gripper_free)

        self.register_action('reach_two_stage', self.reach_two_stage)
        self.register_action('planar_push', self.planar_push)
        self.register_action('push_two_stage', self.push_two_stage)
        self.register_action('pick_and_place', self.pick_and_place)
        self.register_action('reach_two_stage', self.reach_two_stage, interface='franka')
        self.register_action('planar_push', self.planar_push, interface='franka')
        self.register_action('pick_and_place', self.pick_and_place, interface='franka')

    def get_robot_body_id(self) -> int:
        return self.panda

    def get_robot_ee_link_id(self) -> int:
        return self.panda_ee_tip

    def get_qpos(self) -> np.ndarray:
        return self.client.w.get_batched_qpos_by_id(self.panda, self.panda_joints)

    def set_qpos(self, qpos: np.ndarray) -> None:
        self.client.w.set_batched_qpos_by_id(self.panda, self.panda_joints, qpos)

    def set_qpos_with_holding(self, qpos: np.ndarray) -> None:
        self.set_qpos(qpos)
        if self.gripper_constraint is not None:
            # TODO(Jiayuan Mao @ 2023/03/02): should consider the actual link that the gripper is attached to, and self-collision, etc.
            constraint = self.p.getConstraintInfo(self.gripper_constraint)
            other_body_id = constraint[2]
            parent_frame_pos = constraint[6]
            parent_frame_orn = constraint[8]

            body_pose = self.p.getLinkState(self.panda, self.w.link_names[f'{self.body_name}/panda_leftfinger'].link_id)
            obj_to_body = (parent_frame_pos, parent_frame_orn)
            obj_pose = p.multiplyTransforms(body_pose[0], body_pose[1], *obj_to_body)

            # print('other_body_id:', other_body_id, 'pose:', self.world.get_body_state_by_id(other_body_id).get_transformation())

            self.w.set_body_state2_by_id(other_body_id, obj_pose[0], obj_pose[1])

    def get_home_qpos(self):
        return type(self).PANDA_JOINT_HOMES

    def reset_home_qpos(self):
        self.client.w.set_batched_qpos_by_id(self.panda, self.panda_joints, type(self).PANDA_JOINT_HOMES)
        self.client.w.set_batched_qpos_by_id(self.panda, self.gripper_joints, [type(self).PANDA_GRIPPER_HOME] * 2)
        self.gripper_activated = None

    def get_ee_pose(self, fk=False) -> Tuple[np.ndarray, np.ndarray]:
        state = self.client.w.get_link_state_by_id(self.panda, self.panda_ee_tip, fk=fk)
        return np.array(state.position), np.array(state.orientation)

    def set_ee_pose(self, pos: np.ndarray, quat: np.ndarray) -> bool:
        qpos = self.ik(pos, quat)
        if qpos is None:
            return False
        self.set_qpos(qpos)
        return True

    def get_ee_home_pos(self):
        return self.ee_home_pos

    def get_ee_home_quat(self):
        return self.ee_home_quat

    def get_gripper_state(self) -> Optional[bool]:
        return self.gripper_activated

    def fk(self, qpos: np.ndarray, link_name_or_id: Optional[Union[str, int]] = None) -> Tuple[np.ndarray, np.ndarray]:
        with BodyFullStateSaver(self.client.w, self.panda):
            self.client.w.set_batched_qpos_by_id(self.panda, self.panda_joints, qpos)
            if link_name_or_id is None:
                return self.get_ee_pose(fk=True)
            elif isinstance(link_name_or_id, str):
                state = self.client.w.get_link_state(link_name_or_id, fk=True)
                return np.array(state.position), np.array(state.orientation)
            elif isinstance(link_name_or_id, int):
                state = self.client.w.get_link_state_by_id(self.panda, link_name_or_id, fk=True)
                return np.array(state.position), np.array(state.orientation)
            else:
                raise TypeError(f'link_name_or_id must be str or int, got {type(link_name_or_id)}.')

    def ik_pybullet(self, pos: np.ndarray, quat: np.ndarray, force: bool = False, verbose: bool = False) -> Optional[np.ndarray]:
        joints = self.client.p.calculateInverseKinematics(
            bodyUniqueId=self.panda,
            endEffectorLinkIndex=self.panda_ee_tip,
            targetPosition=pos,
            targetOrientation=quat,
            lowerLimits=self.panda_joints_lower,
            upperLimits=self.panda_joints_upper,
            jointRanges=self.panda_joints_upper - self.panda_joints_lower,
            restPoses=type(self).PANDA_JOINT_HOMES,
            maxNumIterations=100,
            residualThreshold=1e-5,
        )

        if verbose:
            print('IK (solution, lower, upper):\n', np.stack([joints[:7], self.panda_joints_lower, self.panda_joints_upper], axis=0), sep='')

        if np.all(self.panda_joints_lower[:7] <= joints[:7]) and np.all(joints[:7] <= self.panda_joints_upper[:7]):
            qpos = np.array(joints[:7])
            fk_pos, fk_quat = self.fk(qpos)
            if np.allclose(fk_pos, pos, atol=1e-2) and np.abs(quat_mul(quat, quat_conjugate(fk_quat))[3]) > 0.99:
                return qpos

        if force:
            return np.array(joints[:7])

        return None

    def _init_ikfast(self):
        if self.ik_fast_wrapper is None:
            from concepts.simulator.pybullet.ikfast.ikfast_common import IKFastWrapper
            import concepts.simulator.pybullet.ikfast.franka_panda.ikfast_panda_arm as ikfast_module
            self.ik_fast_wrapper = IKFastWrapper(
                self.world, ikfast_module,
                body_id=self.panda,
                joint_ids=self.panda_joints,
                free_joint_ids=[self.panda_joints[-1]],
                use_xyzw=True
            )

    def ikfast(self, pos: np.ndarray, quat: np.ndarray, last_qpos: Optional[np.ndarray] = None, max_attempts: int = 1000, max_distance: float = float('inf'), error_on_fail: bool = True, verbose: bool = False) -> Optional[np.ndarray]:
        self._init_ikfast()

        if verbose:
            print('Solving IK for pos:', pos, 'quat:', quat)

        # NB(Jiayuan Mao @ 2023/08/09): relative transformation between the base and the end-effector.
        pos_delta = [0, 0, 0.1]
        # quat_delta = (0, 0, 1, 0)
        quat_delta = (0.0, 0.0, 0.9238795325108381, 0.38268343236617297)
        inner_quat = quat_mul(quat, quat_conjugate(quat_delta))
        inner_pos = np.array(pos) - rotate_vector(pos_delta, quat)

        body_state = self.world.get_body_state_by_id(self.panda)
        inner_pos = inner_pos - np.array(body_state.position)
        if not np.allclose(body_state.orientation, [0, 0, 0, 1]):
            # TODO(Jiayuan Mao @ 2022/12/27): solve when orientation != identity.
            raise NotImplementedError('IKFast does not support non-identity orientation for the robot yet.')

        try:
            ik_solution = list(itertools.islice(self.ik_fast_wrapper.gen_ik(inner_pos, inner_quat, last_qpos=last_qpos, max_attempts=max_attempts, max_distance=max_distance, verbose=False), 1))[0]
        except IndexError:
            if error_on_fail:
                raise
            return None

        if verbose:
            print('IK (solution, lower, upper):\n', np.stack([ik_solution, self.panda_joints_lower, self.panda_joints_upper], axis=0), sep='')
            print('FK:', self.fk(ik_solution))

        # fk_pos, fk_quat = self.fk(ik_solution, 'panda/tool_link')
        # link8_pos, link8_quat = self.fk(ik_solution, 'panda/panda_link8')
        # print('query (outer):', pos, quat, 'solution:', ik_solution, 'fk', fk_pos, fk_quat, 'fk_diff', np.linalg.norm(fk_pos - pos), quat_mul(quat_conjugate(fk_quat), quat)[3])
        # print('query (outer):', 'link8 IN', link8_pos, link8_quat)
        # print('query (outer):', 'linkT IN', fk_pos, fk_quat)
        # print('query (outer):', 'linkT OT', link8_pos + rotate_vector((0, 0, 0.1), link8_quat), quat_mul(link8_quat, quat_delta))
        # print(self.w.get_joint_info_by_id(self.panda, 11))
        # print(self.w.get_joint_state_by_id(self.panda, 11))

        return ik_solution

    def get_configuration_space(self) -> BoxConfigurationSpace:
        if self._cspace is None:
            ranges = list()
            for lower, upper in zip(self.panda_joints_lower, self.panda_joints_upper):
                ranges.append(Range(lower, upper))
            self._cspace = BoxConfigurationSpace(ranges, 0.02)

        return self._cspace

    def is_colliding(self, q: Optional[np.ndarray] = None, return_contacts: bool = False):
        if q is not None:
            self.set_qpos_with_holding(q)
        contacts = self.world.get_contact(self.panda, update=True)

        ignore_bodies = [self.panda]
        if self.gripper_constraint is not None:
            constraint = self.p.getConstraintInfo(self.gripper_constraint)
            other_body_id = constraint[2]
            contacts += self.world.get_contact(other_body_id, update=True)
            ignore_bodies.append(other_body_id)
            # print('other_body_id:', other_body_id, 'pose:', self.world.get_body_state_by_id(other_body_id).get_transformation())

        filtered_contacts = list()
        for c in contacts:
            if c.body_b not in ignore_bodies:
                filtered_contacts.append(c)
                # jacinle.log_function.print('Detected collision between', c.body_a_name, 'and', c.body_b_name)
                if not return_contacts:
                    return True

        return False if not return_contacts else filtered_contacts

    def get_collision_free_problem_space(self) -> CollisionFreeProblemSpace:
        if self._cfree_pspace is None:
            cspace = self.get_configuration_space()

            def is_colliding(q):
                self.set_qpos(q)
                contacts = self.world.get_contact(self.panda, update=True)

                ignore_bodies = [self.panda]
                if self.gripper_constraint is not None:
                    # TODO(Jiayuan Mao @ 2023/03/02): should consider the actual link that the gripper is attached to, and self-collision, etc.
                    constraint = self.p.getConstraintInfo(self.gripper_constraint)
                    other_body_id = constraint[2]
                    parent_frame_pos = constraint[6]
                    parent_frame_orn = constraint[8]

                    body_pose = self.p.getLinkState(self.panda, self.w.link_names[f'{self.body_name}/panda_leftfinger'].link_id)
                    obj_to_body = (parent_frame_pos, parent_frame_orn)
                    obj_pose = p.multiplyTransforms(body_pose[0], body_pose[1], *obj_to_body)

                    # print('other_body_id:', other_body_id, 'pose:', self.world.get_body_state_by_id(other_body_id).get_transformation())

                    self.w.set_body_state2_by_id(other_body_id, obj_pose[0], obj_pose[1])

                    contacts += self.world.get_contact(other_body_id, update=True)
                    ignore_bodies.append(other_body_id)

                    # print('other_body_id:', other_body_id, 'pose:', self.world.get_body_state_by_id(other_body_id).get_transformation())

                for c in contacts:
                    if c.body_b not in ignore_bodies:
                        # jacinle.log_function.print('Detected collision between', c.body_a_name, 'and', c.body_b_name)
                        return True
                return False

            self._cfree_pspace = CollisionFreeProblemSpace(cspace, is_colliding)

        return self._cfree_pspace

    def rrt_collision_free(self, pos1: np.ndarray, pos2: Optional[np.ndarray] = None, smooth_fine_path: bool = False, disable_renderer: bool = True, **kwargs):
        if pos2 is None:
            pos2 = pos1
            pos1 = self.get_qpos()

        cfree_pspace = self.get_collision_free_problem_space()
        with self.world.save_world_v2(), jacinle.cond_with(self.client.disable_rendering(suppress_stdout=False), disable_renderer):
            path = birrt(cfree_pspace, pos1, pos2, smooth_fine_path=smooth_fine_path, **kwargs)
            if path[0] is not None:
                return True, path[0]
            return False, None

    def move_qpos(self, target_qpos, speed=0.01, timeout=10, gains=1, local_smoothing: bool = True) -> bool:
        last_qpos = None
        for _ in self.client.timeout(timeout):
            current_qpos = self.client.w.get_batched_qpos_by_id(self.panda, self.panda_joints)
            diff = target_qpos - current_qpos
            norm = np.linalg.norm(diff, ord=2)

            rel_diff, rel_norm = None, 1e9
            if last_qpos is not None:
                rel_diff = last_qpos - current_qpos
                rel_norm = np.linalg.norm(rel_diff, ord=1)
            else:
                last_qpos = current_qpos
                rel_norm = 1e9

            # pos, quat = self.get_ee_pose()
            # print('  current pos', pos, 'quat', quat)
            # print('  current/target joint state\n', np.stack([current_qpos, target_qpos], axis=0), sep='')
            # print('  diff', diff, 'norm', norm)
            # if last_qpos is not None:
            #     print('  rel_diff', rel_diff, 'rel_norm', rel_norm)

            if norm < 0.01:
                return True

            # if rel_norm < 0.001:
            #     return False

            if local_smoothing:
                # If speed is bigger than norm, we clip it to norm.
                step_qpos = current_qpos + diff / norm * min(speed, norm)
            else:
                step_qpos = target_qpos

            vector_gains = np.ones(len(self.panda_joints)) * gains
            self.client.p.setJointMotorControlArray(
                bodyIndex=self.panda,
                jointIndices=self.panda_joints,
                controlMode=p.POSITION_CONTROL,
                targetPositions=step_qpos,
                positionGains=vector_gains,
            )
            self._set_gripper_control()
            self.client.step()

        if not self.warnings_suppressed:
            logger.warning(f'{self.body_name}: Moving timeout ({timeout}s).')
        return False

    def move_qpos_trajectory_v2(
        self, qpos_trajectory: Iterable[np.ndarray],
        step_size: float = 1, gains: float = 0.3,
        atol: float = 0.03, timeout: float = 20,
        verbose: bool = False,
        return_world_states: bool = False,
    ):
        qpos_trajectory = np.array(qpos_trajectory)
        qpos_trajectory_dedup = list()
        last_qpos = None
        for qpos in qpos_trajectory:
            if qpos is None:
                continue
            if last_qpos is None or np.linalg.norm(qpos - last_qpos, ord=2) > 0.01:
                qpos_trajectory_dedup.append(qpos)
                last_qpos = qpos
        qpos_trajectory = np.array(qpos_trajectory_dedup)

        # spl = gen_cubic_spline(qpos_trajectory)
        spl = gen_linear_spline(qpos_trajectory)
        prev_qpos = None
        prev_qpos_not_moving = 0
        next_id = None

        world_states = []
        for _ in self.client.timeout(timeout):
            current_qpos = self.client.w.get_batched_qpos_by_id(self.panda, self.panda_joints)
            # next_target = get_next_target_cubic_spline(spl, current_qpos, step_size, qpos_trajectory)
            next_id, next_target = get_next_target_linear_spline(
                spl, current_qpos, step_size,
                minimum_x=next_id - step_size + 0.2 if next_id is not None else None
            )

            if verbose:
                print('this step size', np.linalg.norm(next_target - current_qpos, ord=1))
                print('next_id', next_id)
                print('this_target (lower, current, next, upper)\n', np.stack([
                    self.panda_joints_lower,
                    current_qpos,
                    next_target,
                    self.panda_joints_upper,
                ]), sep='')

            last_norm = np.linalg.norm(qpos_trajectory[-1] - current_qpos, ord=1)

            if verbose:
                print('last_norm', last_norm)
            if prev_qpos is not None:
                last_moving_dist = np.linalg.norm(prev_qpos - current_qpos, ord=1)
                if last_moving_dist < 0.001:
                    prev_qpos_not_moving += 1
                else:
                    prev_qpos_not_moving = 0
                if prev_qpos_not_moving > 10:
                    if last_norm < atol * 10:
                        return True
                    else:
                        if not self.warnings_suppressed:
                            logger.warning(f'{self.body_name}: No progress for 10 steps.')
                        return False
            prev_qpos = current_qpos

            if last_norm < atol:
                if return_world_states:
                    return world_states
                return True

            vector_gains = np.ones(len(self.panda_joints)) * gains
            self.client.p.setJointMotorControlArray(
                bodyIndex=self.panda,
                jointIndices=self.panda_joints,
                controlMode=p.POSITION_CONTROL,
                targetPositions=next_target,
                positionGains=vector_gains,
            )
            self._set_gripper_control()
            self.client.step()
            world_states.append(self.world.save_world())

        if not self.warnings_suppressed:
            logger.warning(f'{self.body_name}: Moving timeout ({timeout}s).')
        return False

    def move_pose_smooth(self, pos: np.ndarray, quat: np.ndarray, speed: float = 0.01, max_qpos_distance: float = 1.0) -> bool:
        if self.ik_fast:
            # In ik_fast mode, we need to specify the `max_qpos_distance` to avoid the IK solution to be too far away.
            target_qpos = self.ik(pos, quat, max_distance=max_qpos_distance)
        else:
            target_qpos = self.ik(pos, quat)

        rv = self.move_qpos(target_qpos, speed=speed)
        return rv

    def move_home_cfree(self, speed: float = 0.01, timeout: float = 10.0):
        with self.client.disable_rendering(suppress_stdout=False):
            try:
                success, qpath = self.rrt_collision_free(self.get_home_qpos())
            except ValueError:
                success = False

        if not success:
            if not self.warnings_suppressed:
                logger.warning('Cannot find a collision-free path to home. Doing a direct move.')
            return self.move_qpos(self.get_home_qpos(), speed=speed, timeout=timeout)
        return self.move_qpos_trajectory(qpath, speed=speed, timeout=timeout)

    def internal_set_gripper_state(self, activate: bool, constraint_info: Optional[ConstraintInfo] = None) -> None:
        if not activate:
            if self.use_magic_gripper:
                if self.gripper_constraint is not None:
                    self.client.p.removeConstraint(self.gripper_constraint)
                    self.gripper_constraint = None
        else:
            if self.use_magic_gripper:
                if constraint_info is not None:
                    self.create_gripper_constraint(constraint_info.child_body)

        self.gripper_activated = activate

    def open_gripper_free(self, timeout: float = 2, force: bool = False) -> bool:
        return self._change_gripper_state_free(False, timeout=timeout, force=force)

    def close_gripper_free(self, timeout: float = 2, force: bool = False) -> bool:
        return self._change_gripper_state_free(True, timeout=timeout, force=force)

    @contextlib.contextmanager
    def set_next_pand_open(self, open_target: float):
        # TODO(Jiayuan Mao @ 2023/04/21): find a better way to do this.
        type(self).PANDA_GRIPPER_OPEN = open_target
        yield
        type(self).PANDA_GRIPPER_OPEN = 0.04

    def _change_gripper_state_free(self, activate: bool, timeout: float = 2.0, force: bool = False, verbose: bool = False) -> bool:
        """A helper function that change the gripper state assuming no contact with other objects.

        Args:
            activate (bool): True to activate the gripper, False to release it.

        Returns:
            bool: True if the gripper state changed, False otherwise.
        """

        if self.gripper_activated is not None and activate == self.gripper_activated and not force:
            return True

        if verbose:
            logger.info(f'{self.body_name}: Change gripper state to {activate}, use_magic_gripper={self.use_magic_gripper}, gripper_constraint={self.gripper_constraint}')

        if not activate:
            if self.use_magic_gripper:
                if self.gripper_constraint is not None:
                    self.client.p.removeConstraint(self.gripper_constraint)
                    self.gripper_constraint = None

        self.gripper_activated = activate
        target = type(self).PANDA_GRIPPER_CLOSE if activate else type(self).PANDA_GRIPPER_OPEN
        target_qpos = np.array([target, target])

        for _ in self.client.timeout(timeout):
            current_qpos = self.client.w.get_batched_qpos_by_id(self.panda, self.gripper_joints)
            diff = target_qpos - current_qpos

            if all(np.abs(diff) < 1e-3):
                return False

            self._set_gripper_control()
            self.client.step()

        if not self.warnings_suppressed:
            logger.warning(f'{self.body_name}: Moving gripper timeout ({timeout}s).')
        return True

    def _set_gripper_control(self, target: Optional[float] = None):
        if target is None:
            target = type(self).PANDA_GRIPPER_CLOSE if self.gripper_activated else type(self).PANDA_GRIPPER_OPEN

        for joint_index in self.gripper_joints:
            self.client.p.setJointMotorControl2(
                self.panda, joint_index,
                controlMode=p.POSITION_CONTROL,
                targetPosition=target,
                force=40
            )

    def create_gripper_constraint(self, object_id, verbose: bool = False):
        if self.gripper_constraint is not None:
            self.client.p.removeConstraint(self.gripper_constraint)
            self.gripper_constraint = None

        if verbose:
            logger.info(f'{self.body_name}: Creating constraint between {self.panda} and {object_id}.')

        panda_leftfinger = self.world.link_names[f'{self.body_name}/panda_leftfinger'].link_id

        body_pose = self.p.getLinkState(self.panda, panda_leftfinger)  # the base link.
        obj_pose = self.p.getBasePositionAndOrientation(object_id)
        world_to_body = p.invertTransform(body_pose[0], body_pose[1])
        obj_to_body = p.multiplyTransforms(world_to_body[0], world_to_body[1], obj_pose[0], obj_pose[1])
        self.gripper_constraint = self.p.createConstraint(
            parentBodyUniqueId=self.panda,
            parentLinkIndex=panda_leftfinger,
            childBodyUniqueId=object_id,
            childLinkIndex=-1,  # should be the link index of the contact link.
            jointType=p.JOINT_FIXED,
            jointAxis=(0, 0, 0),
            parentFramePosition=obj_to_body[0],
            parentFrameOrientation=obj_to_body[1],
            childFramePosition=(0, 0, 0),
            childFrameOrientation=(0, 0, 0)
        )


class PandaGrasp(RobotActionPrimitive):
    robot: PandaRobot

    def __init__(self, robot: PandaRobot):
        super().__init__(robot)
        self.panda_leftfinger = self.w.link_names[f'{self.robot.body_name}/panda_leftfinger'].link_id
        self.panda_rightfinger = self.w.link_names[f'{self.robot.body_name}/panda_rightfinger'].link_id

    def __call__(self, timeout=2, target_object: Optional[int] = None, force_constraint: bool = False, verbose: bool = False) -> bool:
        target_qpos = np.array([type(self.robot).PANDA_GRIPPER_CLOSE] * 2)
        self.robot.gripper_activated = True

        for _ in self.client.timeout(timeout):
            current_qpos = self.w.get_batched_qpos_by_id(self.robot.panda, self.robot.gripper_joints)
            diff = target_qpos - current_qpos

            if self.detect_contact(verbose=verbose):
                return True

            if all(np.abs(diff) < 1e-3):
                return False

            for joint_index, target in zip(self.robot.gripper_joints, target_qpos):
                self.p.setJointMotorControl2(
                    self.robot.panda, joint_index,
                    controlMode=p.POSITION_CONTROL,
                    targetPosition=target,
                    force=40
                )
            # self.robot._set_gripper_control()
            self.client.step()

        if force_constraint:
            if target_object is not None and self.robot.use_magic_gripper and self.robot.gripper_constraint is None:
                self.robot.create_gripper_constraint(target_object)

        if not self.warnings_suppressed:
            logger.warning(f'{self.robot.body_name}: Grasping timeout ({timeout}s).')
        return False

    def detect_contact(self, verbose: bool = False):
        points1 = self.robot.client.p.getContactPoints(bodyA=self.robot.panda, linkIndexA=self.panda_leftfinger)
        points2 = self.robot.client.p.getContactPoints(bodyA=self.robot.panda, linkIndexA=self.panda_rightfinger)

        objects1 = {point[2] for point in points1 if point[2] != self.robot.panda}
        objects2 = {point[2] for point in points2 if point[2] != self.robot.panda}
        objects_inter = list(set.intersection(objects1, objects2))

        if len(objects_inter) > 0:
            if verbose:
                for obj_id in objects_inter:
                    logger.info(f'{self.robot.body_name}: Contact detected: {obj_id}, name = {self.w.body_names[obj_id]}.')

            if self.robot.use_magic_gripper:
                obj_id = objects_inter[0]
                self.robot.create_gripper_constraint(obj_id)

            return True
        return False


class PandaReachTwoStage(RobotActionPrimitive):
    robot: PandaRobot

    def __call__(self, pos: np.ndarray, quat: np.ndarray, height: float = 0.2, speed: float = 0.01) -> bool:
        logger.info(f'{self.robot.body_name}: Moving to pose (two-stage): {pos}, {quat}.')

        self.robot.move_pose(pos + rotate_vector([0, 0, -height], quat), quat, speed=speed)
        for i in reversed(range(10)):
            self.robot.move_pose(pos + rotate_vector([0, 0, -height / 10 * i], quat), quat, speed=speed)


class PandaPushTwoStage(RobotActionPrimitive):
    robot: PandaRobot

    def __call__(self, start_pos: np.ndarray, normal: np.ndarray, distance: float, timeout: float = 10, prepush_height: float = 0.1, speed: float = 0.01) -> bool:
        self.robot.close_gripper_free()

        pos, quat = self.robot.get_ee_pose()
        self.robot.move_pose(start_pos - prepush_height * normal / np.linalg.norm(normal), quat, speed=speed)

        pos, quat = self.robot.get_ee_pose()
        end_pos = start_pos + normal * distance

        for _ in self.client.timeout(timeout):
            diff = end_pos - pos
            if np.linalg.norm(diff) < 0.01:
                return True
            self.robot.move_pose_smooth(pos + diff / np.linalg.norm(diff) * speed, quat, speed=speed)
            pos, _ = self.robot.get_ee_pose()

        if not self.warnings_suppressed:
            logger.info(f'{self.robot.body_name}: Pushing timeout ({timeout}s).')
        return False


class PandaPlanarPush(RobotActionPrimitive):
    robot: PandaRobot

    def __call__(self, target_pos: np.ndarray, timeout: float = 10, speed: float = 0.01, tol: float = 0.01) -> bool:
        target_pos = np.array(target_pos)
        pos, quat = self.robot.get_ee_pose()
        for _ in self.client.timeout(timeout):
            diff = target_pos - pos

            if np.linalg.norm(diff) < tol:
                return True

            self.robot.move_pose_smooth(pos + diff / np.linalg.norm(diff) * speed, quat, speed=speed)
            pos, _ = self.robot.get_ee_pose()

        if not self.warnings_suppressed:
                logger.warning(f'{self.robot.body_name}: Planar push timeout ({timeout}s).')
        return False


class PandaPickAndPlace(RobotActionPrimitive):
    """A helpful pick-and-place primitive."""

    robot: PandaRobot

    def __call__(
        self,
        pick_pos: np.ndarray, pick_quat: np.ndarray,
        place_pos: np.ndarray, place_quat: np.ndarray,
        reach_kwargs: Optional[Dict[str, Any]] = None, grasp_kwargs: Optional[Dict[str, Any]] = None
    ) -> bool:
        pick_pos = np.array(pick_pos)
        pick_quat = np.array(pick_quat)
        pick_after_pos = pick_pos + np.array([0, 0, 0.2])
        place_pos = np.array(place_pos)
        place_quat = np.array(place_quat)
        place_after_pos = place_pos + np.array([0, 0, 0.2])

        if reach_kwargs is None:
            reach_kwargs = {}
        if grasp_kwargs is None:
            grasp_kwargs = {}

        self.robot.do('open_gripper_free')
        self.robot.do('reach_two_stage', pick_pos, pick_quat, **reach_kwargs)
        self.robot.do('grasp', **grasp_kwargs)
        self.robot.do('move_pose', pick_after_pos, pick_quat)
        self.robot.do('move_pose', place_pos, place_quat)
        self.robot.do('open_gripper_free')
        self.robot.do('move_pose', place_after_pos, place_quat)


concepts/simulator/pybullet/components/panda/panda_robot.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : robot_base.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/24/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""This file contains basic classes for defining robots and their action primitives."""

import contextlib

import numpy as np
from typing import Optional, Union, Iterable, Tuple, List, Dict, Callable

from jacinle.logging import get_logger
from jacinle.utils.enum import JacEnum
from concepts.simulator.pybullet.client import BulletClient
from concepts.simulator.pybullet.components.component_base import BulletComponent

logger = get_logger(__file__)

__all__ = ['Gripper', 'GripperObjectIndices', 'IKMethod', 'Robot', 'RobotActionPrimitive']


class Gripper(BulletComponent):
    """Base gripper class."""

    def __init__(self, client: BulletClient):
        super().__init__(client)
        self.activated = False

    def step(self):
        """This function can be used to create gripper-specific behaviors."""
        return

    def activate(self, objects):
        return

    def release(self):
        return


GripperObjectIndices = Dict[str, List[int]]


class IKMethod(JacEnum):
    PYBULLET = 'pybullet'
    IKFAST = 'ikfast'


class Robot(BulletComponent):
    def __init__(self, client: BulletClient, body_name: Optional[str] = None, gripper_objects: Optional[GripperObjectIndices] = None, current_interface='pybullet', ik_method: Union[str, IKMethod] = 'pybullet'):
        super().__init__(client)
        self.body_name = body_name if body_name is not None else self.__class__.__name__
        self.gripper_objects = gripper_objects
        if self.gripper_objects is None:
            self.gripper_objects = self.client.w.body_groups

        self.current_interface = current_interface
        self.primitive_actions: Dict[Tuple[str, str], Callable] = dict()
        self.ik_method = IKMethod.from_string(ik_method)
        self.warnings_suppressed = False

    @contextlib.contextmanager
    def suppress_warnings(self):
        self.warnings_suppressed = True
        yield
        self.warnings_suppressed = False

    def register_action(self, name: str, func: Callable, interface: Optional[str] = None) -> None:
        """Register an action primitive.

        Args:
            name (str): Name of the action primitive.
            func (Callable): Function that implements the action primitive.
            interface (Optional[str], optional): Interface of the action primitive. Defaults to None.
                If None, the action primitive is registered for the current interface.

        Raises:
            ValueError: If the action primitive is already registered.
        """
        if interface is None:
            interface = self.current_interface
        if (name, interface) in self.primitive_actions:
            raise ValueError(f'Action primitive {name} for interface {interface} already registered.')
        self.primitive_actions[(interface, name)] = func

    def do(self, action_name: str, *args, **kwargs) -> bool:
        """Execute an action primitive.

        Args:
            action_name (str): Name of the action primitive.

        Returns:
            bool: True if the action primitive is successful.
        """
        assert (self.current_interface, action_name) in self.primitive_actions, f'Action primitive {action_name} for interface {self.current_interface} not registered.'
        return self.primitive_actions[(self.current_interface, action_name)](*args, **kwargs)

    def set_ik_method(self, ik_method: Union[str, IKMethod]):
        self.ik_method = IKMethod.from_string(ik_method)

    def get_robot_body_id(self) -> int:
        """Get the pybullet body ID of the robot.

        Returns:
            int: Body ID of the robot.
        """

        raise NotImplementedError()

    def get_robot_ee_link_id(self) -> int:
        """Get the pybullet link ID of the robot end effector.

        Returns:
            int: Link ID of the robot end effector.
        """
        raise NotImplementedError()

    def get_qpos(self) -> np.ndarray:
        """Get the current joint configuration.

        Returns:
            np.ndarray: Current joint configuration.
        """
        raise NotImplementedError()

    def set_qpos(self, qpos: np.ndarray) -> None:
        """Set the joint configuration.

        Args:
            qpos (np.ndarray): Joint configuration.
        """
        raise NotImplementedError()

    def set_qpos_with_holding(self, qpos: np.ndarray) -> None:
        """Set the joint configuration with the gripper holding the object.

        Args:
            qpos (np.ndarray): Joint configuration.
        """
        return self.set_qpos(qpos)

    def get_home_qpos(self) -> np.ndarray:
        """Get the home joint configuration."""
        raise NotImplementedError()

    def reset_home_qpos(self):
        """Reset the home joint configuration."""
        raise NotImplementedError()

    def get_ee_pose(self) -> Tuple[np.ndarray, np.ndarray]:
        """Get the pose of the end effector.

        Returns:
            Tuple[np.ndarray, np.ndarray]: 3D position and 4D quaternion of the end effector.
        """
        raise NotImplementedError()

    def get_ee_home_pos(self) -> np.ndarray:
        """Get the home position of the end effector."""
        raise NotImplementedError()

    def get_ee_home_quat(self) -> np.ndarray:
        """Get the home orientation of the end effector."""
        raise NotImplementedError()

    def get_gripper_state(self) -> Optional[bool]:
        """Get the gripper state.

        Returns:
            Optional[bool]: True if the gripper is activated.
        """
        raise NotImplementedError()

    def fk(self, qpos: np.ndarray, link_name_or_id: Optional[Union[str, int]] = None) -> Tuple[np.ndarray, np.ndarray]:
        """Forward kinematics.

        Args:
            qpos (np.ndarray): Joint configuration.
            link_name_or_id (Union[str, int], optional): name or id of the link. If not specified, the pose of the end effector is returned.

        Returns:
            Tuple[np.ndarray, np.ndarray]: 3D position and 4D quaternion of the end effector.
        """
        raise NotImplementedError()

    def ik(self, pos: np.ndarray, quat: np.ndarray, force: bool = False, max_distance: float = float('inf'), max_attempts: int = 1000, verbose: bool = False) -> Optional[np.ndarray]:
        """Inverse kinematics.

        Args:
            pos (np.ndarray): 3D position of the end effector.
            quat (np.ndarray): 4D quaternion of the end effector.
            force (bool, optional): Whether to force the IK solver to return a solution. Defaults to False.
                If set, the IK solve may return a solution even if the end effector is not at the given pose.
                This function is useful for debugging and for moving towards a certain direction.
            max_distance (float, optional): Maximum distance between the last qpos and the solution (Only used for IKFast). Defaults to float('inf').
            max_attempts (int, optional): Maximum number of attempts (only used for IKFast). Defaults to 1000.
            verbose (bool, optional): Whether to print debug information.

        Returns:
            np.ndarray: Joint configuration.
        """

        if self.ik_method is IKMethod.PYBULLET:
            assert max_distance == float('inf'), 'max_distance is not supported in PyBullet IK'
            return self.ik_pybullet(pos, quat, force=force, verbose=verbose)
        elif self.ik_method is IKMethod.IKFAST:
            assert force is False, 'force is not supported in ik_fast'
            return self.ikfast(pos, quat, max_distance=max_distance, max_attempts=max_attempts, verbose=verbose)

    def ik_pybullet(self, pos: np.ndarray, quat: np.ndarray, force: bool = False, verbose: bool = False) -> Optional[np.ndarray]:
        """Inverse kinematics using pybullet.

        Args:
            pos (np.ndarray): 3D position of the end effector.
            quat (np.ndarray): 4D quaternion of the end effector.
            force (bool, optional): Whether to force the IK solver to return a solution. Defaults to False.
                If set, the IK solve may return a solution even if the end effector is not at the given pose.
                This function is useful for debugging and for moving towards a certain direction.
            verbose (bool, optional): Whether to print debug information.

        Returns:
            np.ndarray: Joint configuration.
        """
        raise NotImplementedError()

    def ikfast(self, pos: np.ndarray, quat: np.ndarray, last_qpos: Optional[np.ndarray] = None, max_attempts: int = 1000, max_distance: float = float('inf'), error_on_fail: bool = True, verbose: bool = False) -> Optional[np.ndarray]:
        """Inverse kinematics using IKFast.

        Args:
            pos (np.ndarray): 3D position of the end effector.
            quat (np.ndarray): 4D quaternion of the end effector.
            last_qpos (Optional[np.ndarray], optional): Last joint configuration. Defaults to None.
                If None, the current joint configuration is used.
            max_attempts (int, optional): Maximum number of IKFast attempts. Defaults to 1000.
            max_distance (float, optional): Maximum distance between the target pose and the end effector. Defaults to float('inf').
            error_on_fail (bool, optional): Whether to raise an error if the IKFast solver fails. Defaults to True.
            verbose (bool, optional): Whether to print debug information.

        Returns:
            np.ndarray: Joint configuration.
        """
        raise NotImplementedError(f'IKFast is not supported for this {self.__class__.__name__}.')

    def move_qpos(self, target_qpos: np.ndarray, speed: float = 0.01, timeout: float = 10.0, local_smoothing: bool = True) -> bool:
        """Move the robot to the given joint configuration.

        Args:
            target_qpos (np.ndarray): Target joint configuration.
            speed (float, optional): Speed of the movement. Defaults to 0.01.
            timeout (float, optional): Timeout of the movement. Defaults to 10.

        Returns:
            bool: True if the movement is successful.
        """
        raise NotImplementedError()

    def move_home(self, timeout: float = 10.0) -> bool:
        """Move the robot to the home configuration."""
        return self.move_qpos(self.get_home_qpos(), timeout=timeout)

    def move_pose(self, pos, quat, speed=0.01, force: bool = False, verbose: bool = False) -> bool:
        """Move the end effector to the given pose.

        Args:
            pos (np.ndarray): 3D position of the end effector.
            quat (np.ndarray): 4D quaternion of the end effector.
            speed (float, optional): Speed of the movement. Defaults to 0.01.
            force (bool, optional): Whether to force the IK solver to return a solution. Defaults to False.
                If set, the IK solve may return a solution even if the end effector is not at the specified pose.

        Returns:
            bool: True if the movement is successful.
        """

        if verbose:
            logger.info(f'{self.body_name}: Moving to pose: {pos}, {quat}.')
        target_qpos = self.ik(pos, quat, force=force)
        rv = self.move_qpos(target_qpos, speed)
        return rv

    def move_qpos_trajectory(self, qpos_trajectory: Iterable[np.ndarray], speed: float = 0.01, timeout: float = 10, first_timeout: float = None, local_smoothing: bool = True, verbose: bool = False) -> bool:
        """Move the robot along the given joint configuration trajectory.

        Args:
            qpos_trajectory (Iterable[np.ndarray]): Joint configuration trajectory.
            speed (float, optional): Speed of the movement. Defaults to 0.01.
            timeout (float, optional): Timeout of the movement. Defaults to 10.
            first_timeout (float, optional): Timeout of the first movement. Defaults to None (same as timeout).

        Returns:
            bool: True if the movement is successful.
        """

        this_timeout = timeout if first_timeout is None else first_timeout
        for target_qpos in qpos_trajectory:
            if verbose:
                logger.info(f'{self.body_name}: Moving to qpos: {target_qpos}.')
            rv = self.move_qpos(target_qpos, speed, this_timeout, local_smoothing=local_smoothing)
            this_timeout = timeout
            # if verbose:
            #     logger.info(f'{self.body_name}: Moved to qpos:  {self.get_qpos()}.')
        if not rv:
            return False
        return True

    def replay_qpos_trajectory(self, qpos_trajectory: Iterable[np.ndarray], verbose: bool = True):
        """Replay a joint confifguartion trajectory. This function is useful for debugging a generated joint trajectory.

        Args:
            qpos_trajectory: joint configuration trajectory.
            verbose: whether to print debug information.
        """
        for i, target_qpos in enumerate(qpos_trajectory):
            if verbose:
                logger.info(f'{self.body_name}: Moving to qpos: {target_qpos}. Trajectory replay {i + 1}/{len(qpos_trajectory)}.')
            self.set_qpos_with_holding(target_qpos)
            self.client.wait_for_duration(0.1)
            # input('Press enter to continue.')


class RobotActionPrimitive(object):
    def __init__(self, robot: Robot):
        self.robot = robot

    def __call__(self, *args, **kwargs) -> bool:
        raise NotImplementedError()

    @property
    def client(self):
        return self.robot.client

    @property
    def p(self):
        return self.robot.client.p

    @property
    def world(self):
        return self.robot.client.w

    @property
    def w(self):
        return self.robot.client.w

    @property
    def warnings_suppressed(self):
        return self.robot.warnings_suppressed
concepts/simulator/pybullet/components/robot_base.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/23/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.



concepts/simulator/pybullet/components/ur5/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : ur5_gripper.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/08/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""
Classes to handle gripper dynamics.
This file is directly copied from https://github.com/cliport/cliport/blob/master/cliport/tasks/grippers.py
"""

import os.path as osp
import numpy as np
import pybullet as p

from jacinle.utils.enum import JacEnum
from concepts.simulator.pybullet.client import BulletClient
from concepts.simulator.pybullet.components.robot_base import Gripper, GripperObjectIndices
from concepts.simulator.pybullet.rotation_utils import get_ee_to_tool, solve_tool_from_ee

SPATULA_BASE_URDF = 'robots/ur5/spatula/spatula-base.urdf'
SUCTION_BASE_URDF = 'robots/ur5/suction/suction-base.urdf'
SUCTION_HEAD_URDF = 'robots/ur5/suction/suction-head.urdf'


class UR5GripperType(JacEnum):
    NONE = 'none'
    SPATULA = 'spatula'
    SUCTION = 'suction'

    def get_class(self):
        if self is UR5GripperType.NONE:
            return None
        elif self is UR5GripperType.SPATULA:
            return Spatula
        elif self is UR5GripperType.SUCTION:
            return Suction
        else:
            raise ValueError(f'Unknown gripper type: {self}.')


class Spatula(Gripper):
    """Simulate simple spatula for pushing."""

    def __init__(self, client: BulletClient, robot: int, ee: int, obj_ids: GripperObjectIndices):  # pylint: disable=unused-argument
        """Creates spatula and 'attaches' it to the robot."""
        super().__init__(client)

        # Load spatula model.
        pose = ((0.487, 0.109, 0.438), p.getQuaternionFromEuler((np.pi, 0, 0)))
        base = self.client.load_urdf(osp.join(client.assets_root, SPATULA_BASE_URDF), pose[0], pose[1])
        self.p.createConstraint(
            parentBodyUniqueId=robot,
            parentLinkIndex=ee,
            childBodyUniqueId=base,
            childLinkIndex=-1,
            jointType=p.JOINT_FIXED,
            jointAxis=(0, 0, 0),
            parentFramePosition=(0, 0, 0),
            childFramePosition=(0, 0, 0.01)
        )


class Suction(Gripper):
    """Simulate simple suction dynamics."""

    def __init__(self, client, robot, ee, obj_ids):
        """Creates suction and 'attaches' it to the robot.

        Has special cases when dealing with rigid vs deformables. For rigid,
        only need to check contact_constraint for any constraint. For soft
        bodies (i.e., cloth or bags), use cloth_threshold to check distances
        from gripper body (self.head_id) to any vertex in the cloth mesh. We
        need correct code logic to handle gripping potentially a rigid or a
        deformable (and similarly for releasing).

        To be clear on terminology: 'deformable' here should be interpreted
        as a PyBullet 'softBody', which includes cloths and bags. There's
        also cables, but those are formed by connecting rigid body beads, so
        they can use standard 'rigid body' grasping code.

        To get the suction gripper pose, use p.getLinkState(self.head_id, 0),
        and not p.getBasePositionAndOrientation(self.head_id) as the latter is
        about z=0.03m higher and empirically seems worse.

        Args:
          client: the BulletClient instance.
          robot: int representing PyBullet ID of robot.
          ee: int representing PyBullet ID of end effector link.
          obj_ids: list of PyBullet IDs of all suctionable objects in the env.
        """
        super().__init__(client)

        self.robot_id = robot
        self.ee_id = ee

        # Load suction gripper base model (visual only).
        pose = ((0.487, 0.109, 0.438), p.getQuaternionFromEuler((np.pi, 0, 0)))
        self.base_id = self.client.load_urdf(osp.join(client.assets_root, SUCTION_BASE_URDF), pose[0], pose[1])
        self.suction_base_constraint = self.p.createConstraint(
            parentBodyUniqueId=robot,
            parentLinkIndex=ee,
            childBodyUniqueId=self.base_id,
            childLinkIndex=-1,
            jointType=p.JOINT_FIXED,
            jointAxis=(0, 0, 0),
            parentFramePosition=(0, 0, 0),
            childFramePosition=(0, 0, 0.01)
        )

        # Load suction tip model (visual and collision) with compliance.
        # urdf = 'assets/ur5/suction/suction-head.urdf'
        pose = ((0.487, 0.109, 0.347), p.getQuaternionFromEuler((np.pi, 0, 0)))
        self.head_id = self.client.load_urdf(osp.join(client.assets_root, SUCTION_HEAD_URDF), pose[0], pose[1])
        self.suction_head_constraint = self.p.createConstraint(
            parentBodyUniqueId=robot,
            parentLinkIndex=ee,
            childBodyUniqueId=self.head_id,
            childLinkIndex=-1,
            jointType=p.JOINT_FIXED,
            jointAxis=(0, 0, 0),
            parentFramePosition=(0, 0, 0),
            childFramePosition=(0, 0, -0.08)
        )
        self.p.changeConstraint(self.suction_head_constraint, maxForce=50)

        # Reference to object IDs in environment for simulating suction.
        self.obj_ids = obj_ids

        # Indicates whether gripper is gripping anything (rigid or def).
        self.activated = False

        # For gripping and releasing rigid objects.
        self.contact_constraint = None
        self.contact_ee_to_obj = None

        # Defaults for deformable parameters, and can override in tasks.
        self.def_ignore = 0.035  # TODO(daniel) check if this is needed
        self.def_threshold = 0.030
        self.def_nb_anchors = 1

        # Track which deformable is being gripped (if any), and anchors.
        self.def_grip_item = None
        self.def_grip_anchors = []

        # Determines release when gripped deformable touches a rigid/def.
        # TODO(daniel) should check if the code uses this -- not sure?
        self.def_min_vetex = None
        self.def_min_distance = None

        # Determines release when a gripped rigid touches defs (e.g. cloth-cover).
        self.init_grip_distance = None
        self.init_grip_item = None

    def activate(self):
        """Simulate suction using a rigid fixed constraint to contacted object."""
        # TODO(andyzeng): check deformables logic.
        # del def_ids

        if not self.activated:
            points = self.p.getContactPoints(bodyA=self.head_id, linkIndexA=0)
            if points:
                # Handle contact between suction with a rigid object.
                for point in points:
                    obj_id, contact_link = point[2], point[4]
                if obj_id in self.obj_ids['rigid']:
                    body_pose = self.p.getLinkState(self.head_id, 0)  # the base link.
                    obj_pose = self.p.getBasePositionAndOrientation(obj_id)
                    world_to_body = p.invertTransform(body_pose[0], body_pose[1])
                    obj_to_body = p.multiplyTransforms(world_to_body[0], world_to_body[1], obj_pose[0], obj_pose[1])
                    self.contact_constraint = self.p.createConstraint(
                        parentBodyUniqueId=self.head_id,
                        parentLinkIndex=0,
                        childBodyUniqueId=obj_id,
                        childLinkIndex=contact_link,
                        jointType=p.JOINT_FIXED,
                        jointAxis=(0, 0, 0),
                        parentFramePosition=obj_to_body[0],
                        parentFrameOrientation=obj_to_body[1],
                        childFramePosition=(0, 0, 0),
                        childFrameOrientation=(0, 0, 0)
                    )
                    self.contact_ee_to_object = get_ee_to_tool(self.client, self.head_id, 0, obj_id)

                self.activated = True

    def sync_gripper_qpos(self):
        """Sync gripper qpos with PyBullet simulation."""

        # recompute the pose for the suction body
        link_state = self.client.world.get_link_state_by_id(self.robot_id, self.ee_id)
        pos, quat = link_state.pos, link_state.orientation

        # update the suction base
        self.world.set_body_state2_by_id(self.base_id, pos + np.array((0, 0, 0.01)), quat)

        # update the suction head
        self.world.set_body_state2_by_id(self.head_id, pos + np.array((0, 0, -0.08)), quat)

        # update the object being gripped
        if self.contact_constraint is not None:
            constraint = self.world.get_constraint(self.contact_constraint)
            pos, quat = self.world.get_link_state_by_id(self.head_id, 0).get_transformation()
            new_pos, new_quat = solve_tool_from_ee(pos, quat, self.contact_ee_to_object)
            self.world.set_body_state2_by_id(constraint.child_body, new_pos, new_quat)
        self.world.update_contact()

    def release(self):
        """Release gripper object, only applied if gripper is 'activated'.

        If suction off, detect contact between gripper and objects.
        If suction on, detect contact between picked object and other objects.

        To handle deformables, simply remove constraints (i.e., anchors).
        Also reset any relevant variables, e.g., if releasing a rigid, we
        should reset init_grip values back to None, which will be re-assigned
        in any subsequent grasps.
        """
        if self.activated:
            self.activated = False
            print('Releasing gripper.')

            # Release gripped rigid object (if any).
            if self.contact_constraint is not None:
                print('Releasing gripped rigid object.')
                try:
                    self.p.removeConstraint(self.contact_constraint)
                    self.contact_constraint = None
                except:  # noqa: E772
                    print('Failed to release gripped rigid object.')
                    pass
                else:
                    print('Released gripped rigid object.')
                self.init_grip_distance = None
                self.init_grip_item = None
            self.contact_ee_to_object = None

            # Release gripped deformable object (if any).
            if self.def_grip_anchors:
                for anchor_id in self.def_grip_anchors:
                    self.p.removeConstraint(anchor_id)
                self.def_grip_anchors = []
                self.def_grip_item = None
                self.def_min_vetex = None
                self.def_min_distance = None

    def detect_contact(self):
        """Detects a contact with a rigid object."""
        body, link = self.head_id, 0
        if self.activated and self.contact_constraint is not None:
            try:
                info = self.p.getConstraintInfo(self.contact_constraint)
                body, link = info[2], info[3]
            except:  # noqa: E772
                self.contact_constraint = None
                pass

        # Get all contact points between the suction and a rigid body.
        points = self.p.getContactPoints(bodyA=body, linkIndexA=link)
        if self.activated:
            points = [point for point in points if point[2] != self.head_id]
        # print(body, link, [self.client.w.body_names[point[2]] for point in points])

        # We know if len(points) > 0, contact is made with SOME rigid item.
        if points:
            return True

        return False

    def check_grasp(self):
        """Check a grasp (object in contact?) for picking success."""

        suctioned_object = None
        if self.contact_constraint is not None:
            suctioned_object = self.p.getConstraintInfo(self.contact_constraint)[2]
        return suctioned_object is not None


concepts/simulator/pybullet/components/ur5/ur5_gripper.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : ur5_robot.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/08/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import time
import contextlib
import numpy as np
import pybullet as p
import jacinle

from typing import Optional, Tuple, Union
from concepts.simulator.pybullet.client import BulletClient
from concepts.simulator.pybullet.components.robot_base import Robot, Gripper, GripperObjectIndices, RobotActionPrimitive
from concepts.simulator.pybullet.components.ur5.ur5_gripper import UR5GripperType
from concepts.simulator.pybullet.rotation_utils import compose_transformation

__all__ = ['UR5Robot', 'UR5ReachAndPick', 'UR5ReachAndPlace', 'UR5PlanarMove']

logger = jacinle.get_logger(__file__)

g_debug_options = jacinle.FileOptions(__file__, show_timeout_warning=True)


class UR5Robot(Robot):
    UR5_FILE = 'assets://robots/ur5/ur5.urdf'
    UR5_JOINT_HOMES = np.array([-1, -0.5, 0.5, -0.5, -0.5, 0]) * np.pi

    def __init__(
        self,
        client: BulletClient,
        body_name='ur5',
        gripper_type: str = 'suction',
        gripper_objects: Optional[GripperObjectIndices] = None,
        pos: Optional[np.ndarray] = None,
        quat: Optional[np.ndarray] = None,
    ):
        super().__init__(client, body_name, gripper_objects)

        self.gripper_type = UR5GripperType.from_string(gripper_type)

        self.ur5 = client.load_urdf(type(self).UR5_FILE, pos=pos, quat=quat, static=True, body_name=body_name, group=None)

        ur5_joints = client.w.get_joint_info_by_body(self.ur5)
        self.ur5_joints = [joint.joint_index for joint in ur5_joints if joint.joint_type == p.JOINT_REVOLUTE]
        self.ur5_ee_tip = 10

        gripper_cls = self.gripper_type.get_class()
        self.gripper: Optional[Gripper] = None
        if gripper_cls is not None:
            self.gripper = gripper_cls(client, self.ur5, 9, self.gripper_objects)

        self.reset_home_qpos()
        self.ee_home_pos, self.ee_home_quat = self.get_ee_pose()

        self.reach_and_pick = UR5ReachAndPick(self)
        self.reach_and_place = UR5ReachAndPlace(self)
        self.planar_move = UR5PlanarMove(self)

        self._ignore_physics = False

    @contextlib.contextmanager
    def ignore_physics(self, ignore_physics: bool = True):
        old_ignore_physics = self._ignore_physics
        self._ignore_physics = ignore_physics
        yield
        self._ignore_physics = old_ignore_physics

    def set_ignore_physics(self, ignore_physics: bool = True):
        self._ignore_physics = ignore_physics

    def get_robot_body_id(self) -> int:
        return self.ur5

    def get_qpos(self) -> np.ndarray:
        return np.array([self.client.p.getJointState(self.ur5, i)[0] for i in self.ur5_joints], dtype=np.float32)

    def set_qpos(self, qpos: np.ndarray) -> None:
        self.world.set_batched_qpos_by_id(self.ur5, self.ur5_joints, qpos)

    def get_home_qpos(self) -> np.ndarray:
        return type(self).UR5_JOINT_HOMES

    def reset_home_qpos(self):
        self.client.w.set_batched_qpos_by_id(self.ur5, self.ur5_joints, type(self).UR5_JOINT_HOMES)
        if self.gripper is not None:
            self.gripper.release()

    def get_ee_pose(self) -> Tuple[np.ndarray, np.ndarray]:
        state = self.world.get_link_state_by_id(self.ur5, self.ur5_ee_tip)
        return state.position, state.orientation

    def get_ee_home_pos(self) -> np.ndarray:
        return self.ee_home_pos

    def get_ee_home_quat(self) -> np.ndarray:
        return self.ee_home_quat

    def get_gripper_state(self) -> Optional[bool]:
        return self.gripper.activated if self.gripper is not None else None

    def fk(self, qpos: np.ndarray, link_name_or_id: Optional[Union[str, int]] = None) -> Tuple[np.ndarray, np.ndarray]:
        with self.world.save_body(self.ur5):
            self.set_qpos(qpos)
            state = self.world.get_link_state_by_id(self.ur5, self.ur5_ee_tip)
            return state.position, state.orientation

    def ik_pybullet(self, pos: np.ndarray, quat: np.ndarray, force: bool = False, verbose: bool = False) -> Optional[np.ndarray]:
        """Calculate joint configuration with inverse kinematics."""
        joints = self.client.p.calculateInverseKinematics(
            bodyUniqueId=self.ur5,
            endEffectorLinkIndex=self.ur5_ee_tip,
            targetPosition=pos,
            targetOrientation=quat,
            lowerLimits=[-3 * np.pi / 2, -2.3562, -17, -17, -17, -17],
            upperLimits=[-np.pi / 2, 0, 17, 17, 17, 17],
            jointRanges=[np.pi, 2.3562, 34, 34, 34, 34],  # * 6,
            restPoses=type(self).UR5_JOINT_HOMES.tolist(),
            maxNumIterations=100,
            residualThreshold=1e-5,
        )
        joints = np.float32(joints)
        joints[2:] = (joints[2:] + np.pi) % (2 * np.pi) - np.pi
        return joints

    def ikfast(
        self,
        pos: np.ndarray,
        quat: np.ndarray,
        last_qpos: Optional[np.ndarray] = None,
        max_attempts: int = 1000,
        max_distance: float = float('inf'),
        error_on_fail: bool = True,
        verbose: bool = False
    ) -> Optional[np.ndarray]:
        raise NotImplementedError('IKFast is not implemented for UR5.')

    def move_qpos(self, targj, speed=0.01, timeout=10):
        """Move UR5 to target joint configuration."""
        for _ in jacinle.timeout(timeout):
            currj = [self.client.p.getJointState(self.ur5, i)[0] for i in self.ur5_joints]
            currj = np.asarray(currj)
            diffj = targj - currj

            # pos, quat = self.client.w.get_link_state_by_id(self.ur5, self.ur5_ee_tip)

            if all(np.abs(diffj) < 5e-2):
                return True

            # Move with constant velocity
            norm = np.linalg.norm(diffj)
            v = diffj / norm if norm > 0 else 0
            stepj = currj + v * min(speed, norm)

            if self._ignore_physics:
                self.set_qpos(stepj)
                self.gripper.sync_gripper_qpos()
                time.sleep(1 / self.client.fps)
            else:
                gains = np.ones(len(self.ur5_joints))
                self.client.p.setJointMotorControlArray(
                    bodyIndex=self.ur5,
                    jointIndices=self.ur5_joints,
                    controlMode=p.POSITION_CONTROL,
                    targetPositions=stepj,
                    positionGains=gains
                )
                self.client.step()

        if g_debug_options.show_timeout_warning:
            logger.warning(f'{self.body_name}: move qpos exceeded {timeout} second timeout.')
        return False

    def activate_gripper(self):
        if self.gripper is not None:
            return self.gripper.activate()

    def release_gripper(self):
        if self.gripper is not None:
            return self.gripper.release()


class UR5ReachAndPick(RobotActionPrimitive):
    robot: UR5Robot

    def __init__(self, robot: UR5Robot, height=0.32, speed=0.01):
        super().__init__(robot)
        self.height = height
        self.speed = speed

    def __call__(self, pos: np.ndarray, quat: np.ndarray, timeout: float = 10, speed: Optional[float] = None) -> bool:
        """Execute reach and pick primitive.

        Returns:
            success: if the pick was successful.
        """
        if speed is None:
            speed = self.speed

        # Execute picking primitive.
        prepick_pos, prepick_quat = compose_transformation(pos, quat, (0, 0, self.height), (0, 0, 0, 1))
        postpick_pos, postpick_quat = compose_transformation(pos, quat, (0, 0, self.height), (0, 0, 0, 1))

        with self.robot.client.with_fps(30):
            succ = self.robot.move_pose(prepick_pos, prepick_quat, speed=self.speed)
        if not succ:
            return False

        pos_delta = np.array([0, 0, -0.001], dtype=np.float32)
        target_pos = prepick_pos
        for _ in jacinle.timeout(timeout, self.client.fps):
            if self.robot.gripper.detect_contact():
                break
            target_pos = target_pos + pos_delta
            succ = self.robot.move_pose(target_pos, prepick_quat, speed=self.speed)
            if not succ:
                return False

        self.robot.activate_gripper()

        succ = self.robot.move_pose(postpick_pos, postpick_quat, self.speed)
        pick_success = self.robot.gripper.check_grasp()
        if not succ or not pick_success:
            return False

        return True


class UR5ReachAndPlace(RobotActionPrimitive):
    robot: UR5Robot

    def __init__(self, robot: UR5Robot, height=0.32, speed=0.01):
        super().__init__(robot)
        self.height = height
        self.speed = speed

    def __call__(self, pos: np.ndarray, quat: np.ndarray, release: bool = True, timeout: float = 10, speed: Optional[float] = None) -> bool:
        """Execute reach and place primitive.

        Returns:
            success: if the place was successful.
        """
        if speed is None:
            speed = self.speed

        self.robot.activate_gripper()

        preplace_pos, preplace_quat = compose_transformation(pos, quat, (0, 0, self.height), (0, 0, 0, 1))
        postplace_pos, postplace_quat = compose_transformation(pos, quat, (0, 0, self.height), (0, 0, 0, 1))

        target_pos = preplace_pos
        pos_delta = np.array([0, 0, -0.001], dtype=np.float32)

        for _ in jacinle.timeout(timeout, self.client.fps):
            if self.robot.gripper.detect_contact():
                break
            target_pos = target_pos + pos_delta
            succ = self.robot.move_pose(target_pos, preplace_quat, self.speed)
            if not succ:
                return False

        succ = True
        if release:
            self.robot.release_gripper()
            succ = self.robot.move_pose(postplace_pos, postplace_quat, self.speed)

        return succ


class UR5PlanarMove(RobotActionPrimitive):
    robot: UR5Robot

    def __init__(self, robot: UR5Robot, speed=0.01):
        super().__init__(robot)
        self.speed = speed

    def __call__(self, pos_xy: np.ndarray, speed=None, timeout: float = 10) -> bool:
        """Execute planar move primitive.

        Returns:
            success: if the move was successful.
        """

        if speed is None:
            speed = self.speed

        robot = self.robot
        link_state = self.robot.client.w.get_link_state_by_id(robot.ur5, robot.ur5_ee_tip)

        target_pos = np.array([pos_xy[0], pos_xy[1], link_state.position[2]], dtype=np.float32)
        target_quat = link_state.orientation

        for _ in jacinle.timeout(timeout, self.client.fps):
            current_link_state = self.robot.client.w.get_link_state_by_id(robot.ur5, robot.ur5_ee_tip)
            current_pos = current_link_state.position

            diff_pos = target_pos - current_pos
            if all(np.abs(diff_pos) < 3e-2):
                return True

            norm = np.linalg.norm(diff_pos)
            v = diff_pos / norm if norm > 0 else 0
            step_pos = current_pos + v * speed
            succ = self.robot.move_pose(step_pos, target_quat, self.speed)
            if not succ:
                return False

        return True

concepts/simulator/pybullet/components/ur5/ur5_robot.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : cspace.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/17/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
import contextlib
from typing import List

from concepts.simulator.pybullet.client import BulletClient

__all__ = ['BulletConfigurationSpace']


class BulletConfigurationSpace(object):
    def __init__(self, client: BulletClient, cspace_names: List[str]):
        self.client = client
        self.cspace_names = cspace_names

    @property
    def world(self):
        return self.client.world

    def get_qpos(self):
        return np.array([self.world.get_qpos(name) for name in self.cspace_names])

    def set_qpos(self, qpos, forward=False):
        for name, v in zip(self.cspace_names, qpos):
            self.world.set_qpos(name, v)
        if forward:
            self.client.p.stepSimulation()

    @contextlib.contextmanager
    def qpos_context(self, forward=False):
        backup = self.get_qpos()
        yield
        self.set_qpos(backup, forward=forward)

    def get_xpos(self, name, type=None):
        return self.world.get_xpos(name, type=type)

    def get_xquat(self, name, type=None):
        return self.world.get_xquat(name, type=type)

    def get_xmat(self, name, type=None):
        return self.world.get_xmat(name, type=type)

    def fk(self, end_effector_name, qpos):
        with self.qpos_context():
            self.set_qpos(qpos)
            return self.get_xpos(end_effector_name), self.get_xmat(end_effector_name)

    def ik(self, end_effector_name, goal_xpos, goal_xquat=None):
        raise NotImplementedError()

    def get_collisions(self, qpos):
        raise NotImplementedError()


concepts/simulator/pybullet/cspace.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 07/10/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/simulator/pybullet/ikfast/__init__.py
concepts-ai/Concepts
import fnmatch
import importlib
import os
import shutil

from distutils.core import setup
from distutils.dir_util import copy_tree
from distutils.extension import Extension

# Build C++ extension by running: 'python setup.py build'
# see: https://docs.python.org/3/extending/building.html

# http://openrave.org/docs/0.8.2/openravepy/ikfast/
# https://github.com/rdiankov/openrave/blob/master/python/ikfast.py#L92
# http://wiki.ros.org/Industrial/Tutorials/Create_a_Fast_IK_Solution

# Yijiang
# https://github.com/yijiangh/ikfast_pybind
# https://github.com/yijiangh/conrob_pybullet/tree/master/utils/ikfast
# https://github.com/yijiangh/choreo/blob/bc777069b8eb7283c74af26e5461532aec3d9e8a/framefab_robot/abb/framefab_irb6600/framefab_irb6600_support/doc/ikfast_tutorial.rst


def compile_ikfast(module_name, cpp_filename, remove_build=False):
    ikfast_module = Extension(module_name, sources=[cpp_filename])
    setup(
        name=module_name,
        version='1.0',
        description="ikfast module {}".format(module_name),
        ext_modules=[ikfast_module],
        script_args=['build'],
    )

    build_lib_path = None
    for root, dirnames, filenames in os.walk(os.getcwd()):
        if fnmatch.fnmatch(root, os.path.join(os.getcwd(), "*build", "lib*")):
            build_lib_path = root
            break

    assert build_lib_path is not None

    copy_tree(build_lib_path, os.getcwd())
    if remove_build:
        # TODO: error when compiling multiple arms for python2
        # error: unable to open output file 'build/temp.macosx-10.15-x86_64-2.7/movo_right_arm_ik.o': 'No such file or directory'
        shutil.rmtree(os.path.join(os.getcwd(), 'build'))

    try:
        importlib.import_module(module_name)
        print('\nikfast module {} imported successful'.format(module_name))
    except ImportError as e:
        print('\nikfast module {} imported failed'.format(module_name))
        raise e

    return True


concepts/simulator/pybullet/ikfast/compile.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/24/2022
#
# This file is part of HACL-PyTorch.
# Distributed under terms of the MIT license.

from jacinle.jit.cext import auto_travis

auto_travis(__file__)

concepts/simulator/pybullet/ikfast/franka_panda/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : setup.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/24/2022
#
# This file is part of HACL-PyTorch.
# Distributed under terms of the MIT license.

import sys
from concepts.simulator.pybullet.ikfast.compile import compile_ikfast


def main():
    robot_name = 'panda_arm'
    compile_ikfast(
        module_name='ikfast_{}'.format(robot_name),
        cpp_filename='ikfast_{}.cpp'.format(robot_name),
        remove_build=True
    )


if __name__ == '__main__':
    main()


concepts/simulator/pybullet/ikfast/franka_panda/setup.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : ikfast_common.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/26/2022
#
# This file is part of HACL-PyTorch.
# Distributed under terms of the MIT license.

import itertools
from typing import Optional, Iterable, Tuple, List
from collections import namedtuple

import random
import numpy as np
import numpy.random as npr

from jacinle.logging import get_logger

from concepts.utils.rotationlib import mat2quat, quat2mat

from concepts.simulator.pybullet.world import BulletWorld
from concepts.simulator.pybullet.rotation_utils import quat_conjugate, quat_mul

logger = get_logger(__file__)


class IKFastWrapper(object):
    def __init__(
        self,
        world: BulletWorld, module,
        body_id, joint_ids: List[int], free_joint_ids: List[int],
        use_xyzw: bool = True,  # PyBullet uses xyzw.
        max_attempts: int = 1000
    ):
        self.world = world
        self.module = module

        self.body_id = body_id
        self.joint_ids = joint_ids
        self.free_joint_ids = free_joint_ids

        self.use_xyzw = use_xyzw
        self.max_attempts = max_attempts

        joint_info = [self.world.get_joint_info_by_id(self.body_id, joint_id) for joint_id in joint_ids]
        self.joints_lower = np.array([info.joint_lower_limit for info in joint_info])
        self.joints_upper = np.array([info.joint_upper_limit for info in joint_info])
        self.free_joints_lower = np.array([info.joint_lower_limit for info in joint_info if info.joint_index in free_joint_ids])
        self.free_joints_upper = np.array([info.joint_upper_limit for info in joint_info if info.joint_index in free_joint_ids])

        assert len(self.free_joint_ids) + 6 == len(self.joint_ids)

    def fk(self, qpos: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        pos, mat = self.module.get_fk(qpos)
        quat = mat2quat(mat)
        if self.use_xyzw:
            return pos, quat[[1, 2, 3, 0]]
        return pos, quat

    def ik_internal(self, pos: np.ndarray, quat: np.ndarray, sampled: Optional[np.ndarray] = None) -> List[np.ndarray]:
        if self.use_xyzw:
            quat = quat[[3, 0, 1, 2]]
        mat = quat2mat(quat)
        if sampled is None:
            solutions = self.module.get_ik(mat.tolist(), pos.tolist())
        else:
            solutions = self.module.get_ik(mat.tolist(), pos.tolist(), list(sampled))

        if solutions is None:
            return list()
        return [np.array(solution) for solution in solutions]

    def get_current_joint_positions(self) -> np.ndarray:
        return np.array([self.world.get_joint_state_by_id(self.body_id, joint_id).position for joint_id in self.joint_ids])

    def get_current_free_joint_positions(self) -> np.ndarray:
        return np.array([self.world.get_joint_state_by_id(self.body_id, joint_id).position for joint_id in self.free_joint_ids])

    def gen_ik(self, pos: np.ndarray, quat: np.ndarray, last_qpos: Optional[np.ndarray], max_attempts: Optional[int] = None, max_distance: float = float('inf'), verbose: bool = False) -> Iterable[np.ndarray]:
        if last_qpos is None:
            current_joint_positions = self.get_current_joint_positions()
            current_free_joint_positions = self.get_current_free_joint_positions()
        else:
            current_joint_positions = last_qpos
            current_free_joint_positions = [last_qpos[i] for i, joint_id in enumerate(self.joint_ids) if joint_id in self.free_joint_ids]

        generator = itertools.chain([current_free_joint_positions], gen_uniform_sample_joints(self.free_joints_lower, self.free_joints_upper))
        if max_attempts is not None:
            generator = itertools.islice(generator, max_attempts)
        else:
            generator = itertools.islice(generator, self.max_attempts)

        succeeded = False
        for sampled in generator:
            solutions = self.ik_internal(pos, quat, sampled)
            random.shuffle(solutions)
            for solution in solutions:
                if check_joint_limits(solution, self.joints_lower, self.joints_upper):
                    if distance_fn(solution, current_joint_positions) < max_distance:
                        succeeded = True

                        # fk_pos, fk_quat = self.fk(solution.tolist())
                        # print('query (inside): ', pos, quat, 'solution: ', solution, 'fk', fk_pos, fk_quat, 'fk_diff', np.linalg.norm(fk_pos - pos), quat_mul(quat_conjugate(fk_quat), quat)[3])

                        yield solution
                    elif verbose:
                        print(f'IK solution is too far from current joint positions: {solution} vs {current_joint_positions}')

        if not succeeded and max_attempts is None:
            logger.warning(f'Failed to find IK solution for {pos} {quat} after {self.max_attempts} attempts.')


def check_joint_limits(qpos: np.ndarray, lower_limits: np.ndarray, upper_limits: np.ndarray) -> bool:
    return np.all(np.logical_and(qpos >= lower_limits, qpos <= upper_limits))


def uniform_sample_joints(lower_limits: np.ndarray, upper_limits: np.ndarray) -> np.ndarray:
    return np.array([npr.uniform(lower, upper) for lower, upper in zip(lower_limits, upper_limits)])


def gen_uniform_sample_joints(lower_limits: np.ndarray, upper_limits: np.ndarray) -> Iterable[np.ndarray]:
    while True:
        yield uniform_sample_joints(lower_limits, upper_limits)


def random_select_solution(solutions: List[np.ndarray]) -> np.ndarray:
    return random.choice(solutions)


def distance_fn(qpos1: np.ndarray, qpos2: np.ndarray) -> float:
    return np.linalg.norm(np.array(qpos1) - np.array(qpos2), ord=2)


def closest_select_solution(solutions: List[np.ndarray], current_qpos: np.ndarray) -> np.ndarray:
    return min(solutions, key=lambda qpos: distance_fn(qpos, current_qpos))


concepts/simulator/pybullet/ikfast/ikfast_common.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.



concepts/simulator/pybullet/manipulation_utils/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : contact_samplers.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/05/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import itertools
import random
from dataclasses import dataclass
from typing import Optional, Iterable, Sequence, Tuple, List

import numpy as np
import numpy.random as npr
import open3d as o3d

import jacinle
from jacinle.logging import get_logger
from concepts.simulator.pybullet.rotation_utils import get_quaternion_from_axes, find_orthogonal_vector, quat_mul, rotate_vector, enumerate_quaternion_from_vectors
from concepts.simulator.pybullet.components.robot_base import Robot

logger = get_logger(__file__)

__all__ = [
    'GraspParameter', 'sample_grasp', 'gen_grasp_trajectory',
    'PushParameter', 'sample_push_with_support', 'gen_push_trajectory',
    'IndirectPushParameter', 'sample_indirect_push_with_support',
    'PlacementParameter',
    'qpos_trajectory_from_poses', 'qpos_trajectory_from_poses_pybullet',
    'collision_free_qpos', 'collision_free_pos', 'collision_free_qpos_trajectory',
    'mesh_line_intersect', 'get_single_contact_normal'
]


@dataclass
class GraspParameter(object):
    point1: np.ndarray
    normal1: np.ndarray
    point2: np.ndarray
    normal2: np.ndarray
    ee_pos: np.ndarray
    ee_quat: np.ndarray
    qpos: np.ndarray


def sample_grasp(
    robot: Robot, object_id: int,
    gripper_distance: float, max_intersection_distance: float = 10,
    mesh_filename: Optional[str] = None, mesh_scale: float = 1.0,
    verbose: bool = False, max_test_points_before_first: int = 250, max_test_points: int = 100000000, batch_size: int = 100,
    surface_pointing_tol: float = 0.9, min_point_distance: float = 0.0001
) -> Iterable[GraspParameter]:
    """Given the name of the object, sample a 6D grasp pose. Before calling this function, we should make sure that the gripper is open."""

    mesh = robot.world.get_mesh(object_id, zero_center=False, mesh_filename=mesh_filename, mesh_scale=mesh_scale)
    t_mesh = o3d.t.geometry.TriangleMesh.from_legacy(mesh)

    # print the x, y, z center of the mesh
    # print('mesh center', mesh.get_center())
    # bounding box
    # print('mesh bounding box', mesh.get_axis_aligned_bounding_box())

    found = False
    nr_test_points_before_first = 0

    for _ in range(int(max_test_points / batch_size)):
        # TODO(Jiayuan Mao @ 2023/01/02): accelerate the computation.
        pcd = mesh.sample_points_uniformly(batch_size, use_triangle_normal=True)
        # pcd = mesh.sample_points_poisson_disk(batch_size, use_triangle_normal=True)

        indices = list(range(len(pcd.points)))
        random.shuffle(indices)
        for i in indices:
            if not found:
                nr_test_points_before_first += 1

            point = np.asarray(pcd.points[i])
            normal = np.asarray(pcd.normals[i])

            if verbose:
                print('sample_grasp_v2_gen', 'point', point, 'normal', normal)

            point2 = point - normal * max_intersection_distance
            other_intersection = mesh_line_intersect(t_mesh, point2, normal)

            if verbose:
                print('  other_intersection', other_intersection)

            # if no intersection, try the next point.
            if other_intersection is None:
                if verbose:
                    print('  skip: no intersection')
                continue

            other_point, other_normal = other_intersection

            # if two intersection points are too close, try the next point.
            if np.linalg.norm(other_point - point) < min_point_distance:
                if verbose:
                    print('  skip: too close')
                continue

            # if the surface normals are too different, try the next point.
            if np.abs(np.dot(normal, other_normal)) < surface_pointing_tol:
                if verbose:
                    print('  skip: normal too different')
                continue

            grasp_center = (point + other_point) / 2
            grasp_distance = np.linalg.norm(point - other_point)
            grasp_normal = normal

            if grasp_distance > gripper_distance:
                if verbose:
                    print('  skip: too far')
                continue

            ee_d = grasp_normal
            # ee_u and ee_v are two vectors that are perpendicular to ee_d
            ee_u = find_orthogonal_vector(ee_d)
            ee_v = np.cross(ee_u, ee_d)

            # if verbose:
            #     print('  grasp_center', grasp_center, 'grasp_distance', grasp_distance)
            #     print('  grasp axes:\n', np.array([ee_d, ee_u, ee_v]))

            # enumerate four possible grasp orientations
            for ee_norm1 in [ee_u, ee_v, -ee_u, -ee_v]:
                ee_norm2 = np.cross(ee_d, ee_norm1)
                ee_quat = get_quaternion_from_axes(ee_norm2, ee_d, ee_norm1)

                qpos = robot.ikfast(grasp_center, ee_quat, max_attempts=100, error_on_fail=False)
                # qpos = ctx.robot._solve_ik(grasp_center, ee_quat, force=True)
                if qpos is not None:
                    rv = collision_free_qpos(robot, qpos, verbose=verbose)
                    # with robot.world.save_body(robot.get_robot_body_id()):
                    #     robot.set_qpos(qpos)
                    #     robot.client.wait_for_duration(0.1)
                    #     # visualize the mesh and point and other_point as two red spheres
                    #     o3d.visualization.draw_geometries([
                    #         mesh,
                    #         o3d.geometry.TriangleMesh.create_sphere(0.1).translate(point),
                    #         o3d.geometry.TriangleMesh.create_sphere(0.1).translate(other_point),
                    #     ])

                    #     input(f'testing grasp, cfree={rv}')

                    if rv:
                        found = True
                        yield GraspParameter(
                            point1=point, normal1=normal,
                            point2=other_point, normal2=other_normal,
                            ee_pos=grasp_center, ee_quat=ee_quat,
                            qpos=qpos
                        )
                    elif verbose:
                        print('    gripper pos', grasp_center)
                        print('    gripper quat', ee_quat)
                        print('    skip: collision')

        if not found and nr_test_points_before_first > max_test_points_before_first:
            if verbose:
                logger.warning(f'Failed to find a grasp after {nr_test_points_before_first} points tested.')
            return


def gen_grasp_trajectory(grasp_pos: np.ndarray, grasp_quat: np.ndarray, height: float = 0.1) -> List[Tuple[np.ndarray, np.ndarray]]:
    """Given the name of the object, the grasp pose, generate a trajectory of 6D poses that the robot should follow to grasp the object."""

    trajectory = [(grasp_pos + rotate_vector([0, 0, -height], grasp_quat), grasp_quat)]
    for i in reversed(range(10)):
        trajectory.append((
            grasp_pos + rotate_vector([0, 0, -height / 10 * i], grasp_quat),
            grasp_quat
        ))

    return trajectory


@dataclass
class PushParameter(object):
    push_pos: np.ndarray
    push_dir: np.ndarray
    distance: float
    ee_quat: np.ndarray


def sample_push_with_support(robot: Robot, object_id: int, support_object_id: int, max_attempts: int = 1000, batch_size: int = 100, verbose=False) -> Iterable[PushParameter]:
    mesh = robot.world.get_mesh(object_id, zero_center=False)

    nr_batches = int(max_attempts / batch_size)
    feasible_point_indices = list()
    for _ in range(nr_batches):
        pcd = mesh.sample_points_uniformly(batch_size, use_triangle_normal=True)

        # get the contact points between the object and the support object
        contact_normal = get_single_contact_normal(robot, object_id, support_object_id)

        # filter out the points that are not on the contact plane
        feasible_point_cond = np.abs(np.asarray(pcd.normals).dot(contact_normal)) < 0.02
        feasible_point_indices = np.where(feasible_point_cond)[0]

        # print(f'Found {len(feasible_point_indices)} feasible points.')

        if len(feasible_point_indices) == 0:
            continue

        # o3d.visualization.draw([pcd.select_by_index(feasible_point_indices), mesh])

        npr.shuffle(feasible_point_indices)
        rows = list()
        for index in feasible_point_indices:
            rows.append((index, pcd.points[index], -pcd.normals[index]))

        if verbose:
            import tabulate
            jacinle.log_function.print(tabulate.tabulate(rows, headers=['index', 'point', 'normal']))

        # create a new point cloud
        for index in feasible_point_indices:
            if verbose:
                jacinle.log_function.print('sample_push_with_support', 'point', pcd.points[index], 'normal', -pcd.normals[index])
            yield PushParameter(np.asarray(pcd.points[index]), -np.asarray(pcd.normals[index]), 0.1, robot.get_ee_home_quat())

    if len(feasible_point_indices) == 0:
        raise ValueError(f'No feasible points for {object_id} on {support_object_id} after {nr_batches * batch_size} attempts.')


def gen_push_trajectory(push_pos: np.ndarray, push_dir: np.ndarray, push_distance: float, quat: np.ndarray, prepush_distance: float = 0.1) -> List[Tuple[np.ndarray, np.ndarray]]:
    """Given the name of the object, the push pose, the push direction, and the push distance, generate a trajectory of 6D poses that the robot should follow to push the object."""

    push_pos += (0, 0, 0.010)
    real_start_pos = push_pos - prepush_distance * push_dir / np.linalg.norm(push_dir)
    real_end_pos = push_pos + push_dir * push_distance

    trajectory = list()
    for i in range(10 + 1):
        pos = real_start_pos + (real_end_pos - real_start_pos) * i / 10
        trajectory.append((pos, quat))

    return trajectory


@dataclass
class IndirectPushParameter(object):
    target_push_pos: np.ndarray
    target_push_dir: np.ndarray

    tool_pos: np.ndarray
    tool_quat: np.ndarray

    tool_point_pos: np.ndarray
    tool_point_normal: np.ndarray

    ee_pos: np.ndarray = None
    ee_quat: np.ndarray = None

    prepush_distance: float = 0.05
    push_distance: float = 0.1

    @property
    def total_push_distance(self):
        return self.prepush_distance + self.push_distance


def sample_indirect_push_with_support(
    robot: Robot, target_id: int, tool_id: int, support_object_id: int,
    prepush_distance: float = 0.05,
    max_attempts: int = 10000, batch_size: int = 100, filter_push_dir: Optional[np.ndarray] = None,
    push_distance_distribution: Sequence[float] = (0.1, 0.15), push_distance_sample: bool = False,
    contact_normal_tol: float = 0.01, verbose: bool = False
) -> Iterable[IndirectPushParameter]:
    """Sample a push of the target object using the tool object. The target object should be placed on the support object.

    Args:
        robot: the robot.
        target_id: the PyBullet body id of the target object.
        tool_id: the PyBullet body id of the tool object.
        support_object_id: the PyBullet body id of the support object.
        prepush_distance: the distance to pre-push the tool object before the actual push.
        max_attempts: the maximum number of attempts for sampling.
        batch_size: the number of samples in batch processing.
        filter_push_dir: if specified, the push direction will be filtered to be within 0.2 cosine distance from this direction.
        push_distance_distribution: the distribution of the push distance. If `push_distance_sample` is True, this will be used to sample the push distance.
        push_distance_sample: whether to sample the push distance from the distribution. If False, the function will enumerate the push distances in the distribution.
        contact_normal_tol: the tolerance for the contact normal. We will sample the push direction such that the contact normal is within `contact_normal_tol` degree from the gravity direction.
        verbose: whether to print the sampling information.
    """
    target_mesh = robot.world.get_mesh(target_id, zero_center=False)
    tool_mesh = robot.world.get_mesh(tool_id, zero_center=False)

    nr_batches = int(max_attempts / batch_size)
    contact_normal = get_single_contact_normal(robot, target_id, support_object_id)

    for _ in range(nr_batches):
        target_pcd = target_mesh.sample_points_uniformly(batch_size, use_triangle_normal=True)
        tool_pcd = tool_mesh.sample_points_uniformly(batch_size, use_triangle_normal=True)

        # feasible_target_point_cond = np.abs(np.asarray(target_pcd.normals).dot(contact_normal)) < 0.01
        # 0.1 for real demo.
        feasible_target_point_cond = np.abs(np.asarray(target_pcd.normals).dot(contact_normal)) < contact_normal_tol

        if filter_push_dir is not None:
            feasible_target_point_cond = np.logical_and(
                feasible_target_point_cond,
                np.asarray(target_pcd.normals, dtype=np.float32).dot(-filter_push_dir) > 0.8
            )

        feasible_target_point_indices = np.where(feasible_target_point_cond)[0]

        all_index_pairs = list(itertools.product(feasible_target_point_indices, range(batch_size)))
        random.shuffle(all_index_pairs)
        for target_index, tool_index in all_index_pairs:
            target_point_pos = np.asarray(target_pcd.points[target_index])
            target_point_normal = -np.asarray(target_pcd.normals[target_index])  # point inside

            tool_point_pos = np.asarray(tool_pcd.points[tool_index])
            tool_point_normal = np.asarray(tool_pcd.normals[tool_index])  # point outside (towards the tool)

            current_tool_pos, current_tool_quat = robot.world.get_body_state_by_id(tool_id).get_transformation()

            # Solve for a quaternion that aligns the tool normal with the target normal
            for rotation_quat in enumerate_quaternion_from_vectors(tool_point_normal, target_point_normal, 4):
                # This is the world coordinate for the tool point after rotation.
                new_tool_point_pos = current_tool_pos + rotate_vector(tool_point_pos - current_tool_pos, rotation_quat)
                # Now compute the displacement for the tool object
                final_tool_pos = target_point_pos - new_tool_point_pos + current_tool_pos
                final_tool_pos -= target_point_normal * prepush_distance
                final_tool_quat = quat_mul(rotation_quat, current_tool_quat)

                success = True
                with robot.world.save_body(tool_id):
                    robot.world.set_body_state2_by_id(tool_id, final_tool_pos, final_tool_quat)
                    contacts = robot.world.get_contact(tool_id, update=True)

                    for contact in contacts:
                        if contact.body_b != tool_id:
                            success = False
                            break

                if success:
                    if push_distance_sample:
                        distances = [npr.choice(push_distance_distribution)]
                    else:
                        distances = push_distance_distribution
                    kwargs = dict(
                        target_push_pos=target_point_pos,
                        target_push_dir=target_point_normal,
                        tool_pos=final_tool_pos,
                        tool_quat=final_tool_quat,
                        tool_point_pos=rotate_vector(tool_point_pos - current_tool_pos, rotation_quat) + final_tool_pos,
                        tool_point_normal=rotate_vector(tool_point_normal, rotation_quat),
                        prepush_distance=prepush_distance
                    )
                    for distance in distances:
                        yield IndirectPushParameter(**kwargs, push_distance=distance)


@dataclass
class PlacementParameter(object):
    target_pos: np.ndarray
    target_quat: np.ndarray

    support_normal: np.ndarray


def sample_placement(
    robot: Robot, target_id: int, support_id: int,
    retain_target_orientation: bool = False,
    batch_size: int = 100, max_attempts: int = 10000,
    placement_tol: float = 0.03, support_dir_tol: float = 30,
    verbose: bool = False
) -> Iterable[PlacementParameter]:
    """Sample a placement of the target object on the support object.

    Args:
        robot: the robot.
        target_id: the id of the target object.
        support_id: the id of the support object.
        retain_target_orientation: if True, the target object will be placed with the same orientation as the current one.
        batch_size: the number of samples in batch processing.
        max_attempts: the maximum number of attempts for sampling.
        placement_tol: the tolerance for the placement. To check collision, we will place the object at a position that is `placement_tol` away from the support surface and check for collision.
        support_dir_tol: the tolerance for the support direction. We will sample the placement direction such that the support direction is within `support_dir_tol` degree from the gravity direction.

    Yields:
        the placement parameters.
    """
    target_mesh = robot.world.get_mesh(target_id, zero_center=False)
    support_mesh = robot.world.get_mesh(support_id, zero_center=False)

    nr_batches = int(max_attempts / batch_size)
    for _ in range(nr_batches):
        support_points = support_mesh.sample_points_uniformly(batch_size, use_triangle_normal=True)
        normals = np.asarray(support_points.normals)
        support_indices = np.where(normals.dot(np.array([0, 0, 1])) > np.cos(np.deg2rad(support_dir_tol)))[0]

        # support_points_inner = o3d.geometry.PointCloud()
        # support_points_inner.points = o3d.utility.Vector3dVector(np.asarray(support_points.points)[support_indices])
        # support_points_inner.normals = o3d.utility.Vector3dVector(np.asarray(support_points.normals)[support_indices])
        # o3d.visualization.draw_geometries([support_mesh, support_points_inner])

        target_points = target_mesh.sample_points_uniformly(batch_size, use_triangle_normal=True)

        if retain_target_orientation:
            normals = np.asarray(target_points.normals)
            target_indices = np.where(normals.dot(np.array([0, 0, -1])) > np.cos(np.deg2rad(support_dir_tol)))[0]
        else:
            target_indices = list(range(batch_size))

        # o3d.visualization.draw_geometries([target_mesh, target_points])

        all_pair_indices = list(itertools.product(target_indices, support_indices))
        random.shuffle(all_pair_indices)
        for target_index, support_index in all_pair_indices:
            target_point_pos = np.asarray(target_points.points[target_index])
            target_point_normal = -np.asarray(target_points.normals[target_index])

            support_point_pos = np.asarray(support_points.points[support_index])
            support_point_normal = np.asarray(support_points.normals[support_index])

            current_target_pos, current_target_quat = robot.world.get_body_state_by_id(target_id).get_transformation()

            # Solve for a quaternion that aligns the tool normal with the target normal
            for rotation_quat in enumerate_quaternion_from_vectors(target_point_normal, support_point_normal, 4):
                # This is the world coordinate for the tool point after rotation.
                new_target_point_pos = current_target_pos + rotate_vector(target_point_pos - current_target_pos, rotation_quat)
                # Now compute the displacement for the tool object
                final_target_pos = support_point_pos - new_target_point_pos + current_target_pos
                final_target_pos += support_point_normal * placement_tol
                final_target_quat = quat_mul(rotation_quat, current_target_quat)

                # print('final_target_pos', final_target_pos, 'final_target_quat', final_target_quat)

                success = True
                with robot.world.save_body(target_id):
                    robot.world.set_body_state2_by_id(target_id, final_target_pos, final_target_quat)
                    contacts = robot.world.get_contact(target_id, update=True)

                    for contact in contacts:
                        if contact.body_b != target_id:
                            # print('collision ...', 'between', robot.world.get_body_name(contact.body_a), 'and', robot.world.get_body_name(contact.body_b))
                            # input('press enter to continue')
                            success = False
                            break

                if success:
                    yield PlacementParameter(
                        target_pos=final_target_pos,
                        target_quat=final_target_quat,
                        support_normal=support_point_normal
                    )


def qpos_trajectory_from_poses(robot: Robot, trajectory: List[Tuple[np.ndarray, np.ndarray]]) -> Tuple[bool, List[np.ndarray]]:
    """Given a list of 6D poses, generate a list of qpos that the robot should follow to follow the trajectory."""

    qpos_trajectory = list()
    last_qpos = robot.ikfast(trajectory[0][0], trajectory[0][1], error_on_fail=False)
    if last_qpos is None:
        return False, qpos_trajectory
    qpos_trajectory.append(last_qpos)

    success = True
    for this_pos, this_quat in trajectory:
        qpos = robot.ikfast(this_pos, this_quat, max_distance=0.1, last_qpos=last_qpos, error_on_fail=False)
        if qpos is None:
            return False, qpos_trajectory
        qpos_trajectory.append(qpos)
        last_qpos = qpos

    return success, qpos_trajectory


def qpos_trajectory_from_poses_pybullet(robot: Robot, trajectory: List[Tuple[np.ndarray, np.ndarray]]) -> Tuple[bool, List[np.ndarray]]:
    qpos_trajectory = list()
    for pos, quat in trajectory:
        qpos = robot.ik_pybullet(pos, quat)
        if qpos is None:
            return False, qpos_trajectory
        qpos_trajectory.append(qpos)
    return True, qpos_trajectory


def collision_free_qpos(robot: Robot, qpos: np.ndarray, exclude: Optional[List[int]] = None, verbose: bool = False) -> bool:
    """Check whether the given qpos is collision free. The function also accepts a list of object ids to exclude (e.g., the object in hand).

    Args:
        robot: the robot.
        qpos: the qpos to check.
        exclude: the object ids to exclude.
        verbose: whether to print the collision information.

    Returns:
        True if the qpos is collision free.
    """
    with robot.world.save_body(robot.get_robot_body_id()):
        robot.set_qpos(qpos)
        contacts = robot.world.get_contact(robot.get_robot_body_id(), update=True)
        if exclude is not None:
            for c in contacts:
                if c.body_b not in exclude and c.body_b != c.body_a:
                    if verbose:
                        print(f'  collision_free_qpos: collide bewteen {robot.world.body_names[c.body_a]} and {robot.world.body_names[c.body_b]}')
                    return False
        else:
            for c in contacts:
                if c.body_b != c.body_a:
                    if verbose:
                        print(f'  collision_free_qpos: collide bewteen {robot.world.body_names[c.body_a]} and {robot.world.body_names[c.body_b]}')
                    return False
        return True


def collision_free_pos(robot: Robot, pos: np.ndarray, quat: Optional[np.ndarray] = None, exclude: Optional[List[int]] = None, verbose: bool = False) -> bool:
    """Check whether there is a qpos at the givne pose that is collision free. The function also accepts a list of object ids to exclude (e.g., the object in hand).
    This function is not recommended to use, since it is not fully reproducible (there will be multiple qpos that can achieve the same pose).

    Args:
        robot: the robot.
        pos: the position to check.
        quat: the quaternion to check. If not specified, the home quaternion will be used.
        exclude: the object ids to exclude.
        verbose: whether to print the collision information.

    Returns:
        True if there is a collision free qpos.
    """

    if quat is None:
        quat = robot.get_ee_home_quat()
    qpos = robot.ikfast(pos, quat)
    if qpos is None:
        return False
    return collision_free_qpos(robot, qpos, exclude=exclude, verbose=verbose)


def collision_free_qpos_trajectory(robot: Robot, qpos_trajectory: List[np.ndarray], exclude: Optional[List[int]] = None, verbose: bool = False) -> bool:
    for qpos in qpos_trajectory:
        if not collision_free_qpos(robot, qpos, exclude=exclude, verbose=verbose):
            return False
    return True


def mesh_line_intersect(t_mesh: o3d.t.geometry.TriangleMesh, ray_origin: np.ndarray, ray_direction: np.ndarray) -> Optional[Tuple[np.ndarray, np.ndarray]]:
    """Intersects a ray with a mesh.

    Args:
        t_mesh: the mesh to intersect with.
        ray_origin: the origin of the ray.
        ray_direction: the direction of the ray.

    Returns:
        A tuple of (point, normal) if an intersection is found, None otherwise.
    """

    scene = o3d.t.geometry.RaycastingScene()
    scene.add_triangles(t_mesh)

    ray = o3d.core.Tensor.from_numpy(np.array(
        [[ray_origin[0], ray_origin[1], ray_origin[2], ray_direction[0], ray_direction[1], ray_direction[2]]],
        dtype=np.float32
    ))
    result = scene.cast_rays(ray)

    # no intersection.
    if result['geometry_ids'][0] == scene.INVALID_ID:
        return None

    inter_point = np.asarray(ray_origin) + np.asarray(ray_direction) * result['t_hit'][0].item()
    inter_normal = result['primitive_normals'][0].numpy()
    return inter_point, inter_normal


def get_single_contact_normal(robot: Robot, object_id: int, support_object_id: int, deviation_tol: float = 0.05) -> np.ndarray:
    body_names = robot.world.body_names.int_to_string
    object_name, support_name = body_names[object_id], body_names[support_object_id]
    contacts = robot.world.get_contact(object_id, support_object_id)

    if len(contacts) == 0:
        # TODO(Jiayuan Mao @ 2023/03/15): find a better way to configure this.
        robot.client.step(1)
        contacts = robot.world.get_contact(object_id, support_object_id)

    if len(contacts) == 0:
        raise ValueError(f'No contact between {object_name} and {support_name}.')

    contact_normals = np.array([c.contact_normal_on_b for c in contacts])
    contact_normal_avg = np.mean(contact_normals, axis=0)
    contact_normal_avg /= np.linalg.norm(contact_normal_avg)

    deviations = np.abs(1 - contact_normals.dot(contact_normal_avg) / np.linalg.norm(contact_normals, axis=1))
    if np.max(deviations) > deviation_tol:
        raise ValueError(
            f'Contact normals of {object_name} and {support_name} are not consistent. This is likely due to multiple contact points.\n'
            f'  Contact normals: {contact_normals}\n  Deviations: {deviations}.'
        )

    return contact_normal_avg


concepts/simulator/pybullet/manipulation_utils/contact_samplers.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : grasping_samplers.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/06/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
import trimesh
import open3d as o3d

from dataclasses import dataclass
from typing import Union, Optional, Tuple

from concepts.simulator.pybullet.rotation_utils import get_quaternion_from_matrix, find_orthogonal_vector
from concepts.simulator.pybullet.components.robot_base import Robot


@dataclass
class GraspGeneratorConfig(object):
    pass


@dataclass
class GraspReturn(object):
    point1: np.ndarray
    normal1: np.ndarray
    point2: np.ndarray
    normal2: np.ndarray
    ee_pos: np.ndarray
    ee_quat: np.ndarray
    qpos: np.ndarray


class GraspSampler(object):
    """Generate grasps for the input mesh."""

    def __init__(self, robot: Robot, gripper_size: float, gripper_gap: float, reachability_min_bound: Optional[np.ndarray] = None, reachability_max_bound: Optional[np.ndarray] = None):
        """Initialize the grasp generator.

        Args:
            robot: the robot instance.
            gripper_size: the size of the gripper (assumed to be a square contact surface).
            gripper_gap: the gap between the gripper fingers.
            reachability_min_bound: the minimum bound of the reachable workspace. If specified, will be used to filter out grasps that are not reachable.
            reachability_max_bound: the maximum bound of the reachable workspace. If specified, will be used to filter out grasps that are not reachable.
        """
        self.robot = robot
        self.gripper_size = gripper_size
        self.gripper_gap = gripper_gap

        self.reachability_min_bound = reachability_min_bound
        self.reachability_max_bound = reachability_max_bound

    def sample_grasp(self, input_shape: Union[np.ndarray, o3d.geometry.PointCloud, o3d.geometry.TriangleMesh]) -> Optional[GraspReturn]:
        """Sample a grasp from the input point cloud, or mesh

        Args:
            input_pc: the input point cloud.

        Returns:
            GraspReturn: The sampled grasp.
        """

        if isinstance(input_shape, np.ndarray):
            input_shape = o3d.geometry.PointCloud()
            input_shape.points = o3d.utility.Vector3dVector(input_shape)
            input_shape.estimate_normals()

        if self.reachability_min_bound is not None and self.reachability_max_bound is not None:
            if isinstance(input_shape, o3d.geometry.TriangleMesh):
                aabb = o3d.geometry.AxisAlignedBoundingBox(self.reachability_min_bound, self.reachability_max_bound)
                input_shape = input_shape.crop(aabb)
                if len(input_shape.triangles) < 10:
                    return None
            elif isinstance(input_shape, o3d.geometry.PointCloud):
                aabb = o3d.geometry.AxisAlignedBoundingBox(self.reachability_min_bound, self.reachability_max_bound)
                input_shape = input_shape.crop(aabb)
                if len(input_shape.points) < 10:
                    return None
            else:
                raise TypeError(f'Unknown supported type {type(input_shape)}.')

        return self._sample_grasp_pcd(input_shape)

    def _sample_grasp_pcd(self, input_pc: o3d.geometry.PointCloud) -> Optional[GraspReturn]:
        raise TypeError('PCD grasp sampling is not implemented.')

    def _sample_grasp_mesh(self, input_mesh: o3d.geometry.TriangleMesh) -> Optional[GraspReturn]:
        raise TypeError('Mesh grasp sampling is not implemented.')


@dataclass
class LocalConvexificationGraspGeneratorConfig(GraspGeneratorConfig):
    """The minimum number of points in the point cloud to perform local convexification."""
    min_pointcloud_size: int = 10

    """The minimum gap between two gripper attachment points."""
    min_gripper_gap: float = 0.01

    """The minimum difference between the two normals of the gripper attachment points."""
    min_normal_difference: float = 0.9

    """The minimum quaternion norm."""
    min_quat_norm: float = 0.5

    """The maximum difference between the ik solution and the target."""
    max_ik_diff: float = 0.02

    """If true, the z-axis of the gripper should point to the front."""
    constraint_pos_x: bool = True

    """Visualize the configuration returned by the grasp generator."""
    debug_vis_return: bool = False

    """Visualize the each configuration tested by the grasp generator."""
    debug_vis_trials: bool = False


class LocalConvexificationGraspSampler(GraspSampler):
    config: LocalConvexificationGraspGeneratorConfig

    def __init__(
        self, robot: Robot,
        gripper_size: float, gripper_gap: float,
        reachability_min_bound: Optional[np.ndarray] = None, reachability_max_bound: Optional[np.ndarray] = None,
        config: Optional[LocalConvexificationGraspGeneratorConfig] = None,
        **kwargs
    ):
        super().__init__(robot, gripper_size, gripper_gap, reachability_min_bound, reachability_max_bound)

        if config is None:
            self.config = LocalConvexificationGraspGeneratorConfig(**kwargs)
        else:
            assert len(kwargs) == 0, 'Cannot specify both config and kwargs.'
            self.config = config

    def _sample_grasp_pcd(self, pcd: o3d.geometry.PointCloud, trials: int = 100) -> Optional[GraspReturn]:
        for _ in range(trials):
            # Sample a point in the point cloud.
            index = np.random.randint(0, len(pcd.points))
            point1 = pcd.points[index]
            normal1 = -pcd.normals[index]

            # Make a local crop of the point cloud.
            min_bound = point1 - self.gripper_size * 2
            max_bound = point1 + self.gripper_size * 2
            pcd_crop = pcd.crop(o3d.geometry.AxisAlignedBoundingBox(min_bound, max_bound))

            if len(pcd_crop.points) < self.config.min_pointcloud_size:
                continue

            # Perform local convexification.
            # print(np.asarray(pcd_crop.points))
            # print(pcd_crop.get_axis_aligned_bounding_box())
            hull, _ = pcd_crop.compute_convex_hull()

            # Create the triangular mesh with the vertices and faces from open3d.
            tmesh = trimesh.Trimesh(np.asarray(hull.vertices), np.asarray(hull.triangles), vertex_normals=np.asarray(hull.vertex_normals))
            # Ray casting the point cloud to find the other point on the mesh.
            intersector = trimesh.ray.ray_triangle.RayMeshIntersector(tmesh)
            index, point2 = _trimesh_ray_casting(intersector, point1, normal1)

            if index is not None:
                normal2 = tmesh.face_normals[index]

                # Filter the gap.
                gap = np.linalg.norm(point1 - point2)
                if gap < self.config.min_gripper_gap:
                    continue
                if gap > self.gripper_gap:
                    continue

                # Filter the normal.
                if np.abs(np.dot(normal1, normal2)) < self.config.min_normal_difference:
                    continue

                # Construct three axes for the gripper.
                ee_pos = (point1 + point2) / 2
                ee_d = (point2 - point1) / np.linalg.norm(point2 - point1)
                ee_u = find_orthogonal_vector(point2 - point1)
                ee_v = np.cross(ee_d, ee_u)

                with self.robot.client.w.save_world() as world_saver:
                    for ee_norm1 in [ee_u, ee_v, -ee_u, -ee_v]:
                        if self.config.debug_vis_trials:
                            world_saver.restore()

                        ee_norm2 = np.cross(ee_d, ee_norm1)

                        # if self.config.constraint_pos_x and ee_norm2[0] < 0:
                        #     continue
                        # if self.config.constraint_pos_x and ee_norm2[0] > 0:
                        #     ee_norm2 = -ee_norm2
                        #     ee_norm1 = -ee_norm1

                        ee_quat = self._construct_quat(ee_norm2, ee_d, ee_norm1)

                        if self.config.debug_vis_trials:
                            print(f'Testing x={ee_norm2}, y={ee_d}, z={ee_norm1}, point1={point1}, point2={point2}.')

                        ee_grasp = self._check_collision(ee_pos, ee_quat)

                        if self.config.debug_vis_trials:
                            self.robot.client.step(2)
                            input('Press enter to continue.')

                        if ee_grasp is not None:
                            return GraspReturn(point1, normal1, point2, normal2, *ee_grasp)
            else:
                continue
        return None

    def _construct_quat(self, x: np.ndarray, y: np.ndarray, z: np.ndarray) -> np.ndarray:
        """Construct a quaternion from the new x-y-z coordinate axes of the gripper."""
        return get_quaternion_from_matrix(np.stack([x, y, z], axis=1))

    def _check_collision(self, ee_pos, ee_quat) -> Optional[Tuple[np.ndarray, np.ndarray, np.ndarray]]:
        """Check if the grasping is collision-free."""

        if np.linalg.norm(ee_quat) < self.config.min_quat_norm:
            if self.config.debug_vis_trials:
                print('Quaternion checking failed.')
            return None

        qpos = self.robot.ik(ee_pos, ee_quat, force=True)

        if qpos is None:
            if self.config.debug_vis_trials:
                print('IK failed.')
            return None

        self.robot.set_qpos(qpos)
        if np.linalg.norm(self.robot.get_ee_pose()[0] - ee_pos) > self.config.max_ik_diff:
            if self.config.debug_vis_trials:
                print('IK post checking failed: ', np.linalg.norm(self.robot.get_ee_pose()[0] - ee_pos))
            return None

        self.robot.client.p.performCollisionDetection()

        # self.robot.client.step(1)
        # input('Press enter to continue')

        contacts = self.robot.client.w.get_contact(a=self.robot.get_robot_body_id())

        all_contact_bodies = {c.body_b for c in contacts} - {self.robot.get_robot_body_id()}
        all_contact_body_names = [self.robot.client.w.body_names[b] for b in all_contact_bodies]

        if self.config.debug_vis_return:
            if len(all_contact_bodies) == 0:
                print('No contact')
                print('-' * 80)
                print(*contacts, sep='\n')
                print(all_contact_body_names)
                self.robot.client.step(5)
        if self.config.debug_vis_trials:
            print('Contacts')
            for c in contacts:
                print(c.body_a_name, c.body_b_name, c.link_a_name, c.link_b_name, c.position_on_a, c.position_on_b)

        if len(all_contact_body_names) > 0:
            return None

        return ee_pos, ee_quat, qpos


def _trimesh_ray_casting(intersector, point, normal):
    """Ray casting using trimesh ray interesector."""

    index, _, locations = intersector.intersects_id([point], [normal], return_locations=True, multiple_hits=True)
    if len(index) > 0:
        if len(index) > 1:
            distances = np.linalg.norm(locations - point, axis=1)
            i = np.argmax(distances)
            return index[i], locations[i]
        return index[0], locations[0]
    return None, None


concepts/simulator/pybullet/manipulation_utils/grasping_samplers.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : mjcf_utils.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/17/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import jacinle.io as io
import os.path as osp

__all__ = ['MJCFCanonizer', 'canonize_mjcf']


class MJCFCanonizer(object):
    def __init__(self, filename, xml_object=None):
        self.filename = filename
        self.dirname = osp.dirname(self.filename)
        self.xml_object = xml_object

        if self.xml_object is None:
            self.xml_object = io.load_xml(self.filename)

        self.directories = dict()

    def resolve_path(self, path, type=None):
        if osp.isfile(path):
            return path
        if type is not None:
            if type in self.directories:
                asset_path = osp.join(self.directories[type], path)
                if osp.isfile(asset_path):
                    return path
                else:
                    rel_path = osp.realpath(osp.join(self.dirname, path))
                    rel_path = osp.relpath(rel_path, start=self.directories[type])
                    return rel_path
        rel_path = osp.realpath(osp.join(self.dirname, path))
        return rel_path

    def _canonize_include_inner(self, object):
        found = False
        for k in list(object.keys()):
            if k == 'include':
                found = True
                v = object.pop(k)
                if isinstance(v, list):
                    for sub_v in v:
                        self._canonize_include_append(object, sub_v)
                else:
                    self._canonize_include_append(object, v)
            elif not k.startswith('__'):
                v = object[k]
                if isinstance(v, list):
                    for sub_v in v:
                        found = found or self._canonize_include_inner(sub_v)
                else:
                    found = found or self._canonize_include_inner(v)
        return found

    def _canonize_include_append(self, object, include_obj):
        sub_xml_object = io.load_xml(self.resolve_path(include_obj['__attribute__']['file']))
        for name, child in sub_xml_object.items():
            if not name.startswith('__'):
                _xml_add_to(object, name, child)

    def _canonize_include(self, object):
        for i in range(100):
            found = self._canonize_include_inner(object)
            if not found:
                break

    def _canonize_directories(self, object):
        if 'compiler' in object:
            attributes = object['compiler']['__attribute__']
            if 'meshdir' in attributes:
                self.directories['mesh'] = attributes['meshdir'] = self.resolve_path(attributes['meshdir']) + '/'
            if 'texturedir' in attributes:
                self.directories['texture'] = attributes['texturedir'] = (
                    self.resolve_path(attributes['texturedir']) + '/'
                )

    def _canonize_attributes(self, object):
        for k, v in object.items():
            if not k.startswith('__'):
                if isinstance(v, list):
                    for sub_v in v:
                        self._canonize_attributes(sub_v)
                else:
                    self._canonize_attributes(v)

        tag, attributes = object['__name__'], object['__attribute__']

        if tag == 'geom':
            if 'type' not in attributes:
                attributes['type'] = 'sphere'
        elif tag == 'joint':
            if 'type' not in attributes:
                attributes['type'] = 'hinge'
        elif tag == 'mesh':
            if 'file' in attributes:
                attributes['file'] = self.resolve_path(attributes['file'], type='mesh')
        elif tag == 'texture':
            if 'file' in attributes:
                attributes['file'] = self.resolve_path(attributes['file'], type='texture')

    def canonize(self):
        self._canonize_include(self.xml_object)
        self._canonize_directories(self.xml_object)
        self._canonize_attributes(self.xml_object)
        return self.xml_object


def canonize_mjcf(filename, xml_object=None):
    return MJCFCanonizer(filename, xml_object=xml_object).canonize()


def _xml_add_to(object, name, child):
    if name not in object:
        object[name] = child
    else:
        if isinstance(object[name], list):
            if isinstance(child, list):
                object[name].extend(child)
            else:
                object[name].append(child)
        else:
            if isinstance(child, list):
                object[name] = [object[name]] + child
            else:
                object[name] = [object[name], child]
    return object

concepts/simulator/pybullet/mjcf_utils.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : rotation_utils.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/18/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
from typing import Tuple, TYPE_CHECKING
from concepts.utils.rotationlib import euler2quat

if TYPE_CHECKING:
    from concepts.simulator.pybullet.client import BulletClient


# For testing whether a number is close to zero
_FLOAT_EPS = np.finfo(np.float64).eps
_EPS4 = _FLOAT_EPS * 4.0


def normalize_vector(v: np.ndarray) -> np.ndarray:
    """Normalize a vector."""
    return v / np.linalg.norm(v)


def find_orthogonal_vector(v: np.ndarray) -> np.ndarray:
    """Find an orthogonal vector to the given vector.

    The returned vector is guaranteed to be normalized.
    """
    v = v / np.linalg.norm(v, axis=-1, keepdims=True)

    mask = np.abs(v[..., 0]) < 0.5
    return np.cross(v, np.array([1, 0, 0])) * mask + np.cross(v, np.array([0, 1, 0])) * (1 - mask)


def getQuaternionFromMatrix(mat):
    """ Convert Rotation Matrix to Quaternion.  See rotation.py for notes """
    mat = np.asarray(mat, dtype=np.float64)
    assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)

    Qxx, Qyx, Qzx = mat[..., 0, 0], mat[..., 0, 1], mat[..., 0, 2]
    Qxy, Qyy, Qzy = mat[..., 1, 0], mat[..., 1, 1], mat[..., 1, 2]
    Qxz, Qyz, Qzz = mat[..., 2, 0], mat[..., 2, 1], mat[..., 2, 2]
    # Fill only lower half of symmetric matrix
    K = np.zeros(mat.shape[:-2] + (4, 4), dtype=np.float64)
    K[..., 0, 0] = Qxx - Qyy - Qzz
    K[..., 1, 0] = Qyx + Qxy
    K[..., 1, 1] = Qyy - Qxx - Qzz
    K[..., 2, 0] = Qzx + Qxz
    K[..., 2, 1] = Qzy + Qyz
    K[..., 2, 2] = Qzz - Qxx - Qyy
    K[..., 3, 0] = Qyz - Qzy
    K[..., 3, 1] = Qzx - Qxz
    K[..., 3, 2] = Qxy - Qyx
    K[..., 3, 3] = Qxx + Qyy + Qzz
    K /= 3.0
    # TODO: vectorize this -- probably could be made faster
    q = np.empty(K.shape[:-2] + (4,))
    it = np.nditer(q[..., 0], flags=['multi_index'])
    while not it.finished:
        # Use Hermitian eigenvectors, values for speed
        vals, vecs = np.linalg.eigh(K[it.multi_index])
        # Select largest eigenvector, reorder to w,x,y,z quaternion
        q[it.multi_index] = vecs[:, np.argmax(vals)]
        # Prefer quaternion with positive w
        # (q * -1 corresponds to same rotation as q)
        if q[it.multi_index][3] < 0:
            q[it.multi_index] *= -1
        it.iternext()
    return q


def get_quaternion_from_matrix(mat):
    """Convert a rotation matrix to a quaternion.

    The quaternion is represented as a 4-element array in the order of (x, y, z, w),
    which follows the convention of pybullet.
    """
    mat = np.asarray(mat, dtype=np.float64)
    assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)
    return getQuaternionFromMatrix(mat)


def get_quaternion_from_axes(x: np.ndarray, y: np.ndarray, z: np.ndarray) -> np.ndarray:
    """Converts a rotation matrix to a quaternion."""
    m = np.stack([x, y, z], axis=1)
    return get_quaternion_from_matrix(m)


def compose_transformation(pos1, quat1, pos2, quat2):
    """Compsoe two transformations.

    The transformations are represented as a tuple of position and quaternion.
    The quaternion is represented as a 4-element array in the order of (x, y, z, w), which follows the convention of pybullet.
    """

    pos1 = np.asarray(pos1, dtype=np.float32)
    quat1 = np.asarray(quat1, dtype=np.float32)
    pos2 = np.asarray(pos2, dtype=np.float32)
    quat2 = np.asarray(quat2, dtype=np.float32)

    return pos1 + rotate_vector(pos2, quat1), quat_mul(quat1, quat2)


def inverse_transformation(pos, quat):
    """Inverse a transformation.

    The transformation is represented as a tuple of position and quaternion.
    The quaternion is represented as a 4-element array in the order of (x, y, z, w), which follows the convention of pybullet.
    """
    pos = np.asarray(pos, dtype=np.float32)
    quat = np.asarray(quat, dtype=np.float32)

    inv_pos = -pos
    inv_quat = quat_conjugate(quat)
    return inv_pos, inv_quat


def get_transform_a_to_b(pos1, quat1, pos2, quat2):
    """Get the transformation from frame A to frame B.

    The transformations are represented as a tuple of position and quaternion.
    The quaternion is represented as a 4-element array in the order of (x, y, z, w), which follows the convention of pybullet.
    """
    pos1 = np.asarray(pos1, dtype=np.float32)
    quat1 = np.asarray(quat1, dtype=np.float32)
    pos2 = np.asarray(pos2, dtype=np.float32)
    quat2 = np.asarray(quat2, dtype=np.float32)

    inv_quat1 = quat_conjugate(quat1)
    a_to_b_pos = rotate_vector(pos2 - pos1, inv_quat1)
    a_to_b_quat = quat_mul(inv_quat1, quat2)
    return a_to_b_pos, a_to_b_quat


def frame_mul(pos_a, quat_a, a_to_b):
    """Multiply a frame with a transformation.

    The frame is represented as a tuple of position and quaternion.
    The transformation is represented as a tuple of position and quaternion.
    The quaternion is represented as a 4-element array in the order of (x, y, z, w), which follows the convention of pybullet.
    """
    pos_a = np.asarray(pos_a, dtype=np.float32)
    quat_a = np.asarray(quat_a, dtype=np.float32)
    transform_pos = np.asarray(a_to_b[0], dtype=np.float32)
    transform_quat = np.asarray(a_to_b[1], dtype=np.float32)

    pos_b = pos_a + rotate_vector(transform_pos, quat_a)
    quat_b = quat_mul(quat_a, transform_quat)
    return pos_b, quat_b


def frame_inv(pos_b, quat_b, a_to_b):
    """Inverse a frame with a transformation.

    The frame is represented as a tuple of position and quaternion.
    The transformation is represented as a tuple of position and quaternion.
    The quaternion is represented as a 4-element array in the order of (x, y, z, w), which follows the convention of pybullet.
    """
    pos_b = np.asarray(pos_b, dtype=np.float32)
    quat_b = np.asarray(quat_b, dtype=np.float32)
    transform_pos = np.asarray(a_to_b[0], dtype=np.float32)
    transform_quat = np.asarray(a_to_b[1], dtype=np.float32)

    inv_transform_quat = quat_conjugate(transform_quat)
    quat_a = quat_mul(quat_b, inv_transform_quat)
    pos_a = pos_b - rotate_vector(transform_pos, quat_a)
    return pos_a, quat_a


def quat_mul(q0, q1, *args):
    if len(args) > 0:
        return quat_mul(q0, quat_mul(q1, *args))

    q0 = np.asarray(q0)
    q1 = np.asarray(q1)
    assert q0.shape == q1.shape
    assert q0.shape[-1] == 4
    assert q1.shape[-1] == 4

    w0 = q0[..., 3]
    x0 = q0[..., 0]
    y0 = q0[..., 1]
    z0 = q0[..., 2]

    w1 = q1[..., 3]
    x1 = q1[..., 0]
    y1 = q1[..., 1]
    z1 = q1[..., 2]

    w = w0 * w1 - x0 * x1 - y0 * y1 - z0 * z1
    x = w0 * x1 + x0 * w1 + y0 * z1 - z0 * y1
    y = w0 * y1 + y0 * w1 + z0 * x1 - x0 * z1
    z = w0 * z1 + z0 * w1 + x0 * y1 - y0 * x1
    q = np.array([x, y, z, w])
    if q.ndim == 2:
        q = q.swapaxes(0, 1)
    assert q.shape == q0.shape

    if w0 > 0:
        return q
    else:
        return -q


def quat_conjugate(q):
    q = np.asarray(q)
    assert q.shape[-1] == 4
    q = np.array([-q[..., 0], -q[..., 1], -q[..., 2], q[..., 3]])
    return q


def rpy(r, p, y, degree=True):
    """Create a quaternion from euler angles."""
    if degree:
        r = np.deg2rad(r)
        p = np.deg2rad(p)
        y = np.deg2rad(y)
    quat_wxyz = euler2quat((r, p, y))
    return quat_wxyz[..., [1, 2, 3, 0]]


def rotate_vector(vec, quat):
    """Rotate a vector by a quaternion."""
    vec = np.asarray(vec)
    quat = np.asarray(quat)
    u = quat[:3]
    s = quat[3]
    return 2.0 * np.dot(u, vec) * u + (s * s - np.dot(u, u)) * vec + 2.0 * s * np.cross(u, vec)


def rotate_vector_batch(batch_vector, quat):
    """Rotate a batch of vectors by a quaternion."""
    vec = np.asarray(batch_vector)
    quat = np.asarray(quat)

    u = quat[:3]
    s = quat[3]

    return 2.0 * np.dot(vec, u)[..., np.newaxis] * u + (s * s - np.dot(u, u)) * vec + 2.0 * s * np.cross(u, vec)


def quaternion_from_vectors(vec1, vec2):
    """Create a rotation quaternion q such that q * vec1 = vec2."""
    vec1 = np.asarray(vec1)
    vec2 = np.asarray(vec2)
    assert vec1.shape == vec2.shape, f'Vector shapes do not match: {vec1.shape} != {vec2.shape}'
    assert vec1.shape[-1] == 3, f'Vector shape must be 3, got {vec1.shape[-1]}'
    assert vec2.shape[-1] == 3, f'Vector shape must be 3, got {vec2.shape[-1]}'

    vec1 = vec1 / np.linalg.norm(vec1, axis=-1, keepdims=True)
    vec2 = vec2 / np.linalg.norm(vec2, axis=-1, keepdims=True)

    u = np.cross(vec1, vec2)
    s = np.dot(vec1, vec2)

    if np.linalg.norm(u) < 1e-6:
        return np.array([0, 0, 0, 1])

    opposite_pairs = (s < -1 + 1e-6)
    u = u * (1 - opposite_pairs) + opposite_pairs * find_orthogonal_vector(u)

    if len(u.shape) == 1:
        s = np.array([s], dtype=u.dtype)
    q = np.concatenate([u, s + 1], axis=-1)

    q = q / np.linalg.norm(q, axis=-1, keepdims=True)
    return q


def enumerate_quaternion_from_vectors(input_normal, target_normal, nr_samples: int = 4):
    base_quat = quaternion_from_vectors(input_normal, target_normal)

    yaw_quat = quaternion_from_vectors(target_normal, np.array([0, 0, 1]))
    for yaw in np.arange(0, 2 * np.pi, 2 * np.pi / nr_samples):
        quat = quat_mul(
            quat_conjugate(yaw_quat),
            rpy(0, 0, yaw, degree=False),
            yaw_quat,
            base_quat
        )
        yield quat


def patch_pybullet():
    import pybullet as p

    p.getQuaternionFromMatrix = getQuaternionFromMatrix


def get_ee_to_tool(client: 'BulletClient', robot_id, ee_id, tool_id: int) -> Tuple[np.ndarray, np.ndarray]:
    """Get the transformation from the end-effector to the tool that the robot is currently holding."""

    # if robot.gripper_constraint is not None:
    #     constraint = robot.world.get_constraint(robot.gripper_constraint)
    #     assert constraint.child_body == tool_id
    #     from concepts.simulator.pybullet.rotation_utils import quat_conjugate
    #     return -np.array(constraint.parent_frame_pos), quat_conjugate(np.array(constraint.parent_frame_orn))

    robot_pos, robot_quat = client.world.get_link_state_by_id(robot_id, ee_id, fk=True).get_transformation()
    tool_pos, tool_quat = client.world.get_body_state_by_id(tool_id).get_transformation()
    return get_transform_a_to_b(robot_pos, robot_quat, tool_pos, tool_quat)


def solve_tool_from_ee(ee_pos, ee_quat, ee_to_tool):
    return frame_mul(ee_pos, ee_quat, ee_to_tool)


def solve_ee_from_tool(target_tool_pos, target_tool_quat, ee_to_tool):
    """Solve for the end-effector position and orientation given the tool position and orientation."""
    return frame_inv(target_tool_pos, target_tool_quat, ee_to_tool)


concepts/simulator/pybullet/rotation_utils.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : rotation_utils_torch.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/22/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import torch

__all__ = ['th_quat_mul', 'th_quat_conjugate', 'th_compose_transformation', 'th_inverse_transformation']


def th_quat_mul(quat1: torch.Tensor, quat2: torch.Tensor, *args: torch.Tensor) -> torch.Tensor:
    """Multiply two quaternions. The quaternions are represented as 4D vectors,
    in the form of [x, y, z, w].

    Args:
        quat1: a tensor of shape (..., 4).
        quat2: a tensor of shape (..., 4).
        *args: additional tensors of shape (..., 4).

    Returns:
        A tensor of shape (..., 4).
    """

    if len(args) > 0:
        return th_quat_mul(quat1, th_quat_mul(quat2, *args))

    assert quat1.shape == quat2.shape and quat1.shape[-1] == 4 and quat2.shape[-1] == 4

    x1, y1, z1, w1 = quat1[..., 0], quat1[..., 1], quat1[..., 2], quat1[..., 3]
    x2, y2, z2, w2 = quat2[..., 0], quat2[..., 1], quat2[..., 2], quat2[..., 3]

    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 + y1 * w2 + z1 * x2 - x1 * z2
    z = w1 * z2 + z1 * w2 + x1 * y2 - y1 * x2
    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2

    q = torch.stack([x, y, z, w], dim=-1)
    sign = -1 + 2 * (w > 0)

    return q * sign


def th_quat_conjugate(quat: torch.Tensor) -> torch.Tensor:
    """Compute the conjugate of a quaternion. The quaternion is represented as a 4D vector,
    in the form of [x, y, z, w].

    Args:
        quat: the quaternion, as a 4D vector.

    Returns:
        the conjugate, as a 4D vector.
    """

    return torch.stack([-quat[..., 0], -quat[..., 1], -quat[..., 2], quat[..., 3]], dim=-1)


def th_compose_transformation(pose1: torch.Tensor, pose2: torch.Tensor, *args: torch.Tensor) -> torch.Tensor:
    """Compose two transformations.

    Args:
        pose1: the first transformation, as a 7D vector.
        pose2: the second transformation, as a 7D vector.
        *args: additional transformations, as 7D vectors.

    Returns:
        the composed transformation, as a 7D vector.
    """

    if len(args) > 0:
        return th_compose_transformation(pose1, th_compose_transformation(pose2, *args))

    pos1, quat1 = pose1[..., :3], pose1[..., 3:]
    pos2, quat2 = pose2[..., :3], pose2[..., 3:]

    pos = pos1 + th_rotate_vector(pos2, quat1)
    quat = th_quat_mul(quat1, quat2)
    return torch.cat([pos, quat], dim=-1)


def th_inverse_transformation(pose1: torch.Tensor) -> torch.Tensor:
    """Compute the inverse of a transformation.

    Args:
        pose1 the transformation, as a 7D vector.

    Returns:
        the inverse transformation, as a 7D vector.
    """

    pos1, quat1 = pose1[..., :3], pose1[..., 3:]
    quat = th_quat_conjugate(quat1)
    return torch.cat([-pos1, quat], dim=-1)


def th_rotate_vector(vec: torch.Tensor, quat: torch.Tensor) -> torch.Tensor:
    """Rotate a vector by a quaternion."""
    u = quat[..., :3]
    s = quat[..., 3]
    return 2.0 * torch.einsum('...i,...i->...', u, vec) * u + (s * s - torch.einsum('...i,...i->...', u, u)) * vec + 2.0 * s * torch.cross(u, vec)

concepts/simulator/pybullet/rotation_utils_torch.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 07/10/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/simulator/pybullet/urdf_utils/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : obj2urdf.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 08/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
import tempfile
import os
import copy
import trimesh
import xml.etree.ElementTree as ET

from scipy.spatial.transform import Rotation

"""
Convert obj file to urdf file.

Copied from https://github.com/harvard-microrobotics/object2urdf
"""


class ObjectUrdfBuilder:
    def __init__(self, object_folder, log_file=None, urdf_prototype=None):
        self.object_folder = os.path.abspath(object_folder)

        if log_file is not None:
            self.log_file = os.path.abspath(log_file)
        else:
            f = tempfile.NamedTemporaryFile('r', suffix='log', prefix='vhacd', delete=False)
            f.close()
            self.log_file = f.name

        self.suffix = "vhacd"

        if urdf_prototype is None:
            self.urdf_base = self._read_xml(os.path.join(os.path.dirname(__file__), "urdf_prototype.urdf"))
        else:
            self.urdf_base = self._read_xml(os.path.join(object_folder, urdf_prototype))

    # Recursively get all files with a specific extension, excluding a certain suffix
    def _get_files_recursively(self, start_directory, filter_extension, exclude_suffix):
        for root, dirs, files in os.walk(start_directory):
            for file in files:
                if file.lower().endswith(filter_extension):
                    if not file.lower().endswith(exclude_suffix + filter_extension):
                        yield (root, file, os.path.abspath(os.path.join(root, file)))

    # Read and parse a URDF from a file
    def _read_xml(self, filename):
        root = ET.parse(filename).getroot()
        return root

    # Convert a list to a space-separated string
    def _list2str(self, in_list):
        out = ""
        for el in in_list:
            out += str(el) + " "
        return out[:-1]

    # Convert a space-separated string to a list
    def _str2list(self, in_str):
        out = in_str.split(' ')
        out = [float(el) for el in out]
        return out

    # Find the center of mass of the object
    def get_center_of_mass(self, filename):
        mesh = trimesh.load(filename)
        if isinstance(mesh, trimesh.Scene):
            print("Imported combined mesh: using centroid rather than center of mass")
            return mesh.centroid
        else:
            return mesh.center_mass

    # Find the geometric center of the object
    def get_geometric_center(self, filename):
        mesh = trimesh.load(filename)
        return copy.deepcopy(mesh.centroid)

    # Get the middle of a face of the bounding box
    def get_face(self, filename, edge):
        mesh = trimesh.load(filename)
        bounds = mesh.bounds
        face = copy.deepcopy(mesh.centroid)
        if edge in ['top', 'xy_pos']:
            face[2] = bounds[1][2]
        elif edge in ['bottom', 'xy_neg']:
            face[2] = bounds[0][2]
        elif edge in ['xz_pos']:
            face[1] = bounds[1][1]
        elif edge in ['xz_neg']:
            face[1] = bounds[0][1]
        elif edge in ['yz_pos']:
            face[0] = bounds[1][0]
        elif edge in ['yz_neg']:
            face[0] = bounds[0][0]

        return face

    # Do a convex decomposition
    def do_vhacd(self, filename, outfile, debug=False, **kwargs):
        try:
            mesh = trimesh.load(filename)
            convex_list = trimesh.interfaces.vhacd.convex_decomposition(mesh, debug=debug, **kwargs)

            convex = trimesh.util.concatenate(convex_list)
            convex.export(outfile)
        except ValueError:
            print("No direct VHACD backend available, trying pybullet")
            pass

        try:
            import pybullet as p
            p.vhacd(filename, outfile, self.log_file, **kwargs)
        except ModuleNotFoundError:
            print(
                '\n' + "ERROR - pybullet module not found: If you want to do convex decomposisiton, make sure you install pybullet (https://pypi.org/project/pybullet) or install VHACD directly (https://github.com/mikedh/trimesh/issues/404)" + '\n'
            )
            raise

    # Find the center of mass of the object
    def save_to_obj(self, filename):
        name, ext = os.path.splitext(filename)
        obj_filename = name + '.obj'
        mesh = trimesh.load(filename)
        mesh.export(obj_filename)
        return obj_filename

    # Replace an attribute in a feild of a URDF
    def replace_urdf_attribute(self, urdf, feild, attribute, value):
        urdf = self.replace_urdf_attributes(urdf, feild, {attribute: value})
        return urdf

    # Replace several attributes in a feild of a URDF
    def replace_urdf_attributes(self, urdf, feild, attribute_dict, sub_feild=None):

        if sub_feild is None:
            sub_feild = []

        field_obj = urdf.find(feild)

        if field_obj is not None:
            if len(sub_feild) > 0:
                for child in reversed(sub_feild):
                    field_obj = ET.SubElement(field_obj, child)
            field_obj.attrib.update(attribute_dict)
            # field_obj.attrib = attribute_dict
        else:
            feilds = feild.split("/")
            new_feild = "/".join(feilds[0:-1])
            sub_feild.append(feilds[-1])
            self.replace_urdf_attributes(urdf, new_feild, attribute_dict, sub_feild)

    # Make an updated copy of the URDF for the current object
    def update_urdf(self, object_file, object_name, collision_file=None, override=None, mass_center=None):
        # If no separate collision geometry is provided, use the object file
        if collision_file is None:
            collision_file = object_file

        # Update the filenames and object name
        new_urdf = copy.deepcopy(self.urdf_base)
        self.replace_urdf_attribute(new_urdf, './/visual/geometry/mesh', 'filename', object_file)
        self.replace_urdf_attribute(new_urdf, './/collision/geometry/mesh', 'filename', collision_file)
        new_urdf.attrib['name'] = object_name

        # Update the overrides
        if override is not None:
            for orverride_el in override:
                # Update attributes
                out_el_all = new_urdf.findall('.//' + orverride_el.tag)

                for out_el in out_el_all:

                    for key in orverride_el.attrib:
                        out_el.set(key, orverride_el.attrib[key])

                    # Remove fields that will be overwritten
                    for child in orverride_el:
                        el = out_el.find(child.tag)
                        if el is not None:
                            out_el.remove(el)
                    # Add updated feilds
                    out_el.extend(orverride_el)

        # Output the center of mass if provided
        if mass_center is not None:
            # Check if there's a geometry offset
            offset_ob = new_urdf.find('.//collision/origin')
            if offset_ob is not None:
                offset_str = offset_ob.attrib.get('xyz', '0 0 0')
                rot_str = offset_ob.attrib.get('rpy', '0 0 0')
                offset = self._str2list(offset_str)
                rpy = self._str2list(rot_str)
            else:
                offset = [0, 0, 0]
                rpy = [0, 0, 0]

            # Check if there's a scale factor and apply it
            scale_ob = new_urdf.find('.//collision/geometry/mesh')
            if scale_ob is not None:
                scale_str = scale_ob.attrib.get('scale', '1 1 1')
                scale = self._str2list(scale_str)
            else:
                scale = [1, 1, 1]

            for idx, axis in enumerate(mass_center):
                mass_center[idx] = -mass_center[idx] * scale[idx] + offset[idx]

            rot = Rotation.from_euler('xyz', rpy)
            rot_matrix = rot.as_matrix()
            mass_center = np.matmul(rot_matrix, np.vstack(np.asarray(mass_center))).squeeze()

            self.replace_urdf_attributes(
                new_urdf,
                './/visual/origin',
                {'xyz': self._list2str(mass_center), 'rpy': self._list2str(rpy)}
            )
            self.replace_urdf_attributes(
                new_urdf,
                './/collision/origin',
                {'xyz': self._list2str(mass_center), 'rpy': self._list2str(rpy)}
            )

        return new_urdf

    # Save a URDF to a file
    def save_urdf(self, new_urdf, filename, overwrite=False):
        out_file = os.path.join(self.object_folder, filename)

        # Do not overwrite the file unless the option is True
        if os.path.exists(out_file) and not overwrite:
            return

        # Save the file
        mydata = ET.tostring(new_urdf)
        with open(out_file, "wb") as f:
            f.write(mydata)

    # Build a URDF from an object file
    def build_urdf(
        self, filename, output_folder=None,
        force_overwrite=False, decompose_concave=False, force_decompose=False,
        center='mass', **kwargs
    ):

        # If no output folder is specified, use the base object folder
        if output_folder is None:
            output_folder = self.object_folder

        # Generate a relative path from the output folder to the geometry files
        filename = os.path.abspath(filename)
        common = os.path.commonprefix([output_folder, filename])
        rel = os.path.join(filename.replace(common, ''))
        if rel[0] == os.path.sep:
            rel = rel[1:]
        name = rel.split(os.path.sep)[0]
        rel = rel.replace(os.path.sep, '/')

        file_name_raw, file_extension = os.path.splitext(filename)

        # If an override file exists, include its data in the URDF
        override_file = filename.replace(file_extension, '.ovr')
        if os.path.exists(override_file):
            overrides = self._read_xml(override_file)
        else:
            overrides = None

        # Calculate the center of mass
        if center == 'mass':
            mass_center = self.get_center_of_mass(filename)

        elif center == 'geometric':
            mass_center = self.get_geometric_center(filename)

        elif center in ['top', 'bottom', 'xy_pos', 'xy_neg', 'xz_pos', 'xz_neg', 'yz_pos', 'yz_neg']:
            mass_center = self.get_face(filename, center)

        else:
            mass_center = None

        mesh = trimesh.load(filename)
        print(mesh.bounds)

        # If the user wants to run convex decomposition on concave objects, do it.
        if decompose_concave:
            if file_extension == '.stl':
                obj_filename = self.save_to_obj(filename)
                visual_file = rel.replace(file_extension, '.obj')
            elif file_extension == '.obj':
                obj_filename = filename
                visual_file = rel
            else:
                raise ValueError("Your filetype needs to be an STL or OBJ to perform concave decomposition")

            outfile = obj_filename.replace('.obj', '_' + self.suffix + '.obj')
            collision_file = visual_file.replace('.obj', '_' + self.suffix + '.obj')

            # Only run a decomposition if one does not exist, or if the user forces an overwrite
            if not os.path.exists(outfile) or force_decompose:
                self.do_vhacd(obj_filename, outfile, **kwargs)

            urdf_out = self.update_urdf(visual_file, name, collision_file=collision_file, override=overrides, mass_center=mass_center)
        else:
            urdf_out = self.update_urdf(rel, name, override=overrides, mass_center=mass_center)

        self.save_urdf(urdf_out, name + '.urdf', force_overwrite)

    # Build the URDFs for all objects in your library.
    def build_library(self, **kwargs):
        print("\nFOLDER: %s" % (self.object_folder))

        # Get all OBJ files
        obj_files = self._get_files_recursively(self.object_folder, filter_extension='.obj', exclude_suffix=self.suffix)
        stl_files = self._get_files_recursively(self.object_folder, filter_extension='.stl', exclude_suffix=self.suffix)

        obj_folders = []
        for root, _, full_file in obj_files:
            obj_folders.append(root)
            self.build_urdf(full_file, **kwargs)

            common = os.path.commonprefix([self.object_folder, full_file])
            rel = os.path.join(full_file.replace(common, ''))
            print('\tBuilding: %s' % (rel))

        for root, _, full_file in stl_files:
            if root not in obj_folders:
                self.build_urdf(full_file, **kwargs)

                common = os.path.commonprefix([self.object_folder, full_file])
                rel = os.path.join(full_file.replace(common, ''))
                print('Building: %s' % (rel))
concepts/simulator/pybullet/urdf_utils/obj2urdf.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : world.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/17/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import time
import inspect
import collections
import six
from typing import Any, Optional, Union, Iterable, Tuple, List, Dict

import numpy as np
import open3d as o3d
import pybullet as p

from concepts.simulator.pybullet.camera import CameraConfig
from concepts.simulator.pybullet.rotation_utils import quat_mul, quat_conjugate, rotate_vector_batch, compose_transformation

__all__ = ['BodyState', 'JointState', 'LinkState', 'DebugCameraState', 'JointInfo', 'ContactInfo', 'BulletSaver', 'GroupSaver', 'BodyStateSaver', 'JointStateSaver', 'BodyFullStateSaver', 'WorldSaver', 'BulletWorld']


class _NameToIdentifier(object):
    def __init__(self):
        self.string_to_int = dict()
        self.int_to_string = dict()

    def add_int_to_string(self, i, s):
        self.int_to_string[i] = s
        self.string_to_int[s] = i

    def __getitem__(self, item):
        if isinstance(item, six.string_types):
            return self.string_to_int[item]
        return self.int_to_string[item]

    def __len__(self):
        return len(self.int_to_string)

    def __iter__(self):
        yield from self.int_to_string.items()


GlobalIdentifier = collections.namedtuple('GlobalIdentifier', ['type', 'body_id', 'joint_or_link_id'])
JointIdentifier = collections.namedtuple('JointIdentifier', ['body_id', 'joint_id'])
LinkIdentifier = collections.namedtuple('LinkIdentifier', ['body_id', 'link_id'])


class BodyState(collections.namedtuple('_BodyState', ['position', 'orientation', 'linear_velocity', 'angular_velocity'])):
    @property
    def pos(self):
        return self.position

    @property
    def quat_xyzw(self):
        return self.orientation

    @property
    def quat_wxyz(self):
        return (self.orientation[3], self.orientation[0], self.orientation[1], self.orientation[2])

    def get_transformation(self) -> Tuple[np.ndarray, np.ndarray]:
        return self.position, self.orientation

    def get_7dpose(self) -> np.ndarray:
        return np.concatenate([self.position, self.orientation])


class JointState(collections.namedtuple('_JointState', ['position', 'velocity'])):
    pass


class LinkState(collections.namedtuple('_LinkState', ['position', 'orientation'])):
    @property
    def pos(self):
        return self.position

    @property
    def quat_xyzw(self):
        return self.orientation

    @property
    def quat_wxyz(self):
        return (self.orientation[3], self.orientation[0], self.orientation[1], self.orientation[2])

    def get_transformation(self) -> Tuple[np.ndarray, np.ndarray]:
        return np.array(self.position), np.array(self.orientation)

    def get_7dpose(self) -> np.ndarray:
        return np.concatenate([self.position, self.orientation])


class DebugCameraState(collections.namedtuple('_DebugCameraState', [
    'width',
    'height',
    'viewMatrix',
    'projectionMatrix',
    'cameraUp',
    'cameraForward',
    'horizontal',
    'vertical',
    'yaw',
    'pitch',
    'distance',
    'target',
])):

    @property
    def dist(self) -> float:
        return self.distance


class JointInfo(collections.namedtuple('_JointInfo', [
    'joint_index',
    'joint_name',
    'joint_type',
    'qindex',
    'uindex',
    'flags',
    'joint_damping',
    'joint_friction',
    'joint_lower_limit',
    'joint_upper_limit',
    'joint_max_force',
    'joint_max_velocity',
    'link_name',
    'joint_axis',
    'parent_frame_pos',
    'parent_frame_orn',
    'parent_index',
])):
    pass


class CollisionShapeData(collections.namedtuple('_CollisionShapeInfo', [
    'object_unique_id',
    'link_index',
    'shape_type',
    'dimensions',
    'filename',
    'local_pos',
    'local_orn',
    'world_pos',  # the position relative to the world frame
    'world_orn',  # the orientation relative to the world frame
])):
    pass


class ContactInfo(collections.namedtuple('_ContactInfo', [
    'world',
    'contact_flag',
    'body_a',
    'body_b',
    'link_a',
    'link_b',
    'position_on_a',
    'position_on_b',
    'contact_normal_on_b',
    'contact_distance',
    'contact_normal_force',
    'lateral_friction_1',
    'lateral_friction_dir1',
    'lateral_friction_2',
    'lateral_friction_dir2',
])):
    @property
    def body_a_name(self):
        return self.world.body_names[self.body_a]

    @property
    def link_a_name(self):
        return self.world.link_names[self.body_a, self.link_a]

    @property
    def body_b_name(self):
        return self.world.body_names[self.body_b]

    @property
    def link_b_name(self):
        return self.world.link_names[self.body_b, self.link_b]

    @property
    def a_name(self):
        if self.link_a != -1:
            return '@link/' + self.link_a_name
        else:
            return '@body/' + self.body_a_name

    @property
    def b_name(self):
        if self.link_b != -1:
            return '@link/' + self.link_b_name
        else:
            return '@body/' + self.body_b_name

    def hash(self):
        return (self.body_a, self.body_b, self.link_a, self.link_b)


class ConstraintInfo(collections.namedtuple('_ConstraintInfo', [
    'parent_body',
    'parent_joint',
    'child_body',
    'child_link',
    'constraint_type',
    'joint_axis',
    'parent_frame_pos',
    'child_frame_pos',
    'parent_frame_orn',
    'child_frame_orn',
    'max_force',
    'gear_ratio',
    'gear_aux_link',
    'relative_position_target',
    'erp'
])):
    pass


class BulletSaver(object):
    def __init__(self, world: 'BulletWorld'):
        self.world = world

    def save(self):
        pass

    def restore(self):
        raise NotImplementedError()

    def __enter__(self):
        self.save()
        return self

    def __exit__(self, type, value, traceback):
        self.restore()


class GroupSaver(BulletSaver):
    def __init__(self, savers: Iterable[BulletSaver]):
        savers = tuple(savers)
        assert len(savers) > 0, 'savers should not be empty'
        super().__init__(savers[0].world)
        self.savers = savers

    def save(self):
        for saver in self.savers:
            saver.save()

    def restore(self):
        for saver in self.savers:
            saver.restore()


class BodyStateSaver(BulletSaver):
    def __init__(self, world: 'BulletWorld', body_id: int, state: Optional[BodyState] = None, save: bool = True):
        super().__init__(world)
        self.body_id = body_id
        self.body_name = self.world.body_names.int_to_string.get(body_id, None)

        if save:
            self.save(state)

    def save(self, state: Optional[BodyState] = None):
        if state is None:
            state = self.world.get_body_state_by_id(self.body_id)
        self.state = state

    def restore(self):
        self.world.set_body_state_by_id(self.body_id, self.state)

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__, self.body_name or self.body_id)


class JointStateSaver(BulletSaver):
    def __init__(self, world: 'BulletWorld', body_id: int, joint_ids: Optional[List[int]] = None, states: Optional[List[JointState]] = None, save: bool = True):
        super().__init__(world)
        self.body_id = body_id
        self.body_name = self.world.body_names.int_to_string.get(body_id, None)

        if save:
            self.save(joint_ids, states)

    def save(self, joint_ids: Optional[List[int]] = None, states: Optional[List[JointState]] = None):
        if joint_ids is None:
            joint_ids = self.world.get_free_joints(self.body_id)
        self.joint_ids = joint_ids
        if states is None:
            states = [self.world.get_joint_state_by_id(self.body_id, joint_id) for joint_id in joint_ids]
        self.states = states

    def restore(self):
        for joint_id, state in zip(self.joint_ids, self.states):
            self.world.set_joint_state_by_id(self.body_id, joint_id, state)

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__, self.body_name or self.body_id)


class BodyFullStateSaver(BulletSaver):
    """Save and restore the full state of a body. That is, including both the state of the base link and all joints."""

    def __init__(self, world: 'BulletWorld', body_id: int, save: bool = True):
        super().__init__(world)
        self.body_id = body_id
        self.body_name = self.world.body_names.int_to_string.get(body_id, None)

        self.pose_saver = BodyStateSaver(world, body_id, save=save)
        self.joint_saver = JointStateSaver(world, body_id, save=save)

    def save(self, body_state: Optional[BodyState] = None, joint_ids: Optional[List[int]] = None, joint_states: Optional[List[JointState]] = None):
        self.pose_saver.save(body_state)
        self.joint_saver.save(joint_ids, joint_states)

    def restore(self):
        self.pose_saver.restore()
        self.joint_saver.restore()

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__, self.body_name or self.body_id)


class WorldSaver(BulletSaver):
    def __init__(self, world: 'BulletWorld', body_ids: Optional[List[int]] = None, save: bool = True):
        super().__init__(world)
        self.body_savers = list()
        self.body_ids = body_ids

        if save:
            self.save()

    def save(self):
        if len(self.body_savers) == 0:
            if self.body_ids is None:
                self.body_ids = list(self.world.body_names.int_to_string.keys())
            self.body_savers = [BodyFullStateSaver(self.world, body_id, save=True) for body_id in self.body_ids]
        else:
            for body_saver in self.body_savers:
                body_saver.save()

    def restore(self):
        for body_saver in self.body_savers:
            body_saver.restore()


class WorldSaverV2(BulletSaver):
    def __init__(self, world: 'BulletWorld', save: bool = True):
        super().__init__(world)

        self.saved_id = None
        if save:
            self.save()

    def save(self):
        self.saved_id = p.saveState(physicsClientId=self.world.client_id)

    def restore(self):
        if self.saved_id is None:
            raise RuntimeError('No state has been saved.')
        p.restoreState(self.saved_id, physicsClientId=self.world.client_id)
        p.removeState(self.saved_id, physicsClientId=self.world.client_id)
        self.saved_id = None


def _geometry_cache(type_id: int, geom_name_template: str):
    def wrapper(func):
        sig = inspect.signature(func)

        def wrapped_func(self: 'BulletWorld', *args, **kwargs):
            bound_args = sig.bind(self, *args, **kwargs)
            bound_args.apply_defaults()
            geom_name = geom_name_template.format(**bound_args.arguments)

            if (type_id, geom_name) in self.cached_geometries:
                return self.cached_geometries[type_id, geom_name]
            else:
                value = func(self, *args, **kwargs)
                self.cached_geometries[(type_id, geom_name)] = value
                return value
        return wrapped_func
    return wrapper


class BulletWorld(object):
    def __init__(self, client_id=None):
        self.client_id = client_id

        self.body_names = _NameToIdentifier()
        self.link_names = _NameToIdentifier()
        self.joint_names = _NameToIdentifier()
        self.global_names = _NameToIdentifier()
        self.body_base_link: Dict[str, str] = dict()
        self.body_groups: Dict[str, List[int]] = dict()

        self.cached_geometries: Dict[Tuple[int, str], Any] = dict()

        if self.client_id is not None:
            self.refresh_names()

    def set_client_id(self, client_id):
        self.client_id = client_id
        self.refresh_names()

    @property
    def nr_bodies(self):
        return len(self.body_names)

    @property
    def nr_joints(self):
        return len(self.joint_names)

    @property
    def nr_links(self):
        return len(self.link_names)

    def notify_update(self, body_id, body_name=None, group=None):
        body_info = p.getBodyInfo(body_id, physicsClientId=self.client_id)
        base_link_name = body_info[0].decode('utf-8')
        if body_name is None:
            body_name = body_info[1].decode('utf-8')
        self.body_base_link[body_name] = base_link_name

        self.body_names.add_int_to_string(body_id, body_name)
        self.link_names.add_int_to_string(LinkIdentifier(body_id, -1), body_name + '/' + base_link_name)
        self.global_names.add_int_to_string(GlobalIdentifier('body', body_id, -1), body_name)
        self.global_names.add_int_to_string(GlobalIdentifier('link', body_id, -1), body_name + '/' + base_link_name)

        for j in range(p.getNumJoints(body_id, physicsClientId=self.client_id)):
            joint_info = JointInfo(*p.getJointInfo(body_id, j, physicsClientId=self.client_id))
            joint_name = joint_info[1].decode('utf-8')
            link_name = joint_info[12].decode('utf-8')
            self.joint_names.add_int_to_string(JointIdentifier(body_id, j), joint_name)
            self.joint_names.add_int_to_string(JointIdentifier(body_id, j), body_name + '/' + joint_name)
            self.global_names.add_int_to_string(GlobalIdentifier('joint', body_id, j), joint_name)
            self.global_names.add_int_to_string(GlobalIdentifier('joint', body_id, j), body_name + '/' + joint_name)

            self.link_names.add_int_to_string(LinkIdentifier(body_id, j), link_name)
            self.link_names.add_int_to_string(LinkIdentifier(body_id, j), body_name + '/' + link_name)
            self.global_names.add_int_to_string(GlobalIdentifier('link', body_id, j), link_name)
            self.global_names.add_int_to_string(GlobalIdentifier('link', body_id, j), body_name + '/' + link_name)

        if group is not None:
            if group not in self.body_groups:
                self.body_groups[group] = list()
            self.body_groups[group].append(body_id)

    def refresh_names(self):
        self.global_names = _NameToIdentifier()
        self.body_names = _NameToIdentifier()
        self.joint_names = _NameToIdentifier()
        self.link_names = _NameToIdentifier()
        self.body_base_link = dict()

        for i in range(p.getNumBodies(physicsClientId=self.client_id)):
            body_info = p.getBodyInfo(i, physicsClientId=self.client_id)
            base_link_name, body_name = body_info[0].decode('utf-8'), body_info[1].decode('utf-8')
            self.body_base_link[body_name] = base_link_name

            self.body_names.add_int_to_string(i, body_name)
            self.global_names.add_int_to_string(('body', i), body_name)
            self.link_names.add_int_to_string((i, -1), body_name + '/' + base_link_name)
            self.global_names.add_int_to_string(('link', i, -1), body_name + '/' + base_link_name)

        for body_id in range(self.nr_bodies):
            body_name = self.body_names[body_id]
            for j in range(p.getNumJoints(body_id, physicsClientId=self.client_id)):
                joint_info = JointInfo(*p.getJointInfo(body_id, j, physicsClientId=self.client_id))
                joint_name = joint_info[1].decode('utf-8')
                link_name = joint_info[12].decode('utf-8')
                self.joint_names.add_int_to_string((body_id, j), joint_name)
                self.joint_names.add_int_to_string((body_id, j), body_name + '/' + joint_name)
                self.global_names.add_int_to_string(('joint', body_id, j), joint_name)
                self.global_names.add_int_to_string(('joint', body_id, j), body_name + '/' + joint_name)

                self.link_names.add_int_to_string((body_id, j), link_name)
                self.link_names.add_int_to_string((body_id, j), body_name + '/' + link_name)
                self.global_names.add_int_to_string(('link', body_id, j), link_name)
                self.global_names.add_int_to_string(('link', body_id, j), body_name + '/' + link_name)

    def get_body_name(self, body_id: int) -> str:
        return self.body_names.int_to_string[body_id]

    def get_link_name(self, body_id: int, link_id: int) -> str:
        return self.link_names.int_to_string[(body_id, link_id)]

    def get_body_index(self, body_name: str) -> int:
        return self.body_names.string_to_int[body_name]

    def get_link_index(self, link_name: str) -> LinkIdentifier:
        return self.link_names.string_to_int[link_name]

    def get_link_index_with_body(self, body_id: int, link_name: str) -> int:
        body_name = self.body_names.int_to_string[body_id]
        return self.link_names.string_to_int[body_name + '/' + link_name].link_id

    def get_xpos(self, name, type=None):
        info = self.get_state(name, type)
        assert isinstance(info, (BodyState, LinkState))
        return info.position

    def get_xquat(self, name, type=None):
        info = self.get_state(name, type)
        assert isinstance(info, (BodyState, LinkState))
        return info.orientation

    def get_xmat(self, name, type=None):
        return p.getMatrixFromQuaternion(self.get_xquat(name, type=type), physicsClientId=self.client_id).reshape((3, 3))

    def get_qpos(self, name):
        info = self.get_joint_state(name)
        return info.position

    def get_qpos_by_id(self, body_id, joint_id):
        info = self.get_joint_state_by_id(body_id, joint_id)
        return info.position

    def set_qpos(self, name, qpos):
        return p.resetJointState(*self.joint_names[name], qpos, physicsClientId=self.client_id)

    def set_qpos_by_id(self, body_id, joint_id, qpos):
        return p.resetJointState(body_id, joint_id, qpos, physicsClientId=self.client_id)

    def get_batched_qpos(self, names, numpy=True):
        rv = [self.get_qpos(name) for name in names]
        if numpy:
            return np.array(rv)
        return rv

    def get_batched_qpos_by_id(self, body_id, joint_ids, numpy=True):
        rv = [self.get_qpos_by_id(body_id, joint_id) for joint_id in joint_ids]
        if numpy:
            return np.array(rv)
        return rv

    def set_batched_qpos(self, names, qpos):
        for name, q in zip(names, qpos):
            self.set_qpos(name, q)

    def set_batched_qpos_by_id(self, body_index, joint_indices, qpos):
        for index, q in zip(joint_indices, qpos):
            p.resetJointState(body_index, index, q, physicsClientId=self.client_id)

    def get_state(self, name: str, type: Optional[str] = None) -> Union[BodyState, LinkState, JointState]:
        if ':' in name:
            type, name = name.split(':')

        try:
            if type is None:
                record = self.global_names[name]
                type, args = record[0], record[1:]
            elif type == 'body':
                args = (self.body_names[type],)
            elif type == 'joint':
                args = self.joint_names[type]
            elif type == 'link':
                args = self.link_names[type]
            else:
                raise ValueError('Unknown object type: {}.'.format(type))
        except KeyError:
            raise ValueError('Unknown name: {}.'.format(name))

        if type == 'body':
            return self.get_body_state_by_id(*args)
        elif type == 'joint':
            return self.get_joint_state_by_id(*args)
        elif type == 'link':
            return self.get_link_state_by_id(*args)
        else:
            raise ValueError('Unknown object type: {}.'.format(type))

    def __getitem__(self, item):
        return self.get_state(item)

    def get_debug_camera(self) -> DebugCameraState:
        return DebugCameraState(*p.getDebugVisualizerCamera(physicsClientId=self.client_id))

    def set_debug_camera(self, distance: float, yaw: float, pitch: float, target: Tuple[float, float, float]):
        p.resetDebugVisualizerCamera(distance, yaw, pitch, target, physicsClientId=self.client_id)

    def change_visual_color(self, body_id, rgba, link_id=None):
        if link_id is None:
            for i_body_id, link_id in self.link_names.int_to_string:
                if i_body_id == body_id:
                    p.changeVisualShape(body_id, link_id, rgbaColor=rgba, physicsClientId=self.client_id)
        else:
            p.changeVisualShape(body_id, link_id, rgbaColor=rgba, physicsClientId=self.client_id)

    def get_body_state_by_id(self, body_id: int) -> BodyState:
        state = p.getBasePositionAndOrientation(body_id, physicsClientId=self.client_id)
        vel = p.getBaseVelocity(body_id, physicsClientId=self.client_id)
        return BodyState(np.array(state[0]), np.array(state[1]), np.array(vel[0]), np.array(vel[1]))

    def get_body_state(self, body_name: str) -> BodyState:
        return self.get_body_state_by_id(self.body_names[body_name])

    def set_body_state_by_id(self, body_id: int, state: BodyState):
        p.resetBasePositionAndOrientation(body_id, tuple(state.position), tuple(state.orientation), physicsClientId=self.client_id)
        p.resetBaseVelocity(body_id, tuple(state.linear_velocity), tuple(state.angular_velocity), physicsClientId=self.client_id)

    def set_body_state(self, body_name: str, state: BodyState):
        self.set_body_state_by_id(self.body_names[body_name], state)

    def set_body_state2_by_id(self, body_id: int, position: np.ndarray, orientation: np.ndarray):
        p.resetBasePositionAndOrientation(body_id, tuple(position), tuple(orientation), physicsClientId=self.client_id)

    def set_body_state2(self, body_name: str, position: np.ndarray, orientation: np.ndarray):
        self.set_body_state2_by_id(self.body_names[body_name], position, orientation)

    def get_joint_info_by_id(self, body_id: int, joint_id: int) -> JointInfo:
        info = p.getJointInfo(body_id, joint_id, physicsClientId=self.client_id)
        return JointInfo(*info)

    def get_joint_info_by_body(self, body_id: int) -> List[JointInfo]:
        return [self.get_joint_info_by_id(body_id, i) for i in range(p.getNumJoints(body_id, physicsClientId=self.client_id))]

    def get_joint_info(self, joint_name: str) -> JointInfo:
        return self.get_joint_info_by_id(*self.joint_names[joint_name])

    def get_joint_state_by_id(self, body_id: int, joint_id: int) -> JointState:
        state = p.getJointState(body_id, joint_id, physicsClientId=self.client_id)
        return JointState(state[0], state[1])

    def get_joint_state(self, joint_name: str) -> JointState:
        return self.get_joint_state_by_id(*self.joint_names[joint_name])

    def set_joint_state_by_id(self, body_id: int, joint_id: int, state: JointState):
        p.resetJointState(body_id, joint_id, state.position, state.velocity, physicsClientId=self.client_id)

    def set_joint_state(self, joint_name: str, state: JointState):
        self.set_joint_state_by_id(*self.joint_names[joint_name], state)

    def get_link_state_by_id(self, body_id: int, link_id: int, fk=False) -> LinkState:
        if link_id == -1:
            state = p.getBasePositionAndOrientation(body_id, physicsClientId=self.client_id)
            return LinkState(np.array(state[0]), np.array(state[1]))
        state = p.getLinkState(body_id, link_id, computeForwardKinematics=fk, physicsClientId=self.client_id)
        return LinkState(np.array(state[0]), np.array(state[1]))

    def get_link_state(self, link_name: str, fk=False) -> LinkState:
        return self.get_link_state_by_id(*self.link_names[link_name], fk=fk)

    def get_collision_shape_data(self, body_name: str) -> List[CollisionShapeData]:
        body_id = self.body_names[body_name]
        return self.get_collision_shape_data_by_id(body_id)

    def get_collision_shape_data_by_id(self, body_id: int) -> List[CollisionShapeData]:
        output_data = list()
        for i in range(-1, p.getNumJoints(body_id, physicsClientId=self.client_id)):
            link_state = self.get_link_state_by_id(body_id, i, fk=True)

            data_list = p.getCollisionShapeData(body_id, i, physicsClientId=self.client_id)
            for shape_data in data_list:
                world_pos, world_orn = compose_transformation(link_state.position, link_state.orientation, shape_data[5], shape_data[6])
                # world_pos, world_orn = p.multiplyTransforms(link_state.position, link_state.orientation, shape_data[5], shape_data[6], physicsClientId=self.client_id)
                output_data.append(CollisionShapeData(*shape_data, world_pos, world_orn))
        return output_data

    def get_free_joints(self, body_id: int) -> List[int]:
        """Get a list of indices corresponding to all non-fixed joints in the body."""
        all_joints = list()
        for joint_id in range(p.getNumJoints(body_id, physicsClientId=self.client_id)):
            if self.get_joint_info_by_id(body_id, joint_id).joint_type != p.JOINT_FIXED:
                all_joints.append(joint_id)
        return all_joints

    def get_constraint(self, constraint_id: int) -> ConstraintInfo:
        return ConstraintInfo(*p.getConstraintInfo(constraint_id, physicsClientId=self.client_id))

    def get_contact(self, a: Optional[Union[int, LinkIdentifier, str]] = None, b: Optional[Union[int, LinkIdentifier, str]] = None, update: bool = False) -> List[ContactInfo]:
        if update:
            p.performCollisionDetection(physicsClientId=self.client_id)

        kwargs = dict(physicsClientId=self.client_id)

        def update_kwargs(name, value):
            if value is not None:
                if isinstance(value, int):
                    kwargs['body' + name] = value
                elif isinstance(value, (tuple, list)):
                    kwargs['body' + name], kwargs['linkIndex' + name] = value
                else:
                    assert isinstance(value, six.string_types)
                    value = self.global_names[value]
                    if value[0] == 'body':
                        kwargs['body' + name] = value[1]
                    elif value[0] == 'link':
                        kwargs['body' + name], kwargs['linkIndex' + name] = value[1:]
                    else:
                        raise ValueError('get_contact API only allows the specification of body or link.')

        update_kwargs('A', a)
        update_kwargs('B', b)
        while True:
            contacts = p.getContactPoints(**kwargs)
            if contacts is not None:
                break
            # TODO(Jiayuan Mao @ 2023/04/13): sometimes PyBullet returns None, not sure why.
            time.sleep(0.001)
        return [ContactInfo(self, *c) for c in contacts]

    def update_contact(self):
        p.performCollisionDetection(physicsClientId=self.client_id)

    def is_collision_free(self, a: Union[int, LinkIdentifier, str], b: Union[int, LinkIdentifier, str]) -> bool:
        return len(self.get_contact(a, b)) == 0

    def save_world(self) -> WorldSaver:
        return WorldSaver(self)

    def save_world_v2(self) -> WorldSaverV2:
        return WorldSaverV2(self)

    def save_body(self, body_identifier: Union[str, int]) -> BodyFullStateSaver:
        if isinstance(body_identifier, int):
            return BodyFullStateSaver(self, body_identifier)
        else:
            return BodyFullStateSaver(self, self.body_names[body_identifier])

    def save_bodies(self, body_identifiers: List[Union[str, int]]) -> GroupSaver:
        return GroupSaver([self.save_body(b) for b in body_identifiers])

    def render_image(self, config: CameraConfig, image_size: Optional[Tuple[int, int]] = None) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        view_matrix, proj_matrix = config.get_view_and_projection_matricies(image_size)
        image_size = image_size if image_size is not None else config.image_size
        znear, zfar = config.zrange

        # Render with OpenGL camera settings.
        _, _, color, depth, segm = p.getCameraImage(
            width=image_size[1],
            height=image_size[0],
            viewMatrix=view_matrix,
            projectionMatrix=proj_matrix,
            shadow=1,
            flags=p.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX,
            renderer=p.ER_BULLET_HARDWARE_OPENGL
        )

        # Get color image.
        color_image_size = (image_size[0], image_size[1], 4)
        color = np.asarray(color, dtype=np.uint8).reshape(color_image_size)
        color = color[:, :, :3]  # remove alpha channel

        # Get depth image.
        depth_image_size = (image_size[0], image_size[1])
        zbuffer = np.asarray(depth).reshape(depth_image_size)
        depth = zbuffer
        # depth = (zfar + znear - (2. * zbuffer - 1.) * (zfar - znear))
        # depth = (2. * znear * zfar) / depth

        # Get segmentation image.
        segm = np.uint8(segm).reshape(depth_image_size)

        return color, depth, segm

    def get_pointcloud(self, body_id: int, points_per_geom: int = 1000, zero_center: bool = True) -> np.ndarray:
        """Get the point cloud of a body."""

        all_pcds = list()

        body_state = self.get_body_state_by_id(body_id)
        for shape_info in self.get_collision_shape_data_by_id(body_id):
            if shape_info.shape_type == p.GEOM_BOX:
                pcd = self.get_pointcloud_box(shape_info.dimensions, points_per_geom)
            # elif shape_info.shape_type == client.p.GEOM_SPHERE:
            #     return get_point_cloud_sphere(shape_info.dimensions, points_per_geom)
            elif shape_info.shape_type == p.GEOM_MESH:
                pcd = self.get_pointcloud_mesh(shape_info.filename, shape_info.dimensions[0], points_per_geom)
            else:
                raise ValueError(f'Unsupported shape type: {shape_info.shape_type}.')

            if zero_center:
                pos = shape_info.world_pos - body_state.pos
                orn = quat_mul(quat_conjugate(body_state.quat_xyzw), shape_info.world_orn)
            else:
                pos, orn = shape_info.world_pos, shape_info.world_orn

            pcd = self.transform_pcd(pcd, pos, orn)
            all_pcds.append(pcd)

        if len(all_pcds) == 0:
            return np.zeros((0, 3), dtype=np.float32)
        return np.concatenate(all_pcds, axis=0)

    @_geometry_cache(type_id=0, geom_name_template='box_{dimensions[0]}_{dimensions[1]}_{dimensions[2]}_{points_per_geom}')
    def get_pointcloud_box(self, dimensions, points_per_geom=1000) -> np.ndarray:
        """Get a point cloud for a box."""

        total_volume = dimensions[0] * dimensions[1] * dimensions[2]
        density = points_per_geom / total_volume
        linear_density = density ** (1. / 3.)

        x, y, z = np.meshgrid(
            np.linspace(-dimensions[0] / 2, dimensions[0] / 2, int(np.ceil(dimensions[0] * linear_density))),
            np.linspace(-dimensions[1] / 2, dimensions[1] / 2, int(np.ceil(dimensions[1] * linear_density))),
            np.linspace(-dimensions[2] / 2, dimensions[2] / 2, int(np.ceil(dimensions[2] * linear_density))),
        )
        return np.stack([x, y, z], axis=-1).reshape(-1, 3)

    @_geometry_cache(type_id=0, geom_name_template='mesh_{mesh_file}_{mesh_scale}_{points_per_geom}')
    def get_pointcloud_mesh(self, mesh_file, mesh_scale, points_per_geom=1000) -> np.ndarray:
        """Get a point cloud for a mesh."""
        mesh = o3d.io.read_triangle_mesh(mesh_file)
        pcd = mesh.sample_points_poisson_disk(points_per_geom)
        return np.asarray(pcd.points, dtype=np.float32) * mesh_scale

    def transform_pcd(self, raw_pcd, pos, quat_xyzw):
        return rotate_vector_batch(raw_pcd, quat_xyzw) + pos

    def get_mesh(self, body_id: int, zero_center: bool = True, verbose: bool = False, mesh_filename: Optional[str] = None, mesh_scale: float = 1.0) -> o3d.geometry.TriangleMesh:
        """Get the point cloud of a body."""

        base_mesh = o3d.geometry.TriangleMesh()
        body_state = self.get_body_state_by_id(body_id)
        for shape_info in self.get_collision_shape_data_by_id(body_id):
            if shape_info.shape_type == p.GEOM_BOX:
                mesh = self.get_mesh_box(shape_info.dimensions)
            # elif shape_info.shape_type == client.p.GEOM_SPHERE:
            #     return get_point_cloud_sphere(shape_info.dimensions, points_per_geom)
            elif shape_info.shape_type == p.GEOM_MESH:
                if mesh_filename is not None:
                    mesh = self.get_mesh_mesh(mesh_filename, mesh_scale)
                else:
                    mesh = self.get_mesh_mesh(shape_info.filename, shape_info.dimensions[0])
            else:
                raise ValueError(f'Unsupported shape type: {shape_info.shape_type}.')

            if zero_center:
                pos = shape_info.world_pos - body_state.pos
                orn = quat_mul(quat_conjugate(body_state.quat_xyzw), shape_info.world_orn)
            else:
                pos, orn = shape_info.world_pos, shape_info.world_orn

            if verbose:
                print(shape_info)
                print('Zero-center:', zero_center)
                print(f'Transforming the mesh: pos: {pos}, orn: {orn}')

            mesh = self.transform_mesh(mesh, pos, orn)
            base_mesh += mesh

        return base_mesh

    @_geometry_cache(type_id=1, geom_name_template='box_{dimensions[0]}_{dimensions[1]}_{dimensions[2]}')
    def get_mesh_box(self, dimensions) -> o3d.geometry.TriangleMesh:
        """Get a triangle mesh for a box primitive."""
        mesh = o3d.geometry.TriangleMesh.create_box(dimensions[0], dimensions[1], dimensions[2])
        mesh = mesh.translate([-dimensions[0] / 2, -dimensions[1] / 2, -dimensions[2] / 2])
        return mesh

    @_geometry_cache(type_id=1, geom_name_template='mesh_{mesh_file}_{mesh_scale}')
    def get_mesh_mesh(self, mesh_file, mesh_scale) -> o3d.geometry.TriangleMesh:
        """Get a triangle mesh for a mesh primitive."""
        mesh = o3d.io.read_triangle_mesh(mesh_file)
        mesh = mesh.scale(mesh_scale, np.array([0, 0, 0], dtype=np.float64))
        return mesh

    def transform_mesh(self, mesh: o3d.geometry.TriangleMesh, pos, quat_xyzw):
        rotation_matrix = o3d.geometry.get_rotation_matrix_from_quaternion([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]])
        mesh = o3d.geometry.TriangleMesh(mesh)
        mesh = mesh.rotate(rotation_matrix, center=(0, 0, 0))
        mesh = mesh.translate(pos)
        return mesh

    def transform_mesh2(self, mesh: o3d.geometry.TriangleMesh, pos: np.ndarray, quat_xyzw: np.ndarray, current_pos: np.ndarray, current_quat_xyzw: np.ndarray):
        rotation_matrix1 = o3d.geometry.get_rotation_matrix_from_quaternion([current_quat_xyzw[3], -current_quat_xyzw[0], -current_quat_xyzw[1], -current_quat_xyzw[2]])
        rotation_matrix2 = o3d.geometry.get_rotation_matrix_from_quaternion([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]])
        mesh = o3d.geometry.TriangleMesh(mesh)
        mesh.translate(-current_pos)
        mesh.rotate(rotation_matrix1, center=np.array([0, 0, 0]))
        mesh.rotate(rotation_matrix2, center=np.array([0, 0, 0]))
        mesh.translate(pos)
        return mesh

concepts/simulator/pybullet/world.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 06/21/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.



concepts/simulator/pymunk/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : body_utils.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 06/21/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import random
import pymunk
from typing import Optional, Iterable, Tuple

from .world import PymunkWorld
from .constants import color_consts


__all__ = [
    'get_screen_size', 'get_body_bbox',
    'add_ball', 'add_box',
    'add_shape_I', 'add_shape_L', 'add_shape_T', 'add_shape_C',
    'random_body_pos',
    'select_body'
]


def get_screen_size() -> Tuple[int, int]:
    import pygame
    screen_width, screen_height = pygame.display.get_window_size()
    return screen_width, screen_height


def get_body_bbox(body: pymunk.Body) -> pymunk.BB:
    bb = pymunk.BB()
    for s in body.shapes:
        bb = bb.merge(s.cache_bb())
    return bb


def add_ball(world: PymunkWorld, mass: float = 1.0, radius: float = 14.0, pos: Optional[Tuple[float, float]] = None, **kwargs):
    inertia = pymunk.moment_for_circle(mass, 0, radius, (0, 0))
    body = pymunk.Body(mass, inertia, body_type=pymunk.Body.DYNAMIC)
    shape = pymunk.Circle(body, radius, (0, 0))
    shape.color = color_consts.BLUE
    shape.friction = 1.0
    shape.elasticity = 0.8

    body.position = pos if pos is not None else random_body_pos(world, body)
    return world.add_body(body, **kwargs)


def add_box(world: PymunkWorld, mass: float = 1.0, half_length: float = 14.0, pos: Optional[Tuple[float, float]] = None, **kwargs):
    inertia = pymunk.moment_for_box(mass, (half_length * 2, half_length * 2))
    body = pymunk.Body(mass, inertia)
    shape = pymunk.Poly(body, [(-half_length, -half_length), (half_length, -half_length), (half_length, half_length), (-half_length, half_length)])
    shape.color = color_consts.BLUE
    shape.friction = 1.0
    shape.elasticity = 0.8

    body.position = pos if pos is not None else random_body_pos(world, body)
    return world.add_body(body, **kwargs)


def add_shape_I(world: PymunkWorld, length: float = 200, thickness: float = 3, pos: Optional[Tuple[float, float]] = None, **kwargs):
    body = pymunk.Body(body_type=pymunk.Body.KINEMATIC)
    shape = pymunk.Segment(body, (0, 0), (0, length), thickness)
    shape.color = color_consts.BLACK
    shape.friction = 1.0

    body.position = pos if pos is not None else random_body_pos(world, body)
    return world.add_body(body, **kwargs)


def add_shape_L(world: PymunkWorld, length1: float = 200, length2: float = 50, thickness: float = 3, pos: Optional[Tuple[float, float]] = None, **kwargs):
    body = pymunk.Body(body_type=pymunk.Body.KINEMATIC)
    shape1 = pymunk.Segment(body, (0, 0), (0, length1), thickness)
    shape1.color = color_consts.BLACK
    shape1.friction = 1.0
    shape2 = pymunk.Segment(body, (0, 0), (length2, 0), thickness)
    shape2.color = color_consts.BLACK
    shape2.friction = 1.0

    body.position = pos if pos is not None else random_body_pos(world, body)
    return world.add_body(body, **kwargs)


def add_shape_T(world: PymunkWorld, length1: float = 200, length2: float = 50, thickness: float = 3, pos: Optional[Tuple[float, float]] = None, dynamic=False, **kwargs):
    if dynamic:
        body = pymunk.Body(body_type=pymunk.Body.DYNAMIC, mass=1.0, moment=100)
    else:
        body = pymunk.Body(body_type=pymunk.Body.KINEMATIC)

    shape1 = pymunk.Segment(body, (0, 0), (0, length1), thickness)
    shape1.color = color_consts.BLACK
    shape1.friction = 1.0
    shape2 = pymunk.Segment(body, (-length2, 0), (+length2, 0), thickness)
    shape2.color = color_consts.BLACK
    shape2.friction = 1.0

    body.position = pos if pos is not None else random_body_pos(world, body)
    return world.add_body(body, **kwargs)


def add_shape_C(world: PymunkWorld, length1: float = 200, length2: float = 50, length3: float = 25, thickness: float = 3, pos: Optional[Tuple[float, float]] = None, **kwargs):
    body = pymunk.Body(body_type=pymunk.Body.KINEMATIC)
    shape1 = pymunk.Segment(body, (0, 0), (0, length1), thickness)
    shape1.color = color_consts.BLACK
    shape2 = pymunk.Segment(body, (-length2, 0), (+length2, 0), thickness)
    shape2.color = color_consts.BLACK
    shape3 = pymunk.Segment(body, (-length2, 0), (-length2, -length3), thickness)
    shape3.color = color_consts.BLACK
    shape4 = pymunk.Segment(body, (+length2, 0), (+length2, -length3), thickness)
    shape4.color = color_consts.BLACK

    body.position = pos if pos is not None else random_body_pos(world, body)
    return world.add_body(body, **kwargs)


def random_body_pos(world: PymunkWorld, body: pymunk.Body) -> Tuple[float, float]:
    bb = get_body_bbox(body)
    obj_width = int(bb.right - bb.left) // 2 + 1
    obj_height = int(bb.top - bb.bottom) // 2 + 1
    return (random.randint(obj_width, world.screen_width - obj_width), random.randint(obj_height, world.screen_height - obj_height))


def select_body(world: PymunkWorld, pos: Tuple[float, float], selectable_bodies: Optional[Iterable[pymunk.Body]] = None):
    if selectable_bodies is None:
        selectable_bodies = world.bodies

    for b in selectable_bodies:
        for s in b.shapes:
            info = s.point_query(pos)
            if info.distance < 0:
                return b
    return None


concepts/simulator/pymunk/body_utils.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : collision.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 07/01/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import pymunk
from typing import Optional, Tuple, List, Set, Dict
from .world import PymunkWorld

__all__ = ['SpacePositionRestorer', 'collision_test']


class SpacePositionRestorer(object):
    def __init__(self, world: PymunkWorld):
        self.world = world

        self.positions = dict()
        for body in self.world.bodies:
            self.positions[body] = body.position

    def restore(self):
        for body, position in self.positions.items():
            body.position = position


def collision_test(world: PymunkWorld, body_positions: Optional[Dict[str, Tuple[float, float]]] = None, bodies: Optional[List[pymunk.Body]] = None) -> Set[Tuple[str, str]]:
    if body_positions is None:
        body_positions = dict()

    restorer = SpacePositionRestorer(world)
    for body, position in body_positions.items():
        body = world.get_body_by_label(body)
        body.position = position

    ret = set()
    shape2body = dict()
    for body in world.bodies:
        for shape in body.shapes:
            shape2body[shape] = body

    if bodies is None:
        bodies = world.bodies
    for body in bodies:
        for shape in body.shapes:
            all_collisions = world.shape_query(shape)
            for other_shape, _ in all_collisions:
                if other_shape is not None and shape2body[other_shape] is not body:
                    ret.add((shape2body[shape].label, shape2body[other_shape].label))

    restorer.restore()
    return ret


concepts/simulator/pymunk/collision.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : constants.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 06/21/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Constants for pymunk.

The following is a table of the constants defined in pymunk.

.. csv-table::
    :header: "Name", "Value", "Description"

    "color_consts.RED", "255, 32, 32, 255", "Red color."
    "color_consts.BLACK", "0, 0, 0, 255", "Black color."
    "color_consts.BLUE", "32, 128, 255, 255", "Blue color."
"""

import argparse

color_consts = argparse.Namespace()
color_consts.RED = (255, 32, 32, 255)
color_consts.BLACK = (0, 0, 0, 255)
color_consts.BLUE = (32, 128, 255, 255)


concepts/simulator/pymunk/constants.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : world.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 06/21/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import pymunk
import random
from typing import Any, Optional, Tuple, Dict

from concepts.algorithm.configuration_space import BoxConfigurationSpace, CollisionFreeProblemSpace
from concepts.utils.range import Range

__all__ = ['PymunkWorld', 'PymunkSingleObjectConfigurationSpace', 'PymunkCollisionFreeProblemSpace']


class PymunkWorld(pymunk.Space):
    """A wrapper :class:`pymunk.Space` providing a manager for screen size and body labels."""

    def __init__(self, *args, screen_width: Optional[int] = None, screen_height: Optional[int] = None, **kwargs):
        super().__init__(*args, **kwargs)

        self.screen_width = screen_width
        self.screen_height = screen_height
        self.body2label = dict()
        self.label2body = dict()
        self.body_selectable = dict()
        self.selectable_bodies = list()

    def add_body(self, body, selectable=False, label=None) -> pymunk.Body:
        """Add a body to the space.

        Args:
            body: the body to be added.
            selectable: whether the body is selectable.
            label: the label of the body.

        Returns:
            The body added.
        """
        self.add(body, *body.shapes)
        self.label2body[label] = body
        self.body2label[body] = label
        self.body_selectable[body] = selectable

        body.label = label
        body.selectable = selectable
        if selectable:
            self.selectable_bodies.append(body)

        return body

    @property
    def bodies_extra(self):
        for body in self.bodies:
            yield body, self.body_selectable[body], self.body2label[body]

    def select_body(self, point: Tuple[float, float]) -> Optional[pymunk.Body]:
        import concepts.simulator.pymunk.body_utils as body_utils
        return body_utils.select_body(self, point, self.selectable_bodies)

    def get_body_by_label(self, name: str) -> pymunk.Body:
        assert name in self.label2body, f'Body "{name}" is not in the space.'
        return self.label2body[name]

    def random_body_pos(self, body: Optional[pymunk.Body] = None) -> Tuple[float, float]:
        import concepts.simulator.pymunk.body_utils as body_utils
        if body is None:
            return random.randint(0, self.screen_width), random.randint(0, self.screen_height)
        else:
            return body_utils.random_body_pos(self, body)

    def get_body_poses(self) -> Dict[str, Tuple[float, float]]:
        return {label: tuple(body.position) for label, body in self.label2body.items()}

    def get_body_states(self) -> Dict[str, Dict[str, Any]]:
        return {label: {
            'position': tuple(body.position),
            'velocity': tuple(body.velocity),
            'angle': body.angle,
            'angular_velocity': body.angular_velocity
        } for label, body in self.label2body.items()}

    def get_collision_free_pspace(self, controlling_object: str, ignore_collision_filter=None, max_diff=2) -> 'PymunkCollisionFreeProblemSpace':
        return PymunkCollisionFreeProblemSpace(
            PymunkSingleObjectConfigurationSpace(self, controlling_object, max_diff=max_diff),
            ignore_collision_filter=ignore_collision_filter
        )


class PymunkSingleObjectConfigurationSpace(BoxConfigurationSpace):
    def __init__(self, space: PymunkWorld, controlling_object: str, max_diff: float = 2):
        super().__init__([Range(0, space.screen_width), Range(space.screen_height // 2 - 200, space.screen_height)], max_diff)
        self.controlling_object = controlling_object
        self.pymunk_space = space


class PymunkCollisionFreeProblemSpace(CollisionFreeProblemSpace):
    def __init__(self, cspace: PymunkSingleObjectConfigurationSpace, ignore_collision_filter=None):
        super().__init__(cspace)
        self.ignore_collision_filter = ignore_collision_filter

    @property
    def space(self) -> PymunkWorld:
        return self.cspace.pymunk_space

    @property
    def controlling_object(self) -> str:
        return self.cspace.controlling_object

    def collide(self, configuration):
        from .collision import collision_test
        all_collisions = collision_test(self.space, {self.cspace.controlling_object: configuration}, bodies=[self.space.get_body_by_label(self.cspace.controlling_object)])
        if self.ignore_collision_filter is not None:
            all_collisions = [c for c in all_collisions if not self.ignore_collision_filter(*c)]
        return len(all_collisions) > 0


concepts/simulator/pymunk/world.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 04/29/2021
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

concepts/simulator/tk/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : drawing_window.py
# Author : Zhezheng Luo
# Email  : ezzluo@mit.edu
# Date   : 04/23/2021
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

# DrawingWindow.py - Michael Haimes, Leslie Kaelbling                      #
# Hacked by LPK to be independent of SoaR                                  #
############################################################################
#    MIT SoaR 2.0 - A Python abstraction layer for MobileRobots            #
#    Pioneer3DX robots, and a simulator to simulate their operation        #
#    in a python environment (for testing)                                 #
#                                                                          #
#    Copyright (C) 2006-2007 Michael Haimes <mhaimes@mit.edu>              #
#                                                                          #
#   This program is free software; you can redistribute it and/or modify   #
#   it under the terms of the GNU General Public License as published by   #
#   the Free Software Foundation; either version 2 of the License, or      #
#   (at your option) any later version.                                    #
#                                                                          #
#   This program is distributed in the hope that it will be useful,        #
#   but WITHOUT ANY WARRANTY; without even the implied warranty of         #
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the          #
#   GNU General Public License for more details.                           #
#                                                                          #
#   You should have received a copy of the GNU General Public License along#
#   with this program; if not, write to the Free Software Foundation, Inc.,#
#   51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.            #
############################################################################

__all__ = ['DrawingWindow']


from tkinter import *
import math


class Thing:
    pass


class DrawingWindow:
    def __init__(self, windowWidth, windowHeight, xMin, xMax, yMin, yMax, title, parent=None):
        self.title = title
        if parent:
            self.parent = parent
            self.top = parent.getWindow(title)
        else:
            self.tk = Tk()
            self.tk.withdraw()
            self.top = Toplevel(self.tk)
            self.top.wm_title(title)
            self.top.protocol('WM_DELETE_WINDOW', self.top.destroy)

        self.windowWidth = windowWidth
        self.windowHeight = windowHeight
        self.canvas = Canvas(self.top, width=self.windowWidth, height=self.windowHeight, background="white")
        self.canvas.pack()

        # multiply an input value by this to get pixels
        self.xScale = windowWidth / float(xMax - xMin)
        self.yScale = windowHeight / float(yMax - yMin)

        self.drawingframe = Thing()
        self.drawingframe.width_px = windowWidth

        self.xMin = xMin
        self.yMin = yMin
        self.xMax = xMax
        self.yMax = yMax

    def scaleX(self, x):
        return self.xScale * (x - self.xMin)

    def scaleY(self, y):
        return self.windowHeight - self.yScale * (y - self.yMin)

    def drawPoint(self, x, y, color="blue", radius=1):
        windowX = self.scaleX(x)
        windowY = self.scaleY(y)
        return self.canvas.create_rectangle(
            windowX - radius, windowY - radius, windowX + radius, windowY + radius, fill=color, outline=color
        )

    def drawRobotWithNose(self, x, y, theta, color="blue", size=6):
        rawx = math.cos(theta)
        rawy = math.sin(theta)
        hx, hy = 0.15, 0.0
        noseX = x + rawx * hx - rawy * hy
        noseY = y + rawy * hx + rawx * hy
        return self.drawRobot(x, y, noseX, noseY, color=color, size=size)

    def drawRobot(self, x, y, noseX, noseY, color="blue", size=8):
        windowX = self.scaleX(x)
        windowY = self.scaleY(y)
        hsize = int(size) / 2  # For once, we want the int division!
        return (
            self.canvas.create_rectangle(
                windowX - hsize, windowY - hsize, windowX + hsize, windowY + hsize, fill=color, outline=color
            ),
            self.canvas.create_line(
                windowX, windowY, self.scaleX(noseX), self.scaleY(noseY), fill=color, width=2, arrow="last"
            ),
        )

    def drawText(self, x, y, label):
        windowX = self.scaleX(x)
        windowY = self.scaleY(y)
        # font="Arial 20",fill="#ff0000"
        return self.canvas.create_text(windowX, windowY, text=label)

    def drawPolygon(self, verts, color="black", outline="black"):
        return self.canvas.create_polygon(
            [(self.scaleX(point.x), self.scaleY(point.y)) for point in verts], fill=color, outline=outline
        )

    def drawRect(self, x1, y1, x2, y2, color="black"):
        return self.canvas.create_rectangle(
            self.scaleX(x1), self.scaleY(y1), self.scaleX(x2), self.scaleY(y2), fill=color
        )

    def drawOval(self, x1, y1, x2, y2, color="black"):
        return self.canvas.create_oval(self.scaleX(x1), self.scaleY(y1), self.scaleX(x2), self.scaleY(y2), fill=color)

    def drawLineSeg(self, x1, y1, x2, y2, color="black", width=2):
        return self.canvas.create_line(
            self.scaleX(x1), self.scaleY(y1), self.scaleX(x2), self.scaleY(y2), fill=color, width=width
        )

    def drawUnscaledLineSeg(self, x1, y1, xproj, yproj, color="black", width=1):
        return self.canvas.create_line(
            self.scaleX(x1), self.scaleY(y1), self.scaleX(x1) + xproj, self.scaleY(y1) - yproj, fill=color, width=width
        )

    def drawUnscaledRect(self, x1, y1, xproj, yproj, color="black"):
        return self.canvas.create_rectangle(
            self.scaleX(x1) - xproj,
            self.scaleY(y1) + yproj,
            self.scaleX(x1) + xproj,
            self.scaleY(y1) - yproj,
            fill=color,
        )

    def drawLine(self, a, b, c, color="black"):
        if abs(b) < 0.001:
            startX = self.scaleX(-c / a)
            endX = self.scaleX(-c / a)
            startY = self.scaleY(self.yMin)
            endY = self.scaleY(self.yMax)
        else:
            startX = self.scaleX(self.xMin)
            startY = self.scaleY(-(a * self.xMin + c) / b)
            endX = self.scaleX(self.xMax)
            endY = self.scaleY(-(a * self.xMax + c) / b)
        return self.canvas.create_line(startX, startY, endX, endY, fill=color)

    def delete(self, thing):
        self.canvas.delete(thing)

    def clear(self):
        self.canvas.delete("all")

    def update(self):
        self.canvas.update()

concepts/simulator/tk/drawing_window.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : toy_1drobot_visualizer.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 04/29/2021
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from typing import Optional
from collections import namedtuple
from concepts.simulator.tk.drawing_window import DrawingWindow


class Toy1DRobotVState(namedtuple('_Toy1DRobotVState', ['robot', 'object', 'goal'])):
    pass


class Toy1DRobotVisualizer(object):
    def __init__(self, name='Toy1DRobot'):
        self.window = DrawingWindow(512 * 2, 116 * 2, -11, 11, 3, 8, title=name)

    def draw(
        self,
        current: Toy1DRobotVState,
        predicted: Optional[Toy1DRobotVState] = None,
        text: Optional[str] = None,
        wait_key=True
    ):
        current = Toy1DRobotVState(*[float(x) if x is not None else None for x in current])
        if predicted is not None:
            predicted = Toy1DRobotVState(*[float(x) if x is not None else None for x in predicted])

        self.window.clear()

        self.window.drawText(0, 7.5, label="Blue: Robot; Red: Object; Green: Goal")
        if text is None:
            self.window.drawText(0, 7, label=text)

        self.window.drawText(0, 6, label="Observation")
        x = current.robot
        self.window.drawRect(x - 0.3, 5, x + 0.3, 5 + 0.3 * 2, color='blue')
        x = current.object
        self.window.drawRect(x - 0.3, 5, x + 0.3, 5 + 0.3 * 2, color='red')
        x = current.goal
        self.window.drawRect(x - 1, 5 - 0.2, x + 1, 5, color='green')

        if predicted is not None:
            self.window.drawText(0, 4.5, label="Prediction")
            x = predicted.object
            self.window.drawRect(x - 0.3, 5 - 1.5, x + 0.3, 5 - 1.5 + 0.3 * 2, color='#FF8888')
            x = predicted.goal
            self.window.drawRect(x - 1, 5 - 1.5 - 0.2, x + 1, 5 - 1.5, color='#88FF88')

        if wait_key:
            char = input('Press <Enter> to continue...')
            if char == 'q':
                raise KeyboardInterrupt

concepts/simulator/tk/toy_1drobot_visualizer.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : __init__.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 10/25/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.


concepts/utils/__init__.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : interpolation_utils.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 05/26/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from typing import Optional

import numpy as np
from scipy.interpolate import CubicSpline
from scipy.optimize import minimize
from scipy.interpolate import interp1d


class LinearSpline(object):
    """Linear spline interpolation.

    This class is a wrapper of scipy.interpolate.interp1d that mimics the CubicSpline interface.
    """
    def __init__(self, xs: np.ndarray, ys: np.ndarray):
        assert len(xs) == len(ys)
        self.xs = np.asarray(xs, dtype=np.float32)
        self.ys = np.asarray(ys, dtype=np.float32)
        assert len(self.xs.shape) == 1, 'xs should be a 1D array.'
        self.interpolator = interp1d(xs, ys, axis=0, kind='linear', fill_value='extrapolate')

    def __call__(self, x: float):
        return self.interpolator(x)


def gen_cubic_spline(ys: np.ndarray):
    """Generate a cubic spline interpolation from an array of points."""
    xs = list(range(len(ys)))
    xs = np.asarray(xs, dtype=np.float32)
    ys = np.asarray(ys, dtype=np.float32)
    return CubicSpline(xs, ys)


def gen_linear_spline(ys):
    """Generate a linear spline interpolation from an array of points."""
    xs = list(range(len(ys)))
    return LinearSpline(xs, ys)


def project_to_cubic_spline(spl: CubicSpline, y: np.ndarray, ys: np.ndarray):
    y = np.asarray(y, dtype=np.float32)
    dspl = spl.derivative()

    def f(x):
        return ((spl(float(x)) - y) ** 2).sum()

    def df(x):
        return 2 * ((spl(float(x)) - y) * dspl(float(x))).sum()

    x0 = np.argmin(np.abs(ys - y).sum(axis=-1))
    # res = minimize(f, x0, jac=df, method='L-BFGS-B', bounds=[(0, len(ys) - 1)])
    res = minimize(f, x0, jac=df, method='L-BFGS-B', bounds=[(max(0, x0 - 1), min(len(ys) - 1, x0 + 1))])

    # print('estimated time x0', x0)
    # print('bound', [(max(0, x0 - 1), min(len(ys) - 1, x0 + 1))])
    # print(res)
    return float(res.x[0])


def get_next_target_cubic_spline(spl: CubicSpline, y: np.ndarray, step_size: float, ys: np.ndarray):
    x = project_to_cubic_spline(spl, y, ys)
    # print('current time x', x, '/', len(ys) - 1)
    x = min(max(x + step_size, 0), len(ys) - 1)
    # print('next time x', x, '/', len(ys) - 1)
    return spl(x)


def project_to_linear_spline(spl: LinearSpline, y: np.ndarray, minimum_x=None):
    y = np.asarray(y, dtype=np.float32)

    def f(x):
        y_prime = spl(x)
        return ((y - y_prime) ** 2).sum()

    x0 = np.argmin(np.abs(spl.ys - y).sum(axis=-1))
    min_x = max(0, x0 - 1)
    if minimum_x is not None:
        min_x = max(min_x, minimum_x)
    res = minimize(f, x0, bounds=[(min_x, max(min_x + 0.1, min(len(spl.ys) - 1, x0 + 1)))])
    return float(res.x[0])


def get_next_target_linear_spline(spl: LinearSpline, y: np.ndarray, step_size: float, minimum_x: Optional[float] = None):
    x = project_to_linear_spline(spl, y, minimum_x=minimum_x)
    x = min(max(x + step_size, 0), len(spl.ys) - 1)
    return x, spl(x)

concepts/utils/interpolation_utils.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : range.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 04/23/2021
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

# Below are classes for defining the ranges of values of the joints.
# These classes are generic, they don't depend on the particular
# robot.

import random

__all__ = ['Range']


class Range(object):
    """A range of values, handles wraparound."""

    def __init__(self, low, high, wrap_around=False):
        self.low = low
        self.high = high
        self.wrap_around = wrap_around

    def difference(self, config1, config2):  # difference (with wraparound)
        if self.wrap_around:
            if config1 < config2:
                if abs(config2 - config1) < abs((self.low - config1) + (config2 - self.high)):
                    return config2 - config1
                else:
                    return (self.low - config1) + (config2 - self.high)
            else:
                if abs(config2 - config1) < abs((self.high - config1) + (config2 - self.low)):
                    return config2 - config1
                else:
                    return (self.high - config1) + (config2 - self.low)
        else:
            return config2 - config1

    def in_range(self, value):  # test if in range
        if self.wrap_around:
            altered = value
            while not (self.low <= altered <= self.high):
                if altered > self.high:
                    altered -= self.high - self.low
                else:
                    altered += self.high - self.low
            return altered
        else:
            if self.contains(value):
                return value
            else:
                return None

    def sample(self):  # sample random value
        return (self.high - self.low) * random.random() + self.low

    def contains(self, x):
        return self.wrap_around or (self.low <= x <= self.high)

concepts/utils/range.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : rotationlib.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 12/04/2019
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""
Tools for converting between rotation representations.

Conventions
-----------

- All functions accept batches as well as individual rotations.
- All rotation conventions match respective MuJoCo defaults (e.g., quaternions use wxyz convention).
    Note that this is DIFFERENT from PyBullet (which uses xyzw).
- All angles are in radians.
- Matricies follow LR convention.
- Euler Angles are all relative with 'xyz' axes ordering.
- See specific representation for more information.

Representations
---------------

Euler
    There are many euler angle frames -- here we will strive to use the default
    in MuJoCo, which is eulerseq='xyz'.

    This frame is a relative rotating frame, about x, y, and z axes in order.
    Relative rotating means that after we rotate about x, then we use the
    new (rotated) y, and the same for z.

Quaternions
    These are defined in terms of rotation (angle) about a unit vector (x, y, z)
    We use the following <q0, q1, q2, q3> convention:

    .. code-block:: python

        q0 = cos(angle / 2)
        q1 = sin(angle / 2) * x
        q2 = sin(angle / 2) * y
        q3 = sin(angle / 2) * z

    This is also sometimes called qw, qx, qy, qz.

    Note that quaternions are ambiguous, because we can represent a rotation by
    angle about vector <x, y, z> and -angle about vector <-x, -y, -z>.
    To choose between these, we pick "first nonzero positive", where we
    make the first nonzero element of the quaternion positive.

    This can result in mismatches if you're converting an quaternion that is not
    "first nonzero positive" to a different representation and back.

Axis Angle
    .. warning::

        (Not currently implemented)
        These are very straightforward.  Rotation is angle about a unit vector.

XY Axes
    .. warning::

        (Not currently implemented)
        We are given x axis and y axis, and z axis is cross product of x and y.

Z Axis
    .. warning::

        This is NOT RECOMMENDED.  Defines a unit vector for the Z axis,
        but rotation about this axis is not well defined.

    Instead pick a fixed reference direction for another axis (e.g. X)
    and calculate the other (e.g. Y = Z cross-product X),
    then use XY Axes rotation instead.
SO3
    .. warning::

        (Not currently implemented)
        While not supported by MuJoCo, this representation has a lot of nice features.
        We expect to add support for these in the future.

TODOs/Missings
    - Rotation integration or derivatives (e.g. velocity conversions)
    - More representations (SO3, etc)
    - Random sampling (e.g. sample uniform random rotation)
    - Performance benchmarks/measurements
    - (Maybe) define everything as to/from matricies, for simplicity
"""

# Copyright (c) 2009-2017, Matthew Brett and Christoph Gohlke
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
#    IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
#    THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
#    PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
#    CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#    EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#    PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#    PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#    LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#    NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#    SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# Many methods borrow heavily or entirely from transforms3d:
# https://github.com/matthew-brett/transforms3d
# They have mostly been modified to support batched operations.

import itertools
import numpy as np

# For testing whether a number is close to zero
_FLOAT_EPS = np.finfo(np.float64).eps
_EPS4 = _FLOAT_EPS * 4.0


def as_rotation(r):
    """Convert a 3x3 matrix or a quaternion into a standard 3x3 matrix representation."""
    if isinstance(r, np.ndarray) and r.shape == (3, 3):
        return r
    if isinstance(r, np.ndarray) and r.shape == (4,):
        return quat2mat(r)
    raise ValueError('Invalid rotation: {}.'.format(r))


def euler2mat(euler):
    """ Convert Euler Angles to Rotation Matrix.  See rotation.py for notes """
    euler = np.asarray(euler, dtype=np.float64)
    assert euler.shape[-1] == 3, "Invalid shaped euler {}".format(euler)

    ai, aj, ak = -euler[..., 2], -euler[..., 1], -euler[..., 0]
    si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak)
    ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak)
    cc, cs = ci * ck, ci * sk
    sc, ss = si * ck, si * sk

    mat = np.empty(euler.shape[:-1] + (3, 3), dtype=np.float64)
    mat[..., 2, 2] = cj * ck
    mat[..., 2, 1] = sj * sc - cs
    mat[..., 2, 0] = sj * cc + ss
    mat[..., 1, 2] = cj * sk
    mat[..., 1, 1] = sj * ss + cc
    mat[..., 1, 0] = sj * cs - sc
    mat[..., 0, 2] = -sj
    mat[..., 0, 1] = cj * si
    mat[..., 0, 0] = cj * ci
    return mat


def euler2quat(euler):
    """ Convert Euler Angles to Quaternions.  See rotation.py for notes """
    euler = np.asarray(euler, dtype=np.float64)
    assert euler.shape[-1] == 3, "Invalid shape euler {}".format(euler)

    ai, aj, ak = euler[..., 2] / 2, -euler[..., 1] / 2, euler[..., 0] / 2
    si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak)
    ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak)
    cc, cs = ci * ck, ci * sk
    sc, ss = si * ck, si * sk

    quat = np.empty(euler.shape[:-1] + (4,), dtype=np.float64)
    quat[..., 0] = cj * cc + sj * ss
    quat[..., 3] = cj * sc - sj * cs
    quat[..., 2] = -(cj * ss + sj * cc)
    quat[..., 1] = cj * cs - sj * sc
    return quat


def mat2euler(mat):
    """ Convert Rotation Matrix to Euler Angles.  See rotation.py for notes """
    mat = np.asarray(mat, dtype=np.float64)
    assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)

    cy = np.sqrt(mat[..., 2, 2] * mat[..., 2, 2] + mat[..., 1, 2] * mat[..., 1, 2])
    condition = cy > _EPS4
    euler = np.empty(mat.shape[:-1], dtype=np.float64)
    euler[..., 2] = np.where(
        condition, -np.arctan2(mat[..., 0, 1], mat[..., 0, 0]), -np.arctan2(-mat[..., 1, 0], mat[..., 1, 1])
    )
    euler[..., 1] = np.where(condition, -np.arctan2(-mat[..., 0, 2], cy), -np.arctan2(-mat[..., 0, 2], cy))
    euler[..., 0] = np.where(condition, -np.arctan2(mat[..., 1, 2], mat[..., 2, 2]), 0.0)
    return euler


def mat2quat(mat):
    """ Convert Rotation Matrix to Quaternion.  See rotation.py for notes """
    mat = np.asarray(mat, dtype=np.float64)
    assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)

    Qxx, Qyx, Qzx = mat[..., 0, 0], mat[..., 0, 1], mat[..., 0, 2]
    Qxy, Qyy, Qzy = mat[..., 1, 0], mat[..., 1, 1], mat[..., 1, 2]
    Qxz, Qyz, Qzz = mat[..., 2, 0], mat[..., 2, 1], mat[..., 2, 2]
    # Fill only lower half of symmetric matrix
    K = np.zeros(mat.shape[:-2] + (4, 4), dtype=np.float64)
    K[..., 0, 0] = Qxx - Qyy - Qzz
    K[..., 1, 0] = Qyx + Qxy
    K[..., 1, 1] = Qyy - Qxx - Qzz
    K[..., 2, 0] = Qzx + Qxz
    K[..., 2, 1] = Qzy + Qyz
    K[..., 2, 2] = Qzz - Qxx - Qyy
    K[..., 3, 0] = Qyz - Qzy
    K[..., 3, 1] = Qzx - Qxz
    K[..., 3, 2] = Qxy - Qyx
    K[..., 3, 3] = Qxx + Qyy + Qzz
    K /= 3.0
    # TODO: vectorize this -- probably could be made faster
    q = np.empty(K.shape[:-2] + (4,))
    it = np.nditer(q[..., 0], flags=['multi_index'])
    while not it.finished:
        # Use Hermitian eigenvectors, values for speed
        vals, vecs = np.linalg.eigh(K[it.multi_index])
        # Select largest eigenvector, reorder to w,x,y,z quaternion
        q[it.multi_index] = vecs[[3, 0, 1, 2], np.argmax(vals)]
        # Prefer quaternion with positive w
        # (q * -1 corresponds to same rotation as q)
        if q[it.multi_index][0] < 0:
            q[it.multi_index] *= -1
        it.iternext()
    return q


def quat2euler(quat):
    """ Convert Quaternion to Euler Angles.  See rotation.py for notes """
    return mat2euler(quat2mat(quat))


def subtract_euler(e1, e2):
    assert e1.shape == e2.shape
    assert e1.shape[-1] == 3
    q1 = euler2quat(e1)
    q2 = euler2quat(e2)
    q_diff = quat_mul(q1, quat_conjugate(q2))
    return quat2euler(q_diff)


def quat2mat(quat):
    """ Convert Quaternion to Euler Angles.  See rotation.py for notes """
    quat = np.asarray(quat, dtype=np.float64)
    assert quat.shape[-1] == 4, "Invalid shape quat {}".format(quat)

    w, x, y, z = quat[..., 0], quat[..., 1], quat[..., 2], quat[..., 3]
    Nq = np.sum(quat * quat, axis=-1)
    s = 2.0 / Nq
    X, Y, Z = x * s, y * s, z * s
    wX, wY, wZ = w * X, w * Y, w * Z
    xX, xY, xZ = x * X, x * Y, x * Z
    yY, yZ, zZ = y * Y, y * Z, z * Z

    mat = np.empty(quat.shape[:-1] + (3, 3), dtype=np.float64)
    mat[..., 0, 0] = 1.0 - (yY + zZ)
    mat[..., 0, 1] = xY - wZ
    mat[..., 0, 2] = xZ + wY
    mat[..., 1, 0] = xY + wZ
    mat[..., 1, 1] = 1.0 - (xX + zZ)
    mat[..., 1, 2] = yZ - wX
    mat[..., 2, 0] = xZ - wY
    mat[..., 2, 1] = yZ + wX
    mat[..., 2, 2] = 1.0 - (xX + yY)
    return np.where((Nq > _FLOAT_EPS)[..., np.newaxis, np.newaxis], mat, np.eye(3))


def quat_conjugate(q):
    inv_q = -q
    inv_q[..., 0] *= -1
    return inv_q


def quat_mul(q0, q1):
    assert q0.shape == q1.shape
    assert q0.shape[-1] == 4
    assert q1.shape[-1] == 4

    w0 = q0[..., 0]
    x0 = q0[..., 1]
    y0 = q0[..., 2]
    z0 = q0[..., 3]

    w1 = q1[..., 0]
    x1 = q1[..., 1]
    y1 = q1[..., 2]
    z1 = q1[..., 3]

    w = w0 * w1 - x0 * x1 - y0 * y1 - z0 * z1
    x = w0 * x1 + x0 * w1 + y0 * z1 - z0 * y1
    y = w0 * y1 + y0 * w1 + z0 * x1 - x0 * z1
    z = w0 * z1 + z0 * w1 + x0 * y1 - y0 * x1
    q = np.array([w, x, y, z])
    if q.ndim == 2:
        q = q.swapaxes(0, 1)
    assert q.shape == q0.shape
    return q


def quat_rot_vec(q, v0):
    q_v0 = np.array([0, v0[0], v0[1], v0[2]])
    q_v = quat_mul(q, quat_mul(q_v0, quat_conjugate(q)))
    v = q_v[1:]
    return v


def quat_identity():
    return np.array([1, 0, 0, 0])


def axisangle2quat(axis, angle):
    quat = np.zeros(4, dtype='float64')
    quat[0] = np.cos(angle / 2)
    quat[1:] = np.sin(angle / 2) * axis
    return quat


def quat2axisangle(quat):
    theta = 0
    axis = np.array([0, 0, 1])
    sin_theta = np.linalg.norm(quat[1:])

    if sin_theta > 0.0001:
        theta = 2 * np.arcsin(sin_theta)
        theta *= 1 if quat[0] >= 0 else -1
        axis = quat[1:] / sin_theta

    return axis, theta


def euler2point_euler(euler):
    _euler = euler.copy()
    if len(_euler.shape) < 2:
        _euler = np.expand_dims(_euler, 0)
    assert _euler.shape[1] == 3
    _euler_sin = np.sin(_euler)
    _euler_cos = np.cos(_euler)
    return np.concatenate([_euler_sin, _euler_cos], axis=-1)


def point_euler2euler(euler):
    _euler = euler.copy()
    if len(_euler.shape) < 2:
        _euler = np.expand_dims(_euler, 0)
    assert _euler.shape[1] == 6
    angle = np.arctan(_euler[..., :3] / _euler[..., 3:])
    angle[_euler[..., 3:] < 0] += np.pi
    return angle


def quat2point_quat(quat):
    # Should be in qw, qx, qy, qz
    _quat = quat.copy()
    if len(_quat.shape) < 2:
        _quat = np.expand_dims(_quat, 0)
    assert _quat.shape[1] == 4
    angle = np.arccos(_quat[:, [0]]) * 2
    xyz = _quat[:, 1:]
    xyz[np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5] = (xyz / np.sin(angle / 2))[
        np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5
    ]
    return np.concatenate([np.sin(angle), np.cos(angle), xyz], axis=-1)


def point_quat2quat(quat):
    _quat = quat.copy()
    if len(_quat.shape) < 2:
        _quat = np.expand_dims(_quat, 0)
    assert _quat.shape[1] == 5
    angle = np.arctan(_quat[:, [0]] / _quat[:, [1]])
    qw = np.cos(angle / 2)

    qxyz = _quat[:, 2:]
    qxyz[np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5] = (qxyz * np.sin(angle / 2))[
        np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5
    ]
    return np.concatenate([qw, qxyz], axis=-1)


def normalize_angles(angles):
    """Puts angles in [-pi, pi] range."""
    angles = angles.copy()
    if angles.size > 0:
        angles = (angles + np.pi) % (2 * np.pi) - np.pi
        assert -np.pi - 1e-6 <= angles.min() and angles.max() <= np.pi + 1e-6
    return angles


def round_to_straight_angles(angles):
    """Returns closest angle modulo 90 degrees """
    angles = np.round(angles / (np.pi / 2)) * (np.pi / 2)
    return normalize_angles(angles)


def get_parallel_rotations():
    mult90 = [0, np.pi / 2, -np.pi / 2, np.pi]
    parallel_rotations = []
    for euler in itertools.product(mult90, repeat=3):
        canonical = mat2euler(euler2mat(euler))
        canonical = np.round(canonical / (np.pi / 2))
        if canonical[0] == -2:
            canonical[0] = 2
        if canonical[2] == -2:
            canonical[2] = 2
        canonical *= np.pi / 2
        if all([(canonical != rot).any() for rot in parallel_rotations]):
            parallel_rotations += [canonical]
    assert len(parallel_rotations) == 24
    return parallel_rotations

concepts/utils/rotationlib.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : typing_utils.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 07/18/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

"""Utilities for typing.

Table of Custom Types

.. csv-table::
    :header: "Type", "Description"

    ``Vec2f``, "A 2D vector of float."
    ``Vec2i``, "A 2D vector of int."
    ``Vec2``, "A 2D vector of float or int."
    ``BoardcastableVec2f``, "A 2D vector of float or a float."
    ``Vec3f``, "A 3D vector of float."
    ``Vec3i``, "A 3D vector of int."
    ``Vec3``, "A 3D vector of float or int."
    ``BoardcastableVec3f``, "A 3D vector of float or a float."
    ``Vec4f``, "A 4D vector of float."
    ``Vec4i``, "A 4D vector of int."
    ``Vec4``, "A 4D vector of float or int."
    ``Vec6f``, "A 6D vector of float."
    ``Vec6i``, "A 6D vector of int."
    ``Vec6``, "A 6D vector of float or int."
    ``Vec7f``, "A 7D vector of float."
    ``Vec7i``, "A 7D vector of int."
    ``Vec7``, "A 7D vector of float or int."
    ``Color``, "A color, can be a 3D or 4D vector of float or int."
    ``Pos2D``, "A 2D position, can be a 2D vector of float or int."
    ``Pos3D``, "A 3D position, can be a 3D vector of float or int."
    ``Pos6D``, "A 6D position (usually xyz + rpy), can be a 6D vector of float or int."
    ``Pos7D``, "A 7D position (usually xyz + quat), can be a 7D vector of float or int."
"""

import numpy as np
from typing import Union, Tuple, List


Vec2f = Union[Tuple[float, float], List[float], np.ndarray]
Vec2i = Union[Tuple[int, int], List[int], np.ndarray]
Vec2 = Union[Vec2f, Vec2i]

BraodcastableVec2f = Union[float, Vec2f]

Vec3f = Union[Tuple[float, float, float], List[float], np.ndarray]
Vec3i = Union[Tuple[int, int, int], List[int], np.ndarray]
Vec3 = Union[Vec3f, Vec3i]

BroadcastableVec3f = Union[float, Vec3f]

Vec4f = Union[Tuple[float, float, float, float], List[float], np.ndarray]
Vec4i = Union[Tuple[int, int, int, int], List[int], np.ndarray]
Vec4 = Union[Vec4f, Vec4i]

Vec6f = Union[Tuple[float, float, float, float, float, float], List[float], np.ndarray]
Vec6i = Union[Tuple[int, int, int, int, int, int], List[int], np.ndarray]
Vec6 = Union[Vec6f, Vec6i]

Vec7f = Union[Tuple[float, float, float, float, float, float, float], List[float], np.ndarray]
Vec7i = Union[Tuple[int, int, int, int, int, int, int], List[int], np.ndarray]
Vec7 = Union[Vec7f, Vec7i]

Color = Union[Vec3f, Vec3i, Vec4f, Vec4i]
Pos2D = Union[Vec2f, Vec2i]
Pos3D = Union[Vec3f, Vec3i]
Pos6D = Union[Vec6f, Vec6i]
Pos7D = Union[Vec7f, Vec7i]


concepts/utils/typing_utils.py
concepts-ai/Concepts
# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
import sys
from os.path import dirname

sys.path.insert(0, dirname(dirname(dirname(__file__))))


# -- Project information -----------------------------------------------------

project = 'Project Concept'
copyright = '2022, Jiayuan Mao'
author = 'Jiayuan Mao'


# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
    'sphinx.ext.autodoc',
    'sphinx.ext.autosummary',
    'sphinx.ext.intersphinx',
    'sphinx.ext.viewcode',
    'sphinx.ext.napoleon',
    'sphinx.ext.todo',
    'nbsphinx',
]

# Mappings for sphinx.ext.intersphinx. Projects have to have Sphinx-generated doc! (.inv file)
intersphinx_mapping = {
    'python': ('https://docs.python.org/', None),
    'numpy': ('http://docs.scipy.org/doc/numpy/', None),
    'torch': ('https://pytorch.org/docs/stable/', None),
    'jacinle': ('https://jacinle.jiayuanm.com/', None),
}

# autosummary
# autosummary_generate = True  # Turn on sphinx.ext.autosummary

# autodoc
autodoc_mock_imports = ['pyrealsense2']
autodoc_inherit_docstrings = True
autodoc_class_signature = 'separated'
# autodoc_member_order = 'bysource'
autodoc_member_order = 'groupwise'
autodoc_typehints = 'both'
autodoc_typehints_format = 'short'
autodoc_typehints_description_target = 'all'
autodoc_preserve_defaults = True
# We will use autosummary to generate per-function and per-class docs.
# autodoc_default_options = {'members': True, 'undoc-members': True, 'inherited-members': True, 'show-inheritance': True}

autoclass_content = "class"

add_module_names = False  # Remove namespaces from class/method signatures
set_type_checking_flag = True  # Enable 'expensive' imports for sphinx_autodoc_typehints
autodoc_typehints = "description"  # Sphinx-native method. Not as good as sphinx_autodoc_typehints

# napoleaon
napolean_use_rtype = False

# Add any paths that contain templates here, relative to this directory.
templates_path = ['_templates']

# -- Options for HTML output -------------------------------------------------

# Pydata theme
html_theme = "pydata_sphinx_theme"
# html_logo = "_static/logo-company.png"
html_theme_options = {
    'show_prev_next': False,
    "show_nav_level": 0
}
html_css_files = ['custom.css']
html_show_sourcelink = False  # Remove 'view source code' from top of page (for html, not python)

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ['_static']


docs/source/conf.py
concepts-ai/Concepts
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : setup.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/09/2023
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.
# Available at setup time due to pyproject.toml

# from pybind11.setup_helpers import Pybind11Extension, build_ext
from setuptools import setup

__version__ = "0.5.1"

# The main interface is through Pybind11Extension.
# * You can add cxx_std=11/14/17, and then build_ext can be removed.
# * You can set include_pybind11=false to add the include directory yourself,
#   say from a submodule.
#
# Note:
#   Sort input source files if you glob sources to ensure bit-for-bit
#   reproducible builds (https://github.com/pybind/python_example/pull/53)

# ext_modules = [
#     Pybind11Extension(
#         "concepts._C",
#         ["concepts/cc/main.cc"],
#         define_macros=[('VERSION_INFO', __version__)],
#     ),
# ]

with open("README.md", "r", encoding="utf-8") as fh:
    long_description = fh.read()

requirements = [
    'six',
    'Cython',
    'pkgconfig',
    'ipdb',
    'tabulate',
    'numpy',
    'scipy',
    'matplotlib',
    'torch',
    'torchvision',
    'tqdm',
    'lark',
    'nltk',
    'Pillow',
]

setup(
    name="concepts",
    version=__version__,
    author="Jiayuan Mao",
    author_email="maojiayuan@gmail.com",
    install_requires=requirements,
    url="https://concepts.jiayuanm.com",
    description="Concepts",
    long_description=long_description,
    # ext_modules=ext_modules,
    license='MIT',
    classifiers=[
        'Development Status :: 3 - Alpha',
        'Intended Audience :: Science/Research',
        'License :: OSI Approved :: MIT License',
        'Programming Language :: Python :: 3.10',
    ],
    # Currently, build_ext only provides an optional "highest supported C++
    # level" feature, but in the future it may provide more features.
    # cmdclass={"build_ext": build_ext},
    packages=['concepts'],
    package_data={},
    zip_safe=False,
    python_requires=">=3.10",
)

setup.py
concepts-ai/Concepts