"""
This file implements the functionalities of a Mors Gym environment using pybullet physical engine.
"""
import copy
import numpy as np
import os
import inspect
currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
parentdir = os.path.dirname(os.path.dirname(currentdir))
os.sys.path.insert(0, parentdir)

import time
import random
import gym
from gym import spaces
from gym.utils import seeding
import pybullet
from pybullet_utils import bullet_client as bc
import pybullet_data
from pkg_resources import parse_version
from mors_sim.motor_simple import MotorSimple
from mors_sim.motor_accurate import MotorAccurate
from mors_sim.world_creator import WorldCreator
from mors_sim.lidar import Lidar
from mors_sim.camera import Camera

import math

INIT_POSITION = [0, 0, .76]
INIT_ORIENTATION = [0, 0, 0, 1]
LEG_POSITION = ["R1", "L1", "R2", "L2"]
MOTOR_NAMES = [
    "abad_joint_R1", "hip_joint_R1", "knee_joint_R1",
    "abad_joint_L1", "hip_joint_L1", "knee_joint_L1",
    "abad_joint_R2", "hip_joint_R2", "knee_joint_R2",
    "abad_joint_L2", "hip_joint_L2", "knee_joint_L2",
]
LEG_LINK_ID = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
MOTOR_LINK_ID = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] # shoulders: 0, 3, 6, 9
SHOULDER_MOTOR_ID = [0, 3, 6, 9]
FOOT_LINK_ID = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] # it is needed for friction setting
BASE_LINK_ID = -1

NUM_SUBSTEPS = 5
NUM_MOTORS = 12
MOTOR_ANGLE_OBSERVATION_INDEX = 0
MOTOR_VELOCITY_OBSERVATION_INDEX = MOTOR_ANGLE_OBSERVATION_INDEX + NUM_MOTORS
MOTOR_TORQUE_OBSERVATION_INDEX = MOTOR_VELOCITY_OBSERVATION_INDEX + NUM_MOTORS
BASE_ORIENTATION_OBSERVATION_INDEX = MOTOR_TORQUE_OBSERVATION_INDEX + NUM_MOTORS
ACTION_EPS = 0.01
OBSERVATION_EPS = 0.01
RENDER_HEIGHT = 720
RENDER_WIDTH = 960
K_SCALE = 2

MOTOR_ANGLE_MAX = 2.62
MOTOR_VELOCITY_MAX = 12
MOTOR_TORQUE_MAX = 110
BASE_ORIENTATION_MAX = 1.57

# LIDAR_JOINT = 27

