#!/usr/bin/env python3
# necessary to register envs
import math
import time

from ament_index_python import get_package_share_directory
from bitbots_utils.utils import get_parameters_from_ros_yaml
from humanoid_league_msgs.msg import RobotControlState
from parallel_parameter_search.utils import load_yaml_to_param
from rclpy.duration import Duration

import deep_quintic
import argparse
import glob
import os
from collections import defaultdict
from typing import Tuple, Dict, Any, Optional

import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.action import ActionClient
import stable_baselines3
import yaml
from bitbots_moveit_bindings import get_position_fk
from bitbots_msgs.msg import JointCommand, FootPressure
from biped_interfaces.msg import Phase
from geometry_msgs.msg import PoseStamped, Twist
from sb3_contrib import QRDQN, TQC
from sensor_msgs.msg import Imu, JointState
from stable_baselines3 import PPO, DQN, DDPG, A2C, SAC, TD3
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.vec_env import VecEnv, DummyVecEnv, VecNormalize, VecFrameStack
from transforms3d.affines import compose, decompose
from transforms3d.euler import quat2euler, mat2euler, euler2mat, euler2quat
from transforms3d.quaternions import quat2mat
from urdf_parser_py.urdf import URDF

from deep_quintic import WolfgangWalkEnv
from deep_quintic.robot import Robot
from deep_quintic.state import CartesianState, BaseState
from deep_quintic.utils import Rot, compute_ik
from deep_quintic.reward import CartesianActionReward

ALGOS = {
    "a2c": A2C,
    "ddpg": DDPG,
    "dqn": DQN,
    "ppo": PPO,
    "sac": SAC,
    "td3": TD3,
    # SB3 Contrib,
    "qrdqn": QRDQN,
    "tqc": TQC,
}

class DummyPressureSensor:
    pressure = 0

    def get_force(self):
        return None, self.pressure

    def get_value(self, type):
        if type == "raw":
            return self.pressure
        elif type == "filtered":
            return self.pressure
        elif type == "binary":
            print("not implemented")
            exit()

class ExecuteEnv(WolfgangWalkEnv):
    def __init__(self, simulator_type="pybullet", reward_function="CartesianActionVelReward", step_freq=30,
                 ros_debug=False,
                 gui=False, trajectory_file=None, ep_length_in_s=10, use_engine=False,
                 cartesian_state=True, cartesian_action=True, relative=False, use_state_buffer=False,
                 state_type="full", cyclic_phase=True, rot_type="rpy", filter_actions=False, terrain_height=0,
                 phase_in_state=True, foot_sensors_type="", leg_vel_in_state=False, use_rt_in_state=False,
                 randomize=False, use_complementary_filter=True, random_head_movement=True, adaptive_phase=False,
                 use_gyro=True, use_imu_orientation=True, node: Node = None):
        super().__init__(simulator_type=simulator_type + "_off", reward_function=reward_function, step_freq=step_freq,
                         ros_debug=True, gui=gui,
                         trajectory_file=trajectory_file, state_type=state_type, ep_length_in_s=ep_length_in_s,
                         use_engine=False, cartesian_state=cartesian_state,
                         cartesian_action=cartesian_action, relative=relative,
                         use_state_buffer=use_state_buffer, cyclic_phase=cyclic_phase, rot_type=rot_type,
                         use_rt_in_state=use_rt_in_state, filter_actions=filter_actions,
                         terrain_height=terrain_height, foot_sensors_type=foot_sensors_type,
                         phase_in_state=phase_in_state, randomize=randomize, leg_vel_in_state=leg_vel_in_state,
                         use_complementary_filter=False, random_head_movement=False, adaptive_phase=adaptive_phase,
                         use_gyro=use_gyro, use_imu_orientation=use_imu_orientation, node=node)
        # use dummy pressure sensors since we are not connected to a simulation
        self.robot.pressure_sensors = defaultdict(lambda: DummyPressureSensor())
        self.node = node

        # additional class variables we need since we are not running the walking at the same time
        self.refbot.phase = 0
        self.last_time = None
        walk_parameters = get_parameters_from_ros_yaml("walking",
                                                       f"{get_package_share_directory('bitbots_quintic_walk')}"
                                                       f"/config/deep_quintic_{simulator_type}.yaml",
                                                       use_wildcard=True)
        for parameter in walk_parameters:
            if parameter.name == "engine.freq":
                self.freq = parameter.value.double_value
        # todo we dont initialize action filter correctly. history is empty

        # initialize ROS stuff
        self.got_foot_pressure = False
        self.current_joint_positions = None
        self.current_joint_velocities = None
        self.imu_rpy = None
        self.ang_vel = None
        self.last_joint_state_time = 0
        self.current_command_speed = (0, 0, 0)
        self.got_stop_command = True
        self.performing_stop_step = False
        self.is_stopped = True
        self.robot_state = RobotControlState.CONTROLLABLE

        self.joint_publisher = self.node.create_publisher(JointCommand, 'walking_motor_goals', 1)
        self.support_publisher = self.node.create_publisher(Phase, 'walk_support_state', 1)
        self.current_command_speed_sub = self.node.create_subscription(Twist, 'cmd_vel', self.current_command_speed_cb,
                                                                       1)
        self.imu_sub = self.node.create_subscription(Imu, '/imu/data', self.imu_cb, 1)
        self.joint_state_sub = self.node.create_subscription(JointState, '/joint_states', self.joint_state_cb, 1)
        self.robot_state_sub = self.node.create_subscription(RobotControlState, 'robot_state', self.robot_state_cb, 1)
        if foot_sensors_type == "filtered":
            self.left_pressure_sub = self.node.create_subscription(FootPressure, 'foot_pressure_left/filtered',
                                                                   lambda msg: self.foot_pressure_cb(msg, True), 1)
            self.right_pressure_sub = self.node.create_subscription(FootPressure, 'foot_pressure_right/filtered',
                                                                    lambda msg: self.foot_pressure_cb(msg, False), 1)
        elif foot_sensors_type in ["raw", "binary"]:
            self.left_pressure_sub = self.node.create_subscription(FootPressure, 'foot_pressure_left/raw',
                                                                   lambda msg: self.foot_pressure_cb(msg, True), 1)
            self.right_pressure_sub = self.node.create_subscription(FootPressure, 'foot_pressure_right/raw',
                                                                    lambda msg: self.foot_pressure_cb(msg, False), 1)
        elif foot_sensors_type == "":
            self.got_foot_pressure = True
        else:
            print(f"Problem: use foot sensors is {foot_sensors_type}")
            exit()

        # urdf = URDF.from_xml_string(self.node.get_parameter('robot_description'))
        # self.joint_limits = {
        #    joint.name: (joint.limit.lower, joint.limit.upper) for joint in urdf.joints if joint.limit
        # }

        # see that we got all data from the various subscribers
        for i in range(5):
            rclpy.spin_once(self.node)
        while self.current_joint_positions is None or self.ang_vel is None or not self.got_foot_pressure:
            self.node.get_logger().info("Waiting for data from subscribers", throttle_duration_sec=10)
            time.sleep(0.01)
            rclpy.spin_once(self.node)

    def apply_action(self, action):
        # only do something if we are not stopped
        if not self.is_stopped:
            action = self.action_filter.filter(action)
            msg = JointCommand()
            msg.header.stamp = self.node.get_clock().now().to_msg()
            msg.joint_names = self.robot.used_joint_names
            msg.velocities = [-1.0] * len(self.robot.used_joint_names)
            msg.accelerations = [-1.0] * len(self.robot.used_joint_names)
            msg.max_currents = [-1.0] * len(self.robot.used_joint_names)
            if self.cartesian_action:
                left_foot_pos, left_foot_rpy, right_foot_pos, right_foot_rpy = \
                    self.robot.scale_action_to_pose(action, Rot.RPY)
                left_foot_quat = euler2quat(*left_foot_rpy)
                right_foot_quat = euler2quat(*right_foot_rpy)
                ik_result, success = compute_ik(left_foot_pos, left_foot_quat, right_foot_pos, right_foot_quat,
                                                self.robot.used_joint_names, self.robot.joint_indexes,
                                                collision=False, approximate=True)
                msg.positions = list(ik_result)
            else:
                msg.positions = self.robot.joints_scaled_to_radiant(action)
            self.joint_publisher.publish(msg)

            # support state is necessary for odometry
            msg = Phase()
            msg.header.stamp = self.node.get_clock().now().to_msg()
            if self.refbot.phase > 0.5:
                msg.phase = Phase.LEFT_STANCE
            else:
                msg.phase = Phase.RIGHT_STANCE
            # todo we are not detecting double support state
            # msg.state = SupportState.DOUBLE
            self.support_publisher.publish(msg)

    def compute_observation(self):
        # save current state as previous state
        self.robot.step()

        # compute current state
        self.robot.imu_ang_vel = self.ang_vel
        self.robot.imu_rpy = self.imu_rpy
        self.robot.joint_positions = self.current_joint_positions
        if not self.cartesian_state:
            self.robot.joint_positions = self.current_joint_positions
            self.robot.joint_velocities = self.current_joint_velocities
        if self.cartesian_state:
            self.robot.solve_fk(force=True)
            if self.robot.previous_left_foot_pos is None:
                # need to have two states updates first so that the previous positions are set for velocity computation
                return None
        # calc velocities
        time = self.node.get_clock().now().nanoseconds / 1e9
        if self.last_time is None:
            self.last_time = time - 0.001
        time_diff = time - self.last_time
        if time_diff == 0:
            time_diff = 1 / self.step_freq
        if self.cartesian_state:
            self.robot.compute_velocities(time_diff)

        observation = self.state.get_state_array(scaled=True)
        if len(self.state_buffer) > 0:
            self.state_buffer.pop(0)
        self.state_buffer.append(observation)

        return observation

    def progress_phase(self, time, action):
        if self.adaptive_phase:
            # progress phase based on the policies action
            time_of_single_step = 0.5 / self.freq
            dt = (action[-1] + 1) / 2 * time_of_single_step
        else:
            if self.last_time is None:
                # handle first iteration
                dt = 0.001
            else:
                dt = time - self.last_time
                if dt <= 0:
                    # handle edge case or sim time being resettled
                    dt = 0.001
            self.last_time = time

        # Check for too long dt
        if dt > 0.5 / self.freq:
            self.node.get_logger().error(f"DeepQuintic error too long dt phase={self.refbot.phase} dt={dt}")
            return

        # Update the phase
        last_phase = self.refbot.phase
        self.refbot.phase += dt * self.freq

        if (last_phase < 0.5 and self.refbot.phase >= 0.5) or (last_phase >= 0.5 and self.refbot.phase < 0.5):
            # if we need to stop, this is our chance. we need to do one step with zero vel.
            # this is basically a tiny implicit state machine
            if self.got_stop_command:
                if self.performing_stop_step:
                    self.is_stopped = True
                else:
                    self.performing_stop_step = True
            else:
                self.is_stopped = False
                self.performing_stop_step = False

        # reset if step complete
        if self.refbot.phase > 1.0:
            self.refbot.phase -= 1

    def current_command_speed_cb(self, msg: Twist):
        self.current_command_speed = (msg.linear.x, msg.linear.y, msg.angular.z)
        self.got_stop_command = msg.angular.x < 0

    def imu_cb(self, msg: Imu):
        self.imu_rpy = [*quat2euler((msg.orientation.w, msg.orientation.x, msg.orientation.y, msg.orientation.z))]
        self.ang_vel = [msg.angular_velocity.x, msg.angular_velocity.y, msg.angular_velocity.z]
        # self.imu_rpy = [0,0,0]
        # self.ang_vel = [0,0,0]

    def foot_pressure_cb(self, msg: FootPressure, is_left: bool):
        print("foot_pressure_cb")
        self.got_foot_pressure = True
        # LLB, LLF, LRF, LRB, RLB, RLF, RRF, RRB
        pressures = [msg.left_back, msg.left_front, msg.right_front, msg.right_back]
        prefix = 'L' if is_left else 'R'
        self.robot.pressure_sensors[prefix + 'LB'].pressure = pressures[0]
        self.robot.pressure_sensors[prefix + 'LF'].pressure = pressures[1]
        self.robot.pressure_sensors[prefix + 'RF'].pressure = pressures[2]
        self.robot.pressure_sensors[prefix + 'RB'].pressure = pressures[3]

    def joint_state_cb(self, msg: JointState):
        # sort joint positions in correct order
        self.current_joint_positions = np.empty(self.robot.num_used_joints)
        self.current_joint_velocities = np.empty(self.robot.num_used_joints)
        for i in range(len(msg.name)):
            joint_name = msg.name[i]
            if joint_name in self.robot.used_joint_names:
                index = self.robot.joint_indexes[joint_name]
                self.current_joint_positions[index] = msg.position[i]
                self.current_joint_velocities[index] = msg.velocity[i]

    def robot_state_cb(self, msg: RobotControlState):
        self.robot_state = msg.state

    def run_node(self, model, venv):
        rate = self.step_freq
        state = None
        # start main loop
        while rclpy.ok():
            next_sleep = self.node.get_clock().now() + Duration(seconds=1.0 / rate)
            if self.robot_state == RobotControlState.WALKING or (
                    self.robot_state == RobotControlState.CONTROLLABLE and (
                    not self.current_command_speed == (0, 0, 0) or not self.is_stopped)):
                # only run if we are in the correct state
                obs = self.compute_observation()
                if obs is None:
                    continue
                # we need to normalize the observation
                norm_obs = venv.normalize_obs(obs)
                action, state = model.predict(norm_obs, state=state, deterministic=True)
                self.last_action = action
                self.apply_action(action)
                # progress phase and compute current state
                time = self.node.get_clock().now().nanoseconds / 1.0e9
                self.progress_phase(time, action)
                self.ros_interface.publish_action()
            else:
                # hard stop walking when falling, getting up, etc.
                self.current_command_speed = (0, 0, 0)
                self.is_stopped = True

            self.state.publish_debug()

            #todo improve this, maybe complete asynch spin?
            # need to spin for multiple subscribers
            for i in range(5):
                rclpy.spin_once(self.node)

            self.node.get_clock().sleep_until(next_sleep)

