import tensorflow as tf
import numpy as np
import os
import shutil

# core config parameters
_gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.7)
_initializer = tf.random_uniform_initializer(-0.1,  0.1)
_dropout_keep_prob = 0.6
_max_grad_norm = 0.5
_tensorboard_log_rel_path = "tb_log"
_save_kept_num = 10

class TfBaseConfig(object):
    @property
    def model_save_path(self):
        if self._model_save_path[-1] == "/":
            return self._model_save_path[:-1]
        return self._model_save_path

    def __init__(self, model_save_path, gpu_options=_gpu_options, initializer=_initializer,
                 keep_prob=_dropout_keep_prob,
                 max_grad_norm=_max_grad_norm, save_kept_num=_save_kept_num):
        self._model_save_path = model_save_path
        self.gpu_options = gpu_options
        self.initializer = initializer
        self.keep_prob = keep_prob
        self.max_grad_norm = max_grad_norm
        self.save_kept_num = save_kept_num
        self.tensorboard_log_path = self.model_save_path + "/" + _tensorboard_log_rel_path

class TfBaseModel(object):
    def __init__(self, config, global_reuse, *args, graph=None, sess=None, **kwargs):
        """
        can provide another model's Graph and Session. i.e. when reusing to create a test model, provide the train model's graph and session
        """
        self.error_list = []
        self.graph = tf.Graph() if graph is None else graph
        self.session = tf.Session(graph=self.graph) if graph is None else sess
        self.config = config
        self.tensorboard_watches = {}
        self.merged_summary = None
        self.global_reuse = global_reuse
        with self.graph.as_default(), tf.variable_scope("model", reuse=global_reuse):
            self.build(*args, **kwargs)
            self.saver = tf.train.Saver(max_to_keep=10)

        if not os.path.exists(self.config.tensorboard_log_path):
            os.makedirs(self.config.tensorboard_log_path)

        self.summary_writer = tf.summary.FileWriter(self.config.tensorboard_log_path)

    def build(self, *args, **kwargs):
        """
        TODO: build model
        :return:
        """
        return

    def back_propagate(self, cost, check_finite=False, optimizer=None, clip_grads=True, lr=None, momentum=None):
        """
        optimizer优先度：自定义optimizer > 如果lr=None就Adam > 如果lr<0就动态lr，注意需要feed self.lr > 固定lr的梯度下降
        """

        if optimizer is not None:
            bp_optimizer = optimizer
        elif lr is None:
            bp_optimizer = tf.train.AdamOptimizer()
        elif momentum is None:
            if lr < 0:
                lr = tf.placeholder(tf.float32, shape=(), name="adjustable_lr")
                self.lr = lr

            bp_optimizer = tf.train.GradientDescentOptimizer(lr)
        else:
            bp_optimizer = tf.train.MomentumOptimizer(lr, momentum)

        tvars = tf.trainable_variables()
        unclipped_grads =tf.gradients(cost, tvars)

        if check_finite:
            checked_grads = []
            for one_grad in unclipped_grads:
                grad_check = tf.logical_and(tf.is_finite(one_grad), tf.logical_not(tf.is_nan(one_grad)))
                checked_zero_grads = tf.where(grad_check, one_grad, tf.zeros(one_grad.shape))
                finite_unclipped_grads = checked_zero_grads
                checked_grads.append(finite_unclipped_grads)
        else:
            checked_grads = unclipped_grads

        if clip_grads:
            grads, _ = tf.clip_by_global_norm(checked_grads, self.config.max_grad_norm)
        else:
            grads = unclipped_grads

        self.eval_op = bp_optimizer.apply_gradients(zip(grads, tvars))

    def initialize_variables(self, assign_dict=None):
        """
        :param transfer_dict: {variable_reference/variable_name_with_scope_name : tensor_value, ......}
                                key can be string, or a tf tensor reference
        :return:
        """
        with self.graph.as_default():
            tf.global_variables_initializer().run(session=self.session)

            if assign_dict is not None:
                ref_assign_dict = {}

                for k, v in assign_dict.items():
                    if isinstance(k, str):
                        tf_vars = tf.trainable_variables()
                        matching_refs = list(filter(lambda x: x.name.__contains__(k), tf_vars))
                        if len(matching_refs) == 1:
                            ref_assign_dict[matching_refs[0]] = v
                        else:
                            print("transfer variable not found in tf variables, or found with more than one variable containing the key, please check the variable names")
                    else:
                        ref_assign_dict[k] = v

                assign_ops = [tf.assign(x, y) for x, y in ref_assign_dict.items()]
                self.session.run(assign_ops)
        return

    def add_to_tensorboard(self, var, watch_name):
        with tf.name_scope("summaries/" + watch_name):
            if isinstance(var, tf.Tensor):
                var_tensor = var
                delay_assign = False
            else:
                if isinstance(var, int):
                    var = float(var)
                var_tensor = tf.get_variable("summary_"+watch_name, initializer=var, dtype=tf.float32)
                delay_assign = True

            type_str = "scalar" if var_tensor.shape.ndims == 0 else "histogram"

            if type_str == "histogram":
                summary_var =  tf.summary.histogram(watch_name, var_tensor)
            else:
                summary_var = tf.summary.scalar(watch_name, var_tensor)

            self.tensorboard_watches[watch_name] = (var_tensor, summary_var, delay_assign)

    def save_best_models(self, iter_num, error, folder_path=None):
        if folder_path is None:
            folder_path = self.config.model_save_path

        def save_best():
            with self.graph.as_default():
                self.saver.save(self.session, folder_path + "/model_"+str(iter_num) + "_"+str(error) + "_.ckpt")

        if not os.path.exists(folder_path + "/measure_record.txt"):
            self.error_lsit = []
            measure_recorder = open(folder_path + "/measure_record.txt", mode="w")
        else:
            fr = open(folder_path + "/measure_record.txt", mode="r")
            lines = fr.read().split("\n")
            fr.close()
            self.error_list = [(int(line.split("\t")[0]), float(line.split("\t")[1])) for line in lines if line != ""]
            measure_recorder = open(folder_path + "/measure_record.txt", mode="w")

        if len(self.error_list) < self.config.save_kept_num:
            save_best()
            self.error_list.append((iter_num, error))
        elif error < max([x[1] for x in self.error_list]):
            threashold = max([x[1] for x in self.error_list])
            self.error_list = list(filter(lambda x: x[1] < threashold, self.error_list))
            save_best()
            self.error_list.append((iter_num, error))
        else:
            pass

        self.error_list = list(sorted(self.error_list, key=lambda x: x[1]))
        measure_recorder.write("\n".join([str(x)+"\t"+str(y) for x, y in self.error_list]))
        measure_recorder.close()
        return

    def load_model(self, iter_num, folder_path=None):

        if folder_path is None:
            folder_path = self.config.model_save_path

        import os
        file_list = os.listdir(folder_path)
        candidate_files = list(filter(lambda x: (len(x.split("_")) > 2) and int((x.split("_"))[1])==iter_num, file_list))
        evaluate_measure = candidate_files[0].split("_")[2]
        with self.graph.as_default():
            self.saver.restore(self.session, folder_path+"/model_"+str(iter_num)+"_"+str(evaluate_measure)+"_.ckpt")

        return

    def process_batch(self, iter_num, feed_dict, *args, segmented_num=None, tensors_requires_segmentation=None, tensorboard_extra_assigns=None):
        """
        :param feed_dict: feed_dict for session.run
        :param args: running tf variables
        :param segmented_num: whether feed batch need to be segmented into smaller batches. int
        :param tensors_requires_segmentation: for each value in feed_dict, indicate whether it should be segmented. list of bool
        :param tensorboard_extra_assigns: dict with key=watch_name, val=assign_value for extra tensorboard watches in runtime
        :return:
        """

        if segmented_num is not None:
            assert len(tensors_requires_segmentation) == len(feed_dict)

        if self.merged_summary is None and len(self.tensorboard_watches) > 0:
            merge_item_list = [x[1] for x in self.tensorboard_watches.values()]
            self.merged_summary = tf.summary.merge(merge_item_list)
        if tensorboard_extra_assigns is not None:
            assign_item_list = [(self.tensorboard_watches[kv[0]][0], kv[1]) for kv in tensorboard_extra_assigns.items()
                                if kv[0] in self.tensorboard_watches and self.tensorboard_watches[kv[0]][2]]
            self.session.run([tf.assign(x[0], x[1]) for x in assign_item_list])

        def run_with_merge(_args, feeds):
            if self.merged_summary is None:
                return self.session.run(args, feed_dict=feeds)
            else:
                run_li = list(args)
                run_li.append(self.merged_summary)
                results = self.session.run(run_li, feed_dict=feeds)
                runned_merged = results[-1]
                self.summary_writer.add_summary(runned_merged, iter_num)
                if iter_num > 0 and iter_num % 20 == 0:
                    self.summary_writer.reopen()
                tensor_results = results[:-1]
                return tensor_results

        with self.graph.as_default():
            if segmented_num is None:
                return run_with_merge(args, feed_dict)
            else:
                total_sample_num = (list(filter(lambda x: x[0], zip(tensors_requires_segmentation, list(feed_dict.values()))))[0])[1].shape[0]
                mini_batch_num = total_sample_num // segmented_num
                residual_batch_num = total_sample_num % segmented_num

                results_list = [[] for _ in range(len(args))]
                for batch_i in range(mini_batch_num):
                    feed_dict_mini = dict([(x, y[segmented_num*batch_i:segmented_num*(batch_i+1)]) for x, y in feed_dict.items()])
                    results_mini = run_with_merge(args, feed_dict_mini)
                    for ri in range(len(results_list)):
                        results_list[ri].append(results_mini[ri])

                if residual_batch_num > 0:
                    feed_dict_mini = dict([(x, y[segmented_num*mini_batch_num:]) for x, y in feed_dict.items()])
                    results_mini = self.session.run(args, feed_dict=feed_dict_mini)
                    for ri in range(len(results_list)):
                        results_list[ri].append(results_mini[ri])

                concatnated_results = [np.concatenate(x, axis=0) for x in results_list]
                return tuple(concatnated_results)

import random
class RandomSampler(object):

    @property
    def item_list(self):
        return self._item_list

    def __init__(self, item_list, shuffle=True):
        self._item_list = item_list
        self._pointer = 0
        self._shuffle = shuffle
        if shuffle:
            random.shuffle(self._item_list)

    def get_batch(self, batch_size, func_on_each, *args, **kwargs):
        batch_result = []
        while(len(batch_result) < batch_size):

            result = func_on_each(self._item_list[self._pointer], *args, **kwargs)
            if result is not None:
                batch_result.append(result)

            self._pointer +=1
            if self._pointer >= len(self._item_list):
                self._pointer = 0
                if self._shuffle:
                    random.shuffle(self._item_list)

        return batch_result


