##############################################################################
# PyLipID: A python module for analysing protein-lipid interactions
#
# Author: Wanling Song
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
##############################################################################

"""This module contains the class Durations for using dual-cutoff scheme.
"""
import numpy as np

__all__ = ["cal_contact_residues", "Duration", "cal_occupancy", "cal_lipidcount"]

def cal_contact_residues(dist_matrix, cutoff):
    """Obtain contact residues as a function of time.

    This function takes a distance matrix that records the measured distance for molecules at each trajectory frame,
    then returns the indices of molecules the distance of which are smaller than the provided cutoff at each frame. It
    also returns the molecule indices the frame indices in which the molecule is within the cutoff.

    Parameters
    ----------
    dist_matrix : list or numpy.ndarray, shape=(n_residues, n_frames)
        The measured distance for molecules at each trajectory frame.

    cutoff : scalar
        The distance cutoff to define a contact. A distance to the target
        equal or lower to the ``cutoff`` is considered as in contact.

    Returns
    -------
    contact_list : list
        A list that records the indices of molecules that are within the given cutoff in each frame

    frame_id_set : list
        A list of frame indices for contacting molecules.

    residue_id_set : lsit
        A list of contacting molecules indices.

    Examples
    --------
    >>> dr0 = [0.9, 0.95, 1.2, 1.1, 1.0, 0.9] # the distances of R0 to the target as a function of time
    >>> dr1 = [0.95, 0.9, 0.95, 1.1, 1.2, 1.1] # the distances of R1
    >>> dr2 = [0.90, 0.90, 0.85, 0.95, 1.0, 1.1] # the distances of R2
    >>> dist_matrix = [dr0, dr1, dr2]
    >>> contact_list, frame_id_set, residue_id_set = cal_contact_residues(dist_matrix, 1.0)
    >>> print(contact_list)
    [[0, 1, 2], [0, 1, 2], [1, 2], [2], [0, 2], [0]]
    >>> print(frame_id_set)
    array([0, 1, 4, 5, 0, 1, 2, 0, 1, 2, 3, 4])
    >>> print(residue_id_set)
    array([0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2, 2])

    """
    residue_id_set, frame_id_set = np.where(np.array(dist_matrix) <= cutoff)
    contact_list = [[] for dummy in np.arange(len(dist_matrix[0]))]
    for frame_id, residue_id in zip(frame_id_set, residue_id_set):
        contact_list[frame_id].append(residue_id)
    return contact_list, frame_id_set, residue_id_set

class Duration:
    def __init__(self, contact_low, contact_high, dt):
        """Dual cutoff scheme for calculating the interaction durations.

        In the dual cutoff scheme, a continuous contact starts when a molecule moves closer than the lower distance cutoff
        and ends when the molecule moves out of the upper cutoff. The duration between these two time points is the
        duration of the contact.

        Here, the ``contact_low`` is the lipid index for the lower cutoff and ``contact_high`` is the lipid index
        for the upper cutoff. For calculation of contact durations, a lipid molecule that appears in the ``contact_low``
        is searched in the subsequent frames of the ``contact_high`` and the search then stops if this
        molecule disappears from the ``contact_high``. This lipid molecule is labeled as 'checked', and the duration of
        this contact is calculated from the number of frames in which this lipid molecule appears in the lipid indices.
        This calculation iterates until all lipid molecules in the lower lipid index are labeled as 'checked'.

        Parameters
        ----------
        contact_low : list
            A list that records the indices of lipid molecule within the lower distance cutoff at each trajectory frame.

        contact_high : list
            A list that records the indices of lipid molecule within the upper distance cutoff at each trajectory frame.

        dt : scalar
            The timestep between two adjacent trajectory frames.

        """
        self.contact_low = contact_low
        self.contact_high = contact_high
        self.dt = dt
        self.pointer = [np.zeros_like(self.contact_high[idx], dtype=np.int)
                        for idx in range(len(self.contact_high))]
        return

    def cal_durations(self):
        """Calculate interaction durations using the dual-cutoff scheme.

        Calculate the durations of the appearances that start from the point when a molecule appears
        in ``contact_low`` and ends when it disappears from ``contact_high``.

        Returns
        -------
        durations : list
            A list of durations of the contacts defined by ``contact_low`` and ``contact_high``.

        """

        durations = []
        for i in range(len(self.contact_low)):
            for j in range(len(self.contact_low[i])):
                pos = np.where(self.contact_high[i] == self.contact_low[i][j])[0][0]
                if self.pointer[i][pos] == 0:
                    durations.append(self._get_duration(i, pos))
        if len(durations) == 0:
            return [0]
        else:
            durations.sort()
            return durations

    def _get_duration(self, i, j):
        count = 1
        self.pointer[i][j] = 1
        lipid_to_search = self.contact_high[i][j]
        for k in range(i+1, len(self.contact_high)):
            locations = np.where(self.contact_high[k] == lipid_to_search)[0]
            if len(locations) == 0:
                return count * self.dt
            else:
                pos = locations[0]
                self.pointer[k][pos] = 1
                count +=1
        return (count - 1) * self.dt

def cal_occupancy(contact_list):
    """Calculate the percentage of frames in which a contact is formed.

    ``contact_list`` records a list of residue indices of contact lipid molecules at each trajectory frames. This function
    calculates the percentage of frames that a lipid contact is formed.

    Parameters
    ___________
    contact_list : list
        A list of residue indices of contact lipid molecules at each trajectory frames.

    Returns
    -------
    Ocupancy : scalar
        The percentage of frames in which a contact is formed

    Examples
    --------
    >>> contact_list = [[], [130], [130, 145], [145], [], [], [145], [145]] # contacts are formed in 5 out of the 8 frames
    >>> occupancy = cal_occupancy(contact_list)
    >>> print(occupancy) # percentage
    62.5

    See also
    --------
    pylipid.api.LipidInteraction.compute_residue_occupancy
        Calculate the percentage of frames in which the specified residue formed lipid contacts for residues.
    pylipid.api.LipidInteraction.compute_site_occupancy
        Calculate the percentage of frames in which the specified lipid contacts are formed for binding sites.

    """
    if len(contact_list) == 0:
        return 0
    else:
        contact_counts = [len(item) for item in contact_list]
        mask = np.array(contact_counts) > 0
        return 100 * np.sum(mask)/len(contact_list)

def cal_lipidcount(contact_list):
    """Calculate the average number of contacting molecules.

    This function calculates the average number of contacting molecules when any contact is formed.

    Parameters
    ___________
    contact_list : list
        A list of residue indices of contact lipid molecules at each trajectory frames.

    Returns
    -------
    LipidCount : scalar
        The average number of contacts in frames in which any contact is formed.

    Examples
    --------
    >>> contact_list = [[], [130], [130, 145], [145], [], [], [145], [145]]
    >>> lipidcount = cal_lipidcount(contact_list)
    >>> print(lipidcount) # (1+2+1+1+1)/5
    1.2

    See also
    --------
    pylipid.api.LipidInteraction.compute_residue_lipidcount
        Calculate the average number of contacting lipids for residues.
    pylipid.api.LipidInteraction.compute_site_lipidcount
        Calculate the average number of contacting lipids for binding sites.

    """
    if len(contact_list) == 0:
        return 0
    else:
        contact_counts = np.array([len(item) for item in contact_list])
        mask = contact_counts > 0
        if np.sum(mask) == 0:
            return 0
        else:
            contact_counts_nonzero = contact_counts[mask]
            return np.nan_to_num(contact_counts_nonzero.mean())


"""
Implement the qmix algorithm with tensorflow, also thanks to the pymarl repo.

https://github.com/oxwhirl/pymarl  with pytorch implemented.

"""

import os
import sys
from types import SimpleNamespace as SN

import numpy as np

from absl import logging
# from attrdict import AttrDict
from xt.model.tf_compat import tf
from xt.algorithm import Algorithm
from xt.algorithm.algorithm import ZFILL_LENGTH
from xt.algorithm.qmix.episode_buffer_np import EpisodeBatchNP, ReplayBufferNP
from xt.algorithm.qmix.transforms import OneHotNp
from zeus.common.util.register import Registers

class DecayThenFlatSchedule(object):
    """Schedule from pymarl."""

    def __init__(self, start, finish, time_length, decay="exp"):

        self.start = start
        self.finish = finish
        self.time_length = time_length
        self.delta = (self.start - self.finish) / self.time_length
        self.decay = decay

        if self.decay in ["exp"]:
            self.exp_scaling = ((-1) * self.time_length / np.log(self.finish)
                                if self.finish > 0 else 1)

    def eval(self, t):
        """Schedule with eval times."""
        val = 0
        if self.decay in ["linear"]:
            val = max(self.finish, self.start - self.delta * t)
        elif self.decay in ["exp"]:
            val = min(self.start, max(self.finish, np.exp(-t / self.exp_scaling)))
        else:
            raise KeyError("invalid decay-{} configured".format(self.decay))
        return val

class EpsilonGreedyActionSelector(object):
    """Create epsilon greedy action selector from pymarl."""

    def __init__(self, args):
        self.args = args

        self.schedule = DecayThenFlatSchedule(
            args["epsilon_start"],
            args["epsilon_finish"],
            args["epsilon_anneal_time"],
            decay="linear",
        )
        self.epsilon = self.schedule.eval(0)

    def select_action(self, agent_inputs, avail_actions, t_env, test_mode=False):
        """Assume agent_inputs is a batch of Q-Values for each agent bav."""
        self.epsilon = self.schedule.eval(t_env)

        if test_mode:
            # Greedy action selection only
            self.epsilon = 0.0

        # mask actions that are excluded from selection
        # masked_q_values = agent_inputs.clone()
        # we make the `agent_inputs` value as a numpy !!!
        masked_q_values = agent_inputs.copy()
        # should never be selected!
        # masked_q_values[avail_actions == 0.0] = -float("inf")
        masked_q_values[avail_actions < 1e-6] = -float("inf")

        # random_numbers = th.rand_like(agent_inputs[:, :, 0])
        random_numbers = np.random.rand(*agent_inputs[:, :, 0].shape)
        # pick_random = (random_numbers < self.epsilon).long()
        pick_random = np.array(random_numbers < self.epsilon).astype(np.long)
        _batch_size, agent_num, action_dim = avail_actions.shape
        # avail_actions = avail_actions.reshape(-1, last_dim)
        avail_norm_p = np.array(
            avail_actions / np.expand_dims(avail_actions.sum(-1), -1)).astype(np.float)
        avail_norm_p = avail_norm_p.reshape(-1, action_dim)

        random_actions = np.array(
            [np.random.choice(action_dim, p=pi)
             for pi in avail_norm_p]).reshape((-1, agent_num))

        picked_actions = pick_random * random_actions + \
            (1 - pick_random) * masked_q_values.argmax(axis=2)
        # print(picked_actions, picked_actions.shape)
        # # [[5. 5.]] (1, 2)
        return picked_actions