class StoreDict(argparse.Action):
    """
    From https://github.com/DLR-RM/rl-baselines3-zoo
    Custom argparse action for storing dict.

    In: args1:0.0 args2:"dict(a=1)"
    Out: {'args1': 0.0, arg2: dict(a=1)}
    """

    def __init__(self, option_strings, dest, nargs=None, **kwargs):
        self._nargs = nargs
        super(StoreDict, self).__init__(option_strings, dest, nargs=nargs, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        arg_dict = {}
        for arguments in values:
            key = arguments.split(":")[0]
            value = ":".join(arguments.split(":")[1:])
            # Evaluate the string as python code
            arg_dict[key] = eval(value)
        setattr(namespace, self.dest, arg_dict)

def get_saved_hyperparams(
        stats_path: str,
        norm_reward: bool = False,
        test_mode: bool = False,
) -> Tuple[Dict[str, Any], str]:
    """
    From https://github.com/DLR-RM/rl-baselines3-zoo
    :param stats_path:
    :param norm_reward:
    :param test_mode:
    :return:
    """
    hyperparams = {}
    if not os.path.isdir(stats_path):
        stats_path = None
    else:
        config_file = os.path.join(stats_path, "config.yml")
        if os.path.isfile(config_file):
            # Load saved hyperparameters
            with open(os.path.join(stats_path, "config.yml"), "r") as f:
                hyperparams = yaml.load(f, Loader=yaml.UnsafeLoader)  # pytype: disable=module-attr
            hyperparams["normalize"] = hyperparams.get("normalize", False)
        else:
            obs_rms_path = os.path.join(stats_path, "obs_rms.pkl")
            hyperparams["normalize"] = os.path.isfile(obs_rms_path)

        # Load normalization params
        if hyperparams["normalize"]:
            if isinstance(hyperparams["normalize"], str):
                normalize_kwargs = eval(hyperparams["normalize"])
                if test_mode:
                    normalize_kwargs["norm_reward"] = norm_reward
            else:
                normalize_kwargs = {"norm_obs": hyperparams["normalize"], "norm_reward": norm_reward}
            hyperparams["normalize_kwargs"] = normalize_kwargs
    return hyperparams, stats_path

def create_test_env(
        env_id: str,
        n_envs: int = 1,
        stats_path: Optional[str] = None,
        seed: int = 0,
        log_dir: Optional[str] = None,
        should_render: bool = True,
        hyperparams: Optional[Dict[str, Any]] = None,
        env_kwargs: Optional[Dict[str, Any]] = None,
) -> VecEnv:
    """
    Create environment for testing a trained agent

    :param env_id:
    :param n_envs: number of processes
    :param stats_path: path to folder containing saved running averaged
    :param seed: Seed for random number generator
    :param log_dir: Where to log rewards
    :param should_render: For Pybullet env, display the GUI
    :param hyperparams: Additional hyperparams (ex: n_stack)
    :param env_kwargs: Optional keyword argument to pass to the env constructor
    :return:
    """
    hyperparams = {} if hyperparams is None else hyperparams

    vec_env_kwargs = {}
    vec_env_cls = DummyVecEnv

    env = make_vec_env(
        env_id,
        n_envs=n_envs,
        monitor_dir=log_dir,
        seed=seed,
        env_kwargs=env_kwargs,
        vec_env_cls=vec_env_cls,
        vec_env_kwargs=vec_env_kwargs,
    )

    # Load saved stats for normalizing input and rewards
    # And optionally stack frames
    if stats_path is not None:
        if hyperparams["normalize"]:
            print("Loading running average")
            print(f"with params: {hyperparams['normalize_kwargs']}")
            path_ = os.path.join(stats_path, "vecnormalize.pkl")
            if os.path.exists(path_):
                env = VecNormalize.load(path_, env)
                # Deactivate training and reward normalization
                env.training = False
                env.norm_reward = False
            else:
                print(f"VecNormalize stats {path_} not found")
                exit(1)
        else:
            print("normalize param is False. There is probably something wrong")
            exit(1)

        n_stack = hyperparams.get("frame_stack", 0)
        if n_stack > 0:
            print(f"Stacking {n_stack} frames")
            env = VecFrameStack(env, n_stack)
    return env

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--algo", help="RL Algorithm", default="ppo", type=str, required=False,
                        choices=list(ALGOS.keys()))
    parser.add_argument("-f", "--folder", help="Log folder", type=str, default="rl-trained-agents")
    parser.add_argument("--exp-id", help="Experiment ID", required=True, type=int)
    parser.add_argument(
        "--load-best", action="store_true", default=False, help="Load best model instead of last model if available"
    )
    parser.add_argument(
        "--load-checkpoint",
        type=int,
        help="Load checkpoint instead of last model if available, "
             "you must pass the number of timesteps corresponding to it",
    )
    parser.add_argument(
        "--load-last-checkpoint",
        action="store_true",
        default=False,
        help="Load last checkpoint instead of last model if available",
    )
    parser.add_argument(
        "--env-kwargs", type=str, nargs="+", action=StoreDict,
        help="Optional keyword argument to pass to the env constructor"
    )
    args = parser.parse_args()

    # load model, like it is done in the enjoy.py from rl-baselines zoo
    env_id = "WolfgangBulletEnv-v1"
    algo = args.algo
    folder = args.folder
    log_path = os.path.join(folder, algo, f"{env_id}_{args.exp_id}")
    assert os.path.isdir(log_path), f"The {log_path} folder was not found"

    found = False
    for ext in ["zip"]:
        model_path = os.path.join(log_path, f"{env_id}.{ext}")
        found = os.path.isfile(model_path)
        if found:
            break

    if args.load_best:
        model_path = os.path.join(log_path, "best_model.zip")
        found = os.path.isfile(model_path)
    if args.load_checkpoint is not None:
        model_path = os.path.join(log_path, f"rl_model_{args.load_checkpoint}_steps.zip")
        found = os.path.isfile(model_path)
    if args.load_last_checkpoint:
        checkpoints = glob.glob(os.path.join(log_path, "rl_model_*_steps.zip"))
        if len(checkpoints) == 0:
            raise ValueError(f"No checkpoint found for {algo} on {env_id}, path: {log_path}")

        def step_count(s):
            # path follow the pattern "rl_model_*_steps.zip", we count from the back to ignore any other _ in the path
            return int(s.split("_")[-2])

        checkpoints = sorted(checkpoints, key=step_count)
        model_path = checkpoints[-1]
        found = True

    if not found:
        raise ValueError(f"No model found for {algo} on {env_id}, path: {model_path}")
    print(f"Loading {model_path}")

    stats_path = os.path.join(log_path, env_id)
    hyperparams, stats_path = get_saved_hyperparams(stats_path, norm_reward=False, test_mode=True)

    # load env_kwargs if existing
    env_kwargs = {}
    args_path = os.path.join(log_path, env_id, "args.yml")
    if os.path.isfile(args_path):
        with open(args_path, "r") as f:
            loaded_args = yaml.load(f, Loader=yaml.UnsafeLoader)  # pytype: disable=module-attr
            if loaded_args["env_kwargs"] is not None:
                env_kwargs = loaded_args["env_kwargs"]
    # overwrite with command line arguments
    if args.env_kwargs is not None:
        env_kwargs.update(args.env_kwargs)

    print(env_kwargs)
    venv = create_test_env(
        "ExecuteEnv-v1",
        n_envs=1,
        stats_path=stats_path,
        log_dir=None,
        should_render=False,
        hyperparams=hyperparams,
        env_kwargs=env_kwargs,
    )

    # direct reference to wolfgnag env object
    env = venv.venv.envs[0].env.env

    custom_objects = {
        "learning_rate": 0.0,
        "lr_schedule": lambda _: 0.0,
        "clip_range": lambda _: 0.0,
    }
    model = ALGOS[algo].load(model_path, env=venv, custom_objects=custom_objects)

    env.run_node()


import os
import config
import random
import colorsys
import numpy as np
import tensorflow as tf
from model.yolo3_model import yolo

