# own mini version of dreamer for processed features
import gym
import os
from matplotlib.pyplot import axis, step
import torch
import threading
import numpy as np

from atariari.benchmark.wrapper import AtariARIWrapper
from torch.utils.tensorboard import SummaryWriter
from torch.distributions import Categorical
from tqdm import tqdm
from torch.optim import Adam, SGD
from itertools import count
from rtpt import RTPT
from time import sleep, time

from xrl.algorithms.minidr.dataset import ModelDataset
from xrl.algorithms.minidr.model import WorldPredictor, Policy

from xrl.features import feature_processing

torch.autograd.set_detect_anomaly(True)

# world pred parameter
batch = 10
history_size = 50
min_history = 10
lr_pred = 1e-3
alpha = 1e-1   #alpha to scale down loss of state for loss of reward #old 1e-5

# policy parameter
lr_policy = 1e-3
gamma = 0.97

# misc parameter

# helper function to save all models
def save(
    path,
    world_predictor,
    policy,
    optim_predictor,
    optim_policy,
    episode,
    global_step):
    # save model
    print("Saving {}".format(path))
    torch.save({
            'world_predictor': world_predictor.state_dict(),
            'policy': policy.state_dict(),
            'optim_predictor': optim_predictor.state_dict(),
            'optim_policy': optim_policy.state_dict(),
            'episode': episode,
            'global_step': global_step
            }, path)
    return None

# helper function to load all models
def load(path,
    world_predictor,
    policy,
    optim_predictor,
    optim_policy):
    print("{} does exist, loading ... ".format(path))
    checkpoint = torch.load(path)
    world_predictor.load_state_dict(checkpoint['world_predictor'])
    policy.load_state_dict(checkpoint['policy'])
    if optim_predictor is not None:
        optim_predictor.load_state_dict(checkpoint['optim_predictor'])
    if optim_policy is not None:
        optim_policy.load_state_dict(checkpoint['optim_policy'])
    episode = checkpoint['episode']
    steps_done = checkpoint['global_step']
    return world_predictor, policy, optim_predictor, optim_policy, episode, steps_done

# function to select action
# TODO: fix for planning
def select_action(n_actions, policy):
    action = np.random.randint(0, n_actions)
    return action, None

def train(cfg):
    ### HYPERPARAMETERS ###
    print('Experiment name:', cfg.exp_name)
    torch.manual_seed(cfg.seed)
    print('Seed:', torch.initial_seed())
    print('Env name:', cfg.env_name)
    print('Using device:', cfg.device)
    if 'cuda' in cfg.device:
        print('Using parallel:', cfg.parallel)
    if cfg.parallel:
        print('Device ids:', cfg.device_ids)

    # params for saving loading
    save_path = os.getcwd() + "/xrl/checkpoints/"
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    save_path = save_path + cfg.exp_name + "_model.pth"

    # device to use for training
    device = cfg.device

    # logger
    writer = SummaryWriter(os.getcwd() + cfg.logdir + cfg.exp_name)

    # params for game
    env = AtariARIWrapper(gym.make(cfg.env_name))
    n_actions = env.action_space.n
    env.reset()
    raw_features, features, _, _ = feature_processing.do_step(env)
    len_raw_features = 12
    # init policy net

    ### MODELS ###
    predictor = WorldPredictor(input=len_raw_features).to(device)
    policy = Policy(len(features), cfg.train.hidden_layer_size, n_actions).to(device)
    criterion = torch.nn.MSELoss()
    optim_predictor = Adam(predictor.parameters(), lr=lr_pred)
    optim_policy = Adam(policy.parameters(), lr=lr_policy)

    history = []
    steps_done = [0]
    i_episode = 1

    # load if checkpoint exists
    if os.path.isfile(save_path):
        predictor, policy, optim_predictor, optim_policy, i_episode, steps_done = load(save_path, predictor, policy, optim_predictor, optim_policy)

    # inner function to give it access to shared variables
    # loop to interactively fill history with newest episodes
    def gather_episode():
        with torch.no_grad():
            rtpt_s = RTPT(name_initials='DV', experiment_name=cfg.exp_name + "_interactive",
                    max_iterations=cfg.train.num_episodes)
            rtpt_s.start()
            while steps_done[0] < cfg.train.max_steps:
                obs = env.reset()
                raw_features, features, _, _ = feature_processing.do_step(env)
                episode = []
                done = False
                r_sum = 0
                t = 0
                while t < 50000:
                    a, _ = select_action(n_actions, policy)
                    # create episode entry with last state, action, state and reward
                    episode.append(raw_features)
                    episode.append(a)
                    raw_features, features, reward, done = feature_processing.do_step(env, a, raw_features)
                    # append reward to features to have both
                    episode.append(raw_features)
                    episode.append(reward)
                    r_sum += reward
                    #print("Step: {}, Reward: {}".format(t, rew), end="\r")
                    steps_done[0] += 1
                    if done:
                        break
                    t += 1
                writer.add_scalar('Train/Reward', r_sum, steps_done[0])
                history.append(episode)
                for _ in range(len(history) - history_size):
                    history.pop(0)
                rtpt_s.step()

    #start gathering episode thread
    t = threading.Thread(target=gather_episode)
    t.start()

    print("Dataset init")
    while len(history) < min_history:
        # check every second if first history entry is inside
        print("Wainting until history large enough (current size: {}, needed: {})".format(len(history), min_history), end="\r")
        sleep(1.0)
    print("done")

    ### DATASET ###
    ds = ModelDataset(history, history_size=history_size)
    loader = torch.utils.data.DataLoader(
        ds,
        batch_size=batch,
        shuffle=True,
        num_workers=0,
        drop_last=False,
        pin_memory=True
    )

    rtpt = RTPT(name_initials='DV', experiment_name=cfg.exp_name,
                    max_iterations=cfg.train.num_episodes)
    rtpt.start()
    pbar = tqdm(total=cfg.train.max_steps)
    pbar.update(n=steps_done[0])
    while steps_done[0] < cfg.train.max_steps:
        tmp_steps_done = steps_done[0]
        loss_pred = 0
        loss_s = 0
        loss_r = 0
        len_h = 0

        # create copy of predictor for policy training
        #policy_predictor = WorldPredictor(input=len_raw_features).to(device)
        #policy_predictor.load_state_dict(predictor.state_dict())

        train_s = None

        for ls, a, s, r in loader:
            ls = ls.to(device)
            a = a.unsqueeze(1).to(device)
            s = s.to(device)
            r = r.unsqueeze(1).to(device)

            ### TRAIN PREDICTOR

            # predict
            ps, pr = predictor(ls, a)

            #print(ls[0])
            #print(s[0])
            #print(ps[0])
            #print((ls[0]-ps[0]).detach().numpy())
            #print((s[0]-ps[0]).detach().numpy())
            #print(torch.mean(ls[0]-ps[0]))
            #print(torch.mean(s[0]-ps[0]))
            #print("###")
            #print(r)
            #print(pr)
            #print(ss)

            # calc loss
            optim_predictor.zero_grad()
            loss_s = criterion(ps, s)
            loss_r = criterion(pr, r)
            loss_pred = alpha * loss_s + loss_r
            loss_pred.backward()
            optim_predictor.step()

            ### TRAIN POLICY

            # TODO: planning

            if i_episode % 5 == 0:
                # log all
                writer.add_scalar('Train/Loss World Predictor', loss_pred, steps_done[0])
                writer.add_scalar('Train/Loss World Predictor State', loss_s, steps_done[0])
                writer.add_scalar('Train/Loss World Predictor Reward', loss_r, steps_done[0])
                writer.add_scalar('Train/Diff State Identity', torch.mean(ls-ps), steps_done[0])
                writer.add_scalar('Train/Diff State Next State', torch.mean(s-ps), steps_done[0])
                # set logging variables
                len_h = len(history)
                # display
                pbar.set_postfix(
                    l_pred = loss_pred.item(),
                    len_h = len_h,
                    i_episode = i_episode
                )

        # update tqdm
        pbar.update(n = steps_done[0] - tmp_steps_done)

        # save
        if i_episode % cfg.train.save_every == 0:
            save(save_path, predictor, policy, optim_predictor, optim_policy, i_episode, steps_done)

        # end episode
        i_episode += 1
        rtpt.step()