class MorsMiniBulletEnv(gym.Env):
  """
  The gym environment for Mors.

  It simulates the locomotion of a mors, a quadruped robot. The state space
  include base orientation, the angles, velocities and torques for all the motors and the action
  space is the desired motor angle for each motor. The reward function is based
  on how fast robot can move without twitches and side deviations.

  """
  metadata = {"render.modes": ["human", "rgb_array"], "video.frames_per_second": 50}

  def __init__(self,
              urdf_root="./urdf",
              sim_freq=100,
              world="empty",
              hard_reset=False,
              render=True,
              on_rack=False,
              self_collision_enabled=False,
              debug_mode = False,
              floating_camera = False,
              step_enabled = True,

              accurate_motor_model_enabled=False,
              simple_motor_model_enabled=False,
              torque_control_enabled=False,
              motor_kp=30.0,
              motor_kd=0.1,
              motor_velocity_limit=np.inf,

              max_timesteps = np.inf,
              action_repeat=1,
              rew_scale = 1,
              distance_limit=float("inf"),
              observation_noise_stdev=0.0,
              normalization=True,

              camera_enabled=False,
              lidar_enabled=False,
              ext_disturbance_enabled=False,

              ):

    """Initialize the mors gym environment.

    Args:
      urdf_root: The path to the urdf data folder relative to parent folder.
      sim_freq: The value that used to simulate environment.
      hard_reset: Whether to wipe the simulation and load everything when reset
        is called. If set to false, reset just place the mors back to start
        position and set its pose to initial configuration.
      render: Whether to render the simulation.
      on_rack: Whether to place the mors on rack. This is only used to debug
        walking gaits. In this mode, the mors's base is hanged midair so
        that its walking gait is clearer to visualize.
      self_collision_enabled: Whether to enable self collision in the sim.
      max_timesteps: Time limit during which the simulation is going on.
      debug_mode: Whether to show side panel with debug parameters or not.
      floating_camera: If True, camera follows the robot.
      step_enabled: If True, it is needed to send motor commands as arguments in
        step() function. If False, it is neede to use set_motor_commands() function.
      accurate_motor_model_enabled: Whether to use the accurate BLDC motor model.
      simple_motor_model_enabled: Whether to use the simple version of BLDC motor model.
      torque_control_enabled: Whether to use the torque control, if set to
        False, pose control will be used.
      motor_kp: proportional gain for the accurate and simple motor models.
      motor_kd: derivative gain for the accurate and simple motor models.
      motor_velocity_limit: The velocity limit of each motor.

      action_repeat: The number of simulation steps before actions are applied.
      rew_scale: How much to gain reward. Sometimes it helps to learn a model.
      distance_limit: The maximum distance to terminate the episode.
      observation_noise_stdev: The standard deviation of observation noise.
      normalization: Whether to normalize observations or not.

          """
    self._time_step = 1/sim_freq
    self._num_bullet_solver_iterations = sim_freq*NUM_SUBSTEPS
    self._action_repeat = action_repeat
    self._urdf_root = urdf_root
    self._self_collision_enabled = self_collision_enabled
    self._motor_velocity_limit = motor_velocity_limit
    self._is_render = render
    self._world_name = world
    self._step_enabled = step_enabled
    self._floating_camera = floating_camera
    self._distance_limit = distance_limit
    self._observation_noise_stdev = observation_noise_stdev
    self._accurate_motor_model_enabled = accurate_motor_model_enabled
    self._simple_motor_model_enabled = simple_motor_model_enabled
    self._motor_kp = motor_kp
    self._motor_kd = motor_kd
    self._torque_control_enabled = torque_control_enabled
    self._on_rack = on_rack
    self._max_timesteps = max_timesteps
    self._rew_scale = rew_scale
    self._hard_reset = True
    self._debug_mode = debug_mode
    self._normalization = normalization
    self._ext_disturbance_enabled = ext_disturbance_enabled
    self._lidar_enabled = lidar_enabled
    self._camera_enabled = camera_enabled
    self.force = 136

    self._lateral_friction = 1.0
    self._spinning_friction = 0.0063

    self._ext_force_magn = 1200
    self._ext_force_duration = 6
    self._ext_force_interval = 2000
    self._force_dir = 1

    self._env_step_counter = 0
    self._action_bound = 1

    self._cam_dist = 1.0
    self._cam_yaw = 0
    self._cam_pitch = -30
    self._last_frame_time = 0.0

    self.num_legs = 4
    self.num_motors = int(self.num_legs * 3)
    self._motor_direction = [1]*12
    self._observed_motor_torques = np.zeros(self.num_motors)
    self._applied_motor_torques = np.zeros(self.num_motors)
    self._max_force = 3.5

    self.motor_angles = np.array([0]*12)
    self.motor_velocities = np.array([0]*12)
    self.motor_torques = np.array([0]*12)
    self.base_orientation = np.array([0]*3)

    # reward terms
    self.k_vel_x = 48
    self.k_body = 2
    self.k_vel_y = 3
    self.k_smooth = 0.2
    self.k_torq = 0.0003

    # motor models need smaller time step for stability.
    if self._simple_motor_model_enabled or self._accurate_motor_model_enabled:
      self._time_step /= NUM_SUBSTEPS
      self._num_bullet_solver_iterations /= NUM_SUBSTEPS
      self._action_repeat *= NUM_SUBSTEPS

    # create motor model objects
    if self._accurate_motor_model_enabled:
      self._accurate_motor_model = MotorAccurate(NUM_MOTORS)
      if self._torque_control_enabled:
        self._accurate_motor_model.set_kp(0)
        self._accurate_motor_model.set_kd(self._motor_kd)
      else:
        self._accurate_motor_model.set_kp(self._motor_kp)
        self._accurate_motor_model.set_kd(self._motor_kd)
    elif self._simple_motor_model_enabled:
      self._simple_motor_model = MotorSimple(NUM_MOTORS)
      if self._torque_control_enabled:
        self._simple_motor_model.set_kp(0)
        self._simple_motor_model.set_kd(self._motor_kd)
      else:
        self._simple_motor_model.set_kp(self._motor_kp)
        self._simple_motor_model.set_kd(self._motor_kd)

    if self._is_render:
      self._pybullet_client = bc.BulletClient(connection_mode=pybullet.GUI)
    else:
      self._pybullet_client = bc.BulletClient()
    self._pybullet_client.setAdditionalSearchPath(pybullet_data.getDataPath())

    self.action_dim = NUM_MOTORS
    action_high = np.array([self._action_bound] * self.action_dim)
    self.action_space = spaces.Box(-np.float32(action_high), np.float32(action_high), dtype=np.float32)
    self.action = np.array([0]*self.action_dim, dtype=np.float32)
    self.action_prev = np.array([0]*self.action_dim, dtype=np.float32)
    self.action_space = spaces.Box(low=np.float32(-1), high=np.float32(1), shape=(self.action_dim,), dtype=np.float32)

    self.obs_dim = 39
    self.obs = np.array([0]*self.obs_dim, dtype=np.float32)
    observation_high = (self.get_observation_upper_bound())
    observation_low = (self.get_observation_lower_bound())
    self.observation_space = spaces.Box(np.float32(observation_low), np.float32(observation_high), dtype=np.float32)

    self.world_creator = WorldCreator(self._pybullet_client, spinning_friction=self._spinning_friction, lateral_friction=self._lateral_friction)

    self.set_friction()
    self.seed()
    self.reset()

    if self._lidar_enabled == True:
      self.lidar =  Lidar(pybullet_client=self._pybullet_client,
                          robot_id=self.quadruped,
                          joint_id=18,
                          frequency=20,
                          base_frequency=sim_freq,
                          render=True,

                          angle_min=np.radians(130),
                          angle_max=-np.radians(130),
                          point_num=360,
                          range_min=0.25,
                          range_max=8.0,
                          )

    if self._camera_enabled == True:
      self.camera = Camera(pybullet_client=self._pybullet_client,
                           robot_id=self.quadruped,
                           joint_id=17,
                           frequency=20,
                           base_frequency=sim_freq)

    self._hard_reset = hard_reset # this needs to be after reset

    self.viewer = None

    self.fallen = 0

    # debug parameters
    self.view = True
    if self._debug_mode == True and self._is_render == True:
      self.reset_id = self._pybullet_client.addUserDebugParameter("Reset",1,0,1)
      self.view_id = self._pybullet_client.addUserDebugParameter("View",1,0,1)
      self.reset_clicked_prev = self._pybullet_client.readUserDebugParameter(self.reset_id)
      self.reset_clicked = self._pybullet_client.readUserDebugParameter(self.reset_id)
      self.view_clicked_prev = self._pybullet_client.readUserDebugParameter(self.view_id)
      self.view_clicked = self._pybullet_client.readUserDebugParameter(self.view_id)

      self.distance_id = self._pybullet_client.addUserDebugParameter("Distance",0.3, 2, 1.2)
      self.yaw_id = self._pybullet_client.addUserDebugParameter("Yaw",-180, 180, 47)
      self.pitch_id = self._pybullet_client.addUserDebugParameter("Pitch",-180, 180, -20)

      self.sl_id = self._pybullet_client.addUserDebugParameter("Step Length", 0, 0.2, 0)
      self.sh_id = self._pybullet_client.addUserDebugParameter("Step Height", 0, 0.1, 0)
      self.omega_id = self._pybullet_client.addUserDebugParameter("Step Frequency", 0, 4, 2.9)
      self.psi_id = self._pybullet_client.addUserDebugParameter("Sidewalk Angle", -3.14, 3.14, 0)
      self.rot_dir_id = self._pybullet_client.addUserDebugParameter("Turning Direction", -1.0, 1.0, 0.0)
      self.rot_r_id = self._pybullet_client.addUserDebugParameter("Turning Radius", 0, 10, 5)

      self.body_tr_x_id = self._pybullet_client.addUserDebugParameter("Body Translation X", -0.0003, 0.0003, 0)
      self.body_tr_y_id = self._pybullet_client.addUserDebugParameter("Body Translation Y", -0.0003, 0.0003, 0)
      self.body_tr_z_id = self._pybullet_client.addUserDebugParameter("Body Translation Z", -0.0003, 0.0003, 0)

      self.body_rot_x_id = self._pybullet_client.addUserDebugParameter("Body Rotation X", -0.5, 0.5, 0)
      self.body_rot_y_id = self._pybullet_client.addUserDebugParameter("Body Rotation Y", -0.4, 0.4, 0)
      self.body_rot_z_id = self._pybullet_client.addUserDebugParameter("Body Rotation Z", -0.4, 0.4, 0)

    self.first = True

  def reset(self):
    """
    Resets all Mors parameters.
    If hard_reset is True, it also resets simulation.
    """
    if self._hard_reset:
      self._pybullet_client.resetSimulation()
      self._pybullet_client.setPhysicsEngineParameter(
          numSolverIterations=int(self._num_bullet_solver_iterations))
      self._pybullet_client.setTimeStep(self._time_step)

      # plane = self._pybullet_client.loadURDF("%s/plane.urdf" % pybullet_data.getDataPath())
      # self._pybullet_client.changeVisualShape(plane, -1, rgbaColor=[1, 1, 1, 0.9])
      # self._pybullet_client.configureDebugVisualizer(self._pybullet_client.COV_ENABLE_PLANAR_REFLECTION, 0)
      self._pybullet_client.configureDebugVisualizer(self._pybullet_client.COV_ENABLE_SHADOWS, 0)
      self._pybullet_client.setGravity(0, 0, -9.81)
      # self.set_world(self._world_name)
      self.world_creator.create_world(self._world_name, lateralFriction=self._lateral_friction, spinningFriction=self._spinning_friction)
      self._reset_robot(reload_urdf=True)
    else:
      self._reset_robot(reload_urdf=False)

    self._env_step_counter = 0
    # self._last_base_position = np.array([0, 0, 0])
    # self._objectives = []
    self._pybullet_client.resetDebugVisualizerCamera(self._cam_dist, self._cam_yaw,
                                                     self._cam_pitch, [0, 0, 0])

    self.obs = self._noisy_observation()
    self.obs = self.norm_obs(self.obs)
    return self.obs

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

  def step(self, action):
    """Step forward the simulation, given the action.

    Args:
      action: A list of desired motor angles for eight motors.

    Returns:
      observations: The angles, velocities and torques of all motors.
      reward: The reward for the current state-action pair.
      done: Whether the episode has ended.
      info: A dictionary that stores diagnostic information.

    Raises:
      ValueError: The action dimension is not the same as the number of motors.
      ValueError: The magnitude of actions is out of bounds.
    """
    # debug parameters

    if self._is_render:
      if self._debug_mode == True:
        # reset button
        self.reset_clicked = self._pybullet_client.readUserDebugParameter(self.reset_id)
        if self.reset_clicked != self.reset_clicked_prev:
            self._reset_robot(reload_urdf=False)
        self.reset_clicked_prev = self.reset_clicked
        # view button
        self.view_clicked = self._pybullet_client.readUserDebugParameter(self.view_id)
        if self.view_clicked != self.view_clicked_prev:
            self.view = not self.view
        self.view_clicked_prev = self.view_clicked
        # camera position
        distance = self._pybullet_client.readUserDebugParameter(self.distance_id)
        pitch = self._pybullet_client.readUserDebugParameter(self.pitch_id)
        yaw = self._pybullet_client.readUserDebugParameter(self.yaw_id)
      else:
        distance = 0.5
        yaw = -10
        pitch = -20
      # Sleep, otherwise the computation takes less time than real time,
      # which will make the visualization like a fast-forward video.
      time_spent = time.time() - self._last_frame_time
      self._last_frame_time = time.time()
      time_to_sleep = self._action_repeat * self._time_step - time_spent
      if time_to_sleep > 0:
        time.sleep(time_to_sleep)
      base_pos, _ = self.get_base_position_and_orientation()
      if self.view == True and self._floating_camera == False:
        self._pybullet_client.resetDebugVisualizerCamera(distance, yaw, pitch, base_pos)

    if self._ext_disturbance_enabled:
      self._apply_force()

    if self._lidar_enabled:
      self.lidar.update(self._env_step_counter)
    # if self._camera_enabled == True:
    #   self.camera.update(self._env_step_counter)

    self.action = action
    for _ in range(self._action_repeat):
      self._apply_action(self.action)
      self._pybullet_client.stepSimulation()

    self._env_step_counter += 1
    reward = self._rew_scale*self._reward()
    done = self._termination()
    self.obs = self._noisy_observation()
    if self._normalization == True:
      self.obs = self.norm_obs(self.obs)
    self.action_prev[:] = self.action[:]
    return self.obs, reward, done, {}

  def render(self, mode="rgb_array", close=False):
    if mode != "rgb_array":
      return np.array([])
    base_pos, _ = self.get_base_position_and_orientation()
    view_matrix = self._pybullet_client.computeViewMatrixFromYawPitchRoll(
        cameraTargetPosition=base_pos,
        distance=self._cam_dist,
        yaw=self._cam_yaw,
        pitch=self._cam_pitch,
        roll=0,
        upAxisIndex=2)
    proj_matrix = self._pybullet_client.computeProjectionMatrixFOV(fov=60,
                                                                   aspect=float(RENDER_WIDTH) /
                                                                   RENDER_HEIGHT,
                                                                   nearVal=0.1,
                                                                   farVal=100.0)
    (_, _, px, _,
     _) = self._pybullet_client.getCameraImage(width=RENDER_WIDTH,
                                               height=RENDER_HEIGHT,
                                               viewMatrix=view_matrix,
                                               projectionMatrix=proj_matrix,
                                               renderer=pybullet.ER_BULLET_HARDWARE_OPENGL)
    rgb_array = np.array(px)
    rgb_array = rgb_array[:, :, :3]
    return rgb_array

  def close(self):
    """
    Close the environment and disconnect physical engine.
    It is recommended to use this function at the end of your program.
    """
    super().close()
    self._pybullet_client.disconnect()

  def is_fallen(self):
    """Decide whether the mors has fallen.

    If the up directions between the base and the world is larger (the dot
    product is smaller than 0.85) or the base is very low on the ground
    (the height is smaller than 0.13 meter), the mors is considered fallen.

    Returns:
      Boolean value that indicates whether the mors has fallen.
    """
    orientation = self.get_base_orientation()
    rot_mat = self._pybullet_client.getMatrixFromQuaternion(orientation)
    local_up = rot_mat[6:]
    pos, _ = self.get_base_position_and_orientation()
    return (np.dot(np.asarray([0, 0, 1]), np.asarray(local_up)) < 0.5 or pos[2] < 0.05)

  def _termination(self):
    position, _ = self.get_base_position_and_orientation()
    distance = np.sqrt(position[0]**2 + position[1]**2)
    if self.is_fallen():
      self.fallen = 1
    else:
      self.fallen = 0
    if self._env_step_counter >= self._max_timesteps:
      self.timeout = 1
    else:
      self.timeout = 0
    return bool(self.timeout or self.fallen or (distance > self._distance_limit))

  def _reward(self):
    """
    Calculate reward function
    Returns:
      reward value during current timestep
    """
    orient = np.abs(self.get_base_orientation())
    ang_vel = self.get_base_ang_vel()
    lin_vel = self.get_base_lin_vel()
    d_theta_ref_offset = self.action[0:12] - self.action_prev[0:12]

    self.vel_x_reward = lin_vel[0]**2
    self.body_reward = np.exp(-5*np.abs(ang_vel[0])) + np.exp(-5*np.abs(ang_vel[1])) + np.exp(-40*orient[0]) + np.exp(-40*orient[1])
    self.vel_y_reward = np.exp(-(lin_vel[1])**2/0.005)
    self.smooth_reward = np.exp(-np.linalg.norm(self.action[0:12], 1)) + np.exp(-np.linalg.norm(d_theta_ref_offset, 1))
    self.torque_penalty = np.linalg.norm(self.get_motor_torques()[:], 1)**2

    reward = (self.k_vel_x*self.vel_x_reward
              + self.k_vel_y *self.vel_y_reward
              + self.k_body*self.body_reward
              + self.k_smooth*self.smooth_reward
              - self.k_torq*self.torque_penalty)
    return reward

  def set_reward_terms(self, k_vel_x=48, k_vel_y=3, k_body=2, k_smooth=0.2, k_torq=0.0003):
    self.k_vel_x = k_vel_x
    self.k_vel_y = k_vel_y
    self.k_body = k_body
    self.k_smooth = k_smooth
    self.k_torq = k_torq

  def get_reward_parts(self):
    return [self.k_vel_x*self.vel_x_reward , self.k_vel_y *self.vel_y_reward,
            self.k_body*self.body_reward, self.k_smooth*self.smooth_reward,
            -self.k_torq*self.torque_penalty]

  def _noisy_observation(self):
    observation = self._get_observation()
    obs = np.array(observation, dtype=np.float32)
    if self._observation_noise_stdev > 0:
      obs += np.random.normal(scale=self._observation_noise_stdev, size=obs.shape)
    return obs

  def _get_observation(self):
    """Get the observations of mors.

    It includes the angles, velocities, torques and the orientation of the base.

    Returns:
      The observation list
      observation[0:11] are motor angles
      observation[12:23] are motor velocities
      observation[24:35] are motor torques
      observation[35:38] is the orientation of the base, in Euler form.
    """
    observation = []
    self.motor_angles = self.get_motor_angles().tolist()
    self.motor_velocities = self.get_motor_velocities().tolist()
    self.motor_torques = self.get_motor_torques().tolist()
    self.base_orientation = list(self.get_base_orientation_euler())
    observation.extend(self.motor_angles)
    observation.extend(self.motor_velocities)
    observation.extend(self.motor_torques)
    observation.extend(self.base_orientation)
    return observation

  def norm_obs(self, obs):
    norm_obs = np.array([0]*self.obs_dim, dtype=np.float32)
    for i in range(0,12):
        norm_obs[i] = (obs[i]+MOTOR_ANGLE_MAX)/(MOTOR_ANGLE_MAX+MOTOR_ANGLE_MAX)
    for i in range(12,24):
        norm_obs[i] = np.clip((obs[i]+MOTOR_VELOCITY_MAX)/(MOTOR_VELOCITY_MAX+MOTOR_VELOCITY_MAX), 0, 1)
    for i in range(24, 36):
        norm_obs[i] = np.clip((obs[i]+MOTOR_TORQUE_MAX)/(MOTOR_TORQUE_MAX+MOTOR_TORQUE_MAX), 0, 1)
    norm_obs[36] = (obs[36]+BASE_ORIENTATION_MAX)/(BASE_ORIENTATION_MAX+BASE_ORIENTATION_MAX)
    norm_obs[37] = (obs[37]+BASE_ORIENTATION_MAX)/(BASE_ORIENTATION_MAX+BASE_ORIENTATION_MAX)
    norm_obs[38] = (obs[38]+BASE_ORIENTATION_MAX)/(BASE_ORIENTATION_MAX+BASE_ORIENTATION_MAX)

    return norm_obs

  def _denorm_obs(self, norm_obs):
    obs = np.array([0]*self.obs_dim, dtype=np.float32)
    for i in range(0,12):
        obs[i] = 2*norm_obs[i]*MOTOR_ANGLE_MAX - MOTOR_ANGLE_MAX
    for i in range(12,24):
        obs[i] = 2*norm_obs[i]*MOTOR_VELOCITY_MAX - MOTOR_VELOCITY_MAX
    for i in range(24, 36):
        obs[i] = 2*norm_obs[i]*MOTOR_TORQUE_MAX - MOTOR_TORQUE_MAX
    obs[36] = 2*BASE_ORIENTATION_MAX*norm_obs[36] - BASE_ORIENTATION_MAX
    obs[37] = 2*BASE_ORIENTATION_MAX*norm_obs[37] - BASE_ORIENTATION_MAX
    obs[38] = 2*BASE_ORIENTATION_MAX*norm_obs[38] - BASE_ORIENTATION_MAX

    return obs

  def set_kpkd(self, kp : float, kd : float):
    self._motor_kp = kp
    self._motor_kd = kd

  def set_motor_commands(self, angles, vels, torques):
    self.ref_angles = angles
    self.ref_vels = vels
    self.ref_torques = torques

  def _apply_action(self, motor_commands):
    """Set the desired motor commands to the motors of the mors.

    The desired motor commands are clipped based on the maximum allowed velocity.
    If the pd_control_enabled is True, a torque is calculated according to
    the difference between current and desired joint angle, as well as the joint
    velocity. This torque is exerted to the motor. For more information about
    PD control, please refer to: https://en.wikipedia.org/wiki/PID_controller.

    Args:
      motor_commands: The twelve desired motor angles or torques depending on torque_control_enabled variable.
    """
    current_motor_angle = self.get_motor_angles()
    if self._motor_velocity_limit < np.inf:

      motor_commands_max = (current_motor_angle + self._time_step * self._motor_velocity_limit)
      motor_commands_min = (current_motor_angle - self._time_step * self._motor_velocity_limit)
      motor_commands = np.clip(motor_commands, motor_commands_min, motor_commands_max)

    motor_commands_with_direction = np.multiply(motor_commands, self._motor_direction)
    self.motor_angles = current_motor_angle[:].tolist()
    self.motor_velocities = self.get_motor_velocities().tolist()

    if self._accurate_motor_model_enabled:
      self._accurate_motor_model.set_sensor_data(self.motor_angles, self.motor_velocities)

      if self._step_enabled == True:
        if self._torque_control_enabled:
          self._accurate_motor_model.set_ref_torque(motor_commands_with_direction)
          self._accurate_motor_model.set_kd(self._motor_kd)
        else:
          self._accurate_motor_model.set_ref_angle(motor_commands_with_direction)
          self._accurate_motor_model.set_kp(self._motor_kp)
          self._accurate_motor_model.set_kd(self._motor_kd)

      else:
        self._accurate_motor_model.set_ref_angle(self.ref_angles)
        self._accurate_motor_model.set_ref_vel(self.ref_vels)
        self._accurate_motor_model.set_ref_torque(self.ref_torques)
        self._accurate_motor_model.set_kp(self._motor_kp)
        self._accurate_motor_model.set_kd(self._motor_kd)

      self.tau = self._accurate_motor_model.step()

      for motor_id, tau in zip(self._motor_id_list, self.tau):
        self._SetMotorTorqueById(motor_id, float(tau))

    elif self._simple_motor_model_enabled:
        self._simple_motor_model.set_sensor_data(self.motor_angles, self.motor_velocities)

        if self._step_enabled == True:
          if self._torque_control_enabled:
            self._simple_motor_model.set_ref_torque(motor_commands_with_direction)
            self._simple_motor_model.set_kd(self._motor_kd)
          else:
            self._simple_motor_model.set_ref_angle(motor_commands_with_direction)
            self._simple_motor_model.set_kp(self._motor_kp)
            self._simple_motor_model.set_kd(self._motor_kd)

        else:
          self._simple_motor_model.set_ref_angle(self.ref_angles)
          self._simple_motor_model.set_ref_vel(self.ref_vels)
          self._simple_motor_model.set_ref_torque(self.ref_torques)
          self._simple_motor_model.set_kp(self._motor_kp)
          self._simple_motor_model.set_kd(self._motor_kd)

        self.tau = self._simple_motor_model.step()

        for motor_id, tau in zip(self._motor_id_list, self.tau):
          self._SetMotorTorqueById(motor_id, float(tau))
    else:

      if self._torque_control_enabled == False:
        if self._step_enabled == True:
          for motor_id, motor_command_with_direction in zip(self._motor_id_list, motor_commands_with_direction):
            self._SetMotorAngleById(motor_id, motor_command_with_direction)
        else:
          for motor_id, motor_command_with_direction in zip(self._motor_id_list, self.ref_angles):
            self._SetMotorAngleById(motor_id, motor_command_with_direction)

      else:
        if self._step_enabled == True:
          for motor_id, motor_command_with_direction in zip(self._motor_id_list, motor_commands_with_direction):
            self._SetMotorTorqueById(motor_id, float(motor_command_with_direction))
        else:
          for motor_id, motor_command_with_direction in zip(self._motor_id_list, self.ref_torques):
            self._SetMotorTorqueById(motor_id, float(motor_command_with_direction))

  def _RecordMassInfoFromURDF(self):
    self._base_mass_urdf = self._pybullet_client.getDynamicsInfo(self.quadruped, BASE_LINK_ID)[0]
    self._leg_masses_urdf = []
    self._leg_masses_urdf.append(
        self._pybullet_client.getDynamicsInfo(self.quadruped, LEG_LINK_ID[0])[0])
    self._leg_masses_urdf.append(
        self._pybullet_client.getDynamicsInfo(self.quadruped, MOTOR_LINK_ID[0])[0])

  def _BuildJointNameToIdDict(self):
    num_joints = self._pybullet_client.getNumJoints(self.quadruped)
    self._joint_name_to_id = {}
    for i in range(num_joints):
      joint_info = self._pybullet_client.getJointInfo(self.quadruped, i)
      self._joint_name_to_id[joint_info[1].decode("UTF-8")] = joint_info[0]

  def _BuildMotorIdList(self):
    self._motor_id_list = [self._joint_name_to_id[motor_name] for motor_name in MOTOR_NAMES]

  def _SetMotorTorqueById(self, motor_id, torque : float):
    self._pybullet_client.setJointMotorControl2(bodyIndex=self.quadruped,
                                                jointIndex=motor_id,
                                                controlMode=pybullet.TORQUE_CONTROL,
                                                force=torque)

  def _SetMotorVelocityById(self, motor_id, desired_velocity):
    self._pybullet_client.setJointMotorControl2(bodyIndex=self.quadruped,
                                                jointIndex=motor_id,
                                                controlMode=self._pybullet_client.VELOCITY_CONTROL,
                                                targetVelocity=desired_velocity,
                                                #positionGain=self._kp,
                                                #velocityGain=self._kd,
                                                maxVelocity=self._motor_velocity_limit,
                                                force=self.force
                                                )

  def _SetMotorAngleById(self, motor_id, desired_angle):
    self._pybullet_client.setJointMotorControl2(bodyIndex=self.quadruped,
                                                jointIndex=motor_id,
                                                controlMode=self._pybullet_client.POSITION_CONTROL,
                                                targetPosition=desired_angle,
                                                #positionGain=self._kp,
                                                #velocityGain=self._kd,
                                                maxVelocity=self._motor_velocity_limit,
                                                force=self.force#self._max_force
                                                )

  def _SetMotorAngleByName(self, motor_name, desired_angle):
    self._SetMotorAngleById(self._joint_name_to_id[motor_name], desired_angle)

  def _ResetPose(self):
    """Reset the pose of the mors.
      constraint: Whether to add a constraint at the joints of two feet.
    """
    for i in range(self.num_legs):
      self._ResetPoseForLeg(i)

    for j in range(16):
            self._pybullet_client.setJointMotorControl2(self.quadruped, j, pybullet.VELOCITY_CONTROL, targetVelocity=0, force=0)

  def _ResetPoseForLeg(self, leg_id):
    """Reset the initial pose for the leg.

    Args:
      leg_id: It should be 0, 1, 2, or 3, which represents the leg at
        front_left, back_left, front_right and back_right.
    """
    knee_angle = 0
    shoulder_angle = 0
    hip_angle = 0

    leg_position = LEG_POSITION[leg_id]
    self._pybullet_client.resetJointState(self.quadruped,
                                          self._joint_name_to_id["abad_joint_" + leg_position],
                                          self._motor_direction[2 * leg_id] * shoulder_angle,
                                          targetVelocity=0)
    self._pybullet_client.resetJointState(self.quadruped,
                                          self._joint_name_to_id["hip_joint_" + leg_position],
                                          self._motor_direction[2 * leg_id] * knee_angle,
                                          targetVelocity=0)
    self._pybullet_client.resetJointState(self.quadruped,
                                          self._joint_name_to_id["knee_joint_" + leg_position],
                                          self._motor_direction[2 * leg_id + 1] * hip_angle,
                                          targetVelocity=0)

  def _reset_robot(self, reload_urdf=True):
    self._urdf_root = self._urdf_root.replace('.', '')
    addr = parentdir + self._urdf_root + "/mors.urdf"
    print(addr)
    if reload_urdf:
      if self._self_collision_enabled:
        self.quadruped = self._pybullet_client.loadURDF(
            addr,
            INIT_POSITION,
            flags=(pybullet.URDF_USE_SELF_COLLISION_INCLUDE_PARENT | pybullet.URDF_USE_SELF_COLLISION))
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, -1, 0, 0)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, -1, 1, 0)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, -1, 2, 1)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, 0, 1, 0)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, 1, 2, 0)

        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, -1, 4, 0)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, -1, 5, 0)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, -1, 6, 1)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, 4, 5, 0)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, 5, 6, 1)

        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, -1, 8, 0)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, -1, 9, 0)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, -1, 10, 1)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, 8, 9, 0)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, 9, 10, 0)

        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, -1, 12, 0)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, -1, 13, 0)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, -1, 14, 1)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, 12, 13, 0)
        self._pybullet_client.setCollisionFilterPair(self.quadruped, self.quadruped, 13, 14, 0)
      else:
        self.quadruped = self._pybullet_client.loadURDF(
            addr, INIT_POSITION)
      if self._on_rack:
        self._pybullet_client.createConstraint(self.quadruped, BASE_LINK_ID, -1, -1,
                                               self._pybullet_client.JOINT_FIXED, [0, 0, 0],
                                               [0, 0, 0], INIT_POSITION)
      self._BuildJointNameToIdDict()
      self._BuildMotorIdList()
      self._RecordMassInfoFromURDF()
      self._ResetPose()

      for joint_number in range(self._pybullet_client.getNumJoints(self.quadruped)):
        self._pybullet_client.enableJointForceTorqueSensor(self.quadruped, joint_number)

    else:
      self._pybullet_client.resetBasePositionAndOrientation(self.quadruped, INIT_POSITION,
                                                            INIT_ORIENTATION)
      self._pybullet_client.resetBaseVelocity(self.quadruped, [0, 0, 0], [0, 0, 0])
      self._ResetPose()

    self._motor_enabled_list = [True] * self.num_motors

  def get_action_dimension(self):
    """Get the length of the action list.

    Returns:
      The length of the action list.
    """
    return self.action_dim

  def get_observation_upper_bound(self):
    """Get the upper bound of the observation.

    Returns:
    The upper bound of an observation. See GetObservation() for the details
        of each element of an observation.
    """
    upper_bound = np.array([0.0] * self.get_observation_dimension())
    upper_bound[:] = 1.0  # Quaternion of base orientation.
    return upper_bound

  def get_observation_lower_bound(self):
    """Get the lower bound of the observation."""
    lower_bound = np.array([0.0] * self.get_observation_dimension())
    lower_bound[:] = 0.0  # Quaternion of base orientation.
    return lower_bound

  def get_observation_dimension(self):
    """Get the length of the observation list.

    Returns:
    The length of the observation list.
    """
    return self.obs.size

  #--------------------------
  # get sensor data
  #--------------------------
  def get_base_orientation(self):
    """Get the orientation of mors's base, represented as quaternion.

    Returns:
      The orientation of mors's base.
    """
    _, orientation = (self._pybullet_client.getBasePositionAndOrientation(self.quadruped))
    return orientation

  def get_base_orientation_euler(self):
    """Get the orientation of mors's base, represented as Euler angles.

    Returns:
      The orientation of mors's base.
    """
    _, orientation = (self._pybullet_client.getBasePositionAndOrientation(self.quadruped))
    euler_orient = self._pybullet_client.getEulerFromQuaternion(orientation)
    return euler_orient

  def get_base_position_and_orientation(self):
    position, quat = self._pybullet_client.getBasePositionAndOrientation(self.quadruped)
    return position, quat

  def get_base_lin_vel(self):
    vel = self._pybullet_client.getBaseVelocity(self.quadruped)
    return vel[0]

  def get_base_ang_vel(self):
    vel = self._pybullet_client.getBaseVelocity(self.quadruped)
    return vel[1]

  def get_motor_states(self):
    joint_number_range = range(self._pybullet_client.getNumJoints(self.quadruped))
    joint_states = self._pybullet_client.getJointStates(self.quadruped, joint_number_range)
    joint_infos = [self._pybullet_client.getJointInfo(self.quadruped, i) for i in joint_number_range]
    joint_states, joint_name = \
        zip(*[(j, i[1]) for j, i in zip(joint_states, joint_infos) if i[2] != self._pybullet_client.JOINT_FIXED])
    joint_positions = [state[0] for state in joint_states]
    joint_velocities = [state[1] for state in joint_states]
    joint_reaction_forces = [state[2] for state in joint_states]
    joint_torques = [0]*12
    for i in range(12):
      joint_torques[i] = joint_reaction_forces[i][5]
    # print(f"{[ '%.3f' % elem for elem in joint_torques ]}")
    # print(joint_name)
    return joint_positions, joint_velocities, joint_torques, joint_name

  def get_motor_angles(self):
    """Get the eight motor angles at the current moment.

    Returns:
      Motor angles.
    """
    motor_angles = [
        self._pybullet_client.getJointState(self.quadruped, motor_id)[0]
        for motor_id in self._motor_id_list
    ]
    motor_angles = np.multiply(motor_angles, self._motor_direction)
    return motor_angles

  def get_motor_velocities(self):
    """Get the velocity of all motors.

    Returns:
      Velocities of all motors.
    """
    motor_velocities = [
        self._pybullet_client.getJointState(self.quadruped, motor_id)[1]
        for motor_id in self._motor_id_list
    ]
    motor_velocities = np.multiply(motor_velocities, self._motor_direction)
    return motor_velocities

  def get_motor_torques(self):
    """Get the amount of torques the motors are exerting.

    Returns:
      Motor torques of all twelve motors.
    """
    motor_torques = [
        self._pybullet_client.getJointState(self.quadruped, motor_id)[3]
        for motor_id in self._motor_id_list
    ]
    motor_torques = np.multiply(motor_torques, self._motor_direction)
    return motor_torques

  def GetCOMPos(self):
    base_pos, base_orientation = self._pybullet_client.getBasePositionAndOrientation(self.quadruped)
    com_pos_loc = [-0.002185, -0.000018, 0.000286]
    com_orient_loc = [0.0, 0.0, 0.0, 1.0]

    com_pos, com_orient = self._pybullet_client.multiplyTransforms(com_pos_loc, com_orient_loc, base_pos, base_orientation)
    return com_pos

  def GetEFPos(self):
    r1_pos, r1_or,_,_,_,_  = self._pybullet_client.getLinkState(self.quadruped, 4)
    l1_pos, l1_or,_,_,_,_  = self._pybullet_client.getLinkState(self.quadruped, 12)
    r2_pos, r2_or,_,_,_,_  = self._pybullet_client.getLinkState(self.quadruped, 8)
    l2_pos, l2_or,_,_,_,_  = self._pybullet_client.getLinkState(self.quadruped, 16)

    ef = [0]*12
    for i in range(3):
      ef[i] = r1_pos[i]
      ef[i+3] = l1_pos[i]
      ef[i+6] = r2_pos[i]
      ef[i+9] = l2_pos[i]

    return ef

  def get_contact_flags(self):
    self.contact_points = self._pybullet_client.getContactPoints(self.quadruped)
    contact_flags = [False]*4

    for contact_point in self.contact_points:
      if contact_point[3] == 3:
          contact_flags[1] = True
      if contact_point[3] == 7:
          contact_flags[0] = True
      if contact_point[3] == 11:
          contact_flags[3] = True
      if contact_point[3] == 15:
          contact_flags[2] = True
    return self.contact_points, contact_flags

  # ----------------------
  # Physical Parameters
  # ----------------------
  def get_base_mass_from_urdf(self):
    """Get the mass of the base from the URDF file."""
    return self._base_mass_urdf

  def get_leg_masses_from_urdf(self):
    """Get the mass of the legs from the URDF file."""
    return self._leg_masses_urdf

  def set_base_mass(self, base_mass):
    self._pybullet_client.changeDynamics(self.quadruped, BASE_LINK_ID, mass=base_mass)

  def set_leg_masses(self, shoulder, hip, knee):
    for i in range(4):
      self._pybullet_client.changeDynamics(self.quadruped, i*4, mass=shoulder)
      self._pybullet_client.changeDynamics(self.quadruped, i*4+1, mass=hip)
      self._pybullet_client.changeDynamics(self.quadruped, i*4+2, mass=knee)

  def set_foot_friction(self, foot_friction):
    """Set the lateral friction of the feet.

    Args:
      foot_friction: The lateral friction coefficient of the foot. This value is
        shared by all four feet.
    """
    # for link_id in FOOT_LINK_ID:
      # self._pybullet_client.changeDynamics(self.quadruped, link_id, lateralFriction=foot_friction)
    self._pybullet_client.changeDynamics(self.quadruped, 3, lateralFriction=self._lateral_friction, spinningFriction=self._spinning_friction)
    self._pybullet_client.changeDynamics(self.quadruped, 7, lateralFriction=self._lateral_friction, spinningFriction=self._spinning_friction)
    self._pybullet_client.changeDynamics(self.quadruped, 11, lateralFriction=self._lateral_friction, spinningFriction=self._spinning_friction)
    self._pybullet_client.changeDynamics(self.quadruped, 15, lateralFriction=self._lateral_friction, spinningFriction=self._spinning_friction)

  def set_friction(self, lateral_friction=1.0, spinning_friction=0.0063):
    self._lateral_friction = lateral_friction
    self._spinning_friction = spinning_friction

  def set_joint_forces(self, force):
    self.force = force

  def set_joint_max_velocity(self, velocity):
    self._motor_velocity_limit = velocity

  def set_control_mode(self, mode):
    self._control_mode = mode
    if mode == pybullet.TORQUE_CONTROL:
      if self.first == True:
        self.first = False
        for j in range(16):
              self._pybullet_client.setJointMotorControl2(self.quadruped, j, pybullet.VELOCITY_CONTROL, force=0)
    else:
      self.first = True

  # ----------------------
  # Terrain Parameters
  # ----------------------
  def set_world(self, world):
    self._world_name = world

  # def set_reset_type(self, type : bool):
  #   self._hard_reset = type

  # ----------------------
  # External Forces
  # ----------------------
  def set_ext_forces_params(self, magn, duration, interval): #duration and interval in timesteps
    self._ext_force_magn = magn
    self._ext_force_duration = duration
    self._ext_force_interval = interval

  def _apply_force(self):
    if 100 < (self._env_step_counter%self._ext_force_interval) <= (100+self._ext_force_duration):
        if (self._env_step_counter%self._ext_force_interval) < 102:
            self._force_dir = -self._force_dir
        force = [0, self._force_dir*self._ext_force_magn, 0]
        print(force)
        self._pybullet_client.applyExternalForce(self.quadruped, -1, force, [0,0,0], pybullet.LINK_FRAME)
        com_pos = self.GetCOMPos()
        self._pybullet_client.addUserDebugLine([com_pos[0], com_pos[1], com_pos[2]],
                                                    [com_pos[0]+force[0]/1000, com_pos[1]+force[1]/1000, com_pos[2]+force[2]/1000],
                                                    [1, 0, 0], 3, 0.4)


