text:
import abc

import pytest
import sympy

from comb_spec_searcher.utils import taylor_expand
from tilings import GriddedPerm, Tiling
from tilings.algorithms import (
    DatabaseEnumeration,
    LocalEnumeration,
    MonotoneTreeEnumeration,
)
from tilings.exception import InvalidOperationError

class CommonTest(abc.ABC):
    @abc.abstractmethod
    @pytest.fixture
    def enum_verified(self):
        raise NotImplementedError

    @abc.abstractmethod
    @pytest.fixture
    def enum_not_verified(self):
        raise NotImplementedError

    def test_verified(self, enum_verified, enum_not_verified):
        assert enum_verified.verified()
        assert not enum_not_verified.verified()

    @abc.abstractmethod
    def test_get_genf(self, enum_verified):
        raise NotImplementedError

    def test_get_genf_not_verified(self, enum_not_verified):
        with pytest.raises(InvalidOperationError):
            enum_not_verified.get_genf()

class TestLocalEnumeration(CommonTest):
    @pytest.fixture
    def enum_verified(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1, 2), ((0, 0),) * 3),
                GriddedPerm((0, 2, 1), ((1, 0),) * 3),
                GriddedPerm((0, 1, 2), ((1, 0),) * 3),
                GriddedPerm((0, 1), ((1, 1),) * 2),
            ],
            requirements=[
                [
                    GriddedPerm((0, 1), ((0, 0),) * 2),
                    GriddedPerm((1, 0), ((0, 0),) * 2),
                ],
                [GriddedPerm((0,), ((1, 0),))],
            ],
        )
        return LocalEnumeration(t)

    @pytest.fixture
    def enum_not_verified(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1, 2), ((0, 0),) * 3),
                GriddedPerm((0, 2, 1), ((1, 0),) * 3),
                GriddedPerm((0, 1, 2), ((1, 0),) * 3),
                GriddedPerm((0, 1), ((1, 1),) * 2),
                GriddedPerm((0, 1), ((0, 0), (1, 1))),
            ],
            requirements=[
                [
                    GriddedPerm((0, 1), ((0, 0),) * 2),
                    GriddedPerm((0, 1), ((1, 0),) * 2),
                ]
            ],
        )
        return LocalEnumeration(t)

    @pytest.fixture
    def onebyone_enum(self):
        return LocalEnumeration(Tiling.from_string("123"))

    @pytest.fixture
    def enum_no_req(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1, 2), ((0, 0),) * 3),
                GriddedPerm((0, 2, 1), ((1, 0),) * 3),
                GriddedPerm((0, 1, 2), ((1, 0),) * 3),
                GriddedPerm((0, 1), ((1, 1),) * 2),
            ],
            requirements=[
                [
                    GriddedPerm((0, 1), ((0, 0),) * 2),
                    GriddedPerm((0, 1), ((1, 0),) * 2),
                ]
            ],
        )
        return LocalEnumeration(t, no_req=True)

    def test_req_is_single_cell(self):
        assert LocalEnumeration._req_is_single_cell([GriddedPerm((0,), ((0, 1),))])
        assert LocalEnumeration._req_is_single_cell(
            [GriddedPerm((0, 1), ((0, 1), (0, 1)))]
        )
        assert not LocalEnumeration._req_is_single_cell(
            [GriddedPerm((0, 1), ((0, 0), (0, 1)))]
        )
        assert not LocalEnumeration._req_is_single_cell(
            [GriddedPerm((0,), ((0, 1),)), GriddedPerm((0,), ((1, 0),))]
        )
        assert LocalEnumeration._req_is_single_cell(
            [
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((1, 0), ((0, 0),) * 2),
            ]
        )
        assert not LocalEnumeration._req_is_single_cell(
            [
                GriddedPerm((0, 1), ((1, 0),) * 2),
                GriddedPerm((1, 0), ((0, 0),) * 2),
            ]
        )

    def test_verified(self, enum_verified, enum_not_verified):
        assert enum_verified.verified()
        assert not enum_not_verified.verified()

    def test_crossing_req_list(self):
        """
        This tiling is not local verified because of the requirement list in multiple
        cells.
        """
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 2, 1), ((0, 1),) * 3),
                GriddedPerm((0, 2, 1), ((1, 0),) * 3),
            ],
            requirements=[
                [
                    GriddedPerm((0,), ((0, 1),)),
                    GriddedPerm((2, 0, 1), ((1, 0),) * 3),
                ],
                [GriddedPerm((0,), ((1, 0),))],
            ],
        )
        assert not LocalEnumeration(t).verified()

    def test_get_genf(self, enum_verified):
        with pytest.raises(NotImplementedError):
            enum_verified.get_genf()

    def test_get_genf_not_verified(self, enum_not_verified):
        with pytest.raises(InvalidOperationError):
            enum_not_verified.get_genf()

    def test_1x1_verified(self, onebyone_enum):
        assert onebyone_enum.verified()

    def test_no_req_option(self, enum_no_req):
        assert not enum_no_req.verified()

class TestMonotoneTreeEnumeration(CommonTest):
    @pytest.fixture
    def enum_verified(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((0, 1),) * 2),
                GriddedPerm((0, 1), ((0, 2),) * 2),
                GriddedPerm((0, 1), ((2, 0),) * 2),
                GriddedPerm((0, 1, 2), ((1, 1),) * 3),
            ]
        )
        return MonotoneTreeEnumeration(t)

    @pytest.fixture
    def enum_not_verified(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((0, 1),) * 2),
                GriddedPerm((0, 1), ((0, 2),) * 2),
                GriddedPerm((0, 1), ((2, 0),) * 2),
                GriddedPerm((0, 1), ((2, 2),) * 2),
                GriddedPerm((0, 1, 2), ((1, 1),) * 3),
            ]
        )
        return MonotoneTreeEnumeration(t)

    @pytest.fixture
    def enum_with_list_req(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0), (0, 0))),
                GriddedPerm((0, 1), ((1, 0), (1, 0))),
            ],
            requirements=[
                [
                    GriddedPerm((0,), ((0, 0),)),
                    GriddedPerm((0,), ((1, 0),)),
                ]
            ],
        )
        return MonotoneTreeEnumeration(t)

    @pytest.fixture
    def onebyone_enum(self):
        return MonotoneTreeEnumeration(Tiling.from_string("123"))

    @pytest.fixture
    def enum_with_crossing(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((0, 1),) * 2),
                GriddedPerm((0, 1), ((0, 2),) * 2),
                GriddedPerm((0, 1), ((2, 0),) * 2),
                GriddedPerm((0, 1), ((0, 0), (0, 1))),
                GriddedPerm((0, 1, 2), ((1, 1),) * 3),
            ]
        )
        return MonotoneTreeEnumeration(t)

    def test_visited_cells_aligned(self, enum_verified):
        visited = {(1, 1), (0, 1)}
        assert sorted(enum_verified._visted_cells_aligned((0, 2), visited)) == [(0, 1)]

    def test_cell_tree_traversal(self, enum_verified):
        order = list(enum_verified._cell_tree_traversal((1, 1)))
        assert len(order) == 4
        assert (1, 1) not in order
        assert order[0] == (0, 1)
        assert order[3] == (2, 0)
        assert set(order[1:3]) == {(0, 0), (0, 2)}

    def test_not_verified(self, enum_with_list_req, onebyone_enum, enum_with_crossing):
        assert not enum_with_crossing.verified()
        assert not enum_with_list_req.verified()
        assert not onebyone_enum.verified()
        forest_tiling = Tiling(
            obstructions=[
                GriddedPerm((0,), ((0, 0),)),
                GriddedPerm((0,), ((1, 1),)),
                GriddedPerm((0,), ((2, 1),)),
                GriddedPerm((0, 1), ((1, 0), (1, 0))),
                GriddedPerm((0, 1), ((2, 0), (2, 0))),
                GriddedPerm((0, 1, 2), ((0, 1), (0, 1), (0, 1))),
            ],
            requirements=[[GriddedPerm((0,), ((0, 1),))]],
        )
        assert not MonotoneTreeEnumeration(forest_tiling).verified()

    def test_get_genf(self, enum_verified):
        x = sympy.Symbol("x")
        expected_gf = -(
            sympy.sqrt(
                -(4 * x ** 3 - 14 * x ** 2 + 8 * x - 1) / (2 * x ** 2 - 4 * x + 1)
            )
            - 1
        ) / (2 * x * (x ** 2 - 3 * x + 1))
        assert sympy.simplify(enum_verified.get_genf() - expected_gf) == 0
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((1, 0),) * 2),
            ]
        )
        enum_no_start = MonotoneTreeEnumeration(t)
        expected_gf = -1 / ((x - 1) * (x / (x - 1) + 1))
        assert sympy.simplify(enum_no_start.get_genf() - expected_gf) == 0

    def test_get_genf_simple(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((1, 0), ((1, 0),) * 2),
            ]
        )
        enum = MonotoneTreeEnumeration(t)
        print(t)
        assert enum.verified()
        assert sympy.simplify(enum.get_genf() - sympy.sympify("1/(1-2*x)")) == 0

    def test_with_finite_monotone_cell(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((1, 0), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((1, 0),) * 2),
                GriddedPerm((1, 0), ((1, 0),) * 2),
            ]
        )
        enum = MonotoneTreeEnumeration(t)
        print(t)
        assert enum.verified()
        assert enum.get_genf().expand() == sympy.sympify("1+2*x+2*x**2")

    def test_with_finite_monotone_cell2(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((1, 0), ((0, 1),) * 2),
                GriddedPerm((0, 1), ((0, 1),) * 2),
                GriddedPerm((1, 0), ((1, 1),) * 2),
            ]
        )
        enum = MonotoneTreeEnumeration(t)
        print(t)
        assert enum.verified()
        assert (
            sympy.sympify("x/(1-x)**4 + 1/(1-x)**2") - enum.get_genf()
        ).simplify() == 0

    def test_interleave_fixed_length(self, enum_verified):
        track_var = MonotoneTreeEnumeration._tracking_var
        cell_var = enum_verified._cell_variable((1, 0))
        dummy_var = enum_verified._cell_variable((0, 0))
        x = sympy.var("x")
        F = x ** 8 * track_var ** 3 * dummy_var ** 3
        assert (
            enum_verified._interleave_fixed_length(F, (1, 0), 1)
            == 4 * x ** 9 * dummy_var ** 3 * cell_var ** 1
        )
        assert (
            enum_verified._interleave_fixed_length(F, (1, 0), 3)
            == 20 * x ** 11 * dummy_var ** 3 * cell_var ** 3
        )
        assert (
            enum_verified._interleave_fixed_length(F, (1, 0), 0)
            == x ** 8 * dummy_var ** 3
        )

    def test_interleave_fixed_lengths(self, enum_verified):
        track_var = MonotoneTreeEnumeration._tracking_var
        cell_var = enum_verified._cell_variable((1, 0))
        dummy_var = enum_verified._cell_variable((0, 0))
        x = sympy.var("x")
        F = x ** 8 * track_var ** 3 * dummy_var ** 3
        assert (
            enum_verified._interleave_fixed_lengths(F, (1, 0), 1, 1)
            == 4 * x ** 9 * dummy_var ** 3 * cell_var ** 1
        )
        assert (
            enum_verified._interleave_fixed_lengths(F, (1, 0), 3, 3)
            == 20 * x ** 11 * dummy_var ** 3 * cell_var ** 3
        )
        assert (
            enum_verified._interleave_fixed_lengths(F, (1, 0), 0, 0)
            == x ** 8 * dummy_var ** 3
        )
        assert (
            enum_verified._interleave_fixed_lengths(F, (1, 0), 0, 2)
            == x ** 8 * dummy_var ** 3
            + 4 * x ** 9 * dummy_var ** 3 * cell_var ** 1
            + 10 * x ** 10 * dummy_var ** 3 * cell_var ** 2
        )
        assert (
            enum_verified._interleave_fixed_lengths(F, (1, 0), 1, 3)
            == 4 * x ** 9 * dummy_var ** 3 * cell_var ** 1
            + 10 * x ** 10 * dummy_var ** 3 * cell_var ** 2
            + 20 * x ** 11 * dummy_var ** 3 * cell_var ** 3
        )

    def test_genf_with_req(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((1, 0),) * 2),
            ],
            requirements=[
                [GriddedPerm((1, 0), ((0, 0),) * 2)],
                [GriddedPerm((0,), ((1, 0),))],
            ],
        )
        enum = MonotoneTreeEnumeration(t)
        print(t)
        assert enum.verified()
        genf = enum.get_genf().expand()
        terms = [0, 0, 0, 3, 10, 25, 56, 119, 246, 501, 1012]
        assert taylor_expand(genf) == terms

    def test_genf_with_big_finite_cell(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((1, 0),) * 2),
                GriddedPerm((3, 2, 1, 0), ((0, 0),) * 4),
                GriddedPerm((3, 2, 1, 0), ((1, 0),) * 4),
            ]
        )
        enum = MonotoneTreeEnumeration(t)
        print(t)
        assert enum.verified()
        genf = enum.get_genf().expand()
        x = sympy.var("x")
        assert (
            genf
            == 1
            + 2 * x
            + 4 * x ** 2
            + 8 * x ** 3
            + 14 * x ** 4
            + 20 * x ** 5
            + 20 * x ** 6
        )

    def test_with_two_reqs(self):
        t = Tiling(
            obstructions=(
                GriddedPerm((0,), ((1, 1),)),
                GriddedPerm((0, 1), ((0, 0), (0, 0))),
                GriddedPerm((0, 1), ((0, 1), (0, 1))),
                GriddedPerm((0, 1), ((1, 0), (1, 0))),
                GriddedPerm((1, 0), ((0, 1), (0, 1))),
            ),
            requirements=(
                (GriddedPerm((0,), ((0, 0),)),),
                (GriddedPerm((0,), ((0, 1),)),),
            ),
        )
        enum = MonotoneTreeEnumeration(t)
        expected_enum = [0, 0, 2, 7, 19, 47, 111, 255, 575, 1279, 2815]
        assert enum.verified()
        assert taylor_expand(enum.get_genf()) == expected_enum

    def test_corner(self):
        t = Tiling(
            obstructions=(
                GriddedPerm((0,), ((1, 1),)),
                GriddedPerm((0, 1), ((0, 0), (0, 0))),
                GriddedPerm((0, 1), ((0, 1), (0, 1))),
                GriddedPerm((0, 1), ((1, 0), (1, 0))),
            ),
            requirements=((GriddedPerm((0,), ((0, 0),)),),),
        )
        enum = MonotoneTreeEnumeration(t)
        expected_enum = [0, 1, 5, 17, 50, 138, 370, 979, 2575, 6755, 17700]
        assert enum.verified()
        assert taylor_expand(enum.get_genf()) == expected_enum