# Copyright 2022 Sony Semiconductors Israel, Inc. 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.
# ==============================================================================
from typing import Callable, List

from model_compression_toolkit import common
from model_compression_toolkit.common import Logger
from model_compression_toolkit.common.constants import PYTORCH
from model_compression_toolkit.common.gptq.gptq_config import GradientPTQConfig
from model_compression_toolkit.common.target_platform import TargetPlatformCapabilities
from model_compression_toolkit.common.mixed_precision.kpi import KPI
from model_compression_toolkit.common.framework_info import FrameworkInfo
from model_compression_toolkit.common.network_editors.actions import EditRule
from model_compression_toolkit.common.quantization.core_config import CoreConfig
from model_compression_toolkit.common.quantization.debug_config import DebugConfig
from model_compression_toolkit.common.mixed_precision.mixed_precision_quantization_config import \
    MixedPrecisionQuantizationConfig, DEFAULT_MIXEDPRECISION_CONFIG, MixedPrecisionQuantizationConfigV2
from model_compression_toolkit.common.post_training_quantization import post_training_quantization
from model_compression_toolkit.common.quantization.quantization_config import QuantizationConfig
from model_compression_toolkit.common.quantization.quantization_config import DEFAULTCONFIG

import importlib

if importlib.util.find_spec("torch") is not None:
    from model_compression_toolkit.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
    from model_compression_toolkit.pytorch.pytorch_implementation import PytorchImplementation
    from model_compression_toolkit.pytorch.constants import DEFAULT_TP_MODEL
    from torch.nn import Module

    from model_compression_toolkit import get_target_platform_capabilities
    DEFAULT_PYTORCH_TPC = get_target_platform_capabilities(PYTORCH, DEFAULT_TP_MODEL)

    def pytorch_post_training_quantization(in_module: Module,
                                           representative_data_gen: Callable,
                                           n_iter: int = 500,
                                           quant_config: QuantizationConfig = DEFAULTCONFIG,
                                           fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
                                           network_editor: List[EditRule] = [],
                                           gptq_config: GradientPTQConfig = None,
                                           analyze_similarity: bool = False,
                                           target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
        """
        Quantize a trained Pytorch module using post-training quantization.
        By default, the module is quantized using a symmetric constraint quantization thresholds
        (power of two) as defined in the default TargetPlatformCapabilities.
        The module is first optimized using several transformations (e.g. BatchNormalization folding to
        preceding layers). Then, using a given dataset, statistics (e.g. min/max, histogram, etc.) are
        being collected for each layer's output (and input, depends on the quantization configuration).
        Thresholds are then being calculated using the collected statistics and the module is quantized
        (both coefficients and activations by default).
        If gptq_config is passed, the quantized weights are optimized using gradient based post
        training quantization by comparing points between the float and quantized modules, and minimizing the
        observed loss.

        Args:
            in_module (Module): Pytorch module to quantize.
            representative_data_gen (Callable): Dataset used for calibration.
            n_iter (int): Number of calibration iterations to run.
            quant_config (QuantizationConfig): QuantizationConfig containing parameters of how the module should be quantized. `Default configuration. <https://github.com/sony/model_optimization/blob/21e21c95ca25a31874a5be7af9dd2dd5da8f3a10/model_compression_toolkit/common/quantization/quantization_config.py#L154>`_
            fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g., kernel channels indices, groups of layers by how they should be quantized, etc.). `Default PyTorch info <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/pytorch/default_framework_info.py>`_
            network_editor (List[EditRule]): List of EditRules. Each EditRule consists of a node filter and an action to change quantization settings of the filtered nodes.
            gptq_config (GradientPTQConfig): Configuration for using gptq (e.g. optimizer).
            analyze_similarity (bool): Whether to plot similarity figures within TensorBoard (when logger is enabled) or not.
            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to. `Default Keras TPC <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/tpc_models/pytorch_tp_models/pytorch_default.py>`_

        Returns:
            A quantized module and information the user may need to handle the quantized module.

        Examples:

            Import a Pytorch module:

            >>> import torchvision.models.mobilenet_v2 as models
            >>> module = models.mobilenet_v2()

            Create a random dataset generator:

            >>> import numpy as np
            >>> def repr_datagen(): return [np.random.random((1,224,224,3))]

            Import mct and pass the module with the representative dataset generator to get a quantized module:

            >>> import model_compression_toolkit as mct
            >>> quantized_module, quantization_info = mct.pytorch_post_training_quantization(module, repr_datagen)

        """

        return post_training_quantization(in_module,
                                          representative_data_gen,
                                          CoreConfig(n_iter, quant_config,
                                                     debug_config=DebugConfig(analyze_similarity=analyze_similarity,
                                                                              network_editor=network_editor)),
                                          fw_info,
                                          PytorchImplementation(),
                                          target_platform_capabilities,
                                          gptq_config)

    def pytorch_post_training_quantization_mixed_precision(in_model: Module,
                                                           representative_data_gen: Callable,
                                                           target_kpi: KPI,
                                                           n_iter: int = 500,
                                                           quant_config: MixedPrecisionQuantizationConfig = DEFAULT_MIXEDPRECISION_CONFIG,
                                                           fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
                                                           network_editor: List[EditRule] = [],
                                                           gptq_config: GradientPTQConfig = None,
                                                           analyze_similarity: bool = False,
                                                           target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
        """
         Quantize a pretrained Pytorch model using post-training quantization. By default, the model is
         quantized using a symmetric constraint quantization thresholds (power of two) as defined in the
         default TargetPlatformCapabilities.
         The model is first optimized using several transformations (e.g. BatchNormalization folding to
         preceding layers). Then, using a given dataset, statistics (e.g. min/max, histogram, etc.) are
         being collected for each layer's output (and input, depends on the quantization configuration).
         For each possible bit width (per layer) a threshold is then being calculated using the collected
         statistics. Then, using an ILP solver we find a mixed-precision configuration, and set a bit width
         for each quantizer (for both activations and weights quantizers, by default).
         In order to limit the maximal model's size, a target KPI need to be passed after weights_memory
         is set (in bytes).
         The model is then quantized (both coefficients and activations by default).
         If gptq_config is passed, the quantized weights are optimized using gradient based post
         training quantization by comparing points between the float and quantized models, and minimizing the
         observed loss.
         Notice that this feature is experimental.

         Args:
             in_model (Model): Pytorch model to quantize.
             representative_data_gen (Callable): Dataset used for calibration.
             target_kpi (KPI): KPI object to limit the search of the mixed-precision configuration as desired.
             n_iter (int): Number of calibration iterations to run.
             quant_config (MixedPrecisionQuantizationConfig): QuantizationConfig containing parameters of how the model should be quantized.
             fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g., kernel channels indices, groups of layers by how they should be quantized, etc.). `Default PyTorch info <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/pytorch/default_framework_info.py>`_
             network_editor (List[EditRule]): List of EditRules. Each EditRule consists of a node filter and an action to change quantization settings of the filtered nodes.
             gptq_config (GradientPTQConfig): Configuration for using GPTQ (e.g. optimizer).
             analyze_similarity (bool): Whether to plot similarity figures within TensorBoard (when logger is enabled) or not.
             target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to. `Default Keras TPC <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/tpc_models/pytorch_tp_models/pytorch_default.py>`_

         Returns:
             A quantized model and information the user may need to handle the quantized model.

         Examples:

             Import MCT:

             >>> import model_compression_toolkit as mct

             Import a Pytorch model:

             >>> import torchvision.models.mobilenet_v2 as models
             >>> module = models.mobilenet_v2()

             Create a random dataset generator:

             >>> import numpy as np
             >>> def repr_datagen(): return [np.random.random((1,224,224,3))]

             Create a mixed-precision configuration, to quantize a model with different bitwidths for different layers.
             The candidates bitwidth for quantization should be defined in the target platform model:

             >>> config = mct.MixedPrecisionQuantizationConfig()

             Create a KPI object to limit our returned model's size. Note that this value affects only coefficients
             that should be quantized (for example, the kernel of Conv2D in Keras will be affected by this value,
             while the bias will not):

             >>> kpi = mct.KPI(sum(p.numel() for p in module.parameters()) * 0.75)  # About 0.75 of the model size when quantized with 8 bits.

             Pass the model, the representative dataset generator, the configuration and the target KPI to get a
             quantized model:

             >>> quantized_model, quantization_info = mct.pytorch_post_training_quantization_mixed_precision(module, repr_datagen, n_iter=10, quant_config=config, target_kpi=kpi)

             For more configuration options, please take a look at our `API documentation <https://sony.github.io/model_optimization/api/api_docs/modules/mixed_precision_quantization_config.html>`_.

         """

        if not isinstance(quant_config, MixedPrecisionQuantizationConfig):
            common.Logger.error("Given quantization config to mixed-precision facade is not of type "
                                "MixedPrecisionQuantizationConfig. Please use pytorch_post_training_quantization API, "
                                "or pass a valid mixed precision configuration.")

        common.Logger.info("Using experimental mixed-precision quantization. "
                           "If you encounter an issue please file a bug.")

        quantization_config, mp_config = quant_config.separate_configs()
        core_config = CoreConfig(n_iter,
                                 quantization_config=quantization_config,
                                 mixed_precision_config=mp_config,
                                 debug_config=DebugConfig(analyze_similarity=analyze_similarity,
                                                          network_editor=network_editor))

        return post_training_quantization(in_model,
                                          representative_data_gen,
                                          core_config,
                                          fw_info,
                                          PytorchImplementation(),
                                          target_platform_capabilities,
                                          gptq_config,
                                          target_kpi)

    def pytorch_post_training_quantization_experimental(in_module: Module,
                                                        representative_data_gen: Callable,
                                                        target_kpi: KPI = None,
                                                        core_config: CoreConfig = CoreConfig(),
                                                        fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
                                                        target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
        """
        Quantize a trained Pytorch module using post-training quantization.
        By default, the module is quantized using a symmetric constraint quantization thresholds
        (power of two) as defined in the default TargetPlatformCapabilities.
        The module is first optimized using several transformations (e.g. BatchNormalization folding to
        preceding layers). Then, using a given dataset, statistics (e.g. min/max, histogram, etc.) are
        being collected for each layer's output (and input, depends on the quantization configuration).
        Thresholds are then being calculated using the collected statistics and the module is quantized
        (both coefficients and activations by default).
        If gptq_config is passed, the quantized weights are optimized using gradient based post
        training quantization by comparing points between the float and quantized modules, and minimizing the
        observed loss.

        Args:
            in_module (Module): Pytorch module to quantize.
            representative_data_gen (Callable): Dataset used for calibration.
            target_kpi (KPI): KPI object to limit the search of the mixed-precision configuration as desired.
            core_config (CoreConfig): Configuration object containing parameters of how the model should be
            quantized, including mixed precision parameters.
            fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g., kernel channels indices, groups of layers by how they should be quantized, etc.). `Default PyTorch info <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/pytorch/default_framework_info.py>`_
            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to. `Default Keras TPC <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/tpc_models/pytorch_tp_models/pytorch_default.py>`_

        Returns:
            A quantized module and information the user may need to handle the quantized module.

        Examples:

            Import a Pytorch module:

            >>> import torchvision.models.mobilenet_v2 as models
            >>> module = models.mobilenet_v2()

            Create a random dataset generator:

            >>> import numpy as np
            >>> def repr_datagen(): return [np.random.random((1,224,224,3))]

            Import mct and pass the module with the representative dataset generator to get a quantized module:

            >>> import model_compression_toolkit as mct
            >>> quantized_module, quantization_info = mct.pytorch_post_training_quantization(module, repr_datagen)

        """

        if core_config.mixed_precision_enable:
            if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfigV2):
                common.Logger.error("Given quantization config to mixed-precision facade is not of type "
                                    "MixedPrecisionQuantizationConfigV2. Please use pytorch_post_training_quantization API,"
                                    "or pass a valid mixed precision configuration.")

            common.Logger.info("Using experimental mixed-precision quantization. "
                               "If you encounter an issue please file a bug.")

        return post_training_quantization(in_module,
                                          representative_data_gen,
                                          core_config,
                                          fw_info,
                                          PytorchImplementation(),
                                          target_platform_capabilities,
                                          target_kpi=target_kpi)

    def pytorch_gradient_post_training_quantization_experimental(in_module: Module,
                                                                 representative_data_gen: Callable,
                                                                 target_kpi: KPI = None,
                                                                 core_config: CoreConfig = CoreConfig(),
                                                                 fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
                                                                 gptq_config: GradientPTQConfig = None,
                                                                 target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
        """
        Quantize a trained Pytorch module using post-training quantization.
        By default, the module is quantized using a symmetric constraint quantization thresholds
        (power of two) as defined in the default TargetPlatformCapabilities.
        The module is first optimized using several transformations (e.g. BatchNormalization folding to
        preceding layers). Then, using a given dataset, statistics (e.g. min/max, histogram, etc.) are
        being collected for each layer's output (and input, depends on the quantization configuration).
        Thresholds are then being calculated using the collected statistics and the module is quantized
        (both coefficients and activations by default).
        If gptq_config is passed, the quantized weights are optimized using gradient based post
        training quantization by comparing points between the float and quantized modules, and minimizing the
        observed loss.

        Args:
            in_module (Module): Pytorch module to quantize.
            representative_data_gen (Callable): Dataset used for calibration.
            target_kpi (KPI): KPI object to limit the search of the mixed-precision configuration as desired.
            core_config (CoreConfig): Configuration object containing parameters of how the model should be
            quantized, including mixed precision parameters.
            fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g., kernel channels indices, groups of layers by how they should be quantized, etc.). `Default PyTorch info <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/pytorch/default_framework_info.py>`_
            gptq_config (GradientPTQConfig): Configuration for using gptq (e.g. optimizer).
            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to. `Default Keras TPC <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/tpc_models/pytorch_tp_models/pytorch_default.py>`_

        Returns:
            A quantized module and information the user may need to handle the quantized module.

        Examples:

            Import a Pytorch module:

            >>> import torchvision.models.mobilenet_v2 as models
            >>> module = models.mobilenet_v2()

            Create a random dataset generator:

            >>> import numpy as np
            >>> def repr_datagen(): return [np.random.random((1,224,224,3))]

            Import mct and pass the module with the representative dataset generator to get a quantized module:

            >>> import model_compression_toolkit as mct
            >>> quantized_module, quantization_info = mct.pytorch_post_training_quantization(module, repr_datagen)

        """

        if core_config.mixed_precision_enable:
            if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfigV2):
                common.Logger.error("Given quantization config to mixed-precision facade is not of type "
                                    "MixedPrecisionQuantizationConfigV2. Please use pytorch_post_training_quantization API,"
                                    "or pass a valid mixed precision configuration.")

            common.Logger.info("Using experimental mixed-precision quantization. "
                               "If you encounter an issue please file a bug.")

        return post_training_quantization(in_module,
                                          representative_data_gen,
                                          core_config,
                                          fw_info,
                                          PytorchImplementation(),
                                          target_platform_capabilities,
                                          gptq_config,
                                          target_kpi=target_kpi)

else:
    # If torch is not installed,
    # we raise an exception when trying to use these functions.
    def pytorch_post_training_quantization(*args, **kwargs):
        Logger.critical('Installing Pytorch is mandatory '
                        'when using pytorch_post_training_quantization. '
                        'Could not find the torch package.')

    def pytorch_post_training_quantization_mixed_precision(*args, **kwargs):
        Logger.critical('Installing tensorflow and tensorflow_model_optimization is mandatory '
                        'when using pytorch_post_training_quantization_mixed_precision. '
                        'Could not find Tensorflow package.')

    def pytorch_post_training_quantization_experimental(*args, **kwargs):
        Logger.critical('Installing Pytorch is mandatory '
                        'when using pytorch_post_training_quantization_experimental. '
                        'Could not find the torch package.')

    def pytorch_gradient_post_training_quantization_experimental(*args, **kwargs):
        Logger.critical('Installing Pytorch is mandatory '
                        'when using pytorch_gradient_post_training_quantization_experimental. '
                        'Could not find the torch package.')


#!/usr/bin/env python2
# Copyright (c) 2015-2017, NVIDIA CORPORATION.  All rights reserved.
"""
Functions for creating temporary LMDBs
Used in test_views
"""
from __future__ import print_function

import argparse
import os
import random
import re
import sys
import time

import lmdb
import numpy as np
import PIL.Image

try:
    xrange          # Python 2
except NameError:
    xrange = range  # Python 3

if __name__ == '__main__':
    dirname = os.path.dirname(os.path.realpath(__file__))
    sys.path.insert(0, os.path.join(dirname, '..', '..'))
    import digits.config  # noqa

from digits import utils  # noqa

# Import digits.config first to set the path to Caffe
import caffe.io  # noqa
import caffe_pb2  # noqa

IMAGE_SIZE = 10
TRAIN_IMAGE_COUNT = 1000
VAL_IMAGE_COUNT = 1000
TEST_IMAGE_COUNT = 10
DB_BATCH_SIZE = 100

def create_lmdbs(folder, file_list, image_count=None, db_batch_size=None):
    """
    Creates LMDBs for generic inference
    Returns the filename for a test image

    Creates these files in "folder":
        train_images/
        train_labels/
        val_images/
        val_labels/
        mean.binaryproto
        test.png
    """

    if image_count is None:
        train_image_count = TRAIN_IMAGE_COUNT
    else:
        train_image_count = image_count
    val_image_count = VAL_IMAGE_COUNT

    if db_batch_size is None:
        db_batch_size = DB_BATCH_SIZE

    # read file list
    images = []
    f = open(file_list)
    for line in f.readlines():
        line = line.strip()
        if not line:
            continue

        path = None
        # might contain a numerical label at the end
        match = re.match(r'(.*\S)\s+(\d+)$', line)
        if match:
            path = match.group(1)
            ground_truth = int(match.group(2))
            images.append([path, ground_truth])

    print("Found %d image paths in image list" % len(images))

    for phase, image_count in [
            ('train', train_image_count),
            ('val', val_image_count)]:

        print("Will create %d pairs of %s images" % (image_count, phase))

        # create DBs
        image_db = lmdb.open(os.path.join(folder, '%s_images' % phase),
                             map_async=True, max_dbs=0)
        label_db = lmdb.open(os.path.join(folder, '%s_labels' % phase),
                             map_async=True, max_dbs=0)

        # add up all images to later create mean image
        image_sum = None
        shape = None

        # save test images (one for each label)
        testImagesSameClass = []
        testImagesDifferentClass = []

        # arrays for image and label batch writing
        image_batch = []
        label_batch = []

        for i in xrange(image_count):
            # pick up random indices from image list
            index1 = random.randint(0, len(images) - 1)
            index2 = random.randint(0, len(images) - 1)
            # label=1 if images are from the same class otherwise label=0
            label = 1 if int(images[index1][1]) == int(images[index2][1]) else 0
            # load images from files
            image1 = np.array(utils.image.load_image(images[index1][0]))
            image2 = np.array(utils.image.load_image(images[index2][0]))
            if not shape:
                # initialize image sum for mean image
                shape = image1.shape
                image_sum = np.zeros((3, shape[0], shape[1]), 'float64')
            assert(image1.shape == shape and image2.shape == shape)

            # create BGR image: blue channel will contain first image,
            # green channel will contain second image
            image_pair = np.zeros(image_sum.shape)
            image_pair[0] = image1
            image_pair[1] = image2

            image_sum += image_pair

            # save test images on first pass
            if label > 0 and len(testImagesSameClass) < TEST_IMAGE_COUNT:
                testImagesSameClass.append(image_pair)
            if label == 0 and len(testImagesDifferentClass) < TEST_IMAGE_COUNT:
                testImagesDifferentClass.append(image_pair)

            # encode into Datum object
            image = image_pair.astype('uint8')
            datum = caffe.io.array_to_datum(image, -1)
            image_batch.append([str(i), datum])

            # create label Datum
            label_datum = caffe_pb2.Datum()
            label_datum.channels, label_datum.height, label_datum.width = 1, 1, 1
            label_datum.float_data.extend(np.array([label]).flat)
            label_batch.append([str(i), label_datum])

            if (i % db_batch_size == (db_batch_size - 1)) or (i == image_count - 1):
                _write_batch_to_lmdb(image_db, image_batch)
                _write_batch_to_lmdb(label_db, label_batch)
                image_batch = []
                label_batch = []

            if i % (image_count / 20) == 0:
                print("%d/%d" % (i, image_count))

        # close databases
        image_db.close()
        label_db.close()

        # save mean
        mean_image = (image_sum / image_count).astype('uint8')
        _save_mean(mean_image, os.path.join(folder, '%s_mean.binaryproto' % phase))
        _save_mean(mean_image, os.path.join(folder, '%s_mean.png' % phase))

        # create test images
        for idx, image in enumerate(testImagesSameClass):
            _save_image(image, os.path.join(folder, '%s_test_same_class_%d.png' % (phase, idx)))
        for idx, image in enumerate(testImagesDifferentClass):
            _save_image(image, os.path.join(folder, '%s_test_different_class_%d.png' % (phase, idx)))

    return

def _write_batch_to_lmdb(db, batch):
    """
    Write a batch of (key,value) to db
    """
    try:
        with db.begin(write=True) as lmdb_txn:
            for key, datum in batch:
                lmdb_txn.put(key, datum.SerializeToString())
    except lmdb.MapFullError:
        # double the map_size
        curr_limit = db.info()['map_size']
        new_limit = curr_limit * 2
        try:
            db.set_mapsize(new_limit)  # double it
        except AttributeError as e:
            version = tuple(int(x) for x in lmdb.__version__.split('.'))
            if version < (0, 87):
                raise ImportError('py-lmdb is out of date (%s vs 0.87)' % lmdb.__version__)
            else:
                raise e
        # try again
        _write_batch_to_lmdb(db, batch)

def _save_image(image, filename):
    # converting from BGR to RGB
    image = image[[2, 1, 0], ...]  # channel swap
    # convert to (height, width, channels)
    image = image.astype('uint8').transpose((1, 2, 0))
    image = PIL.Image.fromarray(image)
    image.save(filename)

def _save_mean(mean, filename):
    """
    Saves mean to file

    Arguments:
    mean -- the mean as an np.ndarray
    filename -- the location to save the image
    """
    if filename.endswith('.binaryproto'):
        blob = caffe_pb2.BlobProto()
        blob.num = 1
        blob.channels = mean.shape[0]
        blob.height = mean.shape[1]
        blob.width = mean.shape[2]
        blob.data.extend(mean.astype(float).flat)
        with open(filename, 'wb') as outfile:
            outfile.write(blob.SerializeToString())

    elif filename.endswith(('.jpg', '.jpeg', '.png')):
        _save_image(mean, filename)
    else:
        raise ValueError('unrecognized file extension')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Create-LMDB tool - DIGITS')

    # Positional arguments
    parser.add_argument('folder', help='Where to save the images')
    parser.add_argument('file_list', help='File list')

    # Optional arguments
    parser.add_argument('-c', '--image_count', type=int, help='How many images')

    args = vars(parser.parse_args())

    if os.path.exists(args['folder']):
        print('ERROR: Folder already exists')
        sys.exit(1)
    else:
        os.makedirs(args['folder'])

    print('Creating images at "%s" ...' % args['folder'])

    start_time = time.time()

    create_lmdbs(
        args['folder'],
        args['file_list'],
        image_count=args['image_count'],
    )

    print('Done after %s seconds' % (time.time() - start_time,))


import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch
from torch.autograd import Variable
from module.self_attention import SelfAttentive, AttentionOneParaPerChan
from module.torch_moji import TorchMoji

import pickle as pkl
import os
from tqdm import tqdm

NUM_EMO = 4

class HierarchicalPredictor(nn.Module):
    """
    A Hierarchical LSTM with for 3 turns dialogue
    """
    def __init__(self, embedding_dim, hidden_dim, vocab_size, USE_ELMO, ADD_LINEAR):
        super(HierarchicalPredictor, self).__init__()
        self.SENT_LSTM_DIM = hidden_dim
        self.bidirectional = True
        self.add_linear = ADD_LINEAR

        self.sent_lstm_directions = 2 if self.bidirectional else 1
        self.cent_lstm_att_fn = AttentionOneParaPerChan
        self.ctx_lstm__att_fn = AttentionOneParaPerChan

        self.deepmoji_model = TorchMoji(nb_classes=None,
                                        nb_tokens=50000,
                                        embed_dropout_rate=0.2,
                                        final_dropout_rate=0.2)
        self.deepmoji_dim = 2304
        self.deepmoji_out = 300
        self.deepmoji2linear = nn.Linear(self.deepmoji_dim, self.deepmoji_out)

        self.elmo_dim = 1024

        self.num_layers = 2
        self.use_elmo = USE_ELMO
        if not self.use_elmo:
            self.elmo_dim = 0

        self.a_lstm = nn.LSTM(embedding_dim + self.elmo_dim, hidden_dim, num_layers=self.num_layers, batch_first=True,
                            bidirectional=self.bidirectional, dropout=0.2)
        self.a_self_attention = self.cent_lstm_att_fn(self.sent_lstm_directions*hidden_dim)
        # self.a_layer_norm = BertLayerNorm(hidden_dim*self.sent_lstm_directions)

        self.b_lstm = nn.LSTM(embedding_dim + self.elmo_dim, hidden_dim, num_layers=self.num_layers, batch_first=True,
                            bidirectional=self.bidirectional, dropout=0.2)
        self.b_self_attention = self.cent_lstm_att_fn(self.sent_lstm_directions*hidden_dim)
        # self.b_layer_norm = BertLayerNorm(hidden_dim*self.sent_lstm_directions)

        self.ctx_bidirectional = True
        self.ctx_lstm_dim = 800
        self.ctx_lstm_directions = 2 if self.ctx_bidirectional else 1

        if not self.add_linear:
            self.deepmoji_out = self.deepmoji_dim

        self.context_lstm = nn.LSTM(self.SENT_LSTM_DIM * self.sent_lstm_directions + self.deepmoji_out, self.ctx_lstm_dim,
                                    num_layers=2, batch_first=True, dropout=0.2, bidirectional=self.ctx_bidirectional)
        self.ctx_self_attention = self.ctx_lstm__att_fn(self.ctx_lstm_directions * self.ctx_lstm_dim)
        # self.ctx_layer_norm = BertLayerNorm(self.ctx_lstm_dim*self.ctx_lstm_directions)
        # self.ctx_pooler = BertPooler(self.ctx_lstm_dim*self.ctx_lstm_directions)

        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)

        self.context_to_emo = nn.Linear(self.ctx_lstm_dim, NUM_EMO)
        self.drop_out = nn.Dropout(0.2)
        self.out2label = nn.Linear(self.ctx_lstm_dim * self.ctx_lstm_directions, NUM_EMO)
        self.out2binary = nn.Linear(self.ctx_lstm_dim * self.ctx_lstm_directions, 2)
        self.out2emo = nn.Linear(self.ctx_lstm_dim * self.ctx_lstm_directions, NUM_EMO-1)

    def init_hidden(self, x):
        batch_size = x.size(0)
        if self.bidirectional:
            h0 = Variable(torch.zeros(2*self.num_layers, batch_size, self.SENT_LSTM_DIM), requires_grad=False).cuda()
            c0 = Variable(torch.zeros(2*self.num_layers, batch_size, self.SENT_LSTM_DIM), requires_grad=False).cuda()
        else:
            h0 = Variable(torch.zeros(1*self.num_layers, batch_size, self.SENT_LSTM_DIM), requires_grad=False).cuda()
            c0 = Variable(torch.zeros(1*self.num_layers, batch_size, self.SENT_LSTM_DIM), requires_grad=False).cuda()
        return (h0, c0)

    @staticmethod
    def sort_batch(batch, lengths):
        seq_lengths, perm_idx = lengths.sort(0, descending=True)
        seq_tensor = batch[perm_idx]
        rever_sort = np.zeros(len(seq_lengths))
        for i, l in enumerate(perm_idx):
            rever_sort[l] = i
        return seq_tensor, seq_lengths, rever_sort.astype(int), perm_idx

    def lstm_forward(self, x, x_len, elmo_x, lstm, hidden=None, attention_layer=None):
        x, x_len_sorted, reverse_idx, perm_idx = self.sort_batch(x, x_len.view(-1))
        max_len = int(x_len_sorted[0])

        emb_x = self.embeddings(x)
        emb_x = self.drop_out(emb_x)
        emb_x = emb_x[:, :max_len, :]

        if self.use_elmo:
            elmo_x = elmo_x[perm_idx]
            elmo_x = self.drop_out(elmo_x)
            emb_x = torch.cat((emb_x, elmo_x), dim=2)

        packed_input = nn.utils.rnn.pack_padded_sequence(emb_x, x_len_sorted.cpu().numpy(), batch_first=True)
        if hidden is None:
            hidden = self.init_hidden(x)
        packed_output, hidden = lstm(packed_input, hidden)
        output, unpacked_len = nn.utils.rnn.pad_packed_sequence(packed_output, batch_first=True)

        attention_layer = None  # testing
        if attention_layer is None:
            seq_len = torch.LongTensor(unpacked_len).view(-1, 1, 1).expand(output.size(0), 1, output.size(2))
            seq_len = Variable(seq_len - 1).cuda()
            output = torch.gather(output, 1, seq_len).squeeze(1)
        else:
            if isinstance(attention_layer, AttentionOneParaPerChan):
                output, alpha = attention_layer(output, unpacked_len)
            else:
                unpacked_len = [int(x.data) for x in unpacked_len]
                # print(unpacked_len)
                max_len = max(unpacked_len)
                mask = [[1] * l + [0] * (max_len - l) for l in unpacked_len]
                mask = torch.FloatTensor(np.asarray(mask)).cuda()
                attention_mask = torch.ones_like(mask)
                extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
                # extended_attention_mask = extended_attention_mask.to(
                #       type=next(self.parameters()).dtype)  # fp16 compatibility
                extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
                output, alpha = attention_layer(output, extended_attention_mask)
                # out, att = self.attention_layer(lstm_out[:, -1:].squeeze(1), lstm_out)
                output = output[:, 0, :]

        return output[reverse_idx], (hidden[0][:, reverse_idx, :], hidden[1][:, reverse_idx, :])

    def forward(self, a, a_len, b, b_len, c, c_len, a_emoji, b_emoji, c_emoji, elmo_a=None, elmo_b=None, elmo_c=None):
        # Sentence LSTM A
        a_out, a_hidden = self.lstm_forward(a, a_len, elmo_a, self.a_lstm,
                                            attention_layer=self.a_self_attention)
        # a_out = a_out[:, 0, :]
        if self.add_linear:
            a_emoji = self.deepmoji_model(a_emoji)
            a_emoji = self.deepmoji2linear(a_emoji)
            a_emoji = F.relu(a_emoji)
            a_emoji = self.drop_out(a_emoji)
        else:
            a_emoji = self.deepmoji_model(a_emoji)
            a_emoji = F.relu(a_emoji)

        # a_out = torch.cat((F.relu(a_out), a_emoji), dim=1)
        a_out = torch.cat((a_out, a_emoji), dim=1)
        # Sentence LSTM B
        b_out, b_hidden = self.lstm_forward(b, b_len, elmo_b, self.a_lstm,
                                            attention_layer=self.b_self_attention)
        # b_out = b_out[:, 0, :]
        if self.add_linear:
            b_emoji = self.deepmoji_model(b_emoji)
            b_emoji = self.deepmoji2linear(b_emoji)
            b_emoji = F.relu(b_emoji)
            b_emoji = self.drop_out(b_emoji)
        else:
            b_emoji = self.deepmoji_model(b_emoji)
            b_emoji = F.relu(b_emoji)

        # b_out = torch.cat((F.relu(b_out), b_emoji), dim=1)
        b_out = torch.cat((b_out, b_emoji), dim=1)
        # Sentence LSTM A
        c_out, c_hidden = self.lstm_forward(c, c_len, elmo_c, self.a_lstm,  # a_hidden,
                                            attention_layer=self.a_self_attention)
        # c_out = c_out[:, 0, :]
        if self.add_linear:
            c_emoji = self.deepmoji_model(c_emoji)
            c_emoji = self.deepmoji2linear(c_emoji)
            c_emoji = F.relu(c_emoji)
            c_emoji = self.drop_out(c_emoji)
        else:
            c_emoji = self.deepmoji_model(c_emoji)
            c_emoji = F.relu(c_emoji)

        # c_out = torch.cat((F.relu(c_out), c_emoji), dim=1)
        c_out = torch.cat((c_out, c_emoji), dim=1)
        # Context LSTM
        context_in = torch.cat((a_out.unsqueeze(1), b_out.unsqueeze(1), c_out.unsqueeze(1)), dim=1)
        ctx_out, _ = self.context_lstm(context_in)
        if isinstance(self.ctx_self_attention, AttentionOneParaPerChan):
            ctx_out, _ = self.ctx_self_attention(ctx_out, torch.LongTensor([3 for _ in range(a.size(0))]))
        else:
            ctx_out, _ = self.ctx_self_attention(ctx_out)
            ctx_out = ctx_out[:, 0, :]

        # ctx_out = self.ctx_layer_norm(ctx_out)
        # ctx_out = self.ctx_pooler(ctx_out)

        # multi-task learning
        out1 = self.out2label(ctx_out)
        out2 = self.out2binary(ctx_out)
        out3 = F.sigmoid(self.out2emo(ctx_out))
        return out1, out2, out3

    def load_embedding(self, emb):
        self.embeddings.weight = nn.Parameter(torch.FloatTensor(emb))


#!/usr/bin/env python
import json
import os
import sys

sys.path.append("waveglow")
import threading
import io
import torch
import time
import numpy as np
from collections import OrderedDict

from flask import Flask, flash, request, redirect, url_for, send_file
import numpy as np

import librosa
import librosa.display

import torch

from scipy.io import wavfile
from pydub import AudioSegment

from hparams import create_hparams
from model import Tacotron2
from layers import TacotronSTFT, STFT
from audio_processing import griffin_lim
from train import load_model
from text import text_to_sequence
from denoiser import Denoiser

import audio

semaphore = threading.Semaphore()

hparams = create_hparams()
hparams.sampling_rate = 22050

checkpoint_path = "tacotron2_statedict.pt"
model = load_model(hparams)
model.load_state_dict(torch.load(checkpoint_path)["state_dict"])
_ = model.cuda().eval().half()

waveglow_path = "waveglow_256channels.pt"
waveglow = torch.load(waveglow_path)["model"]
waveglow.cuda().eval().half()
for k in waveglow.convinv:
    k.float()
denoiser = Denoiser(waveglow)

def infer(text):

    # NOTE: Outdated pytorch usage?

    sequence = np.array(text_to_sequence(text, ["english_cleaners"]))[None, :]
    sequence = torch.autograd.Variable(torch.from_numpy(sequence)).cuda().long()

    with torch.no_grad():
        mel_outputs, mel_outputs_postnet, _, alignments = model.inference(sequence)
        audio = waveglow.infer(mel_outputs_postnet, sigma=0.666)
        # NOTE: Do we want the denoiser?
        audio_denoised = denoiser(audio, strength=0.01)[:, 0]
        return audio_denoised.cpu().numpy()

def np_to_wav(arr):
    arr = arr[0] * (32767 / max(0.01, np.max(np.abs(arr))))
    as_int16 = arr.astype(np.int16)
    out = io.BytesIO()
    wavfile.write(out, hparams.sampling_rate, as_int16)
    return out

app = Flask(__name__)

@app.route("/", methods=["POST"])
def generate():

    args = request.json
    text = args["text"]
    print(f"Text: {text}")

    try:
        with semaphore:
            waveform = infer(text)

    except Exception as exc:
        print(f"Exception with text: {text}")
        print(exc)
        return "Bad Request", 400

    wav = np_to_wav(waveform)

    seg = AudioSegment(wav)
    seg = audio.de_esser(seg)
    wav = io.BytesIO()
    seg.export(wav, format="wav")
    wav.seek(0)

    return send_file(wav, mimetype="audio/wav")

def main():
    app.run(host="0.0.0.0", debug=True, port=6000)

if __name__ == "__main__":
    main()


import cv2
import numpy as np
import argparse

ap = argparse.ArgumentParser()

ap.add_argument("-p", "--prototxt", required=True,
                help="path to Caffe 'deploy' prototxt file")
ap.add_argument("-m", "--model", required=True, help="path to Caffe pre-treined model")
ap.add_argument("-c", "--confidence", type=float, default=0.5,
                help="minimum probability to filter weak detections")

args = vars(ap.parse_args())

# load our model from file
print("[INFO] loading model...")
net = cv2.dnn.readNetFromCaffe(args["prototxt"], args["model"])

video = cv2.VideoCapture('/dev/video0')

assert(video.isOpened() == True)
_, image = video.read()
(h,w) = image.shape[:2]

while True:
    tick = cv2.getTickCount()
    ok, image = video.read()
    assert(ok == True)
    # construct the input blob for the image
    # by resizing to a fixed 300x300 and normalizing it
    blob  = cv2.dnn.blobFromImage(cv2.resize(image, (300,300)),
                                  1.0,
                                  (300,300),
                                  (104.0, 177.0, 123.0))
    # print("[INFO] computing object detections...")
    net.setInput(blob)
    detections = net.forward()

    fps = cv2.getTickFrequency() / (cv2.getTickCount() - tick)

    # loop over the detections
    for i in range(detections.shape[2]):
        confidence = detections[0, 0, i,2]

        if confidence > args["confidence"]:
            box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
            (startX, startY, endX, endY) = box.astype("int")

            text = "{:.2f}%".format(confidence * 100)
            y = startY - 10 if startY - 10 > 10 else startY + 10
            cv2.rectangle(image, (startX, startY), (endX, endY), (255, 0, 255), 2)
            cv2.putText(image, text, (startX, y),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 0, 255), 2)
    cv2.putText(image, f'FPS: {int(fps)}', (20, 20),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 255), 2)
    cv2.imshow('Output', image)
    if cv2.waitKey(1) == (0xFF & ord('q')):
        break


