#!/usr/bin/python3
import numpy as np
import gym
from gym import spaces
from gym.utils import seeding
from gym.envs.classic_control import rendering
import math
from pdb import set_trace
import reeds_shepp

try:
    # When running as a package
    from gym.envs.classic_control.demo import DemoTrack
    from gym.envs.classic_control.path_planner import TruckPlanner
    from gym.envs.classic_control.path_planner import ExcavatorPlanner
    from gym.envs.classic_control.collision_check import Rectangle
    from gym.envs.classic_control.collision_check import CollisionCheck
    from gym.envs.classic_control.collision_check import Contains
except BaseException:
    # When running current file as a file
    from demo import DemoTrack
    from path_planner import TruckPlanner
    from path_planner import ExcavatorPlanner
    from collision_check import Rectangle
    from collision_check import CollisionCheck
    from collision_check import Contains

class MiningEnv(gym.Env):
    """
    Description:
        A truck and an excavator is tasked to move a pile of minerals out of
        the load site. The truck and the excavator both starts with a random
        position and orientation in a designated region. The goal is to
        coordinate both the truck and the excavator to reach a position where
        the hand of the excavator is right above the container of the truck.

    Observations:
        Type: Box(14)
        Num        Observation                              Min            Max
        0          Truck x-position for truck               50             650
        1          Truck y-position for truck               250            700
        2          Truck orientation for truck(deg)         0              360
        3          Excav. x-position for truck              300            600
        4          Excav. y-position for truck              100            200
        5          Excav. orientation for truck(deg)        -90            90
        6          Excav. hand orientation for truck(deg)   -90            90
        7          Truck x-position for excav.              50             650
        8          Truck y-position for excav.              250            700
        9          Truck orientation for excav.(deg)        0              360
        10         Excav. x-position for excav.             300            600
        11         Excav. y-position for excav.             100            200
        12         Excav. orientation for excav.(deg)       -90            90
        13         Excav. hand orientation for excav.(deg)  -90            90

    Actions:
        Type: Box(7)
        Num        Action                              Min            Max
        0          Truck loading x-position            250            650
        1          Truck loading y-position            300            480
        2          Truck loading orientation(deg)      135            225
        3          Excav. loading x-position           300            600
        4          Excav. loading y-position           100            200
        5          Excav. loading orientation(deg)     -90            90
        6          Excav. hand orientation(deg)        -90            90

    Reward:
        Reward is -reeds_shepp_curve_length for every step taken, including
        the termination step, and 7000 for successfully loading on the truck.

    Starting State:
        All observations are assigned a uniform random value corresponding to
        the random initial state.

    Episode Termination:
        The hand of the excavator is above the truck bin. Then the system
        executes dumping and signals the truck to exit the loading site.

        or

        Episode length is greater than 20

    Solved Requirements(beta):
        Considered solved when the average reward is greater than or equal to
        0 over 100 consecutive trials.

    Confusions:
        1. Does the excavator and the truck operate in completely different
           regions of the loading site, thus diminishes the possibility of them
           colliding.
        2. How should I construct the path planner.

    Places that may cause confusions:
        1. The truck in gym environment is pointing down when no rotation
           (0 deg) is added, however the excavator is pointing up when (0 deg)
           no rotation is added. The code for Reeds-Shepp path generation uses
           a different set of angle definitions, in their definition 0 deg is
           pointing right.
    """

    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second': 30
    }

    def __init__(self):

        # [truck_x, truck_y, truck_sin, truck_cos]
        # truck_low = np.array([250, 300, math.sqrt(3)/2, -1/2])
        # truck_high = np.array([650, 480, 1, 1/2])
        # truck_range = truck_high - truck_low

        # [e_x, e_y, ep_sin, ep_cos, ea_sin, ea_cos]
        # excvtr_low = np.array([300, 100, -1/2, math.sqrt(3)/2, math.sqrt(2)/2, -math.sqrt(2)/2])
        # excvtr_high = np.array([600, 200, 1/2, 1, 1, math.sqrt(2)/2])
        # excvtr_range = excvtr_high - excvtr_low

        # [truck_x, truck_y, truck_sin, truck_cos, e_x, e_y, ep_sin, ep_cos, ea_sin, ea_cos]
        # e = excavator, ep = excavator platform, ea = excavator arm
        self.observation_low = np.array([50, 250, -1, -1, 300, 100, -1, -1, -1, -1])
        self.observation_high = np.array([650, 700, 1, 1, 600, 200, 1, 1, 1, 1])
        self.observation_range = self.observation_high - self.observation_low
        self.observation_space = [spaces.Box(np.zeros(10), np.ones(10), dtype=np.float32),
                                  spaces.Box(np.zeros(10), np.ones(10), dtype=np.float32)]

        # [truck_x, truck_y]
        self.truck_low = np.array([250, 300])
        self.truck_high = np.array([650, 320])
        self.truck_range = self.truck_high - self.truck_low
        # [e_x, e_y]
        self.excvtr_low = np.array([300, 180])
        self.excvtr_high = np.array([350, 200])
        self.excvtr_range = self.excvtr_high - self.excvtr_low
        self.action_space = [spaces.Box(np.zeros(2), np.ones(2), dtype=np.float32),
                             spaces.Box(np.zeros(2), np.ones(2), dtype=np.float32)]

        self.seed()
        self.viewer = None
        self.state = None
        self.if_render = 0

        # Demo track
        self.track = DemoTrack()
        self.n = 2
        self.truck_planner = TruckPlanner()
        self.excvtr_planner = ExcavatorPlanner()
        self.success_counter = 0

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def step(self, action):
        """
        Input: list(np.array([1x2]), np.array([1x2]))
        Output: None
        Use: Perform a step in the simulation environment.
        """
        # scaled target state
        scl_truck_trgt, scl_excvtr_trgt = action

        # unscaled target state
        truck_trgt = self.truck_low + np.multiply(scl_truck_trgt, self.truck_range)
        excvtr_trgt = self.excvtr_low + np.multiply(scl_excvtr_trgt, self.excvtr_range)

        # radian state
        rd_state = self._get_radian_state()
        truck_state = rd_state[:3].tolist()
        truck_trgt = truck_trgt.tolist()
        excvtr_state = rd_state[3:-1].tolist()
        excvtr_trgt = excvtr_trgt.tolist()

        # get reeds-shepp path
        self.truck_path = self.truck_planner.get_path(truck_state, truck_trgt)
        self.excvtr_path = self.excvtr_planner.get_path(excvtr_state, excvtr_trgt)
        # set_trace()

        self.step_counter = 0
        # self.counter += 1
        info = [None, None]

        # execute first 60 points on the path if collides terminate the episode
        for _ in range(60):
            if self.if_render:
                self.render()
            obsv, flag = self._update_state()
            if flag:
                # set_trace()
                # print('collision occurred')
                return obsv, [-1000, -1000], [True, True], info

        # get current scaled state
        scaled_state = self._get_scaled_state()
        observation = [scaled_state, scaled_state]

        # get current radian state
        rd_state = self._get_radian_state()
        truck_state = rd_state[:3].tolist()

        excvtr_pos = rd_state[3:5]
        arm_ori = rd_state[-1]
        arm_pos = excvtr_pos + np.array([130*math.cos(arm_ori), 130*math.sin(arm_ori)])
        arm_state = np.concatenate((arm_pos, np.pi/2), axis=None)

        # calculate reward
        reward = -reeds_shepp.path_length(tuple(truck_state), tuple(
            arm_state), self.truck_planner.rho)

        # check for termination
        done = False
        if self._can_load_on_truck():
            done = True
            reward = 8000
            # self.success_counter += 1
            # print('loads on truck, count: {}'.format(self.success_counter))

        return observation, [reward, reward], [done, done], info

    def reset(self):
        # Define unscaled initial state region
        # Initial state for truck is between [50, 650 -100deg] ~ [100, 700, -80deg]
        # For the excavator it is [300, 100, 80deg, 170deg] ~ [350, 120, 100deg, 190deg]
        low = np.array([50, 650, -1, math.cos(math.pi*(-100)/180),
                        300, 100, math.sin(math.pi*(80)/180),
                        math.cos(math.pi*(100)/180),
                        math.sin(math.pi*(190)/180), -1])
        high = np.array([100, 700, math.sin(math.pi*(-80)/180),
                         math.cos(math.pi * (-80)/180), 350, 120, 1,
                         math.cos(math.pi*(80)/180),
                         math.sin(math.pi*(170)/180),
                         math.cos(math.pi*(190)/180)])

        # Sample initial state
        self.state = self.np_random.uniform(low=low, high=high)
        # Scale sampled state
        scaled_state = self._get_scaled_state()
        obsv = [scaled_state, scaled_state]

        # records step number for termination check
        # self.counter = 0
        self.i = 0

        return obsv

    def render(self, mode='human', demo=0, render=1):
        screen_width = 600
        screen_height = 800

        truck_container_width = 60
        truck_container_length = 100
        truck_head_width = 50
        truck_head_length = 20

        excvtr_platform_width = 80
        excvtr_platform_length = 100
        excvtr_track_width = 15
        excvtr_track_length = 80
        excvtr_bin_size = 40  # Square
        excvtr_arm_width = 12
        excvtr_arm_length = 100
        excvtr_hand_width = 40
        excvtr_hand_length = 20

        if self.viewer is None:
            self.viewer = rendering.Viewer(screen_width, screen_height)

            minerals = rendering.make_circle(10)
            minerals.set_color(1, 171/255, 64/255)
            minerals.add_attr(rendering.Transform(translation=(150, 100)))
            self.viewer.add_geom(minerals)

            self.trucktrans = rendering.Transform()
            self.excvtrplattrans = rendering.Transform()
            self.excvtrbintrans = rendering.Transform()

            cntnr = self._render_polygon(truck_container_length, truck_container_width)
            cntnr.set_color(102/255, 102/255, 102/255)
            cntnr.add_attr(rendering.Transform(translation=(0, -truck_container_width/2)))
            cntnr.add_attr(self.trucktrans)
            self.viewer.add_geom(cntnr)

            hd = self._render_polygon(truck_head_length, truck_head_width)
            hd.set_color(191/255, 144/255, 0)
            hd.add_attr(rendering.Transform(translation=(truck_head_length /
                                                         2 + truck_container_length/2, -truck_head_width/2)))
            hd.add_attr(self.trucktrans)
            self.viewer.add_geom(hd)

            excvtr_plat = self._render_polygon(excvtr_platform_length, excvtr_platform_width)
            excvtr_plat.set_color(241/255, 194/255, 50/255)
            excvtr_plat.add_attr(rendering.Transform(
                translation=(excvtr_bin_size/2, -excvtr_platform_width/2)))
            excvtr_plat.add_attr(self.excvtrplattrans)
            self.viewer.add_geom(excvtr_plat)

            excvtr_track_l = self._render_polygon(excvtr_track_length, excvtr_track_width)
            excvtr_track_l.set_color(67/255, 67/255, 67/255)
            excvtr_track_l.add_attr(rendering.Transform(translation=(
                excvtr_bin_size/2,
                excvtr_platform_width/2)))
            excvtr_track_l.add_attr(self.excvtrplattrans)
            self.viewer.add_geom(excvtr_track_l)

            excvtr_track_r = self._render_polygon(excvtr_track_length, excvtr_track_width)
            excvtr_track_r.set_color(67/255, 67/255, 67/255)
            excvtr_track_r.add_attr(rendering.Transform(translation=(
                excvtr_bin_size/2,
                -excvtr_platform_width/2-excvtr_track_width)))
            excvtr_track_r.add_attr(self.excvtrplattrans)
            self.viewer.add_geom(excvtr_track_r)

            excvtr_bin = self._render_polygon(excvtr_bin_size, excvtr_bin_size)
            excvtr_bin.set_color(164/255, 194/255, 244/255)
            excvtr_bin.add_attr(rendering.Transform(translation=(0, -excvtr_bin_size/2)))
            excvtr_bin.add_attr(self.excvtrbintrans)
            self.viewer.add_geom(excvtr_bin)

            excvtr_arm = self._render_polygon(excvtr_arm_length, excvtr_arm_width)
            excvtr_arm.set_color(204/255, 204/255, 204/255)
            excvtr_arm.add_attr(rendering.Transform(
                translation=(excvtr_arm_length/2+excvtr_bin_size/2, -excvtr_arm_width/2)))
            excvtr_arm.add_attr(self.excvtrbintrans)
            self.viewer.add_geom(excvtr_arm)

            excvtr_hand = self._render_polygon(excvtr_hand_length, excvtr_hand_width)
            excvtr_hand.set_color(90/255, 90/255, 90/255)
            excvtr_hand.add_attr(rendering.Transform(
                translation=(excvtr_arm_length+excvtr_bin_size/2+excvtr_hand_length/2, -excvtr_hand_width/2)))
            excvtr_hand.add_attr(self.excvtrbintrans)
            self.viewer.add_geom(excvtr_hand)

        if render:
            # Render planned path
            truck_xys, truck_trgt = self._get_trajectory_list(self.truck_path)
            excvtr_xys, excvtr_trgt = self._get_trajectory_list(self.excvtr_path)

            self.truck_track = rendering.make_polyline(truck_xys)
            self.truck_track.set_color(0, 0, 1)
            self.truck_track.set_linewidth(4)
            self.viewer.add_geom(self.truck_track)

            truck_destination = rendering.make_circle(10)
            truck_destination.set_color(0, 0, 1)
            truck_destination.add_attr(rendering.Transform(translation=(
                truck_trgt[0], truck_trgt[1])))
            self.viewer.add_geom(truck_destination)

            self.excvtr_track = rendering.make_polyline(excvtr_xys)
            self.excvtr_track.set_color(1, 0, 0)
            self.excvtr_track.set_linewidth(4)
            self.viewer.add_geom(self.excvtr_track)

            excvtr_destination = rendering.make_circle(10)
            excvtr_destination.set_color(1, 0, 0)
            excvtr_destination.add_attr(rendering.Transform(translation=(
                excvtr_trgt[0], excvtr_trgt[1])))
            self.viewer.add_geom(excvtr_destination)

        rd_state = self._get_radian_state()

        self.trucktrans.set_translation(rd_state[0], rd_state[1])
        self.trucktrans.set_rotation(rd_state[2])  # in radian
        self.excvtrplattrans.set_translation(rd_state[3], rd_state[4])
        self.excvtrplattrans.set_rotation(rd_state[5])
        self.excvtrbintrans.set_translation(rd_state[3], rd_state[4])
        self.excvtrbintrans.set_rotation(rd_state[6])

        return self.viewer.render(return_rgb_array=mode == 'rgb_array')

    def close(self):
        if self.viewer:
            self.viewer.close()
            self.viewer = None

    def _render_polygon(self, width, length):
        l, r, t, b = -width/2, width/2, length, 0
        return rendering.FilledPolygon([(l, b), (l, t), (r, t), (r, b)])

    def _get_trajectory_list(self, trajectory):
        x = []
        y = []

        for point in trajectory:
            x.append(point[0])
            y.append(point[1])

        xys = list(zip(x, y))
        destination = [x[-1], y[-1]]
        return xys, destination

    def _update_state(self):
        old_state = self._get_radian_state()
        try:
            rd_state = np.asarray(self.truck_path[self.step_counter])
        except BaseException:
            rd_state = np.asarray(self.truck_path[-1])

        try:
            rd_state = np.concatenate((rd_state, np.asarray(
                self.excvtr_path[self.step_counter])), axis=None)
        except BaseException:
            rd_state = np.concatenate((rd_state, np.asarray(
                self.excvtr_path[-1])), axis=None)

        if old_state[6] > np.pi/2:
            rd_state = np.concatenate((rd_state, old_state[6]-np.pi/200), axis=None)
        else:
            rd_state = np.concatenate((rd_state, np.pi/2), axis=None)

        self.state = self._get_sin_cos_state(rd_state)
        scaled_state = self._get_scaled_state()
        obsv = [scaled_state, scaled_state]
        flag = self._check_collision()

        self.step_counter += 1

        return obsv, flag

    def _can_load_on_truck(self):
        """
        Input: None
        Output: bool(1x1)
        Use: Gives a true-false response for whether the current state of the
             truck and the excavator is suitable for loading.
        """

        rd_state = self._get_radian_state()
        truck_pos = rd_state[:2]
        truck_ori = rd_state[2]
        rect = Rectangle(truck_pos, 60, 40, truck_ori)
        excvtr_pos = rd_state[3:5]
        arm_ori = rd_state[-1]
        arm_pos = excvtr_pos + np.array([130*math.cos(arm_ori), 130*math.sin(arm_ori)])
        contains = Contains(arm_pos[0], arm_pos[1], rect)
        return contains.contains()

    def _get_radian_state(self):
        """
        Input: None
        Output: np.array([1x7])
        Use: Returns the state of the environment where the orientation of the
             truck and the excavator are given in radians.
        """

        truck_pos = [self.state[0], self.state[1]]
        truck_ang_rad = math.atan2(self.state[2], self.state[3])
        excvtr_pos = [self.state[4], self.state[5]]
        excvtr_plat_ang_rad = math.atan2(self.state[6], self.state[7])
        excvtr_arm_ang_rad = math.atan2(self.state[8], self.state[9])

        radian_state = truck_pos + [truck_ang_rad] + \
            excvtr_pos + [excvtr_plat_ang_rad, excvtr_arm_ang_rad]
        return np.array(radian_state)

    def _get_scaled_state(self):
        """
        Input: None
        Output: np.array([1x10])
        Use: Returns the state of the environment after scaling them all
             between 0 and 1.
        """
        scaled_state = (self.state - self.observation_low)/self.observation_range
        return scaled_state

    def _get_sin_cos_state(self, rd_state):
        """
        Input: np.array([1x7])
        Output: np.array([1x10])
        Use: Returns the state of the environment after changing the radian
             states into a sin-cos representation.
        """
        truck_pos = rd_state[:2]
        truck_sin_cos = np.array([np.sin(rd_state[2]), np.cos(rd_state[2])])
        excvtr_pos = rd_state[3:5]
        excvtr_plat_sin_cos = np.array([np.sin(rd_state[5]), np.cos(rd_state[5])])
        excvtr_arm_sin_cos = np.array([np.sin(rd_state[6]), np.cos(rd_state[6])])
        sin_cos_state = np.concatenate(
            (truck_pos, truck_sin_cos, excvtr_pos, excvtr_plat_sin_cos, excvtr_arm_sin_cos), axis=None)
        return sin_cos_state

    def _check_collision(self):
        rd_state = self._get_radian_state()
        rect1 = Rectangle(rd_state[:2], 120, 60, rd_state[2], 10)
        rect2 = Rectangle(rd_state[3:5], 100, 80, rd_state[5], 20)
        ccheck = CollisionCheck(rect1, rect2)
        return ccheck.if_collide()