class TestDatabaseEnumeration(CommonTest):
    @pytest.fixture
    def enum_verified(self):
        t = Tiling.from_string("123_132_231")
        return DatabaseEnumeration(t)

    @pytest.fixture
    def enum_not_verified(self):
        t = Tiling.from_string("1324")
        return DatabaseEnumeration(t)

    def test_get_genf(self, enum_verified):
        assert enum_verified.get_genf() == sympy.sympify(
            "(x**2 - x + 1)/(x**2 - 2*x + 1)"
        )

    @pytest.mark.slow
    def test_load_verified_tilings(self):
        DatabaseEnumeration.load_verified_tiling()
        assert DatabaseEnumeration.all_verified_tilings
        sample = next(iter(DatabaseEnumeration.all_verified_tilings))
        Tiling.from_bytes(sample)

    def test_verification_with_cache(self):
        t = Tiling.from_string("123_132_231")
        DatabaseEnumeration.all_verified_tilings = frozenset()
        assert DatabaseEnumeration(t).verified()
        DatabaseEnumeration.all_verified_tilings = frozenset([1, 2, 3, 4])
        assert not DatabaseEnumeration(t).verified()
        DatabaseEnumeration.all_verified_tilings = frozenset([t.to_bytes()])
        assert DatabaseEnumeration(t).verified()

max_stars_repo_name: PermutaTriangle/Tilings

max_issues_repo_name: PermutaTriangle/Tilings

max_forks_repo_name: PermutaTriangle/Tilings

reason: import sympy


text:
from copy import copy, deepcopy
from datetime import datetime
import numpy as np

from opentamp.src.policy_hooks.sample import Sample
from opentamp.src.policy_hooks.utils.policy_solver_utils import *

class MixedPolicy:
    def __init__(self, pol, dU, action_inds, state_inds, opt_traj, opt_strength):
        self.pol = pol
        self.dU = dU
        self.action_inds = action_inds
        self.state_inds = state_inds
        self.opt_traj = opt_traj
        self.opt_strength = opt_strength

    def act(self, X, O, t, noise):
        if self.opt_strength == 0: return self.pol.act(X, O, t, noise)
        opt_u = np.zeros(self.dU)
        for param, attr in self.action_inds:
            opt_u[self.action_inds[param, attr]] = self.opt_traj[t, self.action_inds[param, attr]]

        if self.opt_strength == 1: return opt_u

        return self.opt_strength * opt_u + (1 - self.opt_strength) * self.pol.act(X, O, t, noise)

class MCTSNode():
    def __init__(self, label, value, parent, num_tasks, prim_dims):
        self.label = label
        self.value = value
        self.num_tasks = num_tasks
        self.prim_dims = prim_dims
        self.prim_order = list(prim_dims.keys())
        self.num_prims = list(prim_dims.values())
        self.is_leaf = True
        self.children = [None for _ in range(num_tasks*np.prod([list(prim_dims.values())]))]
        self.parent = parent
        self.n_explored = 1.0
        self.sample_links = {}
        self.sample_to_traj = {}
        self.depth = parent.depth + 1 if parent != None else 0
        if parent is not None:
            parent.add_child(self)

    def is_leaf(self):
        return self.is_leaf()

    def get_task(self):
        return self.label[0]

    def get_prim(self, prim_name):
        return self.label[self.prim_order.index(prim_name)+1]

    def update_value(self, new_value):
        # self.value = (self.value*self.n_explored + new_value) / (self.n_explored + 1)
        # if new_value == 0:
        #     new_value = 1
        # else:
        #     new_value = 0

        self.value = (self.value*self.n_explored + new_value) / (self.n_explored + 1)
        self.n_explored += 1

    def get_child(self, label):
        ind = 0
        m = 1
        for i in range(len(self.prim_order), 0, -1):
            ind += m * label[i]
            m *= self.num_prims[i-1]
        ind += self.num_tasks * label[0]
        return self.children[ind]

    def add_child(self, child):
        ind = 0
        m = 1
        for i in range(len(self.prim_order), 0, -1):
            ind += m * child.label[i]
            m *= self.num_prims[i-1]
        ind += self.num_tasks * child.label[0]
        self.children[ind] = child
        self.is_leaf = False
        child.parent = self

    def get_explored_children(self):
        return [n for n in self.children if n is not None]

    def has_unexplored(self):
        for child in self.child:
            if child is None: return True
        return False

    def __repr__(self):
        return str(self.label)