import argparse

from rlzoo.distributed.dis_components import *
import tensorflow as tf
import numpy as np

def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument('-l', type=int, default=1)
    p.add_argument('-a', type=int, default=1)
    p.add_argument('-s', type=int, default=1)
    p.add_argument('-f', type=str, default='')  # config.json

    args = p.parse_args()
    return args

def run_learner(agent, args, training_conf, env_conf, agent_conf):
    agent_generator = agent_conf['agent_generator']
    total_step, traj_len, train_n_traj = training_conf['total_step'], training_conf['traj_len'], training_conf['train_n_traj'],
    obs_shape, act_shape = env_conf['obs_shape'], env_conf['act_shape']

    if agent.role_rank() == 0:
        param_q = agent.new_queue((Role.Learner, 0), (Role.Server, 0))

    traj_q = agent.new_queue((Role.Server, 0), (Role.Learner, agent.role_rank()))

    rl_agent = agent_generator()
    rl_agent.init_components()

    # init model
    rl_agent.update_model([agent.role_all_reduce(weights) for weights in rl_agent.all_weights])

    if agent.role_rank() == 0:
        for weight in rl_agent.all_weights:
            param_q.put(tf.Variable(weight, dtype=tf.float32))

    n_update = total_step // (traj_len * agent.role_size(Role.Learner) * train_n_traj)
    for i in range(n_update):
        traj_list = [[traj_q.get(dtype=tf.float32, shape=(traj_len, *shape)) for shape in [
            obs_shape, act_shape, (), (), obs_shape, (), (1,)]] for _ in range(train_n_traj)]

        rl_agent.train(traj_list, dis_agent=agent)

        # send weights to server
        if agent.role_rank() == 0:
            for weight in rl_agent.all_weights:
                param_q.put(tf.Variable(weight, dtype=tf.float32))
    print('learner finished')

