<commit_before>import numpy as np

class Environment(object):
    """
    An environment skeleton. Defaults to simple MAB
    H = 1, K=2, rewards are bernoulli, sampled from dirichlet([1,1]) prior.
    """

    def __init__(self):
        self.horizon = 1
        self.actions = [0,1]
        self.reward_dists = np.random.dirichlet(np.array([1,1]))
        self.reward = lambda x,a: np.random.binomial(1,self.reward_dists[a])

        self.transition = lambda x,a: "s"
        self.start = lambda: "s"

        self.state = None
        self.h = 0

    def start_episode(self):
        self.h=0
        self.state = self.start()
        return self.make_obs(self.state)

    def get_actions(self):
        if self.state is None:
            raise Exception("Episode not started")
        if self.h == self.horizon:
            return None
        return self.actions

    def make_obs(self, s):
        return s

    def act(self, a):
        if self.state is None:
            raise Exception("Episode not started")
        r = np.random.binomial(1, self.reward(self.state,a))
        self.h += 1
        if self.h == self.horizon:
            self.state = None
        else:
            self.state = self.transition(self.state,a)
        return(self.make_obs(self.state), r)

    def get_num_actions(self):
        return len(self.actions)

    def is_tabular(self):
        return True

    def get_dimension(self):
        assert not self.is_tabular(), "Not a featurized environment"
        return self.dim

class CombinationLock(Environment):
    def __init__(self, horizon=5):
        self.horizon=horizon
        self.actions = [0,1]
        self.opt = np.random.choice(self.actions,size=self.horizon)

    def transition(self, x, a):
        if x is None:
            raise Exception("Not in any state")
        if x[0] == "g" and a==self.opt[x[1]]:
            return("g",x[1]+1)
        return ("b",x[1]+1)

    def start(self):
        return ("g",0)

    def reward(self, x, a):
        if x == ("g",self.horizon-1) and a == self.opt[x[1]]:
            return 1
        return 0

class StochasticCombinationLock(Environment):
    def __init__(self, horizon=5, swap=0.1):
        self.horizon=horizon
        self.swap = swap
        self.actions = [0,1]
        self.opt_a = np.random.choice(self.actions,size=self.horizon)
        self.opt_b = np.random.choice(self.actions,size=self.horizon)

    def transition(self,x,a):
        if x is None:
            raise Exception("Not in any state")
        b = np.random.binomial(1,self.swap)
        if x[0] == "a" and a == self.opt_a[x[1]]:
            if b == 0:
                return(("a", x[1]+1))
            else:
                return(("b", x[1]+1))
        if x[0] == "b" and a == self.opt_b[x[1]]:
            if b == 0:
                return(("b", x[1]+1))
            else:
                return(("a", x[1]+1))
        else:
            return(("c", x[1]+1))

    def start(self):
        return ("a",0)

    def reward(self, x, a):
        if (x == ("a",self.horizon-1) and a == self.opt_a[x[1]]) or (x == ("b",self.horizon-1) and a == self.opt_b[x[1]]):
            return 1
        return 0

class RandomGridWorld(Environment):
    """
    A M x M grid with walls and a trembling hand.
    Horizon is always 2 M
    """
    def __init__(self, M, swap=0.1, dim=2, noise=0.0):
        self.M = M
        self.swap = swap
        self.noise = noise
        self.dim = dim
        self.seed = 147
        np.random.seed(self.seed)
        self.goal = None

        self.maze = self.generate_maze(self.M)
        self.state = None

        self.actions = [(0,1), (0,-1), (1,0), (-1,0)]
        print("ENV: Generated Random Grid World")
        print("Size: %dx%d, Start: [%d,%d], Goal: [%d,%d], H: %d, Cells: %d" % (self.M, self.M, 0, 0, self.goal[1], self.goal[0], self.horizon, np.count_nonzero(self.maze)))
        self.print_maze()

    def is_tabular(self):
        return False

    def generate_maze(self, M):
        """
        Adapted from http://code.activestate.com/recipes/578356-random-maze-generator/
        """
        mx = M; my = M
        maze = np.matrix(np.zeros((mx,my)))
        dx = [0, 1, 0, -1]; dy = [-1, 0, 1, 0] # 4 directions to move in the maze
        stack = [(0,0)]

        while len(stack) > 0:
            (cx, cy) = stack[-1]
            maze[cy,cx] = 1
            # find a new cell to add
            nlst = [] # list of available neighbors
            for i in range(4):
                nx = cx + dx[i]; ny = cy + dy[i]
                if nx >= 0 and nx < mx and ny >= 0 and ny < my:
                    if maze[ny,nx] == 0:
                        # of occupied neighbors must be 1
                        ctr = 0
                        for j in range(4):
                            ex = nx + dx[j]; ey = ny + dy[j]
                            if ex >= 0 and ex < mx and ey >= 0 and ey < my:
                                if maze[ey,ex] == 1: ctr += 1
                        if ctr == 1: nlst.append(i)
            # if 1 or more neighbors available then randomly select one and move
            if len(nlst) > 0:
                ir = nlst[np.random.randint(0, len(nlst))]
                cx += dx[ir]; cy += dy[ir]
                stack.append((cx, cy))
            elif self.goal is None:
                self.horizon = 2*len(stack)
                (gx,gy) = stack.pop()
                self.goal = [gx,gy]
            else:
                stack.pop()

        return(maze)

    def make_obs(self,s):
        if s is None:
            return None
        tmp = s.copy()
        v = np.random.normal(0, self.noise, self.dim)
        if self.dim > 2:
            mult = np.random.choice(range(self.M), size=self.dim-2)
            tmp.extend(mult)
        return tmp+v

    def start(self):
        return [0,0]

    def reward(self,x,a):
        if x == self.goal:
            return 1
        return 0

    def transition(self, x, a):
        if x == self.goal:
            return None
        nx = x[0]+a[0]
        ny = x[1]+a[1]
        if nx < 0 or nx >= self.M or ny < 0 or ny >= self.M:
            ## Cannot go off the grid
            return x
        if self.maze[ny,nx] == 0:
            ## Cannot enter a wall
            return x
        else:
            z = np.random.binomial(1, self.swap)
            if z == 1:
                return x
            else:
                return [nx,ny]

    def print_maze(self):
        maze = self.maze
        for i in range(self.M):
            for j in range(self.M):
                if maze[i,j] == 0:
                    print(" ",end="")
                elif self.state is not None and self.state[0]==j and self.state[1]==i:
                    print("A",end="")
                elif self.goal[0]==j and self.goal[1]==i:
                    print("G",end="")
                elif maze[i,j] == 1:
                    print(".",end="")
            print("")

if __name__=='__main__':
    import sys, argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--env', action='store', choices=['MAB', 'combolock', 'stochcombolock', 'maze'])
    Args = parser.parse_args(sys.argv[1:])

    if Args.env == 'MAB':
        E = Environment()
        rewards = [0,0]
        counts = [0,0]
        for t in range(1000):
            x = E.start_episode()
            while x != None:
                actions = E.get_actions()
                a = np.random.choice(actions)
                (x,r) = E.act(a)
                rewards[a] += r
                counts[a] += 1
        for a in [0,1]:
            assert (np.abs(np.float(rewards[a])/counts[a] -E.reward_dists[a]) < 0.1)

    if Args.env == 'combolock':
        E = CombinationLock(horizon=3)
        print (E.opt)
        for t in range(10):
            x = E.start_episode()
            while x != None:
                actions = E.get_actions()
                a = np.random.choice(actions)
                old = x
                (x,r) = E.act(a)
                print(old, a, r, x)

    if Args.env == 'stochcombolock':
        E = StochasticCombinationLock(horizon=3, swap=0.5)
        print (E.opt_a)
        print (E.opt_b)
        for t in range(10):
            x = E.start_episode()
            while x != None:
                actions = E.get_actions()
                a = np.random.choice(actions)
                old = x
                (x,r) = E.act(a)
                print(old, a, r, x)

    if Args.env == 'maze':
        E = RandomGridWorld(M=3,swap=0.1, dim=2, noise=0.0)
        T = 0
        while True:
            T += 1
            x = E.start_episode()
            if T % 100 == 0:
                print("Iteration t = %d" % (T))
            while x is not None:
                E.print_maze()
                print(x)
                actions = E.get_actions()
                a = np.random.choice(len(actions))
                (x,r) = E.act(actions[a])
                if r == 1:
                    print("Success: T = %d" % (T))
                    sys.exit()
<commit_msg>adding spurious features to combo locks<commit_after>import numpy as np

class Environment(object):
    """
    An environment skeleton. Defaults to simple MAB
    H = 1, K=2, rewards are bernoulli, sampled from dirichlet([1,1]) prior.
    """

    def __init__(self):
        self.horizon = 1
        self.actions = [0,1]
        self.reward_dists = np.random.dirichlet(np.array([1,1]))
        self.reward = lambda x,a: np.random.binomial(1,self.reward_dists[a])

        self.transition = lambda x,a: "s"
        self.start = lambda: "s"

        self.state = None
        self.h = 0

    def start_episode(self):
        self.h=0
        self.state = self.start()
        return self.make_obs(self.state)

    def get_actions(self):
        if self.state is None:
            raise Exception("Episode not started")
        if self.h == self.horizon:
            return None
        return self.actions

    def make_obs(self, s):
        return s

    def act(self, a):
        if self.state is None:
            raise Exception("Episode not started")
        r = np.random.binomial(1, self.reward(self.state,a))
        self.h += 1
        if self.h == self.horizon:
            self.state = None
        else:
            self.state = self.transition(self.state,a)
        return(self.make_obs(self.state), r)

    def get_num_actions(self):
        return len(self.actions)

    def is_tabular(self):
        return True

    def get_dimension(self):
        assert not self.is_tabular(), "Not a featurized environment"
        return self.dim

class CombinationLock(Environment):
    def __init__(self, horizon=5, dim=None):
        self.horizon=horizon
        self.actions = [0,1]
        self.opt = np.random.choice(self.actions,size=self.horizon)
        if dim is None:
            self.dim = None
        else:
            self.dim = 2*horizon+dim

    def transition(self, x, a):
        if x is None:
            raise Exception("Not in any state")
        if x[0] == 1 and a==self.opt[x[1]]:
            return [1,x[1]+1]
        return [0,x[1]+1]

    def make_obs(self,x):
        if x is None:
            return x
        if self.dim == None:
            return x
        else:
            v = np.zeros(self.dim,dtype=int)
            v[2*self.horizon:] = np.random.binomial(1,0.5, self.dim-2*self.horizon)
            v[2*x[1]+x[0]] = 1
            return (v)

    def start(self):
        return [1,0]

    def reward(self, x, a):
        if x == [1,self.horizon-1] and a == self.opt[x[1]]:
            return np.random.binomial(1, 0.5)
        return 0

    def is_tabular(self):
        return (self.dim == None)