@Registers.algorithm
class QMixAlg(Algorithm):
    """Implemente q-mix algorithm with tensorflow framework."""

    def __init__(
            self,
            model_info: dict,
            alg_config: dict,
            **kwargs,
    ):
        """
        Initialize.

        Consider the compatibility between trainer and explorer,
        QMix Algorithm could support the two scene
        """
        # avail_actions vary with env.map
        # alg_config = AttrDict(alg_config)
        logging.debug("get alg_config: {}".format(alg_config))

        env_info = alg_config["env_attr"]
        alg_config.update({
            "n_agents": env_info["n_agents"],
            "n_actions": env_info["n_actions"],
            "state_shape": env_info["state_shape"],
        })

        self.n_agents = alg_config["n_agents"]  # from env

        self.scheme = {
            "state": {
                "vshape": env_info["state_shape"]
            },
            "obs": {
                "vshape": env_info["obs_shape"],
                "group": "agents"
            },
            "actions": {
                "vshape": (1, ),
                "group": "agents",
                "dtype": np.int64
            },
            "avail_actions": {
                "vshape": (env_info["n_actions"], ),
                "group": "agents",
                "dtype": np.int32,
            },
            "reward": {
                "vshape": (1, )
            },
            "terminated": {
                "vshape": (1, ),
                "dtype": np.uint8
            },
            "actions_onehot": {
                "vshape": (env_info["n_actions"], ),
                "dtype": np.float32,
                "group": "agents",
            },
            # "filled": {"vshape": (1,), "dtype": tf.int64},
        }
        self.obs_shape = self._get_input_shape(alg_config, self.scheme)
        logging.debug("obs_shape: {}".format(self.obs_shape))

        logging.debug("update obs shape: {} --> {}".format(
            model_info["actor"]["model_config"]["obs_shape"], self.obs_shape))
        model_info["actor"]["model_config"]["obs_shape"] = self.obs_shape

        # NOTE: set graph scene, train as default
        model_info["actor"].update({"scene": kwargs.get("scene", "train")})

        super(QMixAlg, self).__init__(alg_name="QMixAlg",
                                      model_info=model_info["actor"],
                                      alg_config=alg_config)

        self.async_flag = False

        self.previous_state = None
        self.ph_hidden_states_in = None
        self.hidden_states_out = None

        self.params = None
        self.inputs = None
        self.out_actions = None

        self.avail_action_num = env_info["n_actions"]

        # use the episode limit as fix shape.
        self.fix_seq_length = env_info["episode_limit"]

        self.schedule = DecayThenFlatSchedule(
            alg_config["epsilon_start"],
            alg_config["epsilon_finish"],
            alg_config["epsilon_anneal_time"],
            decay="linear",
        )
        self.epsilon = self.schedule.eval(0)

        # select action
        self.selector = EpsilonGreedyActionSelector(alg_config)

        # # mix
        # self.state_dim = int(np.prod(alg_config.state_shape))
        # self.embed_dim = alg_config.mixing_embed_dim

        # update target with period, while update explorer each times.
        self.last_target_update_episode = -9999.0

        self.groups = {"agents": env_info["n_agents"]}
        preprocess_np = {"actions": ("actions_onehot",
                                     [OneHotNp(out_dim=alg_config["n_actions"])])}

        self.buffer = ReplayBufferNP(
            self.scheme,
            self.groups,
            alg_config["buffer_size"],
            env_info["episode_limit"] + 1,
            preprocess=preprocess_np,
        )
        self.train_batch = None
        self.train_times = 0

    @staticmethod
    def _get_input_shape(alg_config, scheme):
        """Assemble input shape with alg_config, vary with last_action/agent_id."""
        input_shape = scheme["obs"]["vshape"]
        if alg_config["obs_last_action"]:
            input_shape += scheme["actions_onehot"]["vshape"][0]
        if alg_config["obs_agent_id"]:
            input_shape += alg_config["n_agents"]
        return input_shape

    def reset_hidden_state(self):
        self.actor.reset_hidden_state()

    def build_inputs(self, batch, t):
        """
        Build inputs.

        # Assumes homogenous agents with flat observations.
        # Other MACs might want to e.g. delegate building inputs to each agent
        1. inference stage, use batch = 1,
        2. train stage, use batch = episode.limit

        Also, use numpy for combine the inputs data
        """
        _batch_size = batch.batch_size
        inputs = list()
        inputs.append(batch["obs"][:, t])  # b1av

        if self.alg_config["obs_last_action"]:
            if t == 0:
                inputs.append(np.zeros_like(batch["actions_onehot"][:, t]))
            else:
                inputs.append(batch["actions_onehot"][:, t - 1])

        if self.alg_config["obs_agent_id"]:
            _ag_id = np.expand_dims(np.eye(self.n_agents), axis=0)  # add axis 0
            inputs.append(np.tile(_ag_id, (_batch_size, 1, 1)))  # broadcast_to

        inputs = np.expand_dims(np.concatenate(inputs, axis=-1), axis=1)

        return inputs

    def predict_with_selector(self, ep_batch, t_ep, t_env, test_mode):
        avail_actions = ep_batch["avail_actions"][:, t_ep]
        agent_inputs = self.build_inputs(ep_batch, t_ep)

        out_val = self.actor.infer_actions(agent_inputs)

        select_actions = self.selector.select_action(
            out_val, avail_actions, t_env, test_mode=test_mode)
        # print("out_val: {}, select action: {}, avail_actions, {}, t_env:{}".format(
        #     out_val, select_actions, avail_actions, t_env))
        return select_actions

    def save(self, model_path, model_index):
        """Save qmix explore agent weight with saver."""
        model_name = os.path.join(
            model_path, "actor{}".format(str(model_index).zfill(ZFILL_LENGTH)))
        self.actor.save_explore_agent_weights(model_name)
        return [model_name]

    def restore(self, model_name=None, model_weights=None):
        """
        Restore the model with the priority: model_weight > model_name.
        owing to actor.set_weights would be faster than load model from disk.

        if user used multi model in one algorithm,
        they need overwrite this function.
        And, Caller make the name/weights valid.
        """
        if model_weights is not None:
            self.actor.set_weights(model_weights)
        else:
            logging.debug("{} try load model: {}".format(self.alg_name, model_name))
            self.actor.restore_explorer_variable(model_name)

    @staticmethod
    def _new_data_sn():
        new_data = SN()
        new_data.transition_data = {}
        new_data.episode_data = {}
        return new_data

    def prepare_data(self, train_data, **kwargs):
        """Insert trajectory into buffer, and sample batch if meet required."""
        new_data = self._new_data_sn()
        # received train data batch size ==1, seq_length==limit episode
        for k, val in train_data.items():
            new_data.transition_data[k] = val

        deliver_batch = EpisodeBatchNP(
            self.scheme, self.groups, 1, self.fix_seq_length + 1, data=new_data)
        self.buffer.insert_episode_batch(deliver_batch)

        if self.buffer.can_sample(self.alg_config["batch_size"]):
            self.train_batch = self.buffer.sample(self.alg_config["batch_size"])
        else:
            self.train_batch = None

    def train(self, **kwargs):
        """Train with buffer sampled."""
        # , batch: EpisodeBatchNP, t_env: int, episode_num: int
        if not self.train_batch:
            return np.nan

        episode_num = kwargs.get("episode_num")
        if not episode_num:
            raise KeyError("need episode num to update target network")

        batch = self.train_batch

        # Truncate batch to only filled timesteps
        max_ep_t = batch.max_t_filled()
        logging.debug("episode sample with max_ep_t: {}".format(max_ep_t))
        # batch = batch[:, :max_ep_t]

        # Get the relevant quantities
        rewards = batch["reward"][:, :-1]
        actions = batch["actions"][:, :-1]
        terminated = batch["terminated"][:, :-1].astype(np.float32)
        mask = batch["filled"][:, :-1].astype(np.float32)
        mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1])
        avail_actions = batch["avail_actions"]

        # [bs, seq_len, n_agents, obs_size] [32, 1, 2, 26] --> [32, 301, 2, 26]
        _inputs = [self.build_inputs(batch, t) for t in range(batch.max_seq_length)]

        batch_trajectories = np.concatenate(_inputs, axis=1)

        logging.debug("batch_trajectories.shape: {}".format(batch_trajectories.shape))
        logging.debug("rewards.shape: {}".format(rewards.shape))
        logging.debug("actions.shape: {}".format(actions.shape))
        logging.debug("terminated.shape: {}".format(terminated.shape))
        logging.debug("mask.shape: {}".format(mask.shape))
        logging.debug("avail_actions.shape: {}".format(avail_actions.shape))
        logging.debug("batch.max_seq_length: {}".format(batch.max_seq_length))
        logging.debug("batch.batch_size: {}".format(batch.batch_size))

        # to get action --> [32, 300, 2, 7]
        # [32*301*2, 26] --> [32*301*2, 7] --> [32, 301, 2, 7] --> [32, 300, 2, 7]
        # batch4train = batch_trajectories.reshape([-1, batch_trajectories.shape[-1]])

        self.train_times += 1
        loss_val = self.actor.train(
            batch_trajectories,
            list([max_ep_t for _ in range(batch.batch_size * self.n_agents)]),
            avail_actions,
            actions,
            batch["state"][:, :-1],
            batch["state"][:, 1:],
            rewards,
            terminated,
            mask,
        )

        # update explore agent after each train processing.
        self.actor.assign_explore_agent()

        # update target network as required.
        if (episode_num - self.last_target_update_episode) / self.alg_config["target_update_interval"] >= 1.0:
            self.actor.assign_targets()
            logging.debug("episode-{}, target Q network params replaced!".format(episode_num))
            print("episode-{}, target Q network params replaced!".format(episode_num))
            print(">>> train-{} use seq-len-{}".format(self.train_times, max_ep_t))
            sys.stdout.flush()
            self.last_target_update_episode = episode_num

        return loss_val

    def train_ready(self, elapsed_episode, **kwargs):
        """
        Support custom train logic.

        :return: train ready flag
        """
        # we set train ready as default
        # if elapsed_episode < self.alg_config["batch_size"]:
        if not self.buffer.can_sample(self.alg_config["batch_size"]):
            self._train_ready = False
            if not kwargs.get("dist_dummy_model"):
                raise KeyError("qmix need to dist dummy model.")
            # dist dummy model
            kwargs["dist_dummy_model"]()
        else:
            self._train_ready = True

        return self._train_ready


import os, time
from collections import deque
from typing import Sequence

import numpy as np

from . import agent
from . import spaces
from .bullet_env import BulletEnv
from .utils import quatmultiply, orient2heading, quat2mat, euler_zyx2quat, axis_angle2quat

from .reference_motion import MOTION_DIR, ReferenceMotion

class Observation(object):
    def __init__(self, history=-1, with_velocity=True):
        self.history = history
        self.with_velocity = with_velocity