class MCTS:
    def __init__(self, tasks, prim_dims, gmms, value_f, condition, agent, branch_factor, num_samples, num_distilled_samples, choose_next=None, soft_decision=1.0, C=2, max_depth=20, explore_depth=5, opt_strength=0.0):
        self.tasks = tasks
        self.prim_dims = prim_dims
        self.prim_order = list(prim_dims.keys())
        self.num_prims = list(prim_dims.values())
        self.root = MCTSNode((-1, -1, -1), 0, None, len(tasks), prim_dims)
        self.max_depth = max_depth
        self.explore_depth = explore_depth
        self.condition = condition
        self.agent = agent
        self.soft_decision = soft_decision
        self.C = C
        self.branch_factor = branch_factor
        self.num_samples = 1
        self.num_distilled_samples = num_distilled_samples
        self._choose_next = choose_next if choose_next != None else self._default_choose_next
        self._value_f = value_f
        # self.node_check_f = lambda n: n.value/n.n_explored+self.C*np.sqrt(np.log(n.parent.n_explored)/n.n_explored) if n != None else -np.inf

        self._opt_cache = {}
        self.opt_strength = opt_strength

        self.n_success = 0
        self.n_runs = 0

    def prob_func(self, prim_obs):
        prim_obs = prim_obs.reshape((1, -1))
        return self._prob_func(prim_obs)

    def value_func(self, obs):
        obs = obs.reshape((1, -1))
        return self._value_f(obs)

    def update_vals(self, path, success):
        node = self.root
        for step in path:
            node = node.get_child(*step)
            if node is None:
                node = MCTSNode(step,
                                int(success),
                                node,
                                len(self.tasks),
                                self.prim_dims)

            else:
                node.update_value(int(success))

    def node_check_f(self, label, state, parent):
        child = parent.get_child(label)
        sample = Sample(self.agent)
        sample.set_X(state.copy(), 0)
        # sample.set(TARGETS_ENUM, self.agent.target_vecs[self.condition].copy(), 0)
        sample.set(TRAJ_HIST_ENUM, np.array(self.agent.traj_hist).flatten(), 0)
        task_vec = np.zeros((len(self.tasks)), dtype=np.float32)
        task_vec[label[0]] = 1.
        sample.set(TASK_ENUM, task_vec, 0)
        for i in range(len(list(self.prim_dims.keys()))):
            prim = self.prim_order[i]
            vec = np.zeros(self.prim_dims[prim], dtype='float32')
            vec[label[i+1]] = 1.0
            sample.set(prim, vec, 0)

        self.agent.fill_sample(self.condition, sample, sample.get(STATE_ENUM, 0), 0, label, fill_obs=True)

        prim_obs = sample.get_prim_obs(t=0)
        val_obs = sample.get_val_obs(t=0)
        q_value = self.value_func(val_obs)[1] if child is None else child.value
        # policy_distrs = self.prob_func(prim_obs)
        # prob = np.product([policy_distrs[ind][label[ind]] for ind in range(len(label))])
        child_explored = child.n_explored if child is not None else 0
        # return self.value_func(val_obs)[1] + self.C * np.sqrt(parent.n_explored) / (1 + child_explored)
        # return q_value + self.C * self.value_func(obs)[1] / (1 + child_explored)
        return q_value + self.C * np.sqrt(np.log(parent.n_explored) / (1 + child_explored))

    def print_run(self, state, use_distilled=True):
        path = self.simulate(state.copy(), use_distilled, debug=False)
        print('Testing rollout of MCTS')
        for sample in path:
            task = self.tasks[np.argmax(sample.get(TASK_ENUM, t=0))]
            targ = self.agent.targ_list[np.argmax(sample.get(TARG_ENUM, t=0))]
            print(task, targ)
            print(sample.get_X())
        print('End of MCTS rollout.\n\n')

    def run(self, state, num_rollouts=20, use_distilled=True, hl_plan=None, new_policies=None, debug=False):
        if new_policies != None:
            self.rollout_policy = new_policies
        opt_val = np.inf
        paths = []
        for n in range(num_rollouts):
            self.n_runs += 1
            # if not self.n_runs % 10 and self.n_success == 0:
            #     self.max_depth += 1
            self.agent.reset_hist()
            print("MCTS Rollout {0} for condition {1}.\n".format(n, self.condition))
            next_path = self.simulate(state.copy(), use_distilled, debug=debug)
            print("Finished Rollout {0} for condition {1}.\n".format(n, self.condition))
            if len(next_path):
                end = next_path[-1]
                new_opt_value = self.agent.goal_f(self.condition, state)
                if new_opt_value == 0:
                    paths.append(next_path)
                    self.n_success += 1
                opt_val = np.minimum(new_opt_value, opt_val)

        self.agent.add_task_paths(paths)
        return opt_val

    def _simulate_from_unexplored(self, state, node, prev_sample, exclude_hl=[], use_distilled=True, label=None, debug=False):
        if debug:
            print('Simulating from unexplored children.')
        if label is None:
            sample = Sample(self.agent)
            sample.set_X(state.copy(), 0)
            sample.set(TRAJ_HIST_ENUM, np.array(self.agent.traj_hist).flatten(), 0)
            dummy_label = tuple(np.zeros(len(self.num_prims)+1, dtype='int32'))
            self.agent.fill_sample(self.condition, sample, sample.get(STATE_ENUM, 0), 0, dummy_label, fill_obs=True)
            distrs = self.prob_func(sample.get_prim_obs(t=0))
            distr = np.ones(1)
            for i in range(len(distrs)):
                distr = np.outer(distr, distrs[i]).flatten()
        else:
            distr = np.zeros(len(self.tasks)*np.product(self.num_prims))
            ind = 0
            m = 1
            for i in range(len(label), 1, -1):
                ind += m * label[i-1]
                m *= self.num_prims[i-2]
            ind += m * label[0]
            distr[ind] = 1.0

        next_node = None
        old_traj_hist = self.agent.get_hist()
        while np.any(distr > 0):
            ind = np.argmax(distr)
            distr[ind] = 0.
            label = []
            m = np.product(self.num_prims)
            label.append(int(ind / m))
            ind = ind % m
            for i in range(len(self.num_prims)):
                m /= self.num_prims[i]
                label.append(int(ind / m))
                ind = ind % m
            label = tuple(label)

            if node.get_child(label) is not None:
                continue

            precond_cost = self.agent.cost_f(state, label, self.condition, active_ts=(0,0), debug=debug)
            if precond_cost > 0:
                continue

            self.agent.reset_hist(deepcopy(old_traj_hist))
            next_node = MCTSNode(tuple(label),
                                 0,
                                 node,
                                 len(self.tasks),
                                 self.prim_dims)
            cost, _ = self.simulate_from_next(next_node, state, prev_sample, num_samples=5, use_distilled=use_distilled, save=True, exclude_hl=exclude_hl, debug=debug)
            next_node.update_value(int(cost==0))
            node.add_child(next_node)
            while node != self.root:
                node.update_value(int(cost==0))
                node = node.parent
            return None

        self.agent.reset_hist(old_traj_hist)

        if debug:
            print('Rejected all unexplored child nodes.')

        return None

    def _select_from_explored(self, state, node, exclude_hl=[], label=None, debug=False):
        if debug:
            print("Selecting from explored children.")
            print("State: ", state)

        if label is None:
            children = node.get_explored_children()
            children_distr = list(map(self.node_check_f, children))
        else:
            children = [node.get_child(label)]
            assert children[0] is not None
            children_distr = np.ones(1)

        while np.any(children_distr > -np.inf):
            next_ind = np.argmax(children_distr)
            children_distr[next_ind] = -np.inf
            next_node = children[next_ind]
            label = next_node.label
            plan = self.agent.plans[label]
            if self.agent.cost_f(state, label, self.condition, active_ts=(0,0), debug=debug) == 0:
                if debug:
                    print('Chose explored child.')
                return children[next_ind]
        return None

    def _default_choose_next(self, state, node, prev_sample, exclude_hl=[], use_distilled=True, debug=False):
        if debug:
            print('Choosing next node.')
        parameterizations, values = [], []
        for i in range(len(self.tasks) * np.product(self.num_prims)):
            ind = i
            label = []
            m = np.product(self.num_prims)
            label.append(int(ind / m))
            ind = ind % m
            for j in range(0, len(self.num_prims)):
                m /= self.num_prims[j]
                label.append(int(ind / m))
                ind = ind % m
            label = tuple(label)
            parameterizations.append(label)
            values.append(self.node_check_f(label, state, node))

        values = np.array(values)
        p = parameterizations[np.argmax(values)]
        values[np.argmax(values)] = -np.inf
        cost = self.agent.cost_f(state, p, self.condition, active_ts=(0,0), debug=False)
        while cost > 0 and np.any(values > -np.inf):
            p = parameterizations[np.argmax(values)]
            values[np.argmax(values)] = -np.inf
            cost = self.agent.cost_f(state, p, self.condition, active_ts=(0,0), debug=False)

        child = node.get_child(p)
        if debug:
            print('Chose to explore ', p)
        if child is None:
            new_node = self._simulate_from_unexplored(state, node, prev_sample, exclude_hl, use_distilled, label=p, debug=debug)
        else:
            new_node = self._select_from_explored(state, node, exclude_hl, label=p, debug=debug)
        if new_node is None:
            return new_node, -np.inf
        return new_node, new_node.value

    def sample(self, task, cur_state, plan, num_samples, use_distilled=True, node=None, save=True, debug=False):
        samples = []
        old_traj_hist = self.agent.get_hist()
        task_name = self.tasks[task[0]]

        s, success = None, False
        pol = MixedPolicy(self.rollout_policy[task_name], self.agent.dU, self.agent.action_inds, self.agent.state_inds, None, self.opt_strength)
        if self.opt_strength > 0:
            gmm = self.gmms[task_name] if self.gmms is not None else None
            inf_f = None if gmm is None or gmm.sigma is None else lambda s: gmm.inference(np.concatenate[s.get(STATE_ENUM), s.get_U()])
            s, failed, success = self.agent.solve_sample_opt_traj(cur_state, task, self.condition, inf_f=inf_f)
            if success:
                pol.opt_traj = s.get(ACTION_ENUM)
            else:
                pol.opt_strength = 0

        for n in range(self.num_samples):
            self.agent.reset_hist(deepcopy(old_traj_hist))
            samples.append(self.agent.sample_task(pol, self.condition, cur_state, task, noisy=True))
            if success:
                samples[-1].set_ref_X(s.get_ref_X())
                samples[-1].set_ref_U(s.get_ref_U())

        lowest_cost_sample = samples[0]

        opt_fail = False

        if save:
            self.agent.add_sample_batch(samples, task)
        cur_state = lowest_cost_sample.get_X(t=lowest_cost_sample.T-1)
        self.agent.reset_hist(lowest_cost_sample.get_U()[-self.agent.hist_len:].tolist())

        return lowest_cost_sample, cur_state

    def simulate(self, state, use_distilled=True, early_stop_prob=0.0, debug=False):
        current_node = self.root
        path = []
        samples = []

        success = True
        cur_state = state.copy()
        prev_sample = None
        terminated = False
        iteration = 0
        exclude_hl = []
        path_value = None
        while True:
            if debug:
                print("Taking simulation step")
            if self.agent.goal_f(self.condition, state) == 0: # or current_node.depth >= self.max_depth:
                break

            next_node, _ = self._choose_next(cur_state, current_node, prev_sample, exclude_hl, use_distilled, debug=debug)

            if next_node == None:
                break

            label = next_node.label
            if self.agent.cost_f(cur_state, label, self.condition, active_ts=(0,0)) > 0:
                break

            plan = self.agent.plans[label]
            next_sample, cur_state = self.sample(label, cur_state, plan, self.num_samples, use_distilled, debug=debug)

            if next_sample is None:
                break

            current_node.sample_links[next_sample] = prev_sample # Used to retrace paths
            prev_sample = next_sample

            current_node = next_node
            path.append(current_node)
            # exclude_hl += [self._encode_f(cur_state, plan, self.agent.targets[self.condition])]

            iteration += 1
            if np.random.uniform() < early_stop_prob:
                break

        if path_value is None:
            path_value = self.agent.goal_f(self.condition, cur_state)
        path = []
        while current_node is not self.root:
            path.append(prev_sample)
            prev_sample.task_cost = path_value
            prev_sample = current_node.parent.sample_links[prev_sample]
            current_node.update_value(int(path_value==0))
            current_node = current_node.parent

        path.reverse()
        return path

    def simulate_from_next(self, node, state, prev_sample, num_samples=1, save=False, exclude_hl=[], use_distilled=True, debug=False):
        if debug:
            print("Running simulate from next")
        label = node.label
        cost, samples = self._simulate_from_next(label, node.depth, node.depth, state, self.prob_func, [], num_samples, save, exclude_hl, use_distilled, [], debug=debug)
        if len(samples):
            node.sample_links[samples[0]] = prev_sample
        else:
            samples.append(None)
        return cost, samples[0]

    def _simulate_from_next(self, label, depth, init_depth, state, prob_func, samples, num_samples=1, save=True, exclude_hl=[], use_distilled=True, exclude=[], debug=False):
        # print 'Entering simulate call:', datetime.now()
        task = self.tasks[label[0]]
        new_samples = []

        # if self._encode_f(state, self.agent.plans.values()[0], self.agent.targets[self.condition], (task, obj.name, targ.name)) in exclude_hl:
        #     return self._goal_f(state, self.agent.targets[self.condition], self.agent.plans[task, obj.name]), samples

        plan = self.agent.plans[label]
        if self.agent.cost_f(state, label, self.condition, active_ts=(0,0)) > 0:
            return self.agent.goal_f(self.condition, state), samples

        next_sample, end_state = self.sample(label, state, plan, num_samples=num_samples, use_distilled=use_distilled, debug=debug)

        if next_sample is None:
            path_value = self.agent.goal_f(self.condition, state)
            for sample in samples:
                sample.task_cost = path_value
                sample.success = SUCCESS_LABEL if path_value == 0 else FAIL_LABEL
            return path_value, samples
        samples.append(next_sample)

        path_value = self.agent.goal_f(self.condition, end_state)
        # hl_encoding = self._encode_f(end_state, self.agent.plans.values()[0], self.agent.targets[self.condition])
        if path_value == 0 or depth >= init_depth + self.explore_depth or depth >= self.max_depth: # or hl_encoding in exclude_hl:
            for sample in samples:
                sample.task_cost = path_value
                sample.success = SUCCESS_LABEL if path_value == 0 else FAIL_LABEL
            return path_value, samples

        # exclude_hl = exclude_hl + [hl_encoding]

        sample = Sample(self.agent)
        sample.set_X(end_state.copy(), t=0)
        self.agent.fill_sample(self.condition, sample, sample.get(STATE_ENUM, 0), 0, label, fill_obs=True)
        distrs = self.prob_func(sample.get_prim_obs(t=0))
        next_label = []
        if self.soft_decision:
            for i in range(len(label)):
                if sum(distrs[i]) > 0:
                    distrs[i] = distrs[i] / sum(distrs[i])
                    next_label.append(np.random.choice(list(range(len(distrs[i]))), p=distrs[i]))
                else:
                    next_label.append(np.argmax(distrs[i]))
        else:
            for distr in distrs:
                next_label.append(np.argmax(distr))

        next_label = tuple(next_label)
        return self._simulate_from_next(next_label, depth+1, init_depth, end_state, prob_func, samples, num_samples=num_samples, save=False, use_distilled=use_distilled, debug=debug)