def run_actor(agent, args, training_conf, env_conf):  # sampler
    env_maker, total_step = env_conf['env_maker'], training_conf['total_step']

    from gym import spaces

    env = env_maker()
    action_q, step_data_q = agent.new_queue_pair((Role.Server, 0), (Role.Actor, agent.role_rank()))

    state, reward, done = env.reset(), 0, False
    each_total_step = int(total_step/agent.role_size(Role.Actor))
    action_dtype = tf.int32 if isinstance(env.action_space, spaces.Discrete) else tf.float32
    for i in range(each_total_step):
        step_data_q.put(tf.Variable(state, dtype=tf.float32))
        a = action_q.get(dtype=action_dtype, shape=env.action_space.shape).numpy()
        next_state, reward, done, _ = env.step(a)
        for data in (reward, done, next_state):
            step_data_q.put(tf.Variable(data, dtype=tf.float32))
        if done:
            state = env.reset()
        else:
            state = next_state
    print('actor finished')

def run_server(agent, args, training_conf, env_conf, agent_conf):
    total_step, traj_len, train_n_traj, save_interval = training_conf['total_step'], training_conf['traj_len'], \
                                                        training_conf['train_n_traj'], training_conf['save_interval'],
    obs_shape, env_name = env_conf['obs_shape'], env_conf['env_name']
    agent_generator = agent_conf['agent_generator']

    from rlzoo.algorithms.dppo_clip_distributed.dppo_clip import DPPO_CLIP
    from rlzoo.distributed.dis_components import Role
    from gym import spaces

    learner_size = agent.role_size(Role.Learner)
    rl_agent: DPPO_CLIP = agent_generator()
    rl_agent.init_components()

    # queue to actor
    q_list = [agent.new_queue_pair((Role.Server, 0), (Role.Actor, i)) for i in
              range(agent.role_size(Role.Actor))]
    action_q_list, step_data_q_list = zip(*q_list)

    # queue to learner
    param_q = agent.new_queue((Role.Learner, 0), (Role.Server, 0))
    traj_q_list = [agent.new_queue((Role.Server, 0), (Role.Learner, i)) for i in
                   range(agent.role_size(Role.Learner))]

    # syn net weights from learner
    all_weights = [param_q.get(dtype=weight.dtype, shape=weight.shape) for weight in rl_agent.all_weights]
    rl_agent.update_model(all_weights)

    train_cnt = 0
    action_dtype = tf.int32 if isinstance(rl_agent.actor.action_space, spaces.Discrete) else tf.float32

    curr_step = 0

    total_reward_list = []
    curr_reward_list = []
    tmp_eps_reward = 0
    while curr_step < total_step:
        # tmp_eps_reward = 0  # todo env with no end
        for _ in range(traj_len):
            curr_step += agent.role_size(Role.Actor)

            state_list = []
            for step_data_q in step_data_q_list:
                state_list.append(step_data_q.get(dtype=tf.float32, shape=obs_shape))

            action_list, log_p_list = rl_agent.get_action(state_list, batch_data=True)

            for action_q, action in zip(action_q_list, action_list):
                action_q.put(tf.Variable(action, dtype=action_dtype))
            reward_list, done_list, next_state_list = [], [], [],
            for i, step_data_q in enumerate(step_data_q_list):
                reward = step_data_q.get(dtype=tf.float32, shape=())
                if i == 0:
                    tmp_eps_reward += reward
                reward_list.append(reward)
                done = step_data_q.get(dtype=tf.float32, shape=())
                if i == 0 and done:
                    curr_reward_list.append(tmp_eps_reward)
                    tmp_eps_reward = 0
                done_list.append(done)
                next_state_list.append(step_data_q.get(dtype=tf.float32, shape=obs_shape))
            rl_agent.collect_data(state_list, action_list, reward_list, done_list, next_state_list, log_p_list, True)

        rl_agent.update_traj_list()

        # send traj to each learner and update weight
        learn_traj_len = learner_size * train_n_traj
        if len(rl_agent.traj_list) >= learn_traj_len:
            train_cnt += 1

            # todo env with end
            avg_eps_reward = None
            if curr_reward_list:
                avg_eps_reward = np.mean(curr_reward_list)
                curr_reward_list.clear()
                total_reward_list.append(avg_eps_reward)

            # todo env with no end
            # avg_eps_reward = tmp_eps_reward
            # total_reward_list.append(np.array(avg_eps_reward))

            print('Training iters: {}, steps so far: {}, average eps reward: {}'.format(
                train_cnt, curr_step, np.array(avg_eps_reward)))

            rl_agent.plot_save_log(total_reward_list, env_name)

            traj_iter = iter(rl_agent.traj_list[:learn_traj_len])
            rl_agent.traj_list = rl_agent.traj_list[learn_traj_len:]

            # send traj data to each learner
            for i, traj_q in enumerate(traj_q_list):
                for _ in range(train_n_traj):
                    try:
                        traj_data = next(traj_iter)
                    except StopIteration:
                        break
                    for data in traj_data:
                        traj_q.put(tf.Variable(data, dtype=tf.float32))

            # syn net weights from learner
            all_weights = [param_q.get(dtype=weight.dtype, shape=weight.shape) for weight in rl_agent.all_weights]
            rl_agent.update_model(all_weights)

            # save model
            if not train_cnt % save_interval:
                rl_agent.save_ckpt(env_name)

    # save the final model
    rl_agent.save_ckpt(env_name)
    print('Server Finished.')

