from collections import OrderedDict
from functools import partial

import numpy as np
import torch
from gym.envs.robotics import FetchEnv
from matplotlib.ticker import ScalarFormatter
from torch import optim
from torch import nn

import rlkit.samplers.rollout_functions as rf
import rlkit.torch.pytorch_util as ptu
from multiworld.envs.mujoco.cameras import (
    sawyer_door_env_camera_v0,
    sawyer_init_camera_zoomed_in,
    sawyer_pick_and_place_camera,
)
# from multiworld.envs.mujoco.classic_mujoco.hopper import \
    # HopperFullPositionGoalEnv
from multiworld.envs.mujoco.sawyer_xyz.sawyer_door_hook import SawyerDoorHookEnv
from multiworld.envs.mujoco.sawyer_xyz.sawyer_pick_and_place import \
    (
    SawyerPickAndPlaceEnv, SawyerPickAndPlaceEnvYZ,
)
from multiworld.envs.mujoco.sawyer_xyz.sawyer_push_nips import SawyerPushAndReachXYEnv
from multiworld.envs.mujoco.classic_mujoco.ant import (
    AntXYGoalEnv,
    AntFullPositionGoalEnv,
)
from multiworld.envs.pygame import PickAndPlaceEnv
from rlkit.core.logging import add_prefix
from rlkit.data_management.contextual_replay_buffer import (
    ContextualRelabelingReplayBuffer,
    RemapKeyFn,
)
from rlkit.envs.contextual import (
    ContextualEnv, ContextualRewardFn,
    delete_info,
)
from rlkit.envs.contextual.contextual_env import batchify
from rlkit.envs.contextual.goal_conditioned import (
    AddImageDistribution,
    GoalConditionedDiagnosticsToContextualDiagnostics,
    GoalDictDistributionFromMultitaskEnv,
    PresampledDistribution,
    ThresholdDistanceReward,
    L2Distance,
    IndexIntoAchievedGoal,
    NegativeL2Distance,
)
from rlkit.envs.contextual.gym_goal_envs import (
    GoalDictDistributionFromGymGoalEnv,
    GenericGoalConditionedContextualDiagnostics,
)
from rlkit.envs.images import EnvRenderer, GymEnvRenderer
from rlkit.launchers.contextual.util import (
    get_save_video_function,
    get_gym_env,
)
from rlkit.launchers.experiments.vitchyr.probabilistic_goal_reaching.diagnostics import (
    AntFullPositionGoalEnvDiagnostics,
    SawyerPickAndPlaceEnvAchievedFromObs,
    HopperFullPositionGoalEnvDiagnostics,
)
from rlkit.launchers.experiments.vitchyr.probabilistic_goal_reaching.env import \
    NormalizeAntFullPositionGoalEnv
from rlkit.launchers.experiments.vitchyr.probabilistic_goal_reaching.stochastic_env import \
    NoisyAction
from rlkit.launchers.experiments.vitchyr.probabilistic_goal_reaching.visualize import (
    DynamicsModelEnvRenderer, InsertDebugImagesEnv,
    DynamicNumberEnvRenderer,
    DiscountModelRenderer,
    ProductRenderer,
    ValueRenderer,
)
from rlkit.launchers.rl_exp_launcher_util import create_exploration_policy
from rlkit.misc import eval_util
from rlkit.misc.ml_util import create_schedule
from rlkit.samplers.data_collector.contextual_path_collector import (
    ContextualPathCollector
)
from rlkit.samplers.data_collector.joint_path_collector import \
    JointPathCollector
from rlkit.torch.distributions import (
    Distribution, MultivariateDiagonalNormal,
    IndependentLaplace,
)
from rlkit.torch.networks.basic import MultiInputSequential, Concat
from rlkit.torch.networks.mlp import MultiHeadedMlp, ConcatMlp, ParallelMlp
from rlkit.torch.networks.stochastic.distribution_generator import (
    TanhGaussian, DistributionGenerator,
    Gaussian,
    IndependentLaplaceGen,
)
from rlkit.torch.pgr.dynamics_model import EnsembleToGaussian
from rlkit.torch.pgr.pgr import PGRTrainer
from rlkit.torch.sac.policies import (
    MakeDeterministic,
    PolicyFromDistributionGenerator,
)
from rlkit.torch.supervised_learning.discount_model_trainer import (
    DiscountModelTrainer
)
from rlkit.torch.supervised_learning.dynamics_model_trainer import (
    GenerativeGoalDynamicsModelTrainer
)
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm, JointTrainer

class ProbabilisticGoalRewardFn(ContextualRewardFn):
    def __init__(
            self,
            dynamics_model: DistributionGenerator,
            state_key,
            context_key,
            reward_type='log_prob',
            discount_factor=0.99,
    ):
        if reward_type not in {'log_prob', 'prob'}:
            raise ValueError(reward_type)
        self._model = dynamics_model
        self._state_key = state_key
        self._context_key = context_key
        self._reward_type = reward_type
        self._discount_factor = discount_factor

    def __call__(self, states, actions, next_states, contexts):
        states_torch = ptu.from_numpy(states[self._state_key])
        actions_torch = ptu.from_numpy(actions)
        contexts_torch = ptu.from_numpy(contexts[self._context_key])
        dist = self._model(states_torch, actions_torch)
        log_prob_torch = dist.log_prob(contexts_torch)
        log_prob = ptu.get_numpy(log_prob_torch)

        if self._reward_type == 'log_prob':
            reward = log_prob
        elif self._reward_type == 'prob':
            reward = np.exp(log_prob)
        else:
            raise ValueError(self._reward_type)

        return reward

class IsotropicGaussianAroundState(DistributionGenerator):
    def __init__(self, state_to_goal_fn):
        super().__init__()
        self.state_to_goal_fn = state_to_goal_fn or (lambda x: x)

    def forward(self, states, actions) -> Distribution:
        goal = self.state_to_goal_fn(states)  # hack
        return MultivariateDiagonalNormal(goal, 1)

class StandardLaplaceAroundState(DistributionGenerator):
    def __init__(self, state_to_goal_fn):
        super().__init__()
        self.state_to_goal_fn = state_to_goal_fn or (lambda x: x)

    def forward(self, states, actions) -> Distribution:
        goal = self.state_to_goal_fn(states)  # hack
        return IndependentLaplace(goal, 1)

class DeltaModel(MultiInputSequential):
    def __init__(
            self,
            *args,
            state_to_goal_fn=None,
            log_std_max=2,
            log_std_min=-2,
            outputted_log_std_is_tanh=False,
            global_log_std_dim=None,
    ):
        super().__init__(*args)
        self.state_to_goal_fn = state_to_goal_fn or (lambda x: x)
        self.log_std_max = log_std_max
        self.log_std_min = log_std_min
        self.middle_log_std = (log_std_max + log_std_min) / 2
        self.half_range = (log_std_max - log_std_min) / 2
        self.outputted_log_std_is_tanh = outputted_log_std_is_tanh
        if global_log_std_dim is not None:
            global_log_std = torch.zeros(global_log_std_dim)
            self.global_log_std = nn.Parameter(global_log_std)
        else:
            self.global_log_std = None

    def forward(self, state, action):
        if self.global_log_std is None:
            mean, raw_log_std = super().forward(state, action)
            if self.outputted_log_std_is_tanh:
                log_std = raw_log_std * self.half_range + self.middle_log_std
            else:
                log_std = raw_log_std
        else:
            mean = super().forward(state, action)
            log_std = self.global_log_std
        return mean + self.state_to_goal_fn(state), log_std

class DeterministicEnsembleDeltaModel(MultiInputSequential):
    def __init__(
            self,
            *args,
            state_to_goal_fn=None,
    ):
        super().__init__(*args)
        self.state_to_goal_fn = state_to_goal_fn or (lambda x: x)

    def forward(self, state, action):
        delta = super().forward(state, action)
        current_subset = self.state_to_goal_fn(state).unsqueeze(-1)
        return delta + current_subset

def create_goal_dynamics_model(
        ob_space,
        action_space,
        goal_space,
        version,
        state_to_goal_fn,
        model_kwargs,
        delta_model_kwargs,
        ensemble_model_kwargs,
):
    if version == 'fixed_standard_gaussian':
        model = IsotropicGaussianAroundState(
            state_to_goal_fn=state_to_goal_fn,
        )
    elif version == 'fixed_standard_laplace':
        model = IsotropicGaussianAroundState(
            state_to_goal_fn=state_to_goal_fn,
        )
    elif version == 'learned_model':
        ob_dim = ob_space.high.size
        action_dim = action_space.high.size
        goal_dim = goal_space.high.size
        network = DeltaModel(
            Concat(),
            MultiHeadedMlp(
                input_size=ob_dim + action_dim,
                output_sizes=[goal_dim, goal_dim],
                **model_kwargs
            ),
            state_to_goal_fn=state_to_goal_fn,
            **delta_model_kwargs,
        )
        model = Gaussian(network)
    elif version == 'learned_model_laplace':
        ob_dim = ob_space.high.size
        action_dim = action_space.high.size
        goal_dim = goal_space.high.size
        network = DeltaModel(
            Concat(),
            MultiHeadedMlp(
                input_size=ob_dim + action_dim,
                output_sizes=[goal_dim, goal_dim],
                **model_kwargs
            ),
            state_to_goal_fn=state_to_goal_fn,
            **delta_model_kwargs,
        )
        model = IndependentLaplaceGen(network)
    elif version == 'learned_model_gaussian_global_variance':
        ob_dim = ob_space.high.size
        action_dim = action_space.high.size
        goal_dim = goal_space.high.size
        network = DeltaModel(
            Concat(),
            ConcatMlp(
                input_size=ob_dim + action_dim,
                output_size=goal_dim,
                **model_kwargs
            ),
            global_log_std_dim=goal_dim,
            state_to_goal_fn=state_to_goal_fn,
            **delta_model_kwargs,
        )
        model = Gaussian(network)
    elif version == 'learned_model_laplace_global_variance':
        ob_dim = ob_space.high.size
        action_dim = action_space.high.size
        goal_dim = goal_space.high.size
        network = DeltaModel(
            Concat(),
            ConcatMlp(
                input_size=ob_dim + action_dim,
                output_size=goal_dim,
                **model_kwargs
            ),
            global_log_std_dim=goal_dim,
            state_to_goal_fn=state_to_goal_fn,
            **delta_model_kwargs,
        )
        model = IndependentLaplaceGen(network)
    elif version == 'learned_model_ensemble':
        ob_dim = ob_space.high.size
        action_dim = action_space.high.size
        goal_dim = goal_space.high.size

        network = DeterministicEnsembleDeltaModel(
            Concat(),
            ParallelMlp(
                input_size=ob_dim + action_dim,
                output_size_per_mlp=goal_dim,
                **ensemble_model_kwargs
            ),
            state_to_goal_fn=state_to_goal_fn,
        )

        model = EnsembleToGaussian(network)
    else:
        raise ValueError(version)
    return model

class DiscountModel(ConcatMlp):
    def forward(self, *args, return_logits=False):
        logits = super().forward(*args)
        if return_logits:
            return logits
        else:
            return torch.sigmoid(logits)

def create_discount_model(ob_space, goal_space, action_space, model_kwargs):
    ob_dim = ob_space.high.size
    action_dim = action_space.high.size
    goal_dim = goal_space.high.size
    return DiscountModel(
        input_size=ob_dim + action_dim + goal_dim,
        output_size=1,
        **model_kwargs
    )