ext: py

max_issues_repo_name: Algorithmic-Alignment-Lab/OpenTAMP

max_forks_repo_name: Algorithmic-Alignment-Lab/OpenTAMP

reason: import numpy


text:
from random import randint
import sympy
import argparse
import sys

size = 31
maximum = 2**31

def mod_inverse(prime,max):
    x, y, m0 = 1, 0, max

    if max == 1:
        return 0

    while prime > 1:
        quotient = prime // max
        temp = max
        max = prime % max
        prime, temp = temp, y

        # Update x and y
        y = x - quotient * y
        x = temp

    if x < 0 :
        x = x + m0

    return x

def get_random():
    return randint(2**18,maximum-10)

def main():
    parser = argparse.ArgumentParser(description='Generate constructor values for your prime')
    parser.add_argument('p',type=int,default = 0)
    args = parser.parse_args()
    if args.p:
        if not sympy.ntheory.primetest.isprime(args.p):
            sys.exit('the given integer is not a prime number')
        elif args.p > 2147483647:
            sys.exit('The given prime number should be less than 2147483647')
        else:
            prime = args.p
    else:
        prime = sympy.ntheory.generate.randprime(2**18,maximum-1)

    print('prime : {} \ninverse : {} \nrandom : {}'.format(prime,mod_inverse(prime,maximum),get_random()))

if __name__ == '__main__':
    main()

max_stars_repo_name: elisbyberi/optimus

max_issues_repo_name: elisbyberi/optimus

max_forks_repo_name: elisbyberi/optimus


text:
import logging
import time

from ROAR.utilities_module.utilities import get_ip
import qrcode
import cv2
import numpy as np
import socket
import struct

def showIPUntilAck():
    img = np.array(qrcode.make(f"{get_ip()}").convert('RGB'))
    success = False
    addr = None

    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    try:
        s.bind((get_ip(), 8008))
        s.settimeout(1)
        while True:
            try:
                s.listen()

                cv2.imshow("Scan this code to connect to phone", img)
                k = cv2.waitKey(1) & 0xff
                if k == ord('q') or k == 27:
                    s.close()
                    break
                conn, addr = s.accept()
                addr = addr[0]
                if conn:
                    s.close()
                    success = True
                    break
            except socket.timeout as e:
                logging.info(f"Please tap on the ip address to scan QR code. ({get_ip()}:{8008}). {e}")
    except Exception as e:
        logging.error(f"Unable to bind socket: {e}")
    finally:
        s.close()
        cv2.destroyWindow("Scan this code to connect to phone")
    return success, addr

def recv(s: socket.socket):
    buffer_num = -1
    log = dict()
    while True:
        seg, addr = s.recvfrom(9620)
        prefix_num = int(seg[0:3].decode('ascii'))
        total_num = int(seg[3:6].decode('ascii'))
        curr_buffer = int(seg[6:9].decode('ascii'))

        # print(f"BEFORE curr_buff = {curr_buffer} | prefix_num = {prefix_num} "
        #       f"| total_num = {total_num} | len(log) = {len(log)}")
        if buffer_num == -1:
            # initializing
            buffer_num = curr_buffer
            if prefix_num != 0:
                # if the first one is not the starting byte, dump it.
                print("the first one is not the starting byte")
                buffer_num = -1
                log = dict()
            else:
                # if the first one is the starting byte, start recording
                log[prefix_num] = seg[9:]
        else:
            if prefix_num in log:
                # if i received a frame from another sequence
                print("i received a frame from another sequence")
                buffer_num = -1
                log = dict()
            else:
                log[prefix_num] = seg[9:]
        # print(f"AFTER curr_buff = {curr_buffer} | prefix_num = {prefix_num} | total_num = {total_num} "
        #       f"| len(log) = {len(log)} | log.keys = {list(sorted(log.keys()))}")

        if len(log) - 1 == total_num:
            data = b''
            for k in sorted(log.keys()):
                data += log[k]
            # print()
            return data

def client_test():
    msgFromClient = "Hello UDP Server"
    bytesToSend = str.encode(msgFromClient)
    serverAddressPort = ("10.0.0.26", 8001)
    bufferSize = 1024
    # Create a UDP socket at client side
    UDPClientSocket = socket.socket(family=socket.AF_INET, type=socket.SOCK_DGRAM)
    while True:
        # Send to server using created UDP socket
        start = time.time()
        UDPClientSocket.sendto(b"hello", serverAddressPort)

        # print("message sent")
        data = recv(UDPClientSocket)

        img_data = data[16:]
        intrinsics = data[:16]
        fx, fy, cx, cy = struct.unpack('f', intrinsics[0:4])[0], \
                         struct.unpack('f', intrinsics[4:8])[0], \
                         struct.unpack('f', intrinsics[8:12])[0], \
                         struct.unpack('f', intrinsics[12:16])[0]
        intrinsics = np.array([
            [fx, 0, cx],
            [0, fy, cy],
            [0, 0, 1]
        ])
        # print(intrinsics)
        img = np.frombuffer(img_data, dtype=np.uint8)
        img = cv2.imdecode(img, cv2.IMREAD_UNCHANGED)
        cv2.imshow("img", img)
        cv2.waitKey(1)
        print(f"fps = {1 / (time.time() - start)}")

if __name__ == '__main__':
    client_test()


text:
"""
Excel to Dfs0 Tool:

Created on Wed June 1, 2021

@author: Shubhneet Singh
ssin@dhigroup.com
DHI,US

"""
# Dependencies
import os
import clr
import sys
import time
import openpyxl
import numpy as np #
import pandas as pd #
import datetime
import math

from winreg import ConnectRegistry, OpenKey, HKEY_LOCAL_MACHINE, QueryValueEx

def get_mike_bin_directory_from_registry():
    x86 = False
    dhiRegistry = "SOFTWARE\Wow6432Node\DHI\\"
    aReg = ConnectRegistry(None, HKEY_LOCAL_MACHINE)
    try:
        _ = OpenKey(aReg, dhiRegistry)
    except FileNotFoundError:
        x86 = True
        dhiRegistry = "SOFTWARE\Wow6432Node\DHI\\"
        aReg = ConnectRegistry(None, HKEY_LOCAL_MACHINE)
        try:
            _ = OpenKey(aReg, dhiRegistry)
        except FileNotFoundError:
            raise FileNotFoundError
    year = 2030
    while year > 2010:
        try:
            mikeHomeDirKey = OpenKey(aReg, dhiRegistry + str(year))
        except FileNotFoundError:
            year -= 1
            continue
        if year > 2020:
            mikeHomeDirKey = OpenKey(aReg, dhiRegistry + "MIKE Zero\\" + str(year))

        mikeBin = QueryValueEx(mikeHomeDirKey, "HomeDir")[0]
        mikeBin += "bin\\"

        if not x86:
            mikeBin += "x64\\"

        if not os.path.exists(mikeBin):
            print(f"Cannot find MIKE ZERO in {mikeBin}")
            raise NotADirectoryError
        return mikeBin

    print("Cannot find MIKE ZERO")
    return ""

sys.path.append(get_mike_bin_directory_from_registry())
clr.AddReference("DHI.Generic.MikeZero.DFS")
clr.AddReference("DHI.Generic.MikeZero.EUM")
clr.AddReference("DHI.Projections")

# import xlrd
from  mikeio import *
from mikeio.eum import ItemInfo, EUMType, EUMUnit
from DHI.Generic.MikeZero.DFS import DataValueType

from tkinter import Frame, Label, Button, Entry, Tk, W, END
from tkinter import messagebox as tkMessageBox
# from tkinter.filedialog import askdirectory
from tkinter.filedialog import askopenfilename
from tkinter.filedialog import asksaveasfilename

#------------------------------------------------------------------------------
# UI for this tool:

class interface(Frame):
    def __init__(self, master = None):
        """ Initialize Frame. """
        Frame.__init__(self,master)
        self.grid()
        self.createWidgets()

    def message(self):
        tkMessageBox.showinfo("Task Complete", "Dfs0 Created!")

    def run(self):

        # input1 - Data in excel:
        filename1 = self.file_name1.get()
        # filename1 = r"C:\Users\ssin\OneDrive - DHI\Documents\Shuby's Toolbox\2 Excel2Dfs0(Large) Tool\Excel2Dfs0_Example.xlsx"
        # Output:
        outputFile = self.file_name5.get()
        # outputFile = r"C:\Users\ssin\OneDrive - DHI\Documents\Shuby's Toolbox\2 Excel2Dfs0(Large) Tool\Test1.dfs0"
        # Tool
        begin_time = datetime.datetime.now()

        if str(os.path.splitext(filename1)[1])== '.csv':
            df = pd.read_csv(filename1, index_col=0)
        elif str(os.path.splitext(filename1)[1])== '.xlsx':
            df = pd.read_excel(filename1, index_col=0,engine='openpyxl')

        print('Dataframe created. Time taken > ' + str(datetime.datetime.now() - begin_time))

        # df.index = [df.index[i].round('1s') for i in range(len(df.index))]
        item_names = df.columns
        items_type = df.iloc[0]
        item_units =  df.iloc[1]
        item_datatype = df.iloc[2]

        if type(item_names[0]) != str:
            item_names = [str(item_names[i]) for i in range(len(item_names))]
        df=df[3:]

        data=[df.iloc[:,i].values for i in range(len(df.columns))]
        time=df.index

        if pd.isnull(item_units.iloc[0]) != True:
            items=[ItemInfo(item_names[i], EUMType[items_type.iloc[i]], EUMUnit[item_units.iloc[i]]) for i in range(len(item_names))]
        else:
            items=[ItemInfo(item_names[i], EUMType[items_type.iloc[i]]) for i in range(len(item_names))]

        data_type = [item_datatype.iloc[i] for i in range(len(item_datatype))]
        dtypes = ['Instantaneous', 'Accumulated', 'StepAccumulated' , 'MeanStepBackward', 'MeanStepForward']

        data_value_type=[None]*len(data_type)

        for i in range(len(data_type)):
            for a in range(len(dtypes)):
                if data_type[i]== dtypes[a]:
                    data_value_type[i] = a

        ds = Dataset(data, time,items)

        qq = exec('DataValueType' + "." + item_datatype.iloc[0])

        outputFile_dropExt = str(os.path.splitext(outputFile)[0])
        dfs_title = outputFile_dropExt.rsplit("/", 1)[1]
        print('Dfs0 writing commences > ' + str(datetime.datetime.now() - begin_time))

        dfs = Dfs0()

        dfs.write(filename= outputFile, data=ds, title= dfs_title, data_value_type=data_value_type)

        print('Congrats, dfs0 created! Time taken > ' + str(datetime.datetime.now() - begin_time))

        self.message()

    def createWidgets(self):

        # set all labels of inputs:

        Label(self, text = "Data (*.csv or *.xlsx) :")\
            .grid(row=0, column=0, sticky=W)

        Label(self, text = "Output File (*.dfs0) :")\
            .grid(row=1, column=0, sticky=W)

        # set buttons
        Button(self, text = "Browse", command=self.load_file1, width=10)\
            .grid(row=0, column=6, sticky=W)

        Button(self, text = "Save As", command=self.load_file5, width=10)\
            .grid(row=1, column=6, sticky=W)
        Button(self, text = "Run", command=self.run, width=20)\
            .grid(row=4, column=3, sticky=W)

        # set entry field
        self.file_name1 = Entry(self, width=65)
        self.file_name1.grid(row=0, column=1, columnspan=4, sticky=W)

        self.file_name5 = Entry(self, width=65)
        self.file_name5.grid(row=1, column=1, columnspan=4, sticky=W)

    def load_file1(self):
        self.filename = askopenfilename(initialdir=os.path.curdir, defaultextension=".xlsx", filetypes=(("csv", "*.csv"),("xlsx File", "*.xlsx"),("All Files", "*.*") ))
        if self.filename:
            try:
                #self.settings.set(self.filename)
                self.file_name1.delete(0, END)
                self.file_name1.insert(0, self.filename)
                self.file_name1.xview_moveto(1.0)
            except IOError:
                tkMessageBox.showerror("Error","Failed to read file \n'%s'"%self.filename)

    def load_file5(self):
        self.filename = asksaveasfilename(initialdir=os.path.curdir,defaultextension=".dfs0", filetypes=(("Dfs0 File", "*.dfs0"),("All Files", "*.*") ))
        if self.filename:
            try:
                #self.settings.set(self.filename)
                self.file_name5.delete(0, END)
                self.file_name5.insert(0, self.filename)
                self.file_name5.xview_moveto(1.0)
            except IOError:
                tkMessageBox.showerror("Error","Failed to read file \n'%s'"%self.filename)

##### main program

root = Tk()
UI = interface(master=root)
UI.master.title("Excel to Dfs0 Tool")
UI.master.geometry('625x120')
for child in UI.winfo_children():
    child.grid_configure(padx=4, pady =6)

UI.mainloop()


text:
from numpy.linalg import norm
from scipy.optimize import linear_sum_assignment
import numpy as np
import torch

def cosine(x, y):
    """cosine distance between two 1-d vectors"""
    return(1 - (x.dot(y.T) / (norm(x) + norm(y))))

def assignment(assigned_clusters, correct_labels):
    """duplicate of match_labels"""
    assert assigned_clusters.shape == correct_labels.shape
    unique_classes = np.unique(correct_labels)
    unique_clusters = np.unique(assigned_clusters)
    cost_matrix = np.zeros((len(unique_classes), len(unique_clusters)))
    for i in range(len(unique_classes)):
        boolean_array_labels = correct_labels == unique_classes[i]
        clusters_in_label = assigned_clusters[boolean_array_labels]
        for j in range(len(unique_clusters)):
            cost_matrix[i, j] = len(clusters_in_label[clusters_in_label == unique_clusters[j]])
    cost_matrix = 1 - cost_matrix/np.sum(cost_matrix, axis=0)
    classes_index, clusters_index = linear_sum_assignment(cost_matrix)
    transformation = dict(zip(unique_clusters[list(clusters_index)], unique_classes[list(classes_index)]))
    return transformation

def match_labels(a, b):
    """
    Reorders a to match b, assuming they are two sets of labels, with the same
    number of labels, but they are not nessicarily the same.

    Minimizes the distance between the input label vectors. Inputs must be the
    same length. For details see section 2 of Lange T et al. 2004.

    E.g,
    a = [1,1,2,3,3,4,4,4,2]
    b = [2,2,3,1,1,4,4,4,3]
    optimal: 1 -> 2; 2 -> 3; 3 -> 1; 4 -> 4

    Inspired by http://things-about-r.tumblr.com/post/36087795708/matching-clustering-solutions-using-the-hungarian
    """
    a = np.array(a)
    b = np.array(b)
    ids_a = np.unique(a)
    ids_b = np.unique(b)

    assert len(a) == len(b)
    assert len(ids_a) == len(ids_b)

    # construct a distance matrix D between a and b
    n = len(ids_a)
    D = np.zeros((n, n)) # distance matrix

    for x in np.arange(n):
        for y in np.arange(n):
            idx_a = np.where(a == x)[0]
            idx_b = np.where(b == y)[0]
            n_int = len(np.intersect1d(idx_a, idx_b))

            # distance = (# in label) - 2*sum(# in intersection)
            D[x,y] = (len(idx_a) + len(idx_b) - 2*n_int)

    # permute labels w/ minimum weighted bipartite matching (hungarian method)
    idx_D_x, idx_D_y = linear_sum_assignment(D)
    b_out = np.zeros(len(b))

    # assign values in b so they resemble a (as much as is possible)
    for i, c in enumerate(idx_D_y):
        idx_map = np.where(b == c)
        b_out[idx_map] = i

    return(b_out)

def convert_ds_to_np(D):
    """torch ds --> numpy matrix. x becomes a (n x m) matrix."""
    X, y = torch.load(D)
    X = X.numpy().swapaxes(1,3).squeeze()            # make (n x row x col)
    X = X.reshape(X.shape[0], X.shape[1]*X.shape[2]) # make (n x m)

    return(X, y.numpy())

if __name__ == '__main__':
    a = np.array([1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4])
    b = np.array([37, 37, 37, 99, 150, 150, 150, 99, 99, 120, 120, 120, 150, 99, 99, 150])
    print(assignment(b, a))

reason: import numpy,from numpy,from scipy


text:
#!/usr/bin/python

#######################################################
# Copyright (c) 2015, ArrayFire
# All rights reserved.
#
# This file is distributed under 3-clause BSD license.
# The complete license agreement can be obtained at:
# http://arrayfire.com/licenses/BSD-3-Clause
########################################################

from random import random
from time import time
import arrayfire as af
import sys

try:
    import numpy as np
except ImportError:
    np = None

#alias range / xrange because xrange is faster than range in python2
try:
    frange = xrange  #Python2
except NameError:
    frange = range   #Python3

# Having the function outside is faster than the lambda inside
def in_circle(x, y):
    return (x*x + y*y) < 1

def calc_pi_device(samples):
    x = af.randu(samples)
    y = af.randu(samples)
    return 4 * af.sum(in_circle(x, y)) / samples

def calc_pi_numpy(samples):
    np.random.seed(1)
    x = np.random.rand(samples).astype(np.float32)
    y = np.random.rand(samples).astype(np.float32)
    return 4. * np.sum(in_circle(x, y)) / samples

def calc_pi_host(samples):
    count = sum(1 for k in frange(samples) if in_circle(random(), random()))
    return 4 * float(count) / samples

def bench(calc_pi, samples=1000000, iters=25):
    func_name = calc_pi.__name__[8:]
    print("Monte carlo estimate of pi on %s with %d million samples: %f" % \
          (func_name, samples/1e6, calc_pi(samples)))

    start = time()
    for k in frange(iters):
        calc_pi(samples)
    end = time()

    print("Average time taken: %f ms" % (1000 * (end - start) / iters))

if __name__ == "__main__":
    if (len(sys.argv) > 1):
        af.set_device(int(sys.argv[1]))
    af.info()

    bench(calc_pi_device)
    if np:
        bench(calc_pi_numpy)
    bench(calc_pi_host)

max_stars_repo_name: miketrumpis/arrayfire-python

max_issues_repo_name: miketrumpis/arrayfire-python

max_forks_repo_name: miketrumpis/arrayfire-python


text:
import random
import numpy as np
import matplotlib.pyplot as plt
import os
import cv2
import torch
import dense_correspondence_manipulation.utils.utils as utils
import yaml
utils.add_dense_correspondence_to_python_path()

import dense_correspondence
from dense_correspondence.evaluation.evaluation import *
import dense_correspondence.correspondence_tools.correspondence_plotter as correspondence_plotter
from dense_correspondence.dataset.dense_correspondence_dataset_masked import ImageType

config_filename = os.path.join(utils.getDenseCorrespondenceSourceDir(), 'config',
                               'dense_correspondence', 'evaluation', 'evaluation.yaml')
config = utils.getDictFromYamlFilename(config_filename)
default_config = utils.get_defaults_config()

# utils.set_cuda_visible_devices([0])
dce = DenseCorrespondenceEvaluation(config)
DCE = DenseCorrespondenceEvaluation

network_name = "caterpillar_3"
dcn = dce.load_network_from_config(network_name)
dataset = dcn.load_training_dataset()
# DenseCorrespondenceEvaluation.evaluate_network_qualitative(dcn, dataset=dataset, randomize=True)

dataset_config_filename = os.path.join(utils.getDenseCorrespondenceSourceDir(), 'config', 'dense_correspondence',
                                       'dataset', 'composite',
                                       'caterpillar_upright.yaml')
# dataset_config_filename = os.path.join(utils.getDenseCorrespondenceSourceDir(), 'config', 'dense_correspondence',
#                                        'training', 'training.yaml')
dataset_config = utils.getDictFromYamlFilename(dataset_config_filename)

dataset = SpartanDataset(debug=True, config=dataset_config)
dataset.debug=False

def get_relative_pose(dataset,metadata):
    X_OCa = dataset.get_pose_from_scene_name_and_idx(metadata['scene_name'], metadata['image_a_idx'])
    X_OCb = dataset.get_pose_from_scene_name_and_idx(metadata['scene_name'], metadata['image_b_idx'])
    X_CbCa = np.matmul(np.linalg.inv(X_OCb), X_OCa)
    return X_CbCa