class StochasticCombinationLock(Environment):
    def __init__(self, horizon=5, swap=0.1, dim=None):
        self.horizon=horizon
        self.swap = swap
        self.actions = [0,1]
        self.opt_a = np.random.choice(self.actions,size=self.horizon)
        self.opt_b = np.random.choice(self.actions,size=self.horizon)
        if dim is None:
            self.dim = None
        else:
            self.dim = 3*self.horizon+dim

    def transition(self,x,a):
        if x is None:
            raise Exception("Not in any state")
        b = np.random.binomial(1,self.swap)
        if x[0] == 0 and a == self.opt_a[x[1]]:
            if b == 0:
                return([0, x[1]+1])
            else:
                return([1, x[1]+1])
        if x[0] == 1 and a == self.opt_b[x[1]]:
            if b == 0:
                return([1, x[1]+1])
            else:
                return([0, x[1]+1])
        else:
            return([2, x[1]+1])

    def make_obs(self,x):
        if x is None or self.dim == None:
            return x
        else:
            v = np.zeros(self.dim,dtype=int)
            v[3*self.horizon:] = np.random.binomial(1, 0.5, self.dim-3*self.horizon)
            v[3*x[1]+x[0]] = 1
            return (v)

    def start(self):
        return [0,0]

    def reward(self, x, a):
        if (x == [0,self.horizon-1] and a == self.opt_a[x[1]]) or (x == [1,self.horizon-1] and a == self.opt_b[x[1]]):
            return np.random.binomial(1, 0.5)
        return 0

    def is_tabular(self):
        return (self.dim == None)

class RandomGridWorld(Environment):
    """
    A M x M grid with walls and a trembling hand.
    Horizon is always 2 M
    """
    def __init__(self, M, swap=0.1, dim=2, noise=0.0):
        self.M = M
        self.swap = swap
        self.noise = noise
        self.dim = dim
        self.seed = 147
        np.random.seed(self.seed)
        self.goal = None

        self.maze = self.generate_maze(self.M)
        self.state = None

        self.actions = [(0,1), (0,-1), (1,0), (-1,0)]
        print("ENV: Generated Random Grid World")
        print("Size: %dx%d, Start: [%d,%d], Goal: [%d,%d], H: %d, Cells: %d" % (self.M, self.M, 0, 0, self.goal[1], self.goal[0], self.horizon, np.count_nonzero(self.maze)))
        self.print_maze()

    def is_tabular(self):
        return False

    def generate_maze(self, M):
        """
        Adapted from http://code.activestate.com/recipes/578356-random-maze-generator/
        """
        mx = M; my = M
        maze = np.matrix(np.zeros((mx,my)))
        dx = [0, 1, 0, -1]; dy = [-1, 0, 1, 0] # 4 directions to move in the maze
        stack = [(0,0)]

        while len(stack) > 0:
            (cx, cy) = stack[-1]
            maze[cy,cx] = 1
            # find a new cell to add
            nlst = [] # list of available neighbors
            for i in range(4):
                nx = cx + dx[i]; ny = cy + dy[i]
                if nx >= 0 and nx < mx and ny >= 0 and ny < my:
                    if maze[ny,nx] == 0:
                        # of occupied neighbors must be 1
                        ctr = 0
                        for j in range(4):
                            ex = nx + dx[j]; ey = ny + dy[j]
                            if ex >= 0 and ex < mx and ey >= 0 and ey < my:
                                if maze[ey,ex] == 1: ctr += 1
                        if ctr == 1: nlst.append(i)
            # if 1 or more neighbors available then randomly select one and move
            if len(nlst) > 0:
                ir = nlst[np.random.randint(0, len(nlst))]
                cx += dx[ir]; cy += dy[ir]
                stack.append((cx, cy))
            elif self.goal is None:
                self.horizon = 2*len(stack)
                (gx,gy) = stack.pop()
                self.goal = [gx,gy]
            else:
                stack.pop()

        return(maze)

    def make_obs(self,s):
        if s is None:
            return None
        tmp = s.copy()
        v = np.random.normal(0, self.noise, self.dim)
        if self.dim > 2:
            mult = np.random.choice(range(self.M), size=self.dim-2)
            tmp.extend(mult)
        return tmp+v

    def start(self):
        return [0,0]

    def reward(self,x,a):
        if x == self.goal:
            return 1
        return 0

    def transition(self, x, a):
        if x == self.goal:
            return None
        nx = x[0]+a[0]
        ny = x[1]+a[1]
        if nx < 0 or nx >= self.M or ny < 0 or ny >= self.M:
            ## Cannot go off the grid
            return x
        if self.maze[ny,nx] == 0:
            ## Cannot enter a wall
            return x
        else:
            z = np.random.binomial(1, self.swap)
            if z == 1:
                return x
            else:
                return [nx,ny]

    def print_maze(self):
        maze = self.maze
        for i in range(self.M):
            for j in range(self.M):
                if maze[i,j] == 0:
                    print(" ",end="")
                elif self.state is not None and self.state[0]==j and self.state[1]==i:
                    print("A",end="")
                elif self.goal[0]==j and self.goal[1]==i:
                    print("G",end="")
                elif maze[i,j] == 1:
                    print(".",end="")
            print("")

if __name__=='__main__':
    import sys, argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--env', action='store', choices=['MAB', 'combolock', 'stochcombolock', 'maze'])
    Args = parser.parse_args(sys.argv[1:])

    if Args.env == 'MAB':
        E = Environment()
        rewards = [0,0]
        counts = [0,0]
        for t in range(1000):
            x = E.start_episode()
            while x != None:
                actions = E.get_actions()
                a = np.random.choice(actions)
                (x,r) = E.act(a)
                rewards[a] += r
                counts[a] += 1
        for a in [0,1]:
            assert (np.abs(np.float(rewards[a])/counts[a] -E.reward_dists[a]) < 0.1)

    if Args.env == 'combolock':
        E = CombinationLock(horizon=3)
        print (E.opt)
        for t in range(10):
            x = E.start_episode()
            while x != None:
                actions = E.get_actions()
                a = np.random.choice(actions)
                old = x
                (x,r) = E.act(a)
                print(old, a, r, x)

    if Args.env == 'stochcombolock':
        E = StochasticCombinationLock(horizon=3, swap=0.5)
        print (E.opt_a)
        print (E.opt_b)
        for t in range(10):
            x = E.start_episode()
            while x != None:
                actions = E.get_actions()
                a = np.random.choice(actions)
                old = x
                (x,r) = E.act(a)
                print(old, a, r, x)

    if Args.env == 'maze':
        E = RandomGridWorld(M=3,swap=0.1, dim=2, noise=0.0)
        T = 0
        while True:
            T += 1
            x = E.start_episode()
            if T % 100 == 0:
                print("Iteration t = %d" % (T))
            while x is not None:
                E.print_maze()
                print(x)
                actions = E.get_actions()
                a = np.random.choice(len(actions))
                (x,r) = E.act(actions[a])
                if r == 1:
                    print("Success: T = %d" % (T))
                    sys.exit()
<|endoftext|>


<commit_before>"""Module for Twitch bot and threaded logging.
"""
import logging
import time
import traceback
from collections import defaultdict

import bot.commands
import bot.emotecounter
import bot.ranking
from bot.data_sources.config import ConfigSource
from bot.data_sources.emotes import EmoteSource
from bot.data_sources.twitch import TwitchSource
from bot.utilities.permission import Permission
from bot.utilities.tools import replace_vars
from bot.utilities.tools import sanitize_user_name
from bot.utilities.webcache import WebCache

CACHE_DURATION = 10800