def main(training_conf, env_conf, agent_conf):
    args = parse_args()
    agent = Agent(n_learners=args.l, n_actors=args.a, n_servers=args.s)

    print('%s : %d/%d' % (agent.role(), agent.role_rank(), agent.role_size()))

    agent.barrier()

    if agent.role() == Role.Learner:
        run_learner(agent, args, training_conf, env_conf, agent_conf)
    elif agent.role() == Role.Actor:
        run_actor(agent, args, training_conf, env_conf)
    elif agent.role() == Role.Server:
        run_server(agent, args, training_conf, env_conf, agent_conf)
    else:
        raise RuntimeError('Invalid Role.')

    agent.barrier()


#!/usr/bin/env python
# Copyright 2021 Google LLC
#
# 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.

"""Predict runner."""

import io
import logging
import os
import numpy as np
import sys
import pickle
import haiku as hk
import time

from typing import Any, Mapping, MutableMapping, Optional, Sequence, Union, Tuple

from alphafold.model import config
from alphafold.model import model
from alphafold.model import utils
from alphafold.common import residue_constants
from alphafold.common import protein

def _get_model_haiku_params(model_name: str,
                            params_dir: str):
    """Get the Haiku parameters from a model name."""

    path = os.path.join(params_dir, f'params_{model_name}.npz')
    with open(path, 'rb') as f:
        params = np.load(io.BytesIO(f.read()), allow_pickle=False)

    return utils.flat_params_to_haiku(params)