def get_intrinsic(dataset, metadata):
    scene_directory = dataset.get_full_path_for_scene(metadata['scene_name'])
    yaml_filename = images_dir = os.path.join(scene_directory, 'images','camera_info.yaml')
    config = yaml.load(open(yaml_filename), Loader=yaml.CLoader)

    fx = config['camera_matrix']['data'][0]
    cx = config['camera_matrix']['data'][2]

    fy = config['camera_matrix']['data'][4]
    cy = config['camera_matrix']['data'][5]
    K = np.array([[fx, 0, cx], [0, fy, cy], [0,0,1]])

    dist_coeffs = np.array(config['distortion_coefficients']['data'][:])
    return K, dist_coeffs

def get_matches(dcn, dataset, rgb_a, rgb_b, mask_a, mask_b,num_matches):
    mask_a = np.asarray(mask_a)
    mask_b = np.asarray(mask_b)

    # compute dense descriptors
    rgb_a_tensor = dataset.rgb_image_to_tensor(rgb_a)
    rgb_b_tensor = dataset.rgb_image_to_tensor(rgb_b)

    # these are Variables holding torch.FloatTensors, first grab the data, then convert to numpy
    res_a = dcn.forward_single_image_tensor(rgb_a_tensor).data.cpu().numpy()
    res_b = dcn.forward_single_image_tensor(rgb_b_tensor).data.cpu().numpy()

    # sample points on img_a. Compute best matches on img_b
    # note that this is in (x,y) format
    # TODO: if this mask is empty, this function will not be happy
    # de-prioritizing since this is only for qualitative evaluation plots
    sampled_idx_list = random_sample_from_masked_image(mask_a, num_matches)

    # list of cv2.KeyPoint
    kp1 = []
    kp2 = []
    matches = []  # list of cv2.DMatch
    diffs = []

    # placeholder constants for opencv
    diam = 0.01
    dist = 0.01

    try:
        descriptor_image_stats = dcn.descriptor_image_stats
    except:
        print "Could not find descriptor image stats..."
        print "Only normalizing pairs of images!"
        descriptor_image_stats = None

    for i in xrange(0, num_matches):
        # convert to (u,v) format
        pixel_a = [sampled_idx_list[1][i], sampled_idx_list[0][i]]
        best_match_uv, best_match_diff, norm_diffs =\
            DenseCorrespondenceNetwork.find_best_match(pixel_a, res_a,res_b)

        # be careful, OpenCV format is  (u,v) = (right, down)
        kp1.append(cv2.KeyPoint(pixel_a[0], pixel_a[1], diam))
        kp2.append(cv2.KeyPoint(best_match_uv[0], best_match_uv[1], diam))
        diffs.append(best_match_diff)
        matches.append(cv2.DMatch(i, i, dist))
    return kp1,kp2, matches, diffs

def epipolar_constraint(R,t,kp1, kp2,K,dist_coeffs, mask=None, return_vector=False):
    t=t.reshape(-1)
    #     print t
    t_x = np.array([[0, -t[2],t[1]],[t[2],0,-t[0]],[-t[1],t[0],0]])
    #     print t_x
    #     print R
    E =np.dot(t_x,R)

    kp1_array = np.array([kp1[idx].pt for idx in range(len(kp1))]).reshape(-1, 1,2)
    kp2_array = np.array([kp2[idx].pt for idx in range(len(kp2))]).reshape(-1,1, 2)
    pts_1_norm = cv2.undistortPoints(kp1_array, cameraMatrix=K, distCoeffs=dist_coeffs)
    pts_2_norm = cv2.undistortPoints(kp2_array, cameraMatrix=K, distCoeffs=dist_coeffs)

#     kp1_array = np.array([kp1[idx].pt for idx in range(len(kp1))]).reshape(-1, 2)
#     kp2_array = np.array([kp2[idx].pt for idx in range(len(kp2))]).reshape(-1, 2)

    # kp1_array = np.hstack([kp1_array, np.ones(kp1_array.shape[0],1)])
    # kp2_array = np.hstack([kp2_array, np.ones(kp1_array.shape[0],1)])
    mat_e = E.reshape(-1,1)
    u1 = pts_1_norm[:,0,0]
    u2 = pts_1_norm[:,0,1]
    v1 = pts_2_norm[:,0,0]
    v2 = pts_2_norm[:,0,1]

    A = np.stack([u1*u2, u1*v2, u1, v1*u2, v1*v2, v1, u2, v2,np.ones_like(u1)]).T
    # print A.shape
    error = np.abs(np.matmul(A,mat_e))
    # print 'error:',error
    if return_vector:
    	return error.reshape(-1)
    # print error
    # error = np.matmul(kp1_array, E, kp2_array.T)
    # return error.mean()
    if mask is not None and mask.sum():
    	return (error*mask).sum()/mask.sum()
    else:
    	return error.mean()

def get_pose_essential_opencv(kp1,kp2,K,dist_coeffs):
	pp = K[:2,2]
	focal_length = (K[0][0]+K[1][1])/2

	kp1_array = np.array([kp1[idx].pt for idx in range(len(kp1))]).reshape(-1, 1,2)
	kp2_array = np.array([kp2[idx].pt for idx in range(len(kp2))]).reshape(-1,1, 2)
	pts_1_norm = cv2.undistortPoints(kp1_array, cameraMatrix=K, distCoeffs=dist_coeffs)
	pts_2_norm = cv2.undistortPoints(kp2_array, cameraMatrix=K, distCoeffs=dist_coeffs)

	pts1 = np.int32(kp1_array)
	pts2 = np.int32(kp2_array)
	F, mask = cv2.findFundamentalMat(pts1,pts2,cv2.FM_RANSAC) #,ransacReprojThreshold=

	E = np.matmul(K.T, np.matmul(F, K))

	# E, mask = cv2.findEssentialMat(pts_1_norm, pts_2_norm,cameraMatrix=K, method=cv2.RANSAC, threshold=2/focal_length)
	_, R_est, t_est, mask2 = cv2.recoverPose(E, kp1_array, kp2_array, cameraMatrix=K, mask=mask)
	return R_est, t_est, mask

def homography_method(kp1,kp2,K,dist_coeffs):
	pp = K[:2,2]
	focal_length = (K[0][0]+K[1][1])/2

	kp1_array = np.array([kp1[idx].pt for idx in range(len(kp1))]).reshape(-1, 1,2)
	kp2_array = np.array([kp2[idx].pt for idx in range(len(kp2))]).reshape(-1,1, 2)
	pts_1_norm = cv2.undistortPoints(kp1_array, cameraMatrix=K, distCoeffs=dist_coeffs)
	pts_2_norm = cv2.undistortPoints(kp2_array, cameraMatrix=K, distCoeffs=dist_coeffs)

	H, mask = cv2.findEssentialMat(pts_1_norm, pts_2_norm,method=cv2.RANSAC)
	# print H
	_, R_est, t_est, _ = cv2.decomposeHomographyMat(H, K)
	# print R_est, t_est
	return R_est, t_est, mask

def dense_ransac(image_a,image_b,mask_a,mask_b,K,dist_coeffs=None, num_matches = 100, plot_match=True):
	pp = K[:2,2]
	focal_length = (K[0][0]+K[1][1])/2

	kp1,kp2, matches,diffs = get_matches(dcn, dataset, image_a,image_b, mask_a,mask_b,num_matches)

	R_est, t_est, mask = get_pose_essential_opencv(kp1,kp2,K,dist_coeffs)
	# R_est, t_est, mask = homography_method(kp1,kp2,K,dist_coeffs)

	print 'accept rate: ', mask.sum()/mask.shape[0]
	print 'epipolar constraints: ', epipolar_constraint(R_est, t_est, kp1,kp2,K,dist_coeffs,mask)
	# X_real = get_relative_pose(dataset, metadata)
	# # X_real = np.linalg.inv(X_real)
	# R_real = X_real[:3,:3]
	# t_real = X_real[:3,3]
	# R_real,t_real
	# print(np.matmul(R_est, np.linalg.inv(R_real)))
	# print(t_real/np.linalg.norm(t_real), t_est.reshape(-1))
	if plot_match:
		gray_a_numpy = cv2.cvtColor(np.asarray(image_a), cv2.COLOR_BGR2GRAY)
		gray_b_numpy = cv2.cvtColor(np.asarray(image_b), cv2.COLOR_BGR2GRAY)
		draw_params = dict( singlePointColor = (255,0,0),matchesMask = mask.reshape(-1).tolist(), flags=0)
		img3 = cv2.drawMatches(gray_a_numpy, kp1, gray_b_numpy, kp2, matches, outImg=gray_b_numpy, **draw_params)
		# fig, axes = plt.subplots(nrows=1, ncols=1)
		# fig.set_figheight(10)
		# fig.set_figwidth(15)
		# axes.imshow(img3)
		# plt.show()
		return R_est,t_est,img3
	return R_est,t_est

# .

def cal_essential_cov(pts_1_norm,pts_2_norm,diffs,mask, use_cov=False):
	'''
	This method is currently not working
	'''
	mask = mask!=0
	pts_1_norm = pts_1_norm[mask[:,0],:,:]
	pts_2_norm = pts_2_norm[mask[:,0],:,:]
	diffs = np.array(diffs)[mask[:,0]]
	Q = np.diag(1/diffs**2)

	if not use_cov:
		Q = np.eye(diffs.shape[0])

	u1 = pts_1_norm[:,0,0]
	u2 = pts_1_norm[:,0,1]
	v1 = pts_2_norm[:,0,0]
	v2 = pts_2_norm[:,0,1]

	A = np.stack([u1*u2, u1*v2, u1, v1*u2, v1*v2, v1, u2, v2,np.ones_like(u1)]).T

	print A.shape
	print Q.shape

	w,v = np.linalg.eigh(np.matmul(A.T,Q).dot(A))
	# print 'eigen value:',w

	E = v[:,0].reshape(3,3)

	U, s, Vt = np.linalg.svd(E)

	E = np.matmul(U,np.diag([(s[0]+s[1])/2.,(s[0]+s[1])/2.,0 ]),Vt)

	return E