class StateToGoalFn(object):
    def __init__(self, stub_env):
        self.env_type = type(stub_env)

    def __call__(self, state):
        if issubclass(self.env_type, FetchEnv):
            return state[..., 3:6]
        elif self.env_type == AntXYGoalEnv:
            return state[..., :2]
        elif self.env_type == AntFullPositionGoalEnv:
            return state[..., :15]
        # elif self.env_type == HopperFullPositionGoalEnv:
        #     return state[..., :6]
        elif self.env_type == SawyerPickAndPlaceEnvYZ:
            return state[..., 1:]
        else:
            return state

def probabilistic_goal_reaching_experiment(
        max_path_length,
        qf_kwargs,
        policy_kwargs,
        pgr_trainer_kwargs,
        replay_buffer_kwargs,
        algo_kwargs,
        env_id,
        discount_factor,
        reward_type,
        # Dynamics model
        dynamics_model_version,
        dynamics_model_config,
        dynamics_delta_model_config=None,
        dynamics_adam_config=None,
        dynamics_ensemble_kwargs=None,
        # Discount model
        learn_discount_model=False,
        discount_adam_config=None,
        discount_model_config=None,
        prior_discount_weight_schedule_kwargs=None,
        # Environment
        env_class=None,
        env_kwargs=None,
        observation_key='state_observation',
        desired_goal_key='state_desired_goal',
        exploration_policy_kwargs=None,
        action_noise_scale=0.,
        num_presampled_goals=4096,
        success_threshold=0.05,
        # Video / visualization parameters
        save_video=True,
        save_video_kwargs=None,
        video_renderer_kwargs=None,
        plot_renderer_kwargs=None,
        eval_env_ids=None,
        # Debugging params
        visualize_dynamics=False,
        visualize_discount_model=False,
        visualize_all_plots=False,
        plot_discount=False,
        plot_reward=False,
        plot_bootstrap_value=False,
        # env specific-params
        normalize_distances_for_full_state_ant=False,
):
    if dynamics_ensemble_kwargs is None:
        dynamics_ensemble_kwargs = {}
    if eval_env_ids is None:
        eval_env_ids = {'eval': env_id}
    if discount_model_config is None:
        discount_model_config = {}
    if dynamics_delta_model_config is None:
        dynamics_delta_model_config = {}
    if dynamics_adam_config is None:
        dynamics_adam_config = {}
    if discount_adam_config is None:
        discount_adam_config = {}
    if exploration_policy_kwargs is None:
        exploration_policy_kwargs = {}
    if not save_video_kwargs:
        save_video_kwargs = {}
    if not video_renderer_kwargs:
        video_renderer_kwargs = {}
    if not plot_renderer_kwargs:
        plot_renderer_kwargs = video_renderer_kwargs.copy()
        plot_renderer_kwargs['dpi'] = 48
    context_key = desired_goal_key

    stub_env = get_gym_env(
        env_id, env_class=env_class, env_kwargs=env_kwargs,
        unwrap_timed_envs=True,
    )
    is_gym_env = (
        isinstance(stub_env, FetchEnv)
        or isinstance(stub_env, AntXYGoalEnv)
        or isinstance(stub_env, AntFullPositionGoalEnv)
        # or isinstance(stub_env, HopperFullPositionGoalEnv)
    )
    is_ant_full_pos = isinstance(stub_env, AntFullPositionGoalEnv)

    if is_gym_env:
        achieved_goal_key = desired_goal_key.replace('desired', 'achieved')
        ob_keys_to_save_in_buffer = [observation_key, achieved_goal_key]
    elif isinstance(stub_env, SawyerPickAndPlaceEnvYZ):
        achieved_goal_key = desired_goal_key.replace('desired', 'achieved')
        ob_keys_to_save_in_buffer = [observation_key, achieved_goal_key]
    else:
        achieved_goal_key = observation_key
        ob_keys_to_save_in_buffer = [observation_key]
    # TODO move all env-specific code to other file
    if isinstance(stub_env, SawyerDoorHookEnv):
        init_camera = sawyer_door_env_camera_v0
    elif isinstance(stub_env, SawyerPushAndReachXYEnv):
        init_camera = sawyer_init_camera_zoomed_in
    elif isinstance(stub_env, SawyerPickAndPlaceEnvYZ):
        init_camera = sawyer_pick_and_place_camera
    else:
        init_camera = None

    full_ob_space = stub_env.observation_space
    action_space = stub_env.action_space
    state_to_goal = StateToGoalFn(stub_env)
    dynamics_model = create_goal_dynamics_model(
        full_ob_space[observation_key],
        action_space,
        full_ob_space[achieved_goal_key],
        dynamics_model_version,
        state_to_goal,
        dynamics_model_config,
        dynamics_delta_model_config,
        ensemble_model_kwargs=dynamics_ensemble_kwargs,
    )
    sample_context_from_obs_dict_fn = RemapKeyFn({context_key: achieved_goal_key})

    def contextual_env_distrib_reward(
            _env_id, _env_class=None, _env_kwargs=None
    ):
        base_env = get_gym_env(
            _env_id, env_class=env_class, env_kwargs=env_kwargs,
            unwrap_timed_envs=True,
        )
        if init_camera:
            base_env.initialize_camera(init_camera)
        if (isinstance(stub_env, AntFullPositionGoalEnv)
                and normalize_distances_for_full_state_ant):
            base_env = NormalizeAntFullPositionGoalEnv(base_env)
            normalize_env = base_env
        else:
            normalize_env = None
        env = NoisyAction(base_env, action_noise_scale)
        diag_fns = []
        if is_gym_env:
            goal_distribution = GoalDictDistributionFromGymGoalEnv(
                env,
                desired_goal_key=desired_goal_key,
            )
            diag_fns.append(GenericGoalConditionedContextualDiagnostics(
                desired_goal_key=desired_goal_key,
                achieved_goal_key=achieved_goal_key,
                success_threshold=success_threshold,
            ))
        else:
            goal_distribution = GoalDictDistributionFromMultitaskEnv(
                env,
                desired_goal_keys=[desired_goal_key],
            )
            diag_fns.append(
                GoalConditionedDiagnosticsToContextualDiagnostics(
                    env.goal_conditioned_diagnostics,
                    desired_goal_key=desired_goal_key,
                    observation_key=observation_key,
                )
            )
        if isinstance(stub_env, AntFullPositionGoalEnv):
            diag_fns.append(
                AntFullPositionGoalEnvDiagnostics(
                    desired_goal_key=desired_goal_key,
                    achieved_goal_key=achieved_goal_key,
                    success_threshold=success_threshold,
                    normalize_env=normalize_env,
                )
            )
        # if isinstance(stub_env, HopperFullPositionGoalEnv):
        #     diag_fns.append(
        #         HopperFullPositionGoalEnvDiagnostics(
        #             desired_goal_key=desired_goal_key,
        #             achieved_goal_key=achieved_goal_key,
        #             success_threshold=success_threshold,
        #         )
        #     )
        achieved_from_ob = IndexIntoAchievedGoal(
            achieved_goal_key,
        )
        if reward_type == 'sparse':
            distance_fn = L2Distance(
                achieved_goal_from_observation=achieved_from_ob,
                desired_goal_key=desired_goal_key,
            )
            reward_fn = ThresholdDistanceReward(distance_fn, success_threshold)
        elif reward_type == 'negative_distance':
            reward_fn = NegativeL2Distance(
                achieved_goal_from_observation=achieved_from_ob,
                desired_goal_key=desired_goal_key,
            )
        else:
            reward_fn = ProbabilisticGoalRewardFn(
                dynamics_model,
                state_key=observation_key,
                context_key=context_key,
                reward_type=reward_type,
                discount_factor=discount_factor,
            )
        goal_distribution = PresampledDistribution(
            goal_distribution, num_presampled_goals)
        final_env = ContextualEnv(
            env,
            context_distribution=goal_distribution,
            reward_fn=reward_fn,
            observation_key=observation_key,
            contextual_diagnostics_fns=diag_fns,
            update_env_info_fn=delete_info,
        )
        return final_env, goal_distribution, reward_fn

    expl_env, expl_context_distrib, reward_fn = contextual_env_distrib_reward(
        env_id, env_class, env_kwargs,
    )
    obs_dim = (
            expl_env.observation_space.spaces[observation_key].low.size
            + expl_env.observation_space.spaces[context_key].low.size
    )
    action_dim = expl_env.action_space.low.size

    def create_qf():
        return ConcatMlp(
            input_size=obs_dim + action_dim,
            output_size=1,
            **qf_kwargs
        )
    qf1 = create_qf()
    qf2 = create_qf()
    target_qf1 = create_qf()
    target_qf2 = create_qf()

    def create_policy():
        obs_processor = MultiHeadedMlp(
            input_size=obs_dim,
            output_sizes=[action_dim, action_dim],
            **policy_kwargs
        )
        return PolicyFromDistributionGenerator(
            TanhGaussian(obs_processor)
        )
    policy = create_policy()

    def concat_context_to_obs(batch, replay_buffer, obs_dict, next_obs_dict, new_contexts):
        obs = batch['observations']
        next_obs = batch['next_observations']
        batch['original_observations'] = obs
        batch['original_next_observations'] = next_obs
        context = batch[context_key]
        batch['observations'] = np.concatenate([obs, context], axis=1)
        batch['next_observations'] = np.concatenate([next_obs, context], axis=1)
        return batch

    replay_buffer = ContextualRelabelingReplayBuffer(
        env=expl_env,
        context_keys=[context_key],
        observation_keys=ob_keys_to_save_in_buffer,
        context_distribution=expl_context_distrib,
        sample_context_from_obs_dict_fn=sample_context_from_obs_dict_fn,
        reward_fn=reward_fn,
        post_process_batch_fn=concat_context_to_obs,
        **replay_buffer_kwargs
    )

    def create_trainer():
        trainers = OrderedDict()
        if learn_discount_model:
            discount_model = create_discount_model(
                ob_space=stub_env.observation_space[observation_key],
                goal_space=stub_env.observation_space[context_key],
                action_space=stub_env.action_space,
                model_kwargs=discount_model_config)
            optimizer = optim.Adam(
                discount_model.parameters(),
                **discount_adam_config
            )
            discount_trainer = DiscountModelTrainer(
                discount_model,
                optimizer,
                observation_key='observations',
                next_observation_key='original_next_observations',
                goal_key=context_key,
                state_to_goal_fn=state_to_goal,
            )
            trainers['discount_trainer'] = discount_trainer
        else:
            discount_model = None
        if prior_discount_weight_schedule_kwargs is not None:
            schedule = create_schedule(
                **prior_discount_weight_schedule_kwargs
            )
        else:
            schedule = None
        pgr_trainer = PGRTrainer(
            env=expl_env,
            policy=policy,
            qf1=qf1,
            qf2=qf2,
            target_qf1=target_qf1,
            target_qf2=target_qf2,
            discount=discount_factor,
            discount_model=discount_model,
            prior_discount_weight_schedule=schedule,
            **pgr_trainer_kwargs
        )
        trainers[''] = pgr_trainer
        optimizers = [
            pgr_trainer.qf1_optimizer,
            pgr_trainer.qf2_optimizer,
            pgr_trainer.alpha_optimizer,
            pgr_trainer.policy_optimizer,
        ]
        if dynamics_model_version in {
            'learned_model',
            'learned_model_ensemble',
            'learned_model_laplace',
            'learned_model_laplace_global_variance',
            'learned_model_gaussian_global_variance',
        }:
            model_opt = optim.Adam(
                dynamics_model.parameters(),
                **dynamics_adam_config
            )
        elif dynamics_model_version in {
            'fixed_standard_laplace',
            'fixed_standard_gaussian',
        }:
            model_opt = None
        else:
            raise NotImplementedError()
        model_trainer = GenerativeGoalDynamicsModelTrainer(
            dynamics_model,
            model_opt,
            state_to_goal=state_to_goal,
            observation_key='original_observations',
            next_observation_key='original_next_observations',
        )
        trainers['dynamics_trainer'] = model_trainer
        optimizers.append(model_opt)
        return JointTrainer(trainers), pgr_trainer

    trainer, pgr_trainer = create_trainer()

    eval_policy = MakeDeterministic(policy)

    def create_eval_path_collector(some_eval_env):
        return ContextualPathCollector(
            some_eval_env,
            eval_policy,
            observation_key=observation_key,
            context_keys_for_policy=[context_key],
        )
    path_collectors = dict()
    eval_env_name_to_env_and_context_distrib = dict()
    for name, extra_env_id in eval_env_ids.items():
        env, context_distrib, _ = contextual_env_distrib_reward(extra_env_id)
        path_collectors[name] = create_eval_path_collector(env)
        eval_env_name_to_env_and_context_distrib[name] = (
            env, context_distrib
        )
    eval_path_collector = JointPathCollector(path_collectors)
    exploration_policy = create_exploration_policy(
        expl_env, policy, **exploration_policy_kwargs)
    expl_path_collector = ContextualPathCollector(
        expl_env,
        exploration_policy,
        observation_key=observation_key,
        context_keys_for_policy=[context_key],
    )

    def get_eval_diagnostics(key_to_paths):
        stats = OrderedDict()
        for eval_env_name, paths in key_to_paths.items():
            env, _ = eval_env_name_to_env_and_context_distrib[eval_env_name]
            stats.update(add_prefix(
                    env.get_diagnostics(paths),
                    eval_env_name,
                    divider='/',
            ))
            stats.update(add_prefix(
                eval_util.get_generic_path_information(paths),
                eval_env_name,
                divider='/',
            ))
        return stats

    algorithm = TorchBatchRLAlgorithm(
        trainer=trainer,
        exploration_env=expl_env,
        evaluation_env=None,
        exploration_data_collector=expl_path_collector,
        evaluation_data_collector=eval_path_collector,
        replay_buffer=replay_buffer,
        max_path_length=max_path_length,
        evaluation_get_diagnostic_functions=[get_eval_diagnostics],
        **algo_kwargs
    )
    algorithm.to(ptu.device)

    if normalize_distances_for_full_state_ant and is_ant_full_pos:
        qpos_weights = expl_env.unwrapped.presampled_qpos.std(axis=0)
    else:
        qpos_weights = None

    if save_video:
        if is_gym_env:
            video_renderer = GymEnvRenderer(**video_renderer_kwargs)

            def set_goal_for_visualization(env, policy, o):
                goal = o[desired_goal_key]
                if normalize_distances_for_full_state_ant and is_ant_full_pos:
                    unnormalized_goal = goal * qpos_weights
                    env.unwrapped.goal = unnormalized_goal
                else:
                    env.unwrapped.goal = goal

            rollout_function = partial(
                rf.contextual_rollout,
                max_path_length=max_path_length,
                observation_key=observation_key,
                context_keys_for_policy=[context_key],
                reset_callback=set_goal_for_visualization,
            )
        else:
            video_renderer = EnvRenderer(**video_renderer_kwargs)
            rollout_function = partial(
                rf.contextual_rollout,
                max_path_length=max_path_length,
                observation_key=observation_key,
                context_keys_for_policy=[context_key],
                reset_callback=None,
            )

        renderers = OrderedDict(
            image_observation=video_renderer,
        )
        state_env = expl_env.env
        state_space = state_env.observation_space[observation_key]
        low = state_space.low.min()
        high = state_space.high.max()
        y = np.linspace(low, high, num=video_renderer.image_chw[1])
        x = np.linspace(low, high, num=video_renderer.image_chw[2])
        all_xy_np = np.transpose([np.tile(x, len(y)), np.repeat(y, len(x))])
        all_xy_torch = ptu.from_numpy(all_xy_np)
        num_states = all_xy_torch.shape[0]
        if visualize_dynamics:
            def create_dynamics_visualizer(show_prob, vary_state=False):
                def get_prob(obs_dict, action):
                    obs = obs_dict['state_observation']
                    obs_torch = ptu.from_numpy(obs)[None]
                    action_torch = ptu.from_numpy(action)[None]
                    if vary_state:
                        action_repeated = torch.zeros((num_states, 2))
                        dist = dynamics_model(all_xy_torch, action_repeated)
                        goal = ptu.from_numpy(
                            obs_dict['state_desired_goal'][None])
                        log_probs = dist.log_prob(goal)
                    else:
                        dist = dynamics_model(obs_torch, action_torch)
                        log_probs = dist.log_prob(all_xy_torch)
                    if show_prob:
                        return log_probs.exp()
                    else:
                        return log_probs
                return get_prob

            renderers['log_prob'] = ValueRenderer(
                create_dynamics_visualizer(False), **video_renderer_kwargs
            )
            # renderers['prob'] = ValueRenderer(
            #     create_dynamics_visualizer(True), **video_renderer_kwargs
            # )
            renderers['log_prob_vary_state'] = ValueRenderer(
                create_dynamics_visualizer(False, vary_state=True),
                only_get_image_once_per_episode=True,
                max_out_walls=isinstance(stub_env, PickAndPlaceEnv),
                **video_renderer_kwargs)
            # renderers['prob_vary_state'] = ValueRenderer(
            #     create_dynamics_visualizer(True, vary_state=True),
            #     **video_renderer_kwargs)
        if visualize_discount_model and pgr_trainer.discount_model:
            def get_discount_values(obs, action):
                obs = obs['state_observation']
                obs_torch = ptu.from_numpy(obs)[None]
                combined_obs = torch.cat([
                    obs_torch.repeat(num_states, 1),
                    all_xy_torch,
                ], dim=1)

                action_torch = ptu.from_numpy(action)[None]
                action_repeated = action_torch.repeat(num_states, 1)
                return pgr_trainer.discount_model(combined_obs, action_repeated)
            renderers['discount_model'] = ValueRenderer(
                get_discount_values,
                states_to_eval=all_xy_torch,
                **video_renderer_kwargs)
        if 'log_prob' in renderers and 'discount_model' in renderers:
            renderers['log_prob_time_discount'] = ProductRenderer(
                renderers['discount_model'],
                renderers['log_prob'],
                **video_renderer_kwargs)

        def get_reward(obs_dict, action, next_obs_dict):
            o = batchify(obs_dict)
            a = batchify(action)
            next_o = batchify(next_obs_dict)
            reward = reward_fn(o, a, next_o, next_o)
            return reward[0]

        def get_bootstrap(obs_dict, action, next_obs_dict, return_float=True):
            context_pt = ptu.from_numpy(obs_dict[context_key][None])
            o_pt = ptu.from_numpy(obs_dict[observation_key][None])
            next_o_pt = ptu.from_numpy(next_obs_dict[observation_key][None])
            action_torch = ptu.from_numpy(action[None])
            bootstrap, *_ = pgr_trainer.get_bootstrap_stats(
                torch.cat((o_pt, context_pt), dim=1),
                action_torch,
                torch.cat((next_o_pt, context_pt), dim=1),
            )
            if return_float:
                return ptu.get_numpy(bootstrap)[0, 0]
            else:
                return bootstrap

        def get_discount(obs_dict, action, next_obs_dict):
            bootstrap = get_bootstrap(obs_dict, action, next_obs_dict, return_float=False)
            reward_np = get_reward(obs_dict, action, next_obs_dict)
            reward = ptu.from_numpy(reward_np[None, None])
            context_pt = ptu.from_numpy(obs_dict[context_key][None])
            o_pt = ptu.from_numpy(obs_dict[observation_key][None])
            obs = torch.cat((o_pt, context_pt), dim=1)
            actions = ptu.from_numpy(action[None])
            discount = pgr_trainer.get_discount_factor(
                bootstrap,
                reward,
                obs,
                actions,
            )
            if isinstance(discount, torch.Tensor):
                discount = ptu.get_numpy(discount)[0, 0]
            return np.clip(discount, a_min=1e-3, a_max=1)

        def create_modify_fn(title, set_params=None, scientific=True,):
            def modify(ax):
                ax.set_title(title)
                if set_params:
                    ax.set(**set_params)
                if scientific:
                    scaler = ScalarFormatter(useOffset=True)
                    scaler.set_powerlimits((1, 1))
                    ax.yaxis.set_major_formatter(scaler)
                    ax.ticklabel_format(axis='y', style='sci')
            return modify

        def add_left_margin(fig):
            fig.subplots_adjust(left=0.2)
        if visualize_all_plots or plot_discount:
            renderers['discount'] = DynamicNumberEnvRenderer(
                dynamic_number_fn=get_discount,
                modify_ax_fn=create_modify_fn(
                    title='discount',
                    set_params=dict(
                        # yscale='log',
                        ylim=[-0.05, 1.1],
                    ),
                    # scientific=False,
                ),
                modify_fig_fn=add_left_margin,
                # autoscale_y=False,
                **plot_renderer_kwargs)

        if visualize_all_plots or plot_reward:
            renderers['reward'] = DynamicNumberEnvRenderer(
                dynamic_number_fn=get_reward,
                modify_ax_fn=create_modify_fn(
                    title='reward',
                    # scientific=False,
                ),
                modify_fig_fn=add_left_margin,
                **plot_renderer_kwargs)
        if visualize_all_plots or plot_bootstrap_value:
            renderers['bootstrap-value'] = DynamicNumberEnvRenderer(
                dynamic_number_fn=get_bootstrap,
                modify_ax_fn=create_modify_fn(
                    title='bootstrap value',
                    # scientific=False,
                ),
                modify_fig_fn=add_left_margin,
                **plot_renderer_kwargs)

        def add_images(env, state_distribution):
            state_env = env.env
            if is_gym_env:
                goal_distribution = state_distribution
            else:
                goal_distribution = AddImageDistribution(
                    env=state_env,
                    base_distribution=state_distribution,
                    image_goal_key='image_desired_goal',
                    renderer=video_renderer,
                )
            context_env = ContextualEnv(
                state_env,
                context_distribution=goal_distribution,
                reward_fn=reward_fn,
                observation_key=observation_key,
                update_env_info_fn=delete_info,
            )
            return InsertDebugImagesEnv(
                context_env,
                renderers=renderers,
            )
        img_expl_env = add_images(expl_env, expl_context_distrib)
        if is_gym_env:
            imgs_to_show = list(renderers.keys())
        else:
            imgs_to_show = ['image_desired_goal'] + list(renderers.keys())
        img_formats = [video_renderer.output_image_format]
        img_formats += [r.output_image_format for r in renderers.values()]
        expl_video_func = get_save_video_function(
            rollout_function,
            img_expl_env,
            exploration_policy,
            tag="xplor",
            imsize=video_renderer.image_chw[1],
            image_formats=img_formats,
            keys_to_show=imgs_to_show,
            **save_video_kwargs
        )
        algorithm.post_train_funcs.append(expl_video_func)
        for eval_env_name, (env, context_distrib) in (
                eval_env_name_to_env_and_context_distrib.items()
        ):
            img_eval_env = add_images(env, context_distrib)
            eval_video_func = get_save_video_function(
                rollout_function,
                img_eval_env,
                eval_policy,
                tag=eval_env_name,
                imsize=video_renderer.image_chw[1],
                image_formats=img_formats,
                keys_to_show=imgs_to_show,
                **save_video_kwargs
            )
            algorithm.post_train_funcs.append(eval_video_func)

    algorithm.train()


