import mdtraj as md
import numpy as np
from openmmtools import integrators
from simtk import unit
from simtk.openmm import app
from tqdm import tqdm

import rlmm.environment.openmmWrappers.utils as mmWrapperUtils
from rlmm.utils.config import Config
from rlmm.utils.loggers import DCDReporter
from rlmm.utils.loggers import make_message_writer

class OpenMMSimulationWrapper:
    class Config(Config):
        def __init__(self, args):
            self.verbose = None
            self.parameters = mmWrapperUtils.SystemParams(args['params'])
            self.systemloader = None
            if args is not None:
                self.__dict__.update(args)

        def get_obj(self, system_loader, *args, **kwargs):
            self.systemloader = system_loader
            return OpenMMSimulationWrapper(self, *args, **kwargs)

    def __init__(self, config_: Config, old_sampler_state=None):
        """

        :param systemLoader:
        :param config:
        """
        self._times = None
        self.config = config_
        self.logger = make_message_writer(self.config.verbose, self.__class__.__name__)
        if self.config.systemloader.system is None:
            self.system = self.config.systemloader.get_system(self.config.parameters.createSystem)
            self.topology = self.config.systemloader.topology
            positions, velocities = self.config.systemloader.get_positions(), None

        else:
            self.system = self.config.systemloader.system
            self.topology = self.config.systemloader.topology
            positions, velocities = self.config.systemloader.get_positions(), None

        self.topology = self.config.systemloader.topology
        self.explicit = self.config.systemloader.explicit

        integrator = integrators.GeodesicBAOABIntegrator(temperature=self.config.parameters.integrator_params['temperature'],
                                                    timestep=self.config.parameters.integrator_params['timestep'],
                                                    collision_rate=self.config.parameters.integrator_params[
                                                        'collision_rate'],
                                                    constraint_tolerance=self.config.parameters.integrator_setConstraintTolerance)
        self.simulation = app.Simulation(self.topology, self.system, integrator,
                                         self.config.parameters.platform, self.config.parameters.platform_config)
        self.simulation.context.setPositions(self.config.systemloader.positions)
        self.simulation.minimizeEnergy(self.config.parameters.minMaxIters)
        self.simulation.context.setVelocitiesToTemperature(self.config.parameters.integrator_params['temperature'])

    def get_sim_time(self):
        return self.config.parameters.integrator_params['timestep'] *self.config.n_steps

    def get_coordinates(self):
        return mmWrapperUtils.get_coordinates(self.topology,
                                              self.simulation.context.getState(getPositions=True).getPositions(),
                                              self.simulation.system.getDefaultPeriodicBoxVectors(),
                                              self.simulation.system.getNumParticles(),
                                              self.explicit)

    def get_velocities(self):
        return self.simulation.context.getState(getVelocities=True).getVelocities(asNumpy=True)

    def run_amber_mmgbsa(self, run_decomp=False):
        return mmWrapperUtils.run_amber_mmgbsa(self.logger, self.explicit, self.config.tempdir(), run_decomp=run_decomp)

    def get_pdb(self, file_name=None):
        return mmWrapperUtils.get_pdb(self.simulation.topology,
                                      self.simulation.context.getState(getPositions=True).getPositions(),
                                      file_name=file_name)

    def run(self, iters, steps_per_iter):
        """

        :param steps:
        """
        with self.logger("run") as logger:

            if 'cur_sim_steps' not in self.__dict__:
                self.cur_sim_steps = 0.0 * unit.picosecond

            pbar = tqdm(range(iters), desc="running {} steps per sample".format(steps_per_iter))
            self._trajs = np.zeros((iters, self.system.getNumParticles(), 3))
            self._times = np.zeros((iters))
            dcdreporter = DCDReporter(f"{self.config.tempdir()}/traj.dcd", 1, append=False)
            for i in pbar:
                self.simulation.step(self.config.n_steps * steps_per_iter)
                self.cur_sim_steps += (steps_per_iter * self.get_sim_time())
                _state = self.simulation.context.getState(getPositions=True)
                dcdreporter.report(self.topology, _state, (i + 1), 0.5 * unit.femtosecond)

                # log trajectory
                self._trajs[i] = np.array(_state.getPositions(asNumpy=True).value_in_unit(unit.nanometer)).reshape(
                    (self.system.getNumParticles(), 3))
                self._times[i] = self.cur_sim_steps.value_in_unit(unit.picosecond)
            pbar.close()

    def writetraj(self):
        if self.explicit:
            lengths, angles = mmWrapperUtils.get_mdtraj_box(
                boxvec=self.simulation.system.getDefaultPeriodicBoxVectors(),
                iterset=self._trajs.shape[0])
            traj = md.Trajectory(self._trajs, md.Topology.from_openmm(self.topology),
                                 unitcell_lengths=lengths,
                                 unitcell_angles=angles, time=self._times)
            traj.image_molecules(inplace=True)
        else:
            traj = md.Trajectory(self._trajs, md.Topology.from_openmm(self.topology), time=self._times)

        # traj.save_pdb(f'{self.config.tempdir()}/mdtraj_traj.pdb')
        traj.save_hdf5(f'{self.config.tempdir()}/mdtraj_traj.h5')
        # traj.save_dcd(f'{self.config.tempdir()}/mdtraj_traj.dcd')
        return traj