class MotionImitate(BulletEnv):
    AGENT = agent.Humanoid

    def __init__(self, action: Sequence[str], observation: (Observation or dict[str, Observation]), **kwargs):
        def arg_parse(name, def_val):
            return kwargs[name] if kwargs is not None and name in kwargs else def_val

        self.fps = arg_parse("fps", 30.0)
        self.frame_skip = arg_parse("frame_skip", 20)
        kwargs["time_step"] = arg_parse("time_step", 1.0/(self.fps*self.frame_skip))
        super().__init__(**kwargs)
        self.overtime = arg_parse("overtime", 500)
        self.random_init_pose = arg_parse("random_init_pose", True)
        self.init_pose_with_noise = arg_parse("init_pose_with_noise", False)
        self.observation = observation

        self.agent = self.AGENT(self)
        self.ref_motion = ReferenceMotion(self.agent, action, motion_dir=arg_parse("motion_dir", MOTION_DIR))
        self.ground = None

        self.init_sim_state = None
        self.rng = np.random.RandomState()
        self.init()

    def init(self):
        self.simulation_steps = 0
        self.elapsed_time = 0
        self.init_time = 0
        if self.init_sim_state is not None:
            self.remove_state(self.init_sim_state)
            self.init_sim_state = None
        super().connect()

        self.agent.init()
        self.up_dir = self.agent.UP_AXIS
        self.ground = self.build_ground()
        g = [0, 0, 0]
        g[self.up_dir] = -9.8
        self.gravity = g
        if self.up_dir == 1:
            self.configure_debug_visualizer(self.COV_ENABLE_Y_AXIS_UP,1)
        elif self.up_dir == 0:
            raise ValueError("Unsupported to use X-axis as the up axis.")

        self.init_sim_state = self.save_state()
        if self.ref_motion is not None:
            self.ref_motion.init()
            self.init_action_space()
            self.init_observation_space()

    def build_ground(self):
        import pybullet_data
        self.setAdditionalSearchPath(pybullet_data.getDataPath())
        if self.up_dir == 2:
            orient = [0, 0, 0, 1]
        elif self.up_dir == 1:
            orient = euler_zyx2quat(0, 0, -np.pi * 0.5)
        else:
            assert(self.up_dir in [1, 2])
        ground = self.load_urdf("plane_implicit.urdf", [0, 0, 0], orient,
            useMaximalCoordinates=True)
        self.change_dynamics(ground, -1, lateralFriction=0.9)
        return ground

    def init_observation_space(self):
        obs = self.observe(self.agent)
        if isinstance(obs, dict):
            self.observation_space = {}
            for k, ob in obs.items():
                self.observation_space[k] = spaces.Box()
                self.observation_space[k].shape = np.array(ob).shape
        else:
            self.observation_space = spaces.Box()
            self.observation_space.shape = np.array(obs).shape

    def init_action_space(self):
        self.action_mean, self.action_std = [], []
        lower_bound, upper_bound = [], []
        control_range = 2 # control range for the position controller, 2 means 2 times of the oringal movement range
        for jid, lim in zip(self.agent.motors, self.agent.movement_lim):
            info = self.agent.joint_info(jid)
            joint_type = info[2]
            if joint_type == self.JOINT_REVOLUTE:
                self.action_mean.append(0.5*(lim[1] + lim[0]))
                self.action_std.append((lim[1] - lim[0])*control_range)
                lower_bound.append(-1.0)
                upper_bound.append(1.0)
            elif joint_type == self.JOINT_SPHERICAL:
                y_offset, z_offset = 0, 0.2 # self.up_dir == 1
                if self.up_dir == 2:
                    y_offset, z_offset = -z_offset, y_offset
                self.action_mean.extend([0, y_offset, z_offset, 0])  # in order of axis, angle
                self.action_std.extend([1, 1, 1, (lim[1] - lim[0])*control_range])
                lower_bound.extend([-1.0, -1.0-y_offset, -1.0-z_offset, -1.0])
                upper_bound.extend([1.0, 1.0-y_offset, 1.0-z_offset, 1.0])
            else:
                assert(joint_type in [self.JOINT_REVOLUTE, self.JOINT_SPHERICAL])
        self.action_space = spaces.Box()
        self.action_space.shape = [len(lower_bound)]
        self.action_space.low, self.action_space.high = lower_bound, upper_bound

    def pre_process_action(self, action):
        action = np.add(self.action_mean, np.multiply(action, self.action_std))
        i = 0
        for jid in self.agent.motors:
            joint_type = self.agent.joint_info(jid)[2]
            if joint_type == self.JOINT_REVOLUTE:
                i += 1
            else: # joint_type == self.JOINT_SPHERICAL
                q = axis_angle2quat((action[i+0], action[i+1], action[i+2]), action[i+3])
                action[i+0], action[i+1], action[i+2], action[i+3] = q
                i += 4
        return action

    def seed(self, s):
        self.rng.seed(s)

    def reset(self):
        self.restore_state(self.init_sim_state)
        self.simulation_steps = 0

        if len(self.ref_motion) > 1:
            target_ref_motion = self.rng.randint(len(self.ref_motion.motions))
            self.ref_motion.set_motion(target_ref_motion)
        self.ref_motion.reset()
        self.contactable_links = None if self.ref_motion.contactable_links is None else [
            self.agent.links[l] for l in self.ref_motion.contactable_links
        ]

        if self.random_init_pose and self.ref_motion.random_init_pose:
            phase = self.rng.rand()
        else:
            phase = 0.0
        if self.ref_motion.loopable:
            self.init_time = phase * self.ref_motion.duration
        else:
            self.init_time = phase * max(0, self.ref_motion.duration - 100/self.fps)

        self.elapsed_time = self.init_time
        ref_pose = self.ref_motion.set_sim_time(self.elapsed_time, add_noise=self.init_pose_with_noise)
        self.agent.pose = ref_pose

        dist = np.inf
        for lid in range(self.agent.n_links):
            aabb_min, aabb_max = self.agent.aabb(lid)
            dist = min(dist, aabb_min[self.up_dir])
        dist -= 0.001
        if dist < 0:
            ref_pose["base_position"][self.up_dir] -= dist
            self.agent.base_position_and_orientation = ref_pose["base_position"], ref_pose["base_orientation"]
            self.ref_motion.sync(self.elapsed_time, ref_pose["base_position"], None)
            self.ref_motion.set_sim_time(self.elapsed_time)

        state = self.observe(self.agent)
        self.info["expert_state"] = self.observe(self.ref_motion.agent)
        if self._render: self._render_delay = time.time()
        return state

    def step(self, action):
        action = self.pre_process_action(action)
        assert(not np.any(np.isnan(action)))

        self.simulation_steps += 1
        phase = self.ref_motion.phase(self.elapsed_time)

        for frame in range(self.frame_skip):
            self.elapsed_time += self.time_step

            self.agent.target_position = action
            self.do_simulation()

            if self._render:
                t = time.time()
                time.sleep(max(0, self.time_step - (t-self._render_delay)/1000))
                pos, orient = self.agent.base_position_and_orientation
                self.reset_debug_visualizer_camera(2, 180, 0, [pos[0], 0.84, pos[2]])
                self._render_delay = t

        if self.ref_motion.loopable and self.ref_motion.phase(self.elapsed_time) < phase:
            self.ref_motion.accumulate_cycle_offset()
        self.ref_motion.set_sim_time(self.elapsed_time)

        if self.contactable_links is None:
            fail = False
        else:
            fail = self.agent.has_contact(self.ground, exclusive_links=self.contactable_links)
        terminal = fail
        if not terminal and self.ref_motion and not self.ref_motion.loopable:
            terminal = self.elapsed_time >= self.ref_motion.duration
        if not terminal and self.overtime:
            self.info["TimeLimit.truncated"] = self.simulation_steps >= self.overtime
            terminal = self.info["TimeLimit.truncated"]
        reward = self.reward(fail)
        state = self.observe(self.agent)
        self.info["expert_state"] = self.observe(self.ref_motion.agent)

        return state, reward, terminal, self.info

    def observe(self, agent):
        base_pos, base_orient = agent.base_position_and_orientation
        base_height = base_pos[self.up_dir]

        up_dir_vec = [0, 0, 0]
        up_dir_vec[self.up_dir] = 1
        root_p_inv = [-p for p in base_pos]
        heading = orient2heading(base_orient, self.up_dir)
        root_q_inv = axis_angle2quat(up_dir_vec, -heading)
        root_R_inv = quat2mat(root_q_inv)

        obs = {}
        for key, opt in self.observation.items():
            with_velocity = opt.with_velocity
            if not hasattr(agent, "state_history"):
                agent.state_history = {}
            if key not in agent.state_history:
                agent.state_history[key] = deque(maxlen=opt.history)
            buffer = agent.state_history[key]
            if self.simulation_steps == 0:
                buffer.clear()

            joints = list(range(agent.n_links))
            if with_velocity:
                full_states = [
                    (s[0], s[1], s[6], s[7]) # p, q, v, w
                    for s in agent.link_state(joints,
                    compute_forward_kinematics=True, compute_link_velocity=True)
                ]
                buffer.append(full_states)
                s = []
                for _ in buffer:
                    for p, q, v, w in _:
                        p = np.matmul(root_R_inv, np.add(p, root_p_inv))
                        q = quatmultiply(root_q_inv, q)
                        if q[3] < 0: q = [-e for e in q]
                        s.extend(p)
                        s.extend(q)
                        v = np.matmul(root_R_inv, v)
                        w = np.matmul(root_R_inv, w)
                        s.extend(v)
                        s.extend(w)
            else:
                full_states = [
                    (s[0], s[1]) # p, q
                    for s in agent.link_state(joints,
                    compute_forward_kinematics=True)
                ]
                buffer.append(full_states)
                s = []
                for _ in buffer:
                    for p, q in _:
                        p = np.matmul(root_R_inv, np.add(p, root_p_inv))
                        q = quatmultiply(root_q_inv, q)
                        if q[3] < 0: q = [-e for e in q]
                        s.extend(p)
                        s.extend(q)

            s = np.array(s, dtype=np.float32)
            s.shape = (len(buffer), -1)
            obs[key] = s
        return next(iter(obs.values())) if len(obs) == 1 else obs

    def reward(self, terminal):
        if terminal:
            reward = 0
        else:
            reward = 1
        return reward


import os
import glob
import tqdm
import torch
import random
import librosa
import argparse
import numpy as np
from multiprocessing import Pool, cpu_count

from utils.audio import Audio
from utils.hparams import HParam

#python generator.py -c config/config.yaml -d ../datasets/raw_libri/LibriSpeech/ -o training_prepared --no_need_spec (for convtasnet)
#python generator.py -c config/config.yaml -d ../datasets/raw_libri/librispeech -o training_libri_mel (for voicefilter)
#python generator.py -c config/config.yaml -d ../datasets/raw_libri/librispeech -o test_libri_mel --need_phase (for eval voicefilter)

generate_new_field_for_prepared=False
generate_own_data=False

def formatter(dir_, form, num):
    return os.path.join(dir_, form.replace('*', '%06d' % num))

def vad_merge(w):
    intervals = librosa.effects.split(w, top_db=20)
    temp = list()
    for s, e in intervals:
        temp.append(w[s:e])
    return np.concatenate(temp, axis=None)

def mix_wrapper(hp, args, audio, num, s1_dvec, s1_target, s2_target, s2_dvec, us_dvec, train):
    try:
        mix(hp, args, audio, num, s1_dvec, s1_target, s2_target, s2_dvec, us_dvec, train)
    except Exception as e:
        print(s1_dvec, s1_target, s2_target, s2_dvec, us_dvec)
        print(e)

def get_text_from_wav(wav_path):
    parent=os.path.dirname(wav_path)
    wav_id=os.path.splitext(os.path.basename(wav_path))[0][:-5]
    clip_id='-'.join(wav_id.split('-')[:2])
    trans_txt="%s/%s.trans.txt" % (parent, clip_id)
    with open(trans_txt) as f:
        for line in f:
            line=line.strip()
            fields=line.split(' ', 1)
            if fields[0]==wav_id:
                return fields[1].lower()
    assert False, "there's no transcript for the file %s" % wav_path