#!/usr/bin/env python
# coding: utf-8

# # Word Embeddings First Steps: Data Preparation
#
# In this series of ungraded notebooks, you'll try out all the individual techniques that you learned about in the lectures. Practicing on small examples will prepare you for the graded assignment, where you will combine the techniques in more advanced ways to create word embeddings from a real-life corpus.
#
# This notebook focuses on data preparation, which is the first step of any machine learning algorithm. It is a very important step because models are only as good as the data they are trained on and the models used require the data to have a particular structure to process it properly.
#
# To get started, import and initialize all the libraries you will need.

# In[1]:

import re
import nltk

nltk.download('punkt')

import emoji
import numpy as np
from nltk.tokenize import word_tokenize
from utils2 import get_dict

# # Data preparation

# In the data preparation phase, starting with a corpus of text, you will:
#
# - Clean and tokenize the corpus.
#
# - Extract the pairs of context words and center word that will make up the training data set for the CBOW model. The context words are the features that will be fed into the model, and the center words are the target values that the model will learn to predict.
#
# - Create simple vector representations of the context words (features) and center words (targets) that can be used by the neural network of the CBOW model.

# ## Cleaning and tokenization
#
# To demonstrate the cleaning and tokenization process, consider a corpus that contains emojis and various punctuation signs.