import pandas as pd
import glob
from torch import nn
import json
import pickle
import torch
import os
import sys
from src.models.contrastive_model import Model
from src.utilities.config_parser import get_config
import argparse
from src.models.contrastive_model import get_pretrained_weights_for_transfer
from src.data.raw_transforms import get_raw_transforms_v2, simple_supervised_transforms, PadToSize, PeakNormalization, Compose
from src.data.raw_dataset import RawWaveformDataset
from torch.utils.data import DataLoader
from src.models.model_helper import get_feature_extractor
from src.data.utils import _collate_fn_raw, _collate_fn_raw_multiclass
import tqdm
import numpy as np
from src.utilities import metrics_helper
from sklearn.metrics import average_precision_score
from torch.nn import functional as F
torch.multiprocessing.set_sharing_strategy('file_system')

class BaselineModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        num_classes = cfg['num_classes']
        self.features = get_feature_extractor(cfg)
        self.fc = nn.Linear(cfg['proj_out_dim'], num_classes)

    def forward(self, x):
        x, _ = self.features(x)
        x = F.relu(x, inplace=True)
        x = self.fc(x)
        return x

class FinetunedModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        pretrained_hparams_path = cfg['pretrained_hparams_path']
        pretrained_ckpt_path = cfg['pretrained_ckpt_path']
        num_classes = cfg['num_classes']
        self.features, output_dims = get_pretrained_weights_for_transfer(pretrained_hparams_path,
                                                                         pretrained_ckpt_path)
        self.fc = nn.Linear(output_dims, num_classes)

    def forward(self, x):
        x, _ = self.features(x)
        x = F.relu(x, inplace=True)
        x = self.fc(x)
        return x

def calculate_mAP(preds, gts, mode="weighted"):
    preds = torch.cat(preds, 0).numpy()
    gts = torch.cat(gts, 0).numpy()
    map_value = average_precision_score(gts, preds, average=mode)
    return map_value

parser = argparse.ArgumentParser()
parser.add_argument("--exp_dir", type=str)
parser.add_argument("--contrastive_pretrained", action="store_true")
parser.add_argument('--results_file', type=str)

def get_val_acc(x):
    x = x.split("/")[-1]
    x = x.replace(".pth","")
    x = x.split("val_acc=")[-1]
    return float(x)

if __name__ == '__main__':
    args = parser.parse_args()
    hparams_path = os.path.join(args.exp_dir, "hparams.pickle")
    ckpts = sorted(glob.glob(os.path.join(args.exp_dir, "ckpts", "*")), key=get_val_acc)
    print(ckpts)
    # metrics_path = os.path.join(args.exp_dir, "sorted_metrics.csv")
    # filename = pd.read_csv(metrics_path).iloc[0]['filename']
    # print(filename)
    # filename = filename.split("/")[-1]
    ckpt_path = ckpts[-1]
    print(ckpt_path)
    res_file = os.path.join(args.exp_dir, "results.txt")
    if os.path.exists(res_file):
        print("exists...")
        exit()
    checkpoint = torch.load(ckpt_path)
    with open(hparams_path, "rb") as fp:
        hparams = pickle.load(fp)
    if args.contrastive_pretrained:
        model = FinetunedModel(hparams.cfg['model'])
    else:
        print("making baseline model")
        model = BaselineModel(hparams.cfg['model'])
    model.load_state_dict(checkpoint['model_state_dict'])
    model = model.cuda().eval()
    # val_tfs = get_raw_transforms_v2(False, 8000*10, center_crop_val=False)
    ac = hparams.cfg['audio_config']
    print(ac)
    val_clip_size = int(ac['val_clip_size'] * ac['sample_rate'])
    val_tfs = Compose([
        PadToSize(val_clip_size, 'wrap'),
        PeakNormalization(sr=ac['sample_rate'])
    ])
    # val_tfs = simple_supervised_transforms(False, val_clip_size,
    #                                        sample_rate=ac['sample_rate'])
    val_set = RawWaveformDataset("/media/user/nvme/datasets/fsd50k/fsd50k_8000/meta/eval.csv",
                                 "/media/user/nvme/datasets/fsd50k/fsd50k_8000/meta/lbl_map.json",
                                 hparams.cfg['audio_config'], mode='multilabel',
                                 transform=val_tfs, is_val=True
                                 )

    with open("/media/user/nvme/datasets/fsd50k/fsd50k_8000/meta/lbl_map.json", "r") as fd:
        lbl_map = json.load(fd)

    inv_map = {v: k for k, v in lbl_map.items()}

    val_loader = DataLoader(val_set, sampler=None, num_workers=4,
                            collate_fn=_collate_fn_raw,
                            shuffle=False, batch_size=1,
                            pin_memory=False)

    val_preds = []
    val_gts = []
    for batch in tqdm.tqdm(val_loader):
        x, _, y = batch
        x = x.cuda()
        with torch.no_grad():
            pred = model(x)
            y_pred_sigmoid = torch.sigmoid(pred)
        val_preds.append(y_pred_sigmoid.detach().cpu().float())
        val_gts.append(y.detach().cpu().float())

    macro_mAP = calculate_mAP(val_preds, val_gts, mode="macro")
    weighted_mAP = calculate_mAP(val_preds, val_gts, mode="weighted")
    print("macro_mAP: {} | weighted_mAP: {}".format(macro_mAP, weighted_mAP))

    val_preds = torch.cat(val_preds, 0).numpy()
    val_gts = torch.cat(val_gts, 0).numpy()

    stats = metrics_helper.calculate_stats(val_preds, val_gts)
    mAP = np.mean([stat['AP'] for stat in stats])
    mAUC = np.mean([stat['auc'] for stat in stats])
    dp = metrics_helper.d_prime(mAUC)
    print("mAP: {} | mAUC: {} | dprime: {}".format(mAP, mAUC, dp))
    fname = ckpt_path.split("/")[-3]

    ckpt_ext = "/".join(ckpt_path.split("/")[-3:])
    with open(res_file, "w") as fd:
        fd.writelines("model,mAP,mAUC,dprime,ckpt_ext\n")
        fd.writelines("{},{},{},{},{}\n".format(fname, macro_mAP, mAUC, dp, ckpt_ext))


# # Unity ML-Agents Toolkit
# ## ML-Agents Learning (POCA)
# Contains an implementation of MA-POCA.

from collections import defaultdict
from typing import cast, Dict

import numpy as np

from mlagents_envs.side_channel.stats_side_channel import StatsAggregationMethod
from mlagents_envs.logging_util import get_logger
from mlagents_envs.base_env import BehaviorSpec
from mlagents.trainers.buffer import BufferKey, RewardSignalUtil
from mlagents.trainers.trainer.rl_trainer import RLTrainer
from mlagents.trainers.policy import Policy
from mlagents.trainers.policy.torch_policy import TorchPolicy
from mlagents.trainers.poca.optimizer_torch import TorchPOCAOptimizer
from mlagents.trainers.trajectory import Trajectory
from mlagents.trainers.behavior_id_utils import BehaviorIdentifiers
from mlagents.trainers.settings import TrainerSettings, POCASettings

logger = get_logger(__name__)