def cal_essential_cov2(kp1,kp2,K,dist_coeffs,diffs,mask):
	'''
	This method is currently not working
	'''
	# only for 8-point method
	assert np.array(mask).sum()==8
	pp = K[:2,2]
	focal_length = (K[0][0]+K[1][1])/2

	kp1_array = np.array([kp1[idx].pt for idx in range(len(kp1))]).reshape(-1, 1,2)
	kp2_array = np.array([kp2[idx].pt for idx in range(len(kp2))]).reshape(-1,1, 2)
	pts_1_norm = cv2.undistortPoints(kp1_array, cameraMatrix=K, distCoeffs=dist_coeffs)
	pts_2_norm = cv2.undistortPoints(kp2_array, cameraMatrix=K, distCoeffs=dist_coeffs)

	mask = mask!=0
	pts_1_norm = pts_1_norm[mask[:,0],:,:]
	pts_2_norm = pts_2_norm[mask[:,0],:,:]
	diffs = np.array(diffs)[mask[:,0]]
	# Q = np.diag(1/diffs**2)

	pts1 = np.int32(kp1_array[mask[:,0],:,:])
	pts2 = np.int32(kp2_array[mask[:,0],:,:])
	F, mask = cv2.findFundamentalMat(pts1,pts2,cv2.FM_RANSAC) #,ransacReprojThreshold=

	E = np.matmul(K.T, np.matmul(F, K))

	E = E*diffs.reshape(3,3)**2
	U, s, Vt = np.linalg.svd(E)
	E = np.matmul(U,np.diag([(s[0]+s[1])/2.,(s[0]+s[1])/2.,0 ]),Vt)

	# E, mask = cv2.findEssentialMat(pts_1_norm, pts_2_norm,cameraMatrix=K, method=cv2.RANSAC, threshold=2/focal_length)
	_, R_est, t_est, mask2 = cv2.recoverPose(E, kp1_array, kp2_array, cameraMatrix=K, mask=mask)
	return R_est, t_est, mask

	return E

def revised_ransac(kp1, kp2, K, dist_coeffs,diffs, num_iter = 500): #sigma_0=sigma_0, epsilon=epsilon
	kp1_array = np.array([kp1[idx].pt for idx in range(len(kp1))]).reshape(-1, 1,2)
	kp2_array = np.array([kp2[idx].pt for idx in range(len(kp2))]).reshape(-1,1, 2)
	pts_1_norm = cv2.undistortPoints(kp1_array, cameraMatrix=K, distCoeffs=dist_coeffs)
	pts_2_norm = cv2.undistortPoints(kp2_array, cameraMatrix=K, distCoeffs=dist_coeffs)
	diffs = np.array(diffs)

	max_score = -np.inf

	for i in range(num_iter):
		selected_idx = np.random.choice(list(range(kp1_array.shape[0])),8)
		kp1_array_new = kp1_array[selected_idx,:,:]
		kp2_array_new = kp2_array[selected_idx,:,:]

		# E, mask_temp = cv2.findEssentialMat(pts_1_norm_new, pts_2_norm_new,cameraMatrix=K, method=cv2.LMEDS)
		pts1 = np.int32(kp1_array_new)
		pts2 = np.int32(kp2_array_new)
		# F, mask_temp = cv2.findFundamentalMat(pts1,pts2,cv2.FM_RANSAC)
		F, mask_temp = cv2.findFundamentalMat(pts1,pts2,cv2.LMEDS)
		if F is None:
			# print 'skip:',i
			continue

		E = np.matmul(K.T, np.matmul(F, K))

		# print , mask_temp.shape
		mask = np.zeros((kp1_array.shape[0],1)).astype(mask_temp.dtype)
		mask[selected_idx] = 1
		_, R_est, t_est, mask2 = cv2.recoverPose(E, kp1_array, kp2_array, cameraMatrix=K, mask=mask)
		error = epipolar_constraint(R_est, t_est, kp1,kp2,K,dist_coeffs,mask=None, return_vector = True)

		p = 1/np.sqrt(np.pi*2)/(sigma_0*diffs) * np.exp(-0.5*(error/(sigma_0*diffs))**2)
		# p  = np.exp(-0.5*(error/(sigma_0*diffs))**2)
		p=np.clip(p,epsilon,1.)
		# print p.shape
		score = np.log(p).sum()-p.shape[0]*np.log(epsilon)
		# score = p.sum() - p.shape[0]*epsilon#np.log(p).sum()-p.shape[0]*np.log(epsilon)
		# print i,score

		if score>max_score:
			best_E = E
			# best_mask = (p>epsilon).astype(mask_temp.dtype)
			# best_mask = mask#(p>epsilon).astype(mask_temp.dtype)
			best_mask = get_inliers(error).reshape((-1,1)).astype(mask_temp.dtype)

	return best_E,best_mask

def get_inliers(error,rate=0.1, threshold=0.01):
	# print error.mean()

	#
	mask = error<0.01
	if mask.sum()<error.shape[0]*rate:
		cutoff_idx = error.argsort()[int(error.shape[0]*rate)]
		mask = error<=error[cutoff_idx]
	return mask

def get_pose_essential2(kp1,kp2,K,dist_coeffs, diffs):
	pp = K[:2,2]
	focal_length = (K[0][0]+K[1][1])/2

	kp1_array = np.array([kp1[idx].pt for idx in range(len(kp1))]).reshape(-1, 1,2)
	kp2_array = np.array([kp2[idx].pt for idx in range(len(kp2))]).reshape(-1,1, 2)
	pts_1_norm = cv2.undistortPoints(kp1_array, cameraMatrix=K, distCoeffs=dist_coeffs)
	pts_2_norm = cv2.undistortPoints(kp2_array, cameraMatrix=K, distCoeffs=dist_coeffs)

	# E, mask = cv2.findEssentialMat(pts_1_norm, pts_2_norm,cameraMatrix=K, method=cv2.RANSAC, threshold=2/focal_length)
	# _, R_est, t_est, mask2 = cv2.recoverPose(E, kp1_array, kp2_array, cameraMatrix=K, mask=mask)
	# print R_est, t_est

	E, mask = revised_ransac(kp1,kp2,K, dist_coeffs,diffs)
	_, R_est, t_est, mask2 = cv2.recoverPose(E, kp1_array, kp2_array, cameraMatrix=K, mask=mask)
	# print R_est, t_estf

	# E = cal_essential_cov(pts_1_norm, pts_2_norm, diffs, mask)
	# _, R_est, t_est, mask2 = cv2.recoverPose(E, kp1_array, kp2_array, cameraMatrix=K, mask=mask)
	# print R_est, t_est
	# raw_input()
	return R_est, t_est, mask
	pass

def dense_ransac2(image_a,image_b,mask_a,mask_b,K,dist_coeffs=None, num_matches = 100, plot_match=True):
	pp = K[:2,2]
	focal_length = (K[0][0]+K[1][1])/2

	kp1,kp2, matches,diffs = get_matches(dcn, dataset, image_a,image_b, mask_a,mask_b,num_matches)

	R_est, t_est, mask = get_pose_essential2(kp1,kp2,K,dist_coeffs, diffs)
	# R_est, t_est, mask = homography_method(kp1,kp2,K,dist_coeffs)

	print 'accept rate: ', mask.sum()/mask.shape[0]
	print 'epipolar constraints: ', epipolar_constraint(R_est, t_est, kp1,kp2,K,dist_coeffs,mask)
	# X_real = get_relative_pose(dataset, metadata)
	# # X_real = np.linalg.inv(X_real)
	# R_real = X_real[:3,:3]
	# t_real = X_real[:3,3]
	# R_real,t_real
	# print(np.matmul(R_est, np.linalg.inv(R_real)))
	# print(t_real/np.linalg.norm(t_real), t_est.reshape(-1))
	if plot_match:
		gray_a_numpy = cv2.cvtColor(np.asarray(image_a), cv2.COLOR_BGR2GRAY)
		gray_b_numpy = cv2.cvtColor(np.asarray(image_b), cv2.COLOR_BGR2GRAY)
		img3 = cv2.drawMatches(gray_a_numpy, kp1, gray_b_numpy, kp2, matches, flags=2, outImg=gray_b_numpy,  matchesMask = mask.reshape(-1).tolist())
		# fig, axes = plt.subplots(nrows=1, ncols=1)
		# fig.set_figheight(10)
		# fig.set_figwidth(15)
		# axes.imshow(img3)
		# plt.show()
		return R_est,t_est,img3
	return R_est,t_est

def dense_ransac_comp(image_a,image_b,mask_a,mask_b,K,dist_coeffs, metadata, num_matches = 100, plot_match=True, num_iter=10):
	pp = K[:2,2]
	focal_length = (K[0][0]+K[1][1])/2

	min_error2 = np.inf
	min_error1 = np.inf

	for i in range(num_iter):
		# print 'round: ',i

		kp1,kp2, matches,diffs = get_matches(dcn, dataset, image_a,image_b, mask_a,mask_b,num_matches)

		R_est, t_est, mask = get_pose_essential2(kp1,kp2,K,dist_coeffs, diffs)
		if mask.sum()==0:
			# pass
			print '2,masked'

		# dtheta, dt = cal_error(dataset, metadata,R_est, t_est)
		error_vector = epipolar_constraint(R_est, t_est, kp1,kp2,K,dist_coeffs,mask,return_vector=True)
		diffs = np.array(diffs)
		weigtht = (1./np.array(diffs+0.01)**2).reshape(-1)
		error2 = (error_vector*weigtht*mask).sum()/(weigtht*mask).sum()
		if error2<min_error2:
			min_error2 = error2
			best_R_est2, best_t_est2 = R_est, t_est
		# print 'old ransac'

		R_est, t_est, mask = get_pose_essential_opencv(kp1,kp2,K,dist_coeffs)
		if mask.sum()==0:
			pass
			# print '1,masked '
		# print 'accept rate: ', mask.sum()/mask.shape[0]
		# raw_input()
		error1 = epipolar_constraint(R_est, t_est, kp1,kp2,K,dist_coeffs,mask)
		if error1<min_error1:
			min_error1 = error1
			best_R_est1, best_t_est1 = R_est, t_est
		# dtheta, dt = cal_error(dataset, metadata,R_est, t_est)

	return best_R_est1, best_t_est1, best_R_est2, best_t_est2
		# print 'accept rate: ', mask.sum()/mask.shape[0]
		# print 'epipolar constraints: ', epipolar_constraint(R_est, t_est, kp1,kp2,K,dist_coeffs,mask)