def _load_features(features_path):
    with open(features_path, 'rb') as f:
        features = pickle.load(f)
    return features

def predict(
    model_features_path: str,
    model_params_path: str,
    model_name: str,
    num_ensemble: int,
    random_seed: int,
) -> Tuple[Mapping, Mapping]:

    model_config = config.model_config(model_name)

    # we assume  a monomer pipeline in a POC
    model_config.data.eval.num_ensemble = num_ensemble
    model_params = _get_model_haiku_params(
        model_name=model_name, params_dir=model_params_path)
    model_runner = model.RunModel(model_config, model_params)

    features = _load_features(model_features_path)
    processed_feature_dict = model_runner.process_features(
        raw_features=features,
        random_seed=random_seed
    )
    prediction_result = model_runner.predict(
        feat=processed_feature_dict,
        random_seed=random_seed
    )

    plddt = prediction_result['plddt']
    plddt_b_factors = np.repeat(
        plddt[:, None], residue_constants.atom_type_num, axis=-1)
    unrelaxed_structure= protein.from_prediction(
        features=processed_feature_dict,
        result=prediction_result,
        b_factors=plddt_b_factors,
        remove_leading_feature_dimension=not model_runner.multimer_mode)
    unrelaxed_pdbs = protein.to_pdb(unrelaxed_structure)

    return prediction_result, unrelaxed_pdbs