"""Define some functions to sort substructures of parameter by importance.
"""
# 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 logging
import numpy as np
from ..core import GraphWrapper
from ..common import get_logger
from ..core import Registry

__all__ = ["IDX_SELECTOR"]

IDX_SELECTOR = Registry('idx_selector')

@IDX_SELECTOR.register
def default_idx_selector(group, scores, ratios):
    """Get the pruned indices by scores of master tensor.

    This function return a list of parameters' pruned indices on given axis.
    Each element of list is a tuple with format (name, axis, indices)
    in which 'name' is parameter's name and 'axis' is the axis pruning on and
    `indices` is indices to be pruned.

    Args:
       group(Group): A group of pruning operations.
       scores(dict): The key is name of tensor, the value is a dict with axis as key and scores as value.
       ratios(dict): The pruned ratio of each tensor. The key is name of tensor and the value is the pruned ratio.

    Returns:

       list: pruned indices with format (name, axis, pruned_indices).

    """
    # sort channels by the master convolution's score
    name = group.master["name"]
    axis = group.master["axis"]
    score = scores[name][axis]

    # get max convolution groups attribution
    max_groups = 1
    for pruning_details in group.all_pruning_details():
        groups = pruning_details.op.attr("groups")
        if groups is not None and groups > max_groups:
            max_groups = groups
    if max_groups > 1:
        score = score.reshape([max_groups, -1])
        group_size = score.shape[1]
        # get score for each group of channels
        score = np.mean(score, axis=1)
    sorted_idx = score.argsort()
    ratio = ratios[name]
    pruned_num = int(round(len(sorted_idx) * ratio))
    pruned_idx = sorted_idx[:pruned_num]
    # convert indices of channel groups to indices of channels.
    if max_groups > 1:
        correct_idx = []
        for idx in pruned_idx:
            for offset in range(group_size):
                correct_idx.append(idx * group_size + offset)
        pruned_idx = correct_idx[:]
    ret = []
    for _pruning_details in group.all_pruning_details():
        ret.append((_pruning_details.name, _pruning_details.axis, pruned_idx,
                    _pruning_details.transform))
    return ret

@IDX_SELECTOR.register
def optimal_threshold(group, scores, ratios):
    """Get the pruned indices by scores of master tensor.

    This function return a list of parameters' pruned indices on given axis.
    Each element of list is a tuple with format (name, axis, indices)
    in which 'name' is parameter's name and 'axis' is the axis pruning on and
    `indices` is indices to be pruned.

    Args:
       group(Group): A group of pruning operations.
       scores(dict): The key is name of tensor, the value is a dict with axis as key and scores as value.
       ratios(dict): The pruned ratio of each tensor. The key is name of tensor and the value is the pruned ratio.

    Returns:
       list: pruned indices with format (name, axis, pruned_indices).
    """
    # sort channels by the master tensor
    name = group.master["name"]
    axis = group.master["axis"]
    score = scores[name][axis]
    ratio = ratios[name]

    score[score < 1e-18] = 1e-18
    score_sorted = np.sort(score)
    score_square = score_sorted**2
    total_sum = score_square.sum()
    acc_sum = 0
    for i in range(score_square.size):
        acc_sum += score_square[i]
        if acc_sum / total_sum > ratio:
            break
    th = (score_sorted[i - 1] + score_sorted[i]) / 2 if i > 0 else 0

    pruned_idx = np.squeeze(np.argwhere(score < th))

    idxs = []
    for _pruning_details in group.all_pruning_details():
        idxs.append((_pruning_details.name, _pruning_details.axis, pruned_idx,
                     _pruning_details.transform))
    return idxs


import logging
from typing import Optional, Sequence, Union

import numpy as np
import pandas as pd
from anndata import AnnData

from scvi import _CONSTANTS
from scvi._compat import Literal
from scvi.core.data_loaders import AnnotationDataLoader
from scvi.core.models import ArchesMixin, BaseModelClass, RNASeqMixin, VAEMixin
from scvi.core.modules import SCANVAE, VAE
from scvi.core.trainers import SemiSupervisedTrainer, UnsupervisedTrainer

from .scvi import SCVI

logger = logging.getLogger(__name__)