#def mix(hp, args, audio, num, s1_dvec, s1_target, s2, train):
#def mix(hp, args, audio, num, s1_dvec, s1_target, s2_target, s2_dvec, train):
def mix(hp, args, audio, num, s1_dvec, s1_target, s2_target, s2_dvec, us_dvec, train):
    srate = hp.audio.sample_rate
    dir_ = os.path.join(args.out_dir, 'train' if train else 'test')

    d, _ = librosa.load(s1_dvec, sr=srate)
    w1, _ = librosa.load(s1_target, sr=srate)
    w2, _ = librosa.load(s2_target, sr=srate)
    d2, _ = librosa.load(s2_dvec, sr=srate)
    d3, _ = librosa.load(us_dvec, sr=srate)
    assert len(d.shape) == len(w1.shape) == len(w2.shape) == len(d2.shape) == len(d3.shape) == 1, \
        'wav files must be mono, not stereo'

    d, _ = librosa.effects.trim(d, top_db=20)
    w1, _ = librosa.effects.trim(w1, top_db=20)
    w2, _ = librosa.effects.trim(w2, top_db=20)
    d2, _ = librosa.effects.trim(d2, top_db=20)
    d3, _ = librosa.effects.trim(d3, top_db=20)

    # run once for test data
    if True and not train:
        candidates=glob.glob(os.path.join(os.path.dirname(us_dvec), '**', hp.form.input), recursive=True)
        ref_utt=''
        while d3.shape[0] < 1.1 * hp.embedder.window * hp.audio.hop_length:
            print("Too short dvec3 wav")
            ref_utt=random.choice(candidates)
            d3, _ = librosa.load(ref_utt, sr=srate)
            d3, _ = librosa.effects.trim(d, top_db=20)
        if ref_utt != '':
            print('changed dvec', num, ref_utt)
    if train:
        # if reference for d-vector is too short, discard it
        if d.shape[0] < 1.1 * hp.embedder.window * hp.audio.hop_length:
            print("Too short dvec wav")
            return
        # if reference for d-vector is too short, discard it
        if d2.shape[0] < 1.1 * hp.embedder.window * hp.audio.hop_length:
            print("Too short dvec2 wav")
            return
        #TODO: for future we assert d3 length.
        candidates=glob.glob(os.path.join(os.path.dirname(us_dvec), '**', hp.form.input), recursive=True)
        ref_utt=''
        while d3.shape[0] < 1.1 * hp.embedder.window * hp.audio.hop_length:
            print("Too short dvec3 wav")
            ref_utt=random.choice(candidates)
            d3, _ = librosa.load(ref_utt, sr=srate)
            d3, _ = librosa.effects.trim(d3, top_db=20)
        if ref_utt != '':
            print('changed dvec3', num, ref_utt)

    # LibriSpeech dataset have many silent interval, so let's vad-merge them
    # VoiceFilter paper didn't do that. To test SDR in same way, don't vad-merge.
    if args.vad == 1:
        w1, w2 = vad_merge(w1), vad_merge(w2)

    if train:
        # I think random segment length will be better, but let's follow the paper first
        # fit audio to `hp.data.audio_len` seconds.
        # if merged audio is shorter than `L`, discard it
        L = int(srate * hp.data.audio_len)
        if w1.shape[0] < L or w2.shape[0] < L:
            return
        w1, w2 = w1[:L], w2[:L]
    else:
        L = max(len(w1), len(w2))
        w1 = np.pad(w1, (0, L-len(w1)), 'constant', constant_values=(0.0, 0.0))
        w2 = np.pad(w2, (0, L-len(w2)), 'constant', constant_values=(0.0, 0.0))

    mixed = w1 + w2

    norm = np.max(np.abs(mixed)) * 1.1
    w1, w2, mixed = w1/norm, w2/norm, mixed/norm
    assert len(mixed) == len(w1) == len(w2)

    # save vad & normalized wav files
    target_wav_path = formatter(dir_, hp.form.target.wav, num)
    mixed_wav_path = formatter(dir_, hp.form.mixed.wav, num)
    librosa.output.write_wav(target_wav_path, w1, srate)
    librosa.output.write_wav(mixed_wav_path, mixed, srate)
    if True:
        dvec_wav_path = formatter(dir_, hp.form.dvec_wav, num)
        librosa.output.write_wav(dvec_wav_path, d, srate)
    flag=True #train
    if True:
        dvec3_wav_path = formatter(dir_, hp.form.dvec3_wav, num)
        librosa.output.write_wav(dvec3_wav_path, d3, srate)
    if flag:
        target2_wav_path = formatter(dir_, hp.form.target2.wav, num)
        librosa.output.write_wav(target2_wav_path, w2, srate)
    if flag:
        dvec2_wav_path = formatter(dir_, hp.form.dvec2_wav, num)
        librosa.output.write_wav(dvec2_wav_path, d2, srate)
    if flag:
        dvec3_wav_path = formatter(dir_, hp.form.dvec3_wav, num)
        librosa.output.write_wav(dvec3_wav_path, d3, srate)

    if not args.no_need_spec:
        # save magnitude spectrograms
        target_mag, target_phase = audio.wav2spec(w1)
        mixed_mag, mixed_phase = audio.wav2spec(mixed)
        target_mag_path = formatter(dir_, hp.form.target.mag, num)
        mixed_mag_path = formatter(dir_, hp.form.mixed.mag, num)
        torch.save(torch.from_numpy(target_mag), target_mag_path)
        torch.save(torch.from_numpy(mixed_mag), mixed_mag_path)

    if args.need_mel:
        # save mel spectrograms
        target_mel = audio.wav2mel(w1)
        mixed_mel = audio.wav2mel(mixed)
        target_mel_path = formatter(dir_, hp.form.target.mel, num)
        mixed_mel_path = formatter(dir_, hp.form.mixed.mel, num)
        torch.save(torch.from_numpy(target_mel), target_mel_path)
        torch.save(torch.from_numpy(mixed_mel), mixed_mel_path)

    if args.need_phase:
        target_phase_path = formatter(dir_, hp.form.target.phase, num)
        mixed_phase_path = formatter(dir_, hp.form.mixed.phase, num)
        torch.save(torch.from_numpy(target_phase), target_phase_path)
        torch.save(torch.from_numpy(mixed_phase), mixed_phase_path)

    s1_text_path = formatter(dir_, '*-trans.txt', num)
    with open(s1_text_path, 'w') as f:
        f.write(get_text_from_wav(s1_target))

    # save selected sample as text file. d-vec will be calculated soon
    dvec_text_path = formatter(dir_, hp.form.dvec, num)
    with open(dvec_text_path, 'w') as f:
        f.write(s1_dvec)
    #assert os.path.exists(dvec_wav_path) and os.path.exists(dvec2_wav_path) \
    #        and os.path.exists(mixed_wav_path) \
    #        and os.path.exists(target_wav_path) and os.path.exists(mixed_mag_path) \
    #        and os.path.exists(target_mag_path) and os.path.exists(dvec_text_path) \
    #        and os.path.exists(target2_wav_path)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', type=str, required=True,
                        help="yaml file for configuration")
    parser.add_argument('-d', '--libri_dir', type=str, default=None,
                        help="Directory of LibriSpeech dataset, containing folders of train-clean-100, train-clean-360, dev-clean.")
    parser.add_argument('-v', '--voxceleb_dir', type=str, default=None,
                        help="Directory of VoxCeleb2 dataset, ends with 'aac'")
    parser.add_argument('-o', '--out_dir', type=str, required=True,
                        help="Directory of output training triplet")
    parser.add_argument('-p', '--process_num', type=int, default=None,
                        help='number of processes to run. default: cpu_count')
    parser.add_argument('--vad', type=int, default=0,
                        help='apply vad to wav file. yes(1) or no(0, default)')
    parser.add_argument('--need_phase', action='store_true',
                        help='apply vad to wav file. yes(1) or no(0, default)')
    parser.add_argument('--need_mel', action='store_true',
                        help='apply vad to wav file. yes(1) or no(0, default)')
    parser.add_argument('--no_need_spec', action='store_true',
                        help='apply vad to wav file. yes(1) or no(0, default)')
    parser.add_argument('--train_csv', type=str, default='datasets/train_tuples_clean_all.csv',
                        help="train csv")
    parser.add_argument('--test_csv', type=str, default='datasets/dev_tuples_clean.csv',
                        help="test csv")
    args = parser.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    os.makedirs(os.path.join(args.out_dir, 'train'), exist_ok=True)
    os.makedirs(os.path.join(args.out_dir, 'test'), exist_ok=True)

    hp = HParam(args.config)

    audio = Audio(hp)

    cpu_num = cpu_count() if args.process_num is None else args.process_num

    if generate_own_data:
        start_index=70000
        total_num=50000
        if args.libri_dir is None and args.voxceleb_dir is None:
            raise Exception("Please provide directory of data")

        if args.libri_dir is not None:
            train_folders = [x for x in glob.glob(os.path.join(args.libri_dir, 'train-clean-100', '*'))
                                if os.path.isdir(x)] + \
                            [x for x in glob.glob(os.path.join(args.libri_dir, 'train-clean-360', '*'))
                                if os.path.isdir(x)]
                            # we recommned to exclude train-other-500
                            # See https://github.com/mindslab-ai/voicefilter/issues/5#issuecomment-497746793
                            # + \
                            #[x for x in glob.glob(os.path.join(args.libri_dir, 'train-other-500', '*'))
                            #    if os.path.isdir(x)]
            test_folders = [x for x in glob.glob(os.path.join(args.libri_dir, 'dev-clean', '*'))]

        elif args.voxceleb_dir is not None:
            all_folders = [x for x in glob.glob(os.path.join(args.voxceleb_dir, '*'))
                                if os.path.isdir(x)]
            train_folders = all_folders[:-20]
            test_folders = all_folders[-20:]

        train_spk = [glob.glob(os.path.join(spk, '**', hp.form.input), recursive=True)
                        for spk in train_folders]
        train_spk = [x for x in train_spk if len(x) >= 2]

        test_spk = [glob.glob(os.path.join(spk, '**', hp.form.input), recursive=True)
                        for spk in test_folders]
        test_spk = [x for x in test_spk if len(x) >= 2]

        def train_wrapper(num):
            num+=start_index
            spk1, spk2 = random.sample(train_spk, 2)
            s1_dvec, s1_target = random.sample(spk1, 2)
            s2 = random.choice(spk2)
            mix_wrapper(hp, args, audio, num, s1_dvec, s1_target, s2, train=True)

        def test_wrapper(num):
            num+=start_index
            spk1, spk2 = random.sample(test_spk, 2)
            s1_dvec, s1_target = random.sample(spk1, 2)
            s2 = random.choice(spk2)
            mix_wrapper(hp, args, audio, num, s1_dvec, s1_target, s2, train=False)

        def generate_train_wrapper(num):
            spk1, spk2 = random.sample(train_spk, 2)
            s1_dvec, s1_target = random.sample(spk1, 2)
            #s2 = random.choice(spk2)
            s2_dvec, s2_target = random.sample(spk2, 2)
            s1_dvec=os.path.splitext(os.path.basename(s1_dvec))[0][:-5]
            s1_target=os.path.splitext(os.path.basename(s1_target))[0][:-5]
            #s2=os.path.splitext(os.path.basename(s2))[0]
            s2_dvec=os.path.splitext(os.path.basename(s2_dvec))[0][:-5]
            s2_target=os.path.splitext(os.path.basename(s2_target))[0][:-5]
            #return ','.join([s1_dvec, s1_target, s2])+'\n'
            return ','.join([s1_dvec, s1_target, s2_target, s2_dvec])+'\n'

        with open("datasets/random%d.csv" % total_num, 'w') as f:
            for i in tqdm.tqdm(range(total_num)):
                f.write(generate_train_wrapper(i))

        ##arr = list(range(10**5))
        #arr = list(range(50000))
        #with Pool(cpu_num) as p:
        #    r = list(tqdm.tqdm(p.imap(train_wrapper, arr), total=len(arr)))

        #arr = list(range(10**2))
        #with Pool(cpu_num) as p:
        #    r = list(tqdm.tqdm(p.imap(test_wrapper, arr), total=len(arr)))
    else:
        def filter_clean(files_ids3):
            file_dir1=os.path.join(args.libri_dir, 'train-clean-100')
            file_dir2=os.path.join(args.libri_dir, 'train-clean-360')
            file_dir3=os.path.join(args.libri_dir, 'dev-clean')
            file_ids=files_ids3.strip().split(',')
            for file_id in file_ids:
                file_paths=file_id.split('-')
                file_dirname=os.path.join(file_dir1, file_paths[0], file_paths[1])
                if not os.path.exists(file_dirname):
                    file_dirname=os.path.join(file_dir2, file_paths[0], file_paths[1])
                if not os.path.exists(file_dirname):
                    file_dirname=os.path.join(file_dir3, file_paths[0], file_paths[1])
                if not os.path.exists(file_dirname):
                    return False
            return True

        if generate_new_field_for_prepared:
            def get_speaker2_ref_utt(files_ids3):
                file_dir1=os.path.join(args.libri_dir, 'train-clean-100')
                file_dir2=os.path.join(args.libri_dir, 'train-clean-360')
                file_dir3=os.path.join(args.libri_dir, 'dev-clean')
                file_ids=files_ids3.strip().split(',')
                file_id=file_ids[2]
                file_paths=file_id.split('-')
                file_dirname=os.path.join(file_dir1, file_paths[0])
                if not os.path.exists(file_dirname):
                    file_dirname=os.path.join(file_dir2, file_paths[0])
                if not os.path.exists(file_dirname):
                    file_dirname=os.path.join(file_dir3, file_paths[0])
                if not os.path.exists(file_dirname):
                    raise Exception('File not found %s' % file_paths[0])
                ref_utt=file_id+'-norm.wav'
                candidates=glob.glob(os.path.join(file_dirname, '**', hp.form.input), recursive=True)
                while ref_utt == file_id+'-norm.wav':
                    ref_utt=random.choice(candidates)
                    ref_utt=os.path.basename(ref_utt)
                return ref_utt[:-9]

            file_dir1=os.path.join(args.libri_dir, 'train-clean-100')
            file_dir2=os.path.join(args.libri_dir, 'train-clean-360')
            file_dir3=os.path.join(args.libri_dir, 'dev-clean')
            all_speakers=[]
            all_speakers+=[speaker for speaker in os.listdir(file_dir1) \
                    if os.path.isdir(os.path.join(file_dir1, speaker))]
            all_speakers+=[speaker for speaker in os.listdir(file_dir2) \
                    if os.path.isdir(os.path.join(file_dir2, speaker))]
            all_speakers+=[speaker for speaker in os.listdir(file_dir3) \
                    if os.path.isdir(os.path.join(file_dir3, speaker))]

            def get_unpresented_ref_utt(files_ids3):
                file_dir1=os.path.join(args.libri_dir, 'train-clean-100')
                file_dir2=os.path.join(args.libri_dir, 'train-clean-360')
                file_dir3=os.path.join(args.libri_dir, 'dev-clean')

                file_ids=files_ids3.strip().split(',')
                speaker1=file_ids[1].split('-')[0]
                speaker2=file_ids[2].split('-')[0]
                unpresented=speaker1
                while unpresented==speaker1 or unpresented==speaker2:
                    unpresented=random.choice(all_speakers)

                file_dirname=os.path.join(file_dir1, unpresented)
                if not os.path.exists(file_dirname):
                    file_dirname=os.path.join(file_dir2, unpresented)
                if not os.path.exists(file_dirname):
                    file_dirname=os.path.join(file_dir3, unpresented)
                if not os.path.exists(file_dirname):
                    raise Exception('File not found %s' % unpresented)
                candidates=glob.glob(os.path.join(file_dirname, '**', hp.form.input), recursive=True)
                ref_utt=random.choice(candidates)
                ref_utt=os.path.basename(ref_utt)
                return ref_utt[:-9]

            with open(args.train_csv) as f:
                train_list=f.readlines()
            print("Prev train list len : %d" % len(train_list))
            train_list=[train_line for train_line in train_list if filter_clean(train_line)]
            print("Post train list len : %d" % len(train_list))
            #clean_train_csv='%s_clean.csv' % os.path.splitext(args.train_csv)[0]
            clean_train_csv='%s_unpresented.csv' % os.path.splitext(args.train_csv)[0]
            with open(clean_train_csv, 'w') as f:
                for line in tqdm.tqdm(train_list):
                    line=line.strip()
                    #ref_utt=get_speaker2_ref_utt(line)
                    ref_utt=get_unpresented_ref_utt(line)
                    line="%s,%s" % (line, ref_utt)
                    f.write(line+'\n')

            with open(args.test_csv) as f:
                test_list=f.readlines()
            print("Prev test list len : %d" % len(test_list))
            test_list=[test_line for test_line in test_list if filter_clean(test_line)]
            print("Post test list len : %d" % len(test_list))
            #clean_test_csv='%s_clean.csv' % os.path.splitext(args.test_csv)[0]
            clean_test_csv='%s_unpresented.csv' % os.path.splitext(args.test_csv)[0]
            with open(clean_test_csv, 'w') as f:
                for line in tqdm.tqdm(test_list):
                    line=line.strip()
                    #ref_utt=get_speaker2_ref_utt(line)
                    ref_utt=get_unpresented_ref_utt(line)
                    line="%s,%s" % (line, ref_utt)
                    f.write(line+'\n')
            exit()

        with open(args.train_csv) as f:
            train_list=f.readlines()
        print("Prev train list len : %d" % len(train_list))
        train_list=[train_line for train_line in train_list if filter_clean(train_line)]
        print("Post train list len : %d" % len(train_list))
        with open(args.test_csv) as f:
            test_list=f.readlines()
        print("Prev test list len : %d" % len(test_list))
        test_list=[test_line for test_line in test_list if filter_clean(test_line)]
        print("Post test list len : %d" % len(test_list))

        def get_libri_full_path_from_id(file_id):
            file_dir1=os.path.join(args.libri_dir, 'train-clean-100')
            file_dir2=os.path.join(args.libri_dir, 'train-clean-360')
            file_dir3=os.path.join(args.libri_dir, 'dev-clean')
            file_paths=file_id.split('-')
            file_suffix='-norm.wav'
            file_dirname=os.path.join(file_dir1, file_paths[0], file_paths[1])
            if not os.path.exists(file_dirname):
                file_dirname=os.path.join(file_dir2, file_paths[0], file_paths[1])
            if not os.path.exists(file_dirname):
                file_dirname=os.path.join(file_dir3, file_paths[0], file_paths[1])
            if not os.path.exists(file_dirname):
                print("Cannot find dir path %s" % file_dirname)
            file_full_path=os.path.join(file_dirname, "%s%s" % (file_id, file_suffix))
            if not os.path.exists(file_full_path):
                print("Cannot find file %s" % file_full_path)
            return file_full_path

        def prepared_train_wrapper(num):
            #s1_dvec, s1_target, s2 = train_list[num].strip().split(',')
            #s1_dvec, s1_target, s2_target, s2_dvec = train_list[num].strip().split(',')
            s1_dvec, s1_target, s2_target, s2_dvec, us_dvec = train_list[num].strip().split(',')
            s1_dvec = get_libri_full_path_from_id(s1_dvec)
            s1_target = get_libri_full_path_from_id(s1_target)
            s2_target = get_libri_full_path_from_id(s2_target)
            s2_dvec = get_libri_full_path_from_id(s2_dvec)
            us_dvec = get_libri_full_path_from_id(us_dvec)
            mix_wrapper(hp, args, audio, num, s1_dvec, s1_target, s2_target, s2_dvec, us_dvec, train=True)

        def prepared_test_wrapper(num):
            #s1_dvec, s1_target, s2 = test_list[num].strip().split(',')
            s1_dvec, s1_target, s2_target, s2_dvec, us_dvec = test_list[num].strip().split(',')
            s1_dvec = get_libri_full_path_from_id(s1_dvec)
            s1_target = get_libri_full_path_from_id(s1_target)
            s2_target = get_libri_full_path_from_id(s2_target)
            s2_dvec = get_libri_full_path_from_id(s2_dvec)
            us_dvec = get_libri_full_path_from_id(us_dvec)
            mix_wrapper(hp, args, audio, num, s1_dvec, s1_target, s2_target, s2_dvec, us_dvec, train=False)

        ##arr = list(range(10**5))
        #arr = list(range(112591))
        #with Pool(cpu_num) as p:
        #    r = list(tqdm.tqdm(p.imap(prepared_train_wrapper, arr), total=len(arr)))

        #arr = list(range(10**3))
        arr = list(range(1352))
        with Pool(20) as p:
            r = list(tqdm.tqdm(p.imap(prepared_test_wrapper, arr), total=len(arr)))