def _main(
    model_features_path: str,
    model_params_path: str,
    model_name: str,
    num_ensemble: int,
    random_seed: int,
    raw_prediction_path: str,
    unrelaxed_protein_path: str,
):
    logging.info(f'Running prediction using model {model_name}')
    prediction_result, unrelaxed_pdbs = predict(
        model_features_path=model_features_path,
        model_params_path=model_params_path,
        model_name=model_name,
        num_ensemble=num_ensemble,
        random_seed=random_seed
    )

    logging.info(f'Writing model {model_name} prediction to {raw_prediction_path}')
    with open(raw_prediction_path, 'wb') as f:
        pickle.dump(prediction_result, f, protocol=4)

    logging.info(f'Writing unrelaxed protein generated by {model_name} to {unrelaxed_protein_path}')
    with open(unrelaxed_protein_path, 'w') as f:
        f.write(unrelaxed_pdbs)

if __name__=='__main__':
    logging.basicConfig(format='%(asctime)s - %(message)s',
                        level=logging.INFO,
                        datefmt='%d-%m-%y %H:%M:%S',
                        stream=sys.stdout)

    random_seed = int(os.getenv('RANDOM_SEED', '0'))

    # TODO: Do something more intelligent with random seed
    if not random_seed:
        random_seed = int(os.environ['MODEL_NAME'][-1])

    _main(
        model_features_path=os.environ['FEATURES_PATH'],
        model_params_path=os.environ['MODEL_PARAMS_PATH'],
        model_name=os.environ['MODEL_NAME'],
        num_ensemble=int(os.getenv('NUM_ENSEMBE', '1')),
        random_seed=random_seed,
        raw_prediction_path=os.environ['RAW_PREDICTION_PATH'],
        unrelaxed_protein_path=os.environ["UNRELAXED_PROTEIN_PATH"]
    )