def main():
    # Test rendering effect
    m = MiningEnv()

    for i in range(60000):
        done = False
        rewards = 0
        m.reset()
        while not done:
            truck_action = m.action_space[0].sample()
            excvtr_action = m.action_space[1].sample()
            action = [truck_action, excvtr_action]
            # action = [np.array([0.5, 0]), np.array([0.5, 0.8])]
            observation, reward, done_n, info = m.step(action)
            rewards += reward[0]
            done = all(done_n)
        # print('Episode {}, Reward {}'.format(i, rewards))
    m.close()

if __name__ == "__main__":
    main()


import logging
from datetime import datetime, timedelta

from aniso8601 import parse_time
from hearthstone.enums import (
	BlockType, ChoiceType, FormatType, GameTag, GameType,
	MetaDataType, Mulligan, OptionType, PlayReq, PowerType
)

from . import packets, tokens
from .exceptions import ParsingError, RegexParsingError
from .player import LazyPlayer, PlayerManager
from .utils import parse_enum, parse_tag

def parse_entity_id(entity):
	if entity.isdigit():
		return int(entity)

	if entity == tokens.GAME_ENTITY:
		# GameEntity is always 1
		return 1

	sre = tokens.ENTITY_RE.match(entity)
	if sre:
		id = sre.groups()[0]
		return int(id)