class POCATrainer(RLTrainer):
    """The POCATrainer is an implementation of the MA-POCA algorithm."""

    def __init__(
        self,
        behavior_name: str,
        reward_buff_cap: int,
        trainer_settings: TrainerSettings,
        training: bool,
        load: bool,
        seed: int,
        artifact_path: str,
    ):
        """
        Responsible for collecting experiences and training POCA model.
        :param behavior_name: The name of the behavior associated with trainer config
        :param reward_buff_cap: Max reward history to track in the reward buffer
        :param trainer_settings: The parameters for the trainer.
        :param training: Whether the trainer is set for training.
        :param load: Whether the model should be loaded.
        :param seed: The seed the model will be initialized with
        :param artifact_path: The directory within which to store artifacts from this trainer.
        """
        super().__init__(
            behavior_name,
            trainer_settings,
            training,
            load,
            artifact_path,
            reward_buff_cap,
        )
        self.hyperparameters: POCASettings = cast(
            POCASettings, self.trainer_settings.hyperparameters
        )
        self.seed = seed
        self.policy: TorchPolicy = None  # type: ignore
        self.collected_group_rewards: Dict[str, int] = defaultdict(lambda: 0)

    def _process_trajectory(self, trajectory: Trajectory) -> None:
        """
        Takes a trajectory and processes it, putting it into the update buffer.
        Processing involves calculating value and advantage targets for model updating step.
        :param trajectory: The Trajectory tuple containing the steps to be processed.
        """
        super()._process_trajectory(trajectory)
        agent_id = trajectory.agent_id  # All the agents should have the same ID

        agent_buffer_trajectory = trajectory.to_agentbuffer()
        # Update the normalization
        if self.is_training:
            self.policy.update_normalization(agent_buffer_trajectory)

        # Get all value estimates
        (
            value_estimates,
            baseline_estimates,
            value_next,
            value_memories,
            baseline_memories,
        ) = self.optimizer.get_trajectory_and_baseline_value_estimates(
            agent_buffer_trajectory,
            trajectory.next_obs,
            trajectory.next_group_obs,
            trajectory.all_group_dones_reached
            and trajectory.done_reached
            and not trajectory.interrupted,
        )

        if value_memories is not None and baseline_memories is not None:
            agent_buffer_trajectory[BufferKey.CRITIC_MEMORY].set(value_memories)
            agent_buffer_trajectory[BufferKey.BASELINE_MEMORY].set(baseline_memories)

        for name, v in value_estimates.items():
            agent_buffer_trajectory[RewardSignalUtil.value_estimates_key(name)].extend(
                v
            )
            agent_buffer_trajectory[
                RewardSignalUtil.baseline_estimates_key(name)
            ].extend(baseline_estimates[name])
            self._stats_reporter.add_stat(
                f"Policy/{self.optimizer.reward_signals[name].name.capitalize()} Baseline Estimate",
                np.mean(baseline_estimates[name]),
            )
            self._stats_reporter.add_stat(
                f"Policy/{self.optimizer.reward_signals[name].name.capitalize()} Value Estimate",
                np.mean(value_estimates[name]),
            )

        self.collected_rewards["environment"][agent_id] += np.sum(
            agent_buffer_trajectory[BufferKey.ENVIRONMENT_REWARDS]
        )
        self.collected_group_rewards[agent_id] += np.sum(
            agent_buffer_trajectory[BufferKey.GROUP_REWARD]
        )
        for name, reward_signal in self.optimizer.reward_signals.items():
            evaluate_result = (
                reward_signal.evaluate(agent_buffer_trajectory) * reward_signal.strength
            )
            agent_buffer_trajectory[RewardSignalUtil.rewards_key(name)].extend(
                evaluate_result
            )
            # Report the reward signals
            self.collected_rewards[name][agent_id] += np.sum(evaluate_result)

        # Compute lambda returns and advantage
        tmp_advantages = []
        for name in self.optimizer.reward_signals:

            local_rewards = np.array(
                agent_buffer_trajectory[RewardSignalUtil.rewards_key(name)].get_batch(),
                dtype=np.float32,
            )

            baseline_estimate = agent_buffer_trajectory[
                RewardSignalUtil.baseline_estimates_key(name)
            ].get_batch()
            v_estimates = agent_buffer_trajectory[
                RewardSignalUtil.value_estimates_key(name)
            ].get_batch()

            lambd_returns = lambda_return(
                r=local_rewards,
                value_estimates=v_estimates,
                gamma=self.optimizer.reward_signals[name].gamma,
                lambd=self.hyperparameters.lambd,
                value_next=value_next[name],
            )

            local_advantage = np.array(lambd_returns) - np.array(baseline_estimate)

            agent_buffer_trajectory[RewardSignalUtil.returns_key(name)].set(
                lambd_returns
            )
            agent_buffer_trajectory[RewardSignalUtil.advantage_key(name)].set(
                local_advantage
            )
            tmp_advantages.append(local_advantage)

        # Get global advantages
        global_advantages = list(
            np.mean(np.array(tmp_advantages, dtype=np.float32), axis=0)
        )
        agent_buffer_trajectory[BufferKey.ADVANTAGES].set(global_advantages)

        self._append_to_update_buffer(agent_buffer_trajectory)

        # If this was a terminal trajectory, append stats and reset reward collection
        if trajectory.done_reached:
            self._update_end_episode_stats(agent_id, self.optimizer)
            # Remove dead agents from group reward recording
            if not trajectory.all_group_dones_reached:
                self.collected_group_rewards.pop(agent_id)

        # If the whole team is done, average the remaining group rewards.
        if trajectory.all_group_dones_reached and trajectory.done_reached:
            self.stats_reporter.add_stat(
                "Environment/Group Cumulative Reward",
                self.collected_group_rewards.get(agent_id, 0),
                aggregation=StatsAggregationMethod.HISTOGRAM,
            )
            self.collected_group_rewards.pop(agent_id)

    def _is_ready_update(self):
        """
        Returns whether or not the trainer has enough elements to run update model
        :return: A boolean corresponding to whether or not update_model() can be run
        """
        size_of_buffer = self.update_buffer.num_experiences
        return size_of_buffer > self.hyperparameters.buffer_size

    def _update_policy(self):
        """
        Uses demonstration_buffer to update the policy.
        The reward signal generators must be updated in this method at their own pace.
        """
        buffer_length = self.update_buffer.num_experiences
        self.cumulative_returns_since_policy_update.clear()

        # Make sure batch_size is a multiple of sequence length. During training, we
        # will need to reshape the data into a batch_size x sequence_length tensor.
        batch_size = (
            self.hyperparameters.batch_size
            - self.hyperparameters.batch_size % self.policy.sequence_length
        )
        # Make sure there is at least one sequence
        batch_size = max(batch_size, self.policy.sequence_length)

        n_sequences = max(
            int(self.hyperparameters.batch_size / self.policy.sequence_length), 1
        )

        advantages = np.array(
            self.update_buffer[BufferKey.ADVANTAGES].get_batch(), dtype=np.float32
        )
        self.update_buffer[BufferKey.ADVANTAGES].set(
            (advantages - advantages.mean()) / (advantages.std() + 1e-10)
        )
        num_epoch = self.hyperparameters.num_epoch
        batch_update_stats = defaultdict(list)
        for _ in range(num_epoch):
            self.update_buffer.shuffle(sequence_length=self.policy.sequence_length)
            buffer = self.update_buffer
            max_num_batch = buffer_length // batch_size
            for i in range(0, max_num_batch * batch_size, batch_size):
                update_stats = self.optimizer.update(
                    buffer.make_mini_batch(i, i + batch_size), n_sequences
                )
                for stat_name, value in update_stats.items():
                    batch_update_stats[stat_name].append(value)

        for stat, stat_list in batch_update_stats.items():
            self._stats_reporter.add_stat(stat, np.mean(stat_list))

        if self.optimizer.bc_module:
            update_stats = self.optimizer.bc_module.update()
            for stat, val in update_stats.items():
                self._stats_reporter.add_stat(stat, val)
        self._clear_update_buffer()
        return True

    def end_episode(self) -> None:
        """
        A signal that the Episode has ended. The buffer must be reset.
        Get only called when the academy resets. For POCA, we should
        also zero out the group rewards.
        """
        super().end_episode()
        self.collected_group_rewards.clear()

    def create_torch_policy(
        self, parsed_behavior_id: BehaviorIdentifiers, behavior_spec: BehaviorSpec
    ) -> TorchPolicy:
        """
        Creates a policy with a PyTorch backend and POCA hyperparameters
        :param parsed_behavior_id:
        :param behavior_spec: specifications for policy construction
        :return policy
        """
        policy = TorchPolicy(
            self.seed,
            behavior_spec,
            self.trainer_settings,
            condition_sigma_on_obs=False,  # Faster training for POCA
            separate_critic=True,  # Match network architecture with TF
        )
        return policy

    def create_poca_optimizer(self) -> TorchPOCAOptimizer:
        return TorchPOCAOptimizer(self.policy, self.trainer_settings)

    def add_policy(
        self, parsed_behavior_id: BehaviorIdentifiers, policy: Policy
    ) -> None:
        """
        Adds policy to trainer.
        :param parsed_behavior_id: Behavior identifiers that the policy should belong to.
        :param policy: Policy to associate with name_behavior_id.
        """
        if not isinstance(policy, TorchPolicy):
            raise RuntimeError(f"policy {policy} must be an instance of TorchPolicy.")
        self.policy = policy
        self.policies[parsed_behavior_id.behavior_id] = policy
        self.optimizer = self.create_poca_optimizer()
        for _reward_signal in self.optimizer.reward_signals.keys():
            self.collected_rewards[_reward_signal] = defaultdict(lambda: 0)

        self.model_saver.register(self.policy)
        self.model_saver.register(self.optimizer)
        self.model_saver.initialize_or_load()

        # Needed to resume loads properly
        self._step = policy.get_current_step()

    def get_policy(self, name_behavior_id: str) -> Policy:
        """
        Gets policy from trainer associated with name_behavior_id
        :param name_behavior_id: full identifier of policy
        """

        return self.policy

def lambda_return(r, value_estimates, gamma=0.99, lambd=0.8, value_next=0.0):
    returns = np.zeros_like(r)
    returns[-1] = r[-1] + gamma * value_next
    for t in reversed(range(0, r.size - 1)):
        returns[t] = (
            gamma * lambd * returns[t + 1]
            + r[t]
            + (1 - lambd) * gamma * value_estimates[t + 1]
        )
    return returns


#!/usr/bin/env python

import os, rospy
import torch
import torch.nn.functional as F
import numpy as np
from buffer import ReplayBuffer
from networks import ActorNetwork, CriticNetwork, ValueNetwork

class Agent():
    def __init__(self, alpha=0.001, beta=0.001, input_dims=[8],
            env=None, gamma=0.99, n_actions=2, max_size=1000000, tau=0.005,
            layer1_size=256, layer2_size=256, batch_size=256, reward_scale=2, auto_entropy=False,
            entr_lr=None, reparam_noise=1e-6):
        self.gamma = gamma
        self.tau = tau
        self.memory = ReplayBuffer(max_size, input_dims, n_actions)
        self.batch_size = batch_size
        self.n_actions = n_actions
        self.max_action = env.max_action
        self.env = env
        self.auto_entropy = auto_entropy
        self.actor = ActorNetwork(alpha, input_dims, fc1_dims=layer1_size, fc2_dims=layer2_size, n_actions=n_actions,
                    name='actor', max_action=self.max_action)
        self.critic_1 = CriticNetwork(beta, input_dims, fc1_dims=layer1_size, fc2_dims=layer2_size, n_actions=n_actions,
                    name='critic_1')
        self.critic_2 = CriticNetwork(beta, input_dims,fc1_dims=layer1_size, fc2_dims=layer2_size, n_actions=n_actions,
                    name='critic_2')
        self.value = ValueNetwork(beta, input_dims, fc1_dims=layer1_size, fc2_dims=layer2_size, name='value')
        self.target_value = ValueNetwork(beta, input_dims,fc1_dims=layer1_size, fc2_dims=layer2_size, name='target_value')

        self.scale = reward_scale
        self.update_network_parameters(tau=1)
        self.target_entropy = -np.prod(self.env.action_space).astype(np.float32)
        self.log_alpha = torch.zeros(1, requires_grad=True, device=('cuda:0' if torch.cuda.is_available() else 'cpu'))
        if entr_lr is None:
            self.entr_lr = alpha
        else:
            self.entr_lr = entr_lr
        self.alpha_optim = torch.optim.Adam([self.log_alpha], lr=alpha)
        self.entropy = 0

    def choose_action(self, observation):
        state = torch.Tensor([observation]).to(self.actor.device)
        actions, _ = self.actor.sample_normal(state, reparameterize=False)
        return actions.cpu().detach().numpy()[0]

    def remember(self, state, action, reward, new_state, done):
        self.memory.store_transition(state, action, reward, new_state, done)
        return 1

    def update_network_parameters(self, tau=None):
        if tau is None:
            tau = self.tau

        target_value_params = self.target_value.named_parameters()
        value_params = self.value.named_parameters()

        target_value_state_dict = dict(target_value_params)
        value_state_dict = dict(value_params)

        for name in value_state_dict:
            value_state_dict[name] = tau*value_state_dict[name].clone() + \
                    (1-tau)*target_value_state_dict[name].clone()

        self.target_value.load_state_dict(value_state_dict)

    def save_models(self):
        print('.... saving models ....')
        self.actor.save_checkpoint()
        self.value.save_checkpoint()
        self.target_value.save_checkpoint()
        self.critic_1.save_checkpoint()
        self.critic_2.save_checkpoint()

    def load_models(self):
        print('.... loading models ....')
        self.actor.load_checkpoint()
        self.value.load_checkpoint()
        self.target_value.load_checkpoint()
        self.critic_1.load_checkpoint()
        self.critic_2.load_checkpoint()

    def learn(self, update_params=True):
        print("LEARNING")
        if(self.auto_entropy):
            if self.memory.mem_cntr < self.batch_size:
                return

            state, action, reward, new_state, done = \
                    self.memory.sample_buffer(self.batch_size)

            reward = torch.tensor(reward, dtype=torch.float).to(self.actor.device)
            done = torch.tensor(done).to(self.actor.device)
            state_ = torch.tensor(new_state, dtype=torch.float).to(self.actor.device)
            state = torch.tensor(state, dtype=torch.float).to(self.actor.device)
            action = torch.tensor(action, dtype=torch.float).to(self.actor.device)

            value = self.value(state).view(-1)
            value_ = self.target_value(state_).view(-1)
            value_[done] = 0.0

            actions, log_probs = self.actor.sample_normal(state, reparameterize=False)
            log_probs = log_probs.view(-1)
            q1_new_policy = self.critic_1.forward(state, actions)
            q2_new_policy = self.critic_2.forward(state, actions)
            critic_value = torch.min(q1_new_policy, q2_new_policy)
            critic_value = critic_value.view(-1)

            self.value.optimizer.zero_grad()
            value_target = critic_value - self.entropy*log_probs
            value_loss = 0.5 * F.mse_loss(value, value_target)
            value_loss.backward(retain_graph=True)
            self.value.optimizer.step()

            actions, log_probs = self.actor.sample_normal(state, reparameterize=True)
            log_probs = log_probs.view(-1)
            q1_new_policy = self.critic_1.forward(state, actions)
            q2_new_policy = self.critic_2.forward(state, actions)
            critic_value = torch.min(q1_new_policy, q2_new_policy)
            critic_value = critic_value.view(-1)

            actor_loss = self.entropy*log_probs - critic_value
            actor_loss = torch.mean(actor_loss)
            self.actor.optimizer.zero_grad()
            actor_loss.backward(retain_graph=True)
            self.actor.optimizer.step()

            self.critic_1.optimizer.zero_grad()
            self.critic_2.optimizer.zero_grad()
            q_hat = self.scale*reward + self.gamma*value_
            q1_old_policy = self.critic_1.forward(state, action).view(-1)
            q2_old_policy = self.critic_2.forward(state, action).view(-1)
            critic_1_loss = 0.5 * F.mse_loss(q1_old_policy, q_hat)
            critic_2_loss = 0.5 * F.mse_loss(q2_old_policy, q_hat)

            critic_loss = critic_1_loss + critic_2_loss
            critic_loss.backward()
            self.critic_1.optimizer.step()
            self.critic_2.optimizer.step()

            alpha_loss = -(self.log_alpha * (log_probs + self.target_entropy).detach()).mean()
            self.alpha_optim.zero_grad()
            alpha_loss.backward()
            self.alpha_optim.step()
            self.entropy = self.log_alpha.exp()

            if update_params:
                self.update_network_parameters()
        else:
            if self.memory.mem_cntr < self.batch_size:
                return

            state, action, reward, new_state, done = \
                    self.memory.sample_buffer(self.batch_size)

            reward = torch.tensor(reward, dtype=torch.float).to(self.actor.device)
            done = torch.tensor(done).to(self.actor.device)
            state_ = torch.tensor(new_state, dtype=torch.float).to(self.actor.device)
            state = torch.tensor(state, dtype=torch.float).to(self.actor.device)
            action = torch.tensor(action, dtype=torch.float).to(self.actor.device)

            value = self.value(state).view(-1)
            value_ = self.target_value(state_).view(-1)
            value_[done] = 0.0

            actions, log_probs = self.actor.sample_normal(state, reparameterize=False)
            log_probs = log_probs.view(-1)
            q1_new_policy = self.critic_1.forward(state, actions)
            q2_new_policy = self.critic_2.forward(state, actions)
            critic_value = torch.min(q1_new_policy, q2_new_policy)
            critic_value = critic_value.view(-1)

            self.value.optimizer.zero_grad()
            value_target = critic_value - log_probs
            value_loss = 0.5 * F.mse_loss(value, value_target)
            value_loss.backward(retain_graph=True)
            self.value.optimizer.step()

            actions, log_probs = self.actor.sample_normal(state, reparameterize=True)
            log_probs = log_probs.view(-1)
            q1_new_policy = self.critic_1.forward(state, actions)
            q2_new_policy = self.critic_2.forward(state, actions)
            critic_value = torch.min(q1_new_policy, q2_new_policy)
            critic_value = critic_value.view(-1)

            actor_loss = log_probs - critic_value
            actor_loss = torch.mean(actor_loss)
            self.actor.optimizer.zero_grad()
            actor_loss.backward(retain_graph=True)
            self.actor.optimizer.step()

            self.critic_1.optimizer.zero_grad()
            self.critic_2.optimizer.zero_grad()
            q_hat = self.scale*reward + self.gamma*value_
            q1_old_policy = self.critic_1.forward(state, action).view(-1)
            q2_old_policy = self.critic_2.forward(state, action).view(-1)
            critic_1_loss = 0.5 * F.mse_loss(q1_old_policy, q_hat)
            critic_2_loss = 0.5 * F.mse_loss(q2_old_policy, q_hat)

            critic_loss = critic_1_loss + critic_2_loss
            critic_loss.backward()
            self.critic_1.optimizer.step()
            self.critic_2.optimizer.step()

            if update_params:
                self.update_network_parameters()