class TwitchBot:
    """TwitchBot extends the IRCClient to interact with Twitch.tv."""

    def __init__(self, root):
        """Initialize bot."""
        self.root = root
        self.irc = None
        self.cache = WebCache(duration=CACHE_DURATION)  # 3 hours

        # Sources
        self.emotes, self.twitch, self.config = self.load_sources()

        # Commands
        self.commands = []
        self.games, self.passivegames = self.load_commands()

        # States
        self.last_warning = defaultdict(int)
        self.host_target = False
        self.pause = False
        self.game_running = False
        self.antispeech = False
        self.pyramid_block = False
        self.last_plebcmd = time.time() - self.config.pleb_cooldowntime
        self.last_plebgame = time.time() - self.config.pleb_gametimer

        # Emote Counter
        self.ecount = bot.emotecounter.EmoteCounterForBot(self)
        self.ecount.start_cpm()
        self.ranking = bot.ranking.Ranking(self)

        # User Groups
        self.mods = set()
        self.trusted_mods = set()
        self.subs = set()
        self.users = self.twitch.get_chatters()

    def load_sources(self):
        """Reloads data sources."""
        config = ConfigSource(self.root, self.cache)
        emotes = EmoteSource(
            config.channel,
            cache=self.cache,
            twitch_api_headers=config.twitch_api_headers,
        )
        twitch = TwitchSource(
            config.channel,
            cache=self.cache,
            twitch_api_headers=config.twitch_api_headers,
        )
        return emotes, twitch, config

    def reload(self):
        """Reloads sources (and therefore entire bot)."""
        logging.warning("Reloading bot!")
        self.close_commands()
        self.emotes, self.twitch, self.config = self.load_sources()
        self.reload_commands()

    def reload_commands(self):
        """Reloads variables."""
        self.games, self.passivegames = self.load_commands()

    def load_commands(self):
        """Reloads reloadable commands.

        For hard reload, set bot.commands = [] prior to calling this.
        """
        logging.warning("Reloading commands")

        # Reload commands
        self.close_commands()
        games = []
        passivegames = []

        if not self.commands:
            for cmd in bot.commands.commands:
                self.commands.append(cmd(self))
        else:
            for i, cmd in enumerate(self.commands):
                reloadable = True
                for non_reloadable_class in bot.commands.non_reload:
                    if cmd.__class__ == non_reloadable_class:
                        reloadable = False
                if reloadable:
                    self.commands[i] = cmd.__class__(self)

        for cmd in self.commands:
            if cmd.__class__ in bot.commands.games:
                games.append(cmd)
            if cmd.__class__ in bot.commands.passivegames:
                passivegames.append(cmd)
        return games, passivegames

    @staticmethod
    def mode_changed(_, channel, added, __, args):
        """Update IRC mod list when mod joins or leaves. Seems not useful."""
        change = "added" if added else "removed"
        info_msg = "[{}] IRC Mod {}: {}".format(channel, change, ", ".join(args))
        logging.warning(info_msg)

    def handle_whisper(self, sender, content):
        """Entry point for all incoming whisper messages to bot."""
        # currently do nothing at all
        print(
            "{} sent me [{}] this in a whisper: {}".format(
                sender, self.config.nickname, content
            )
        )

    def set_host(self, channel, target):
        """React to a channel being hosted."""
        if target == "-":
            self.host_target = None
            logging.warning("[{}] Exited host mode".format(channel))
        else:
            self.host_target = target
            logging.warning("[{}] Now hosting {}".format(channel, target))

    def user_state(self, prefix, tags):
        """Track user tags."""
        # NOTE: params in PRIVMSG() are already processed and does not contains these data,
        # so we have to get them from lineReceived() -> manually called userState() to parse the tags.
        # Also I don't want to crash the original PRIVMSG() functions by modifing them

        twitch_user_tag = prefix.split("!")[0]  # also known as login id

        name = sanitize_user_name(twitch_user_tag)
        if "subscriber" in tags:
            if tags["subscriber"] == "1":
                self.subs.add(name)
            elif name in self.subs:
                self.subs.discard(name)

        if "user-type" in tags:
            # This also works #if tags['user-type'] == 'mod':
            if tags["mod"] == "1":
                self.mods.add(name)
            elif name in self.mods:
                self.mods.discard(name)

    def get_permission(self, user):
        """Return the users permission level."""
        if user in self.config.owner_list:
            return Permission.Admin
        elif user in self.mods:
            return Permission.Moderator
        elif user in self.subs:
            return Permission.Subscriber
        return Permission.User

    def select_commands(self, perm):
        """If a game is active and plebcommands on cooldown, only iterate through game list.

        If no game is active only allow 'passive games' a.k.a PyramidGame
        """
        if perm == 0:
            if time.time() - self.last_plebcmd < self.config.pleb_cooldowntime:
                if self.game_running:
                    return self.games
                else:
                    return self.passivegames
            else:
                return self.commands
        else:
            return self.commands

    def process_command(self, user, msg, tag_info):
        """Process messages and call commands."""
        # Ignore messages by ignored user
        if user in self.config.ignored_users:
            return

        self.ranking.increment_points(user, 1, self)

        # Check if bot is paused
        if (
            self.pause
            and user not in self.config.owner_list
            and user not in self.trusted_mods
        ):
            return

        perm_levels = ["User", "Subscriber", "Moderator", "Owner"]
        perm = self.get_permission(user)
        msg = msg.strip()

        """Emote Count Function"""
        self.ecount.process_message(msg)

        """Limit pleb bot spam. Only allow certain commands to be processed by plebs, if plebcmds on cooldown."""
        cmdlist = self.select_commands(perm)

        # Flip through commands and execute everyone that matches.
        # Check if user has permission to execute command.
        # Also reduce warning message spam by limiting it to one per minute.
        for cmd in cmdlist:
            try:
                match = cmd.match(self, user, msg, tag_info)
                if not match:
                    continue
                cname = cmd.__class__.__name__
                if perm < cmd.perm:
                    if time.time() - self.last_warning[cname] < 60:
                        continue
                    self.last_warning[cname] = int(time.time())
                    reply = "{}: You don't have access to that command. Minimum level is {}."
                    self.write(reply.format(user, perm_levels[cmd.perm]))
                else:
                    if (
                        perm == 0 and cmd not in self.games
                    ):  # Only reset plebtimer if no game was played
                        self.last_plebcmd = time.time()
                    cmd.run(self, user, msg, tag_info)
            except (ValueError, TypeError):  # Not sure which Errors might happen here.
                logging.error(traceback.format_exc())
        """Reset antispeech for next command"""
        self.antispeech = False

    # functions of Twitch customized IRC messages -- check USERNOTICE in multibot_irc_cilent
    def incoming_raid(self, tags):
        """Send a message when channel got raided."""
        # NOTE: raid seems to have no custom message currently, so msg is not passed

        # log raid?
        logging.warning("New raid detected")
        logging.warning(tags["system-msg"])

        amount = int(tags["msg-param-viewerCount"])
        channel = tags["msg-param-displayName"]

        if amount < self.config.raid_announce_treshold:
            return

        responses = self.config.responses["usernotice"]["raid"]
        var = {"<AMOUNT>": amount, "<CHANNEL>": channel}
        reply = replace_vars(responses["msg"], var)

        self.write(reply)

    def incoming_ritual(self, tags, msg):
        """Execute some obscure not well documented ritual."""
        # leave it here for a while, the IRC example message in Twitch API page seems incorrect already
        pass

    def sub_gift(self, tags):
        """Send a message when someone gift a sub to another viewer."""
        # Auto-message generated by Twitch - "XXX gifted a $x.xx sub to xxx!"
        logging.warning("New sub donation detected")
        logging.warning(tags["system-msg"])

        # Setup message to be printed
        responses = self.config.responses["usernotice"]["subgift"]
        plan = responses["subplan"]["msg"]

        donor = tags["display-name"] if tags["display-name"] else tags["login"]
        recipient = (
            tags["msg-param-recipient-display-name"]
            if tags["msg-param-recipient-display-name"]
            else tags["msg-param-recipient-user-name"]
        )

        months = int(tags["msg-param-months"])
        subtype = tags["msg-param-sub-plan"]

        if int(months) <= 1:
            var = {
                "<DONOR>": donor,
                "<RECIPIENT>": recipient,
                "<SUBPLAN>": plan[subtype],
            }
            reply = replace_vars(responses["msg_standard"]["msg"], var)
        else:
            var = {
                "<DONOR>": donor,
                "<RECIPIENT>": recipient,
                "<SUBPLAN>": plan[subtype],
                "<MONTHS>": months,
            }
            reply = replace_vars(responses["msg_with_months"]["msg"], var)

        self.write(reply)

    def sub_message(self, tags, _):
        """Send a message when someone subscribes."""
        # Sub message by user is in subMsg, which is not part of IRC tags
        # Auto-message generated by Twitch - "XXX has subscribed for n months"
        logging.warning("New sub detected")
        logging.warning(tags["system-msg"])

        # Setup message to be printed by bot
        responses = self.config.responses["usernotice"]["sub"]
        plan = responses["subplan"]["msg"]

        # according to Twitch doc, display-name can be empty if it is never set
        user = tags["display-name"] if tags["display-name"] else tags["login"]
        months = tags["msg-param-months"]
        subtype = tags["msg-param-sub-plan"]

        if int(months) <= 1:
            var = {"<USER>": user, "<SUBPLAN>": plan[subtype]}
            reply = replace_vars(responses["msg_standard"]["msg"], var)
        else:
            var = {"<USER>": user, "<SUBPLAN>": plan[subtype], "<MONTHS>": months}
            reply = replace_vars(responses["msg_with_months"]["msg"], var)

        self.write(reply)

    def close_commands(self):
        """Gracefully end commands."""
        for cmd in self.commands:
            try:
                cmd.close(self)
            except (TypeError, ValueError):  # Not sure which Errors might happen here.
                logging.error(traceback.format_exc())

    def terminate(self):
        """Terminate bot."""
        self.close_commands()

    def access_to_emote(self, username, emote):
        """Check if user has access to a certain emote."""
        # Some notes about emotes and IRC:
        # emote tag in IRC message is like emotes=1902:0-4,6-10/1901:12-16; (not sure about order)
        # The message is "Keepo Keepo Kippa"
        # The format is like emotes=[emote_id]:[pos_start]-[pos_end],.../[another emote_id]:[pos_start]-[pos_end]...;
        #
        # Twitch internally parse your message and change them to emotes with the emote tag in IRC message
        #
        user_id = self.twitch.get_user_id(username)
        emotelist = self.emotes.get_user_emotes(user_id)
        for sets in emotelist:
            for key in range(0, len(emotelist[sets])):
                if emote == emotelist[sets][key]["code"]:
                    return True
        return False

    def clear_cache(self):
        """Clear the cache."""
        self.cache = WebCache(duration=CACHE_DURATION)

    def write(self, msg):
        """Write a message."""
        #  print("Fake print message: ", msg)
        #  return
        if self.irc is not None:
            self.irc.write(self.config.channel, msg)
        else:
            logging.warning(
                "The bot {} in channel {} wanted to say something, but irc isn't set.".format(
                    self.config.nickname, self.config.channel
                )
            )

    def whisper(self, msg, user):
        """Whisper a message to a user."""
        whisper = "/w {} {}".format(user, msg)
        if self.irc is not None:
            self.irc.write(self.config.channel, whisper)
        else:
            logging.warning(
                "The bot {} in channel {} wanted to whisper to {}, but irc isn't set.".format(
                    self.config.nickname, self.config.channel, user
                )
            )

    def timeout(self, user, duration):
        """Timout a user for a certain time in the channel."""
        timeout = "/timeout {} {}".format(user, duration)
        if self.irc is not None:
            self.irc.write(self.config.channel, timeout)
        else:
            logging.warning(
                "The bot {} in channel {} wanted to timout {}, but irc isn't set.".format(
                    self.config.nickname, self.config.channel, user
                )
            )

    def ban(self, user):
        """Ban a user from the channel."""
        ban = "/ban {}".format(user)
        if self.irc is not None:
            self.irc.write(self.config.channel, ban)
        else:
            logging.warning(
                "The bot {} in channel {} wanted to ban {}, but irc isn't set.".format(
                    self.config.nickname, self.config.channel, user
                )
            )

    def unban(self, user):
        """Unban a user for the channel."""
        unban = "/unban {}".format(user)
        if self.irc is not None:
            self.irc.write(self.config.channel, unban)
        else:
            logging.warning(
                "The bot {} in channel {} wanted to unban {}, but irc isn't set.".format(
                    self.config.nickname, self.config.channel, user
                )
            )
<commit_msg>fix: check correct trusted_mods list for sleep<commit_after>"""Module for Twitch bot and threaded logging.
"""
import logging
import time
import traceback
from collections import defaultdict

import bot.commands
import bot.emotecounter
import bot.ranking
from bot.data_sources.config import ConfigSource
from bot.data_sources.emotes import EmoteSource
from bot.data_sources.twitch import TwitchSource
from bot.utilities.permission import Permission
from bot.utilities.tools import replace_vars
from bot.utilities.tools import sanitize_user_name
from bot.utilities.webcache import WebCache

CACHE_DURATION = 10800