class yolo_predictor:
    def __init__(self, obj_threshold, nms_threshold, classes_file, anchors_file):
        """
        Introduction
        ------------
            初始化函数
        Parameters
        ----------
            obj_threshold: 目标检测为物体的阈值
            nms_threshold: nms阈值
        """
        self.obj_threshold = obj_threshold
        self.nms_threshold = nms_threshold
        self.classes_path = classes_file
        self.anchors_path = anchors_file
        self.class_names = self._get_class()
        self.anchors = self._get_anchors()
        hsv_tuples = [(x / len(self.class_names), 1., 1.)for x in range(len(self.class_names))]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))
        random.seed(10101)
        random.shuffle(self.colors)
        random.seed(None)

    def _get_class(self):
        """
        Introduction
        ------------
            读取类别名称
        """
        classes_path = os.path.expanduser(self.classes_path)
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

    def _get_anchors(self):
        """
        Introduction
        ------------
            读取anchors数据
        """
        anchors_path = os.path.expanduser(self.anchors_path)
        with open(anchors_path) as f:
            anchors = f.readline()
            anchors = [float(x) for x in anchors.split(',')]
            anchors = np.array(anchors).reshape(-1, 2)
        return anchors

    def eval(self, yolo_outputs, image_shape, max_boxes = 20):
        """
        Introduction
        ------------
            根据Yolo模型的输出进行非极大值抑制，获取最后的物体检测框和物体检测类别
        Parameters
        ----------
            yolo_outputs: yolo模型输出
            image_shape: 图片的大小
            max_boxes:  最大box数量
        Returns
        -------
            boxes_: 物体框的位置
            scores_: 物体类别的概率
            classes_: 物体类别
        """
        anchor_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
        boxes = []
        box_scores = []
        input_shape = tf.shape(yolo_outputs[0])[1 : 3] * 32
        # 对三个尺度的输出获取每个预测box坐标和box的分数，score计算为置信度x类别概率
        for i in range(len(yolo_outputs)):
            _boxes, _box_scores = self.boxes_and_scores(yolo_outputs[i], self.anchors[anchor_mask[i]], len(self.class_names), input_shape, image_shape)
            boxes.append(_boxes)
            box_scores.append(_box_scores)
        boxes = tf.concat(boxes, axis = 0)
        box_scores = tf.concat(box_scores, axis = 0)

        mask = box_scores >= self.obj_threshold
        max_boxes_tensor = tf.constant(max_boxes, dtype = tf.int32)
        boxes_ = []
        scores_ = []
        classes_ = []
        for c in range(len(self.class_names)):
            class_boxes = tf.boolean_mask(boxes, mask[:, c])
            class_box_scores = tf.boolean_mask(box_scores[:, c], mask[:, c])
            nms_index = tf.image.non_max_suppression(class_boxes, class_box_scores, max_boxes_tensor, iou_threshold = self.nms_threshold)
            class_boxes = tf.gather(class_boxes, nms_index)
            class_box_scores = tf.gather(class_box_scores, nms_index)
            classes = tf.ones_like(class_box_scores, 'int32') * c
            boxes_.append(class_boxes)
            scores_.append(class_box_scores)
            classes_.append(classes)
        boxes_ = tf.concat(boxes_, axis = 0)
        scores_ = tf.concat(scores_, axis = 0)
        classes_ = tf.concat(classes_, axis = 0)
        return boxes_, scores_, classes_

    def boxes_and_scores(self, feats, anchors, classes_num, input_shape, image_shape):
        """
        Introduction
        ------------
            将预测出的box坐标转换为对应原图的坐标，然后计算每个box的分数
        Parameters
        ----------
            feats: yolo输出的feature map
            anchors: anchor的位置
            class_num: 类别数目
            input_shape: 输入大小
            image_shape: 图片大小
        Returns
        -------
            boxes: 物体框的位置
            boxes_scores: 物体框的分数，为置信度和类别概率的乘积
        """
        box_xy, box_wh, box_confidence, box_class_probs = self._get_feats(feats, anchors, classes_num, input_shape)
        boxes = self.correct_boxes(box_xy, box_wh, input_shape, image_shape)
        boxes = tf.reshape(boxes, [-1, 4])
        box_scores = box_confidence * box_class_probs
        box_scores = tf.reshape(box_scores, [-1, classes_num])
        return boxes, box_scores

    def correct_boxes(self, box_xy, box_wh, input_shape, image_shape):
        """
        Introduction
        ------------
            计算物体框预测坐标在原图中的位置坐标
        Parameters
        ----------
            box_xy: 物体框左上角坐标
            box_wh: 物体框的宽高
            input_shape: 输入的大小
            image_shape: 图片的大小
        Returns
        -------
            boxes: 物体框的位置
        """
        box_yx = box_xy[..., ::-1]
        box_hw = box_wh[..., ::-1]
        input_shape = tf.cast(input_shape, dtype = tf.float32)
        image_shape = tf.cast(image_shape, dtype = tf.float32)
        new_shape = tf.round(image_shape * tf.reduce_min(input_shape / image_shape))
        offset = (input_shape - new_shape) / 2. / input_shape
        scale = input_shape / new_shape
        box_yx = (box_yx - offset) * scale
        box_hw *= scale

        box_mins = box_yx - (box_hw / 2.)
        box_maxes = box_yx + (box_hw / 2.)
        boxes = tf.concat([
            box_mins[..., 0:1],
            box_mins[..., 1:2],
            box_maxes[..., 0:1],
            box_maxes[..., 1:2]
        ], axis = -1)
        boxes *= tf.concat([image_shape, image_shape], axis = -1)
        return boxes

    def _get_feats(self, feats, anchors, num_classes, input_shape):
        """
        Introduction
        ------------
            根据yolo最后一层的输出确定bounding box
        Parameters
        ----------
            feats: yolo模型最后一层输出
            anchors: anchors的位置
            num_classes: 类别数量
            input_shape: 输入大小
        Returns
        -------
            box_xy, box_wh, box_confidence, box_class_probs
        """
        num_anchors = len(anchors)
        anchors_tensor = tf.reshape(tf.constant(anchors, dtype=tf.float32), [1, 1, 1, num_anchors, 2])
        grid_size = tf.shape(feats)[1:3]
        predictions = tf.reshape(feats, [-1, grid_size[0], grid_size[1], num_anchors, num_classes + 5])
        # 这里构建13*13*1*2的矩阵，对应每个格子加上对应的坐标
        grid_y = tf.tile(tf.reshape(tf.range(grid_size[0]), [-1, 1, 1, 1]), [1, grid_size[1], 1, 1])
        grid_x = tf.tile(tf.reshape(tf.range(grid_size[1]), [1, -1, 1, 1]), [grid_size[0], 1, 1, 1])
        grid = tf.concat([grid_x, grid_y], axis = -1)
        grid = tf.cast(grid, tf.float32)
        # 将x,y坐标归一化为占416的比例
        box_xy = (tf.sigmoid(predictions[..., :2]) + grid) / tf.cast(grid_size[::-1], tf.float32)
        # 将w,h也归一化为占416的比例
        box_wh = tf.exp(predictions[..., 2:4]) * anchors_tensor / tf.cast(input_shape[::-1], tf.float32)
        box_confidence = tf.sigmoid(predictions[..., 4:5])
        box_class_probs = tf.sigmoid(predictions[..., 5:])
        return box_xy, box_wh, box_confidence, box_class_probs

    def predict(self, inputs, image_shape):
        """
        Introduction
        ------------
            构建预测模型
        Parameters
        ----------
            inputs: 处理之后的输入图片
            image_shape: 图像原始大小
        Returns
        -------
            boxes: 物体框坐标
            scores: 物体概率值
            classes: 物体类别
        """
        model = yolo(config.norm_epsilon, config.norm_decay, self.anchors_path, self.classes_path, pre_train = False)
        output = model.yolo_inference(inputs, config.num_anchors // 3, config.num_classes, training = False)
        boxes, scores, classes = self.eval(output, image_shape, max_boxes = 20)
        return boxes, scores, classes


import numpy as np
from image_processing import *
import random
import torch
import os
from model import DQN
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision.transforms as T
import pickle

from logger import Logger
logger = Logger('./logs')

# Global Variables
BATCH_SIZE = 32
TASK = 'MontezumaRevenge-v0'

MAX_FRAMES = 2500000
REPLAY_BUFFER_SIZE = 1000000

TARGET_UPDATE_FREQ = 10000
TEST_FREQ = 10000

REPEAT_ACTION = 4
SKIP_FRAMES = 4

GAMMA = 0.99
LEARNING_FREQ = 4
ALPHA = 0.95
EPS = 0.01
LEARNING_STARTS = 10000
SAVE_RESULTS_N_STEP = 100000

num_param_updates = 0
mean_episode_reward      = -float('nan')
best_mean_episode_reward = -float('inf')
LOG_EVERY_N_STEPS = 1000
SAVE_MODEL_EVERY_N_STEPS = 500000
SUBGOAL_DISCOVERY_FREQ = REPLAY_BUFFER_SIZE // 20
META_CONTROLLER_UPDATE_FREQ = 10000
NOOP_MAX = 30 # no operation at the beggining of the game
MAX_STEPS = 1000
MAX_STEPS_TEST = 250
ALPHA = 0.95
EPS = 0.01
LEARNING_RATE = 0.00025

def sample_from_random_subgoal_set(subgoal_set):
	index = random.randint(0, len(subgoal_set)-1)
	return index, subgoal_set[index]
# Create random subgoals from image features
rec = Recognizer()
base_img = cv2.imread('./templates/base.png')
img_edge = base_img[55:181,10:150,:]
im2, contours, hierarchy = edge_detector(img_edge)
coords = rec.get(base_img)
# random_subgoals_set = create_random_subgoal_set_from_objects(coords)
good_subgoals = create_good_subgoal_set_from_objects(coords)
random_subgoals_set = good_subgoals
# This is how we find if the subgoal is attained
man_mask = rec.get_man_mask(base_img)
subgoal_id, random_subgoal_mask = sample_from_random_subgoal_set(random_subgoals_set)
random_subgoal_frame = create_mask_frame(base_img,random_subgoal_mask)
intrinsic_learning_task_done = are_masks_align(man_mask, random_subgoal_mask)

# BUILD MODEL
USE_CUDA = torch.cuda.is_available()
if torch.cuda.is_available():
	device0 = torch.device("cuda:0")
else:
	device0 = torch.device("cpu")

dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
dlongtype = torch.cuda.LongTensor if torch.cuda.is_available() else torch.LongTensor
duinttype = torch.cuda.ByteTensor if torch.cuda.is_available() else torch.ByteTensor

Qt = DQN(in_channels=5, num_actions=18).type(dtype)
Qt_t = DQN(in_channels=5, num_actions=18).type(dtype)
Qt_t.load_state_dict(Qt.state_dict())
Qt_t.eval()
for param in Qt_t.parameters():
	param.requires_grad = False

if torch.cuda.device_count() > 0:
	Qt.cuda()
	Qt = nn.DataParallel(Qt).to(device0)
	Qt_t = nn.DataParallel(Qt_t).to(device0)
	batch_size = BATCH_SIZE * torch.cuda.device_count()
else:
	batch_size = BATCH_SIZE

# optimizer
optimizer = optim.RMSprop(Qt.parameters(),lr=LEARNING_RATE, alpha=ALPHA, eps=EPS)

# training parameters
# Create environment
import gym
env = gym.make('MontezumaRevenge-v0')
test_env = gym.make('MontezumaRevenge-v0')

actions_meaning = env.unwrapped.get_action_meanings()

# from hrl import MetaLearning
# agent = MetaLearning()

def reset(noop_max=30):
	s = env.reset()
	noop_random = random.randint(1, NOOP_MAX)
	for _ in range(noop_random):
		s,_,_,_ = env.step(0)
	S,_, _, _ = step(0)
	return S

def reset_test_env(noop_max=30):
	s = test_env.reset()
	noop_random = random.randint(1, NOOP_MAX)
	for _ in range(noop_random):
		s,_,_,_ = test_env.step(0)
	S,_, _, _ = step(0)
	return S

def skip_frames(a,skip_frame=SKIP_FRAMES):
	frames = []
	rewards = []
	for _ in range(skip_frame):
		s,r,done,step_info = env.step(a)
		frames.append(s)
		rewards.append(r)
	s_max = np.max(np.stack(frames), axis=0)
	total_rewards = sum(rewards)
	return s_max, total_rewards, done, step_info

def skip_frames_test_env(a,skip_frame=SKIP_FRAMES):
	frames = []
	rewards = []
	for _ in range(skip_frame):
		s,r,done,step_info = test_env.step(a)
		frames.append(s)
		rewards.append(r)
	s_max = np.max(np.stack(frames), axis=0)
	total_rewards = sum(rewards)
	return s_max, total_rewards, done, step_info

def step(a,repeat_action=REPEAT_ACTION):
	frames = []
	rewards = []
	for _ in range(repeat_action):
		s,r,done,step_info = skip_frames(a)
		frames.append(s)
		rewards.append(r)
	total_rewards = sum(rewards)
	return frames, total_rewards, done, step_info

def test_step(a,repeat_action=REPEAT_ACTION):
	frames = []
	rewards = []
	for _ in range(repeat_action):
		s,r,done,step_info = skip_frames_test_env(a)
		frames.append(s)
		rewards.append(r)
	total_rewards = sum(rewards)
	return frames, total_rewards, done, step_info

def get_man_mask(frames):
	s_max = np.max(np.stack(frames), axis=0)
	man_mask = rec.get_man_mask(s_max)
	return man_mask

def get_man_xy_np_coordinate(man_mask):
	x = man_mask.x
	y = man_mask.y
	return np.array([x,y])

def four_frames_to_4_84_84(S):
	""" 0) Atari frames: 210 x 160
		1) Get image grayscale
		2) Rescale image 110 x 84
		3) Crop center 84 x 84 (you can crop top/bottom according to the game)
		"""
	for i, img in enumerate(S):
		gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
		h, w = 110, 84
		gray_resized = cv2.resize(gray,(w,h))
		gray_cropped = gray_resized[13:110-13,:]
		gray_reshaped = gray_cropped.reshape((1,84,84))
		if i == 0:
			s = gray_reshaped
		else:
			s = np.concatenate((s,gray_reshaped),axis=0)
	return s

def single_channel_frame_to_1_84_84(subgoal_frame):
	reshaped = subgoal_frame.reshape((210,160))
	resized =  cv2.resize(reshaped,(84,110))
	cropped = resized[13:110-13,:]
	g = cropped.reshape((1,84,84))
	return g

def epsilon_greedy(Q,epsilon=0.1):
	if random.random() < epsilon:
		return env.action_space.sample()
	else:
		return Q.argmax()

# create an Expereince Memory
from memory import Experience, ExperienceMemory
experience_memory = ExperienceMemory(size=REPLAY_BUFFER_SIZE)

from subgoal_discovery import SubgoalDiscovery
sd = SubgoalDiscovery()

# centroid_low_return = []
# centroid_high_return = []

# outliers_low_return = []
# outliers_high_return = []

### PHASE I: INTRINSIC MOTIVATION LEARNING + SUBGOAL DISCOVERY
### Goal: Training the Controller via intrinsic motivation learning
### Random subgoals and explorations
# just Controller, not Meta-Controller
S = reset()
s = four_frames_to_4_84_84(S)
man_mask = get_man_mask(S)
man_loc = get_man_xy_np_coordinate(man_mask)
subgoal_index, subgoal_mask = sample_from_random_subgoal_set(random_subgoals_set) # random g
subgoal_frame = create_mask_frame(base_img,subgoal_mask)
g = single_channel_frame_to_1_84_84(subgoal_frame)
steps = 0
R = 0 # return
firs_time = True
num_param_updates = 0
epsilon = 1
first_time_kmeans = True

# saving training variables
outliers = []
centroids = []
G = []
episode_rewards = []

mean_reward_episodes_list = []
best_reward_episodes_list = []
episode_rewards_list = []
good_test_states = []
subgoal_success_test = {}
for i in range(len(random_subgoals_set)):
	subgoal_success_test[i] = 0

def test_task():
	S = reset_test_env()
	s = four_frames_to_4_84_84(S)
	man_mask = get_man_mask(S)
	man_loc = get_man_xy_np_coordinate(man_mask)
	success_test = 0
	test_episode_frames = []
	test_episode_total_rewards = 0

	for subgoal_index, subgoal_mask in enumerate(random_subgoals_set[0:4]):

		# sometimes pick the other door
		if random.random() < 0.5 and subgoal_index == 3:
			subgoal_index = 4
			subgoal_mask = random_subgoals_set[4]

		subgoal_frame = create_mask_frame(base_img,subgoal_mask)
		g = single_channel_frame_to_1_84_84(subgoal_frame)
		x = np.concatenate((s,g),axis=0).reshape((1,5,84,84))
		k = 0
		while(k < MAX_STEPS_TEST):
			qt = Qt.forward(torch.Tensor(x).type(dtype)/255)
			a = epsilon_greedy(qt.cpu().detach().numpy(),epsilon=0) # random action
			SP, r, terminal, step_info = test_step(a)
			test_episode_frames = S[3]
			sp = four_frames_to_4_84_84(SP)
			xp = np.concatenate((sp,g),axis=0).reshape((1,5,84,84))
			man_mask = get_man_mask(SP)
			man_loc = get_man_xy_np_coordinate(man_mask)
			# intrinsic_done_task = are_masks_align(man_mask, subgoal_mask)
			intrinsic_done_task = is_man_inside_subgoal_mask(man_mask, subgoal_mask)
			# outlier
			if r > 0:
				print('reward: ', r)
				print('Outler detected at', man_loc)
				outliers.append(man_loc)
				test_episode_total_rewards += r

			g_id = subgoal_index

			if intrinsic_done_task:
				print('succesful intrinsic motivation test to g in ',\
										 (subgoal_mask.x,subgoal_mask.y) )

				subgoal_success_test[subgoal_index] += 1
				break

			s = sp
			k += 1

			if terminal:
				S = reset_test_env() # s is reserved for 4*84*84 input image
				s = four_frames_to_4_84_84(S)
				break

		if r>0 and intrinsic_done_task:
			success_test = 1

	return test_episode_frames, subgoal_success_test ,success_test, test_episode_total_rewards

for t in range(MAX_FRAMES):
	x = np.concatenate((s,g),axis=0).reshape((1,5,84,84))
	if t < LEARNING_STARTS:
		a = env.action_space.sample()
	else:
		qt = Qt.forward(torch.Tensor(x).type(dtype)/255)
		a = epsilon_greedy(qt.cpu().detach().numpy(),epsilon=epsilon) # random action
	SP, r, terminal, step_info = step(a)
	episode_rewards.append(r)
	sp = four_frames_to_4_84_84(SP)
	xp = np.concatenate((sp,g),axis=0).reshape((1,5,84,84))
	man_mask = get_man_mask(SP)
	man_loc = get_man_xy_np_coordinate(man_mask)
	# intrinsic_done_task = are_masks_align(man_mask, subgoal_mask)
	intrinsic_done_task = is_man_inside_subgoal_mask(man_mask, subgoal_mask)
	# outlier
	if r > 0:
		print('Outler detected at', man_loc)
		outliers.append(man_loc)
	R += r

	r = np.clip(r, -1.0, 1.0)

	g_id = subgoal_index

	if intrinsic_done_task:
		print('succesful intrinsic motivation learning to g in ',\
									 (subgoal_mask.x,subgoal_mask.y) )
		intrinsic_done = 1
		tilde_r = +1
		subgoal_index, subgoal_mask = \
			sample_from_random_subgoal_set(random_subgoals_set)
		subgoal_frame = create_mask_frame(base_img,subgoal_mask)
	else:
		intrinsic_done = 0
		tilde_r = -1

	if terminal:
		done = 1
	else:
		done = 0
		r = -1

	experience = Experience(s, g, g_id, a, r, tilde_r, sp, intrinsic_done, done, man_loc)
	experience_memory.push(experience)

	epsilon = max(0.1, 1-(1-0.1)*t/1000000 )
	s = sp
	steps += 1

	if terminal or (steps>MAX_STEPS):
		S = reset() # s is reserved for 4*84*84 input image
		s = four_frames_to_4_84_84(S)
		man_mask = get_man_mask(S)
		man_loc = get_man_xy_np_coordinate(man_mask)
		subgoal_index, subgoal_mask = sample_from_random_subgoal_set(random_subgoals_set) # random g
		subgoal_frame = create_mask_frame(base_img,subgoal_mask)
		g = single_channel_frame_to_1_84_84(subgoal_frame)
		steps = 0
		R = 0
		episode_rewards = []

	if (t > LEARNING_STARTS and t % LEARNING_FREQ == 0):
		states, subgoals, actions, rewards, state_primes, intrinsic_dones = \
						experience_memory.sample_controller(batch_size=batch_size)
		x = np.concatenate((states,subgoals),axis=1)
		x = torch.Tensor(x)
		xp = np.concatenate((state_primes,subgoals),axis=1)
		xp = torch.Tensor(xp)
		actions = torch.Tensor(actions).type(dlongtype)
		rewards = torch.Tensor(rewards).type(dtype)
		intrinsic_dones = torch.Tensor(intrinsic_dones).type(dtype)

		if torch.cuda.is_available():
			with torch.cuda.device(0):
				x = torch.Tensor(x).to(device0).type(dtype)/255
				xp = torch.Tensor(xp).to(device0).type(dtype)/255
				actions = actions.to(device0)
				rewards = rewards.to(device0)
				intrinsic_dones = intrinsic_dones.to(device0)

		# if torch.cuda.device_count() > 0:
		# 	Qt.to(device0)
		# 	Qt_t = Qt_t.to(device0)

		qt_values = Qt.forward(x)
		qt = qt_values.gather(1, actions.unsqueeze(1))
		qt = qt.squeeze()

		qt_p1 = Qt.forward(xp)
		_, a_prime = qt_p1.max(1)

		qt_t_p1 = Qt_t.forward(xp).detach()
		qt_t_p1 = qt_t_p1.gather(1, a_prime.unsqueeze(1))
		qt_t_p1 = qt_t_p1.squeeze()

		# error = rewards + GAMMA * (1 - intrinsic_dones) * qt_t_p1 - qt
		target = rewards + GAMMA * (1 - intrinsic_dones) * qt_t_p1
		# clipped_error = -1.0 * error.clamp(-1, 1)

		# Compute Huber loss
		loss = F.smooth_l1_loss(qt, target)
		optimizer.zero_grad()
		# qt.backward(clipped_error.data)
		loss.backward()

		for param in Qt.parameters():
			param.grad.data.clamp_(-1, 1)

		optimizer.step()
		num_param_updates += 1

		if num_param_updates % TARGET_UPDATE_FREQ == 0:
			Qt_t.load_state_dict(Qt.state_dict())

		if t % LOG_EVERY_N_STEPS == 0:
			for tag, value in Qt.named_parameters():
				tag = tag.replace('.', '/')
				logger.histo_summary(tag, value.data.cpu().numpy(), t+1)
				logger.histo_summary(tag+'/grad', value.grad.cpu().numpy(), t+1)

	if t % SAVE_MODEL_EVERY_N_STEPS == 0:
		if not os.path.exists("models"):
			os.makedirs("models")
		add_str = 'DQN'
		model_save_path = "models/%s_%s_%d_%s.model" %(str(TASK), add_str, t, str(time.ctime()).replace(' ', '_'))
		torch.save(Qt.state_dict(), model_save_path)

	if len(episode_rewards) > 0:
		mean_episode_reward = np.mean(episode_rewards[-100:])
		best_mean_episode_reward = max(best_mean_episode_reward, mean_episode_reward)

	if t % LOG_EVERY_N_STEPS == 0:
		print("---------------------------------")
		print("Timestep %d" % (t,))
		print("learning started? %d" % (t > LEARNING_STARTS))
		print("mean reward (100 episodes) %f" % mean_episode_reward)
		print("best mean reward %f" % best_mean_episode_reward)
		print("episodes %d" % len(episode_rewards))
		print("exploration %f" % epsilon)
		print("learning_rate %f" % LEARNING_RATE)
		sys.stdout.flush()

		best_reward_episodes_list.append(best_mean_episode_reward)
		mean_reward_episodes_list.append(mean_episode_reward)
		episode_rewards_list.append(sum(episode_rewards))
		#============ TensorBoard logging ============#
		# Log the scalar values
		info = {
			'learning_started': (t > LEARNING_STARTS),
			'num_episodes': len(episode_rewards),
			'exploration': epsilon,
			'learning_rate': LEARNING_RATE,
		}

		for tag, value in info.items():
			logger.scalar_summary(tag, value, t+1)

		if len(episode_rewards) > 0:
			info = {
				'last_episode_rewards': episode_rewards[-1],
			}

			for tag, value in info.items():
				logger.scalar_summary(tag, value, t+1)

		if (best_mean_episode_reward != -float('inf')):
			info = {
				'mean_episode_reward_last_100': mean_episode_reward,
				'best_mean_episode_reward': best_mean_episode_reward
			}

			for tag, value in info.items():
				logger.scalar_summary(tag, value, t+1)

	if t % SAVE_RESULTS_N_STEP==0 and t>0:
		results_file_path = './results/performance_results_t_'+ str(t) +'.pkl'
		with open(results_file_path, 'wb') as f:
			pickle.dump([best_reward_episodes_list, mean_reward_episodes_list,episode_rewards_list], f)

	if t % SUBGOAL_DISCOVERY_FREQ == 0 and t > 0:
		if first_time_kmeans:
			X = experience_memory.get_man_positions()
			sd.feed_data(X)
			sd.find_kmeans_clusters()
			C = sd.cluster_centroids()
			first_time_kmeans = False
		else:
			X = experience_memory.get_man_positions()
			sd.feed_data(X)
			sd.find_kmeans_clusters(init=C)
			C = sd.cluster_centroids()
		centroids.append(C)
		results_file_path = './results/clustering_results_t_'+ str(t) +'.pkl'
		with open(results_file_path, 'wb') as f:
			pickle.dump([C,outliers], f)
		print('Current discovered centroids')
		print(C)

	if t % TEST_FREQ == 0 and t>0:
		test_episode_frames, subgoal_success_test ,success_test, test_episode_total_rewards = test_task()
		results_file_path = './results/test_results_t_'+ str(t) +'.pkl'
		with open(results_file_path, 'wb') as f:
			pickle.dump([test_episode_frames, subgoal_success_test ,success_test, test_episode_total_rewards], f)


import oneflow as flow
import networks
import itertools
from image import ImagePool, ndarray2image
import numpy as np
import cv2

class CycleGANModel:
    def __init__(self, opt):
        self.opt = opt
        self.device = "cuda"
        self.netG_A = networks.ResnetGenerator(n_blocks=opt.n_blocks).to(self.device)
        self.netG_B = networks.ResnetGenerator(n_blocks=opt.n_blocks).to(self.device)

        self.netD_A = networks.NLayerDiscriminator().to(self.device)
        self.netD_B = networks.NLayerDiscriminator().to(self.device)

        self.fake_A_pool = ImagePool(opt.pool_size)
        self.fake_B_pool = ImagePool(opt.pool_size)

        self.criterionGAN = networks.GANLoss("lsgan").to(self.device)
        self.criterionCycle = flow.nn.L1Loss()
        self.criterionIdt = flow.nn.L1Loss()

        self.optimizer_G = flow.optim.Adam(
            itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()),
            lr=opt.lr,
            betas=(opt.beta1, opt.beta2),
        )
        self.optimizer_D = flow.optim.Adam(
            itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()),
            lr=opt.lr,
            betas=(opt.beta1, opt.beta2),
        )
        self.optimizers = [self.optimizer_G, self.optimizer_D]
        self.schedulers = [
            flow.optim.lr_scheduler.CosineAnnealingLR(optimizer, steps=100, alpha=0.0)
            for optimizer in self.optimizers
        ]

    def update_learning_rate(self):
        """Update learning rates for all the networks; called at the end of every epoch"""
        old_lr = self.optimizers[0].param_groups[0]["lr"]
        for scheduler in self.schedulers:
            scheduler.step()

        lr = self.optimizers[0].param_groups[0]["lr"]
        print("learning rate %.7f -> %.7f" % (old_lr, lr))

    def set_requires_grad(self, nets, requires_grad=False):
        """Set requies_grad=Fasle for all the networks to avoid unnecessary computations
        Parameters:
            nets (network list)   -- a list of networks
            requires_grad (bool)  -- whether the networks require gradients or not
        """
        if not isinstance(nets, list):
            nets = [nets]
        for net in nets:
            if net is not None:
                for param in net.parameters():
                    param.requires_grad = requires_grad

    def set_input(self, image_A, image_B):
        """Unpack input data from the dataloader and perform necessary pre-processing steps.

        Parameters:
            input (dict): include the data itself and its metadata information.

        The option 'direction' can be used to swap domain A and domain B.
        """
        self.real_A = flow.Tensor(image_A).to(self.device)
        self.real_B = flow.Tensor(image_B).to(self.device)

    def forward(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
        self.fake_B = self.netG_A(self.real_A)  # G_A(A)
        self.rec_A = self.netG_B(self.fake_B)  # G_B(G_A(A))
        self.fake_A = self.netG_B(self.real_B)  # G_B(B)
        self.rec_B = self.netG_A(self.fake_A)  # G_A(G_B(B))

    def backward_D_basic(self, netD, real, fake):
        """Calculate GAN loss for the discriminator

        Parameters:
            netD (network)      -- the discriminator D
            real (tensor array) -- real images
            fake (tensor array) -- images generated by a generator

        Return the discriminator loss.
        We also call loss_D.backward() to calculate the gradients.
        """
        # Real
        pred_real = netD(real)
        loss_D_real = self.criterionGAN(pred_real, True)
        # Fake
        pred_fake = netD(fake.detach())
        loss_D_fake = self.criterionGAN(pred_fake, False)
        # Combined loss and calculate gradients
        loss_D = (loss_D_real + loss_D_fake) * 0.5
        loss_D.backward()
        return loss_D

    def backward_D_A(self):
        """Calculate GAN loss for discriminator D_A"""
        fake_B = self.fake_B_pool.query(self.fake_B)
        self.loss_D_A = self.backward_D_basic(self.netD_A, self.real_B, fake_B)

    def backward_D_B(self):
        """Calculate GAN loss for discriminator D_B"""
        fake_A = self.fake_A_pool.query(self.fake_A)
        self.loss_D_B = self.backward_D_basic(self.netD_B, self.real_A, fake_A)

    def backward_G(self):
        """Calculate the loss for generators G_A and G_B"""
        lambda_idt = self.opt.lambda_identity
        lambda_A = self.opt.lambda_A
        lambda_B = self.opt.lambda_B
        # Identity loss
        if lambda_idt > 0:
            # G_A should be identity if real_B is fed: ||G_A(B) - B||
            self.idt_A = self.netG_A(self.real_B)
            self.loss_idt_A = (
                self.criterionIdt(self.idt_A, self.real_B) * lambda_B * lambda_idt
            )
            # G_B should be identity if real_A is fed: ||G_B(A) - A||
            self.idt_B = self.netG_B(self.real_A)
            self.loss_idt_B = (
                self.criterionIdt(self.idt_B, self.real_A) * lambda_A * lambda_idt
            )
        else:
            self.loss_idt_A = 0
            self.loss_idt_B = 0

        # GAN loss D_A(G_A(A))
        self.loss_G_A = self.criterionGAN(self.netD_A(self.fake_B), True)
        # GAN loss D_B(G_B(B))
        self.loss_G_B = self.criterionGAN(self.netD_B(self.fake_A), True)
        # Forward cycle loss || G_B(G_A(A)) - A||
        self.loss_cycle_A = self.criterionCycle(self.rec_A, self.real_A) * lambda_A
        # Backward cycle loss || G_A(G_B(B)) - B||
        self.loss_cycle_B = self.criterionCycle(self.rec_B, self.real_B) * lambda_B
        # combined loss and calculate gradients
        self.loss_G = (
            self.loss_G_A
            + self.loss_G_B
            + self.loss_cycle_A
            + self.loss_cycle_B
            + self.loss_idt_A
            + self.loss_idt_B
        )
        self.loss_G.backward()

    def optimize_parameters(self):
        """Calculate losses, gradients, and update network weights; called in every training iteration"""
        # forward
        self.forward()  # compute fake images and reconstruction images.
        # G_A and G_B
        self.optimizer_G.zero_grad()  # set G_A and G_B's gradients to zero
        self.backward_G()  # calculate gradients for G_A and G_B
        self.optimizer_G.step()  # update G_A and G_B's weights
        # D_A and D_B
        self.optimizer_D.zero_grad()  # set D_A and D_B's gradients to zero
        self.backward_D_A()  # calculate gradients for D_A
        self.backward_D_B()  # calculate graidents for D_B
        self.optimizer_D.step()  # update D_A and D_B's weights

    def save_result(self, path):
        imageA = ndarray2image(self.real_A.numpy())
        imageB = ndarray2image(self.real_B.numpy())

        image_fake_B = ndarray2image(self.fake_B.numpy())
        image_rec_A = ndarray2image(self.rec_A.numpy())

        image_fake_A = ndarray2image(self.fake_A.numpy())
        image_rec_B = ndarray2image(self.rec_B.numpy())

        result1 = np.concatenate((imageA, image_fake_B, image_rec_A), axis=1)
        result2 = np.concatenate((imageB, image_fake_A, image_rec_B), axis=1)
        result = np.concatenate((result1, result2), axis=0)
        cv2.imwrite(path, result)

    def log_loss(self, epoch, iter):
        print(
            "epoch: %d, iter: %d, d_loss: %f, g_loss: %f"
            % (
                epoch,
                iter,
                self.loss_D_A.numpy()[0] + self.loss_D_B.numpy()[0],
                self.loss_G.numpy()[0],
            )
        )


import os
import re
import sys
import time
import math
from collections import Counter
from functools import partial
from tempfile import mkdtemp, NamedTemporaryFile
import logging
import multiprocessing as mp

# "hidden" features, in development
try:
    import MOODS.tools
    import MOODS.parsers
    import MOODS.scan
except ImportError:
    pass

from genomepy import Genome
from diskcache import Cache
import numpy as np
from sklearn.preprocessing import scale
import pandas as pd
import sqlite3

from gimmemotifs import __version__
from gimmemotifs.background import RandomGenomicFasta, gc_bin_bedfile
from gimmemotifs.config import MotifConfig, CACHE_DIR
from gimmemotifs.fasta import Fasta
from gimmemotifs.c_metrics import pwmscan
from gimmemotifs.motif import read_motifs
from gimmemotifs.utils import parse_cutoff, as_fasta, file_checksum, rc

try:
    import copy_reg
    import types

    def _pickle_method(m):
        if m.im_self is None:
            return getattr, (m.im_class, m.im_func.func_name)
        else:
            return getattr, (m.im_self, m.im_func.func_name)

    copy_reg.pickle(types.MethodType, _pickle_method)
except Exception:
    pass

# only used when using cache, should not be a requirement
try:
    from dogpile.cache import make_region
    import xxhash
except ImportError:
    pass

logger = logging.getLogger("gimme.scanner")
config = MotifConfig()
FPR = 0.01
lock = mp.Lock()

def print_cluster_error_message():
    logger.error("Cache is corrupted.")
    logger.error(
        "This can happen when you try to run a GimmeMotifs tool in parallel on a cluster."
    )
    logger.error(f"To solve this, delete the GimmeMotifs cache directory: {CACHE_DIR}")
    logger.error("and then see here for a workaround:")
    logger.error(
        "https://gimmemotifs.readthedocs.io/en/master/faq.html#sqlite-error-when-running-on-a-cluster"
    )

def _format_line(
    seq, seq_id, motif, score, pos, strand, bed=False, seq_p=None, strandmap=None
):
    if seq_p is None:
        seq_p = re.compile(r"([^\s:]+):(\d+)-(\d+)")
    if strandmap is None:
        strandmap = {-1: "-", 1: "+"}
    if bed:
        m = seq_p.search(seq_id)
        if m:
            chrom = m.group(1)
            start = int(m.group(2))
            return "{}\t{}\t{}\t{}\t{}\t{}".format(
                chrom,
                start + pos,
                start + pos + len(motif),
                motif.id,
                score,
                strandmap[strand],
            )
        else:
            return "{}\t{}\t{}\t{}\t{}\t{}".format(
                seq_id, pos, pos + len(motif), motif.id, score, strandmap[strand]
            )
    else:
        return '{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\tmotif_name "{}" ; motif_instance "{}"'.format(
            seq_id,
            "pfmscan",
            "misc_feature",
            pos + 1,  # GFF is 1-based
            pos + len(motif),
            score,
            strandmap[strand],
            ".",
            motif.id,
            seq[pos : pos + len(motif)],
        )

def scan_regionfile_to_table(
    input_table, genome, scoring, pfmfile=None, ncpus=None, zscore=True, gc=True
):
    """Scan regions in input table with motifs.

    Parameters
    ----------
    input_table : str
        Filename of input table. Can be either a text-separated tab file or a
        feather file.

    genome : str
        Genome name. Can be either the name of a FASTA-formatted file or a
        genomepy genome name.

    scoring : str
        "count" or "score"

    pfmfile : str, optional
        Specify a PFM file for scanning.

    ncpus : int, optional
        If defined this specifies the number of cores to use.

    Returns
    -------
    table : pandas.DataFrame
        DataFrame with motif ids as column names and regions as index. Values
        are either counts or scores depending on the 'scoring' parameter.s
    """
    config = MotifConfig()

    if pfmfile is None:
        pfmfile = config.get_default_params().get("motif_db", None)
        if pfmfile is not None:
            pfmfile = os.path.join(config.get_motif_dir(), pfmfile)

    if pfmfile is None:
        raise ValueError("no pfmfile given and no default database specified")

    logger.info("reading table")
    if input_table.endswith("feather"):
        df = pd.read_feather(input_table)
        idx = df.iloc[:, 0].values
    else:
        df = pd.read_table(input_table, index_col=0, comment="#")
        idx = df.index

    regions = list(idx)
    if len(regions) >= 1000:
        check_regions = np.random.choice(regions, size=1000, replace=False)
    else:
        check_regions = regions

    size = int(
        np.median([len(seq) for seq in as_fasta(check_regions, genome=genome).seqs])
    )
    s = Scanner(ncpus=ncpus)
    s.set_motifs(pfmfile)
    s.set_genome(genome)
    s.set_background(genome=genome, gc=gc, size=size)

    scores = []
    if scoring == "count":
        logger.info("setting threshold")
        s.set_threshold(fpr=FPR)
        logger.info("creating count table")
        for row in s.count(regions):
            scores.append(row)
        logger.info("done")
    else:
        s.set_threshold(threshold=0.0)
        msg = "creating score table"
        if zscore:
            msg += " (z-score"
            if gc:
                msg += ", GC%"
            msg += ")"
        else:
            msg += " (logodds)"
        logger.info(msg)
        for row in s.best_score(regions, zscore=zscore, gc=gc):
            scores.append(row)
        logger.info("done")

    motif_names = [m.id for m in read_motifs(pfmfile)]

    logger.info("creating dataframe")
    dtype = "float16"
    if scoring == "count":
        dtype = int
    df = pd.DataFrame(scores, index=idx, columns=motif_names, dtype=dtype)

    return df

def scan_table(
    s, inputfile, fa, motifs, cutoff, bgfile, nreport, scan_rc, pvalue, moods
):
    # header
    yield "\t{}".format("\t".join([m.id for m in motifs]))
    table = True
    if moods:
        result_it = scan_it_moods(
            inputfile, motifs, cutoff, bgfile, nreport, scan_rc, pvalue, table
        )
        for seq_id, counts in result_it:
            yield "{}\t{}".format(seq_id, "\t".join([str(x) for x in counts]))
    else:
        # get iterator
        result_it = s.count(fa, nreport, scan_rc)
        # counts table
        for i, counts in enumerate(result_it):
            yield "{}\t{}".format(fa.ids[i], "\t".join([str(x) for x in counts]))

def scan_score_table(s, fa, motifs, scan_rc, zscore=False, gcnorm=False):

    s.set_threshold(threshold=0.0, gc=gcnorm)
    # get iterator
    result_it = s.best_score(fa, scan_rc, zscore=zscore, gc=gcnorm)
    # header
    yield "\t{}".format("\t".join([m.id for m in motifs]))
    # score table
    for i, scores in enumerate(result_it):
        yield "{}\t{}".format(fa.ids[i], "\t".join(["{:4f}".format(x) for x in scores]))

def scan_normal(
    s,
    inputfile,
    fa,
    motifs,
    cutoff,
    bgfile,
    nreport,
    scan_rc,
    pvalue,
    moods,
    bed,
    zscore,
    gcnorm,
):

    table = False
    if moods:
        result_it = scan_it_moods(
            inputfile, motifs, cutoff, bgfile, nreport, scan_rc, pvalue, table
        )
        for motif, d in result_it:
            for seq_id, matches in d.items():
                for pos, score, strand in matches:
                    yield _format_line(
                        fa[seq_id], seq_id, motif, score, pos, strand, bed=bed
                    )
    else:
        result_it = s.scan(fa, nreport, scan_rc, zscore, gc=gcnorm)
        for i, result in enumerate(result_it):
            seq_id = fa.ids[i]
            seq = fa[seq_id]
            for motif, matches in zip(motifs, result):
                for (score, pos, strand) in matches:
                    yield _format_line(seq, seq_id, motif, score, pos, strand, bed=bed)

def command_scan(
    inputfile,
    pfmfile,
    nreport=1,
    fpr=0.01,
    cutoff=None,
    bed=False,
    scan_rc=True,
    table=False,
    score_table=False,
    moods=False,
    pvalue=None,
    bgfile=None,
    genome=None,
    ncpus=None,
    zscore=False,
    gcnorm=False,
):
    motifs = read_motifs(pfmfile)

    fa = as_fasta(inputfile, genome)

    # initialize scanner
    s = Scanner(ncpus=ncpus)
    s.set_motifs(pfmfile)

    if genome:
        s.set_genome(genome=genome)

    if genome:
        s.set_background(
            genome=genome, fname=bgfile, size=fa.median_length(), gc=gcnorm
        )
    if bgfile:
        s.set_background(genome=genome, fname=bgfile, size=fa.median_length())

    if not score_table:
        s.set_threshold(fpr=fpr, threshold=cutoff)

    if table:
        it = scan_table(
            s, inputfile, fa, motifs, cutoff, bgfile, nreport, scan_rc, pvalue, moods
        )
    elif score_table:
        it = scan_score_table(s, fa, motifs, scan_rc, zscore=zscore, gcnorm=gcnorm)
    else:
        it = scan_normal(
            s,
            inputfile,
            fa,
            motifs,
            cutoff,
            bgfile,
            nreport,
            scan_rc,
            pvalue,
            moods,
            bed,
            zscore=zscore,
            gcnorm=gcnorm,
        )

    for row in it:
        yield row

def scan_to_file(
    inputfile,
    pfmfile,
    filepath_or_buffer=None,
    nreport=1,
    fpr=0.01,
    cutoff=None,
    bed=False,
    scan_rc=True,
    table=False,
    score_table=False,
    moods=False,
    pvalue=False,
    bgfile=None,
    genome=None,
    ncpus=None,
    zscore=True,
    gcnorm=True,
):
    """Scan an inputfile with motifs."""
    should_close = False
    if filepath_or_buffer is None:
        fo = sys.stdout
    else:
        if hasattr(filepath_or_buffer, "write"):
            fo = filepath_or_buffer
        else:
            try:
                fo = open(os.path.expanduser(filepath_or_buffer), "w")
                should_close = True
            except Exception:
                logger.error(f"Could not open {filepath_or_buffer} for writing")
                sys.exit(1)

    if fpr is None and cutoff is None:
        fpr = 0.01

    print("# GimmeMotifs version {}".format(__version__), file=fo)
    print("# Input: {}".format(inputfile), file=fo)
    print("# Motifs: {}".format(pfmfile), file=fo)
    if fpr and not score_table:
        if genome is not None:
            print("# FPR: {} ({})".format(fpr, genome), file=fo)
        elif bgfile:
            print("# FPR: {} ({})".format(fpr, bgfile), file=fo)
    if cutoff is not None:
        print("# Threshold: {}".format(cutoff), file=fo)
    if zscore:
        if gcnorm:
            print("# Scoring: GC frequency normalized z-score", file=fo)
        else:
            print("# Scoring: normalized z-score", file=fo)
    else:
        print("# Scoring: logodds score", file=fo)

    for line in command_scan(
        inputfile,
        pfmfile,
        nreport=nreport,
        fpr=fpr,
        cutoff=cutoff,
        bed=bed,
        scan_rc=scan_rc,
        table=table,
        score_table=score_table,
        moods=moods,
        pvalue=pvalue,
        bgfile=bgfile,
        genome=genome,
        ncpus=ncpus,
        zscore=zscore,
        gcnorm=gcnorm,
    ):
        print(line, file=fo)

    if should_close:
        try:
            fo.close()
        except Exception:
            pass

def scan_to_best_match(
    fname, motifs, ncpus=None, genome=None, score=False, zscore=False, gc=False
):
    """Scan a FASTA file with motifs.

    Scan a FASTA file and return a dictionary with the best match per motif.

    Parameters
    ----------
    fname : str
        Filename of a sequence file in FASTA format.

    motifs : list
        List of motif instances.

    Returns
    -------
    result : dict
        Dictionary with motif scanning results.
    """
    # Initialize scanner
    s = Scanner(ncpus=ncpus)
    s.set_motifs(motifs)
    s.set_threshold(threshold=0.0)
    if genome:
        s.set_genome(genome)

    if isinstance(motifs, str):
        motifs = read_motifs(motifs)

    logger.debug("scanning %s...", fname)
    result = dict([(m.id, []) for m in motifs])
    if score:
        it = s.best_score(fname, zscore=zscore, gc=gc)
    else:
        it = s.best_match(fname, zscore=zscore, gc=gc)
    for scores in it:
        for motif, score in zip(motifs, scores):
            result[motif.id].append(score)

    # Close the pool and reclaim memory
    del s

    return result

def parse_threshold_values(motif_file, cutoff):
    motifs = read_motifs(motif_file)
    d = parse_cutoff(motifs, cutoff)
    threshold = {}
    for m in motifs:
        c = m.pwm_min_score() + (m.pwm_max_score() - m.pwm_min_score()) * d[m.id]
        threshold[m.id] = c
    return threshold

def scan_sequence(
    seq, seq_gc_bin, motifs, nreport, scan_rc, motifs_meanstd=None, zscore=False
):

    ret = []
    # scan for motifs
    for motif, cutoff in motifs:
        if cutoff is None:
            ret.append([])
        else:
            if zscore:
                m_mean, m_std = motifs_meanstd[seq_gc_bin][motif.id]
                result = pwmscan(
                    seq, motif.logodds, motif.pwm_min_score(), nreport, scan_rc
                )
                result = [[(row[0] - m_mean) / m_std, row[1], row[2]] for row in result]
                result = [row for row in result if row[0] >= cutoff]
            else:
                result = pwmscan(seq, motif.logodds, cutoff, nreport, scan_rc)
            if cutoff <= motif.pwm_min_score() and len(result) == 0:
                result = [[motif.pwm_min_score(), 0, 1]] * nreport

            ret.append(result)

    return ret

def scan_seq_mult(
    seqs, seq_gc_bins, motifs, nreport, scan_rc, motifs_meanstd=None, zscore=False
):
    ret = []
    for seq, seq_gc_bin in zip(seqs, seq_gc_bins):
        result = scan_sequence(
            seq.upper(),
            seq_gc_bin,
            motifs,
            nreport,
            scan_rc,
            motifs_meanstd=motifs_meanstd,
            zscore=zscore,
        )
        ret.append(result)
    return ret

def scan_fa_with_motif_moods(
    fo, motifs, matrices, bg, thresholds, nreport, scan_rc=True
):

    scanner = MOODS.scan.Scanner(7)
    scanner.set_motifs(matrices, bg, thresholds)

    ret = []
    for name, seq in fo.items():
        length = len(seq)

        scan_seq = seq.upper()
        if scan_rc:
            scan_seq = "".join((scan_seq, "N" * 50, rc(scan_seq)))
        results = scanner.scan_max_hits(scan_seq, nreport)
        for motif, result in zip(motifs, results):
            matches = []
            for match in result:
                strand = 1
                pos = match.pos
                if scan_rc:
                    if pos > length:
                        pos = length - (pos - length - 50) - len(motif)
                        strand = -1
                matches.append((pos, match.score, strand))
            ret.append((motif, {name: matches}))

    return ret

def scan_fa_with_motif_moods_count(
    fo, motifs, matrices, bg, thresholds, nreport, scan_rc=True
):
    scanner = MOODS.scan.Scanner(7)
    scanner.set_motifs(matrices, bg, thresholds)

    ret = []
    for name, seq in fo.items():
        scan_seq = seq.upper()
        if scan_rc:
            scan_seq = "".join((scan_seq, "N" * 50, rc(scan_seq)))
        results = scanner.counts_max_hits(scan_seq, nreport)
        ret.append((name, results))

    return ret

def calc_threshold_moods(m, c):
    m_min = MOODS.tools.min_score(m)
    m_max = MOODS.tools.max_score(m)

    return m_min + (m_max - m_min) * c

def scan_it_moods(
    infile, motifs, cutoff, bgfile, nreport=1, scan_rc=True, pvalue=None, count=False
):
    tmpdir = mkdtemp()
    matrices = []
    pseudocount = 1e-3
    # sys.stderr.write("bgfile: {}\n".format(bgfile))
    bg = MOODS.tools.bg_from_sequence_dna("".join(Fasta(bgfile).seqs), 1)

    for motif in motifs:
        pfmname = os.path.join(tmpdir, "{}.pfm".format(motif.id))
        with open(pfmname, "w") as f:
            matrix = np.array(motif.pwm).transpose()
            for line in [" ".join([str(x) for x in row]) for row in matrix]:
                f.write("{}\n".format(line))

        matrices.append(MOODS.parsers.pfm_log_odds(pfmname, bg, pseudocount))

    thresholds = []
    if pvalue is not None:
        thresholds = [
            MOODS.tools.threshold_from_p(m, bg, float(pvalue)) for m in matrices
        ]
        # sys.stderr.write("{}\n".format(thresholds))
    else:
        thresholds = [calc_threshold_moods(m, float(cutoff)) for m in matrices]

    scanner = MOODS.scan.Scanner(7)
    scanner.set_motifs(matrices, bg, thresholds)

    config = MotifConfig()
    ncpus = int(config.get_default_params()["ncpus"])
    fa = Fasta(infile)
    chunk = 500
    if (len(fa) / chunk) < ncpus:
        chunk = len(fa) / (ncpus + 1)

    jobs = []
    func = scan_fa_with_motif_moods
    if count:
        func = scan_fa_with_motif_moods_count

    pool = mp.Pool()
    for i in range(0, len(fa), chunk):
        jobs.append(
            pool.apply_async(
                func,
                (fa[i : i + chunk], motifs, matrices, bg, thresholds, nreport, scan_rc),
            )
        )

    for job in jobs:
        for ret in job.get():
            yield ret

class Scanner(object):
    """
    scan sequences with motifs
    """

    def __init__(self, ncpus=None):
        self.config = MotifConfig()
        self._threshold = None
        self.genome = None
        self.background = None
        self.meanstd = {}
        self.gc_bins = [(0, 1)]

        if ncpus is None:
            self.ncpus = int(MotifConfig().get_default_params()["ncpus"])
        else:
            self.ncpus = ncpus

        if self.ncpus > 1:
            # try:
            #    ctx = mp.get_context('spawn')
            #    self.pool = ctx.Pool(processes=self.ncpus)
            # except AttributeError:
            self.pool = mp.Pool(processes=self.ncpus)

        self.use_cache = False
        if self.config.get_default_params().get("use_cache", False):
            self._init_cache()

    def __del__(self):
        # Close the pool because of memory leak
        if hasattr(self, "pool"):
            self.pool.close()

    def _init_cache(self):
        try:
            self.cache = make_region().configure(
                "dogpile.cache.pylibmc",
                expiration_time=3600,
                arguments={"url": ["127.0.0.1"], "binary": True}
                #    'dogpile.cache.dbm',
                #    expiration_time = 3600,
                #    arguments = {
                #        'filename': 'cache.dbm'
                #    }
            )
            self.use_cache = True
        except Exception as e:
            sys.stderr.write("failed to initialize cache\n")
            sys.stderr.write("{}\n".format(e))

    def set_motifs(self, motifs):
        try:
            # Check if motifs is a list of Motif instances
            motifs[0].to_pwm()
            tmp = NamedTemporaryFile(mode="w", delete=False)
            for m in motifs:
                tmp.write("{}\n".format(m.to_pwm()))
            tmp.close()
            motif_file = tmp.name
        except AttributeError:
            motif_file = motifs

        self.motifs = motif_file
        self.motif_ids = [m.id for m in read_motifs(motif_file)]
        self.checksum = {}
        if self.use_cache:
            chksum = xxhash.xxh64("\n".join(sorted(self.motif_ids))).digest()
            self.checksum[self.motifs] = chksum

    def _meanstd_from_seqs(self, motifs, seqs):
        scan_motifs = [(m, m.pwm_min_score()) for m in motifs]

        table = []
        for x in self._scan_sequences_with_motif(scan_motifs, seqs, 1, True):
            table.append([row[0][0] for row in x])

        for (motif, _), scores in zip(scan_motifs, np.array(table).transpose()):
            yield motif, np.mean(scores), np.std(scores)  # cutoff

    def _threshold_from_seqs(self, motifs, seqs, fpr):
        scan_motifs = [(m, m.pwm_min_score()) for m in motifs]
        table = []
        seq_gc_bins = [self.get_seq_bin(seq) for seq in seqs]
        for gc_bin, result in zip(
            seq_gc_bins, self._scan_sequences_with_motif(scan_motifs, seqs, 1, True)
        ):
            table.append([gc_bin] + [row[0][0] for row in result])

        df = pd.DataFrame(table, columns=["gc_bin"] + [m.id for m in motifs])
        return df

    def set_meanstd(self, gc=False):
        if not self.background:
            self.set_background(gc=gc)

        self.meanstd = {}
        seqs = self.background.seqs
        if gc:
            seq_bins = [s.split(" ")[-1] for s in self.background.ids]
        else:
            seq_bins = ["0.00-1.00"] * len(seqs)
        if gc:
            bins = list(set(seq_bins))
        else:
            bins = ["0.00-1.00"]

        motifs = read_motifs(self.motifs)
        lock.acquire()

        try:
            with Cache(CACHE_DIR) as cache:
                scan_motifs = []
                for bin in bins:
                    if bin not in self.meanstd:
                        self.meanstd[bin] = {}
                    bin_seqs = [s for s, b in zip(seqs, seq_bins) if b == bin]

                    for motif in motifs:
                        k = "e{}|{}|{}".format(motif.hash(), self.background_hash, bin)

                        results = cache.get(k)
                        if results is None:
                            scan_motifs.append(motif)
                        else:
                            self.meanstd[bin][motif.id] = results

                    if len(scan_motifs) > 0:
                        logger.debug("Determining mean and stddev for motifs.")
                        for motif, mean, std in self._meanstd_from_seqs(
                            scan_motifs, bin_seqs
                        ):
                            k = "e{}|{}|{}".format(
                                motif.hash(), self.background_hash, bin
                            )
                            cache.set(k, [mean, std])
                            self.meanstd[bin][motif.id] = mean, std

                # Prevent std of 0
                # This should only happen in testing
                for motif in motifs:
                    stds = np.array(
                        [self.meanstd[gcbin][motif.id][1] for gcbin in bins]
                    )
                    idx = stds == 0
                    if True in idx:
                        std = np.mean(stds[~idx])
                        for gcbin in np.array(bins)[idx]:
                            k = "e{}|{}|{}".format(
                                motif.hash(), self.background_hash, gcbin
                            )
                            mean = self.meanstd[gcbin][motif.id][0]
                            cache.set(k, [mean, std])
                            self.meanstd[gcbin][motif.id] = mean, std
        except sqlite3.DatabaseError:
            print_cluster_error_message()
            sys.exit(1)
        lock.release()

        for gc_bin in self.gc_bins:
            gc_bin = "{:.2f}-{:.2f}".format(*gc_bin)
            if gc_bin not in self.meanstd:
                valid_bins = []
                for b in self.gc_bins:
                    bstr = "{:.2f}-{:.2f}".format(b[0], b[1])
                    if bstr in self.meanstd:
                        valid_bins.append(((b[0] + b[1]) / 2, bstr))

                v = float(gc_bin.split("-")[1])
                _, bstr = sorted(valid_bins, key=lambda x: abs(x[0] - v))[0]
                # logger.warn(f"Using {bstr}")
                self.meanstd[gc_bin] = self.meanstd[bstr]

    def set_background(
        self, fname=None, genome=None, size=200, nseq=None, gc=False, gc_bins=None
    ):
        """Set the background to use for FPR and z-score calculations.

        Background can be specified either as a genome name or as the
        name of a FASTA file.

        Parameters
        ----------
        fname : str, optional
            Name of FASTA file to use as background.

        genome : str, optional
            Name of genome to use to retrieve random sequences.

        size : int, optional
            Size of genomic sequences to retrieve. The default
            is 200.

        nseq : int, optional
            Number of genomic sequences to retrieve.
        """
        if self.background:
            return

        size = int(size)

        if gc_bins is None:
            if gc:
                gc_bins = [(0.0, 0.2), (0.8, 1)]
                for b in np.arange(0.2, 0.799, 0.05):
                    gc_bins.append((b, b + 0.05))
            else:
                gc_bins = [(0, 1)]
        if nseq is None:
            nseq = max(10000, len(gc_bins) * 1000)

        if genome and fname:
            logger.debug("using genome for background")
            fname = None

        if fname:
            if not os.path.exists(fname):
                raise IOError("Background file {} does not exist!".format(fname))

            self.background = Fasta(fname)
            self.background_hash = file_checksum(fname)
            return

        if not genome:
            if self.genome:
                genome = self.genome
            else:
                raise ValueError("Need either genome or filename for background.")

        logger.debug("using background: genome {} with size {}".format(genome, size))
        lock.acquire()
        try:
            with Cache(CACHE_DIR) as cache:
                self.background_hash = "d{}:{}:{}:{}".format(
                    genome, int(size), gc, str(gc_bins)
                )
                c = cache.get(self.background_hash)
                if c:
                    fa, gc_bins = c
                else:
                    fa = None

                if not fa:
                    if gc:
                        with NamedTemporaryFile() as tmp:
                            logger.info("using {} sequences".format(nseq))
                            gc_bin_bedfile(
                                tmp.name, genome, number=nseq, length=size, bins=gc_bins
                            )
                            fa = as_fasta(tmp.name, genome=genome)
                    else:
                        fa = RandomGenomicFasta(genome, size, nseq)
                    cache.set(self.background_hash, (fa, gc_bins))
        except sqlite3.DatabaseError:
            print_cluster_error_message()
            sys.exit(1)
        lock.release()

        self.background = fa
        if gc_bins:
            self.gc_bins = gc_bins

    @property
    def threshold(self):
        if self._threshold is None:
            self.set_threshold()
        return self._threshold

    def set_threshold(self, fpr=None, threshold=None, gc=False):
        """Set motif scanning threshold based on background sequences.

        Parameters
        ----------
        fpr : float, optional
            Desired FPR, between 0.0 and 1.0.

        threshold : float or str, optional
            Desired motif threshold, expressed as the fraction of the
            difference between minimum and maximum score of the PWM.
            Should either be a float between 0.0 and 1.0 or a filename
            with thresholds as created by 'gimme threshold'.

        """
        if threshold and fpr:
            raise ValueError("Need either fpr or threshold.")

        if threshold is None and fpr is None:
            if self.genome:
                fpr = 0.01
                logger.info(f"Using default FPR of {fpr}")
            else:
                threshold = 0.95
                logger.info(
                    f"Genome not specified, using default threshold of {threshold}."
                )
                logger.info("This is likely not ideal.")

        if fpr:
            fpr = float(fpr)
            if not (0.0 < fpr < 1.0):
                raise ValueError("Parameter fpr should be between 0 and 1")

        if not self.motifs:
            raise ValueError("please run set_motifs() first")

        motifs = read_motifs(self.motifs)
        gc_bins = ["{:.2f}-{:.2f}".format(*gc_bin) for gc_bin in self.gc_bins]

        if threshold is not None:
            d = parse_threshold_values(self.motifs, threshold)
            self._threshold = pd.DataFrame(d, index=[0])
            self._threshold = self._threshold.join(
                pd.DataFrame(gc_bins, index=[0] * len(gc_bins), columns=["gc_bin"])
            )
            self._threshold = self._threshold.set_index("gc_bin")
            return

        if not self.background:
            try:
                self.set_background(gc=gc)
            except Exception:
                raise ValueError("please run set_background() first")

        seqs = self.background.seqs

        lock.acquire()
        try:
            with Cache(CACHE_DIR) as cache:
                scan_motifs = []
                self._threshold = None
                for motif in motifs:
                    k = "{}|{}|{:.4f}|{}".format(
                        motif.hash(),
                        self.background_hash,
                        fpr,
                        ",".join(sorted(gc_bins)),
                    )
                    vals = cache.get(k)
                    if vals is None:
                        scan_motifs.append(motif)
                    else:
                        if self._threshold is None:
                            self._threshold = vals.to_frame()
                        else:
                            self._threshold[motif.id] = vals

                if len(scan_motifs) > 0:
                    logger.info("determining FPR-based threshold")
                    df = self._threshold_from_seqs(scan_motifs, seqs, fpr).set_index(
                        "gc_bin"
                    )
                    if self._threshold is None:
                        self._threshold = df
                    else:
                        self._threshold = pd.concat((self._threshold, df), axis=1)
                    for motif in scan_motifs:
                        k = "{}|{}|{:.4f}|{}".format(
                            motif.hash(),
                            self.background_hash,
                            fpr,
                            ",".join(sorted(gc_bins)),
                        )
                        cache.set(k, df[motif.id])
        except sqlite3.DatabaseError:
            print_cluster_error_message()
            sys.exit(1)
        lock.release()
        self.threshold_str = "{}_{}_{}_{}".format(
            fpr, threshold, self.background_hash, ",".join(sorted(gc_bins))
        )

    def set_genome(self, genome):
        """
        set the genome to be used for:
            - converting regions to sequences
            - background for MOODS
        """
        if not genome:
            return

        # raises error if checks fail
        Genome(genome)

        self.genome = genome

    def count(self, seqs, nreport=100, scan_rc=True):
        """
        count the number of matches above the cutoff
        returns an iterator of lists containing integer counts
        """
        for matches in self.scan(seqs, nreport, scan_rc):
            counts = [len(m) for m in matches]
            yield counts

    def total_count(self, seqs, nreport=100, scan_rc=True):
        """
        count the number of matches above the cutoff
        returns an iterator of lists containing integer counts
        """

        count_table = [counts for counts in self.count(seqs, nreport, scan_rc)]
        return np.sum(np.array(count_table), 0)

    def best_score(self, seqs, scan_rc=True, zscore=False, gc=False):
        """
        give the score of the best match of each motif in each sequence
        returns an iterator of lists containing floats
        """
        self.set_threshold(threshold=0.0, gc=gc)
        for matches in self.scan(seqs, 1, scan_rc, zscore=zscore, gc=gc):
            scores = np.array(
                [sorted(m, key=lambda x: x[0])[0][0] for m in matches if len(m) > 0]
            )
            yield scores

    def best_match(self, seqs, scan_rc=True, zscore=False, gc=False):
        """
        give the best match of each motif in each sequence
        returns an iterator of nested lists containing tuples:
        (score, position, strand)
        """
        self.set_threshold(threshold=0.0)
        for matches in self.scan(seqs, 1, scan_rc, zscore=zscore, gc=gc):
            yield [m[0] for m in matches]

    def get_seq_bin(self, seq):
        if len(str(seq)) == 0:
            gc = 0
        else:
            useq = seq.upper()
            gc = round((useq.count("G") + useq.count("C")) / len(useq), 2)
        if gc == 0:
            gc = 0.01
        for b_start, b_end in self.gc_bins:
            if gc > round(b_start, 2) and gc <= round(b_end, 2):
                return "{:.2f}-{:.2f}".format(b_start, b_end)

        logger.error(
            "Error determining seq: {}, bins: {}".format(seq, str(self.gc_bins))
        )
        raise ValueError()

    def get_motif_mean_std(self, gc_bin, motif):
        if gc_bin in self.meanstd:
            if motif not in self.meanstd[gc_bin]:
                raise ValueError("Motif mean and std not initialized")
        else:
            logger.warn(
                "GC% {} not present in genome, setting to closest GC% bin".format(
                    gc_bin
                )
            )
            # Ideally this should not happen. If we get here, this means that
            # a sequence has a GC% that does not occur in the genome. This
            # can happen during tests with small genome, or when a sequence
            # is scanned that does not originate from the genome.
            # We use the closest GC% bin from the genome instead.
            valid_bins = []
            for b in self.gc_bins:
                bstr = "{:.2f}-{:.2f}".format(b[0], b[1])
                if bstr in self.meanstd:
                    valid_bins.append(((b[0] + b[1]) / 2, bstr))

            v = float(gc_bin.split("-")[1])
            _, bstr = sorted(valid_bins, key=lambda x: abs(x[0] - v))[0]
            logger.warn(f"Using {bstr}")
            self.meanstd[gc_bin] = self.meanstd[bstr]
        return self.meanstd[gc_bin][motif]

    def scan(self, seqs, nreport=100, scan_rc=True, zscore=False, gc=False):
        """
        Scan a set of regions or sequences.
        """
        seqs = as_fasta(seqs, genome=self.genome)
        if zscore:
            if gc:
                if len(self.meanstd) <= 1:
                    self.set_meanstd(gc=gc)
            else:
                if len(self.meanstd) != 1:
                    self.set_meanstd(gc=gc)

        batch_size = 50000
        logger.debug("Scanning")
        for batch_idx in range(0, len(seqs), batch_size):
            it = self._scan_sequences(
                seqs.seqs[batch_idx : batch_idx + batch_size],
                nreport,
                scan_rc,
                zscore=zscore,
            )
            for result in it:
                yield result

    def get_gc_thresholds(self, seqs, motifs=None, zscore=False):
        # Simple case, only one threshold
        if np.all(self.threshold.nunique(axis=0) == 1):
            return self.threshold.iloc[0].to_dict()

        if motifs is None:
            motifs = read_motifs(self.motifs)
        seq_gc_bins = [self.get_seq_bin(seq) for seq in seqs]

        gc_bin_count = Counter(seq_gc_bins)

        _threshold = self.threshold
        if zscore:
            grouped = _threshold.groupby(_threshold.index).apply(scale, axis=0)
            _threshold = pd.DataFrame(
                np.vstack(grouped.values),
                index=_threshold.index,
                columns=_threshold.columns,
            )

        nseqs = int(20000 / np.sum(list(gc_bin_count.values())))
        t = {}
        maxt = pd.Series([m.pwm_max_score() for m in motifs], index=_threshold.columns)
        # We do this in a loop as the DataFrame will get too big to fit in memory
        # when the difference between the number of sequences per gc_bin is very
        # high.
        _threshold = _threshold.reset_index()
        idx = np.hstack(
            [
                _threshold[_threshold[_threshold.columns[0]] == gc_bin]
                .sample(nseqs * count, replace=True, random_state=42)
                .index.values
                for gc_bin, count in gc_bin_count.items()
            ]
        )
        for motif in _threshold.columns[1:]:
            val = _threshold.loc[idx, motif].quantile(0.99, interpolation="higher")
            if val < maxt.loc[motif]:
                t[motif] = val
            else:
                t[motif] = None
        return t

    def _scan_sequences_with_motif(self, motifs, seqs, nreport, scan_rc):
        scan_func = partial(
            scan_seq_mult, motifs=motifs, nreport=nreport, scan_rc=scan_rc
        )
        for ret in self._scan_jobs(scan_func, seqs):
            yield ret[1]

    def _scan_sequences(self, seqs, nreport, scan_rc, zscore=False):
        thresholds = self.get_gc_thresholds(seqs, zscore=zscore)
        motifs = [(m, thresholds[m.id]) for m in read_motifs(self.motifs)]
        motifs_meanstd = None
        if zscore:
            motifs_meanstd = self.meanstd

        scan_func = partial(
            scan_seq_mult,
            motifs=motifs,
            nreport=nreport,
            scan_rc=scan_rc,
            motifs_meanstd=motifs_meanstd,
            zscore=zscore,
        )
        for _, ret in self._scan_jobs(scan_func, seqs):
            yield ret

    def _scan_jobs(self, scan_func, scan_seqs):
        if self.ncpus > 1:
            median_len = np.median([len(x) for x in scan_seqs])
            chunksize = 200000 // int(median_len)  # 1000 seqs for len 200
            # prepare for parallel processing
            k = 0
            max_queue_size = 2 * self.ncpus
            jobs = []

            # loop over each job/chunk, and keep adding them to the queue
            for i in range(math.ceil(len(scan_seqs) / chunksize)):
                batch_seqs = scan_seqs[i * chunksize : (i + 1) * chunksize]
                seq_gc_bins = [self.get_seq_bin(seq) for seq in batch_seqs]
                job = self.pool.apply_async(scan_func, (batch_seqs, seq_gc_bins))
                jobs.append(job)

                # if our queue is full, wait until oldest job finishes
                while (len(jobs) >= max_queue_size) and not jobs[0].ready():
                    time.sleep(0.05)

                # resolve oldest job if finished
                if jobs[0].ready():
                    for ret in jobs[0].get():
                        region = scan_seqs[k]
                        k += 1
                        yield region, ret
                    jobs = jobs[1:]

            # cleanup the last jobs that did not get resolved in the for loop
            while len(jobs) > 0:
                for ret in jobs[0].get():
                    region = scan_seqs[k]
                    yield region, ret
                jobs = jobs[1:]
        else:
            # non-parallel job scanning
            batchsize = 1000
            for i in range((len(scan_seqs) - 1) // batchsize + 1):
                batch_seqs = scan_seqs[i * batchsize : (i + 1) * batchsize]
                seq_gc_bins = [self.get_seq_bin(seq) for seq in batch_seqs]

                for _j, ret in enumerate(scan_func(batch_seqs, seq_gc_bins)):
                    yield scan_seqs[i], ret