def parse_initial_tag(data):
	"""
	Parse \a data, a line formatted as tag=FOO value=BAR
	Returns the values as int.
	"""
	sre = tokens.TAG_VALUE_RE.match(data)
	if not sre:
		raise RegexParsingError(data)
	tag, value = sre.groups()
	return parse_tag(tag, value)

def clean_option_errors(error, error_param):
	"""
	As of 8.0.0.18336, all option packets are accompanied by an error and an
	errorParam argument.
	This function turns both into their respective types.
	"""

	if error == "NONE":
		error = None
	else:
		error = parse_enum(PlayReq, error)

	if not error_param:
		error_param = None
	else:
		error_param = int(error_param)

	return error, error_param

class PowerHandler:
	def __init__(self):
		super(PowerHandler, self).__init__()
		self.current_block = None
		self._metadata_node = None
		self._packets = None
		self._creating_game = False
		self.game_meta = {}

	def _check_for_mulligan_hack(self, ts, tag, value):
		# Old game logs didn't handle asynchronous mulligans properly.
		# If we're missing an ACTION_END packet after the mulligan SendChoices,
		# we just close it out manually.
		if tag == GameTag.MULLIGAN_STATE and value == Mulligan.DEALING:
			assert self.current_block
			if isinstance(self.current_block, packets.Block):
				logging.warning("WARNING: Broken mulligan nesting. Working around...")
				self.block_end(ts)

	def find_callback(self, method):
		if method == self.parse_method("DebugPrintPower"):
			return self.handle_data
		elif method == self.parse_method("DebugPrintGame"):
			return self.handle_game

	def handle_game(self, ts, data):
		if data.startswith("PlayerID="):
			sre = tokens.GAME_PLAYER_META.match(data)
			if not sre:
				raise RegexParsingError(data)
			player_id, player_name = sre.groups()
			player_id = int(player_id)
			self._packets.manager.register_player_name_by_player_id(player_name, player_id)
		else:
			key, value = data.split("=")
			key = key.strip()
			value = value.strip()
			if key == "GameType":
				value = parse_enum(GameType, value)
			elif key == "FormatType":
				value = parse_enum(FormatType, value)
			else:
				value = int(value)

			self.game_meta[key] = value

	def handle_data(self, ts, data):
		opcode = data.split()[0]

		if opcode == "ERROR:":
			# Line error... skip
			return

		if opcode in PowerType.__members__:
			return self.handle_power(ts, opcode, data)

		if opcode == "GameEntity":
			self.flush()
			self._creating_game = True
			sre = tokens.GAME_ENTITY_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			id, = sre.groups()
			if int(id) != 1:
				raise ParsingError("GameEntity ID: Expected 1, got %r" % (id))
		elif opcode == "Player":
			self.flush()
			sre = tokens.PLAYER_ENTITY_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			self.register_player(ts, *sre.groups())
		elif opcode.startswith("tag="):
			tag, value = parse_initial_tag(data)
			self._entity_packet.tags.append((tag, value))
			if tag == GameTag.CONTROLLER:
				# We need to know entity controllers for player name registration
				self._packets.manager.register_controller(self._entity_packet.entity, value)
		elif opcode.startswith("Info["):
			if not self._metadata_node:
				logging.warning("Metadata Info outside of META_DATA: %r", data)
				return
			sre = tokens.METADATA_INFO_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			idx, entity = sre.groups()
			entity = self.parse_entity_or_player(entity)
			self._metadata_node.info.append(entity)
		else:
			raise NotImplementedError(data)

	def flush(self):
		super(PowerHandler, self).flush()
		if self._metadata_node:
			self._metadata_node = None

	def handle_power(self, ts, opcode, data):
		self.flush()

		if opcode == "CREATE_GAME":
			regex, callback = tokens.CREATE_GAME_RE, self.create_game
		elif opcode in ("ACTION_START", "BLOCK_START"):
			index = None
			effectid, effectindex = None, None
			suboption, trigger_keyword = None, None
			if " SubOption=" in data:
				if " TriggerKeyword=" in data:
					sre = tokens.BLOCK_START_20457_TRIGGER_KEYWORD_RE.match(data)
					if sre is None:
						raise RegexParsingError(data)
					type, entity, effectid, effectindex, target, suboption, trigger_keyword = sre.groups()
				else:
					sre = tokens.BLOCK_START_20457_RE.match(data)
					if sre is None:
						raise RegexParsingError(data)
					type, entity, effectid, effectindex, target, suboption = sre.groups()
			else:
				if opcode == "ACTION_START":
					sre = tokens.ACTION_START_RE.match(data)
				else:
					sre = tokens.BLOCK_START_12051_RE.match(data)

				if sre is None:
					sre = tokens.ACTION_START_OLD_RE.match(data)
					if not sre:
						raise RegexParsingError(data)
					entity, type, index, target = sre.groups()
				else:
					type, entity, effectid, effectindex, target = sre.groups()

			self.block_start(
				ts, entity, type, index, effectid, effectindex, target, suboption, trigger_keyword
			)
			return
		elif opcode in ("ACTION_END", "BLOCK_END"):
			regex, callback = tokens.BLOCK_END_RE, self.block_end
		elif opcode == "FULL_ENTITY":
			if data.startswith("FULL_ENTITY - Updating"):
				regex, callback = tokens.FULL_ENTITY_UPDATE_RE, self.full_entity_update
			else:
				regex, callback = tokens.FULL_ENTITY_CREATE_RE, self.full_entity
		elif opcode == "SHOW_ENTITY":
			regex, callback = tokens.SHOW_ENTITY_RE, self.show_entity
		elif opcode == "HIDE_ENTITY":
			regex, callback = tokens.HIDE_ENTITY_RE, self.hide_entity
		elif opcode == "CHANGE_ENTITY":
			regex, callback = tokens.CHANGE_ENTITY_RE, self.change_entity
		elif opcode == "TAG_CHANGE":
			regex, callback = tokens.TAG_CHANGE_RE, self.tag_change
		elif opcode == "META_DATA":
			regex, callback = tokens.META_DATA_RE, self.meta_data
		elif opcode == "RESET_GAME":
			regex, callback = tokens.RESET_GAME_RE, self.reset_game
		else:
			raise NotImplementedError(data)

		sre = regex.match(data)
		if not sre:
			logging.warning("Could not correctly parse %r", data)
			return
		return callback(ts, *sre.groups())

	def new_packet_tree(self, ts):
		self._packets = packets.PacketTree(ts)
		self._packets.spectator_mode = self.spectator_mode
		self._packets.manager = PlayerManager()
		self.current_block = self._packets
		self.games.append(self._packets)
		return self._packets

	# Messages
	def create_game(self, ts):
		self.new_packet_tree(ts)
		entity_id = 1
		self._game_packet = self._entity_packet = packets.CreateGame(ts, entity_id)
		self.register_packet(self._game_packet)
		return self._game_packet

	def block_start(
		self, ts, entity, type, index, effectid, effectindex, target, suboption, trigger_keyword
	):
		id = self.parse_entity_or_player(entity)
		type = parse_enum(BlockType, type)
		if index is not None:
			index = int(index)
		target = self.parse_entity_or_player(target)
		if suboption is not None:
			suboption = int(suboption)
		if trigger_keyword is not None:
			trigger_keyword = parse_enum(GameTag, trigger_keyword)

		block = packets.Block(
			ts, id, type, index, effectid, effectindex, target, suboption, trigger_keyword
		)
		block.parent = self.current_block
		self.register_packet(block)
		self.current_block = block
		return block

	def block_end(self, ts):
		if not self.current_block.parent:
			logging.warning("[%s] Orphaned BLOCK_END detected", ts)
			return self.current_block
		self.current_block.end()
		block = self.current_block
		self.current_block = self.current_block.parent
		return block

	def full_entity(self, ts, id, card_id):
		id = int(id)
		self._entity_packet = packets.FullEntity(ts, id, card_id)
		self.register_packet(self._entity_packet)

		if self._creating_game:
			# First packet after create game should always be a FULL_ENTITY
			self._creating_game = False
			# It should always have ID 4
			if id != 4:
				raise ParsingError("Expected entity 4 after creating game, got %r" % (id))

			# While we're at it, we check if we got an abnormal amount of players
			player_count = len(self._game_packet.players)
			if player_count != 2:
				raise ParsingError("Expected exactly 2 players, got %r" % (player_count))

		return self._entity_packet

	def full_entity_update(self, ts, entity, card_id):
		id = parse_entity_id(entity)
		return self.full_entity(ts, id, card_id)

	def show_entity(self, ts, entity, card_id):
		id = parse_entity_id(entity)
		self._entity_packet = packets.ShowEntity(ts, id, card_id)
		self.register_packet(self._entity_packet)
		return self._entity_packet

	def hide_entity(self, ts, entity, tag, value):
		id = parse_entity_id(entity)
		tag, value = parse_tag(tag, value)
		if tag != GameTag.ZONE:
			raise ParsingError("HIDE_ENTITY got non-zone tag (%r)" % (tag))
		packet = packets.HideEntity(ts, id, value)
		self.register_packet(packet)
		return packet

	def change_entity(self, ts, entity, card_id):
		id = self.parse_entity_or_player(entity)
		self._entity_packet = packets.ChangeEntity(ts, id, card_id)
		self.register_packet(self._entity_packet)
		return self._entity_packet

	def meta_data(self, ts, meta, data, info):
		meta = parse_enum(MetaDataType, meta)
		if meta == MetaDataType.JOUST:
			data = parse_entity_id(data)
		count = int(info)
		self._metadata_node = packets.MetaData(ts, meta, data, count)
		self.register_packet(self._metadata_node)
		return self._metadata_node

	def tag_change(self, ts, e, tag, value, def_change):
		id = self.parse_entity_or_player(e)
		tag, value = parse_tag(tag, value)
		self._check_for_mulligan_hack(ts, tag, value)

		if isinstance(id, LazyPlayer):
			id = self._packets.manager.register_player_name_on_tag_change(id, tag, value)

		has_change_def = def_change == tokens.DEF_CHANGE
		packet = packets.TagChange(ts, id, tag, value, has_change_def)
		self.register_packet(packet)
		return packet

	def reset_game(self, ts):
		packet = packets.ResetGame(ts)
		self.register_packet(packet)
		return packet