class TwitchBot:
    """TwitchBot extends the IRCClient to interact with Twitch.tv."""

    def __init__(self, root):
        """Initialize bot."""
        self.root = root
        self.irc = None
        self.cache = WebCache(duration=CACHE_DURATION)  # 3 hours

        # Sources
        self.emotes, self.twitch, self.config = self.load_sources()

        # Commands
        self.commands = []
        self.games, self.passivegames = self.load_commands()

        # States
        self.last_warning = defaultdict(int)
        self.host_target = False
        self.pause = False
        self.game_running = False
        self.antispeech = False
        self.pyramid_block = False
        self.last_plebcmd = time.time() - self.config.pleb_cooldowntime
        self.last_plebgame = time.time() - self.config.pleb_gametimer

        # Emote Counter
        self.ecount = bot.emotecounter.EmoteCounterForBot(self)
        self.ecount.start_cpm()
        self.ranking = bot.ranking.Ranking(self)

        # User Groups
        self.mods = set()
        self.subs = set()
        self.users = self.twitch.get_chatters()

    def load_sources(self):
        """Reloads data sources."""
        config = ConfigSource(self.root, self.cache)
        emotes = EmoteSource(
            config.channel,
            cache=self.cache,
            twitch_api_headers=config.twitch_api_headers,
        )
        twitch = TwitchSource(
            config.channel,
            cache=self.cache,
            twitch_api_headers=config.twitch_api_headers,
        )
        return emotes, twitch, config

    def reload(self):
        """Reloads sources (and therefore entire bot)."""
        logging.warning("Reloading bot!")
        self.close_commands()
        self.emotes, self.twitch, self.config = self.load_sources()
        self.reload_commands()

    def reload_commands(self):
        """Reloads variables."""
        self.games, self.passivegames = self.load_commands()

    def load_commands(self):
        """Reloads reloadable commands.

        For hard reload, set bot.commands = [] prior to calling this.
        """
        logging.warning("Reloading commands")

        # Reload commands
        self.close_commands()
        games = []
        passivegames = []

        if not self.commands:
            for cmd in bot.commands.commands:
                self.commands.append(cmd(self))
        else:
            for i, cmd in enumerate(self.commands):
                reloadable = True
                for non_reloadable_class in bot.commands.non_reload:
                    if cmd.__class__ == non_reloadable_class:
                        reloadable = False
                if reloadable:
                    self.commands[i] = cmd.__class__(self)

        for cmd in self.commands:
            if cmd.__class__ in bot.commands.games:
                games.append(cmd)
            if cmd.__class__ in bot.commands.passivegames:
                passivegames.append(cmd)
        return games, passivegames

    @staticmethod
    def mode_changed(_, channel, added, __, args):
        """Update IRC mod list when mod joins or leaves. Seems not useful."""
        change = "added" if added else "removed"
        info_msg = "[{}] IRC Mod {}: {}".format(channel, change, ", ".join(args))
        logging.warning(info_msg)

    def handle_whisper(self, sender, content):
        """Entry point for all incoming whisper messages to bot."""
        # currently do nothing at all
        print(
            "{} sent me [{}] this in a whisper: {}".format(
                sender, self.config.nickname, content
            )
        )

    def set_host(self, channel, target):
        """React to a channel being hosted."""
        if target == "-":
            self.host_target = None
            logging.warning("[{}] Exited host mode".format(channel))
        else:
            self.host_target = target
            logging.warning("[{}] Now hosting {}".format(channel, target))

    def user_state(self, prefix, tags):
        """Track user tags."""
        # NOTE: params in PRIVMSG() are already processed and does not contains these data,
        # so we have to get them from lineReceived() -> manually called userState() to parse the tags.
        # Also I don't want to crash the original PRIVMSG() functions by modifing them

        twitch_user_tag = prefix.split("!")[0]  # also known as login id

        name = sanitize_user_name(twitch_user_tag)
        if "subscriber" in tags:
            if tags["subscriber"] == "1":
                self.subs.add(name)
            elif name in self.subs:
                self.subs.discard(name)

        if "user-type" in tags:
            # This also works #if tags['user-type'] == 'mod':
            if tags["mod"] == "1":
                self.mods.add(name)
            elif name in self.mods:
                self.mods.discard(name)

    def get_permission(self, user):
        """Return the users permission level."""
        if user in self.config.owner_list:
            return Permission.Admin
        elif user in self.mods:
            return Permission.Moderator
        elif user in self.subs:
            return Permission.Subscriber
        return Permission.User

    def select_commands(self, perm):
        """If a game is active and plebcommands on cooldown, only iterate through game list.

        If no game is active only allow 'passive games' a.k.a PyramidGame
        """
        if perm == 0:
            if time.time() - self.last_plebcmd < self.config.pleb_cooldowntime:
                if self.game_running:
                    return self.games
                else:
                    return self.passivegames
            else:
                return self.commands
        else:
            return self.commands

    def process_command(self, user, msg, tag_info):
        """Process messages and call commands."""
        # Ignore messages by ignored user
        if user in self.config.ignored_users:
            return

        self.ranking.increment_points(user, 1, self)

        # Check if bot is paused
        if (
            self.pause
            and user not in self.config.owner_list
            and user not in self.config.trusted_mods
        ):
            return

        perm_levels = ["User", "Subscriber", "Moderator", "Owner"]
        perm = self.get_permission(user)
        msg = msg.strip()

        """Emote Count Function"""
        self.ecount.process_message(msg)

        """Limit pleb bot spam. Only allow certain commands to be processed by plebs, if plebcmds on cooldown."""
        cmdlist = self.select_commands(perm)

        # Flip through commands and execute everyone that matches.
        # Check if user has permission to execute command.
        # Also reduce warning message spam by limiting it to one per minute.
        for cmd in cmdlist:
            try:
                match = cmd.match(self, user, msg, tag_info)
                if not match:
                    continue
                cname = cmd.__class__.__name__
                if perm < cmd.perm:
                    if time.time() - self.last_warning[cname] < 60:
                        continue
                    self.last_warning[cname] = int(time.time())
                    reply = "{}: You don't have access to that command. Minimum level is {}."
                    self.write(reply.format(user, perm_levels[cmd.perm]))
                else:
                    if (
                        perm == 0 and cmd not in self.games
                    ):  # Only reset plebtimer if no game was played
                        self.last_plebcmd = time.time()
                    cmd.run(self, user, msg, tag_info)
            except (ValueError, TypeError):  # Not sure which Errors might happen here.
                logging.error(traceback.format_exc())
        """Reset antispeech for next command"""
        self.antispeech = False

    # functions of Twitch customized IRC messages -- check USERNOTICE in multibot_irc_cilent
    def incoming_raid(self, tags):
        """Send a message when channel got raided."""
        # NOTE: raid seems to have no custom message currently, so msg is not passed

        # log raid?
        logging.warning("New raid detected")
        logging.warning(tags["system-msg"])

        amount = int(tags["msg-param-viewerCount"])
        channel = tags["msg-param-displayName"]

        if amount < self.config.raid_announce_treshold:
            return

        responses = self.config.responses["usernotice"]["raid"]
        var = {"<AMOUNT>": amount, "<CHANNEL>": channel}
        reply = replace_vars(responses["msg"], var)

        self.write(reply)

    def incoming_ritual(self, tags, msg):
        """Execute some obscure not well documented ritual."""
        # leave it here for a while, the IRC example message in Twitch API page seems incorrect already
        pass

    def sub_gift(self, tags):
        """Send a message when someone gift a sub to another viewer."""
        # Auto-message generated by Twitch - "XXX gifted a $x.xx sub to xxx!"
        logging.warning("New sub donation detected")
        logging.warning(tags["system-msg"])

        # Setup message to be printed
        responses = self.config.responses["usernotice"]["subgift"]
        plan = responses["subplan"]["msg"]

        donor = tags["display-name"] if tags["display-name"] else tags["login"]
        recipient = (
            tags["msg-param-recipient-display-name"]
            if tags["msg-param-recipient-display-name"]
            else tags["msg-param-recipient-user-name"]
        )

        months = int(tags["msg-param-months"])
        subtype = tags["msg-param-sub-plan"]

        if int(months) <= 1:
            var = {
                "<DONOR>": donor,
                "<RECIPIENT>": recipient,
                "<SUBPLAN>": plan[subtype],
            }
            reply = replace_vars(responses["msg_standard"]["msg"], var)
        else:
            var = {
                "<DONOR>": donor,
                "<RECIPIENT>": recipient,
                "<SUBPLAN>": plan[subtype],
                "<MONTHS>": months,
            }
            reply = replace_vars(responses["msg_with_months"]["msg"], var)

        self.write(reply)

    def sub_message(self, tags, _):
        """Send a message when someone subscribes."""
        # Sub message by user is in subMsg, which is not part of IRC tags
        # Auto-message generated by Twitch - "XXX has subscribed for n months"
        logging.warning("New sub detected")
        logging.warning(tags["system-msg"])

        # Setup message to be printed by bot
        responses = self.config.responses["usernotice"]["sub"]
        plan = responses["subplan"]["msg"]

        # according to Twitch doc, display-name can be empty if it is never set
        user = tags["display-name"] if tags["display-name"] else tags["login"]
        months = tags["msg-param-months"]
        subtype = tags["msg-param-sub-plan"]

        if int(months) <= 1:
            var = {"<USER>": user, "<SUBPLAN>": plan[subtype]}
            reply = replace_vars(responses["msg_standard"]["msg"], var)
        else:
            var = {"<USER>": user, "<SUBPLAN>": plan[subtype], "<MONTHS>": months}
            reply = replace_vars(responses["msg_with_months"]["msg"], var)

        self.write(reply)

    def close_commands(self):
        """Gracefully end commands."""
        for cmd in self.commands:
            try:
                cmd.close(self)
            except (TypeError, ValueError):  # Not sure which Errors might happen here.
                logging.error(traceback.format_exc())

    def terminate(self):
        """Terminate bot."""
        self.close_commands()

    def access_to_emote(self, username, emote):
        """Check if user has access to a certain emote."""
        # Some notes about emotes and IRC:
        # emote tag in IRC message is like emotes=1902:0-4,6-10/1901:12-16; (not sure about order)
        # The message is "Keepo Keepo Kippa"
        # The format is like emotes=[emote_id]:[pos_start]-[pos_end],.../[another emote_id]:[pos_start]-[pos_end]...;
        #
        # Twitch internally parse your message and change them to emotes with the emote tag in IRC message
        #
        user_id = self.twitch.get_user_id(username)
        emotelist = self.emotes.get_user_emotes(user_id)
        for sets in emotelist:
            for key in range(0, len(emotelist[sets])):
                if emote == emotelist[sets][key]["code"]:
                    return True
        return False

    def clear_cache(self):
        """Clear the cache."""
        self.cache = WebCache(duration=CACHE_DURATION)

    def write(self, msg):
        """Write a message."""
        #  print("Fake print message: ", msg)
        #  return
        if self.irc is not None:
            self.irc.write(self.config.channel, msg)
        else:
            logging.warning(
                "The bot {} in channel {} wanted to say something, but irc isn't set.".format(
                    self.config.nickname, self.config.channel
                )
            )

    def whisper(self, msg, user):
        """Whisper a message to a user."""
        whisper = "/w {} {}".format(user, msg)
        if self.irc is not None:
            self.irc.write(self.config.channel, whisper)
        else:
            logging.warning(
                "The bot {} in channel {} wanted to whisper to {}, but irc isn't set.".format(
                    self.config.nickname, self.config.channel, user
                )
            )

    def timeout(self, user, duration):
        """Timout a user for a certain time in the channel."""
        timeout = "/timeout {} {}".format(user, duration)
        if self.irc is not None:
            self.irc.write(self.config.channel, timeout)
        else:
            logging.warning(
                "The bot {} in channel {} wanted to timout {}, but irc isn't set.".format(
                    self.config.nickname, self.config.channel, user
                )
            )

    def ban(self, user):
        """Ban a user from the channel."""
        ban = "/ban {}".format(user)
        if self.irc is not None:
            self.irc.write(self.config.channel, ban)
        else:
            logging.warning(
                "The bot {} in channel {} wanted to ban {}, but irc isn't set.".format(
                    self.config.nickname, self.config.channel, user
                )
            )

    def unban(self, user):
        """Unban a user for the channel."""
        unban = "/unban {}".format(user)
        if self.irc is not None:
            self.irc.write(self.config.channel, unban)
        else:
            logging.warning(
                "The bot {} in channel {} wanted to unban {}, but irc isn't set.".format(
                    self.config.nickname, self.config.channel, user
                )
            )