class SCANVI(RNASeqMixin, VAEMixin, ArchesMixin, BaseModelClass):
    """
    Single-cell annotation using variational inference [Xu19]_.

    Inspired from M1 + M2 model, as described in (https://arxiv.org/pdf/1406.5298.pdf).

    Parameters
    ----------
    adata
        AnnData object that has been registered via :func:`~scvi.data.setup_anndata`.
    unlabeled_category
        Value used for unlabeled cells in `labels_key` used to setup AnnData with scvi.
    pretrained_model
        Instance of SCVI model that has already been trained.
    n_hidden
        Number of nodes per hidden layer.
    n_latent
        Dimensionality of the latent space.
    n_layers
        Number of hidden layers used for encoder and decoder NNs.
    dropout_rate
        Dropout rate for neural networks.
    dispersion
        One of the following:

        * ``'gene'`` - dispersion parameter of NB is constant per gene across cells
        * ``'gene-batch'`` - dispersion can differ between different batches
        * ``'gene-label'`` - dispersion can differ between different labels
        * ``'gene-cell'`` - dispersion can differ for every gene in every cell
    gene_likelihood
        One of:

        * ``'nb'`` - Negative binomial distribution
        * ``'zinb'`` - Zero-inflated negative binomial distribution
        * ``'poisson'`` - Poisson distribution
    use_cuda
        Use the GPU or not.
    **model_kwargs
        Keyword args for :class:`~scvi.core.modules.VAE` and :class:`~scvi.core.modules.SCANVAE`

    Examples
    --------
    >>> adata = anndata.read_h5ad(path_to_anndata)
    >>> scvi.data.setup_anndata(adata, batch_key="batch", labels_key="labels")
    >>> vae = scvi.model.SCANVI(adata, "Unknown")
    >>> vae.train()
    >>> adata.obsm["X_scVI"] = vae.get_latent_representation()
    >>> adata.obs["pred_label"] = vae.predict()
    """

    def __init__(
        self,
        adata: AnnData,
        unlabeled_category: Union[str, int, float],
        pretrained_model: Optional[SCVI] = None,
        n_hidden: int = 128,
        n_latent: int = 10,
        n_layers: int = 1,
        dropout_rate: float = 0.1,
        dispersion: Literal["gene", "gene-batch", "gene-label", "gene-cell"] = "gene",
        gene_likelihood: Literal["zinb", "nb", "poisson"] = "zinb",
        use_cuda: bool = True,
        **model_kwargs,
    ):
        super(SCANVI, self).__init__(adata, use_cuda=use_cuda)
        self.unlabeled_category_ = unlabeled_category

        if pretrained_model is not None:
            if pretrained_model.is_trained is False:
                raise ValueError("pretrained model has not been trained")
            self._base_model = pretrained_model.model
            self._is_trained_base = True
        else:
            self._base_model = VAE(
                n_input=self.summary_stats["n_vars"],
                n_batch=self.summary_stats["n_batch"],
                n_hidden=n_hidden,
                n_latent=n_latent,
                n_layers=n_layers,
                dropout_rate=dropout_rate,
                dispersion=dispersion,
                gene_likelihood=gene_likelihood,
                **model_kwargs,
            )
            self._is_trained_base = False
        self.model = SCANVAE(
            n_input=self.summary_stats["n_vars"],
            n_batch=self.summary_stats["n_batch"],
            n_labels=self.summary_stats["n_labels"],
            n_hidden=n_hidden,
            n_latent=n_latent,
            n_layers=n_layers,
            dropout_rate=dropout_rate,
            dispersion=dispersion,
            gene_likelihood=gene_likelihood,
            **model_kwargs,
        )

        # get indices for labeled and unlabeled cells
        key = self.scvi_setup_dict_["data_registry"][_CONSTANTS.LABELS_KEY]["attr_key"]
        self._label_mapping = self.scvi_setup_dict_["categorical_mappings"][key][
            "mapping"
        ]
        original_key = self.scvi_setup_dict_["categorical_mappings"][key][
            "original_key"
        ]
        labels = np.asarray(self.adata.obs[original_key]).ravel()
        self._code_to_label = {i: l for i, l in enumerate(self._label_mapping)}
        self._unlabeled_indices = np.argwhere(
            labels == self.unlabeled_category_
        ).ravel()
        self._labeled_indices = np.argwhere(labels != self.unlabeled_category_).ravel()
        self.unsupervised_history_ = None
        self.semisupervised_history_ = None

        self._model_summary_string = (
            "ScanVI Model with the following params: \nunlabeled_category: {}, n_hidden: {}, n_latent: {}"
            ", n_layers: {}, dropout_rate: {}, dispersion: {}, gene_likelihood: {}"
        ).format(
            unlabeled_category,
            n_hidden,
            n_latent,
            n_layers,
            dropout_rate,
            dispersion,
            gene_likelihood,
        )
        self.init_params_ = self._get_init_params(locals())

    @property
    def _trainer_class(self):
        return SemiSupervisedTrainer

    @property
    def _scvi_dl_class(self):
        return AnnotationDataLoader

    @property
    def history(self):
        """Returns computed metrics during training."""
        return {
            "unsupervised_trainer_history": self.unsupervised_history_,
            "semisupervised_trainer_history": self.semisupervised_history_,
        }

    def train(
        self,
        n_epochs_unsupervised: Optional[int] = None,
        n_epochs_semisupervised: Optional[int] = None,
        train_base_model: bool = True,
        train_size: float = 0.9,
        test_size: float = None,
        lr: float = 1e-3,
        n_epochs_kl_warmup: int = 400,
        n_iter_kl_warmup: Optional[int] = None,
        frequency: Optional[int] = None,
        unsupervised_trainer_kwargs: dict = {},
        semisupervised_trainer_kwargs: dict = {},
        unsupervised_train_kwargs: dict = {},
        semisupervised_train_kwargs: dict = {},
    ):
        """
        Train the model.

        Parameters
        ----------
        n_epochs_unsupervised
            Number of passes through the dataset for unsupervised pre-training.
        n_epochs_semisupervised
            Number of passes through the dataset for semisupervised training.
        train_base_model
            Pretrain an SCVI base model first before semisupervised training.
        train_size
            Size of training set in the range [0.0, 1.0].
        test_size
            Size of the test set. If `None`, defaults to 1 - `train_size`. If
            `train_size + test_size < 1`, the remaining cells belong to a validation set.
        lr
            Learning rate for optimization.
        n_epochs_kl_warmup
            Number of passes through dataset for scaling term on KL divergence to go from 0 to 1.
        n_iter_kl_warmup
            Number of minibatches for scaling term on KL divergence to go from 0 to 1.
            To use, set to not `None` and set `n_epochs_kl_warmup` to `None`.
        frequency
            Frequency with which metrics are computed on the data for train/test/val sets for both
            the unsupervised and semisupervised trainers. If you'd like a different frequency for
            the semisupervised trainer, set frequency in semisupervised_train_kwargs.
        unsupervised_trainer_kwargs
            Other keyword args for :class:`~scvi.core.trainers.UnsupervisedTrainer`.
        semisupervised_trainer_kwargs
            Other keyword args for :class:`~scvi.core.trainers.SemiSupervisedTrainer`.
        semisupervised_train_kwargs
            Keyword args for the train method of :class:`~scvi.core.trainers.SemiSupervisedTrainer`.
        """
        unsupervised_trainer_kwargs = dict(unsupervised_trainer_kwargs)
        semisupervised_trainer_kwargs = dict(semisupervised_trainer_kwargs)
        unsupervised_train_kwargs = dict(unsupervised_train_kwargs)
        semisupervised_train_kwargs = dict(semisupervised_train_kwargs)

        if n_epochs_unsupervised is None:
            n_epochs_unsupervised = np.min(
                [round((20000 / self.adata.shape[0]) * 400), 400]
            )
        if n_epochs_semisupervised is None:
            n_epochs_semisupervised = int(
                np.min([10, np.max([2, round(n_epochs_unsupervised / 3.0)])])
            )
        logger.info(
            "Training Unsupervised Trainer for {} epochs.".format(n_epochs_unsupervised)
        )
        logger.info(
            "Training SemiSupervised Trainer for {} epochs.".format(
                n_epochs_semisupervised
            )
        )

        if self._is_trained_base is not True and train_base_model:
            self._unsupervised_trainer = UnsupervisedTrainer(
                self._base_model,
                self.adata,
                train_size=train_size,
                test_size=test_size,
                n_iter_kl_warmup=n_iter_kl_warmup,
                n_epochs_kl_warmup=n_epochs_kl_warmup,
                frequency=frequency,
                use_cuda=self.use_cuda,
                **unsupervised_trainer_kwargs,
            )
            self._unsupervised_trainer.train(
                n_epochs=n_epochs_unsupervised, lr=lr, **unsupervised_train_kwargs
            )
            self.unsupervised_history_ = self._unsupervised_trainer.history
            self._is_trained_base = True

            self.model.load_state_dict(self._base_model.state_dict(), strict=False)

        if "frequency" not in semisupervised_trainer_kwargs and frequency is not None:
            semisupervised_trainer_kwargs["frequency"] = frequency

        self.trainer = SemiSupervisedTrainer(
            self.model,
            self.adata,
            use_cuda=self.use_cuda,
            indices_labelled=self._labeled_indices,
            indices_unlabelled=self._unlabeled_indices,
            **semisupervised_trainer_kwargs,
        )
        self.semisupervised_history_ = self.trainer.history
        self.trainer.train(
            n_epochs=n_epochs_semisupervised,
            **semisupervised_train_kwargs,
        )

        self.is_trained_ = True

    def predict(
        self,
        adata: Optional[AnnData] = None,
        indices: Optional[Sequence[int]] = None,
        soft: bool = False,
        batch_size: int = 128,
    ) -> Union[np.ndarray, pd.DataFrame]:
        """
        Return cell label predictions.

        Parameters
        ----------
        adata
            AnnData object that has been registered via :func:`~scvi.data.setup_anndata`.
        indices
            Indices of cells in adata to use. If `None`, all cells are used.
        soft
            Return probabilities for each class label.
        batch_size
            Minibatch size to use.
        """
        adata = self._validate_anndata(adata)

        if indices is None:
            indices = np.arange(adata.n_obs)

        scdl = self._make_scvi_dl(adata=adata, indices=indices, batch_size=batch_size)

        _, pred = scdl.sequential().compute_predictions(soft=soft)

        if not soft:
            predictions = []
            for p in pred:
                predictions.append(self._code_to_label[p])

            return np.array(predictions)
        else:
            pred = pd.DataFrame(
                pred,
                columns=self._label_mapping,
                index=adata.obs_names[indices],
            )
            return pred