# In[2]:

# Define a corpus
corpus = 'Who ❤ "word embeddings" in 2020? I do!!!'

# First, replace all interrupting punctuation signs — such as commas and exclamation marks — with periods.

# In[3]:

# Print original corpus
print(f'Corpus:  {corpus}')

# Do the substitution
data = re.sub(r'[,!?;-]+', '.', corpus)

# Print cleaned corpus
print(f'After cleaning punctuation:  {data}')

# Next, use NLTK's tokenization engine to split the corpus into individual tokens.

# In[4]:

# Print cleaned corpus
print(f'Initial string:  {data}')

# Tokenize the cleaned corpus
data = nltk.word_tokenize(data)

# Print the tokenized version of the corpus
print(f'After tokenization:  {data}')

# Finally, as you saw in the lecture, get rid of numbers and punctuation other than periods, and convert all the remaining tokens to lowercase.

# In[5]:

# Print the tokenized version of the corpus
print(f'Initial list of tokens:  {data}')

# Filter tokenized corpus using list comprehension
data = [ ch.lower() for ch in data
         if ch.isalpha()
         or ch == '.'
         or emoji.get_emoji_regexp().search(ch)
       ]

# Print the tokenized and filtered version of the corpus
print(f'After cleaning:  {data}')

# Note that the heart emoji is considered as a token just like any normal word.
#
# Now let's streamline the cleaning and tokenization process by wrapping the previous steps in a function.

# In[6]:

# Define the 'tokenize' function that will include the steps previously seen
def tokenize(corpus):
    data = re.sub(r'[,!?;-]+', '.', corpus)
    data = nltk.word_tokenize(data)  # tokenize string to words
    data = [ ch.lower() for ch in data
             if ch.isalpha()
             or ch == '.'
             or emoji.get_emoji_regexp().search(ch)
           ]
    return data

# Apply this function to the corpus that you'll be working on in the rest of this notebook: "I am happy because I am learning"

# In[7]:

# Define new corpus
corpus = 'I am happy because I am learning'

# Print new corpus
print(f'Corpus:  {corpus}')

# Save tokenized version of corpus into 'words' variable
words = tokenize(corpus)

# Print the tokenized version of the corpus
print(f'Words (tokens):  {words}')

# **Now try it out yourself with your own sentence.**

# In[8]:

# Run this with any sentence
tokenize("Now it's your turn: try with your own sentence!")

# ## Sliding window of words

# Now that you have transformed the corpus into a list of clean tokens, you can slide a window of words across this list. For each window you can extract a center word and the context words.
#
# The `get_windows` function in the next cell was introduced in the lecture.

# In[9]:

# Define the 'get_windows' function
def get_windows(words, C):
    i = C
    while i < len(words) - C:
        center_word = words[i]
        context_words = words[(i - C):i] + words[(i+1):(i+C+1)]
        yield context_words, center_word
        i += 1

# The first argument of this function is a list of words (or tokens). The second argument, `C`, is the context half-size. Recall that for a given center word, the context words are made of `C` words to the left and `C` words to the right of the center word.
#
# Here is how you can use this function to extract context words and center words from a list of tokens. These context and center words will make up the training set that you will use to train the CBOW model.

# In[10]:

# Print 'context_words' and 'center_word' for the new corpus with a 'context half-size' of 2
for x, y in get_windows(['i', 'am', 'happy', 'because', 'i', 'am', 'learning'], 2):
    print(f'{x}\t{y}')

# The first example of the training set is made of:
#
# - the context words "i", "am", "because", "i",
#
# - and the center word to be predicted: "happy".
#
# **Now try it out yourself. In the next cell, you can change both the sentence and the context half-size.**

# In[11]:

# Print 'context_words' and 'center_word' for any sentence with a 'context half-size' of 1
for x, y in get_windows(tokenize("Now it's your turn: try with your own sentence!"), 1):
    print(f'{x}\t{y}')

# ## Transforming words into vectors for the training set

# To finish preparing the training set, you need to transform the context words and center words into vectors.
#
# ### Mapping words to indices and indices to words
#
# The center words will be represented as one-hot vectors, and the vectors that represent context words are also based on one-hot vectors.
#
# To create one-hot word vectors, you can start by mapping each unique word to a unique integer (or index). We have provided a helper function, `get_dict`, that creates a Python dictionary that maps words to integers and back.

# In[12]:

# Get 'word2Ind' and 'Ind2word' dictionaries for the tokenized corpus
word2Ind, Ind2word = get_dict(words)

# Here's the dictionary that maps words to numeric indices.

# In[13]:

# Print 'word2Ind' dictionary
word2Ind

# You can use this dictionary to get the index of a word.

# In[14]:

# Print value for the key 'i' within word2Ind dictionary
print("Index of the word 'i':  ",word2Ind['i'])

# And conversely, here's the dictionary that maps indices to words.

# In[15]:

# Print 'Ind2word' dictionary
Ind2word

# In[16]:

# Print value for the key '2' within Ind2word dictionary
print("Word which has index 2:  ",Ind2word[2] )

# Finally, get the length of either of these dictionaries to get the size of the vocabulary of your corpus, in other words the number of different words making up the corpus.

# In[17]:

# Save length of word2Ind dictionary into the 'V' variable
V = len(word2Ind)

# Print length of word2Ind dictionary
print("Size of vocabulary: ", V)

# ### Getting one-hot word vectors
#
# Recall from the lecture that you can easily convert an integer, $n$, into a one-hot vector.
#
# Consider the word "happy". First, retrieve its numeric index.

# In[18]:

# Save index of word 'happy' into the 'n' variable
n = word2Ind['happy']

# Print index of word 'happy'
n

# Now create a vector with the size of the vocabulary, and fill it with zeros.

# In[19]:

# Create vector with the same length as the vocabulary, filled with zeros
center_word_vector = np.zeros(V)

# Print vector
center_word_vector

# You can confirm that the vector has the right size.

# In[20]:

# Assert that the length of the vector is the same as the size of the vocabulary
len(center_word_vector) == V

# Next, replace the 0 of the $n$-th element with a 1.

# In[22]:

# Replace element number 'n' with a 1
center_word_vector[n] = 1

# And you have your one-hot word vector.

# In[23]:

# Print vector
center_word_vector

# **You can now group all of these steps in a convenient function, which takes as parameters: a word to be encoded, a dictionary that maps words to indices, and the size of the vocabulary.**

# In[24]:

# Define the 'word_to_one_hot_vector' function that will include the steps previously seen
def word_to_one_hot_vector(word, word2Ind, V):
    one_hot_vector = np.zeros(V)
    one_hot_vector[word2Ind[word]] = 1
    return one_hot_vector

# Check that it works as intended.

# In[25]:

# Print output of 'word_to_one_hot_vector' function for word 'happy'
word_to_one_hot_vector('happy', word2Ind, V)

# **What is the word vector for "learning"?**

# In[26]:

# Print output of 'word_to_one_hot_vector' function for word 'learning'
word_to_one_hot_vector('learning', word2Ind, V)

# Expected output:
#
#     array([0., 0., 0., 0., 1.])

# ### Getting context word vectors

# To create the vectors that represent context words, you will calculate the average of the one-hot vectors representing the individual words.
#
# Let's start with a list of context words.

# In[27]:

# Define list containing context words
context_words = ['i', 'am', 'because', 'i']

# Using Python's list comprehension construct and the `word_to_one_hot_vector` function that you created in the previous section, you can create a list of one-hot vectors representing each of the context words.

# In[28]:

# Create one-hot vectors for each context word using list comprehension
context_words_vectors = [word_to_one_hot_vector(w, word2Ind, V) for w in context_words]

# Print one-hot vectors for each context word
context_words_vectors

# And you can now simply get the average of these vectors using numpy's `mean` function, to get the vector representation of the context words.

# In[29]:

# Compute mean of the vectors using numpy
np.mean(context_words_vectors, axis=0)

# Note the `axis=0` parameter that tells `mean` to calculate the average of the rows (if you had wanted the average of the columns, you would have used `axis=1`).
#
# **Now create the `context_words_to_vector` function that takes in a list of context words, a word-to-index dictionary, and a vocabulary size, and outputs the vector representation of the context words.**

# In[30]:

# Define the 'context_words_to_vector' function that will include the steps previously seen
def context_words_to_vector(context_words, word2Ind, V):
    context_words_vectors = [word_to_one_hot_vector(w, word2Ind, V) for w in context_words]
    context_words_vectors = np.mean(context_words_vectors, axis=0)
    return context_words_vectors

# And check that you obtain the same output as the manual approach above.

# In[31]:

# Print output of 'context_words_to_vector' function for context words: 'i', 'am', 'because', 'i'
context_words_to_vector(['i', 'am', 'because', 'i'], word2Ind, V)

# **What is the vector representation of the context words "am happy i am"?**

# In[ ]:

# Print output of 'context_words_to_vector' function for context words: 'am', 'happy', 'i', 'am'
context_words_to_vector(['am', 'happy', 'i', 'am'], word2Ind, V)

# Expected output:
#
#     array([0.5 , 0.  , 0.25, 0.25, 0.  ])

#
# ## Building the training set

# You can now combine the functions that you created in the previous sections, to build a training set for the CBOW model, starting from the following tokenized corpus.

# In[ ]:

# Print corpus
words

# To do this you need to use the sliding window function (`get_windows`) to extract the context words and center words, and you then convert these sets of words into a basic vector representation using `word_to_one_hot_vector` and `context_words_to_vector`.

# In[ ]:

# Print vectors associated to center and context words for corpus
for context_words, center_word in get_windows(words, 2):  # reminder: 2 is the context half-size
    print(f'Context words:  {context_words} -> {context_words_to_vector(context_words, word2Ind, V)}')
    print(f'Center word:  {center_word} -> {word_to_one_hot_vector(center_word, word2Ind, V)}')
    print()

# In this practice notebook you'll be performing a single iteration of training using a single example, but in this week's assignment you'll train the CBOW model using several iterations and batches of example.
# Here is how you would use a Python generator function (remember the `yield` keyword from the lecture?) to make it easier to iterate over a set of examples.

# In[ ]:

# Define the generator function 'get_training_example'
def get_training_example(words, C, word2Ind, V):
    for context_words, center_word in get_windows(words, C):
        yield context_words_to_vector(context_words, word2Ind, V), word_to_one_hot_vector(center_word, word2Ind, V)

# The output of this function can be iterated on to get successive context word vectors and center word vectors, as demonstrated in the next cell.

# In[ ]:

# Print vectors associated to center and context words for corpus using the generator function
for context_words_vector, center_word_vector in get_training_example(words, 2, word2Ind, V):
    print(f'Context words vector:  {context_words_vector}')
    print(f'Center word vector:  {center_word_vector}')
    print()

# Your training set is ready, you can now move on to the CBOW model itself which will be covered in the next lecture notebook.
#
# **Congratulations on finishing this lecture notebook!** Hopefully you now have a better understanding of how to prepare your data before feeding it to a continuous bag-of-words model.
#
# **Keep it up!**


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

from patched_smplx.vertex_ids import vertex_ids
from patched_smplx.vertex_joint_selector import VertexJointSelector
from patched_smplx.lbs import lbs, vertices2landmarks, \
    batch_rodrigues, find_dynamic_lmk_idx_and_bcoords_projective

from .index_mapping import get_kinect_smpl_joints_full, get_face_smplx_landmarks, get_shifted_op_joint_ids

def shift_j_abs(body, joints_output, comp_device, kinect_smpl, shifts):
    v_num = kinect_smpl.shape[0]
    if shifts is None:
        body.kjj = joints_output[:, kinect_smpl[:, 1], :]
    else:
        n = joints_output.shape[0]
        shift_j = shifts[0:v_num, :, 0].reshape(1, -1, 3, 1).repeat(n, 1, 1, 1)
        j_h = torch.cat([shift_j, torch.ones((n, v_num, 1, 1), dtype=torch.float, requires_grad=False, device=comp_device)], dim=2)
        j_h = j_h.reshape(n, v_num, 1, 4)
        j_shift = torch.matmul(j_h, body.A[:,  kinect_smpl[:, 2], :, :]).reshape(n, v_num, 4)[:, :, 0:3]
#         for i in range(0, 100):
#             print('shape kin {} kin {} j_shift {} '.format(kinect_smpl.shape[0], joints_output.shape[1], j_shift.shape[1]))
        jos = joints_output[:, kinect_smpl[:, 1], :]
        body.kjj = jos + j_shift
#         exit()

def shift_op_j_abs(joints, A, shifts, inds=None):
    if shifts is None:
        return
    if inds is None:
        inds = get_shifted_op_joint_ids()
    inds = np.asarray(inds, dtype=np.int) #14
    inds_t = torch.tensor(inds, dtype=torch.long, requires_grad=False, device=joints.device)
    b = joints.shape[0]
    shift_j = shifts.reshape(1, -1, 3, 1).repeat(b, 1, 1, 1)
    j_num = inds.shape[0]
    j_h = torch.cat([shift_j, torch.ones((b, j_num, 1, 1), dtype=torch.float, requires_grad=False, device=joints.device)], dim=2)
    j_h = j_h.reshape(b, j_num, 1, 4)
    j_shift = torch.matmul(j_h, A[:,  inds_t, :, :]).reshape(b, j_num, 4)[:, :, 0:3]
    joints[:, inds_t] = joints[:, inds_t] + j_shift

def shift_v_abs(body, verts_output, comp_device, kinect_smplvert, shifts):
    v_num = kinect_smplvert.shape[0]
    if shifts is None:
        body.kjv = verts_output[:, kinect_smplvert[:, 1], :]
    else:
        n = verts_output.shape[0]
        shift_v = shifts[-v_num:, :, 0].reshape(1, -1, 3, 1).repeat(n, 1, 1, 1)
        v_h = torch.cat(
            [shift_v, torch.ones((n, v_num, 1, 1), dtype=torch.float, requires_grad=False, device=comp_device)], dim=2)
        v_h = v_h.reshape(n, v_num, 1, 4)
        v_shift = torch.matmul(v_h, body.A[:, kinect_smplvert[:, 2], :, :]).reshape(n, v_num, 4)[:, :, 0:3]

    # joints_output[:, kinect_smpl[:, 1], :] = joints_output[:, kinect_smpl[:, 1], :] + torch.matmul(
    #     shifts[0:v_num, :, :].reshape(1, -1, 3, 10).repeat(n, 1, 1, 1), A_h.reshape(n, -1, 10, 1)).reshape(n, -1, 3)

        body.kjv = verts_output[:, kinect_smplvert[:, 1], :] + v_shift

class ExpBodyModel(nn.Module):

    def __init__(self,
                 bm_path,
                 params=None,
                 num_betas=10,
                 v_template=None,
                 num_expressions=10,
                 use_posedirs=True,
                 dtype=torch.float32,
                 comp_device=None,
                 is_hand_pca=False,
                 num_hand_pca=6,
                 w_add=None, v_add=None, fe_scale=1, s2v=None):

        super(ExpBodyModel, self).__init__()

        '''
        :param bm_path: path to a SMPL model as pkl file
        :param num_betas: number of shape parameters to include.
                if betas are provided in params, num_betas would be overloaded with number of thoes betas
        :param batch_size: number of smpl vertices to get
        :param device: default on gpu
        :param dtype: float precision of the compuations
        :return: verts, trans, pose, betas
        '''
        # Todo: if params the batchsize should be read from one of the params

        self.dtype = dtype
        self.fe_scale = fe_scale
        self.s2v = s2v

        if params is None: params = {}

        # -- Load SMPL params --
        if '.npz' in bm_path:
            smpl_dict = np.load(bm_path, encoding='latin1')
        else:
            raise ValueError('bm_path should be either a .pkl nor .npz file')

        njoints = smpl_dict['posedirs'].shape[2] // 3

        if s2v is not None:
            s2v = s2v.reshape(-1, 3, 11)

        # Mean template vertices
        if v_template is None:
            v_template = np.repeat(smpl_dict['v_template'][np.newaxis], 1, axis=0)
        else:
            v_template = np.repeat(v_template[np.newaxis], 1, axis=0)
        if s2v is not None:
            s2v = s2v.reshape(-1, 3, 11)
#             print(s2v.shape)
            v_template = np.concatenate([v_template, s2v[:, :, 10].reshape(1, -1, 3)], axis=1)

        v_template = torch.tensor(v_template, dtype=dtype, device=comp_device)
        if v_add is not None:
            n_add = v_add.shape[0]
            v_template = torch.cat([v_template, v_add.reshape(1, -1, 3)], dim=1)

        self.register_buffer('v_template', v_template)

        self.register_buffer('f', torch.tensor(smpl_dict['f'].astype(np.int32), dtype=torch.int32))

        if len(params):
            if 'betas' in params.keys():
                num_betas = params['betas'].shape[1]

        num_total_betas = smpl_dict['shapedirs'].shape[-1]
        if num_betas < 1:
            num_betas = num_total_betas

        shapedirs = smpl_dict['shapedirs'][:, :, :num_betas]

        if v_add is not None or w_add is not None:
            n_add = w_add.shape[0]
            if s2v is None:
                shapedirs = np.concatenate([shapedirs, np.zeros((n_add, 3, 10))])
            else:
                shapedirs = np.concatenate([shapedirs, s2v[:, :, 0:10]])

        self.register_buffer('shapedirs', torch.tensor(shapedirs, dtype=dtype))

        begin_shape_id = 300 if smpl_dict['shapedirs'].shape[-1] > 300 else 10
        exprdirs = smpl_dict['shapedirs'][:, :, begin_shape_id:(begin_shape_id + num_expressions)]

        if v_add is not None or w_add is not None:
            exprdirs = np.concatenate([exprdirs, np.zeros((n_add, 3, 10))])

        self.register_buffer('exprdirs', torch.tensor(exprdirs, dtype=dtype))

        J_regressor = smpl_dict['J_regressor']
        if v_add is not None or w_add is not None:
            J_regressor = np.concatenate([J_regressor, np.zeros((J_regressor.shape[0], n_add))], axis=1)

        self.register_buffer('J_regressor', torch.tensor(J_regressor, dtype=dtype))

        if use_posedirs:
            posedirs = smpl_dict['posedirs']
            posedirs = posedirs.reshape([posedirs.shape[0] * 3, -1]).T
            if v_add is not None or w_add is not None:
                posedirs = np.concatenate([posedirs, np.zeros((posedirs.shape[0], n_add*3))], axis=1)

            self.register_buffer('posedirs', torch.tensor(posedirs, dtype=dtype))
        else:
            self.posedirs = None

        # indices of parents for each joints
        kintree_table = smpl_dict['kintree_table'].astype(np.int32)
        self.register_buffer('kintree_table', torch.tensor(kintree_table, dtype=torch.long, device=comp_device))

        # LBS weights
        # weights = np.repeat(smpl_dict['weights'][np.newaxis], batch_size, axis=0)
        weights = smpl_dict['weights']

        if w_add is not None:
#             print(weights.shape)
#             print(w_add.shape)
#             exit()
            weights = np.concatenate([weights, w_add])

        self.register_buffer('weights', torch.tensor(weights, dtype=dtype))

        lmk_faces_idx = smpl_dict['lmk_faces_idx']
        self.lmk_faces_idx = torch.tensor(lmk_faces_idx, dtype=torch.long, device=comp_device)
        lmk_bary_coords = smpl_dict['lmk_bary_coords']
        self.lmk_bary_coords = torch.tensor(lmk_bary_coords, dtype=torch.float, device=comp_device)
        dynamic_lmk_faces_idx = smpl_dict['dynamic_lmk_faces_idx']
        self.dynamic_lmk_faces_idx = torch.tensor(dynamic_lmk_faces_idx, dtype=torch.long, device=comp_device)

        dynamic_lmk_bary_coords = smpl_dict['dynamic_lmk_bary_coords']
        self.dynamic_lmk_bary_coords = torch.tensor(
            dynamic_lmk_bary_coords, dtype=dtype, device=comp_device)

        self.neck_kin_chain = []
        NECK_IDX = 12
        curr_idx = torch.tensor(NECK_IDX, dtype=torch.long, device=comp_device)
        parents = self.kintree_table[0]
        while curr_idx != -1:
            self.neck_kin_chain.append(curr_idx)
            curr_idx = parents[curr_idx]

        self.neck_kin_chain = torch.stack(self.neck_kin_chain)

        # SMPL and SMPL-H share the same topology, so any extra joints can
        # be drawn from the same place

        self.dtype = dtype

        self.vertex_joint_selector = VertexJointSelector(
            vertex_ids=vertex_ids['smplx'])

        self.is_hand_pca = is_hand_pca
        self.num_hand_pca = num_hand_pca
        if is_hand_pca:
            self.left_hand_mean = torch.tensor(smpl_dict['hands_meanl'], dtype=torch.float, device=comp_device,
                                               requires_grad=False)
            self.right_hand_mean = torch.tensor(smpl_dict['hands_meanr'], dtype=torch.float, device=comp_device,
                                                requires_grad=False)
            self.left_hand_components = torch.tensor(smpl_dict['hands_componentsl'][:num_hand_pca], dtype=torch.float,
                                                     device=comp_device, requires_grad=False)
            self.right_hand_components = torch.tensor(smpl_dict['hands_componentsr'][:num_hand_pca], dtype=torch.float,
                                                      device=comp_device, requires_grad=False)