import pandas as pd
from cassandra.cluster import Cluster
import re
import os
import glob
import numpy as np
import json
import csv
from cql_queries import *

def process_app_history_files():
    """
    This function reads all csv files from the event_data folder,
    and writes the following columns into event_datafile_new.csv file.

    Inputs:
    -------
    This function does not have any input.

    Output:
    -------
    This function does not return anything.
    """

    # Get your current folder and subfolder event data
    filepath = os.getcwd() + '/event_data'

    # Create a for loop to create a list of files and collect each filepath
    for root, dirs, files in os.walk(filepath):

        # join the file path and roots with the subdirectories using glob
        file_path_list = glob.glob(os.path.join(root,'*'))

    # initiating an empty list of rows that will be generated from each file
    full_data_rows_list = []

    # for every filepath in the file path list
    for f in file_path_list:

    # reading csv file
        with open(f, 'r', encoding = 'utf8', newline='') as csvfile:
            # creating a csv reader object
            csvreader = csv.reader(csvfile)
            next(csvreader)

             # extracting each data row one by one and append it
            for line in csvreader:
                #print(line)
                full_data_rows_list.append(line)

    # creating a smaller event data csv file called event_datafile_full csv that will be used to   insert data into the \
    # Apache Cassandra tables
    csv.register_dialect('myDialect', quoting=csv.QUOTE_ALL, skipinitialspace=True)

    with open('event_datafile_new.csv', 'w', encoding = 'utf8', newline='') as f:
        writer = csv.writer(f, dialect='myDialect')
        writer.writerow(['artist','firstName','gender','itemInSession','lastName','length',\
                            'level','location','sessionId','song','userId'])
        for row in full_data_rows_list:
            if (row[0] == ''):
                continue
            writer.writerow((row[0], row[2], row[3], row[4], row[5], row[6], row[7], row[8], row[12], row[13], row[16]))