import multiprocessing

multiprocessing.set_start_method("spawn", True)
import os
import sys
import time
from collections import deque

import numpy as np
import torch
import torch.nn as nn
import wandb

from a2c_ppo_acktr import algo, utils
from a2c_ppo_acktr.model import NNBase, Policy
from a2c_ppo_acktr.storage import RolloutStorage
from aari.envs import make_vec_envs
from src.encoders import ImpalaCNN, NatureCNN
from src.utils import get_argparser

def get_envs(
    env_name, seed=42, num_processes=1, num_frame_stack=1, downsample=False, color=False
):
    return make_vec_envs(
        env_name, seed, num_processes, num_frame_stack, downsample, color
    )

def get_encoder(args, observation_shape, device):
    if args.encoder_type == "Nature":
        encoder = NatureCNN(observation_shape[0], args)
    elif args.encoder_type == "Impala":
        encoder = ImpalaCNN(observation_shape[0], args)

    if args.weights_path == "None":
        sys.stderr.write(
            "Training without loading in encoder weights! Are sure you want to do that??"
        )
    else:
        print(
            "Print loading in encoder weights from probe of type {} from the following path: {}".format(
                args.method, args.weights_path
            )
        )
        encoder.load_state_dict(torch.load(args.weights_path, map_location=device))
        encoder.eval()
    return encoder