#         self.kinect_smpl = get_kinect_smpl_joints()
        self.kinect_smpl = get_kinect_smpl_joints_full()
#         self.kinect_smplvert = get_kinect_smplx_vert()
        self.kinect_smplvert = np.zeros((0, 3), dtype=int)
        self.face_smplx = get_face_smplx_landmarks()
        self.device = comp_device

        self.J = None

    def default_forward(
        self,
        global_trans=None, global_rot=None,
        pose_body=None,
        face_expression=None, pose_jaw=None, pose_eyes=None,
        pose_hand=None,
        beta=None,
        allow_beta_none=False
    ):
        if beta is None:
            assert allow_beta_none

        def to_t(d, device, dtype=torch.float32):
            return torch.tensor(d, device=device, dtype=dtype)

        def to_batch_t(x, size=None):
            if x is None:
                if size is None:
                    raise Exception()
                x = np.zeros(size)
            return to_t(x, device=self.device).reshape(1, -1)

        global_trans = to_batch_t(global_trans, 3)
        global_rot = to_batch_t(global_rot, 3)
        pose_body = to_batch_t(pose_body, 63)
        face_expression = to_batch_t(face_expression, 10)
        pose_jaw = to_batch_t(pose_jaw, 3)
        pose_eyes = to_batch_t(pose_eyes, 6)
        pose_hand = to_batch_t(pose_hand, 2 * self.num_hand_pca)
        beta = to_batch_t(beta, 10)

        forward_result = self(
            root_orient=global_rot,
            pose_body=pose_body,
            pose_hand=pose_hand,
            pose_jaw=pose_jaw,
            pose_eye=pose_eyes,
            betas=beta,
            trans=global_trans,
            expression=face_expression
        )
        return forward_result

    def forward(
        self,
        root_orient, pose_body,
        pose_hand,
        pose_jaw, pose_eye,
        betas,
        trans,
        expression,
        shifts=None,
        shifts_op=None, shift_op_inds=None,
        return_dict=False,
        **kwargs
    ):
        # assert not (v_template  is not None and betas  is not None), ValueError('vtemplate and betas could not be used jointly.')

        if self.is_hand_pca:
            n_h = int(pose_hand.shape[1] / 2)
            left_hand_pose = pose_hand[:, 0:n_h]
            right_hand_pose = pose_hand[:, n_h:]
            b = pose_hand.shape[0]
            if n_h > 0:
                left_hand_pose = self.left_hand_mean.reshape(1, -1).repeat(b, 1) + torch.einsum(
                    'bi,ij->bj', [left_hand_pose, self.left_hand_components])
                right_hand_pose = self.right_hand_mean.reshape(1, -1).repeat(b, 1) + torch.einsum(
                    'bi,ij->bj', [right_hand_pose, self.right_hand_components])
            else:
                left_hand_pose = self.left_hand_mean.reshape(1, -1).repeat(b, 1)
                right_hand_pose = self.right_hand_mean.reshape(1, -1).repeat(b, 1)
            pose_hand = torch.cat([left_hand_pose, right_hand_pose], dim=1)
        v_template = self.v_template
#         if self.s2v is not None:
#             betas_h = torch.cat([betas, torch.ones((b, 1), device=self.device)], dim=1)
#             v_add = torch.bmm(self.s2v.reshape(1, -1, 11).repeat(b,1,1), betas_h.reshape(b, -1, 1)).reshape(b, -1, 3)
#             v_template = torch.cat([v_template, v_add[0]], dim=1)

        if len(pose_body.shape) > 2:
            #rotation matrices as input
            b = pose_hand.shape[0]
            pose_hand = batch_rodrigues(pose_hand.reshape(-1, 3)).reshape(b, -1, 3, 3)
#             print(b)
#             print(root_orient.shape)
#             print(pose_body.shape)
#             print(pose_jaw.shape)
#             print(pose_eye.shape)
#             print(pose_hand.shape)
            full_pose = torch.cat([root_orient.reshape(-1, 1, 3, 3), pose_body.reshape(b, -1, 3, 3), pose_jaw.reshape(-1, 1, 3, 3),
                                   pose_eye.reshape(-1, 2, 3, 3), pose_hand],
                                  dim=1)  # orient:3, body:63, jaw:3, eyel:3, eyer:3, handl, handr
            b = full_pose.shape[0]
            full_pose = full_pose.reshape(b, -1, 9)
            pose2rot = False
        else:
            full_pose = torch.cat([root_orient, pose_body, pose_jaw, pose_eye, pose_hand],
                                  dim=1)  # orient:3, body:63, jaw:3, eyel:3, eyer:3, handl, handr
            pose2rot = True

        shape_components = torch.cat([betas, self.fe_scale * expression], dim=-1)
        shapedirs = torch.cat([self.shapedirs, self.exprdirs], dim=-1)

        verts, joints, A, joints_init = lbs(betas=shape_components, pose=full_pose, v_template=v_template,
                                            shapedirs=shapedirs, posedirs=self.posedirs.unsqueeze(0),
                                            J_regressor=self.J_regressor, parents=self.kintree_table[0].long(),
                                            lbs_weights=self.weights, dtype=self.dtype, pose2rot=pose2rot)

        bs = shape_components.shape[0]
        lmk_faces_idx = self.lmk_faces_idx.unsqueeze(
            dim=0).expand(bs, -1).contiguous()
        lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat(
            bs, 1, 1)

        # dyn_lmk_faces_idx, dyn_lmk_bary_coords = find_dynamic_lmk_idx_and_bcoords(
        #     verts, full_pose, self.dynamic_lmk_faces_idx,
        #     self.dynamic_lmk_bary_coords,
        #     self.neck_kin_chain, dtype=self.dtype)

        neck_joint_loc = joints[:, 12] + trans
#         print(neck_joint_loc.shape)
        dyn_lmk_faces_idx, dyn_lmk_bary_coords = find_dynamic_lmk_idx_and_bcoords_projective(
                verts, neck_joint_loc.reshape(-1, 3), full_pose, self.dynamic_lmk_faces_idx,
                self.dynamic_lmk_bary_coords,
                self.neck_kin_chain, dtype=self.dtype)

        lmk_faces_idx = torch.cat([lmk_faces_idx,
                                   dyn_lmk_faces_idx], 1)
        lmk_bary_coords = torch.cat(
            [lmk_bary_coords.expand(bs, -1, -1),
             dyn_lmk_bary_coords], 1)

        landmarks = vertices2landmarks(verts, self.f.long(),
                                       lmk_faces_idx,
                                       lmk_bary_coords)

        # Add any extra joints that might be needed
        joints = self.vertex_joint_selector(verts, joints)
        # Add the landmarks to the joints
        joints = torch.cat([joints, landmarks], dim=1)
        # Map the joints to the current dataset

        # if self.joint_mapper is not None:
        #     joints = self.joint_mapper(joints=joints, vertices=verts)

        Jtr = joints + trans.unsqueeze(dim=1)
        verts = verts + trans.unsqueeze(dim=1)
#         print('shifts op shape is ')
#         if shifts_op is not None:
#             print(shifts_op.shape)
#         else:
#             print('None')
        shift_op_j_abs(Jtr, A, shifts_op, shift_op_inds)

        res = {}
        res['v'] = verts
        res['f'] = self.f
        res['betas'] = betas
        res['Jtr'] = Jtr  # Todo: ik can be made with vposer
        res['A'] = A
        res['J'] = joints_init

        res['pose_body'] = pose_body
        res['pose_hand'] = pose_hand
        res['pose_jaw'] = pose_jaw
        res['pose_eye'] = pose_eye
        res['full_pose'] = full_pose
        if shifts_op is not None:
            res['shifts_op'] = shifts_op

        if not return_dict:
            class result_meta(object):
                pass

            res_class = result_meta()
            for k, v in res.items():
                res_class.__setattr__(k, v)
            res = res_class

        shift_j_abs(res, res.Jtr, verts.device, self.kinect_smpl, shifts)
        shift_v_abs(res, res.v, verts.device, self.kinect_smplvert, None)

        res.kj = torch.cat([res.kjj, res.kjv], dim=1)
        res.kj = torch.cat([res.kj, res.Jtr[:, self.face_smplx[:, 1]]], dim=1)
        res.v_a = res.v

        return res


from __future__ import division
from qm.qchem.qchem import QChem
from misc import au_to_A, eV_to_au
import os, shutil, re, textwrap, subprocess
import numpy as np