#!/usr/bin/env python
"""
The :mod:`hotspots.calculation` handles the main Fragment Hotspot Maps algorithm. In addition, an alternative pocket burial method, Ghecom, is provided.

The main classes of the :mod:`hotspots.calculation` module are:

- :class:`hotspots.calculation.Buriedness`
- :class:`hotspots.calculation.Runner`

More information about the Fragment Hotspot Maps method is available from:
    - Radoux, C.J. et. al., Identifying the Interactions that Determine Fragment Binding at Protein Hotspots J. Med. Chem. 2016, 59 (9), 4314-4325 [dx.doi.org/10.1021/acs.jmedchem.5b01980]

More information about the Ghecom method is available from:
    - Kawabata T, Go N. Detection of pockets on protein surfaces using small and large probe spheres to find putative ligand binding sites. Proteins 2007; 68: 516-529
"""
from __future__ import print_function, division

import multiprocessing
import operator
import random
import shutil
import sys
import tempfile
import time
from functools import reduce
from os import system, environ
from os.path import join

import numpy as np
import pkg_resources

from scipy import ndimage
from skimage.morphology import ball
from skimage.transform import resize
# from hotspots.protoss import Protoss
from tqdm import tqdm

# Development note. We import the CCDC modules after 3rd part modules where possible as
# we generally find that the ccdc package is less fussy about the underlying compiled
# libraries used.
#

from ccdc.cavity import Cavity
from ccdc.io import MoleculeWriter, MoleculeReader
from ccdc.molecule import Molecule, Coordinates
from ccdc.protein import Protein
from ccdc.utilities import PushDir

from hotspots.atomic_hotspot_calculation import _AtomicHotspot
from hotspots.grid_extension import Grid
from hotspots.hs_utilities import Helper
from hotspots.wrapper_pdb import PDBResult
from hotspots.wrapper_ghecom import Ghecom
from hotspots.result import Results

class ExpBuriedness(object):

    def __init__(self, prot, out_grid, max_probe_radius=11):
        self.prot = prot
        self.out_grid = out_grid
        self.max_probe = max_probe_radius
        self.probe_selem_dict = self._generate_probe_selem()
        self.protein_grid = self._from_molecule(prot)
        if self.out_grid is None:
            self.out_grid = self.protein_grid.copy_and_clear()
        print("buriedness init")

    def _generate_probe_selem(self):
        """
        Generates a set of structuring elements (i.e. probe spheres) starting from a 3x3 cross (0.5 A radius) probe up
        to the max probe radius
        :return:
        """
        probe_selem_dict = {}

        # Small probe

        probe_selem_dict[2] = ball(2)
        # r= probe radius
        for r in range(3,self.max_probe +1):
            probe_selem_dict[r] = ball(r)
        print("generating probes")
        return probe_selem_dict

    def _from_molecule(self, mol, scaling=1):
        """
        generate a molecule mask where gp within the vdw radius of the molecule heavy atoms are set to 1.0
        :param mol: `ccdc.molecule.Molecule`
        :param padding: int
        :param scaling: float
        :return: `hotspots.grid_extension.Grid`
        """

        coords = [a.coordinates for a in mol.atoms]
        g = Grid.initalise_grid(coords=coords, padding= 10, spacing=1)

        for probe in sorted(self.probe_selem_dict.keys(), reverse=True):
            for a in mol.heavy_atoms:
                g.set_sphere(point=a.coordinates,
                             radius=probe * scaling,
                             value=probe,
                             mode='replace',
                             scaling='None')

        for a in mol.heavy_atoms:
            g.set_sphere(point=a.coordinates,
                         radius=a.vdw_radius,
                         value=100,
                         mode='replace',
                         scaling='None')

        try:
            out_bound_box = self.out_grid.bounding_box
            print(out_bound_box)
            origin_indices = g.point_to_indices(out_bound_box[0])
            far_indices = g.point_to_indices(out_bound_box[1])
            region = origin_indices + far_indices
            print(region)
            return g.sub_grid(region)
        except AttributeError:
            return g

    def _close_grid(self, g, probe):

        g_array = g.get_array().astype(int)
        closed_array = ndimage.binary_erosion(g_array, structure=self.probe_selem_dict[probe])
        return Grid.array_to_grid(closed_array.astype(int), g)

    def _multiscale_closing(self,g):

        #probe_sizes.reverse()
        all_g = None
        prot_mask = g > 90

        for probe in sorted(self.probe_selem_dict.keys())[1:]:
            if all_g is None:
                all_g = prot_mask.copy()

            probe_mask = ((g < probe+0.1) & (g>0)) | (prot_mask)

            probe_mask = self._close_grid(probe_mask,probe)
            novel = (-all_g) * (g>0)
            all_g += (probe_mask * novel * (self.max_probe - probe))

        return all_g * (prot_mask < 1)

    def _open_grid(self, g, probe):

        g_array = g.get_array().astype(int)
        opened_array = ndimage.binary_opening(g_array, structure=self.probe_selem_dict[probe])
        return Grid.array_to_grid(opened_array.astype(int), g)

    def buriedness_grid(self):

        closed_g = self._multiscale_closing(self.protein_grid)
        print("close_g")
        out_g = self._open_grid(closed_g,2) * closed_g
        print("out_g")
        out_array = out_g.get_array()
        print("out array")
        scaled_g = Grid.initalise_grid(self.out_grid.bounding_box, padding=0, spacing=0.5)
        print("scaled_g")
        scaled_array = resize(out_array, scaled_g.nsteps ,anti_aliasing=False)
        print(scaled_array)

        # Future tweaking here
        final_array = scaled_array
        print("array to grid")
        return Grid.array_to_grid(final_array.astype(int), scaled_g)

class _WeightedResult(object):
    """
    private class

    A class to handle the weighted grid results

    :param str identifier: identifier, default is the probe identifier assigned at the Atomic Hotspot Calculation stage
    :param `ccdc.utilities.Grid` grid: result grid
    """

    def __init__(self, identifier, grid):
        self.identifier = identifier
        self.grid = grid