class SimpleBase(NNBase):
    def __init__(self, num_inputs, recurrent=False, hidden_size=256, encoder=None):
        super().__init__(recurrent, num_inputs, hidden_size)
        init_ = lambda m: utils.init(
            m, nn.init.orthogonal_, lambda x: nn.init.constant_(x, 0), np.sqrt(2)
        )

        if recurrent:
            num_inputs = hidden_size

        # self.actor = init_(nn.Linear(num_inputs, hidden_size)), nn.Tanh(),
        self.encoder = encoder
        self.critic_linear = init_(nn.Linear(num_inputs, 1))
        self.train()

    def forward(self, inputs, rnn_hxs, masks):
        if args.weights_path == "None":
            x = self.encoder(inputs)
        else:
            with torch.no_grad():
                x = self.encoder(inputs)
        if self.is_recurrent:
            x, rnn_hxs = self._forward_gru(x, rnn_hxs, masks)
        return self.critic_linear(x), x, rnn_hxs

def get_agent(args, envs, encoder, device):

    actor_critic = Policy(
        [encoder.feature_size], envs.action_space, base=SimpleBase, base_kwargs={"encoder": encoder}
    )
    actor_critic.to(device)
    agent = algo.PPO(
        actor_critic,
        args.ppo_clip_param,
        args.ppo_epoch,
        args.ppo_num_mini_batch,
        args.ppo_value_loss_coef,
        args.ppo_entropy_coef,
        lr=args.ppo_lr,
        eps=args.ppo_eps,
        max_grad_norm=args.ppo_max_grad_norm,
    )
    return agent, actor_critic

def train(args, envs, encoder, agent, actor_critic, device):
    rollouts = RolloutStorage(
        args.num_steps,
        args.num_processes,
        envs.observation_space.shape,
        envs.action_space,
        actor_critic.recurrent_hidden_state_size,
    )

    obs = envs.reset()
    rollouts.obs[0].copy_(obs)
    rollouts.to(device)

    episode_rewards = deque(maxlen=10)

    start = time.time()
    num_updates = int(args.num_env_steps) // args.num_steps // args.num_processes
    for j in range(num_updates):

        if args.ppo_use_linear_lr_decay:
            # decrease learning rate linearly
            utils.update_linear_schedule(agent.optimizer, j, num_updates, args.ppo_lr)

        for step in range(args.num_steps):
            # Sample actions
            with torch.no_grad():
                value, action, action_log_probs, recurrent_hidden_states, actor_features, dist_entropy = actor_critic.act(
                    rollouts.obs[step],
                    rollouts.recurrent_hidden_states[step],
                    rollouts.masks[step],
                )

            # Obser reward and next obs
            obs, reward, done, infos = envs.step(action)

            # TODO: Check that the encoder is not updated
            # TODO: Analyze features of vae and infonce-st encoder

            for info in infos:
                if "episode" in info.keys():
                    episode_rewards.append(info["episode"]["r"])

            # If done then clean the history of observations.
            masks = torch.FloatTensor([[0.0] if done_ else [1.0] for done_ in done])
            bad_masks = torch.FloatTensor(
                [[0.0] if "bad_transition" in info.keys() else [1.0] for info in infos]
            )
            rollouts.insert(
                obs,
                recurrent_hidden_states,
                action,
                action_log_probs,
                value,
                reward,
                masks,
                bad_masks,
            )

        with torch.no_grad():
            next_value = actor_critic.get_value(
                rollouts.obs[-1],
                rollouts.recurrent_hidden_states[-1],
                rollouts.masks[-1],
            )

        rollouts.compute_returns(
            next_value, False, args.ppo_gamma, 0.0, args.use_proper_time_limits
        )

        value_loss, action_loss, dist_entropy = agent.update(rollouts)

        rollouts.after_update()

        # save for every interval-th episode or for the last epoch
        if j % args.save_interval == 0 or j == num_updates - 1:
            torch.save(
                [actor_critic, getattr(utils.get_vec_normalize(envs), "ob_rms", None)],
                os.path.join(wandb.run.dir, args.env_name + ".pt"),
            )

        if j % args.log_interval == 0 and len(episode_rewards) > 1:
            total_num_steps = (j + 1) * args.num_processes * args.num_steps
            end = time.time()
            print(
                "Updates {}, num timesteps {}, FPS {} \n Last {} training episodes: mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}\n".format(
                    j,
                    total_num_steps,
                    int(total_num_steps / (end - start)),
                    len(episode_rewards),
                    np.mean(episode_rewards),
                    np.median(episode_rewards),
                    np.min(episode_rewards),
                    np.max(episode_rewards),
                )
            )
            wandb.log(
                {
                    "updates": j,
                    "total_num_steps": total_num_steps,
                    "fps": int(total_num_steps / (end - start)),
                    "episode_rewards_mean": np.mean(episode_rewards),
                    "episode_rewards_median": np.median(episode_rewards),
                    "episode_rewards_min": np.min(episode_rewards),
                    "episode_rewards_max": np.max(episode_rewards),
                    "entropy": dist_entropy,
                    "value_loss": value_loss,
                    "policy_loss": action_loss,
                }
            )