class DFT(QChem):
    """ Class for DFT method of Q-Chem 5.2

        :param object molecule: Molecule object
        :param string basis_set: Basis set information
        :param integer memory: Allocatable memory in the calculation
        :param integer nthreads: Number of threads in the calculation
        :param string functional: Exchange-correlation functional
        :param integer scf_max_iter: Maximum number of SCF iterations
        :param integer scf_wf_tol: Wave function convergence for SCF iterations
        :param integer cis_max_iter: Maximum number of CIS iterations
        :param integer cis_en_tol: Energy convergence for CIS iterations
        :param integer cpscf_max_iter: Maximum number of CP iterations
        :param integer cpscf_grad_tol: Gradient convergence of CP iterations
        :param string root_path: Path for Q-Chem root directory
        :param string version: Q-Chem version
    """
    def __init__(self, molecule, basis_set="sto-3g", memory=2000, nthreads=1, \
        functional="blyp", scf_max_iter=50, scf_wf_tol=8, cis_max_iter=30, cis_en_tol=6, \
        cpscf_max_iter=30, cpscf_grad_tol=6, root_path="./", version="5.2"):
        # Initialize Q-Chem common variables
        super(DFT, self).__init__(basis_set, memory, root_path, nthreads, version)

        self.functional = functional
        self.scf_max_iter = scf_max_iter
        self.nthreads = nthreads
        self.scf_wf_tol = scf_wf_tol
        self.cis_max_iter = cis_max_iter
        self.cis_en_tol = cis_en_tol
        self.cpscf_max_iter = cpscf_max_iter
        self.cpscf_grad_tol = cpscf_grad_tol

        # Q-Chem can provide NACs
        molecule.l_nacme = False
        self.re_calc = True

    def get_data(self, molecule, base_dir, bo_list, dt, istep, calc_force_only):
        """ Extract energy, gradient and nonadiabatic couplings from (TD)DFT method

            :param object molecule: Molecule object
            :param string base_dir: Base directory
            :param integer,list bo_list: List of BO states for BO calculation
            :param double dt: Time interval
            :param integer istep: Current MD step
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        super().get_data(base_dir, calc_force_only)
        self.write_xyz(molecule)
        self.get_input(molecule, bo_list, calc_force_only)
        self.run_QM(base_dir, istep, bo_list)
        self.extract_QM(molecule, bo_list, calc_force_only)
        self.move_dir(base_dir)

    def get_input(self, molecule, bo_list, calc_force_only):
        """ Generate Q-Chem input files: qchem.in

            :param object molecule: Molecule object
            :param integer istep: Current MD step
            :param integer,list bo_list: List of BO states for BO calculation
        """
        # Make Q-Chem input file
        input_qc = ""

        # Molecular information such as charge, geometry
        input_molecule = textwrap.dedent(f"""\
        $molecule
        {int(molecule.charge)}  1
        """)

        for iat in range(molecule.nat_qm):
            input_molecule += f"{molecule.symbols[iat]}"\
                + "".join([f"{i:15.8f}" for i in molecule.pos[iat]]) + "\n"
        input_molecule += "$end\n\n"
        input_qc += input_molecule

        # Job control to calculate NAC
        if (not calc_force_only and self.calc_coupling):
            # Arguments about SCF, xc functional and basis set
            input_nac = textwrap.dedent(f"""\
            $rem
            JOBTYPE SP
            INPUT_BOHR TRUE
            METHOD {self.functional}
            BASIS {self.basis_set}
            SCF_CONVERGENCE {self.scf_wf_tol}
            SYMMETRY FALSE
            SYM_IGNORE TRUE
            """)

            # Arguments about TDDFT and NAC
            input_nac += textwrap.dedent(f"""\
            CIS_N_ROOTS {molecule.nst-1}
            CIS_TRIPLETS FALSE
            CIS_CONVERGENCE {self.cis_en_tol}
            MAX_CIS_CYCLES {self.cis_max_iter}
            CALC_NAC TRUE
            CIS_DER_NUMSTATE {molecule.nst}
            SET_ITER {self.cpscf_max_iter}
            SET_CONV {self.cpscf_grad_tol}
            $end

            $derivative_coupling
            This is comment line
            """)

            for ist in range(molecule.nst):
                input_nac += f"{ist}  "
            input_nac += "\n$end\n\n"
            input_qc += input_nac

        # Job control to calculate force
        input_force = ""

        # BOMD: calc_force_only = F, self_calc_coupling = F
        # In BOMD, read in molecule section, scf_guess and skip_scf are not valid
        guess = "SAD"; skip = "FALSE"
        for ist in bo_list:
            if (not calc_force_only and self.calc_coupling):
                guess = "READ"; skip = "TRUE"
                input_force = textwrap.dedent(f"""\
                @@@

                $molecule
                read
                $end

                """)

            input_force += textwrap.dedent(f"""\
            $rem
            JOBTYPE force
            INPUT_BOHR TRUE
            METHOD {self.functional}
            BASIS {self.basis_set}
            SCF_GUESS {guess}
            SYMMETRY FALSE
            SYM_IGNORE TRUE
            """)

            # When ground state force is calculated, Q-Chem doesn't need CIS option.
            if (ist != 0):
                input_force += textwrap.dedent(f"""\
                CIS_N_ROOTS {molecule.nst-1}
                CIS_STATE_DERIV {ist}
                CIS_TRIPLETS FALSE
                CIS_CONVERGENCE {self.cis_en_tol}
                MAX_CIS_CYCLES {self.cis_max_iter}
                SET_ITER {self.cpscf_max_iter}
                SET_CONV {self.cpscf_grad_tol}
                """)

                # CIS solution isn't saved in scratch.
                if (not calc_force_only and self.calc_coupling):
                    input_force += textwrap.dedent(f"""\
                    CIS_GUESS_DISK TRUE
                    CIS_GUESS_DISK_TYPE 2
                    SKIP_CIS_RPA TRUE
                    """)
            input_force += "$end\n\n"

            input_qc += input_force

        file_name = "qchem.in"
        with open(file_name, "w") as f:
            f.write(input_qc)

    def run_QM(self, base_dir, istep, bo_list):
        """ Run (TD)DFT calculation and save the output files to qm_log directory

            :param string base_dir: Base directory
            :param integer istep: Current MD step
            :param integer,list bo_list: List of BO states for BO calculation
        """
        # Set environment variable
        os.environ["QC"] = self.root_path
        path_qcenv = os.path.join(self.root_path, "qcenv.sh")
        command = f'env -i sh -c "source {path_qcenv} && env"'
        for line in subprocess.getoutput(command).split("\n"):
            key, value = line.split("=")
            os.environ[key] = value
        os.environ["QCSCRATCH"] = self.scr_qm_dir
        os.environ["QCLOCALSCR"] = self.scr_qm_dir

        #TODO: MPI binary
        qm_exec_command = f"$QC/bin/qchem -nt {self.nthreads} qchem.in log save > qcprog.info "

        # Run Q-Chem
        os.system(qm_exec_command)

        tmp_dir = os.path.join(base_dir, "qm_log")
        if (os.path.exists(tmp_dir)):
            log_step = f"log.{istep + 1}.{bo_list[0]}"
            shutil.copy("log", os.path.join(tmp_dir, log_step))

    def extract_QM(self, molecule, bo_list, calc_force_only):
        """ Read the output files to get BO information

            :param object molecule: Molecule object
            :param integer,list bo_list: List of BO states for BO calculation
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        file_name = "log"
        with open(file_name, "r") as f:
            log = f.read()

        if (not calc_force_only):
            # Ground state energy
            energy = re.findall('Total energy in the final basis set =\s*([-]*\S*)', log)
            energy = np.array(energy, dtype=np.float64)
            molecule.states[0].energy = energy[0]

            # Excited state energy
            if (molecule.nst > 1):
                energy = re.findall('Total energy for state\s*\d*.\s*([-]*\S*)', log)
                energy = np.array(energy, dtype=np.float64)

                for ist, en in enumerate(energy):
                    if ist < molecule.nst - 1:
                        molecule.states[ist + 1].energy = en

        # Adiabatic force
        tmp_f = "Gradient of\D*\s*"
        num_line = int(molecule.nat_qm / 6)
        if (num_line >= 1):
            tmp_f += ("\s*\d*\s*\d*\s*\d*\s*\d*\s*\d*\s*\d*"
                 + ("\s*\d?\s*" + "([-]*\S*)\s*" * 6) * 3) * num_line

        dnum = molecule.nat_qm % 6
        tmp_f += "\s*\d*" * dnum
        tmp_f += ("\s*\d?\s*" + "([-]*\S*)\s*" * dnum) * 3

        force = re.findall(tmp_f, log)
        force = np.array(force, dtype=np.float64)

        # Q-Chem provides energy gradient not force
        force = -force

        for index, ist in enumerate(bo_list):
            iline = 0; iiter = 0
            for iiter in range(num_line):
                tmp_force = np.transpose(force[index][18 * iiter:18 * (iiter + 1)].reshape(3, 6, order="C"))
                for iat in range(6):
                    molecule.states[ist].force[6 * iline + iat] = np.copy(tmp_force[iat])
                iline += 1

            if (dnum != 0):
                if (num_line != 0):
                    tmp_force = np.transpose(force[index][18 * (iiter + 1):].reshape(3, dnum, order="C"))
                else:
                    tmp_force = np.transpose(force[index][0:].reshape(3, dnum, order="C"))

                for iat in range(dnum):
                    molecule.states[ist].force[6 * iline + iat] = np.copy(tmp_force[iat])

        # NACs
        if (not calc_force_only and self.calc_coupling):
            tmp_nac = "with ETF[:]*\s*Atom\s*X\s*Y\s*Z\s*[-]*" + ("\s*\d*\s*" + "([-]*\S*)\s*"*3) * molecule.nat_qm
            nac = re.findall(tmp_nac, log)
            nac = np.array(nac, dtype=np.float64)

            kst = 0
            for ist in range(molecule.nst):
                for jst in range(ist + 1, molecule.nst):
                    molecule.nac[ist, jst] = nac[kst].reshape(molecule.nat_qm, 3, order='C')
                    molecule.nac[jst, ist] = - molecule.nac[ist, jst]
                    kst += 1


#!/usr/bin/python
#-*- coding: utf-8 -*-

import sys, time, os, argparse, socket
import yaml
import numpy
import pdb
import torch
import glob
from tuneThreshold import tuneThresholdfromScore
from SpeakerNet import SpeakerNet
from DatasetLoader import get_data_loader

parser = argparse.ArgumentParser(description = "SpeakerNet");

parser.add_argument('--config',         type=str,   default=None,   help='Config YAML file');

## Data loader
parser.add_argument('--max_frames',     type=int,   default=200,    help='Input length to the network for training');
parser.add_argument('--eval_frames',    type=int,   default=300,    help='Input length to the network for testing; 0 uses the whole files');
parser.add_argument('--batch_size',     type=int,   default=200,    help='Batch size, number of speakers per batch');
parser.add_argument('--max_seg_per_spk', type=int,  default=100,    help='Maximum number of utterances per speaker per epoch');
parser.add_argument('--nDataLoaderThread', type=int, default=5,     help='Number of loader threads');
parser.add_argument('--augment',        type=bool,  default=False,  help='Augment input')

## Training details
parser.add_argument('--test_interval',  type=int,   default=10,     help='Test and save every [test_interval] epochs');
parser.add_argument('--max_epoch',      type=int,   default=500,    help='Maximum number of epochs');
parser.add_argument('--trainfunc',      type=str,   default="",     help='Loss function');

## Optimizer
parser.add_argument('--optimizer',      type=str,   default="adam", help='sgd or adam');
parser.add_argument('--scheduler',      type=str,   default="steplr", help='Learning rate scheduler');
parser.add_argument('--lr',             type=float, default=0.001,  help='Learning rate');
parser.add_argument("--lr_decay",       type=float, default=0.95,   help='Learning rate decay every [test_interval] epochs');
parser.add_argument('--weight_decay',   type=float, default=0,      help='Weight decay in the optimizer');

## Loss functions
parser.add_argument("--hard_prob",      type=float, default=0.5,    help='Hard negative mining probability, otherwise random, only for some loss functions');
parser.add_argument("--hard_rank",      type=int,   default=10,     help='Hard negative mining rank in the batch, only for some loss functions');
parser.add_argument('--margin',         type=float, default=1,      help='Loss margin, only for some loss functions');
parser.add_argument('--scale',          type=float, default=15,     help='Loss scale, only for some loss functions');
parser.add_argument('--nPerSpeaker',    type=int,   default=1,      help='Number of utterances per speaker per batch, only for metric learning based losses');
parser.add_argument('--nClasses',       type=int,   default=5994,   help='Number of speakers in the softmax layer, only for softmax-based losses');

## Load and save
parser.add_argument('--initial_model',  type=str,   default="",     help='Initial model weights');
parser.add_argument('--save_path',      type=str,   default="./data/exp1", help='Path for model and logs');

## Training and test data
parser.add_argument('--train_list',     type=str,   default="",     help='Train list');
parser.add_argument('--test_list',      type=str,   default="",     help='Evaluation list');
parser.add_argument('--train_path',     type=str,   default="data/voxceleb2", help='Absolute path to the train set');
parser.add_argument('--test_path',      type=str,   default="data/voxceleb1", help='Absolute path to the test set');
parser.add_argument('--musan_path',     type=str,   default="data/musan_split", help='Absolute path to the test set');
parser.add_argument('--rir_path',       type=str,   default="data/RIRS_NOISES/simulated_rirs", help='Absolute path to the test set');