class OptionsHandler:
	def __init__(self):
		super(OptionsHandler, self).__init__()
		self._option_packet = None
		self._options_packet = None
		self._suboption_packet = None

	def find_callback(self, method):
		if method == self.parse_method("SendOption"):
			return self.handle_send_option
		elif method == self.parse_method("DebugPrintOptions"):
			return self.handle_options

	def _parse_option_packet(self, ts, data):
		if " errorParam=" in data:
			sre = tokens.OPTIONS_OPTION_ERROR_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			optype, id, type, entity, error, error_param = sre.groups()
			error, error_param = clean_option_errors(error, error_param)
		else:
			sre = tokens.OPTIONS_OPTION_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			optype, id, type, entity = sre.groups()
			error, error_param = None, None

		id = int(id)
		type = parse_enum(OptionType, type)
		if entity:
			entity = self.parse_entity_or_player(entity)
		self._option_packet = packets.Option(ts, entity, id, type, optype, error, error_param)
		if not self._options_packet:
			raise ParsingError("Option without a parent option group: %r" % (data))

		self.register_packet(self._option_packet, node=self._options_packet.options)
		self._suboption_packet = None

		return self._option_packet

	def _parse_suboption_packet(self, ts, data):
		if " errorParam=" in data:
			sre = tokens.OPTIONS_SUBOPTION_ERROR_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			optype, id, entity, error, error_param = sre.groups()
			error, error_param = clean_option_errors(error, error_param)
		else:
			sre = tokens.OPTIONS_SUBOPTION_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			optype, id, entity = sre.groups()
			error, error_param = None, None

		id = int(id)
		if entity:
			entity = self.parse_entity_or_player(entity)

		packet = packets.Option(ts, entity, id, None, optype, error, error_param)
		if optype == "subOption":
			self._suboption_packet = packet
			node = self._option_packet
		elif optype == "target":
			node = self._suboption_packet or self._option_packet

		if not node:
			raise ParsingError("SubOption / target without a matching option: %r" % (data))

		self.register_packet(packet, node=node.options)

		return packet

	def handle_options(self, ts, data):
		if data.startswith("id="):
			sre = tokens.OPTIONS_ENTITY_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			id, = sre.groups()
			id = int(id)
			self._options_packet = packets.Options(ts, id)
			self.current_block.packets.append(self._options_packet)
		elif data.startswith("option "):
			return self._parse_option_packet(ts, data)
		elif data.startswith(("subOption ", "target ")):
			return self._parse_suboption_packet(ts, data)

	def handle_send_option(self, ts, data):
		if data.startswith("selectedOption="):
			sre = tokens.SEND_OPTION_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			option, suboption, target, position = sre.groups()
			packet = packets.SendOption(ts, int(option), int(suboption), int(target), int(position))
			self.register_packet(packet)
			return packet
		raise NotImplementedError("Unhandled send option: %r" % (data))