class _SampleGrid(object):
    """
    private class

    class to handle sampled grids

    :param name: str, name of probe (donor, acceptor, apolar, positive, negative)
    :param grid: a :class: `ccdc.utilities.Grid` instance
    :param atom_predicate: atom_predicate will be used to select atoms of a molecule for sampling
    """

    def __init__(self, name, grid, atom_predicate):
        self.name = name
        self.grid = grid
        self.atom_predicate = atom_predicate
        self.mol = None

    @staticmethod
    def add_coordinates(coord, trans):
        """
        provides a coordinate list of atoms to be scored be scored in the SampleGrid
        :param coord: tup, (float(x), float(y), float(z)), set of atomic coordinates for "active" coordinates
        :param trans: tup, (float(x), float(y), float(z)), set of translations to translate probes to points
        above threshold
        :return: list of tup
        """

        return [coord[i] + trans[i] for i in range(0, len(trans))]

    def sample(self, coordinate_list, trans):
        """
        score Molecule in grids for which it has active atoms
        :param coordinate_list: list, set of coordinates of translated and rotated probe atoms to be scored
        :param trans:
        :return:
        """

        try:
            return [self.grid._grid.value(self.add_coordinates(c, trans)) for c in coordinate_list]
        except RuntimeError:
            return [0]

    def set_molecule(self, mol, polar_contribution):
        """
        set which atoms match the grid
        probes with polar atoms contribute do not contribute to apolar maps as this leads to artefacts

        :param mol:
        :param polar_contribution:
        :return:
        """
        self.mol = mol
        if self.name == 'apolar' and not polar_contribution and len([a for a in mol.atoms
                                                                     if a.is_donor or a.is_acceptor]) > 0:
            self._active_atoms = []
        elif not hasattr(self, '_active_atoms'):
            self._active_atoms = [a for a in self.mol.atoms if self.atom_predicate(a)]

    @staticmethod
    def is_donor(a):
        """
        returns true if a given atom is a donor

        :param a: a `ccdc.molecule.Atom` instance
        :return: bool, true if the atom classification is "donor"
        """

        if a.is_donor and a.formal_charge == 0:
            return True
        else:
            return False

    @staticmethod
    def is_acceptor(a):
        """
        returns true if a given atom is a acceptor

        :param a: a `ccdc.molecule.Atom` instance
        :return: bool, true if the atom classification is "acceptor"
        """
        if a.is_acceptor and a.formal_charge == 0:
            return True
        else:
            return False

    @staticmethod
    def is_apolar(a):
        """
        returns true if a given atom is a apolar

        :param a: a `ccdc.molecule.Atom` instance
        :return: bool, true if the atom classification is "apolar"
        """
        if a.is_donor or a.is_acceptor or a.formal_charge != 0 or a.atomic_symbol == "Xe":
            return False
        else:
            return True

    @staticmethod
    def is_positive(a):
        """
        returns true if a given atom is a positively charged

        :param a: a `ccdc.molecule.Atom` instance
        :return: bool, true if the atom is positively charged
        """
        return a.formal_charge > 0

    @staticmethod
    def is_negative(a):
        """
        returns true if a given atom is a negatively charged

        :param a: a `ccdc.molecule.Atom` instance
        :return: bool, true if the atom is negatively charged
        """
        return a.formal_charge < 0

    @staticmethod
    def is_aromatic(a):
        """
        returns true if a given atom is aromatic

        :param a: a `ccdc.molecule.Atom` instance
        :return: bool, true if the atom is aromatic
        """
        return a.atomic_symbol == 'C' and a.is_cyclic and any(b.atom_type == 'aromatic' for b in a.bonds)