def run_rl(args):
    device = torch.device(
        "cuda:" + str(args.cuda_id) if torch.cuda.is_available() else "cpu"
    )
    envs = get_envs(
        env_name=args.env_name,
        seed=args.seed,
        num_processes=args.num_processes,
        num_frame_stack=args.num_frame_stack,
        downsample=not args.no_downsample,
        color=args.color,
    )
    encoder = get_encoder(args, envs.observation_space.shape, device)
    agent, actor_critic = get_agent(args, envs, encoder, device)
    wandb.watch(actor_critic)
    train(args, envs, encoder, agent, actor_critic, device)

if __name__ == "__main__":
    parser = get_argparser()
    args = parser.parse_args()
    tags = ["rl"]
    wandb.init(project=args.wandb_proj, tags=tags)
    config = {}
    config.update(vars(args))
    wandb.config.update(config)
    run_rl(args)


from argparse import ArgumentParser
from pathlib import Path

import numpy as np
import pickle
from pypinyin import lazy_pinyin, Style
from utils.config_manager import Config
from data.audio import Audio
import os
import librosa
import tensorflow as tf
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

def genSpeakerEmbedding(wavpath):
    def load_data(path, sr=16000, win_length=400, hop_length=160, n_fft=512):
        try:
            wav, sr_ret = librosa.load(path, sr=sr)
            intervals = librosa.effects.split(wav, top_db=20)
            wav_output = []
            for sliced in intervals:
                wav_output.extend(wav[sliced[0]:sliced[1]])
            wav = np.array(wav_output)
        except Exception as e:
            print("Exception happened when load_data('{}'): {}".format(path, str(e)))
            return None

        linear_spect = librosa.stft(wav, n_fft=n_fft, win_length=win_length, hop_length=hop_length).T
        mag, _ = librosa.magphase(linear_spect)  # magnitude
        spec_mag = mag.T

        mu = np.mean(spec_mag, 0, keepdims=True)
        std = np.std(spec_mag, 0, keepdims=True)
        spec_mag = (spec_mag - mu) / (std + 1e-5)

        return spec_mag

    model = tf.saved_model.load('spkpb')
    graphModel = model.signatures['serving_default']

    specs = load_data(wavpath)
    specs = np.expand_dims(np.expand_dims(specs, 0), -1)
    embedding = graphModel(tf.convert_to_tensor(specs, tf.float32))['embeddings'][0]
    return embedding

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--config', '-c', dest='config', default='config/session_paths.yaml')
    parser.add_argument('--text', '-t', dest='text', default='直接赋值使用的是引用的方式。而有些情况下需要复制整个对象', type=str)
    parser.add_argument('--file', '-f', dest='file', default=None, type=str)
    parser.add_argument('--checkpoint', '-ckpt', dest='checkpoint', default=None, type=str)
    parser.add_argument('--outdir', '-o', dest='outdir', default=None, type=str)
    parser.add_argument('--store_mel', '-m', dest='store_mel', action='store_true')
    parser.add_argument('--verbose', '-v', dest='verbose', action='store_true')
    parser.add_argument('--all_weights', '-ww', dest='all_weights', action='store_true')
    parser.add_argument('--single', '-s', dest='single', action='store_true')
    args = parser.parse_args()
    # spk_emb_dict = pickle.load(open(str('/root/mydata/Corpus/transformer_tts_data.corpus/spk_emb.pkl'), 'rb'))
    spk_emb_dict = pickle.load(open(str('test_spk_emb.pkl'), 'rb'))
    spk_emb = spk_emb_dict['x_2'][np.newaxis, :]
    spk_emb = genSpeakerEmbedding('ldh.wav')[np.newaxis, :]

    if args.file is not None:
        with open(args.file, 'r') as file:
            text = file.readlines()
        fname = Path(args.file).stem
    elif args.text is not None:
        text = [args.text]
        fname = 'custom_text'
    else:
        fname = None
        text = None
        print(f'Specify either an input text (-t "some text") or a text input file (-f /path/to/file.txt)')
        exit()
    config_loader = Config(config_path=args.config)
    outdir = Path(args.outdir) if args.outdir is not None else config_loader.log_dir
    outdir = outdir / 'outputs' / f'{fname}'
    outdir.mkdir(exist_ok=True, parents=True)
    print('==='*10,outdir)
    audio = Audio(config_loader.config)
    if args.checkpoint is not None:
        all_weights = [args.checkpoint]

    elif args.all_weights:
        all_weights = [(config_loader.weights_dir / x.stem).as_posix() for x in config_loader.weights_dir.iterdir() if
                       x.suffix == '.index']
    else:
        all_weights = [None]  # default

    if args.verbose:
        print(f'\nWeights list: \n{all_weights}\n')
    for weights in all_weights:
        model = config_loader.load_model(weights)
        file_name = f'{fname}_ttsstep{model.step}'
        print(f'Output wav under {outdir}')
        wavs = []
        for i, text_line in enumerate(text):
            out = model.predict(text_line, spk_emb=spk_emb, encode=True, phoneme_max_duration=None, speed_regulator=0.8)
            mel = out['mel'].numpy().T
            wav = audio.reconstruct_waveform(mel)
            wavs.append(wav)
            if args.store_mel:
                np.save((outdir / (file_name + f'_{i}')).with_suffix('.mel'), out['mel'].numpy())
            if args.single:
                audio.save_wav(wav, (outdir / (file_name + f'_{i}')).with_suffix('.wav'))
        audio.save_wav(np.concatenate(wavs), (outdir / file_name).with_suffix('.wav'))