def load_data(session, filepath):
    """
    This function reads a file and load it into the Apache Cassandra tables.

    Inputs:
    -------
    session - object represents connection to an Apache Cassandra cluster
    filepath - access point of the csv file loaded into Apache Cassandra tables

    Output:
    -------
    This function does not return anything.
    """

    with open(filepath, encoding = 'utf8') as f:
        csvreader = csv.reader(f)
        next(csvreader) # skip header
        for line in csvreader:
            # load session data
            session.execute(session_table_insert, (line[0], line[9], float(line[5]), int(line[8]), int(line[3])))
            # load session playlist data
            session.execute(session_playlist_table_insert, (line[0], line[9], line[1], line[4], int(line[10]), int(line[8]), int(line[3])))
            # load song data
            session.execute(song_table_insert, (line[1], line[4], line[9], int(line[10])))

def main():
    print("Connect...\n")
    cluster = Cluster()
    session = cluster.connect()
    session.set_keyspace("sparkifydb")

    print("Extract...\n\nTransform...\n")
    process_app_history_files()

    print("Load...\n")
    load_data(session, filepath="./event_datafile_new.csv")

    print("Disconnect...")
    session.shutdown()
    cluster.shutdown()

if __name__ == "__main__":
    main()


# Напишите программу, удаляющую из текста все слова, содержащие ""абв""
#
# my_text = 'abvgteksb, rtegdjsha, dhfjgbnmxg, ooeoia, rtiop, rats'
# my_text_list = list(filter(lambda i:'a'not in i and 'b' not in i and 'v' not in i, my_text.split()))
# print(my_text_list)