class Runner(object):
    """
    A class for running the Fragment Hotspot Map calculation
    """

    class Settings(object):
        """
        adjusts the default settings for the calculation

        :param int nrotations: number of rotations (keep it below 10**6)
        :param float apolar_translation_threshold: translate probe to grid points above this threshold. Give lower values for greater sampling. Default 15
        :param float polar_translation_threshold: translate probe to grid points above this threshold. Give lower values for greater sampling. Default 15
        :param bool polar_contributions: allow carbon atoms of probes with polar atoms to contribute to the apolar output map.
        :param bool return_probes: Generate a sorted list of molecule objects, corresponding to probe poses
        :param bool sphere_maps: When setting the probe score on the output maps, set it for a sphere (radius 1.5) instead of a single point.
        """

        def __init__(self, nrotations=3000, apolar_translation_threshold=15, polar_translation_threshold=15,
                     polar_contributions=False, return_probes=False, sphere_maps=False):
            self.nrotations = nrotations
            self.apolar_translation_threshold = apolar_translation_threshold
            self.polar_translation_threshold = polar_translation_threshold
            self.polar_contributions = polar_contributions
            self.return_probes = return_probes
            self.sphere_maps = sphere_maps

        @property
        def _num_gp(self):
            """
            number of grid point for a given volume
            :return:
            """
            return int(float(500) / self.nrotations ** 3)

    class _Sampler(object):
        """
        Samples one or more grids with a probe molecule
        """

        def __init__(self, *grids, **kw):
            """
            Settings used to run fragment-hotspot-maps script

            :param grids: list, list of :class:`ccdc.utilities.Grid` instances
            :param kw: 'settings'
            """

            self.grids = grids
            self.grid_dic = {}
            for g in grids:
                self.grid_dic[g.name] = g
            settings = kw.get('settings')
            if settings is None:
                settings = self.Settings()

            self.settings = settings
            self.probe_grids = [_SampleGrid(g.name, g.grid.copy_and_clear(), g.atom_predicate) for g in self.grids]

        def get_priority_atom(self, molecule):
            """
            Select priority atom. Select polar atom. If multiple polar atoms, select the one furthest from the centre of
            geometry. If no polar atoms, select atom furthest from centre of geometry

            :param molecule: a :class: `ccdc.molecule.Molecule` instance
            :return: a :class: `ccdc.molecule.Molecule` instance, str atom type
            """
            c = molecule.centre_of_geometry()
            polar_atoms = [a for a in molecule.atoms if a.is_donor or a.is_acceptor]
            atom_by_distance = {}
            if len(polar_atoms) > 0:
                for a in polar_atoms:
                    d = Helper.get_distance(c, a.coordinates)
                    atom_by_distance[d] = a
            else:
                for a in molecule.atoms:
                    d = Helper.get_distance(c, a.coordinates)
                    atom_by_distance[d] = a

            greatest_distance = sorted(atom_by_distance.keys())[0]
            priority_atom = atom_by_distance[greatest_distance]

            pa_type = None
            if priority_atom.formal_charge != 0:
                if priority_atom.formal_charge < 0:
                    pa_type = "acceptor"
                elif priority_atom.formal_charge > 0:
                    pa_type = "donor"
            else:
                if priority_atom.is_acceptor:
                    pa_type = "acceptor"
                elif priority_atom.is_donor:
                    pa_type = "donor"
                else:
                    pa_type = "apolar"

            return priority_atom, pa_type

        def get_translation_points(self, priority_atom_type):
            """
            returns a list of coordinates that are greater than the threshold, that the probe will be translated to

            :param priority_atom_type: str, atomic interaction type
            :return: list, list of :class: `ccdc.molecule.Molecule` instances
            """
            translate_probe = []
            wg = self.grid_dic[priority_atom_type]

            if priority_atom_type == 'apolar':
                translation_threshold = self.settings.apolar_translation_threshold
            else:
                translation_threshold = self.settings.polar_translation_threshold
            hs = wg.grid.islands(translation_threshold)
            for g in hs:
                nx, ny, nz = g.nsteps

                maxima = [g.indices_to_point(i, j, k)
                          for i in range(nx) for j in range(ny) for k in range(nz)
                          if g.value(i, j, k) >= translation_threshold]

                translate_probe = translate_probe + maxima
            return translate_probe

        def generate_rand_quaternions(self):
            """
            Returns a list of random quaternions. Length matches settings.nrotations

            :return: tup, (a,b,c,d)
            """
            quaternions = []
            i = 1
            if self.settings.nrotations > 1:
                while i <= self.settings.nrotations:
                    r1 = random.uniform(-1, 1)
                    r2 = random.uniform(-1, 1)
                    s1 = r1 * r1 + r2 * r2
                    if s1 < 1:
                        r3 = random.uniform(-1, 1)
                        r4 = random.uniform(-1, 1)

                        s2 = r3 * r3 + r4 * r4
                        if s2 < 1:
                            q = (r1, r2, r3 * (np.sqrt((1 - s1) / s2)), r4 * (np.sqrt((1 - s1) / s2)))
                            quaternions.append(q)

                            i += 1

            return quaternions

        @staticmethod
        def score(values):
            """
            Calculate geometric mean of scores

            :param values: float, scores of atoms in probe
            :return: float, geometric mean of probe atom scores
            """

            return reduce(operator.__mul__, values, 1.0) ** (1. / len(values))
            # return sum(values)

        def sample_pose(self, trans, active_atoms_dic, probe):
            """
            Return a pose score (as defined by the score(self,dic) function) and a dictionary of atom:scores

            :param trans: list of translations
            :param active_atoms_dic: dict {"interaction type": "atoms to be scored"}
            :param probe: str, interaction_type
            :return:
            """
            if probe == "negative" or probe == "positive":
                atoms = [g.mol.atoms for g in self.grids if g.name == "positive"][0]
                weight = int(6 / len([a for a in atoms if str(a.atomic_symbol) == "C"]))

                apolar_values = [g.sample(active_atoms_dic[g.name], trans) for g in self.grids if
                                 len(active_atoms_dic[g.name]) > 0 and g.name == "apolar"] * weight
                charged_values = [g.sample(active_atoms_dic[g.name], trans) for g in self.grids if
                                  len(active_atoms_dic[g.name]) > 0 and g.name != "apolar"]
                values = apolar_values + charged_values

            else:
                values = [g.sample(active_atoms_dic[g.name], trans) for g in self.grids if
                          len(active_atoms_dic[g.name]) > 0]
            scores = [item for sublist in values for item in sublist]
            return self.score(scores)

        def update_out_grids(self, score, active_coordinates_dic, trans):
            """
            For active atoms for a given grid, set closest grid point value to score, unless already set to a higher
            value

            :param score: float, score of a given probe
            :param active_coordinates_dic:
            :param trans: list of translations
            :return:
            """

            for pg in self.probe_grids:
                actives = active_coordinates_dic[pg.name]
                if len(actives) == 0:
                    continue
                for a in actives:
                    coords = pg.add_coordinates(a, trans)
                    i, j, k = pg.grid.point_to_indices(coords)
                    orig_value = pg.grid.value(i, j, k)
                    #
                    if self.settings.sphere_maps:
                        # pg.grid.set_sphere(coords, 2, score, mode='max')
                        if score > orig_value:
                            pg.grid.set_sphere(coords, 1.5, score, mode='max', scaling='None')
                    else:
                        pg.grid.set_value(i, j, k, max(score, orig_value))

        def get_active_coordinates(self):
            """
            Returns a dictionary of {grid_name:[Coordinates]}

            :return:
            """
            active_coords_dic = {
                g.name: [a.coordinates for a in g._active_atoms]
                for g in self.grids
            }

            return active_coords_dic

        def sample(self, molecule, probe, update_grids = True):
            """
            Sample the grids according to the settings

            :param molecule:
            :param probe: str, interaction type, (donor, acceptor, negative, positive, apolar)
            :return:
            """
            priority_atom, priority_atom_type = self.get_priority_atom(molecule)
            translate_points = self.get_translation_points(priority_atom_type)
            molecule.remove_hydrogens()
            quaternions = self.generate_rand_quaternions()
            high_scoring_probes = {}
            print("\n    nRotations:", len(quaternions), "nTranslations:", len(translate_points), "probename:", probe)

            for g in self.grids:
                g.set_molecule(molecule, True)

            for g in self.probe_grids:
                g.set_molecule(molecule, self.settings.polar_contributions)

            for q in tqdm(quaternions):
                molecule.apply_quaternion(q)
                priority_atom_coordinates = priority_atom.coordinates
                active_coordinates_dic = self.get_active_coordinates()

                for priority_atom_point in translate_points:

                    translation = [priority_atom_point[i] - priority_atom_coordinates[i]
                                   for i in range(0, len(priority_atom_coordinates))]

                    score = self.sample_pose(translation, active_coordinates_dic, probe)

                    try:
                        if score < 1:
                            continue
                    except TypeError:
                        continue
                    if update_grids:
                        self.update_out_grids(score, active_coordinates_dic, translation)

                    if self.settings.return_probes is True:
                        if score > 10:
                            m = molecule.copy()
                            m.translate(translation)
                            m.identifier = "{}".format(score)

                            try:
                                high_scoring_probes[score].append(m)
                            except KeyError:
                                high_scoring_probes[score] = [m]

            if self.settings.return_probes is True:
                sampled_probes = []
                for key in sorted(high_scoring_probes.keys(), reverse=True)[:100]:
                    sampled_probes.extend(high_scoring_probes[key])
                print('Returned probes = ', len(sampled_probes))
                return sampled_probes

                # if len(sampled_probes) > 100:
                #     return sampled_probes[:100]
                # else:
                #     return sampled_probes

    def __init__(self, settings=None):
        self.out_grids = {}
        self.super_grids = {}
        self.buriedness = None
        self.sampled_probes = {}

        if settings is None:
            self.sampler_settings = self.Settings()
        else:
            self.sampler_settings = settings

    @property
    def protein(self):
        """
        protein submitted for calculation
        :return:
        """
        return self._protein

    @protein.setter
    def protein(self, prot):
        """
        protein submitted for calculation
        :param `ccdc.protein.Protein` prot: protein
        :return:
        """
        if isinstance(prot, Protein):
            self._protein = prot
        else:
            raise TypeError("`ccdc.protein.Protein` must be supplied. Hint: Use Protein.from_file()")

    @property
    def charged_probes(self):
        """
        optional settings, True if charged features are to be calculated
        :return:
        """
        return self._charged_probes

    @charged_probes.setter
    def charged_probes(self, option):
        """
        optional settings, True if charged features are to be calculated
        :param bool option: (default = True)
        :return:
        """
        if type(option) is bool:
            self._charged_probes = option
        else:
            raise TypeError("Expecting a bool, got {} instead".format(type(option)))

    @property
    def probe_size(self):
        """
        optional settings, change probe size. NB: feature has not been tested.
        :return:
        """
        return self._probe_size

    @probe_size.setter
    def probe_size(self, size):
        """
        optional settings, change probe size. NB: feature has not been tested.
        :param int size: number of heavy atoms contain within the molecular probes (integer between 3 and 8)
        :return:
        """
        if size in range(3, 8):
            self._probe_size = size
        else:
            raise ValueError("Probe size must be an integer between 3-7")

    @property
    def buriedness_method(self):
        """
        optional settings, pocket detection method. (default = "ghecom")
        :return:
        """
        return self._buriedness_method

    @buriedness_method.setter
    def buriedness_method(self, method):
        """
        optional settings, pocket detection method. (default = "ghecom")
        :param str method: either 'ghecom' or 'ligsite'
        :return:
        """
        method = method.lower()
        if method == 'ghecom':
            if sys.platform == 'linux' or sys.platform == 'linux2':
                if 'GHECOM_EXE' in environ:
                    self._buriedness_method = method
                else:
                    raise EnvironmentError("Must set Ghecom environment variable")
            else:
                raise OSError('Ghecom is only supported on linux')

        elif method == 'ghecom_internal':
            self._buriedness_method = method

        elif method == 'ligsite':
            if sys.platform == 'linux' or sys.platform == 'linux2':
                print("RECOMMENDATION: you have chosen LIGSITE as buriedness method, ghecom is recommended")
            self._buriedness_method = method

        else:
            raise ValueError("Buriedness method must be 'ghecom' (default) or 'ligsite")

    @property
    def cavities(self):
        """
        optional settings, cavities supplied to the calculation
        :return:
        """
        return self._cavities

    @cavities.setter
    def cavities(self, obj):
        """
        optional settings, cavities supplied to the calculation
        :param list or Coordinate or `ccdc.molecule.Molecule` or `ccdc.cavity.Cavity`: cavity information provided
        :return:
        """
        if obj is not None:
            if isinstance(obj, list) or isinstance(obj, tuple):
                if isinstance(obj, Coordinates):
                    try:
                        print(obj.x)
                        self._cavities = [obj]
                    except AttributeError:
                        self._cavities = obj

                    self._cavities = [obj]
                elif isinstance(obj[0], Molecule):
                    self._cavities = [m.centre_of_geometry() for m in obj]
                elif isinstance(obj[0], Cavity):
                    self._cavities = [Helper.cavity_centroid(c) for c in obj]
                else:
                    try:
                        self._cavities = [Coordinates(x=obj[0], y=obj[1], z=obj[2])]
                    except RuntimeError:
                        print("WARNING! Failed to detected cavity, Atomic Hotspot detection to run on whole protein")
                        self._cavities = None

            elif isinstance(obj, Molecule):
                self._cavities = [obj.centre_of_geometry()]
            elif isinstance(obj, Cavity):
                self._cavities = [Helper.cavity_centroid(obj)]

            else:
                print("WARNING! Failed to detected cavity, Atomic Hotspot detection to run on whole protein")
                self._cavities = None
        else:
            self._cavities = None

    @property
    def nprocesses(self):
        """
        number of processes used in parallelisation
        :return:
        """
        return self._nprocesses

    @nprocesses.setter
    def nprocesses(self, num):
        """
        number of processes used in parallelisation
        :param int num: number of processor
        :return:
        """
        num = int(num)
        if num in range(0, int(multiprocessing.cpu_count())):
            self._nprocesses = num
        else:
            raise OSError("CPU count = {}".format(multiprocessing.cpu_count()))

    @property
    def sampler_settings(self):
        """
        sampler settings
        :return:
        """
        return self._sampler_settings

    @sampler_settings.setter
    def sampler_settings(self, settings):
        """
        sampler settings
        :param `hotspots.calculation.Runner.Settings` settings: adjust run settings
        :return:
        """
        if isinstance(settings, self.Settings):
            self._sampler_settings = settings
        else:
            self._sampler_settings = None

    def _get_weighted_maps(self):
        """
        private method

        weight superstar output by burriedness
        :return: a list of :class:`WeightedResult` instances
        """
        results = []
        for s in self.superstar_grids:
            g, b = Grid.common_grid([s.grid, self.buriedness], padding=1)
            weighted_grid = g * b
            results.append(_WeightedResult(s.identifier, weighted_grid))

        return results

    def _dock_probe(self,mol, grid_dict):

        donor_grid = _SampleGrid('donor', grid_dict['donor'], _SampleGrid.is_donor)
        acceptor_grid = _SampleGrid('acceptor', grid_dict['acceptor'], _SampleGrid.is_acceptor)
        apolar_grid = _SampleGrid('apolar', grid_dict['apolar'], _SampleGrid.is_apolar)

        kw = {'settings': self.sampler_settings}
        self.sampler = self._Sampler(apolar_grid, donor_grid, acceptor_grid, **kw)

        probes = self.sampler.sample(mol, probe='probe', update_grids=False)
        return probes

    def _get_out_maps(self, probe, grid_dict, return_probes=False):
        """
        private method

        organises the sampling of weighted superstar maps by molecular probes
        :param str probe: probe identifier set in the Atomic Hotspot calculation
        :param dict grid_dict: dictionary with key = probe identifier and value = `hotspots.grid_extension.Grid`
        :param bool return_probes: optional, bool indicating if probe molecules should be returned
        :return:
        """
        donor_grid = _SampleGrid('donor', grid_dict['donor'], _SampleGrid.is_donor)
        acceptor_grid = _SampleGrid('acceptor', grid_dict['acceptor'], _SampleGrid.is_acceptor)
        apolar_grid = _SampleGrid('apolar', grid_dict['apolar'], _SampleGrid.is_apolar)

        if self.charged_probes:
            negative_grid = _SampleGrid('negative', grid_dict['negative'], _SampleGrid.is_negative)
            positive_grid = _SampleGrid('positive', grid_dict['positive'], _SampleGrid.is_positive)

        kw = {'settings': self.sampler_settings}
        if self.charged_probes:
            self.sampler = self._Sampler(apolar_grid, donor_grid, acceptor_grid, negative_grid, positive_grid, **kw)
        else:
            self.sampler = self._Sampler(apolar_grid, donor_grid, acceptor_grid, **kw)

        probe_path = pkg_resources.resource_filename('hotspots', 'probes/')

        if self.charged_probes:
            if probe == "negative" or probe == "positive":
                mol = MoleculeReader(join(probe_path, "rotate-{}_{}_flat.mol2".format(probe, "test")))[0]
            else:
                mol = MoleculeReader(join(probe_path, "rotate-{}_{}_flat.mol2".format(probe, self.probe_size)))[0]
        else:
            mol = MoleculeReader(join(probe_path, "rotate-{}_{}_flat.mol2".format(probe, self.probe_size)))[0]

        probes = self.sampler.sample(mol, probe=probe)

        for pg in self.sampler.probe_grids:
            if pg.name.lower() == probe:
                try:
                    self.out_grids[pg.name].append(pg.grid)
                except KeyError:
                    self.out_grids[pg.name] = [pg.grid]

        if return_probes is True:
            return probes

    def _calc_hotspots(self, return_probes=False):
        """
        handles the organisation of the hotspot calculation
        :param return_probes: optional, bool indicating if probe molecules should be returned
        :return:
        """
        print("Start atomic hotspot detection\n        Processors: {}".format(self.nprocesses))
        a = _AtomicHotspot()
        a.settings.atomic_probes = {"apolar": "AROMATIC CH CARBON",
                                    "donor": "UNCHARGED NH NITROGEN",
                                    "acceptor": "CARBONYL OXYGEN"}
        if self.charged_probes:
            a.settings.atomic_probes = {"negative": "CARBOXYLATE OXYGEN", "positive": "CHARGED NH NITROGEN"}

        probe_types = a.settings.atomic_probes.keys()
        self.superstar_grids = a.calculate(protein=self.protein,
                                           nthreads=self.nprocesses,
                                           cavity_origins=self.cavities)

        if self.clear_tmp == True:
            shutil.rmtree(a.settings.temp_dir)

        print("Atomic hotspot detection complete\n")

        print("Start buriedness calculation")
        if self.buriedness_method.lower() == 'ghecom' and self.buriedness is None:
            print("    method: Ghecom")
            out_grid = self.superstar_grids[0].buriedness.copy_and_clear()

            ghecom = Ghecom()
            path = ghecom.run(self.protein)
            self.buriedness = ghecom.pdb_to_grid(path, out_grid)

            shutil.rmtree(ghecom.temp) # probs not needed

        elif self.buriedness_method.lower() == 'ghecom_internal' and self.buriedness is None:
            print("    method: Internal version Ghecom")
            out_grid = self.superstar_grids[0].buriedness.copy_and_clear()
            b = ExpBuriedness(prot=self.protein, out_grid=None)
            self.buriedness = b.buriedness_grid()

        elif self.buriedness_method.lower() == 'ligsite' and self.buriedness is None:
            print("    method: LIGSITE")
            self.buriedness = Grid.get_single_grid(grd_dict={s.identifier: s.buriedness for s in self.superstar_grids},
                                                   mask=False)

        self.weighted_grids = self._get_weighted_maps()

        print("Buriedness calcualtion complete\n")

        print("Start sampling")
        grid_dict = {w.identifier: w.grid for w in self.weighted_grids}

        for probe in probe_types:
            if return_probes is True:
                ps = self._get_out_maps(probe, grid_dict, return_probes=True)
                print(len(ps))
                self.sampled_probes.update({probe: ps})

            else:
                self._get_out_maps(probe, grid_dict)

        print("Sampling complete\n")

    def _prepare_protein(self, protoss=False):
        """
        default protein preparation settings on the protein
        :return:
        """
        self.protein.remove_all_waters()
        for lig in self.protein.ligands:
            self.protein.remove_ligand(lig.identifier)
        self.protein.remove_all_metals()

        if protoss is False:
            self.protein.add_hydrogens()

    def from_superstar(self, protein, superstar_grids, buriedness, charged_probes=False, probe_size=7,
                       settings=None, clear_tmp=False):
        """
        calculate hotspot maps from precalculated superstar maps. This enables more effective parallelisation and reuse
        of object such as the Buriedness grids

        :param protein: a :class:`ccdc.protein.Protein` instance
        :param superstar_grids: a :class:`hotspots.atomic_hotspot_calculation._AtomicHotspotResult` instance
        :param buriedness: a :class:`hotspots.grid_extension.Grid` instance
        :param bool charged_probes: If True, include positive and negative probes
        :param int probe_size: Size of probe in number of heavy atoms (3-8 atoms)
        :param settings: `hotspots.calculation.Runner.Settings` settings: holds the sampler settings
        :param bool clear_tmp: If True, clear the temporary directory
        :return:
        """
        start = time.time()
        self.super_grids = {}
        self.superstar_grids = superstar_grids
        self.probe_types = [p.identifier for p in self.superstar_grids]
        self.buriedness = buriedness

        self.protein = protein
        self.charged_probes = charged_probes
        self.probe_size = probe_size
        self.clear_tmp = clear_tmp

        if settings is None:
            self.sampler_settings = self.Settings()
        else:
            self.sampler_settings = settings

        self.weighted_grids = self._get_weighted_maps()

        print("Start sampling")
        grid_dict = {w.identifier: w.grid for w in self.weighted_grids}

        for probe in self.probe_types:
            self._get_out_maps(probe, grid_dict)

        self.super_grids = {p: g[0] for p, g in self.out_grids.items()}

        print("Sampling complete\n")
        print("Runtime = {}seconds".format(time.time() - start))

        return Results(super_grids=self.super_grids,
                       protein=self.protein,
                       buriedness=self.buriedness)

    def from_protein(self, protein, charged_probes=False, probe_size=7, buriedness_method='ghecom',
                     cavities=None, nprocesses=1, settings=None, buriedness_grid=None, clear_tmp=False):
        """
        generates a result from a protein

        :param protein: a :class:`ccdc.protein.Protein` instance
        :param bool charged_probes: If True include positive and negative probes
        :param int probe_size: Size of probe in number of heavy atoms (3-8 atoms)
        :param str buriedness_method: Either 'ghecom' or 'ligsite'
        :param cavities: Coordinate or `ccdc.cavity.Cavity` or `ccdc.molecule.Molecule` or list specifying the cavity or cavities on which the calculation should be run
        :param int nprocesses: number of CPU's used
        :param `hotspots.calculation.Runner.Settings` settings: holds the sampler settings
        :param `ccdc.utilities.Grid` buriedness_grid: pre-calculated buriedness grid
        :return: a :class:`hotspots.result.Results` instance

        >>> from ccdc.protein import Protein
        >>> from hotspots.calculation import Runner

        >>> protein = Protein.from_file(<path_to_protein>)

        >>> runner = Runner()
        >>> settings = Runner.Settings()
        >>> settings.nrotations = 1000  # fewer rotations increase speed at the expense of accuracy
        >>> runner.from_protein(protein, nprocesses=3, settings=settings)
        Result()

        """
        start = time.time()
        self.super_grids = {}
        self.buriedness = buriedness_grid
        self.protein = protein
        self.charged_probes = charged_probes
        self.probe_size = probe_size
        self.buriedness_method = buriedness_method
        self.cavities = cavities
        self.clear_tmp = clear_tmp

        print(self.cavities)
        self.nprocesses = nprocesses
        if settings is None:
            self.sampler_settings = self.Settings()
        else:
            self.sampler_settings = settings
        self._calc_hotspots()  # return probes = False by default
        self.super_grids = {p: g[0] for p, g in self.out_grids.items()}
        print("Runtime = {}seconds".format(time.time() - start))

        return Results(super_grids=self.super_grids,
                       protein=self.protein,
                       buriedness=self.buriedness,
                       superstar={x.identifier: x.grid for x in self.superstar_grids},
                       weighted_superstar={x.identifier: x.grid for x in self.weighted_grids})

    def from_pdb(self, pdb_code, charged_probes=False, probe_size=7, buriedness_method='ghecom', nprocesses=3,
                 cavities=False, settings=None, clear_tmp=False):
        """
        generates a result from a pdb code

        :param str pdb_code: PDB code
        :param bool charged_probes: If True include positive and negative probes
        :param int probe_size: Size of probe in number of heavy atoms (3-8 atoms)
        :param str buriedness_method: Either 'ghecom' or 'ligsite'
        :param int nprocesses: number of CPU's used
        :param `hotspots.calculation.Runner.Settings` settings: holds the calculation settings
        :return: a :class:`hotspots.result.Result` instance

        >>> from hotspots.calculation import Runner

        >>> runner = Runner()
        >>> runner.from_pdb("1hcl")
        Result()

        """
        protoss = False

        tmp = tempfile.mkdtemp()
        # if  protoss is True:
        #     protoss = Protoss(out_dir=tmp)
        #     self.protein = protoss.add_hydrogens(pdb_code).protein
        #
        # else:
        PDBResult(identifier=pdb_code).download(out_dir=tmp)
        fname = join(tmp, "{}.pdb".format(pdb_code))
        self.protein = Protein.from_file(fname)
        self._prepare_protein(protoss)
        self.charged_probes = charged_probes
        self.probe_size = probe_size
        self.buriedness_method = buriedness_method
        self.clear_tmp = clear_tmp
        self.cavities = None
        if cavities is True:
            self.cavities = Cavity.from_pdb_file(fname)
        self.nprocesses = nprocesses

        if settings is None:
            self.sampler_settings = self.Settings()
        else:
            self.sampler_settings = settings

        if self.sampler_settings.return_probes is True:
            print('here')
            self._calc_hotspots(return_probes=True)

        else:
            self._calc_hotspots()

        self.super_grids = {p: g[0] for p, g in self.out_grids.items()}

        if clear_tmp == True:
            shutil.rmtree(tmp)

        return Results(super_grids=self.super_grids,
                       protein=self.protein,
                       buriedness=self.buriedness,
                       superstar={x.identifier: x.grid for x in self.superstar_grids},
                       weighted_superstar={x.identifier: x.grid for x in self.weighted_grids})

    def global_dock(self, hr:Results, mol, settings = None):
        self.super_grids = hr.super_grids
        self.buriedness = hr.buriedness
        self.protein = hr.protein
        self.charged_probes = False
        self.probe_size = 7
        self.buriedness_method = 'ligsite'
        self.cavities = None
        self.clear_tmp = True

        if settings is None:
            self.sampler_settings = self.Settings()

        self.sampler_settings.return_probes = True
        self.sampler_settings.nrotations = 2000
        self.apolar_translation_threshold = 17
        self.polar_translation_threshold = 17

        print(hr.weighted_superstar)

        grid_dict = hr.weighted_superstar

        ps = self._dock_probe(mol, grid_dict)
        return ps


# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from scipy.io.wavfile import write
import numpy as np
from tqdm import tqdm
from matplotlib import cm
from tensorboardX import SummaryWriter
from ruamel import yaml
from pathlib import Path
import argparse
from pprint import pprint
import paddle.fluid as fluid
import paddle.fluid.dygraph as dg
from parakeet.g2p.en import text_to_sequence
from parakeet.models.transformer_tts.utils import *
from parakeet import audio
from parakeet.models.transformer_tts import Vocoder
from parakeet.models.transformer_tts import TransformerTTS
from parakeet.utils import io

def add_config_options_to_parser(parser):
    parser.add_argument("--config", type=str, help="path of the config file")
    parser.add_argument("--use_gpu", type=int, default=0, help="device to use")
    parser.add_argument(
        "--max_len",
        type=int,
        default=200,
        help="The max length of audio when synthsis.")

    parser.add_argument(
        "--checkpoint_transformer",
        type=str,
        help="transformer_tts checkpoint to synthesis")
    parser.add_argument(
        "--checkpoint_vocoder",
        type=str,
        help="vocoder checkpoint to synthesis")

    parser.add_argument(
        "--output",
        type=str,
        default="synthesis",
        help="path to save experiment results")

def synthesis(text_input, args):
    local_rank = dg.parallel.Env().local_rank
    place = (fluid.CUDAPlace(local_rank) if args.use_gpu else fluid.CPUPlace())

    with open(args.config) as f:
        cfg = yaml.load(f, Loader=yaml.Loader)

    # tensorboard
    if not os.path.exists(args.output):
        os.mkdir(args.output)

    writer = SummaryWriter(os.path.join(args.output, 'log'))

    fluid.enable_dygraph(place)
    with fluid.unique_name.guard():
        network_cfg = cfg['network']
        model = TransformerTTS(
            network_cfg['embedding_size'], network_cfg['hidden_size'],
            network_cfg['encoder_num_head'], network_cfg['encoder_n_layers'],
            cfg['audio']['num_mels'], network_cfg['outputs_per_step'],
            network_cfg['decoder_num_head'], network_cfg['decoder_n_layers'])
        # Load parameters.
        global_step = io.load_parameters(
            model=model, checkpoint_path=args.checkpoint_transformer)
        model.eval()

    with fluid.unique_name.guard():
        model_vocoder = Vocoder(
            cfg['train']['batch_size'], cfg['vocoder']['hidden_size'],
            cfg['audio']['num_mels'], cfg['audio']['n_fft'])
        # Load parameters.
        global_step = io.load_parameters(
            model=model_vocoder, checkpoint_path=args.checkpoint_vocoder)
        model_vocoder.eval()
    # init input
    text = np.asarray(text_to_sequence(text_input))
    text = fluid.layers.unsqueeze(dg.to_variable(text), [0])
    mel_input = dg.to_variable(np.zeros([1, 1, 80])).astype(np.float32)
    pos_text = np.arange(1, text.shape[1] + 1)
    pos_text = fluid.layers.unsqueeze(dg.to_variable(pos_text), [0])

    pbar = tqdm(range(args.max_len))
    for i in pbar:
        pos_mel = np.arange(1, mel_input.shape[1] + 1)
        pos_mel = fluid.layers.unsqueeze(dg.to_variable(pos_mel), [0])
        mel_pred, postnet_pred, attn_probs, stop_preds, attn_enc, attn_dec = model(
            text, mel_input, pos_text, pos_mel)
        mel_input = fluid.layers.concat(
            [mel_input, postnet_pred[:, -1:, :]], axis=1)

    mag_pred = model_vocoder(postnet_pred)

    _ljspeech_processor = audio.AudioProcessor(
        sample_rate=cfg['audio']['sr'],
        num_mels=cfg['audio']['num_mels'],
        min_level_db=cfg['audio']['min_level_db'],
        ref_level_db=cfg['audio']['ref_level_db'],
        n_fft=cfg['audio']['n_fft'],
        win_length=cfg['audio']['win_length'],
        hop_length=cfg['audio']['hop_length'],
        power=cfg['audio']['power'],
        preemphasis=cfg['audio']['preemphasis'],
        signal_norm=True,
        symmetric_norm=False,
        max_norm=1.,
        mel_fmin=0,
        mel_fmax=None,
        clip_norm=True,
        griffin_lim_iters=60,
        do_trim_silence=False,
        sound_norm=False)

    # synthesis with cbhg
    wav = _ljspeech_processor.inv_spectrogram(
        fluid.layers.transpose(fluid.layers.squeeze(mag_pred, [0]), [1, 0])
        .numpy())
    global_step = 0
    for i, prob in enumerate(attn_probs):
        for j in range(4):
            x = np.uint8(cm.viridis(prob.numpy()[j]) * 255)
            writer.add_image(
                'Attention_%d_0' % global_step,
                x,
                i * 4 + j,
                dataformats="HWC")

    writer.add_audio(text_input + '(cbhg)', wav, 0, cfg['audio']['sr'])

    if not os.path.exists(os.path.join(args.output, 'samples')):
        os.mkdir(os.path.join(args.output, 'samples'))
    write(
        os.path.join(os.path.join(args.output, 'samples'), 'cbhg.wav'),
        cfg['audio']['sr'], wav)

    # synthesis with griffin-lim
    wav = _ljspeech_processor.inv_melspectrogram(
        fluid.layers.transpose(
            fluid.layers.squeeze(postnet_pred, [0]), [1, 0]).numpy())
    writer.add_audio(text_input + '(griffin)', wav, 0, cfg['audio']['sr'])

    write(
        os.path.join(os.path.join(args.output, 'samples'), 'griffin.wav'),
        cfg['audio']['sr'], wav)
    print("Synthesis completed !!!")
    writer.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Synthesis model")
    add_config_options_to_parser(parser)
    args = parser.parse_args()
    # Print the whole config setting.
    pprint(vars(args))
    synthesis("Parakeet stands for Paddle PARAllel text-to-speech toolkit.",
              args)