class ChoicesHandler:
	def __init__(self):
		super(ChoicesHandler, self).__init__()
		self._choice_packet = None
		self._chosen_packet = None
		self._send_choice_packet = None

	def find_callback(self, method):
		if method == self.parse_method("DebugPrintEntityChoices"):
			return self.handle_entity_choices
		elif method == self.parse_method("DebugPrintChoices"):
			return self.handle_entity_choices_old
		elif method == self.parse_method("SendChoices"):
			return self.handle_send_choices
		elif method == self.parse_method("DebugPrintEntitiesChosen"):
			return self.handle_entities_chosen

	def flush(self):
		if self._choice_packet:
			if self._choice_packet.type == ChoiceType.MULLIGAN:
				self._packets.manager.register_player_name_mulligan(self._choice_packet)
			self._choice_packet = None
		if self._chosen_packet:
			self._chosen_packet = None
		if self._send_choice_packet:
			self._send_choice_packet = None

	def handle_entity_choices_old(self, ts, data):
		if data.startswith("id="):
			sre = tokens.CHOICES_CHOICE_OLD_1_RE.match(data)
			if sre:
				self.register_choices_old_1(ts, *sre.groups())
			else:
				sre = tokens.CHOICES_CHOICE_OLD_2_RE.match(data)
				if not sre:
					raise RegexParsingError(data)
				self.register_choices_old_2(ts, *sre.groups())
		else:
			return self.handle_entity_choices(ts, data)

	def handle_entity_choices(self, ts, data):
		if data.startswith("id="):
			sre = tokens.CHOICES_CHOICE_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			return self.register_choices(ts, *sre.groups())
		elif data.startswith("Source="):
			sre = tokens.CHOICES_SOURCE_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			entity, = sre.groups()
			id = self.parse_entity_or_player(entity)
			if not self._choice_packet:
				raise ParsingError("Source Choice Entity outside of choie packet: %r" % (data))
			self._choice_packet.source = id
			return id
		elif data.startswith("Entities["):
			sre = tokens.CHOICES_ENTITIES_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			idx, entity = sre.groups()
			id = self.parse_entity_or_player(entity)
			if not id:
				raise ParsingError("Missing choice entity %r (%r)" % (id, entity))
			if not self._choice_packet:
				raise ParsingError("Choice Entity outside of choice packet: %r" % (data))
			self._choice_packet.choices.append(id)
			return id
		raise NotImplementedError("Unhandled entity choice: %r" % (data))

	def _register_choices(self, ts, id, player, tasklist, type, min, max):
		id = int(id)
		type = parse_enum(ChoiceType, type)
		min, max = int(min), int(max)
		self._choice_packet = packets.Choices(ts, player, id, tasklist, type, min, max)
		self.register_packet(self._choice_packet)
		return self._choice_packet

	def register_choices_old_1(self, ts, id, type):
		player = None
		# XXX: We don't have a player here for old games.
		# Is it safe to assume CURRENT_PLAYER?
		tasklist = None
		min, max = 0, 0
		return self._register_choices(ts, id, player, tasklist, type, min, max)

	def register_choices_old_2(self, ts, id, player_id, type, min, max):
		player_id = int(player_id)
		player = self._packets.manager._players_by_player_id[player_id]
		tasklist = None
		return self._register_choices(ts, id, player, tasklist, type, min, max)

	def register_choices(self, ts, id, player, tasklist, type, min, max):
		player = self.parse_entity_or_player(player)
		if tasklist is not None:
			# Sometimes tasklist is empty
			tasklist = int(tasklist)
		return self._register_choices(ts, id, player, tasklist, type, min, max)

	def handle_send_choices(self, ts, data):
		if data.startswith("id="):
			sre = tokens.SEND_CHOICES_CHOICE_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			id, type = sre.groups()
			id = int(id)
			type = parse_enum(ChoiceType, type)
			self._send_choice_packet = packets.SendChoices(ts, id, type)
			self.register_packet(self._send_choice_packet)
			return self._send_choice_packet
		elif data.startswith("m_chosenEntities"):
			sre = tokens.SEND_CHOICES_ENTITIES_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			idx, entity = sre.groups()
			id = self.parse_entity_or_player(entity)
			if not id:
				raise ParsingError("Missing chosen entity %r (%r)" % (id, entity))
			if not self._send_choice_packet:
				raise ParsingError("Chosen Entity outside of choice packet: %r" % (data))
			self._send_choice_packet.choices.append(id)
			return id
		raise NotImplementedError("Unhandled send choice: %r" % (data))

	def handle_entities_chosen(self, ts, data):
		if data.startswith("id="):
			sre = tokens.ENTITIES_CHOSEN_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			id, player, count = sre.groups()
			id = int(id)
			player = self.parse_entity_or_player(player)
			self._chosen_packet_count = int(count)
			self._chosen_packet = packets.ChosenEntities(ts, player, id)
			self.register_packet(self._chosen_packet)
			return self._chosen_packet
		elif data.startswith("Entities["):
			sre = tokens.ENTITIES_CHOSEN_ENTITIES_RE.match(data)
			if not sre:
				raise RegexParsingError(data)
			idx, entity = sre.groups()
			id = self.parse_entity_or_player(entity)
			if not id:
				raise ParsingError("Missing entity chosen %r (%r)" % (id, entity))
			if not self._chosen_packet:
				raise ParsingError("Entity Chosen outside of choice packet: %r" % (data))
			self._chosen_packet.choices.append(id)
			if len(self._chosen_packet.choices) > self._chosen_packet_count:
				raise ParsingError("Too many choices (expected %r)" % (self._chosen_packet_count))
			return id
		raise NotImplementedError("Unhandled entities chosen: %r" % (data))