def sift_ransac(image_a,image_b,mask_a,mask_b,K,dist_coeffs=None, plot_match=False, use_mask = True):
	sift = cv2.xfeatures2d.SIFT_create()
	# find the keypoints and descriptors with SIFT
	if use_mask:
		kp1, des1 = sift.detectAndCompute(np.array(image_a),np.array(mask_a))
		kp2, des2 = sift.detectAndCompute(np.array(image_b),np.array(mask_b))
	else:
		kp1, des1 = sift.detectAndCompute(np.array(image_a),None)
		kp2, des2 = sift.detectAndCompute(np.array(image_b),None)
	# FLANN parameters
	FLANN_INDEX_KDTREE = 0
	index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)
	search_params = dict(checks=50)
	flann = cv2.FlannBasedMatcher(index_params,search_params)
	matches = flann.knnMatch(des1,des2,k=2)

	# bf = cv2.BFMatcher()
	# matches = bf.knnMatch(des1,des2, k=2)

	pts1 = []
	pts2 = []
	pts1_cv = []
	pts2_cv = []
	good_matches = []
	# ratio test as per Lowe's paper
	# dist = 0.01

	matchesMask = [[0,0] for i in xrange(len(matches))]

	for i,(m,n) in enumerate(matches):
	    if m.distance < 0.8*n.distance:
	    	# print i
	        pts2.append(kp2[m.trainIdx].pt)
	        # print kp2[m.trainIdx].pt

	        pts1.append(kp1[m.queryIdx].pt)
	        # print kp1[m.queryIdx].pt
	        pts2_cv.append(kp2[m.trainIdx])
	        pts1_cv.append(kp1[m.queryIdx])
	        good_matches.append(m)
	        matchesMask[i]=[1,1]

	# kp1_array = np.array([kp1[idx].pt for idx in range(len(kp1))]).reshape(-1, 1,2)
	# kp2_array = np.array([kp2[idx].pt for idx in range(len(kp2))]).reshape(-1,1, 2)
	# pts_1_norm = cv2.undistortPoints(kp1_array, cameraMatrix=K, distCoeffs=dist_coeffs)
	# pts_2_norm = cv2.undistortPoints(kp2_array, cameraMatrix=K, distCoeffs=dist_coeffs)
	pts1 = np.int32(pts1)
	pts2 = np.int32(pts2)
	F, mask = cv2.findFundamentalMat(pts1,pts2,cv2.FM_RANSAC)
	if F is None:
		F, mask = cv2.findFundamentalMat(pts1,pts2,cv2.FM_LMEDS)

	try:
		E = np.matmul(K.T, np.matmul(F, K))
	except:
		# print F.shape, K.shape
		print 'Failed, use random transformation'
		E = np.random.rand(3,3)
		mask = np.ones((pts1.shape[0],1)).astype(np.uint8)

	_, R_est, t_est, _ = cv2.recoverPose(E, pts1, pts2, cameraMatrix=K, mask=mask)

	# print 'accept rate: ', mask.sum()/mask.shape[0]
	# print 'epipolar constraints: ', epipolar_constraint(R_est, t_est, pts1_cv,pts2_cv,K,dist_coeffs,mask)
	# We select only inlier points
	# pts1 = pts1[mask.ravel()==1]
	# pts2 = pts2[mask.ravel()==1]

	# print len(pts1_cv), len(pts2_cv), len(new_matches), len(mask)
	# print pts1_cv
	# print pts2_cv
	# print new_matches
	if plot_match:
		# draw_params = dict(matchColor = (0,255,0), singlePointColor = (255,0,0),matchesMask = mask.reshape(-1).tolist(), flags=2)
		draw_params = dict( singlePointColor = (255,0,0),matchesMask = mask.reshape(-1).tolist(), flags=0)

		gray_a_numpy = cv2.cvtColor(np.asarray(image_a), cv2.COLOR_BGR2GRAY)
		gray_b_numpy = cv2.cvtColor(np.asarray(image_b), cv2.COLOR_BGR2GRAY)
		img3 = cv2.drawMatches(gray_a_numpy, kp1, gray_b_numpy, kp2, good_matches, outImg=gray_b_numpy, **draw_params)
		return R_est, t_est, img3
		# fig, axes = plt.subplots(nrows=1, ncols=1)
		# fig.set_figheight(10)
		# fig.set_figwidth(15)
		# axes.imshow(img3)
		# plt.show()

	return R_est, t_est
	pass

def cal_error(dataset, metadata, R_est, t_est):
	X_real = get_relative_pose(dataset, metadata)
	# X_real = np.linalg.inv(X_real)
	R_real = X_real[:3,:3]

	T0 = np.diag([1,-1,-1])
	# R_real = np.matmul(T0, R_real)
	t_real = X_real[:3,3]
	dR = np.matmul(R_est, np.linalg.inv(R_real))
	# print(dR)
	dtheta = np.arccos((np.trace(dR)-1)/2)
	# print(dtheta)
	# print(t_real/np.linalg.norm(t_real), t_est.reshape(-1))
	dt = 1-abs(np.dot(t_real/np.linalg.norm(t_real), t_est.reshape(-1)))
	# print(dt)
	# t_est = np.dot(T0, t_est)
	# dt = 1-abs(np.dot(t_real/np.linalg.norm(t_real), t_est.reshape(-1)))
	# print(dt)
	return dtheta,dt

def check_valid(mask, threshold=0.1):
	'''
	check if the object is top small
	'''
	mask = np.array(mask)
	ratio = mask.sum()/float(mask.shape[0]*mask.shape[1])
	print 'mask ratio: ', ratio
	if ratio<threshold:
		return False
	else:
		return True

def get_random_two():
	match_type, image_a_rgb, image_b_rgb, \
	matches_a, matches_b, masked_non_matches_a, \
	masked_non_matches_a, non_masked_non_matches_a, \
	non_masked_non_matches_b, blind_non_matches_a, \
	blind_non_matches_b, metadata = dataset.get_single_object_within_scene_data()

	image_a, _, mask_a, pose_a = dataset.get_rgbd_mask_pose(metadata['scene_name'], metadata['image_a_idx'])
	image_b, _, mask_b, pose_b = dataset.get_rgbd_mask_pose(metadata['scene_name'], metadata['image_b_idx'])
	return image_a, mask_a, image_b, mask_b, metadata

def get_consecutive_two():
	metadata = dict()

	# pick a scene
	scene_name = dataset.get_random_scene_name()
	metadata['scene_name'] = scene_name

	# image_a_idx = dataset.get_random_image_index(scene_name)
	pose_data = dataset.get_pose_data(scene_name)
	image_idxs = pose_data.keys()
	image_idxs = sorted(image_idxs, key=lambda x:int(x))
	# print image_idxs
	id1 = np.random.choice(list(range(len(image_idxs)-1)))
	image_a_idx = image_idxs[id1]
	image_b_idx = image_idxs[id1+1]

	image_a, _, mask_a, pose_a = dataset.get_rgbd_mask_pose(scene_name, image_a_idx)
	metadata['image_a_idx'] = image_a_idx

	image_b, _, mask_b, pose_b = dataset.get_rgbd_mask_pose(scene_name, image_b_idx)
	metadata['image_b_idx'] = image_b_idx
	return image_a, mask_a, image_b, mask_b, metadata

def get_fix_two():
	# metadata = {'scene_name': '2018-04-10-16-08-46', 'image_a_idx': 838, 'image_b_idx': 842}
	metadata ={'image_a_idx': 268, 'object_id': 'caterpillar', 'scene_name': '2018-04-16-15-23-41', 'object_id_int': 0, 'type': 0, 'image_b_idx': 2426}
	image_a, _, mask_a, pose_a = dataset.get_rgbd_mask_pose(metadata['scene_name'], metadata['image_a_idx'])
	image_b, _, mask_b, pose_b = dataset.get_rgbd_mask_pose(metadata['scene_name'], metadata['image_b_idx'])
	return image_a, mask_a, image_b, mask_b, metadata

def get_from_metadata(metadata):
	image_a, _, mask_a, pose_a = dataset.get_rgbd_mask_pose(metadata['scene_name'], metadata['image_a_idx'])
	image_b, _, mask_b, pose_b = dataset.get_rgbd_mask_pose(metadata['scene_name'], metadata['image_b_idx'])
	return image_a, mask_a, image_b, mask_b, metadata

def eval_one(num_iters = 10):
	image_a, mask_a, image_b, mask_b, metadata = get_random_two()
	print metadata

	# check_valid(mask_a)
	# check_valid(mask_b)

	K, dist_coeffs = get_intrinsic(dataset, metadata)

	# for i in range(num_iters):
	# 	R_est,t_est = dense_ransac2(image_a,image_b,mask_a,mask_b,K,dist_coeffs)
	# 	cal_error(dataset, metadata,R_est, t_est)
	R_est0,t_est0 = sift_ransac(image_a,image_b,mask_a,mask_b,K,dist_coeffs)
	dtheta0, dt0 = cal_error(dataset, metadata,R_est0, t_est0)

	R_est1, t_est1, R_est2, t_est2 = dense_ransac_comp(image_a,image_b,mask_a,mask_b,K,dist_coeffs,metadata)
	dtheta1, dt1 = cal_error(dataset, metadata,R_est1, t_est1)
	dtheta2, dt2 = cal_error(dataset, metadata,R_est2, t_est2)

	print '0: ', dtheta0, dt0
	print '1: ', dtheta1, dt1
	print '2: ', dtheta2, dt2

	return (dtheta0, dt0), (dtheta1, dt1), (dtheta2, dt2)

def eval_test_cases(filename):
	metadata_list = pkl.load(open(filename,'rb'))

	result = []

	for i,metadata in enumerate(metadata_list):
		image_a, mask_a, image_b, mask_b, metadata = get_from_metadata(metadata)
		print i,metadata

		K, dist_coeffs = get_intrinsic(dataset, metadata)

		# for i in range(num_iters):
		# 	R_est,t_est = dense_ransac2(image_a,image_b,mask_a,mask_b,K,dist_coeffs)
		# 	cal_error(dataset, metadata,R_est, t_est)
		R_est0,t_est0 = sift_ransac(image_a,image_b,mask_a,mask_b,K,dist_coeffs)
		dtheta0, dt0 = cal_error(dataset, metadata,R_est0, t_est0)

		R_est1, t_est1, R_est2, t_est2 = dense_ransac_comp(image_a,image_b,mask_a,mask_b,K,dist_coeffs,metadata)
		dtheta1, dt1 = cal_error(dataset, metadata,R_est1, t_est1)
		dtheta2, dt2 = cal_error(dataset, metadata,R_est2, t_est2)

		print '0: ', dtheta0, dt0
		print '1: ', dtheta1, dt1
		print '2: ', dtheta2, dt2

		result.append((dtheta0, dt0,dtheta1, dt1,dtheta2, dt2))

	dtheta0 = np.mean([x[0] for x in result])
	dt0 = np.mean([x[1] for x in result])
	dtheta1 = np.mean([x[2] for x in result])
	dt1 = np.mean([x[3] for x in result])
	dtheta2 = np.mean([x[4] for x in result])
	dt2 = np.mean([x[5] for x in result])

	print 'final result:'
	print 'sigma_0:',sigma_0
	print 'epsilon',epsilon
	print '0: ', dtheta0, dt0
	print '1: ', dtheta1, dt1
	print '2: ', dtheta2, dt2

def gen_test_set():
	metadata_list = []
	for i in range(100):
		image_a, mask_a, image_b, mask_b, metadata = get_consecutive_two()
		print metadata
		metadata_list.append(metadata)
	pkl.dump(metadata_list, open('consecutive_two_100_test_cases.pkl','wb'))

import pickle as pkl
import time
if __name__ == '__main__':
	sigma_0 = 1.
	epsilon = 0.1
	# eval_test_cases('consecutive_two_100_test_cases.pkl')
	eval_test_cases('random_two_100_test_cases.pkl')
	# sigma_0_list = [1,3]
	# epsilon_list = [0.01,0.03,0.1,0.2]

	# # sigma_0_epsilon_list = [(0.3,0.01)]

	# for sigma_0 in sigma_0_list:
	# 	for epsilon in epsilon_list:
	# 		eval_test_cases('random_two_100_test_cases.pkl')
	# gen_test_set()
	# t1 = time.time()
	# np.random.seed(5)
	# torch.random.manual_seed(3)
	# eval_one()
	# t2 = time.time()

	# print 'time: ', t2-t1