"""
Class for UCCA Embedding operations

  Authors: Jonathan Yellin
  Status: prototype

"""

import itertools
from collections import Counter
import numpy as np
import pickle
import json
from utils import constant

class UccaEmbedding(object):

    def __init__(self, embedding_dim, index_file, embedding_file):
        with open(index_file, 'rb') as index:
            self.encoding_vocab = pickle.load(index)
        self.embedding_matrix = np.load(embedding_file)

        self.__encoding_to_index = {encoding:i for i, encoding in enumerate(self.encoding_vocab)}

    def get_index(self, encoding):
        return self.__encoding_to_index[encoding]

    @staticmethod
    def prepare(embedding_dim, input_files, index_file, embedding_file, source):

        encodings = itertools.chain(*[UccaEmbedding.__load_encodings(input_file, source) for input_file in input_files])
        counter = Counter(t for t in encodings)
        encoding_vocab = sorted([t for t in counter], key=counter.get, reverse=True)
        vocab_size = len(encoding_vocab)
        emb = np.zeros((vocab_size, embedding_dim))

        for i, embedding in enumerate(encoding_vocab):
            array_of_bits = []

            for character in embedding:
                bit_encoding = [1 if digit == '1' else 0 for digit in bin(constant.UCCA_DEP_TO_ID[character])[2:]]
                bit_encoding = [0] * (4-len(bit_encoding)) + bit_encoding
                assert(len(bit_encoding) == 4)

                array_of_bits += bit_encoding

            emb[i][0:len(array_of_bits)] = array_of_bits

        print("dumping to files...")
        with open(index_file, 'wb') as outfile:
            pickle.dump(encoding_vocab, outfile)

        with open(embedding_file, 'wb') as outfile:
            np.save(outfile, emb)

        print("all done.")

    @staticmethod
    def __load_encodings(filename, source):
        all_encodings = []

        with open(filename) as infile:

            data = json.load(infile)
            for d in data:
                num_tokens = len(d['ucca_tokens'])
                encodings_dict = d['ucca_encodings_min_subtree'] if source == 'min_sub_tree' else d['ucca_encodings']
                index_to_encoding = {int(k): v for k, v in encodings_dict.items()} if encodings_dict is not None else {}
                encodings = []

                for index in range(0, num_tokens):
                    if index in index_to_encoding:
                        encodings.append(index_to_encoding[index])
                    else:
                        encodings.append('')

                if not encodings_dict is None:
                    all_encodings += encodings

        print("{} encodings from {} samples loaded from {}.".format(len(all_encodings), len(data), filename))
        return all_encodings