class SpectatorModeHandler:
	def __init__(self):
		super(SpectatorModeHandler, self).__init__()
		self.spectating_first_player = False
		self.spectating_second_player = False

	@property
	def spectator_mode(self):
		return self.spectating_first_player or self.spectating_second_player

	def set_spectating(self, first, second=None):
		self.spectating_first_player = first
		if second is not None:
			self.spectating_second_player = second

	def process_spectator_mode(self, line):
		if line == tokens.SPECTATOR_MODE_BEGIN_GAME:
			self.set_spectating(True)
		elif line == tokens.SPECTATOR_MODE_BEGIN_FIRST:
			self.set_spectating(True, False)
		elif line == tokens.SPECTATOR_MODE_BEGIN_SECOND:
			self.set_spectating(True, True)
		elif line == tokens.SPECTATOR_MODE_END_MODE:
			self.set_spectating(False, False)
		elif line == tokens.SPECTATOR_MODE_END_GAME:
			self.set_spectating(False, False)
		else:
			raise NotImplementedError("Unhandled spectator mode: %r" % (line))

class LogParser(
	PowerHandler, ChoicesHandler, OptionsHandler, SpectatorModeHandler
):
	def __init__(self):
		super(LogParser, self).__init__()
		self.games = []
		self.line_regex = tokens.POWERLOG_LINE_RE
		self._game_state_processor = "GameState"
		self._current_date = None
		self._synced_timestamp = False

	def parse_timestamp(self, ts, method):
		ret = parse_time(ts)

		if not self._synced_timestamp:
			# The first timestamp we parse requires syncing the time
			# (in case _current_date is greater than the start date)
			if self._current_date is not None:
				self._current_date = self._current_date.replace(
					hour=ret.hour,
					minute=ret.minute,
					second=ret.second,
					microsecond=ret.microsecond,
				)
			# Only do it once per parse tree
			self._synced_timestamp = True

		# Logs don't have dates :(
		if self._current_date is None:
			# No starting date is available. Return just the time.
			return ret

		ret = datetime.combine(self._current_date, ret)
		ret = ret.replace(tzinfo=self._current_date.tzinfo)
		if ret < self._current_date:
			# If the new date falls before the last saved date, that
			# means we rolled over and need to increment the day by 1.
			ret += timedelta(days=1)
		self._current_date = ret
		return ret

	def read(self, fp):
		for line in fp:
			self.read_line(line)

	def read_line(self, line):
		sre = tokens.TIMESTAMP_RE.match(line)
		if not sre:
			raise RegexParsingError(line)
		level, ts, line = sre.groups()
		if line.startswith(tokens.SPECTATOR_MODE_TOKEN):
			line = line.replace(tokens.SPECTATOR_MODE_TOKEN, "").strip()
			return self.process_spectator_mode(line)

		sre = self.line_regex.match(line)
		if not sre:
			return
		method, msg = sre.groups()
		msg = msg.strip()
		if not self.current_block and "CREATE_GAME" not in msg:
			# Ignore messages before the first CREATE_GAME packet
			return

		for handler in PowerHandler, ChoicesHandler, OptionsHandler:
			callback = handler.find_callback(self, method)
			if callback:
				ts = self.parse_timestamp(ts, method)
				return callback(ts, msg)

	def register_packet(self, packet, node=None):
		if node is None:
			node = self.current_block.packets
		node.append(packet)
		self._packets._packet_counter += 1
		packet.packet_id = self._packets._packet_counter

	def parse_entity_or_player(self, entity):
		if entity == "-1":
			return

		id = parse_entity_id(entity)
		if id is None:
			# Only case where an id is None is if it's a Player name
			id = self._packets.manager.get_player_by_name(entity)
		return id

	def parse_method(self, m):
		return "%s.%s" % (self._game_state_processor, m)

	def register_player(self, ts, id, player_id, hi, lo):
		id = int(id)
		player_id = int(player_id)
		hi = int(hi)
		lo = int(lo)
		lazy_player = self._packets.manager.new_player(id, player_id, is_ai=lo == 0)
		self._entity_packet = packets.CreateGame.Player(ts, lazy_player, player_id, hi, lo)
		self.register_packet(self._entity_packet, node=self._game_packet.players)
		return lazy_player