<|endoftext|>


<commit_before>import socket
import ncs
import ncs.maagic as mag

def cleanup_or_search(box, to_del):
    """
    A function that deletes all of the unused object groups in the necessary
    order from the inputted device. It returns:
    1. ret: A dictionary with all of the deleted object groups organized by type.
    2. stat: A string that contains the status of the deletion (Success or specific ncs error).
    """
    #Initializing python lists and sets
    og_list = []
    og_typ = []
    acl_list = []
    ret = {}
    used_group_ogs = set()
    orphaned_ogs = set()
    orphaned_dict = {}
    delete_first = set()

    #Creating transaction and setting root to access NSO
    with ncs.maapi.single_write_trans('ncsadmin', 'python', groups=['ncsadmin']) as tran:
        root = ncs.maagic.get_root(tran)

        #Converting access lists and object group list to python lists
        og_list, og_typ, acl_list = nso_to_python(box, root, og_list, og_typ, acl_list)

        #Finding the unused object groups and placing them in sets for prioritization
        orphaned_ogs, delete_first, used_group_ogs, orphaned_dict = find_orphaned_og(root,
        box, og_list, og_typ, acl_list, used_group_ogs, orphaned_ogs, delete_first, orphaned_dict)

        #Prioritize which orphaned object groups need to be deleted first
        delete_first, delete_second = prioritize_del(orphaned_ogs, used_group_ogs, delete_first)

        #Delete the first group and then the second group
        ret = del_1(root, box, delete_first, orphaned_dict, ret, to_del)
        ret = del_2(root, box, delete_second, orphaned_dict, ret, to_del)

        #For cleanup
        if to_del:
            #Apply the changes and show NSO errors to the stat output
            try:
                tran.apply()
                stat = "Success"

            except mag.MaagicError, err:
                stat = mag.MaagicError, err

            return ret, stat
    #For search
    return ret

def nso_to_python(box, root, og_list, og_typ, acl_list):
    """
    A function that converts the inputted device's object group list and access
    lists to python lists for optimization. It returns:
    1. og_list: A python list with object group names.
    2. og_typ:  A python list with the object group type of each object group
                which is alligned with the og_list.
    3. acl_list:A python list that contains access lists that contain the rules
                of its respective access list.
    """
    #Adding all of the object groups and their types to python lists
    for ogtyp in root.devices.device[box].config.asa__object_group:
        for og in root.devices.device[box].config.asa__object_group[ogtyp]:
            og_list.append(og.id)
            og_typ.append(str(og))      #str(og) is the object group type

    #Adding each access list's rules to a python list (temp_rul_list) and
    #then adding those lists as elements of another python list (acl_list)
    for acl in root.devices.device[box].config.asa__access_list.access_list_id:
        temp_rul_list = []
        for rul in root.devices.device[box].config.asa__access_list.access_list_id[acl.id].rule:
            if "object-group" in rul.id:
                temp_rul_list.append(rul.id)
        acl_list.append(temp_rul_list)

    return og_list, og_typ, acl_list

def rec_group_og(used_group_ogs, root, box, og, typ):
    """
    A recursive function that adds a used object group's group-objects (object groups inside
    of object groups) to a set and recursively recalls itself to check if the group objects
    added have any group objects themselves to add as well. It returns:
    1. used_group_ogs: Set of used group-objects that should not be removed.
    """
    if root.devices.device[box].config.asa__object_group[typ][og].group_object:
        for group_og in root.devices.device[box].config.asa__object_group[typ][og].group_object:
            used_group_ogs.add(group_og.id)
            used_group_ogs = rec_group_og(used_group_ogs, root, box, group_og.id, typ)

    return used_group_ogs

def find_orphaned_og(root, box, og_list, og_typ, acl_list, used_group_ogs, orphaned_ogs, delete_first, orphaned_dict):
    """
    A function that finds all of the unused object groups by iterating through all of the access
    lists for each object group, checking if the object group is in any of the rules of any of
    the access lists of the inputted device. It returns:
    1. orphaned_ogs:    A set of all of the unused object groups within the device.
    2. delete_first:    A set of the unused object groups that have group objects.
    3. used_group_ogs:  A set of all of the group-objects that are actually being used.
    4. orphaned_dict:   A dictionary of the unused object groups that each contains an inner dictionary
                        with the object group's type, group-objects, and flag that indicates whether
                        it's been deleted (True) or not (False).

    """
    #Iterating through both object group and object group type lists simultaneously
    for og, typ in zip(og_list, og_typ):
        #flag indicates whether the object group was found in an access list
        flag = False
        for acl in acl_list:
            for rule in acl:
                if og in rule:
                    used_group_ogs = rec_group_og(used_group_ogs, root, box, og, typ)
                    #object group was found, update flag and break
                    flag = True
                    break
            #If found, continue to the next object group
            if flag:
                break
        #If not found in any of the access lists, add to set and create dictionary of orphaned ogs
        if not flag:
            orphaned_ogs.add(og)
            inner_dict = {"og_type" : typ, "group_ogs" : [], "deleted" : False}
            #If the object group has group-objects, add to delete_first set
            if root.devices.device[box].config.asa__object_group[typ][og].group_object:
                delete_first.add(og)
                for group_og in root.devices.device[box].config.asa__object_group[typ][og].group_object:
                    inner_dict["group_ogs"].append(group_og.id)
            #Add inner dictionary to orphaned object group key in dictionary
            orphaned_dict[og] = inner_dict

    return orphaned_ogs, delete_first, used_group_ogs, orphaned_dict

def prioritize_del(orphaned_ogs, used_group_ogs, delete_first):
    """
    A function that removes the used group-objects from a set of all of the orphaned object
    groups and from the high priority deletion set, and creates a delete second set by removing
    the delete first set from the general orphaned object groups. It returns:
    1. delete_first: The orphaned object groups that have group-objects (highest priority to delete)
    2. delete_second: The orphaned object groups that don't have group-objects (NSO error if deleted first)
    """
    orphaned_ogs = orphaned_ogs.difference(used_group_ogs)
    delete_first = delete_first.difference(used_group_ogs)
    delete_second = orphaned_ogs.difference(delete_first)

    return delete_first, delete_second

def del_1(root, box, delete_first, orphaned_dict, ret, to_del):
    """
    A function that deletes the object groups that have group-objects in this specific order:
    object groups that are not group-objects for any other object group are deleted first (avoiding NSO error).
    It also adds all of the deleted object groups to a dictionary. It returns:
    1. ret: A dictionary organized by object group type with all of the deleted object groups.
    """
    #The number of object groups that are left to delete first
    del_f_count = len(delete_first)
    #While there are still object groups to delete within delete first
    while del_f_count:
        for og in delete_first:
            #If that object group has already been deleted, move on to the next object group
            if orphaned_dict[og]["deleted"]:
                continue
            #flag indicates whether the current object group is a group-object within another
            #object group inside delete first that has yet to be deleted
            flag = False
            for og2 in delete_first:
                if orphaned_dict[og2]["deleted"]:
                    continue
                #If the current object group is a group-object, set flag = True
                if og in orphaned_dict[og2]["group_ogs"]:
                    flag = True
                    break
            #If object group is not a group-object in a non-deleted delete first object group,
            #add to dictionary, decrement delete first count, set delete flag to True, and delete
            if not flag:
                if orphaned_dict[og]["og_type"] in ret.keys():
                    ret[orphaned_dict[og]["og_type"]].append(og)
                else:
                    ret[orphaned_dict[og]["og_type"]] = [og]

                del_f_count -= 1
                orphaned_dict[og]["deleted"] = True
                if to_del:
                    del root.devices.device[box].config.asa__object_group[orphaned_dict[og]["og_type"]][og]

    return ret

def del_2(root, box, delete_second, orphaned_dict, ret, to_del):
    """
    A function that deletes the remaining orphaned object groups and adds them to the dictionary
    that will be returned. It returns:
    1. ret: A dictionary organized by object group type with all of the deleted object groups.
    """
    for og in delete_second:
        if orphaned_dict[og]["og_type"] in ret.keys():
            ret[orphaned_dict[og]["og_type"]].append(og)
        else:
            ret[orphaned_dict[og]["og_type"]] = [og]
        if to_del:
            del root.devices.device[box].config.asa__object_group[orphaned_dict[og]["og_type"]][og]

    return ret

def remove_ogs(box, og_type, og_id):
    """
    A function that removes the object group from the object group list using
    the arguments passed: device name, object group name, and object group type.
    It returns:
    1. stat: Status of applying delete to device (Success or NSO error)
    """
    with ncs.maapi.single_write_trans('ncsadmin', 'python', groups=['ncsadmin']) as tran:
        root = ncs.maagic.get_root(tran)
        del root.devices.device[box].config.asa__object_group[og_type][og_id]
        try:
            tran.apply()
            stat = "Success"
        #Provides error message if there is a problem removing an OG
        except Exception, err:
            stat = type(err)

    return stat
<commit_msg>remove function device name list<commit_after>import socket
import ncs
import _ncs
import ncs.maagic as mag