# Реализуйте RLE алгоритм: реализуйте модуль сжатия и восстановления данных.
# Входные и выходные данные хранятся в отдельных текстовых файлах.
# def compress (file):
#     with open (file, 'r') as file:
#         same_str=file.readline()
#         count=1
#         i=0
#         compress_str=''
#         while i< len(same_str)-1:
#             if same_str[i]==same_str[i+1]:
#                 count+=1
#             else:
#                 if count == 1:
#                     compress_str +=same_str[i]
#                 else:
#
#                     compress_str+=str(count)+same_str[i]
#                 count=1
#             i+=1
#         return compress_str
#
#
# compress=compress('Dz5.txt')
#
# def decompress (compress):
#     decompress_str=''
#     for i in range(1,len(compress)):
#         if compress[i-1].isdigit():
#             decompress_str+=(int(compress[i-1])*compress[i])
#         else:
#             if compress[i].isdigit():
#                 i+=1
#             else:
#                 decompress_str += compress[i]
#     return decompress_str
#
# decompress=decompress(compress)
# print (compress)
# print (decompress)

# Создайте программу для игры в ""Крестики-нолики"".

def draw_board(board):
   for i in range(3):
      print("|", board[0+i*3], "|", board[1+i*3], "|", board[2+i*3], "|")

def inpute(board,igrok):
   valid=False
   while not valid:
      ansver=int(input(f'Vvedite nomer kletki dlya {igrok}=> '))
      if ansver >0 and ansver <10:
         if (str(board[ansver-1])) not in 'XO':
            board[ansver-1]= igrok
            valid=True
         else:
            print('Kletka Zanyata')
      else:
         print('Nekoreknyi VVod')

def Pobeda (board):
   win=False

   for i in board:
      if ('X'== board [0] and 'X'== board [1] and 'X'== board [2]) or ('X'== board[3] and 'X'== board [4] and 'X'== board [5]) or ('X'== board[6] and 'X'== board [7] and 'X'== board [8]) or ('X'== board[0] and 'X'== board [3] and 'X'== board [6]) or ('X'== board[1] and 'X'== board [4] and 'X'== board [7]) or ('X'== board[2] and 'X'== board [5] and 'X'== board [8]) or ('X'== board[0] and 'X'== board [4] and 'X'== board [8]) or ('X'== board[6] and 'X'== board [4] and 'X'== board [6]):
         print ('Pobeda X')
         print(draw_board(board))
         win=True
         break
      if ('O'== board [0] and 'O'== board [1] and 'O'== board [2]) or ('O'== board[3] and 'O'== board [4] and 'O'== board [5]) or ('O'== board[6] and 'O'== board [7] and 'O'== board [8]) or ('O'== board[0] and 'O'== board [3] and 'O'== board [6]) or ('O'== board[1] and 'O'== board [4] and 'O'== board [7]) or ('O'== board[2] and 'O'== board [5] and 'O'== board [8]) or ('O'== board[0] and 'O'== board [4] and 'O'== board [8]) or ('O'== board[6] and 'O'== board [4] and 'O'== board [6]):
         print ('Pobeda O')
         print(draw_board(board))
         win = True
         break
   return win

board = list(range(1,10))
igrok='O'
while not Pobeda(board):
   draw_board(board)
   inpute(board, igrok)
   Pobeda(board)
   if igrok=='O':
      igrok='X'
   else:
      igrok = 'O'


#!/usr/bin/env python
import time
import sys
import rospy
from pySerialTransfer import pySerialTransfer as txfer
from messages import RxLog, RxOdometry, RxPlatformData, RxPong, \
     TxBITMode, TxClearOdom, TxDriveMotorsRqst, TxLogLevel, TxPing, TxReboot
from dynamic_reconfigure.server import Server
from mowbot_hardware.cfg import MowbotConfig

pilink_log_lvl = 4
rl500_log_lvl = 4
odom_log_lvl = 4
use_pid = False
pid_p = 1.0
pid_i = 1.0
pid_d = 1.0
rl500_bit_mode = False
esp_reboot = False
clear_odom = False

# initialize link
def init_link(port_name):
    try:
        global link
        link = txfer.SerialTransfer(esp_port_name)
    except:
        import traceback
        traceback.print_exc()
        print("Exception in init_link, exiting")
        sys.exit()

# initialize messages
def init_msgs(link):
    global rx_log, rx_odometry, rx_platform_data, rx_pong, \
        tx_bit_mode, tx_clear_odom, tx_drive_motors_rqst, tx_log_level, tx_ping, tx_reboot
    # instantiate rx & tx message handlers
    rx_log = RxLog(link)
    rx_odometry = RxOdometry(link)
    rx_platform_data = RxPlatformData(link)
    rx_pong = RxPong(link)

    tx_bit_mode = TxBITMode(link)
    tx_clear_odom = TxClearOdom(link)
    tx_drive_motors_rqst = TxDriveMotorsRqst(link)
    tx_log_level = TxLogLevel(link)
    tx_ping = TxPing(link)
    tx_reboot = TxReboot(link)

def tick_link():
    start_time = time.time()
    link.tick()
    end_time = time.time()
    #print('tick_link: {} ms'.format((end_time - start_time)*1000))
    if link.status < 0:
        if link.status == txfer.CRC_ERROR:
            print('ERROR: CRC_ERROR')
        elif link.status == txfer.PAYLOAD_ERROR:
            print('ERROR: PAYLOAD_ERROR')
        elif link.status == txfer.STOP_BYTE_ERROR:
            print('ERROR: STOP_BYTE_ERROR')
        else:
            print('ERROR: {}'.format(link.status))

def ping_callback(event):
    tx_ping.post()

def reconfig_callback(config, level):
    global pilink_log_lvl, rl500_log_lvl, odom_log_lvl

    rospy.logdebug('Reconfig log levels: pilink: {pilink_log_lvl}, rl500: {rl500_log_lvl}, odom: {odom_log_lvl}'.format(**config))
    rospy.logdebug('Reconfig PID: use_pid: {use_pid}, prop gain: {pid_p}, integral: {pid_i}, derivative: {pid_d}'.format(**config))
    rospy.logdebug('Reconfig ESP32: bit_mode: {rl500_bit_mode}, reboot: {esp_reboot}, clear_odometry: {clear_odometry}'.format(**config))

    if config['pilink_log_lvl'] != pilink_log_lvl or config['rl500_log_lvl'] != rl500_log_lvl or config['odom_log_lvl'] != odom_log_lvl:
        pilink_log_lvl = config['pilink_log_lvl']
        rl500_log_lvl = config['rl500_log_lvl']
        odom_log_lvl = config['odom_log_lvl']
        tx_log_level.post(pilink_log_lvl, rl500_log_lvl, odom_log_lvl)

    if config['esp_reboot']:
        tx_reboot.post()
        config['esp_reboot'] = False

    if config['rl500_bit_mode']:
        tx_bit_mode.post()
        config['rl500_bit_mode'] = False

    if config['clear_odometry']:
        tx_clear_odom.post()
        config['clear_odometry'] = False

    return config