#!/usr/bin/env python

from setuptools import setup
from setuptools.command.test import test as TestCommand

# get version
from memsql import __version__

class PyTest(TestCommand):
    user_options = [
        ('watch', 'w',
         "watch tests for changes"),
        ('scan-directory=', 'd',
         "only search for tests in the specified directory"),
    ]
    boolean_options = ['watch']

    def initialize_options(self):
        self.watch = False
        self.scan_directory = None
        self.test_suite = None
        self.test_module = None
        self.test_loader = None

    def finalize_options(self):
        TestCommand.finalize_options(self)

        self.test_suite = True
        self.test_args = []
        if self.watch:
            self.test_args.append('-f')
        if self.scan_directory is not None:
            self.test_args.append(self.scan_directory)

    def run_tests(self):
        import os, sys, glob

        MY_PATH = os.path.dirname(__file__)
        sys.path.append(MY_PATH)
        os.environ['PYTHONPATH'] = os.environ.get('PYTHONPATH', '') + ':' + MY_PATH

        egg_dirs = glob.glob('*.egg')
        ignore_args = ['--ignore=%s' % d for d in egg_dirs]

        import pytest
        errno = pytest.main(ignore_args + self.test_args)
        raise sys.exit(errno)

setup(
    name='memsql',
    version=__version__,
    author='MemSQL',
    author_email='support@memsql.com',
    url='http://github.com/memsql/memsql-python',
    license='LICENSE.txt',
    description='Useful utilities and plugins for MemSQL integration.',
    long_description=open('README.rst').read(),
    classifiers=[
        'License :: OSI Approved :: MIT License',
        'Programming Language :: Python :: 2.6',
        'Programming Language :: Python :: 2.7',
    ],
    packages=[
        'memsql',
        'memsql.common',
        'memsql.perf',
    ],
    zip_safe=False,
    install_requires=['ordereddict', 'MySQL-python>=1.2.4', 'wraptor', 'simplejson'],
    tests_require=['pytest', 'mock', 'pytest-xdist'],
    cmdclass={ 'test': PyTest },
)


import unittest
from appium import webdriver

class TestClass(unittest.TestCase):

    def setUp(self):
        desired_capabilities = {}
        desired_capabilities['testobject_api_key'] = 'your TestObject API key here'
        desired_capabilities['testobject_device'] = 'LG_Nexus_4_E960_real'
        desired_capabilities['testobject_appium_version'] = '1.5.2-patched-chromedriver'
        testobject_endpoint = 'http://appium.testobject.com:80/api/appium/wd/hub'
        self.driver = webdriver.Remote(testobject_endpoint, desired_capabilities)

    # All test methods starts with the word "test_"
    def test_open_testobject_website(self):
        self.driver.get("https://testobject.com/")
        if not "TestObject" in self.driver.title:
            raise Exception("Unable to load TestObject")

    def tearDown(self):
        self.driver.quit()

# This will trigger the test
if __name__ == '__main__':
    unittest.main()