# coding=utf-8

import argparse
import json
import os
import numpy as np
import tensorflow.compat.v1 as tf
import time
import tqdm
from copy import copy
from encode_bpe import BPEEncoder_ja
import model

if int(tf.__version__[0]) > 1:
    from model import HParams as HParams
else:
    from tensorflow.contrib.training import HParams

CHECKPOINT_DIR = 'checkpoint'
SAMPLE_DIR = 'samples'

parser = argparse.ArgumentParser(
    description='Pretraining GPT2-JA on your custom dataset.',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument('--dataset', metavar='PATH', type=str, required=True, help='Input npz file')
parser.add_argument('--base_model', type=str, default='gpt2ja-small', help='a path to a model file')

parser.add_argument('--batch_size', metavar='SIZE', type=int, default=1, help='Batch size')
parser.add_argument('--optim', type=str, default='adam', help='"adam", "adagrad", or "sgd" to use optimizer')
parser.add_argument('--learning_rate', metavar='LR', type=float, default=5e-5, help='Learning rate for optimizer')
parser.add_argument('--warmup_steps', metavar='WR', type=int, default=0, help='Learning rate warming up steps')

parser.add_argument('--run_name', type=str, default='gpt2ja_finetune', help='Run id. Name of subdirectory in checkpoint/')
parser.add_argument('--save_every', metavar='N', type=int, default=1000, help='Write a checkpoint every N steps')

parser.add_argument('--gpu', default='0', help='visible gpu number.')

def maketree(path):
    try:
        os.makedirs(path)
    except:
        pass

with open('ja-bpe.txt', encoding='utf-8') as f:
    bpe = f.read().split('\n')

with open('emoji.json', encoding='utf-8') as f:
    emoji = json.loads(f.read())

enc = BPEEncoder_ja(bpe, emoji)
n_vocab = len(enc)

def main():
    args = parser.parse_args()

    if os.path.isfile(args.base_model+'/hparams.json'):
        with open(args.base_model+'/hparams.json', encoding='utf-8') as f:
            params = json.loads(f.read())
            hparams = HParams(**params)
    elif 'small' in args.base_model:
        hparams = HParams(**{
          "n_vocab": n_vocab,
          "n_ctx": 1024,
          "n_embd": 768,
          "n_head": 12,
          "n_layer": 12
        })
    elif 'medium' in args.base_model:
        hparams = HParams(**{
          "n_vocab": n_vocab,
          "n_ctx": 1024,
          "n_embd": 1024,
          "n_head": 16,
          "n_layer": 24
        })
    elif 'large' in args.base_model:
        hparams = HParams(**{
          "n_vocab": n_vocab,
          "n_ctx": 1024,
          "n_embd": 1280,
          "n_head": 20,
          "n_layer": 36
        })
    else:
        raise ValueError('invalid model name.')

    config = tf.ConfigProto()
    if int(args.gpu) >= 0:
        config.gpu_options.allow_growth = True
        config.gpu_options.visible_device_list = args.gpu
    with tf.Session(config=config,graph=tf.Graph()) as sess:
        context = tf.placeholder(tf.int32, [None, None])
        output = model.model(hparams=hparams, X=context, past=None, reuse=tf.AUTO_REUSE)
        loss = tf.reduce_mean(
            tf.nn.sparse_softmax_cross_entropy_with_logits(
                labels=context[:, 1:], logits=output['logits'][:, :-1]))

        saver = tf.train.Saver()
        ckpt = tf.train.latest_checkpoint(args.base_model)
        saver.restore(sess, ckpt)

        train_vars = tf.trainable_variables()

        global_step = tf.Variable(0, trainable=False)
        if args.warmup_steps > 0:
            learning_rate = tf.compat.v1.train.polynomial_decay(
                    learning_rate=1e-10,
                    end_learning_rate=args.learning_rate,
                    global_step=global_step,
                    decay_steps=args.warmup_steps
                )
        else:
            learning_rate = args.learning_rate

        if args.optim=='adam':
            opt = tf.train.AdamOptimizer(learning_rate=learning_rate,
                                           beta1=0.9,
                                           beta2=0.98,
                                           epsilon=1e-7)
        elif args.optim=='adagrad':
            opt = tf.train.AdagradOptimizer(learning_rate=learning_rate)
        elif args.optim=='sgd':
            opt = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
        else:
            raise ValueError('invalid optimizer name.')

        train_vars = tf.trainable_variables()
        opt_grads = tf.gradients(loss, train_vars)
        opt_grads = list(zip(opt_grads, train_vars))
        opt_apply = opt.apply_gradients(opt_grads)

        summaries = tf.summary.scalar('loss', loss)
        summary_log = tf.summary.FileWriter(
            os.path.join(CHECKPOINT_DIR, args.run_name))

        saver = tf.train.Saver(
            var_list=train_vars,
            max_to_keep=5,
            keep_checkpoint_every_n_hours=2)
        sess.run(tf.global_variables_initializer())

        ckpt = tf.train.latest_checkpoint(args.base_model)
        saver.restore(sess, ckpt)
        print('Loading checkpoint', ckpt)

        print('Loading dataset...')
        global_chunks = []
        with np.load(args.dataset) as npz:
            current_token = []
            for inditem, item in enumerate(npz.files):
                token_chunk = npz[item]
                for ind in range(0,len(token_chunk)):
                    current_token.append(np.uint16(token_chunk[ind]))
                    if len(current_token) == hparams.n_ctx:
                        global_chunks.append(current_token)
                        current_token = []
            if len(current_token) > 1:
                global_chunks.append(current_token + [n_vocab-1]*(hparams.n_ctx-len(current_token)))
                current_token = []
        global_chunk_index = np.random.permutation(len(global_chunks))
        global_chunk_step = 0
        print('Training...')

        def sample_feature():
            nonlocal global_chunks,global_chunk_index,global_chunk_step
            p_input_ids = []

            for b in range(args.batch_size): # FULL-SENTENCES
                idx = global_chunk_index[global_chunk_step]
                global_chunk_step += 1
                if global_chunk_step >= len(global_chunk_index):
                    global_chunk_step = 0
                    global_chunk_index = np.random.permutation(len(global_chunks))
                sampled_token = global_chunks[idx]
                # Make Sequence
                ids = copy(global_chunks[idx])
                p_input_ids.append(ids)

            return {context:p_input_ids}

        counter = 1
        counter_path = os.path.join(CHECKPOINT_DIR, args.run_name, 'counter')
        hparams_path = os.path.join(CHECKPOINT_DIR, args.run_name, 'hparams.json')
        if os.path.exists(counter_path):
            # Load the step number if we're resuming a run
            # Add 1 so we don't immediately try to save again
            with open(counter_path, 'r', encoding='utf-8') as fp:
                counter = int(fp.read()) + 1

        maketree(os.path.join(CHECKPOINT_DIR, args.run_name))

        def save():
            maketree(os.path.join(CHECKPOINT_DIR, args.run_name))
            print(
                'Saving',
                os.path.join(CHECKPOINT_DIR, args.run_name,
                             'model-{}').format(counter))
            saver.save(
                sess,
                os.path.join(CHECKPOINT_DIR, args.run_name, 'model'),
                global_step=counter)
            with open(counter_path, 'w', encoding='utf-8') as fp:
                fp.write(str(counter) + '\n')
            with open(hparams_path, 'w', encoding='utf-8') as fp:
                fp.write(json.dumps({
                      "n_vocab": int(hparams.n_vocab),
                      "n_ctx": int(hparams.n_ctx),
                      "n_embd": int(hparams.n_embd),
                      "n_head": int(hparams.n_head),
                      "n_layer": int(hparams.n_layer),
                }))

        avg_loss = (0.0, 0.0)
        start_time = time.time()

        try:
            for i in range(500):
                if counter % args.save_every == 0:
                    save()

                (_, v_loss, v_summary) = sess.run(
                    (opt_apply, loss, summaries),
                    feed_dict=sample_feature())

                summary_log.add_summary(v_summary, counter)

                avg_loss = (avg_loss[0] * 0.99 + v_loss,
                            avg_loss[1] * 0.99 + 1.0)

                print(
                    '[{counter} | {time:2.2f}] loss={loss:2.2f} avg={avg:2.2f}'
                    .format(
                        counter=counter,
                        time=time.time() - start_time,
                        loss=v_loss,
                        avg=avg_loss[0] / avg_loss[1]))

                counter = counter+1
                if args.warmup_steps > 0:
                    global_step = global_step+1
            save()
        except KeyboardInterrupt:
            print('interrupted')
            save()

if __name__ == '__main__':
    main()