import math
import numpy as np
from paddle.io import BatchSampler
from .data_utils import ordered_indices, num_tokens_vec_fn, get_batches_indices

class DistributedDynamicBatchSamplerV2(BatchSampler):
    ''' 支持多卡训练的动态bsz采样器,与fairseq对齐。 10/2
        修正v1 bug:
        1.v1不支持max_sentence
        *2.v1会丢失部分样本，test得分开写
        *3.v1的src与tgt基本上句长差不多，不符合实际
        4.v1的tokens数容易溢出
    '''

    def __init__(self,
                 dataset,
                 mode='train',
                 max_tokens=4000,
                 max_sentences=None,
                 bsz_factor=1,
                 num_replicas=None,
                 rank=None,
                 drop_last=False):
        self.dataset = dataset
        assert mode in ['train', 'dev', 'test']
        self.shuffle = mode == 'train'
        self.src_sizes = np.array([len(data[0])+1 for data in dataset])
        self.tgt_sizes = np.array([len(data[1])+1 for data in dataset]) if mode != 'test' else None
        # self.num_tokens_fn = lambda idx:self.dataset[idx]+1 # 长度dset,一定要加eos或sos！！
        assert max_tokens is not None or max_sentences is not None, \
            "max_tokens and max_sentences should not be null at the same time, please specify one parameter at least"
        self.max_tokens = max_tokens
        self.max_sentences = max_sentences
        assert isinstance(bsz_factor, int) and bsz_factor > 0, \
            "bsz_factor should be a positive integer"
        self.bsz_factor = bsz_factor
        from paddle.fluid.dygraph.parallel import ParallelEnv
        if num_replicas is not None:
            assert isinstance(num_replicas, int) and num_replicas > 0, \
                "num_replicas should be a positive integer"
            self.nranks = num_replicas
        else:
            self.nranks = ParallelEnv().nranks

        if rank is not None:
            assert isinstance(rank, int) and rank >= 0, \
                "rank should be a non-negative integer"
            self.local_rank = rank
        else:
            self.local_rank = ParallelEnv().local_rank
        assert isinstance(drop_last, bool), \
            "drop_last should be a boolean number"
        self.drop_last = drop_last  # 如果多余了就不变，bool不变，然后删除最后一个；如果没多余
        self.epoch = 0

    def __iter__(self):
        # get indices and shuffle samples
        indices = ordered_indices(src_sizes=self.src_sizes, tgt_sizes=self.tgt_sizes, epoch=self.epoch,
                                  shuffle=self.shuffle)

        # subsample for rank
        batches_indices = self._get_batches_by_max_tokens(indices)

        if self.shuffle:
           self.epoch += 1

        # yield batch indices
        _batch_iter = iter(batches_indices)
        for batch_indices in _batch_iter:
            yield batch_indices

    def __len__(self):
        # we do not know the exactly batch size, so do not call len(dataloader)
        pass

    def set_epoch(self, epoch):
        '''
         Sets the epoch number. When :attr:`shuffle=True`, this number is used
        as seeds of random numbers. By default, users may not set this, all
        replicas (workers) use a different random ordering for each epoch.
        If set same number at each epoch, this sampler will yield the same
        ordering at all epoches.
        '''
        self.epoch = epoch

    def _get_batches_by_max_tokens(self, indices):
        num_tokens_vec = num_tokens_vec_fn(indices, self.src_sizes, self.tgt_sizes)
        batches_indices = get_batches_indices(indices,
                                              num_tokens_vec=num_tokens_vec,
                                              max_tokens=self.max_tokens,
                                              max_sentences=self.max_sentences,
                                              bsz_factor=self.bsz_factor)
        # process last batch
        if self.drop_last and len(batches_indices[-1]) % self.bsz_factor != 0:
            batches_indices.pop()

        # shuffle batches
        if self.shuffle:
            shuffle_idx=np.random.RandomState(self.epoch).permutation(len(batches_indices)).astype(np.int64)
            batches_indices=np.array(batches_indices)[shuffle_idx].tolist()

        # subsample batches_indices for ranks
        if self.nranks > 1:
            local_batches_indices = []
            last_batches = len(batches_indices) % self.nranks  # 多余的batch
            # 补全batches
            if last_batches > 0:
                batches_indices.extend(batches_indices[:(self.nranks - last_batches)])
            assert len(batches_indices) % self.nranks == 0  # 确保batch数是nrank的倍数

            # sabsample for each process
            for i in range(0, len(batches_indices), self.nranks):
                local_batches_indices.append(batches_indices[i])
            return local_batches_indices
        # single process
        return batches_indices