def cleanup_or_search(box, to_del):
    """
    A function that deletes all of the unused object groups in the necessary
    order from the inputted device. It returns:
    1. ret: A dictionary with all of the deleted object groups organized by type.
    2. stat: A string that contains the status of the deletion (Success or specific ncs error).
    """
    #Initializing python lists and sets
    og_list = []
    og_typ = []
    acl_list = []
    ret = {}
    used_group_ogs = set()
    orphaned_ogs = set()
    orphaned_dict = {}
    delete_first = set()

    #Creating transaction and setting root to access NSO
    with ncs.maapi.single_write_trans('ncsadmin', 'python', groups=['ncsadmin']) as tran:
        root = ncs.maagic.get_root(tran)

        #Converting access lists and object group list to python lists
        og_list, og_typ, acl_list = nso_to_python(box, root, og_list, og_typ, acl_list)

        #Finding the unused object groups and placing them in sets for prioritization
        orphaned_ogs, delete_first, used_group_ogs, orphaned_dict = find_orphaned_og(root,
        box, og_list, og_typ, acl_list, used_group_ogs, orphaned_ogs, delete_first, orphaned_dict)

        #Prioritize which orphaned object groups need to be deleted first
        delete_first, delete_second = prioritize_del(orphaned_ogs, used_group_ogs, delete_first)

        #Delete the first group and then the second group
        ret = del_1(root, box, delete_first, orphaned_dict, ret, to_del)
        ret = del_2(root, box, delete_second, orphaned_dict, ret, to_del)

        #For cleanup
        if to_del:
            #Apply the changes and show NSO errors to the stat output
            try:
                tran.apply()
                stat = "Success"

            except mag.MaagicError, err:
                stat = mag.MaagicError, err

            return ret, stat
    #For search
    return ret

def nso_to_python(box, root, og_list, og_typ, acl_list):
    """
    A function that converts the inputted device's object group list and access
    lists to python lists for optimization. It returns:
    1. og_list: A python list with object group names.
    2. og_typ:  A python list with the object group type of each object group
                which is alligned with the og_list.
    3. acl_list:A python list that contains access lists that contain the rules
                of its respective access list.
    """
    #Adding all of the object groups and their types to python lists
    for ogtyp in root.devices.device[box].config.asa__object_group:
        for og in root.devices.device[box].config.asa__object_group[ogtyp]:
            og_list.append(og.id)
            og_typ.append(str(og))      #str(og) is the object group type

    #Adding each access list's rules to a python list (temp_rul_list) and
    #then adding those lists as elements of another python list (acl_list)
    for acl in root.devices.device[box].config.asa__access_list.access_list_id:
        temp_rul_list = []
        for rul in root.devices.device[box].config.asa__access_list.access_list_id[acl.id].rule:
            if "object-group" in rul.id:
                temp_rul_list.append(rul.id)
        acl_list.append(temp_rul_list)

    return og_list, og_typ, acl_list

def rec_group_og(used_group_ogs, root, box, og, typ):
    """
    A recursive function that adds a used object group's group-objects (object groups inside
    of object groups) to a set and recursively recalls itself to check if the group objects
    added have any group objects themselves to add as well. It returns:
    1. used_group_ogs: Set of used group-objects that should not be removed.
    """
    if root.devices.device[box].config.asa__object_group[typ][og].group_object:
        for group_og in root.devices.device[box].config.asa__object_group[typ][og].group_object:
            used_group_ogs.add(group_og.id)
            used_group_ogs = rec_group_og(used_group_ogs, root, box, group_og.id, typ)

    return used_group_ogs

def find_orphaned_og(root, box, og_list, og_typ, acl_list, used_group_ogs, orphaned_ogs, delete_first, orphaned_dict):
    """
    A function that finds all of the unused object groups by iterating through all of the access
    lists for each object group, checking if the object group is in any of the rules of any of
    the access lists of the inputted device. It returns:
    1. orphaned_ogs:    A set of all of the unused object groups within the device.
    2. delete_first:    A set of the unused object groups that have group objects.
    3. used_group_ogs:  A set of all of the group-objects that are actually being used.
    4. orphaned_dict:   A dictionary of the unused object groups that each contains an inner dictionary
                        with the object group's type, group-objects, and flag that indicates whether
                        it's been deleted (True) or not (False).

    """
    #Iterating through both object group and object group type lists simultaneously
    for og, typ in zip(og_list, og_typ):
        #flag indicates whether the object group was found in an access list
        flag = False
        for acl in acl_list:
            for rule in acl:
                if og in rule:
                    used_group_ogs = rec_group_og(used_group_ogs, root, box, og, typ)
                    #object group was found, update flag and break
                    flag = True
                    break
            #If found, continue to the next object group
            if flag:
                break
        #If not found in any of the access lists, add to set and create dictionary of orphaned ogs
        if not flag:
            orphaned_ogs.add(og)
            inner_dict = {"og_type" : typ, "group_ogs" : [], "deleted" : False}
            #If the object group has group-objects, add to delete_first set
            if root.devices.device[box].config.asa__object_group[typ][og].group_object:
                delete_first.add(og)
                for group_og in root.devices.device[box].config.asa__object_group[typ][og].group_object:
                    inner_dict["group_ogs"].append(group_og.id)
            #Add inner dictionary to orphaned object group key in dictionary
            orphaned_dict[og] = inner_dict

    return orphaned_ogs, delete_first, used_group_ogs, orphaned_dict

def prioritize_del(orphaned_ogs, used_group_ogs, delete_first):
    """
    A function that removes the used group-objects from a set of all of the orphaned object
    groups and from the high priority deletion set, and creates a delete second set by removing
    the delete first set from the general orphaned object groups. It returns:
    1. delete_first: The orphaned object groups that have group-objects (highest priority to delete)
    2. delete_second: The orphaned object groups that don't have group-objects (NSO error if deleted first)
    """
    orphaned_ogs = orphaned_ogs.difference(used_group_ogs)
    delete_first = delete_first.difference(used_group_ogs)
    delete_second = orphaned_ogs.difference(delete_first)

    return delete_first, delete_second

def del_1(root, box, delete_first, orphaned_dict, ret, to_del):
    """
    A function that deletes the object groups that have group-objects in this specific order:
    object groups that are not group-objects for any other object group are deleted first (avoiding NSO error).
    It also adds all of the deleted object groups to a dictionary. It returns:
    1. ret: A dictionary organized by object group type with all of the deleted object groups.
    """
    #The number of object groups that are left to delete first
    del_f_count = len(delete_first)
    #While there are still object groups to delete within delete first
    while del_f_count:
        for og in delete_first:
            #If that object group has already been deleted, move on to the next object group
            if orphaned_dict[og]["deleted"]:
                continue
            #flag indicates whether the current object group is a group-object within another
            #object group inside delete first that has yet to be deleted
            flag = False
            for og2 in delete_first:
                if orphaned_dict[og2]["deleted"]:
                    continue
                #If the current object group is a group-object, set flag = True
                if og in orphaned_dict[og2]["group_ogs"]:
                    flag = True
                    break
            #If object group is not a group-object in a non-deleted delete first object group,
            #add to dictionary, decrement delete first count, set delete flag to True, and delete
            if not flag:
                if orphaned_dict[og]["og_type"] in ret.keys():
                    ret[orphaned_dict[og]["og_type"]].append(og)
                else:
                    ret[orphaned_dict[og]["og_type"]] = [og]

                del_f_count -= 1
                orphaned_dict[og]["deleted"] = True
                if to_del:
                    del root.devices.device[box].config.asa__object_group[orphaned_dict[og]["og_type"]][og]

    return ret

def del_2(root, box, delete_second, orphaned_dict, ret, to_del):
    """
    A function that deletes the remaining orphaned object groups and adds them to the dictionary
    that will be returned. It returns:
    1. ret: A dictionary organized by object group type with all of the deleted object groups.
    """
    for og in delete_second:
        if orphaned_dict[og]["og_type"] in ret.keys():
            ret[orphaned_dict[og]["og_type"]].append(og)
        else:
            ret[orphaned_dict[og]["og_type"]] = [og]
        if to_del:
            del root.devices.device[box].config.asa__object_group[orphaned_dict[og]["og_type"]][og]

    return ret

def remove_ogs(box, og_type, og_id):
    """
    A function that removes the object group from the object group list using
    the arguments passed: device name, object group name, and object group type.
    It returns:
    1. stat: Status of applying delete to device (Success or NSO error)
    """
    with ncs.maapi.single_write_trans('ncsadmin', 'python', groups=['ncsadmin']) as tran:
        root = ncs.maagic.get_root(tran)
        del root.devices.device[box].config.asa__object_group[og_type][og_id]
        try:
            tran.apply()
            stat = "Success"
        #Provides error message if there is a problem removing an OG
    except _ncs.error.Error, err:
            stat = err

    return stat
<|endoftext|>


<commit_before>import hoomd
import hoomd.hpmc
from hoomd.hpmc import _hpmc
import numpy as np
import pytest
from hoomd.hpmc.pytest.conftest import *

def check_dict(shape_dict, args):
    """
    check_dict: Function to check that two dictionaries are equivalent

    Arguments: shape_dict and args - dictionaries to test

    Ex: mc = hoomd.hpmc.integrate.Sphere(23456)
        mc.shape["A"] = {"diameter": 1}
        check_dict(mc.shape["A"], {"diameter": 1})

    Useful for more complex nested dictionaries (like the shape key in unions)
    Used to test that the dictionary passed in is what gets passed out
    """
    for key, val in args.items():
        if isinstance(shape_dict[key], list) and len(shape_dict[key]) > 0 \
           and key != 'shapes':
            np.testing.assert_allclose(shape_dict[key], val)
        elif key == 'shapes':
            for i in range(len(shape_dict[key])):
                shape_args = shape_dict[key][i]
                val_args = val[i]
                for shape_key in shape_args:
                    if isinstance(shape_args[shape_key], list) \
                       and len(shape_args[shape_key]) > 0:
                        np.testing.assert_allclose(val_args[shape_key],
                                                   shape_args[shape_key])
                    else:
                        assert shape_args[shape_key] == val_args[shape_key]
        else:
            np.testing.assert_almost_equal(shape_dict[key], val)

def test_dict_conversion(shape_dict_conversion_args):
    for shape_params, args_list in shape_dict_conversion_args():
        if 'poly' not in str(shape_params).lower():
            for args in args_list:
                print("")
                print(args)
                print(shape_params)
                test_shape = shape_params(args)
                test_dict = test_shape.asDict()
                check_dict(test_dict, args)

def test_shape_params(integrator_args):
    for shape_integrator, valid_args, invalid_args in integrator_args():
        mc = shape_integrator(23456)
        for args in valid_args:
            mc.shape["A"] = args
            check_dict(mc.shape["A"], args)
        for args in invalid_args:
            print("")
            print(shape_integrator)
            print(args)
            with pytest.raises(Exception):
                mc.shape["A"] = args