if __name__ == '__main__':
    rospy.init_node('esp_link', disable_signals=True, log_level=rospy.DEBUG)
    #esp_port_name = rospy.get_param('~esp_port_name', "/dev/ttyUSB0")
    esp_port_name = rospy.get_param('~esp_port_name', "/dev/ttyAMA1")
    rospy.loginfo("esp_link using serial port: {}".format(esp_port_name))

    rospy.Timer(rospy.Duration(5.0), ping_callback)

    try:
        # initialize pySerialTransfer
        init_link(esp_port_name)
        init_msgs(link)
        # configure dynamic reconfigure
        srv = Server(MowbotConfig, reconfig_callback)

        '''
        list of callback functions to be called during tick. The index of the function
        reference within this list must correspond to the packet ID. For instance, if
        you want to call the function rx_pong() when you parse a packet with an ID of 0, you
        would write the callback list with "rx_pong" being in the 0th place of the list:
        '''
        callback_list = [ rx_pong.handle_pong, \
                        rx_odometry.handle_odometry, \
                        rx_log.handle_log, \
                        rx_platform_data.handle_platform_data ]

        link.set_callbacks(callback_list)
        link.open()

        time.sleep(2) # allow some time for the Arduino to completely reset

        while True:
            tick_link()
            tx_bit_mode.send_posted()
            tick_link()
            tx_clear_odom.send_posted()
            tick_link()
            tx_drive_motors_rqst.send_posted()
            tick_link()
            tx_log_level.send_posted()
            tick_link()
            tx_ping.send_posted()
            tick_link()
            tx_reboot.send_posted()
            time.sleep(0.01)

    except KeyboardInterrupt:
        try:
            link.close()
        except:
            pass

    except:
        import traceback
        traceback.print_exc()

        try:
            link.close()
        except:
            pass


from __future__ import annotations

import math
import mmap
import sys
from typing import Any, Dict, Optional, Tuple
from types import TracebackType

import numpy as np
import torch
from datasets.dataset_dict import DatasetDict
from torch.utils.data import Dataset

# `multiprocessing.shared_memory` is only available in Python 3.8+`
if sys.version_info.major == 3 and sys.version_info.minor >= 8:
    from multiprocessing.shared_memory import SharedMemory