"""Code for constructing the model."""

# major imports
from absl import logging
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
import os
import time
from functools import partial
from jax.experimental import host_callback as hcb
# following packages for mpi communication will be just-in-time imported if use_mpi is True
# from mpi4py import MPI
# import mpi4jax
# from jax.tree_util import tree_flatten, tree_unflatten

# import type name specifications
from unifold.model.features import FeatureDict
from ml_collections import ConfigDict
from typing import Optional

# import major classes & functions
from unifold.model.modules import AlphaFold
from unifold.train.data_system import cast_to_precision
from unifold.train.optimizer import Optimizer
from unifold.train.mixed_precision import normalize_precision, set_unifold_policy

class Trainer:
  '''
  main class to train the UniFold model.
  '''
  def __init__(
      self,
      global_config: ConfigDict,
      optim_config: ConfigDict,
      model_config: ConfigDict,
      **kwargs):

    self.gc = global_config
    logging.info(f'trainer on device {jax.devices()}')
    self.precision = normalize_precision(self.gc.precision)
    set_unifold_policy(self.precision)

    def _forward_fn(batch):
      model = AlphaFold(model_config.model)
      return model(
          batch,
          is_training=True,
          compute_loss=True,
          ensemble_representations=False)

    if self.gc.debug:
      print("not jitted")
      self._init_fn = (hk.transform(_forward_fn).init)
      self._apply_fn = hk.transform(_forward_fn).apply
    else:
      self._init_fn = jax.jit(hk.transform(_forward_fn).init)
      self._apply_fn = jax.jit(hk.transform(_forward_fn).apply)
    self._loss_fn = None        # has to be initialized by external call on `Trainer.initialize()`
    self._update_fn = None      # has to be initialized by external call on `Trainer.initialize()`

    # optimizer variables, have to be initialized by external call on `Trainer.initialize()`
    self.optim_config = optim_config
    self.optimizer = None       # instance of unifold.trainer.optimizer.Optimizer
    self.optim_state = None     # optimizer state, key variable of maintaining model parameters

    # logging variables organized in format [(step, loss, time), (step, loss, time), ...]
    self.train_losses = []
    self.eval_losses = []

    # timing variables
    self._tic = time.time()      # tic & toc style for timing. handled in `Trainer.logging_on_the_fly`

    # mpi variables
    if self.gc.use_mpi:
      self.mpi_comm = kwargs['mpi_comm']
      self.mpi_rank = self.mpi_comm.Get_rank()
      print("INFO: mpi rank {}".format(self.mpi_rank))
      self.world_size = self.mpi_comm.Get_size()

    # path formatters of ckpts and loss curves
    self.auto_ckpt_name = \
        lambda step, format: f"{self.gc.model_name}_{step:05d}.{format}"
    self.auto_curve_name = \
        lambda is_train: f"{self.gc.model_name}_{'train' if is_train else 'eval'}_curve.npy"

    # step specifier
    self.is_logging_step = lambda i: i % self.gc.logging_freq == 0
    self.is_save_step = lambda i: (i + 1) % self.gc.save_freq == 0
    self.is_eval_step = lambda i: i % self.gc.eval_freq == 0

  @property
  def params(self):
    if self.optim_state is None:
      return None
    else:
      return self.optimizer.get_params(self.optim_state)

  def initialize(
      self,
      batch: Optional[FeatureDict] = None,
      load_format: Optional[str] = "pkl",
      random_seed: Optional[int] = None):

    # create optimizer instance
    self.optimizer = Optimizer(self.optim_config)

    use_autoload = self.gc.start_step > 0
    if use_autoload:
      assert load_format is not None, \
          "must provide `load_format` to auto load models when assigning `start_step` > 0."
      self.autoload(self.gc.start_step, load_format)
    else:
      assert batch is not None, \
          "must provide a batch and a random seed to initialize model from scratch."
      if self.optim_state is not None:
        logging.warning("existed optimizer states are reinitialized.")
      if random_seed is not None:
        logging.warning("external random seed overrides the one in global config.")
      else:
        random_seed = self.gc.random_seed
        rng = jax.random.PRNGKey(random_seed)    # all ranks initialized equally.
      params = self._init_fn(batch=batch, rng=rng)
      params = hk.data_structures.to_mutable_dict(params)
      self.optim_state = self.optimizer.init_state(params)

    # define loss_fn
    def _loss_fn(params, batch, rng):
      # TODO: user external RNG
      _, loss = self._apply_fn(params=params, batch=batch, rng=rng)
      seq_length_weight = jnp.sqrt(jnp.sum(batch['all_atom_mask'][0,:,0]))
      return loss * seq_length_weight

    # define reduce_fn for mpi communication.
    if self.gc.use_mpi:
      # just-in-time imports
      from mpi4py import MPI
      import mpi4jax
      from jax.tree_util import tree_flatten, tree_unflatten
      def _mpi_reduce_value(value):
        value.block_until_ready()
        cpu = jax.devices('cpu')[0]
        device = value.device()
        value = jax.device_put(value, cpu)
        value, _ = mpi4jax.allreduce(value, op=MPI.SUM, comm=self.mpi_comm)
        value /= self.world_size
        value = jax.device_put(value, device)
        return value
      def _mpi_reduce_tree(tree):
        flat_tree, tree_struct = tree_flatten(tree)
        for i, val in enumerate(flat_tree):
          flat_tree[i] = _mpi_reduce_value(val)
        tree = tree_unflatten(tree_struct, flat_tree)
        return tree
      def _hcb_reduce_tree(tree):
        tree = hcb.call(_mpi_reduce_tree, tree, result_shape=tree)
        return tree

    # define eval_fn for validation.
    def _eval_fn(params, batch, rng):
      loss = _loss_fn(params, batch, rng)
      if self.gc.use_mpi:
        loss = _mpi_reduce_value(loss)
      return loss

    self._loss_fn = _loss_fn     # this is not re-jit as loss_fn is much of a wrapped apply_fn.
    self._eval_fn = _eval_fn
    if self.gc.debug:
      self.value_and_grad = jax.value_and_grad(self._loss_fn)
      self.opt_update = self.optimizer.opt_update
    else:
      self.value_and_grad = jax.jit(jax.value_and_grad(self._loss_fn))
      self.opt_update = jax.jit(self.optimizer.opt_update)
    # define update_fn.
    def _update_fn(step, opt_state, batch, rng):
      loss, grads = self.value_and_grad(self.optimizer.get_params(opt_state), batch, rng)
      grads = self.optimizer.clip_grads(grads)
      if self.gc.use_mpi:
        loss = _mpi_reduce_value(loss)
        grads = _mpi_reduce_tree(grads)
      opt_state = self.opt_update(step, grads, opt_state)
      return opt_state, loss

    self._update_fn = _update_fn   # jit transformation of update_fn.
    # start ticking after initialization.
    self._tic = time.time()

  def autosave(self, step, format='pkl'):
    # save ckpt
    save_path = os.path.join(self.gc.save_dir, self.auto_ckpt_name(step + 1, format))
    self.optimizer.save(self.optim_state, save_path)
    # save loss curve
    train_curve_path = os.path.join(self.gc.save_dir, self.auto_curve_name(is_train=True))
    eval_curve_path = os.path.join(self.gc.save_dir, self.auto_curve_name(is_train=False))
    np.save(train_curve_path, np.asarray(self.train_losses))
    np.save(eval_curve_path, np.asarray(self.eval_losses))
    logging.info(f"model autosaved at step {step:05d} successfully.")

  def autoload(self, step, format='pkl'):
    # load ckpt
    load_path = os.path.join(self.gc.load_dir, self.auto_ckpt_name(step, format))
    self.optimizer.load(load_path)
    # load loss curve
    train_curve_path = os.path.join(self.gc.save_dir, self.auto_curve_name(is_train=True))
    eval_curve_path = os.path.join(self.gc.save_dir, self.auto_curve_name(is_train=False))
    def load_loss_curve(loss_curve_path: str):
      try:
        return np.load(loss_curve_path).tolist()
      except:
        logging.warning(f"failed to load curve from {loss_curve_path}. reset loss curve.")
        return []
    self.train_losses = load_loss_curve(train_curve_path)
    self.eval_losses = load_loss_curve(eval_curve_path)
    logging.info(f"model autoloaded at step {step:05d} successfully.")

  def update(self, step, batch, rng):
    # wrapped update_fn for external calls.
    opt_state, loss = self._update_fn(step, self.optim_state, batch, rng)
    self.optim_state = opt_state
    return loss

  def _logging(self, step, loss):
    # print and record training stats at the step.
    toc = time.time()
    step_time = (toc - self._tic) / (
        1 if step == 0 else self.gc.logging_freq)
    self.train_losses.append((step, loss, step_time))
    logging.info(f"step: {step:05d}\ttrain_loss: {loss:3.4f}\tstep_time: {step_time:.2f}s")
    self._tic = time.time()

  def train_step(self, step, batch, rng, silent=True):
    batch = cast_to_precision(batch, self.precision)
    loss = self.update(step, batch, rng)
    if not silent:
      if self.is_logging_step(step):
        self._logging(step, loss)
      if self.is_save_step(step):
        self.autosave(step, format='pkl')

  def eval_step(self, step, batch, rng, silent=True):
    # evaluation on the fly
    tmp_tic = time.time()
    loss = self._eval_fn(self.params, batch, rng)
    eval_time = time.time() - tmp_tic
    if not silent:
      self.eval_losses.append((step, loss, eval_time))
      logging.info(f"step: {step:05d}\teval_loss:  {loss:3.4f}\teval_time: {eval_time:.2f}s")