## Model definition
parser.add_argument('--n_mels',         type=int,   default=40,     help='Number of mel filterbanks');
parser.add_argument('--log_input',      type=bool,  default=False,  help='Log input features')
parser.add_argument('--model',          type=str,   default="",     help='Name of model definition');
parser.add_argument('--encoder_type',   type=str,   default="SAP",  help='Type of encoder');
parser.add_argument('--nOut',           type=int,   default=512,    help='Embedding size in the last FC layer');

## For test only
parser.add_argument('--eval', dest='eval', action='store_true', help='Eval only')

args = parser.parse_args();

## Parse YAML
def find_option_type(key, parser):
    for opt in parser._get_optional_actions():
        if ('--' + key) in opt.option_strings:
           return opt.type
    raise ValueError

if args.config is not None:
    with open(args.config, "r") as f:
        yml_config = yaml.load(f, Loader=yaml.FullLoader)
    for k, v in yml_config.items():
        if k in args.__dict__:
            typ = find_option_type(k, parser)
            args.__dict__[k] = typ(v)
        else:
            sys.stderr.write("Ignored unknown parameter {} in yaml.\n".format(k))

## Initialise directories
model_save_path     = args.save_path+"/model"
result_save_path    = args.save_path+"/result"

if not(os.path.exists(model_save_path)):
    os.makedirs(model_save_path)

if not(os.path.exists(result_save_path)):
    os.makedirs(result_save_path)

## Load models
s = SpeakerNet(**vars(args));

it          = 1;
prevloss    = float("inf");
sumloss     = 0;
min_eer     = [100];

## Load model weights
modelfiles = glob.glob('%s/model0*.model'%model_save_path)
modelfiles.sort()

if len(modelfiles) >= 1:
    s.loadParameters(modelfiles[-1]);
    print("Model %s loaded from previous state!"%modelfiles[-1]);
    it = int(os.path.splitext(os.path.basename(modelfiles[-1]))[0][5:]) + 1
elif(args.initial_model != ""):
    s.loadParameters(args.initial_model);
    print("Model %s loaded!"%args.initial_model);

for ii in range(0,it-1):
    s.__scheduler__.step()

## Evaluation code
if args.eval == True:

    sc, lab, trials = s.evaluateFromList(args.test_list, print_interval=100, test_path=args.test_path, eval_frames=args.eval_frames)
    result = tuneThresholdfromScore(sc, lab, [1, 0.1]);
    print('EER %2.4f'%result[1])

    ## Save scores
    print('Type desired file name to save scores. Otherwise, leave blank.')
    userinp = input()

    while True:
        if userinp == '':
            quit();
        elif os.path.exists(userinp) or '.' not in userinp:
            print('Invalid file name %s. Try again.'%(userinp))
            userinp = input()
        else:
            with open(userinp,'w') as outfile:
                for vi, val in enumerate(sc):
                    outfile.write('%.4f %s\n'%(val,trials[vi]))
            quit();

## Write args to scorefile
scorefile = open(result_save_path+"/scores.txt", "a+");

for items in vars(args):
    print(items, vars(args)[items]);
    scorefile.write('%s %s\n'%(items, vars(args)[items]));
scorefile.flush()

## Initialise data loader
trainLoader = get_data_loader(args.train_list, **vars(args));

while(1):

    clr = [x['lr'] for x in s.__optimizer__.param_groups]

    print(time.strftime("%Y-%m-%d %H:%M:%S"), it, "Training %s with LR %f..."%(args.model,max(clr)));

    ## Train network
    loss, traineer = s.train_network(loader=trainLoader);

    ## Validate and save
    if it % args.test_interval == 0:

        print(time.strftime("%Y-%m-%d %H:%M:%S"), it, "Evaluating...");

        sc, lab, _ = s.evaluateFromList(args.test_list, print_interval=100, test_path=args.test_path, eval_frames=args.eval_frames)
        result = tuneThresholdfromScore(sc, lab, [1, 0.1]);

        min_eer.append(result[1])

        print(time.strftime("%Y-%m-%d %H:%M:%S"), "LR %f, TEER/TAcc %2.2f, TLOSS %f, VEER %2.4f, MINEER %2.4f"%( max(clr), traineer, loss, result[1], min(min_eer)));
        scorefile.write("IT %d, LR %f, TEER/TAcc %2.2f, TLOSS %f, VEER %2.4f, MINEER %2.4f\n"%(it, max(clr), traineer, loss, result[1], min(min_eer)));

        scorefile.flush()

        s.saveParameters(model_save_path+"/model%09d.model"%it);

        with open(model_save_path+"/model%09d.eer"%it, 'w') as eerfile:
            eerfile.write('%.4f'%result[1])

    else:

        print(time.strftime("%Y-%m-%d %H:%M:%S"), "LR %f, TEER/TAcc %2.2f, TLOSS %f"%( max(clr), traineer, loss));
        scorefile.write("IT %d, LR %f, TEER/TAcc %2.2f, TLOSS %f\n"%(it, max(clr), traineer, loss));

        scorefile.flush()

    if it >= args.max_epoch:
        quit();

    it+=1;
    print("");

scorefile.close();


from scipy.spatial import distance	 		#imported distance function from scipy.spatial library
from imutils import face_utils 				#imported face_utils function from imutils library
import imutils								#imported the imputils
import dlib 								#imprted dlib library
import cv2									#import cv2
from time import sleep
import Rpi.GPIO as GPIO
import smtplib
import requests
import json
from urllib.request import urlopen
from geopy.geocoders import Nomination
import traceback
import sys
def eye_aspect_ratio(eye):
	dis1 = distance.euclidean(eye[1], eye[5]) #get the distance between eyes two points"""
	dis2= distance.euclidean(eye[2], eye[4])  #same
	dis3 = distance.euclidean(eye[0], eye[3]) #saME
	ear = (dis1 + dist2) / (2.0 * dis3)		  #gives the eye aspect ration from equeution
	return ear
url ="https://www.fast2sms.com/dev/bulk"
thresh = 0.25								#threshold value for ear comparition
frame_check = 8							#number of frame which is threshold to give alert
GPIO.setmode(GPIO.BCM)
GPIO.setwarnings(False)
GPIO.setup(26,GPIO.OUT)
detect = dlib.get_frontal_face_detector()   #it asign the detector fuction
predict = dlib.shape_predictor(".\shape_predictor_68_face_landmarks.dat")# Dat file is the crux of the code
smtpuser="<sender_email_id>"
password="<sender_password>"
add="reciver_email_id"
fromadd=smtpuser
subject="alert"
header='To:'+add+'\n'+'From: '+fromadd +'\n'+'subject:' +subject
body='hiii I am deiver Drowsyness detection,<user_name> is sriving now and and he is drowsy please call him/her or contact them'

my_data = {
     # Your default Sender ID
    'sender_id': 'FSTSMS',

     # Put your message here!
    'message': 'This is a test message',

    'language': 'english',
    'route': 'p',

    # You can send sms to multiple numbers
    # separated by comma.
    'numbers': '9999999999, 7777777777, 6666666666'
}
headers = {
    'authorization': 'YOUR API KEY HERE',			#enter your api key here
    'Content-Type': "application/x-www-form-urlencoded",
    'Cache-Control': "no-cache"
}
(lStart, lEnd) = face_utils.FACIAL_LANDMARKS_68_IDXS["left_eye"]#savs the left eyes lanmarks(position)
(rStart, rEnd) = face_utils.FACIAL_LANDMARKS_68_IDXS["right_eye"]#savs the right eyes lanmarks(position)
cap=cv2.VideoCapture(0)						#captures the live stream video
#cap=cv2.VideoCapture("first.mp4")			#capture the static camera
flag=0										#flag for counting
f_num=0
while True:
	ret, frame=cap.read()					#it returns the ret(true,false) and frame
	if f_num%3==0:
		frame = imutils.resize(frame, width=450)#it resize the frame
		gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)#it converts frmae's bgr color to gray color
		subjects = detect(gray, 0)				#it detects face from the gray image
		for subject in subjects:
			shape = predict(gray, subject)		#it predicts the shape from the subject
			shape = face_utils.shape_to_np(shape)#converting to NumPy Array
			leftEye = shape[lStart:lEnd]		#it gives the shape co-ordinates of left-eyes
			rightEye = shape[rStart:rEnd]
			leftEAR = eye_aspect_ratio(leftEye) #it saves the ear value for lefteye
			rightEAR = eye_aspect_ratio(rightEye)
			ear = (leftEAR + rightEAR) / 2.0	#it takes the average of ear value
			leftEyeHull = cv2.convexHull(leftEye)#it creates the convexhull shape for lefteye
			rightEyeHull = cv2.convexHull(rightEye)#it creates the convexhull shape for righteye
			cv2.drawContours(frame, [leftEyeHull], -1, (0, 255, 0), 1)#it draws the left eye convexhull on frame
			cv2.drawContours(frame, [rightEyeHull], -1, (0, 255, 0), 1)#it draws the left eye convexhull on frame
			if ear < thresh:					#it compairs the ear and threshlod ratio
				flag += 1
				print (flag)
				if flag >= frame_check:			#it compairs the flag and frame check
					cv2.putText(frame, "****************ALERT!****************", (10, 30),
						cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)#put text on the frame
					cv2.putText(frame, "****************ALERT!****************", (10,325),
						cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
					#print ("Drowsy")
					GPIO.output(26,GPIO.HIGH)
					if(flag==frame_check):
						s=smtplib.SMTP('smtp.gmail.com',587)
						s.ehlo()
						s.starttls()
						s.echo()
						s.login(smtpuser,password)
						s.sendmail(fromadd,add,header+ '\n'+body)
						s.quit()
						try:
							response = requests.request("POST",
                            url,
                            data = my_data,
                            headers = headers)
#
							#load json data from source
							returned_msg = json.loads(response.text)

							# print the send message
							print(returned_msg['message'])
						except:
							exc_type,exc_value,exc_traceback=sys.exc_info()
							s=smtplib.SMTP('smtp.gmail.com',587)
							s.ehlo()
							s.starttls()
							s.ehlo()
							s.login(smtpuser,password)
							s.sendmail(fromadd,add,header+'\n'+str(exc_value))
							s.quit()
					else:
						GPIO.output(26,GPIO.LOW)

			else:
				flag = 0
		cv2.imshow("Frame", frame)				#display the frame
		key = cv2.waitKey(1) & 0xFF				#wait to enter key tostop execution
		if key == ord("q"):
			break
	f_num+=1
cv2.destroyAllWindows()						#distroys all windows
cap.stop()									#stops the camera


import math
import numpy as np
import basis.robot_math as rm
import modeling.meshtools as mt

if __name__ == '__main__':
    '''
    There is an inherent transformation in the official Universal Robot UR3E dae file.
    The old trimesh used by our system does not support applying the transoformation.
    This file was designed to make up the failure.
    As 20201208, I have included the transformation in the io/dae.py file. This conversion is no longer needed.
    It is kept here for reference.
    author: weiwei
    date: 20201207osaka
    '''
    mt.convert_to_stl("base.dae", "base.stl")
    mt.convert_to_stl("shoulder.dae", "shoulder.stl")
    mt.convert_to_stl("upperarm.dae", "upperarm.stl")
    mt.convert_to_stl("forearm.dae", "forearm.stl")
    mt.convert_to_stl("wrist1.dae", "wrist1.stl")
    mt.convert_to_stl("wrist2.dae", "wrist2.stl")
    mt.convert_to_stl("wrist3.dae", "wrist3.stl")