def test_shape_attached(dummy_simulation_factory, integrator_args):
    for shape_integrator, valid_args, invalid_args in integrator_args():
        mc = shape_integrator(23456)
        for args in valid_args:
            mc.shape["A"] = args
            sim = dummy_simulation_factory()
            print("")
            print(shape_integrator)
            print(args)
            sim.operations.add(mc)
            sim.operations.schedule()
            check_dict(mc.shape["A"], args)

def test_moves(device, lattice_simulation_factory, integrator_args):
    dims = 3
    for shape_integrator, valid_args, invalid_args in integrator_args():
        if 'union' not in str(shape_integrator).lower():
            args = valid_args[0]
            if 'polygon' in str(shape_integrator).lower():
                dims = 2
            mc = shape_integrator(23456)
            mc.shape['A'] = args
            sim = lattice_simulation_factory(dimensions=dims)
            sim.operations.add(mc)
            sim.operations.schedule()
            sim.run(100)
            accepted_rejected_trans = sum(sim.operations.integrator.translate_moves)
            assert accepted_rejected_trans > 0
            if 'sphere' not in str(shape_integrator).lower():
                accepted_rejected_rot = sum(sim.operations.integrator.rotate_moves)
                assert accepted_rejected_rot > 0

def test_overlaps_sphere(device, lattice_simulation_factory):
    mc = hoomd.hpmc.integrate.Sphere(23456)
    mc.shape["A"] = {'diameter': 1}
    diameter = mc.shape["A"]["diameter"]

    # Should overlap when spheres are less than one diameter apart
    sim = lattice_simulation_factory(dimensions=2,
                                     n=(2, 1),
                                     a=diameter * 0.9)
    sim.operations.add(mc)
    sim.operations.schedule()
    assert mc.overlaps > 0

    # Should not overlap when spheres are larger than one diameter apart
    s = sim.state.snapshot
    if s.exists:
        s.particles.position[0] = (0, 0, 0)
        s.particles.position[1] = (0, diameter * 1.1, 0)
    sim.state.snapshot = s
    assert mc.overlaps == 0

    # Should barely overlap when spheres are exactly than one diameter apart
    s = sim.state.snapshot
    if s.exists:
        s.particles.position[0] = (0, 0, 0)
        s.particles.position[1] = (0, diameter * 0.9999, 0)
    sim.state.snapshot = s
    assert mc.overlaps == 1

def test_overlaps_ellipsoid(device, lattice_simulation_factory):
    a = 1/4
    b = 1/2
    c = 1
    mc = hoomd.hpmc.integrate.Ellipsoid(23456)
    mc.shape["A"] = {'a': a, 'b': b, 'c': c}

    sim = lattice_simulation_factory(dimensions=2,
                                     n=(2, 1),
                                     a=10)
    sim.operations.add(mc)
    sim.operations.schedule()
    assert mc.overlaps == 0

    abc_list = [(0, 0, c), (0, b, 0), (a, 0, 0)]
    for i in range(len(abc_list)):
        # Should barely overlap when ellipsoids are exactly than one diameter apart
        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (abc_list[i][0]*0.9*2,
                                       abc_list[i][1]*0.9*2,
                                       abc_list[i][2]*0.9*2)
        sim.state.snapshot = s
        assert mc.overlaps == 1

        # Should not overlap when ellipsoids are larger than one diameter apart
        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (abc_list[i][0]*1.15*2,
                                       abc_list[i][1]*1.15*2,
                                       abc_list[i][2]*1.15*2)
        sim.state.snapshot = s
        assert mc.overlaps == 0

    # Line up ellipsoids where they aren't overlapped, and then rotate one so
    # they overlap
    s = sim.state.snapshot
    if s.exists:
        s.particles.position[0] = (0, 0, 0)
        s.particles.position[1] = (a*1.1*2, 0, 0)
        s.particles.orientation[1] = tuple(np.array([1, 0, 0.45, 0])/(1.2025**0.5))
    sim.state.snapshot = s
    assert mc.overlaps > 0

def test_overlaps_polygons(device, lattice_simulation_factory):
    # ConvexPolygon args
    triangle = {'vertices': [(0, (0.75**0.5) / 2),
                             (-0.5, -(0.75**0.5) / 2),
                             (0.5, -(0.75**0.5) / 2)]}

    # SimplePolygon args
    square = {"vertices": np.array([(-1, -1), (1, -1), (1, 1), (-1, 1)])/2}

    shapes = [(triangle, hoomd.hpmc.integrate.ConvexPolygon),
              (square, hoomd.hpmc.integrate.SimplePolygon)]

    for args, integrator in shapes:
        mc = integrator(23456)
        mc.shape['A'] = args

        sim = lattice_simulation_factory(dimensions=2, n=(2, 1), a=10)
        sim.operations.add(mc)
        sim.operations.schedule()
        assert mc.overlaps == 0

        # Place center of shape 2 on each of shape 1's vertices
        for vert in mc.shape["A"]["vertices"]:
            s = sim.state.snapshot
            if s.exists:
                s.particles.position[0] = (0, 0, 0)
                s.particles.position[1] = (vert[0], vert[1], 0)
            sim.state.snapshot = s
            assert mc.overlaps > 0

        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (0, 1.05, 0)
        sim.state.snapshot = s
        assert mc.overlaps == 0

        # Rotate one of the shapes so they will overlap
        s = sim.state.snapshot
        if s.exists:
            s.particles.orientation[1] = tuple(np.array([1, 0, 0, 0.45])/(1.2025**0.5))
        sim.state.snapshot = s
        assert mc.overlaps > 0

def test_overlaps_polyhedra(device, lattice_simulation_factory):
    # ConvexPolyhedron args
    tetrahedron = {"vertices": np.array([(1, 1, 1),
                                         (-1, -1, 1),
                                         (1, -1, -1),
                                         (-1, 1, -1)])/2}

    # Polyhedron args
    cube = {'vertices': [(-0.5, -0.5, -0.5),
                     (-0.5, -0.5, 0.5),
                     (-0.5, 0.5, -0.5),
                     (-0.5, 0.5, 0.5),
                     (0.5, -0.5, -0.5),
                     (0.5, -0.5, 0.5),
                     (0.5, 0.5, -0.5),
                     (0.5, 0.5, 0.5)],
        'faces': [[0, 2, 6],
                  [6, 4, 0],
                  [5, 0, 4],
                  [5, 1, 0],
                  [5, 4, 6],
                  [5, 6, 7],
                  [3, 2, 0],
                  [3, 0, 1],
                  [3, 6, 2],
                  [3, 7, 6],
                  [3, 1, 5],
                  [3, 5, 7]],
        'overlap': [True, True, True, True, True, True,
                    True, True, True, True, True, True]}

    shapes = [(tetrahedron, hoomd.hpmc.integrate.ConvexPolyhedron),
              (cube, hoomd.hpmc.integrate.Polyhedron)]

    for args, integrator in shapes:
        mc = integrator(23456)
        mc.shape['A'] = args

        sim = lattice_simulation_factory(dimensions=2, n=(2, 1), a=10)
        sim.operations.add(mc)
        sim.operations.schedule()
        assert mc.overlaps == 0

        # Place center of shape 2 on each of shape 1's vertices
        for vert in mc.shape["A"]["vertices"]:
            s = sim.state.snapshot
            if s.exists:
                s.particles.position[0] = (0, 0, 0)
                s.particles.position[1] = vert
            sim.state.snapshot = s
            assert mc.overlaps > 0

        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (0, 0.9, 0)
        sim.state.snapshot = s
        assert mc.overlaps > 0

        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (0, 1.1, 0)
        sim.state.snapshot = s
        assert mc.overlaps == 0

        # Rotate one of the polyhedra so they will overlap
        s = sim.state.snapshot
        if s.exists:
            s.particles.orientation[1] = tuple(np.array([1, 1, 1, 0])/(3**0.5))
        sim.state.snapshot = s
        assert mc.overlaps > 0
<commit_msg>Added spheropolygon and spheropolyhedron overlap test<commit_after>import hoomd
import hoomd.hpmc
from hoomd.hpmc import _hpmc
import numpy as np
import pytest
from hoomd.hpmc.pytest.conftest import *

def check_dict(shape_dict, args):
    """
    check_dict: Function to check that two dictionaries are equivalent

    Arguments: shape_dict and args - dictionaries to test

    Ex: mc = hoomd.hpmc.integrate.Sphere(23456)
        mc.shape["A"] = {"diameter": 1}
        check_dict(mc.shape["A"], {"diameter": 1})

    Useful for more complex nested dictionaries (like the shape key in unions)
    Used to test that the dictionary passed in is what gets passed out
    """
    for key, val in args.items():
        if isinstance(shape_dict[key], list) and len(shape_dict[key]) > 0 \
           and key != 'shapes':
            np.testing.assert_allclose(shape_dict[key], val)
        elif key == 'shapes':
            for i in range(len(shape_dict[key])):
                shape_args = shape_dict[key][i]
                val_args = val[i]
                for shape_key in shape_args:
                    if isinstance(shape_args[shape_key], list) \
                       and len(shape_args[shape_key]) > 0:
                        np.testing.assert_allclose(val_args[shape_key],
                                                   shape_args[shape_key])
                    else:
                        assert shape_args[shape_key] == val_args[shape_key]
        else:
            np.testing.assert_almost_equal(shape_dict[key], val)

def test_dict_conversion(shape_dict_conversion_args):
    for shape_params, args_list in shape_dict_conversion_args():
        for args in args_list:
            test_shape = shape_params(args)
            test_dict = test_shape.asDict()
            check_dict(test_dict, args)

def test_shape_params(integrator_args):
    for shape_integrator, valid_args, invalid_args in integrator_args():
        mc = shape_integrator(23456)
        for args in valid_args:
            mc.shape["A"] = args
            check_dict(mc.shape["A"], args)
        for args in invalid_args:
            with pytest.raises(Exception):
                mc.shape["A"] = args

def test_shape_attached(dummy_simulation_factory, integrator_args):
    for shape_integrator, valid_args, invalid_args in integrator_args():
        mc = shape_integrator(23456)
        for args in valid_args:
            mc.shape["A"] = args
            sim = dummy_simulation_factory()
            sim.operations.add(mc)
            sim.operations.schedule()
            check_dict(mc.shape["A"], args)

def test_moves(device, lattice_simulation_factory, integrator_args):
    dims = 3
    for shape_integrator, valid_args, invalid_args in integrator_args():
        if 'union' not in str(shape_integrator).lower():
            args = valid_args[0]
            if 'polygon' in str(shape_integrator).lower():
                dims = 2
            mc = shape_integrator(23456)
            mc.shape['A'] = args
            sim = lattice_simulation_factory(dimensions=dims)
            sim.operations.add(mc)
            sim.operations.schedule()
            sim.run(100)
            accepted_rejected_trans = sum(sim.operations.integrator.translate_moves)
            assert accepted_rejected_trans > 0
            if 'sphere' not in str(shape_integrator).lower():
                accepted_rejected_rot = sum(sim.operations.integrator.rotate_moves)
                assert accepted_rejected_rot > 0