import sys, time, os, importlib, random
import numpy as np
import torch
import torch.nn.functional as F
from scipy.io import wavfile
from config.args import parser
from tools import load_yaml_config, load_speaker_model_parameters, reverb, reverb_np

args = parser.parse_args()
args = load_yaml_config(args, parser)
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

# set random seed for reproducibility
def seed_torch(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.enabled = False
seed_torch(args.seed)

# load wav data with shape (1, SampleNumbers)
def loadWAV(filename):
    sample_rate, audio  = wavfile.read(filename)
    feat = np.stack([audio], axis=0).astype(np.float)
    return feat

# perform attack
def Attack(speaker_model, enroll_speaker_feat, train_speaker, speaker, target, args):
    '''
    speaker_model: the ASV model we want to attack.
    speaker: original speaker name.
    target: targeted speaker name.
    enroll_speaker_feat: all enrolled targeted speaker embeddings.
    train_speaker: original speakers with their train audio set.
    args: hyper-parameters
    '''
    # Set parameter
    num_train = args.num_train
    adv_len = int(args.noise_len * 16000)
    max_steps = args.max_steps
    lr = args.lr
    momentum = args.momentum
    max_delta = args.max_delta
    e1 = args.e1
    e2 = args.e2

    # Load wav data
    wavs = train_speaker[speaker]
    wav_input_np = []
    max_len = 0
    for i in range(num_train):
        wav = os.path.join(args.wav_path, wavs[i])
        wav_data = loadWAV(wav)
        if wav_data.shape[1] > max_len:
            max_len = wav_data.shape[1]
        wav_input_np.append(wav_data)

    # repeat wav data to max audio length in train audio set
    for i in range(num_train):
        wav_input_np[i] = np.tile(wav_input_np[i], (1, max_len // wav_input_np[i].shape[1] + 1))[:, 0:max_len]
    wav_input_np = np.asarray(wav_input_np).squeeze(axis=1)

    # with rir simulation, load train rir path list
    if args.rir:
        rir_wavs_train = os.listdir(args.attack_rir_path)

    # initialize tensors
    adv_noise_np = np.random.uniform(-max_delta, max_delta, adv_len).astype(np.int16)
    target_feat = torch.FloatTensor(enroll_speaker_feat[target]).cuda()
    wav_input = torch.FloatTensor(wav_input_np).cuda()

    # Step1 maximize the attack on ASV model
    batch_num = num_train // args.batch # we do this for GPU memory limit
    grad_pre = 0
    result = adv_noise_np
    for i in range(max_steps):
        loss = 0
        grad = 0
        dist = []
        for b in range(batch_num):
            wav_input_batch = wav_input[b*args.batch:(b+1)*args.batch]

            # optimized tensors
            adv_noise = torch.FloatTensor(adv_noise_np).cuda()
            adv_noise.requires_grad = True

            # pad the adversarial perturbation to have a same length with train audios
            adv_noise_tmp = adv_noise.repeat(max_len // adv_len + 1)[:max_len]

            # if doing rir simulation, we have to reverberate the audio and perturbation
            if not args.rir:
                wav_adv_input = wav_input_batch + adv_noise_tmp.unsqueeze(0)
            else:
                rir_tensor = torch.FloatTensor(loadWAV(os.path.join(args.attack_rir_path, random.choice(rir_wavs_train)))).cuda()
                rir_tensor = rir_tensor / torch.norm(rir_tensor) # Maintain the energy of the original signal
                adv_noise_tmp = reverb(adv_noise_tmp.unsqueeze(0), rir_tensor)
                wav_input_batch = reverb(wav_input_batch, rir_tensor)
                wav_adv_input = wav_input_batch + adv_noise_tmp
            wav_adv_input = torch.clamp(wav_adv_input, -2**15, 2**15-1) # clip the input data to a normal format
            wav_feat = speaker_model.forward(wav_adv_input)
            dist_batch = F.cosine_similarity(wav_feat, target_feat)

            # loss1
            loss1 = torch.sum(torch.clamp(args.thresh - dist_batch + args.margine, min=0.))
            loss1.backward()

            loss += loss1.item()
            grad += adv_noise.grad.data.cpu().numpy()
            dist.extend(list(dist_batch.detach().cpu().numpy()))

        # Collect and average the loss and gradient
        loss /= batch_num
        grad /= batch_num
        print('step1: %02d, %s to %s, loss: % 2.4f' % (i, speaker, target, loss))

        # if convergence break
        result = adv_noise_np
        if loss < e1:
            break

        # PGD can also be interpreted as an iterative algorithm to solve the problem
        grad_new = momentum * grad_pre + grad
        grad_pre = grad_new
        adv_noise_np = adv_noise_np - lr * np.sign(grad_new)
        adv_noise_np = np.clip(adv_noise_np, -max_delta, max_delta)

    # record the iteration steps and best adversarial perturbation for step1
    step1 = i
    result1 = result

    # Step2 minimize impact on ASR model
    # initialize the adversarial perturbation from best result from step1
    grad_pre = 0
    for i in range(max_steps):
        loss1_value = 0
        loss2_value = 0
        grad1 = 0
        grad2 = 0
        dist = []
        for b in range(batch_num):
            wav_input_batch = wav_input[b*args.batch:(b+1)*args.batch]

            # optimized tensors
            adv_noise = torch.FloatTensor(adv_noise_np).cuda()
            adv_noise.requires_grad = True

            # pad the adversarial perturbation to have a same length with train audios
            adv_noise_tmp = adv_noise.repeat(max_len // adv_len + 1)[:max_len]

            # if doing rir simulation, we have to reverberate the audio and perturbation
            if not args.rir:
                wav_adv_input = wav_input_batch + adv_noise_tmp.unsqueeze(0)
            else:
                rir_tensor = torch.FloatTensor(loadWAV(os.path.join(args.attack_rir_path, random.choice(rir_wavs_train)))).cuda()
                rir_tensor = rir_tensor / torch.norm(rir_tensor) # Maintain the energy of the original signal
                adv_noise_tmp = reverb(adv_noise_tmp.unsqueeze(0), rir_tensor)
                wav_input_batch = reverb(wav_input_batch, rir_tensor)
                wav_adv_input = wav_input_batch + adv_noise_tmp
            wav_adv_input = torch.clamp(wav_adv_input, -2**15, 2**15-1) # clip the input data to a normal format
            wav_feat = speaker_model.forward(wav_adv_input)
            dist_batch = F.cosine_similarity(wav_feat, target_feat)

            # loss1
            loss1 = torch.sum(torch.clamp(args.thresh - dist_batch + args.margine, min=0.))

            # loss2
            spec_stft = torch.stft(adv_noise/2**15, n_fft=320, hop_length=160, win_length=320, window=torch.hamming_window(320).cuda())
            loss2 = torch.mean(torch.sqrt(torch.square(spec_stft[:,:,0]) + torch.square(spec_stft[:,:,1])))

            loss1_value += loss1.item()
            loss2_value += loss2.item()

            # caculate gradient separately
            grad1 += torch.autograd.grad(loss1, adv_noise)[0].detach().cpu().numpy()
            grad2 += torch.autograd.grad(loss2, adv_noise)[0].detach().cpu().numpy()
            dist.extend(list(dist_batch.detach().cpu().numpy()))

        # Collect and average the loss and gradient
        loss1_value /= batch_num
        loss2_value /= batch_num
        grad1 /= batch_num
        grad2 /= batch_num
        print('step2: %02d, %s to %s, loss1: %2.4f, loss2: %2.4f' % (i, speaker, target, loss1_value, loss2_value))

        # if convergence break
        result = adv_noise_np
        if loss1_value <= e1 and loss2_value <= e2:
            break

        # PGD can also be interpreted as an iterative algorithm to solve the problem:
        alpha = args.gamma if loss1_value > e1 else 0
        beta = 1 if loss2_value > e2 else 0
        grad_new = momentum * grad_pre + alpha * grad1 / (np.linalg.norm(grad1, 1) + 1e-12) + beta * grad2 / np.linalg.norm(grad2, 1)
        grad_new = np.nan_to_num(grad_new)
        grad_pre = grad_new
        adv_noise_np = adv_noise_np - lr * np.sign(grad_new)
        adv_noise_np = np.clip(adv_noise_np, -max_delta, max_delta)
    step2 = i

    # return the best result, iteration steps, loss value for step1 and step2,
    return result1, step1, loss, result, step2, loss1_value, loss2_value

def Test(speaker_model, enroll_speaker_feat, test_speaker, speaker, target, adv_noise, args):
    '''
    speaker_model: the ASV model we want to attack.
    speaker: original speaker name.
    target: targeted speaker name.
    adv_noise: adversarial perturbation that attack ASV model to verify speaker as target.
    enroll_speaker_feat: all enrolled targeted speaker embeddings.
    train_speaker: original speakers with their train audio set.
    args: hyper-parameters
    '''
    wavs = test_speaker[speaker]
    target_feat = torch.FloatTensor(enroll_speaker_feat[target]).cuda()
    success = 0

    # with rir simulation, load test rir path list
    if args.rir:
        rir_wavs_test = os.listdir(args.attack_rir_path)

    # evaluate test audios one by one
    for wav in wavs:
        wav = os.path.join(args.wav_path, wav)
        wav_data = loadWAV(wav)

        # pad the adversarial perturbation to the same length with one test audio
        adv_noise_tmp = np.tile(adv_noise, wav_data.shape[1] // adv_noise.shape[0] + 1)[0:wav_data.shape[1]]

        # if doing rir simulation, we have to reverberate the audio and perturbation
        if not args.rir:
            wav_data = wav_data + np.expand_dims(adv_noise_tmp, axis=0)
        else:
            rir_np = loadWAV(os.path.join(args.attack_rir_path, random.choice(rir_wavs_test)))
            rir_np = rir_np / np.linalg.norm(rir_np)
            adv_noise_tmp = reverb_np(np.expand_dims(adv_noise_tmp, axis=0), rir_np)
            wav_data = reverb_np(wav_data, rir_np)
            wav_data = wav_data + adv_noise_tmp
        wav_data = np.clip(wav_data, -2**15, 2**15-1)
        wav_data_tensor = torch.FloatTensor(wav_data).cuda()
        wav_feat = speaker_model.forward(wav_data_tensor).detach()

        # use consine similarity or MSE distance metric
        if args.cosine_similarity:
            dist = F.cosine_similarity(wav_feat, target_feat).cpu().numpy()
            score = 1 * np.mean(dist)
        else:
            wav_feat = F.normalize(wav_feat, p=2, dim=1)
            target_feat = F.normalize(target_feat, p=2, dim=1)
            dist = F.pairwise_distance(wav_feat.unsqueeze(-1), target_feat.unsqueeze(-1).transpose(0,2)).detach().cpu().numpy()
            score = -1 * np.mean(dist)
        print(score)
        if score > args.thresh:
            success += 1

    print('%s to %s attack success rate: %2.1f' % (speaker, target, success * 1.0 / len(wavs) * 100))
    return success, len(wavs), success * 1.0 / len(wavs) * 100

####################################################################################################################
####################################################################################################################
### step 1: Set Parameters ###
enroll_file = args.enroll_file
train_file = args.train_file
test_file = args.test_file
wav_path = args.wav_path
enroll_path = args.enroll_path
libri_speaker_file = './datas/gender_librispeech_test_clean.txt'

enroll_speaker = {}
with open(enroll_file, 'r') as f:
    for line in f.readlines():
        speaker = line.split(" ")[0]
        if speaker not in enroll_speaker:
            enroll_speaker[speaker] = []
        enroll_speaker[speaker].append(line.split(" ")[1].strip())

train_speaker = {}
with open(train_file, 'r') as f:
    for line in f.readlines():
        speaker = line.split(" ")[0]
        if speaker not in train_speaker:
            train_speaker[speaker] = []
        train_speaker[speaker].append(line.split(" ")[1].strip())

test_speaker = {}
with open(test_file, 'r') as f:
    for line in f.readlines():
        speaker = line.split(" ")[0]
        if speaker not in test_speaker:
            test_speaker[speaker] = []
        test_speaker[speaker].append(line.split(" ")[1].strip())

### step 2: Load speaker model ###
print('Loading speaker model...')
speaker_model = importlib.import_module('models.' + args.model).__getattribute__('MainModel')
speaker_model = speaker_model(**vars(args)).cuda()
if args.initial_model != "":
    speaker_model = load_speaker_model_parameters(speaker_model, args.initial_model)
    print("Model %s loaded!"%args.initial_model)
speaker_model.eval()

### step 3: enrolling speakers ###
print('Enrolling...')
enroll_speaker_feat = {}
with torch.no_grad():
    for k,v in enroll_speaker.items():
        enroll_speaker_feat[k] = 0
        for i in range(len(v)):
            wav = os.path.join(wav_path, v[i])
            wav_data = torch.FloatTensor(loadWAV(wav)).cuda()
            wav_feat = speaker_model.forward(wav_data)
            enroll_speaker_feat[k] += wav_feat.detach().cpu().numpy()
            del wav_feat, wav_data
        enroll_speaker_feat[k] /= len(v)
print('Enrolling ok!')

### step 4: match original and target speaker ###
female_speaker = []
male_speaker = []
with open(libri_speaker_file, 'r') as f:
    for line in f.readlines():
        if line.split(" ")[1].strip() == 'M':
            male_speaker.append(line.split(" ")[0])
        else:
            female_speaker.append(line.split(" ")[0])

original_speaker = []
target_speaker = []
# intra-genter
for speaker in female_speaker:
    original_speaker.append(speaker)
    while True:
        idx = random.randint(0, len(female_speaker)-1)
        if female_speaker[idx] != speaker:
            target_speaker.append(female_speaker[idx])
            break
for speaker in male_speaker:
    original_speaker.append(speaker)
    while True:
        idx = random.randint(0, len(male_speaker)-1)
        if male_speaker[idx] != speaker:
            target_speaker.append(male_speaker[idx])
            break
# inter-genter
for speaker in female_speaker:
    original_speaker.append(speaker)
    target_speaker.append(male_speaker[random.randint(0, len(male_speaker)-1)])
for speaker in male_speaker:
    original_speaker.append(speaker)
    target_speaker.append(female_speaker[random.randint(0, len(female_speaker)-1)])

# write intra and inter gender matchs into file
with open('./datas/intra_match.txt', 'w') as f:
    for i in range(40):
        f.write(original_speaker[i] + ',' + target_speaker[i] + '\n')
with open('./datas/inter_match.txt', 'w') as f:
    for i in range(40):
        f.write(original_speaker[i+40] + ',' + target_speaker[i+40] + '\n')

### step 5: attack and get universal adversarial noise ###
if not os.path.exists(args.out_path):
    os.mkdir(args.out_path)
    os.mkdir(os.path.join(args.out_path, 'step1'))
    os.mkdir(os.path.join(args.out_path, 'step2'))

f = open(os.path.join(args.out_path, 'result.txt'), 'w')

#intra-gender attack
if args.intra_gender:
    for i in range(len(original_speaker)//2):
        print('Num: %02d' % i)
        speaker = original_speaker[i]
        target = target_speaker[i]
        noise1, step1, loss, noise, step2, loss1_value, loss2_value = Attack(speaker_model, enroll_speaker_feat, train_speaker, speaker, target, args)
        with torch.no_grad():
            success, number, rate = Test(speaker_model, enroll_speaker_feat, test_speaker, speaker, target, noise, args)
        torch.cuda.empty_cache()
        out = 'Num: %02d, %5s to %5s, step1: %4d, loss: %.4f, step2: %4d, loss1: %.4f, loss2: %.4f, success number: %3d, total number: %3d, success rate: %3.2f\n' % (i+1, speaker, target, step1, loss, step2, loss1_value, loss2_value, success, number, rate)
        f.write(out)
        f.flush()
        wavfile.write(os.path.join(args.out_path, 'step1', speaker +'_' + target + '.wav'), 16000, np.asarray(noise1, dtype=np.int16))
        wavfile.write(os.path.join(args.out_path, 'step2', speaker +'_' + target + '.wav'), 16000, np.asarray(noise, dtype=np.int16))

# inter-gender attack
if args.inter_gender:
    for i in range(len(original_speaker)//2):
        print('Num: %02d' % i)
        speaker = original_speaker[i+40]
        target = target_speaker[i+40]
        noise1, step1, loss, noise, step2, loss1_value, loss2_value = Attack(speaker_model, speech_model, enroll_speaker_feat, train_speaker, speaker, target, args)
        with torch.no_grad():
            success, number, rate = Test(speaker_model, speech_model, enroll_speaker_feat, test_speaker, speaker, target, noise, args)
        torch.cuda.empty_cache()
        out = 'Num: %02d, %5s to %5s, step1: %4d, loss: %.4f, step2: %4d, loss1: %.4f, loss2: %.4f, success number: %3d, total number: %3d, success rate: %3.2f\n' % (i+1, speaker, target, step1, loss, step2, loss1_value, loss2_value, success, number, rate)
        f.write(out)
        f.flush()
        wavfile.write(os.path.join(args.out_path, 'step1', speaker +'_' + target + '.wav'), 16000, np.asarray(noise1, dtype=np.int16))
        wavfile.write(os.path.join(args.out_path, 'step2', speaker +'_' + target + '.wav'), 16000, np.asarray(noise, dtype=np.int16))


from functools import lru_cache
from dataclasses import dataclass
from types import SimpleNamespace

import numpy as np
# from skimage.measure import compare_mse
from PIL import Image

from util.richlog import get_logger
from . import imgops
from . import minireco
from . import resources

@dataclass
class RecognizedItem:
    name: str
    quantity: int
    low_confidence: bool = False

logger = get_logger(__name__)

@lru_cache(1)
def load_data():
    _, files = resources.get_entries('items')
    iconmats = {}
    itemmask = np.asarray(resources.load_image('common/itemmask.png', '1'))
    for filename in files:
        if filename.endswith('.png'):
            basename = filename[:-4]
            img = resources.load_image('items/' + filename, 'RGB')
            if img.size != (48, 48):
                img = img.resize((48, 48), Image.BILINEAR)
            mat = np.array(img)
            mat[itemmask] = 0
            iconmats[basename] = mat
    model = resources.load_pickle('minireco/NotoSansCJKsc-DemiLight-nums.dat')
    reco = minireco.MiniRecognizer(model, minireco.compare_ccoeff)
    return SimpleNamespace(itemmats=iconmats, num_recognizer=reco, itemmask=itemmask)

@lru_cache()
def all_known_items():
    _, files1 = resources.get_entries('items')
    _, files2 = resources.get_entries('items/archive')
    _, files3 = resources.get_entries('items/not-loot')
    files = files1 + files2 + files3
    return [filename[:-4] if filename.endswith('.png') else filename for filename in files]

def tell_item(itemimg, with_quantity=True):
    logger.logimage(itemimg)
    cached = load_data()
    # l, t, r, b = scaledwh(80, 146, 90, 28)
    # print(l/itemimg.width, t/itemimg.height, r/itemimg.width, b/itemimg.height)
    # numimg = itemimg.crop(scaledwh(80, 146, 90, 28)).convert('L')
    low_confidence = False
    quantity = None
    if with_quantity:
        numimg = imgops.scalecrop(itemimg, 0.39, 0.71, 0.82, 0.84).convert('L')
        numimg = imgops.crop_blackedge2(numimg, 120)

        if numimg is not None:
            numimg = imgops.clear_background(numimg, 120)
            logger.logimage(numimg)
            numtext, score = cached.num_recognizer.recognize2(numimg, subset='0123456789万')
            logger.logtext('quantity: %s, minscore: %f' % (numtext, score))
            if score < 0.2:
                low_confidence = True
            quantity = int(numtext) if numtext.isdigit() else None

    # scale = 48/itemimg.height
    img4reco = np.array(itemimg.resize((48, 48), Image.BILINEAR).convert('RGB'))
    img4reco[cached.itemmask] = 0

    scores = []
    for name, templ in cached.itemmats.items():
        scores.append((name, imgops.compare_mse(img4reco, templ)))

    scores.sort(key=lambda x: x[1])
    itemname, score = scores[0]
    # maxmatch = max(scores, key=lambda x: x[1])
    logger.logtext(repr(scores[:5]))
    diffs = np.diff([a[1] for a in scores])
    if score < 800 and np.any(diffs > 600):
        logger.logtext('matched %s with mse %f' % (itemname, score))
        name = itemname
    else:
        logger.logtext('no match')
        low_confidence = True
        name = None

    return RecognizedItem(name, quantity, low_confidence)


#---------------------------------------------------------------------------
# IMPORT MODULES
#---------------------------------------------------------------------------

from PIL import Image

import torch
from torchvision import transforms, datasets
import numpy as np
import os
import numpy as np
import torch
from torch.autograd import Variable
from src.get_nets import PNet, RNet, ONet
from src.box_utils import nms, calibrate_box, get_image_boxes, convert_to_square
from src.first_stage import run_first_stage
import torch.nn as nn

#---------------------------------------------------------------------------
# PATH DEFINITIONS
#---------------------------------------------------------------------------

processed_dataset_path = '../Dataset/FaceFeatures/test/'

#---------------------------------------------------------------------------
# MTCNN MODEL DEFINITION for EXTRACTING FACE FEATURES
#---------------------------------------------------------------------------

pnet = PNet()
rnet = RNet()
onet = ONet()
onet.eval()

class OnetFeatures(nn.Module):
    def __init__(self, original_model):
        super(OnetFeatures, self).__init__()
        self.features = nn.Sequential(*list(onet.children())[:-3])

    def forward(self, x):
        x = self.features(x)
        return x

def get_face_features(image, min_face_size=20.0,
                 thresholds=[0.6, 0.7, 0.8],
                 nms_thresholds=[0.7, 0.7, 0.7]):
    """
    Arguments:
        image: an instance of PIL.Image.
        min_face_size: a float number.
        thresholds: a list of length 3.
        nms_thresholds: a list of length 3.

    Returns:
        two float numpy arrays of shapes [n_boxes, 4] and [n_boxes, 10],
        bounding boxes and facial landmarks.
    """

    # LOAD MODELS
    pnet = PNet()
    rnet = RNet()
    onet = ONet()
    onet.eval()

    # BUILD AN IMAGE PYRAMID
    width, height = image.size
    min_length = min(height, width)

    min_detection_size = 12
    factor = 0.707  # sqrt(0.5)

    # scales for scaling the image
    scales = []

    # scales the image so that
    # minimum size that we can detect equals to
    # minimum face size that we want to detect
    m = min_detection_size/min_face_size
    min_length *= m

    factor_count = 0
    while min_length > min_detection_size:
        scales.append(m*factor**factor_count)
        min_length *= factor
        factor_count += 1

    # STAGE 1

    # it will be returned
    bounding_boxes = []

    # run P-Net on different scales
    for s in scales:
        boxes = run_first_stage(image, pnet, scale=s, threshold=thresholds[0])
        bounding_boxes.append(boxes)

    # collect boxes (and offsets, and scores) from different scales
    bounding_boxes = [i for i in bounding_boxes if i is not None]
    bounding_boxes = np.vstack(bounding_boxes)

    keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
    bounding_boxes = bounding_boxes[keep]

    # use offsets predicted by pnet to transform bounding boxes
    bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
    # shape [n_boxes, 5]

    bounding_boxes = convert_to_square(bounding_boxes)
    bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])

    # STAGE 2

    img_boxes = get_image_boxes(bounding_boxes, image, size=24)
    img_boxes = Variable(torch.FloatTensor(img_boxes), volatile=True)
    output = rnet(img_boxes)
    offsets = output[0].data.numpy()  # shape [n_boxes, 4]
    probs = output[1].data.numpy()  # shape [n_boxes, 2]

    keep = np.where(probs[:, 1] > thresholds[1])[0]
    bounding_boxes = bounding_boxes[keep]
    bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
    offsets = offsets[keep]

    keep = nms(bounding_boxes, nms_thresholds[1])
    bounding_boxes = bounding_boxes[keep]
    bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
    bounding_boxes = convert_to_square(bounding_boxes)
    bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])

    # STAGE 3

    img_boxes = get_image_boxes(bounding_boxes, image, size=48)
    if len(img_boxes) == 0:
        return [], []
    img_boxes = Variable(torch.FloatTensor(img_boxes), volatile=True)
    output = onet(img_boxes)

    faceFeatureModel = OnetFeatures(onet)

    featureOutputs = faceFeatureModel(img_boxes)

    landmarks = output[0].data.numpy()  # shape [n_boxes, 10]
    offsets = output[1].data.numpy()  # shape [n_boxes, 4]
    probs = output[2].data.numpy()  # shape [n_boxes, 2]

    keep = np.where(probs[:, 1] > thresholds[2])[0]
    bounding_boxes = bounding_boxes[keep]
    bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
    offsets = offsets[keep]
    landmarks = landmarks[keep]

    bounding_boxes = calibrate_box(bounding_boxes, offsets)
    keep = nms(bounding_boxes, nms_thresholds[2], mode='min')

    featureOutputs = featureOutputs[keep]

    return featureOutputs

#---------------------------------------------------------------------------
# LOAD TEST DATASET
#---------------------------------------------------------------------------

test_data_filelist = sorted(os.listdir('../Dataset/emotiw/test/test_shared/'))

for i in test_data_filelist:
    if i[0] != 't':
        test_data_filelist.remove(i)

print(len(test_data_filelist))

#---------------------------------------------------------------------------
# EXTRACT FACE FEATURES
#---------------------------------------------------------------------------

for i in range(len(test_data_filelist)):
    filename = test_data_filelist[i]
    filename = filename[:-4]
    image = Image.open('../Dataset/emotiw/test/test_shared/'+filename+'.jpg')

    print(filename)
    try:

        if os.path.isfile(processed_dataset_path + filename + '.npz'):
            print(filename + ' Already present')
            continue

        features = get_face_features(image)

        if (type(features)) == tuple:
            with open('hello.text', 'a') as f:
                f.write(filename)
            continue

        features = features.data.numpy()

        if features.size == 0:
            print('MTCNN model handling empty face condition at ' + filename)
        np.savez(processed_dataset_path + filename , a=features)

    except ValueError:
        print('No faces detected for ' + filename + ". Also MTCNN failed.")
        np.savez(processed_dataset_path + filename , a=np.zeros(1))
        continue