class FastHfDataset(Dataset):
    """Fast Hugging Face dataset."""

    def __init__(self, input_ids: torch.Tensor, seq_len: Optional[int] = 1) -> None:
        """Initialize the dataset.

        Args:
            input_ids: Tensor with the inputs (encoded data).
            seq_len: Sequence length.

        """

        super().__init__()

        self.n_input_ids = ((len(input_ids) - 1) // seq_len) * seq_len + 1
        self.seq_len = seq_len

        # `input_ids` should not be sliced since they could be memory mapped
        self.input_ids = input_ids
        self.n_sequences = math.ceil((self.n_input_ids - 1) / self.seq_len)

    def __enter__(self):
        return self

    def __exit__(self, exc_type: type[BaseException], exc_val: BaseException, exc_tb: TracebackType) -> None:
        if isinstance(self.input_ids, np.memmap) and self.input_ids._mmap is not None:
            self.input_ids._mmap.close()

    def __len__(self) -> int:
        return self.n_sequences

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        start_idx = idx * self.seq_len
        seq_len = min(self.seq_len, self.n_input_ids - 1 - start_idx)

        input_ids = torch.as_tensor(self.input_ids[start_idx : (start_idx + seq_len + 1)].astype(np.int64))
        labels = input_ids[1:].clone()

        return input_ids[:-1], labels

class SHMArray(np.ndarray):
    """Numpy array compatible with SharedMemory from `multiprocessing.shared_memory`.

    Reference:
        https://numpy.org/doc/stable/user/basics.subclassing.html#slightly-more-realistic-example-attribute-added-to-existing-array

    """

    def __new__(cls: SHMArray, input_array: np.ndarray, shm: Optional[SharedMemory] = None) -> SHMArray:
        obj = np.asarray(input_array).view(cls)
        obj.shm = shm

        return obj

    def __array_finalize__(self, obj: SHMArray) -> None:
        if obj is None:
            return

        self.shm = getattr(obj, "shm", None)

def process_with_shared_memory(
    dataset_dict: DatasetDict, dtype: np.dtype, num_proc: Optional[int] = 1
) -> Dict[str, SHMArray]:
    """Process the dataset with a shared memory.

    Args:
        dataset_dict: Dataset dictionary.
        dtype: Numpy data type.
        num_proc: Number of processes.

    Returns:
        Dictionary with shared memory-processed datasets.

    """

    def _process_with_shared_memory(example: Dict[str, Any], name, length: int) -> None:
        shared_memory = SharedMemory(name=name)

        shared_memory_array = np.ndarray((length,), dtype=dtype, buffer=shared_memory.buf)
        start_idx = example["offset"] - len(example["input_ids"])
        shared_memory_array[start_idx : example["offset"]] = example["input_ids"]

        shared_memory.close()

    processed_dataset_dict = {}
    for name, ds in dataset_dict.items():
        dataset_dict[name] = ds.add_column("offset", np.cumsum(ds["length"]))
        length = dataset_dict[name][-1]["offset"]

        shared_memory = SharedMemory(create=True, size=length * np.dtype(dtype).itemsize)
        shared_memory_name = shared_memory.name

        dataset_dict[name].map(
            _process_with_shared_memory,
            fn_kwargs={"name": shared_memory_name, "length": length},
            batched=False,
            num_proc=num_proc,
        )

        shared_memory_array = np.ndarray((length,), dtype=dtype, buffer=shared_memory.buf)
        processed_dataset_dict[name] = SHMArray(shared_memory_array, shm=shared_memory)

    return processed_dataset_dict

def process_with_memory_map_files(
    dataset_dict: DatasetDict, cache_dir: str, dtype: np.dtype, num_proc: Optional[int] = 1
) -> Dict[str, np.ndarray]:
    """Process the dataset with memory map files.

    Args:
        dataset_dict: Dataset dictionary.
        cache_dir: Cache directory.
        dtype: Numpy data type.
        num_proc: Number of processes.

    Returns:
        Dictionary with memory map file-processed datasets.

    """

    def _process_with_memory_map_files(example: Dict[str, Any], file_path: str) -> None:
        with open(file_path, "r+b") as f:
            memory_map = mmap.mmap(f.fileno(), 0)

            start_idx = example["offset"] - len(example["input_ids"])
            length = len(example["input_ids"])

            memory_map_array = np.ndarray(
                (length,), dtype=dtype, buffer=memory_map, offset=np.dtype(dtype).itemsize * start_idx
            )
            memory_map_array[:] = example["input_ids"]

            memory_map.flush()

    processed_dataset_dict = {}
    for split, dataset in dataset_dict.items():
        dataset_dict[split] = dataset.add_column("offset", np.cumsum(dataset["length"]))
        length = dataset_dict[split][-1]["offset"]

        file_path = cache_dir / f"{split}.bin"
        with open(file_path.as_posix(), "wb") as f:
            f.truncate(length * np.dtype(dtype).itemsize)

        dataset_dict[split].map(
            _process_with_memory_map_files,
            fn_kwargs={"file_path": file_path},
            batched=False,
            num_proc=num_proc,
        )

        processed_dataset_dict[split] = np.memmap(file_path, dtype=dtype, mode="r", shape=(length,))

    return processed_dataset_dict

def xor(p: Any, q: Any) -> bool:
    """Implements the logical XOR operator.

    Args:
        p: Any instance that may act as `True` or `False`.
        q: Any instance that may act as `True` or `False`.

    Returns:
        Logical value.

    """

    return (p and not q) or (not p and q)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Track Zigbee devices and networks by passing sniffed packeges. Device addresses
and sequence numbers are extracted and are used to modify arbitrary command
packets in a way to spoof devices.
"""

import argparse
import struct
import sys
import time

from killerbee import *
from killerbee.scapy_extensions import *
from scapy.all import *
from scapy.layers.dot15d4 import *

from z3sec import touchlink_crypt

class ZBNetwork():
    """A Zigbee network representation.

    Contains network addresses, the network key (if known) and a list of member
    devices. A ZBNetwork object ist automatically generated and filled when a
    packet of an unknown network is passed via the handle_packets() function.
    However, you can create a ZBNetwork object manually in advance when the
    network key for the network is known in advance and it is unlikely that a
    touchlink keytransport will happen in the future.

    Parameters
    ----------
    pan_id : int
        Short PAN identifier of the network.
    ext_pan_id : int
        Long/IEEE PAN identifier of the network.
    network_key : byte string
        The current network key of the network.

    Attributes
    ----------
    pan_id : int
        Short PAN identifier of the network (2 octets).
    ext_pan_id : int
        Long/IEEE PAN identifier of the network (8 octets).
    network_key : byte string
        The current network key of the network (16 octets)
    devices : list of ZBDevice
        Devices that are associated to the network (member devices).
    """
    def __init__(self, pan_id=None, ext_pan_id=None, network_key=None):
        self.pan_id = pan_id
        self.ext_pan_id = ext_pan_id
        self.network_key = network_key
        self.devices = []

    def show(self):
        """Print the properties of the network in form of a table."""
        print("pan_id:      {}".format(_addr_repr(self.pan_id)))
        print("pan_id_ext:  {}".format(_addr_repr(self.ext_pan_id, ext=True)))
        if self.network_key is not None:
            print("network_key: {}".format(self.network_key.encode('hex')))
        else:
            print("network_key: {}".format(self.network_key))
        print("# Devices:   {}".format(len(self.devices)))
        header = ["#", "addr", "addr_ext", "mac_sqn", "nwk_sqn", "sec_fc", "aps_fc",
                "zcl_sqn", "zdp_sqn"]
        data = []
        for n, dev in enumerate(self.devices):
            data.append([
                n,
                _addr_repr(dev.addr),
                _addr_repr(dev.ext_addr, ext=True),
                dev.mac_sqn,
                dev.nwk_sqn,
                dev.sec_fc,
                dev.aps_fc,
                dev.zcl_sqn,
                dev.zdp_sqn
            ])
        _print_table(header, data)

    def get_device(self, addr, ext=False):
        """Return a member device of the network, or associate a new
        address/device to the network.

        Search and for member devices of the network. If addr is a short
        address and it is not associated to the network yet (not found during
        search), a new device object is created and returned. This new device
        is then already associated to the network.

        Parameters
        ----------
        addr : int
            Either the short or the extended/IEEE address of the device,
            depending on the passed flag `ext`.
        ext : bool
            Specifies if the passed `addr` is assumed to be a short (false) or
            extended/IEEE address (true).

        Returns
        -------
        ZBDevice
            The member device of network (might be newly created) or, if an
            extended address is passed and this address is not asscoiated to
            the network None is returned.
        """
        # check if device is already member of this network
        for dev in self.devices:
            if ext is False:
                if dev.addr == addr:
                    return dev
            else:
                if dev.ext_addr == addr:
                    return dev
        # device is unknown -> creating new
        #   do not create device with ext_addr in order to avoid
        #   doublicate devices in network device list
        if ext is False:
            new_device = ZBDevice()
            new_device.addr = addr
            self.devices.append(new_device)
            return new_device
        return None

class ZBDevice():
    """A Zigbee device representation.

    Contains device addresses, various sequence numbers, the associated network
    and a list of devices to which this device sent packets (destinations).

    Attributes
    ----------
    addr : int
        Short address.
    ext_addr : int
        Extended/IEEE address.
    mac_sqn : int
        802.14.4 mac sequence number (1 octet).
    mac_sqn_offset : int
        Value to increase mac_sqn when spoofing packets.
    nwk_sqn : int
        Zigbee network layer sequence number (1 octet).
    nwk_sqn_offset : int
        Value to increase nwk_sqn when spoofing packets.
    sec_fc : int
        Zigbee security header framecounter (4 octets)
    sec_fc_offset
        Value to increase sec_fc when spoofing packets.
    aps_fc : int
        Zigbee AppData payload counter (1 octet)
    aps_fc_offset
        Value to increase aps_fc when spoofing packets.
    zcl_sqn : int
        Zigbee Cluster Library transaction sequence number (1 octet)
    zcl_sqn_offset : int
        Value to increase zcl_sqn when spoofing packets.
    zdp_sqn : int
        Zigbee Device Profile transaction sequence (1 octet)
    zdp_sqn_offset : int
        Value to increase zdp_sqn when spoofing packets.
    network : ZBNetwork
        The network associated with this device.
    destinations : list of ZBDevice
        The devices to which this device sent packets in the past.
    """
    def __init__(self):
        self.addr = None
        self.ext_addr = None
        # maybe macDSN / macBSN for beacons injection
        self.mac_sqn = None  # Dot15d4.seqnum (1 octet)
        self.mac_sqn_offset = 1
        # self.sec_fc = None  # Dot15d4AuxSecurityHeader.sec_framecounter // no
        # zb_nwk_seqnumber
        self.nwk_sqn = None  # ZigbeeNWK.seqnum (1 octet)
        self.nwk_sqn_offset = 1
        # frame_counter
        self.sec_fc = None  # ZigbeeSecurityHeader.fc (4 octets)
        self.sec_fc_offset = 1
        # zb_zadp_counter
        self.aps_fc = None  # ZigbeeAppDataPayload.counter (1 octet)
        self.aps_fc_offset = 1
        self.zcl_sqn = None  # ZigbeeClusterLibrary.transaction_sequence (1 octet)
        self.zcl_sqn_offset = 1
        self.zdp_sqn = None  # ZigbeeDeviceProfile.transaction_sequence (1 octet)
        self.zdp_sqn_offset = 1
        self.network = None
        self.destinations = []

    def show(self):
        """Print the properties of the device in form of a table."""
        print("> Short Address: {}".format(_addr_repr(self.addr)))
        print("ext addr: {}".format(_addr_repr(self.ext_addr)))
        print("mac sqn:  {}".format(self.mac_sqn))
        print("nwk sqn:  {}".format(self.nwk_sqn))
        print("sec fc:   {}".format(self.sec_fc))
        print("aps fc:   {}".format(self.aps_fc))
        print("zcl sqn:  {}".format(self.zcl_sqn))
        print("zdp sqn:  {}".format(self.zdp_sqn))

    def get_destination(self, addr, ext=False):
        """Get the device object, which is a destination of the device.

        If the address was not previously associated as a destination, then it
        is set as a destination.

        Attributes
        ----------
        addr : int
            The address of the destination device.
        ext : bool
            Specifies if the `addr` attribute is a short or extended address.

        Returns
        -------
        ZBDevice
        """
        # check if dev is already registered as destination
        for dev in self.destinations:
            if ext is False:
                if dev.addr == addr:
                    return dev
            else:
                if dev.ext_addr == addr:
                    return dev
        # get device from network and add to own destinations
        dev = self.network.get_device(addr, ext=ext)
        if dev is not None:
            self.destinations.append(dev)
        return dev

    def update_mac_sqn(self, pkt):
        """Updates the MAC properties with those found in the packet.

        Extracts the MAC seqnum and sec_fc from the packet. Invoke this
        function only the device which is the MAC source device of the
        packet.

        Attributes
        ----------
        pkt : scapy_pkt
            The packet from which the properties are extracted.
        """
        # If ACK -> no update
        if Dot15d4Ack in pkt:
            return
        # update SQNs:
        if Dot15d4 in pkt:
            self.mac_sqn = pkt[Dot15d4].seqnum
        if Dot15d4FCS in pkt:
            self.mac_sqn = pkt[Dot15d4FCS].seqnum
        if ZigbeeSecurityHeader in pkt:
            self.sec_fc = pkt[ZigbeeSecurityHeader].fc

    def update_nwk_sqn(self, pkt):
        """Extracts the ZigBee seqnums / counters / framecounters. Invoke this
        function only on the NWK source device."""
        # If ACK -> no update
        if ZigbeeAppDataPayload in pkt \
            and pkt[ZigbeeAppDataPayload].aps_frametype == 2:
                return
        # update SQNs:
        if ZigbeeNWK in pkt:
            self.nwk_sqn = pkt[ZigbeeNWK].seqnum
        if ZigbeeAppDataPayload in pkt:
            self.aps_fc = pkt[ZigbeeAppDataPayload].counter
        if ZigbeeClusterLibrary in pkt \
                and pkt[ZigbeeClusterLibrary].direction == 0:  # client2server
            self.zcl_sqn = pkt[ZigbeeClusterLibrary].transaction_sequence
        # Extract ZDP counter (only from client2server commands):
        if is_zdp_client2server(pkt):
            self.zdp_sqn = pkt.transaction_sequence

class Observer():
    """Analyse packets and track all networks and devices.

    All packets that are submitted to this object are decrypted (if possible)
    and information of the lower network stack layers are examined.
    Thereby networks, devices, addresses and sequence numbers are extracted.
    A list of all observed Zigbee networks and Zigbee devices is generated.
    These information can be used to impersonate an observed device and spoof
    other devices in the same network.

    Notes
    -----
    Injecting might not work if the network key of a network is not passed
    during initialization or sniffed by observing a touchlink key transport.
    Upper network layer sequence numbers cannot be extracted.

    Attributes
    ----------
    networks : list of ZBNetwork
        A passed network should include a network key, if it is not likely that
        a touchlink key transport is sniffed later.

    """
    def __init__(self, networks=[]):
        self.known_networks = networks
        self.touchlink_scan_responses = []
    def extract_network(self, pkt):
        """Get the network in which the packet was sent.

        If the network is not yet known, a new ZBNetwork object is created and
        returned. The new network is added to the list of known networks.

        Attributes
        ----------
        pkt : scapy_pkt

        Returns
        -------
        ZBNetwork
        """
        # Data or Cmd frame:
        if pkt.haslayer(Dot15d4Data) or pkt.haslayer(Dot15d4Cmd):
            if (pkt.fcf_panidcompress == 0) \
                    and (pkt.fcf_srcaddrmode != 0):
                pan_id = pkt.src_panid
            else:
                pan_id = pkt.dest_panid
        # Beacon frame:
        elif pkt.haslayer(Dot15d4Beacon):
            pan_id = pkt.src_panid
        # No PAN ID found -> no network (should not happen):
        else:
            return None
        if pan_id == 0xffff:  # fake network
            return None

        # check if network is already known
        for network in self.known_networks:
            if network.pan_id == pan_id:
                return network
        # network not yet known -> add new network
        new_network = ZBNetwork()
        new_network.pan_id = pan_id
        self.known_networks.append(new_network)
        return new_network

    def extract_mac_src(self, pkt, network):
        """Get the MAC source device of the packet.

        MAC source device is the device which transmitted the packet most
        recently, even if the packet is relayed and originally sent by an other
        device.

        Attributes
        ----------
        pkt : scapy_pkt
        network : ZBNetwork
            The network to which the MAC source device belongs.

        Returns
        -------
        ZBDevice
        """
        fcf = None
        if Dot15d4 in pkt:
            fcf = pkt[Dot15d4]
        if Dot15d4FCS in pkt:
            fcf = pkt[Dot15d4FCS]
        if fcf is None:
            return None
        # extract source device
        source = None
        if Dot15d4Data in pkt:
            if fcf.fcf_srcaddrmode == 2:  # short
                source = network.get_device(pkt[Dot15d4Data].src_addr)
            if fcf.fcf_srcaddrmode == 3:  # long
                source = network.get_device(pkt[Dot15d4Data].src_addr, ext=True)
        # TODO: extract from Beacon / Cmd frame
        if source is None:
            return None
        # extract destination device
        if Dot15d4Data in pkt:
            if fcf.fcf_destaddrmode == 2:  # short
                dst_addr = pkt[Dot15d4Data].dest_addr
                if dst_addr in range(0xfff8, 0xffff + 1):  # Broadcast / Reserved
                    return source
                destination = source.get_destination(dst_addr)
            if fcf.fcf_destaddrmode == 3:  # long
                destination = source.get_destination(pkt[Dot15d4Data].dest_addr,
                        ext=True)
        return source

    def extract_nwk_src(self, pkt, network):
        """Get the network layer source device of a packet.

        Network layer source device is the device which initiated the
        transmission of a packet. The packet might be send by an other device,
        if the it is relayed.

        Attributes
        ----------
        pkt : scapy_pkt
        network : ZBNetwork
            The network to which the NWK source device belongs.

        Returns
        -------
        ZBDevice
        """
        if ZigbeeNWK in pkt:
            nwk = pkt[ZigbeeNWK]
            # Extract source addresses
            source = network.get_device(nwk.source)
            source.network = network
            if nwk.flags & 16:  # extended source address
                source.ext_addr = nwk.ext_src

            # Extract destination addresses
            dst_addr = nwk.destination
            if dst_addr in range(0xfff8, 0xffff + 1):  # Broadcast / Reserved
                return source
            destination = source.get_destination(nwk.destination)
            if nwk.flags & 8:  # extended destination address
                destination.ext_addr = nwk.ext_dst

            return source
        return None

    def _none_to_zero(self, value):
        if value is None:
            return 0
        else:
            return value

    def make_injectable(self, pkt, mac_src, nwk_src, mac_dst, nwk_dst):
        """Modify the given packet for injecting into a network.

        The packet is modifyed in such a way that is looks like it is
        originated by the specified devices. The addresses and current sequence
        numbers of the passed device objects are inserted into the packet at
        the appropriate places where it is designated in the packet. If the
        network key of the network is known, the packet is additionally
        encrypted. Other modification are not performed.

        Notes
        -----
        In order to impersonate an other device in the network, it is advised
        to send the modified packet immediately. Otherwise it can happen that
        the packet is rejected due to outdated sequence numbers.
        If the network key of the of the network is not known, the packet
        cannot be encrypted. An unencrypted packet will be rejected.

        Attributes
        ----------
        pkt : scapy_pkt
            The unencrypted packet to modify. All address fields and sequence
            numbers are overwritten, in order to fit values of the passed
            device objects.
        mac_src : ZBDevice
            The device which is set as the MAC layer source.
        nwk_src : ZBDevice
            The device which is set as the network layer source.
        mac_dst : ZBDevice
            The device which is set as the MAC layer destination.
        nwk_dst : ZBDevice
            The device which is set as the network layer destination.

        Returns
        -------
        scapy_pkt
            The modified packet; ready for injection.
        """
        # do not modify pkt, so the user can send/update it again later
        pkt = copy.deepcopy(pkt)

        # adjust SQNs:
        if Dot15d4 in pkt:
            pkt[Dot15d4].seqnum = \
                (self._none_to_zero(mac_src.mac_sqn) \
                + mac_src.mac_sqn_offset) % 2**8  # (1 octet)
        if Dot15d4FCS in pkt:
            pkt[Dot15d4FCS].seqnum = \
                (self._none_to_zero(mac_src.mac_sqn) \
                + mac_src.mac_sqn_offset) % 2**8  # (1 octet)
        if ZigbeeNWK in pkt:
            pkt[ZigbeeNWK].seqnum = \
                (self._none_to_zero(nwk_src.nwk_sqn) \
                + nwk_src.nwk_sqn_offset) % 2**8  # (1 octet)
        if ZigbeeSecurityHeader in pkt:
            pkt[ZigbeeSecurityHeader].fc = \
                (self._none_to_zero(mac_src.sec_fc) \
                + mac_src.sec_fc_offset) % 2**32  # (4 octets)
        if ZigbeeAppDataPayload in pkt:
            pkt[ZigbeeAppDataPayload].counter = \
                (self._none_to_zero(nwk_src.aps_fc) \
                + nwk_src.aps_fc_offset) % 2**8  # (1 octet)
        if ZigbeeClusterLibrary in pkt:
            if pkt[ZigbeeClusterLibrary].direction == 0:  # client2server
                pkt[ZigbeeClusterLibrary].transaction_sequence = \
                    (self._none_to_zero(nwk_src.zcl_sqn) \
                    + nwk_src.zcl_sqn_offset) % 2**8  # (1 octet)
            else:  # server2client
                pkt[ZigbeeClusterLibrary].transaction_sequence = \
                    (self._none_to_zero(nwk_dst.zcl_sqn) \
                    + nwk_dst.zcl_sqn_offset) % 2**8  # (1 octet)
        if is_zdp_client2server(pkt):
            pkt.transaction_sequence = \
                (self._none_to_zero(nwk_src.zdp_sqn) \
                + nwk_src.zdp_sqn_offset) % 2**8  # (1 octet)
        if is_zdp_server2client(pkt):
            pkt.transaction_sequence = \
                (self._none_to_zero(nwk_dst.zdp_sqn) \
                + nwk_dst.zdp_sqn_offset) % 2**8  # (1 octet)

        # TODO: adjust SQNs for other application profiles...?

        # set source and destination addresses in MAC layer:
        if (Dot15d4 in pkt or Dot15d4FCS in pkt) and Dot15d4Data in pkt:
            pkt.dest_panid = mac_dst.network.pan_id
            if pkt.fcf_srcaddrmode == 2:  # short
                pkt.src_addr = mac_src.addr
            elif pkt.fcf_srcaddrmode == 3:  # long
                pkt.src_addr = mac_src.ext_addr
            if pkt.fcf_destaddrmode == 2:  # short
                pkt.dest_addr = mac_dst.addr
            if pkt.fcf_destaddrmode == 3:  # long
                pkt.dest_addr = mac_dst.ext_addr
        # set source and destination addresses in NWK layer:
        if ZigbeeNWK in pkt:
            pkt[ZigbeeNWK].source = nwk_src.addr
            pkt[ZigbeeNWK].destination = nwk_dst.addr
            if pkt[ZigbeeNWK].flags & 16:
                pkt[ZigbeeNWK].ext_src = nwk_src.ext_addr
            if pkt[ZigbeeNWK].flags & 8:
                pkt[ZigbeeNWK].ext_dst = nwk_dst.ext_addr
        # set extended source in ZigbeeSecurityHeader:
        if ZigbeeSecurityHeader in pkt:
            if pkt[ZigbeeSecurityHeader].extended_nonce:
                pkt[ZigbeeSecurityHeader].source = nwk_src.ext_addr

        # update SQNs:
        mac_src.update_mac_sqn(pkt)
        nwk_src.update_nwk_sqn(pkt)

        # encrypt packet:
        if ZigbeeSecurityHeader in pkt:
            if pkt[ZigbeeSecurityHeader].key_type == 1:  # network key
                if nwk_src.network.network_key is not None:
                    pkt = self.encrypt(pkt, nwk_src.network.network_key)
                else:
                    print("WARNING: The network key of the source device is "
                            "not known. The packet could not be encrypted.")
        return pkt

    def decrypt(self, pkt, network_key):
        """Decrypt a packet with a network key.

        Attributes
        ----------
        pkt : scapy_pkt
        network_key : byte string

        Returns
        -------
        Decrypted packet.
        """
        payload = kbdecrypt(pkt, network_key, 0)
        pkt.data = ""  # remove encrypted payload
        pkt.mic = ""  # remove mic
        return pkt / payload

    def encrypt(self, pkt, network_key):
        """Encrypt a packet with a network key.

        Attributes
        ----------
        pkt : scapy_pkt
        network_key : byte string

        Returns
        -------
        Encrypted packet.
        """
        payload = pkt[ZigbeeSecurityHeader].payload  # deepcopy instead?
        pkt[ZigbeeSecurityHeader].remove_payload()
        return kbencrypt(pkt, payload, network_key, 0)

    def analyse(self, pkt):
        """Extract information from a packet.

        Extract the network, source devices, destination devices and sequence
        numbers and store the information internally. Even touchlink
        keystransports are analysed in order to obtain the network key of a
        network.

        Attributes
        ----------
        pkt ; scapy_pkt
            The packet to analyse.
        """
        pkt = copy.deepcopy(pkt)

        # Get network; add new ones, if we do not know it yet
        network = self.extract_network(pkt)
        if network is None:
            return pkt

        # extract network_key during touchlink commissioning
        if touchlink_crypt.is_scan_response(pkt):
            self.touchlink_scan_responses.append(pkt)
            return
        if touchlink_crypt.is_keytransport(pkt):
            response_id = touchlink_crypt.get_response_id(pkt,
                    self.touchlink_scan_responses)
            if response_id is None:
                return pkt
            network_key = touchlink_crypt.extract_network_key(pkt, response_id)
            network.network_key = network_key

        # TODO: extract network key from other key transport frames (if master
        # key known).

        # Decrypt packets, if we know the network key
        if network.network_key is not None \
                and pkt.haslayer(ZigbeeSecurityHeader) \
                and pkt.key_type == 1:
            pkt = self.decrypt(pkt, network.network_key)

        # get NWK source device and update SQNs:
        nwk_source = self.extract_nwk_src(pkt, network)
        if nwk_source is not None:
            nwk_source.update_nwk_sqn(pkt)

        # get MAC source device and update SQNs:
        mac_src = self.extract_mac_src(pkt, network)
        if mac_src is not None:
            mac_src.update_mac_sqn(pkt)

        return pkt

def is_zdp_client2server(pkt):
    """Check if the Zigbee device profile command inside the packet is to sent
    from a client to a server.

    Returns
    -------
    bool
    """
    if ZigbeeAppDataPayload in pkt and pkt[ZigbeeAppDataPayload].profile == 0:
        if pkt[ZigbeeAppDataPayload].cluster in \
                range(0x0000, 0x00ff):  # this might be wrong for future clusters
                # range(0x0000, 0x0006 +1) + range(0x0010, 0x001e +1) + \
                # range(0x0020, 0x002a +1) + range(0x0030, 0x0038 +1):
            return True
    return False

def is_zdp_server2client(pkt):
    """Check if the Zigbee device profile command inside the packet is to sent
    from a server to a client.

    Returns
    -------
    bool
    """
    if ZigbeeAppDataPayload in pkt and pkt[ZigbeeAppDataPayload].profile == 0:
        if pkt[ZigbeeAppDataPayload].cluster in \
                range(0x8000, 0x80ff):  # this might be wrong for future clusters
            return True
    return False

def _addr_repr(addr, ext=False):
    """Convert an address value (int) to a hex representation

    Parameters
    ----------
    addr : int
        The address as an integer.
    ext : bool
        Specifies if the passed `addr` is assumed to be a short (false) or
        extended/IEEE address (true).

    Returns
    -------
    string, None
        Human readable hex representation. For short addresses the hex
        representation is used and for extended addresses the bytes (hex) are
        outputted, seperated by ":". If addr is None, then None is returned.
    """
    if addr is None:
        return None
    if ext:
        addr_ext = struct.pack(">Q", addr)
        return ":".join(["{0:>02x}".format(ord(byte)) for byte in addr_ext])
    else:
        return hex(addr)

def _print_table(header, data):
    """Simple table printer.

    Parameters
    ----------
    header : list of str
        Column names.
    data : list of list of str
        Each inner list contains the values for one row.
    """
    # calculate width of columns
    widths = []
    for i, h in enumerate(header):
        max_len = len(str(h))
        for d in data:
            if len(str(d[i])) > max_len:
                max_len = len(str(d[i]))
        widths.append(max_len)
    # generate format string:
    row_format = ""
    for w in widths:
        row_format += "{:<" + str(w+1) + "}|"
    print(row_format.format(*header))
    print("=" * (sum(widths) + 2*len(header)))  # seperator
    for d in data:
        print(row_format.format(*d))