import numpy
import os
import sys

from nematus.nmt import train
from sequence_qe.dataset import mkdir_p

# are src and trg vocab sizes actually required? -- these could be obtained directly from vocab indices
# For AmuNMT WMT 2016 pretrained models
SRC_VOCAB_SIZE=40587
MT_VOCAB_SIZE=40822
PE_VOCAB_SIZE=41101

PATIENCE = 99999999999 #never early stop
SRC = "en"
TGT = "de"

VOCAB_DIR = "/media/1tb_drive/nematus_ape_experiments/amunmt_ape_pretrained/system/models"
SRC_VOCAB = os.path.join(VOCAB_DIR, 'src-pe/vocab.src.json')
MT_VOCAB = os.path.join(VOCAB_DIR, 'mt-pe/vocab.mt.json')
PE_VOCAB = os.path.join(VOCAB_DIR, 'mt-pe/vocab.pe.json')

TRAIN_DATA_DIR = "/media/1tb_drive/Dropbox/data/qe/amunmt_artificial_ape_2016/data/500K_and_20x_task_internal"
SRC_TRAIN = os.path.join(TRAIN_DATA_DIR, 'train.mt.factor_corpus')
TRG_TRAIN = os.path.join(TRAIN_DATA_DIR, 'train.pe.prepped')

# WMT 16 EN-DE QE/APE DEV Data
QE_DATA_DIR = "/media/1tb_drive/Dropbox/data/qe/ape/concat_wmt_2016_2017"
SRC_DEV = os.path.join(QE_DATA_DIR, 'dev.mt.factor_corpus')
TRG_DEV = os.path.join(QE_DATA_DIR, 'dev.pe.prepped')

mkdir_p('model')

# start training from best model from previous experiment
STARTING_MODEL = '/media/1tb_drive/nematus_ape_experiments/ape_qe/en-de/model/model.npz.npz.best_bleu'

if __name__ == '__main__':
    validerr = train(saveto=os.path.join('model/model.npz'),
                     prior_model=STARTING_MODEL,
                     reload_=True,
                     dim_word=256,
                     dim=512,
                     n_words=PE_VOCAB_SIZE,
                     n_words_src=SRC_VOCAB_SIZE,
                     decay_c=0.,
                     clip_c=1.,
                     lrate=0.0001,
                     optimizer='adadelta',
                     maxlen=50,
                     batch_size=32,
                     valid_batch_size=32,
                     datasets=[SRC_TRAIN, TRG_TRAIN],
                     valid_datasets=[SRC_DEV, TRG_DEV],
                     dictionaries=[MT_VOCAB, SRC_VOCAB, PE_VOCAB],
                     tie_encoder_decoder_embeddings=False,
                     factors=2,
                     dim_per_factor=[128, 128],
                     validFreq=1000,
                     dispFreq=500,
                     saveFreq=5000,
                     sampleFreq=1000,
                     use_dropout=False,
                     dropout_embedding=0.2, # dropout for input embeddings (0: no dropout)
                     dropout_hidden=0.2, # dropout for hidden layers (0: no dropout)
                     dropout_source=0.1, # dropout source words (0: no dropout)
                     dropout_target=0.1, # dropout target words (0: no dropout)
                     overwrite=False,
                     external_validation_script='./fine_tune_validate.sh')
    print validerr