def test_overlaps_sphere(device, lattice_simulation_factory):
    mc = hoomd.hpmc.integrate.Sphere(23456)
    mc.shape["A"] = {'diameter': 1}
    diameter = mc.shape["A"]["diameter"]

    # Should overlap when spheres are less than one diameter apart
    sim = lattice_simulation_factory(dimensions=2,
                                     n=(2, 1),
                                     a=diameter * 0.9)
    sim.operations.add(mc)
    sim.operations.schedule()
    assert mc.overlaps > 0

    # Should not overlap when spheres are larger than one diameter apart
    s = sim.state.snapshot
    if s.exists:
        s.particles.position[0] = (0, 0, 0)
        s.particles.position[1] = (0, diameter * 1.1, 0)
    sim.state.snapshot = s
    assert mc.overlaps == 0

    # Should barely overlap when spheres are exactly than one diameter apart
    s = sim.state.snapshot
    if s.exists:
        s.particles.position[0] = (0, 0, 0)
        s.particles.position[1] = (0, diameter * 0.9999, 0)
    sim.state.snapshot = s
    assert mc.overlaps == 1

def test_overlaps_ellipsoid(device, lattice_simulation_factory):
    a = 1/4
    b = 1/2
    c = 1
    mc = hoomd.hpmc.integrate.Ellipsoid(23456)
    mc.shape["A"] = {'a': a, 'b': b, 'c': c}

    sim = lattice_simulation_factory(dimensions=2,
                                     n=(2, 1),
                                     a=10)
    sim.operations.add(mc)
    sim.operations.schedule()
    assert mc.overlaps == 0

    abc_list = [(0, 0, c), (0, b, 0), (a, 0, 0)]
    for i in range(len(abc_list)):
        # Should barely overlap when ellipsoids are exactly than one diameter apart
        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (abc_list[i][0]*0.9*2,
                                       abc_list[i][1]*0.9*2,
                                       abc_list[i][2]*0.9*2)
        sim.state.snapshot = s
        assert mc.overlaps == 1

        # Should not overlap when ellipsoids are larger than one diameter apart
        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (abc_list[i][0]*1.15*2,
                                       abc_list[i][1]*1.15*2,
                                       abc_list[i][2]*1.15*2)
        sim.state.snapshot = s
        assert mc.overlaps == 0

    # Line up ellipsoids where they aren't overlapped, and then rotate one so
    # they overlap
    s = sim.state.snapshot
    if s.exists:
        s.particles.position[0] = (0, 0, 0)
        s.particles.position[1] = (a*1.1*2, 0, 0)
        s.particles.orientation[1] = tuple(np.array([1, 0, 0.45, 0])/(1.2025**0.5))
    sim.state.snapshot = s
    assert mc.overlaps > 0

def test_overlaps_polygons(device, lattice_simulation_factory):
    # ConvexPolygon args
    triangle = {'vertices': [(0, (0.75**0.5) / 2),
                             (-0.5, -(0.75**0.5) / 2),
                             (0.5, -(0.75**0.5) / 2)]}

    # SimplePolygon args
    square = {"vertices": np.array([(-1, -1), (1, -1), (1, 1), (-1, 1)])/2}

    shapes = [(triangle, hoomd.hpmc.integrate.ConvexPolygon),
              (square, hoomd.hpmc.integrate.SimplePolygon)]

    for args, integrator in shapes:
        mc = integrator(23456)
        mc.shape['A'] = args

        sim = lattice_simulation_factory(dimensions=2, n=(2, 1), a=10)
        sim.operations.add(mc)
        sim.operations.schedule()
        assert mc.overlaps == 0

        # Place center of shape 2 on each of shape 1's vertices
        for vert in mc.shape["A"]["vertices"]:
            s = sim.state.snapshot
            if s.exists:
                s.particles.position[0] = (0, 0, 0)
                s.particles.position[1] = (vert[0], vert[1], 0)
            sim.state.snapshot = s
            assert mc.overlaps > 0

        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (0, 1.05, 0)
        sim.state.snapshot = s
        assert mc.overlaps == 0

        # Rotate one of the shapes so they will overlap
        s = sim.state.snapshot
        if s.exists:
            s.particles.orientation[1] = tuple(np.array([1, 0, 0, 0.45])/(1.2025**0.5))
        sim.state.snapshot = s
        assert mc.overlaps > 0

def test_overlaps_polyhedra(device, lattice_simulation_factory):
    # ConvexPolyhedron args
    tetrahedron = {"vertices": np.array([(1, 1, 1),
                                         (-1, -1, 1),
                                         (1, -1, -1),
                                         (-1, 1, -1)])/2}

    # Polyhedron args
    cube = {'vertices': [(-0.5, -0.5, -0.5),
                     (-0.5, -0.5, 0.5),
                     (-0.5, 0.5, -0.5),
                     (-0.5, 0.5, 0.5),
                     (0.5, -0.5, -0.5),
                     (0.5, -0.5, 0.5),
                     (0.5, 0.5, -0.5),
                     (0.5, 0.5, 0.5)],
        'faces': [[0, 2, 6],
                  [6, 4, 0],
                  [5, 0, 4],
                  [5, 1, 0],
                  [5, 4, 6],
                  [5, 6, 7],
                  [3, 2, 0],
                  [3, 0, 1],
                  [3, 6, 2],
                  [3, 7, 6],
                  [3, 1, 5],
                  [3, 5, 7]],
        'overlap': [True, True, True, True, True, True,
                    True, True, True, True, True, True]}

    shapes = [(tetrahedron, hoomd.hpmc.integrate.ConvexPolyhedron),
              (cube, hoomd.hpmc.integrate.Polyhedron)]

    for args, integrator in shapes:
        mc = integrator(23456)
        mc.shape['A'] = args

        sim = lattice_simulation_factory(dimensions=2, n=(2, 1), a=10)
        sim.operations.add(mc)
        sim.operations.schedule()
        assert mc.overlaps == 0

        # Place center of shape 2 on each of shape 1's vertices
        for vert in mc.shape["A"]["vertices"]:
            s = sim.state.snapshot
            if s.exists:
                s.particles.position[0] = (0, 0, 0)
                s.particles.position[1] = vert
            sim.state.snapshot = s
            assert mc.overlaps > 0

        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (0, 0.9, 0)
        sim.state.snapshot = s
        assert mc.overlaps > 0

        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (0, 1.1, 0)
        sim.state.snapshot = s
        assert mc.overlaps == 0

        # Rotate one of the polyhedra so they will overlap
        s = sim.state.snapshot
        if s.exists:
            s.particles.orientation[1] = tuple(np.array([1, 1, 1, 0])/(3**0.5))
        sim.state.snapshot = s
        assert mc.overlaps > 0

def test_overlaps_spheropolygon(device, lattice_simulation_factory):

    triangle = {'vertices': [(0, (0.75**0.5) / 2),
                             (-0.5, -(0.75**0.5) / 2),
                             (0.5, -(0.75**0.5) / 2)],
                'sweep_radius': 0.2}

    square = {"vertices": np.array([(-1, -1), (1, -1), (1, 1), (-1, 1)])/2,
              "sweep_radius": 0.1}

    shapes = [(triangle, hoomd.hpmc.integrate.ConvexSpheropolygon),
              (square, hoomd.hpmc.integrate.ConvexSpheropolygon)]

    for args, integrator in shapes:
        mc = integrator(23456)
        mc.shape['A'] = args

        sim = lattice_simulation_factory(dimensions=2, n=(2, 1), a=10)
        sim.operations.add(mc)
        sim.operations.schedule()
        assert mc.overlaps == 0

        # Place center of shape 2 on each of shape 1's vertices
        for vert in mc.shape["A"]["vertices"]:
            s = sim.state.snapshot
            if s.exists:
                s.particles.position[0] = (0, 0, 0)
                s.particles.position[1] = (vert[0], vert[1], 0)
            sim.state.snapshot = s
            assert mc.overlaps > 0

        # Place shapes where they wouldn't overlap w/o sweep radius
        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (0, 1.2, 0)
        sim.state.snapshot = s
        assert mc.overlaps > 0

        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (0, 1.3, 0)
        sim.state.snapshot = s
        assert mc.overlaps == 0

        # Rotate one of the shapes so they will overlap
        s = sim.state.snapshot
        if s.exists:
            s.particles.orientation[1] = tuple(np.array([1, 0, 0, 0.45])/(1.2025**0.5))
        sim.state.snapshot = s
        assert mc.overlaps > 0

def test_overlaps_spheropolyhedron(device, lattice_simulation_factory):

    tetrahedron = {"vertices": np.array([(1, 1, 1),
                                         (-1, -1, 1),
                                         (1, -1, -1),
                                         (-1, 1, -1)])/2,
                   "sweep_radius": 0.2}

    cube = {"vertices": [(-0.5, -0.5, -0.5),
                         (-0.5, -0.5, 0.5),
                         (-0.5, 0.5, -0.5),
                         (-0.5, 0.5, 0.5),
                         (0.5, -0.5, -0.5),
                         (0.5, -0.5, 0.5),
                         (0.5, 0.5, -0.5),
                         (0.5, 0.5, 0.5)],
            "sweep_radius": 0.2}

    shapes = [(tetrahedron, hoomd.hpmc.integrate.ConvexSpheropolyhedron),
              (cube, hoomd.hpmc.integrate.ConvexSpheropolyhedron)]

    for args, integrator in shapes:
        mc = integrator(23456)
        mc.shape['A'] = args

        sim = lattice_simulation_factory(dimensions=2, n=(2, 1), a=10)
        sim.operations.add(mc)
        sim.operations.schedule()
        assert mc.overlaps == 0

        # Place center of shape 2 on each of shape 1's vertices
        for vert in mc.shape["A"]["vertices"]:
            s = sim.state.snapshot
            if s.exists:
                s.particles.position[0] = (0, 0, 0)
                s.particles.position[1] = vert
            sim.state.snapshot = s
            assert mc.overlaps > 0

        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (0, 1.2, 0)
        sim.state.snapshot = s
        assert mc.overlaps > 0

        s = sim.state.snapshot
        if s.exists:
            s.particles.position[0] = (0, 0, 0)
            s.particles.position[1] = (0, 1.4, 0)
        sim.state.snapshot = s
        assert mc.overlaps == 0

        # Rotate one of the polyhedra so they will overlap
        s = sim.state.snapshot
        if s.exists:
            s.particles.orientation[1] = tuple(np.array([1, 1, 1, 0])/(3**0.5))
        sim.state.snapshot = s
        assert mc.overlaps > 0
<|endoftext|>