# The MIT License (MIT)
#
# Copyright (c) 2016-2018 Albert Kottke
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np
import numba

from scipy.optimize import minimize

from .site import Location, Layer, Profile
from .motion import WaveField, GRAVITY

class AbstractCalculator(object):
    def __init__(self):
        self._loc_input = None
        self._motion = None
        self._profile = None

    def __call__(self, motion, profile, loc_input):
        self._motion = motion
        self._profile = profile
        self._loc_input = loc_input

    @property
    def motion(self):
        return self._motion

    @property
    def profile(self):
        return self._profile

    @property
    def loc_input(self):
        return self._loc_input

    def calc_accel_tf(self, lin, lout):
        raise NotImplementedError

    def calc_stress_tf(self, lin, lout, damped):
        raise NotImplementedError

    def calc_strain_tf(self, lin, lout):
        raise NotImplementedError

@numba.jit
def my_trapz(thickness, property, depth_max):
    total = 0
    depth = 0

    for t, p in zip(thickness, property):
        depth += t
        if depth_max < depth:
            # Partial layer
            total += (t - (depth - depth_max)) * p
            break
        total += t * p
    else:
        # Final infinite layer
        total += (depth_max - depth) * p

    return total / depth_max

class QuarterWaveLenCalculator(AbstractCalculator):
    """Compute quarter-wave length site amplification.

    No consideration for nolninearity is made by this calculator.
    """

    def __init__(self, site_atten=None):
        super().__init__()
        self._site_atten = site_atten

    def __call__(self, motion, profile, loc_input):
        """Perform the wave propagation.

        Parameters
        ----------
        motion: :class:`~.base.motion.Motion`
            Input motion.

        profile: :class:`~.base.site.Profile`
            Site profile.

        loc_input: :class:`~.base.site.Location`
            Location of the input motion.
        """
        super().__call__(motion, profile, loc_input)

        self._crustal_amp, self._site_term = self._calc_amp(
            profile.density, profile.thickness, profile.slowness)

    @property
    def crustal_amp(self):
        return self._crustal_amp

    @property
    def site_term(self):
        return self._site_term

    @property
    def site_atten(self):
        return self._site_atten

    def _calc_amp(self, density, thickness, slowness):
        freqs = self.motion.freqs
        # 1/4 wavelength depth -- estimated for mean slowness
        qwl_depth = 1 / (4 * np.mean(slowness) * freqs)

        def qwl_average(param):
            return np.array(
                [my_trapz(thickness, param, qd) for qd in qwl_depth])

        for _ in range(20):
            qwl_slowness = qwl_average(slowness)
            prev_qwl_depth = qwl_depth
            qwl_depth = 1 / (4 * qwl_slowness * freqs)

            if np.allclose(prev_qwl_depth, qwl_depth, rtol=0.005):
                break

        # FIXME return an error if not converged?

        qwl_density = qwl_average(density)
        crustal_amp = np.sqrt(
            (density[-1] / slowness[-1]) / (qwl_density / qwl_slowness))
        site_term = np.array(crustal_amp)
        if self._site_atten:
            site_term *= np.exp(-np.pi * self.site_atten * freqs)

        return crustal_amp, site_term

    def fit(self,
            target_type,
            target,
            adjust_thickness=False,
            adjust_site_atten=False,
            adjust_source_vel=False):
        """
        Fit to a target crustal amplification or site term.

        The fitting process adjusts the velocity, site attenuation, and layer
        thickness (if enabled) to fit a target values. The frequency range is
        specified by the input motion.

        Parameters
        ----------
        target_type: str
            Options are 'crustal_amp' to only fit to the crustal amplification,
             or 'site_term' to fit both the velocity and the site attenuation
             parameter.
        target: `array_like`
            Target values.
        adjust_thickness: bool (optional)
            If the thickness of the layers is adjusted as well, default: False.
        adjust_site_atten: bool (optional)
            If the site attenuation is adjusted as well, default: False.
        adjust_source_vel: bool (optional)
            If the source velocity should be adjusted, default: False.
        Returns
        -------
        profile: `pyrsa.site.Profile`
            profile optimized to fit a target amplification.
        """
        density = self.profile.density

        nl = len(density)

        # Slowness bounds
        slowness = self.profile.slowness
        thickness = self.profile.thickness
        site_atten = self._site_atten

        # Slowness
        initial = slowness
        bounds = 1 / np.tile((4000, 100), (nl, 1))
        if not adjust_source_vel:
            bounds[-1] = (initial[-1], initial[-1])

        # Thickness bounds
        if adjust_thickness:
            bounds = np.r_[bounds, [[t / 2, 2 * t] for t in thickness]]
            initial = np.r_[initial, thickness]

        # Site attenuation bounds
        if adjust_site_atten:
            bounds = np.r_[bounds, [[0.0001, 0.200]]]
            initial = np.r_[initial, self.site_atten]

        def calc_rmse(this, that):
            return np.mean(((this - that) / that) ** 2)

        def err(x):
            _slowness = x[0:nl]
            if adjust_thickness:
                _thickness = x[nl:(2 * nl)]
            else:
                _thickness = thickness
            if adjust_site_atten:
                self._site_atten = x[-1]

            crustal_amp, site_term = self._calc_amp(density, _thickness,
                                                    _slowness)

            calc = crustal_amp if target_type == 'crustal_amp' else site_term

            err = 10 * calc_rmse(target, calc)
            # Prefer the original values so add the difference to the error
            err += calc_rmse(slowness, _slowness)
            if adjust_thickness:
                err += calc_rmse(thickness, _thickness)
            if adjust_site_atten:
                err += calc_rmse(self._site_atten, site_atten)
            return err

        res = minimize(err, initial, method='L-BFGS-B', bounds=bounds)

        slowness = res.x[0:nl]
        if adjust_thickness:
            thickness = res.x[nl:(2 * nl)]

        profile = Profile([
            Layer(l.soil_type, t, 1 / s)
            for l, t, s in zip(self.profile, thickness, slowness)
        ], self.profile.wt_depth)
        # Update the calculated amplificaiton
        self(self.motion, profile, self.loc_input)

class LinearElasticCalculator(AbstractCalculator):
    """Class for performing linear elastic site response."""

    def __init__(self):
        super().__init__()

        self._waves_a = np.array([])
        self._waves_b = np.array([])
        self._wave_nums = np.array([])

    def __call__(self, motion, profile, loc_input):
        """Perform the wave propagation.

        Parameters
        ----------
        motion: :class:`~.base.motion.Motion`
            Input motion.

        profile: :class:`~.base.site.Profile`
            Site profile.

        loc_input: :class:`~.base.site.Location`
            Location of the input motion.
        """
        super().__call__(motion, profile, loc_input)

        # Set initial properties
        for l in profile:
            l.reset()
            if l.strain is None:
                l.strain = 0.

        self._calc_waves(motion.angular_freqs, profile)

    def _calc_waves(self, angular_freqs, profile):
        """Compute the wave numbers and amplitudes (up- and down-going).

        Parameters
        ----------
        angular_freqs: :class:`numpy.ndarray`
            Angular frequency at which the waves are computed.

        profile: :class:`~.base.site.Profile`
            Site profile.
        """

        # Compute the complex wave numbers of the system
        wave_nums = np.empty((len(profile), len(angular_freqs)), np.complex)
        for i, l in enumerate(profile):
            wave_nums[i, :] = angular_freqs / l.comp_shear_vel

        # Compute the waves. In the top surface layer, the up-going and
        # down-going waves have an amplitude of 1 as they are completely
        # reflected at the surface.
        waves_a = np.ones_like(wave_nums, np.complex)
        waves_b = np.ones_like(wave_nums, np.complex)
        for i, l in enumerate(profile[:-1]):
            # Complex impedance -- wave number can be zero which causes an
            # error.
            with np.errstate(invalid='ignore'):
                cimped = ((wave_nums[i] * l.comp_shear_mod) /
                          (wave_nums[i + 1] * profile[i + 1].comp_shear_mod))

            # Complex term to simplify equations -- uses full layer height
            cterm = 1j * wave_nums[i, :] * l.thickness

            waves_a[i + 1, :] = (
                0.5 * waves_a[i] *
                (1 + cimped) * np.exp(cterm) + 0.5 * waves_b[i] *
                (1 - cimped) * np.exp(-cterm))
            waves_b[i + 1, :] = (
                0.5 * waves_a[i] *
                (1 - cimped) * np.exp(cterm) + 0.5 * waves_b[i] *
                (1 + cimped) * np.exp(-cterm))

            # Set wave amplitudes with zero frequency to 1
            mask = ~np.isfinite(cimped)
            waves_a[i + 1, mask] = 1.
            waves_b[i + 1, mask] = 1.

        # fixme: Better way to handle this?
        # Set wave amplitudes to 1 at frequencies near 0
        mask = np.isclose(angular_freqs, 0)
        waves_a[-1, mask] = 1.
        waves_b[-1, mask] = 1.

        self._waves_a = waves_a
        self._waves_b = waves_b
        self._wave_nums = wave_nums

    def wave_at_location(self, l):
        """Compute the wave field at specific location.

        Parameters
        ----------
        l : site.Location
            :class:`site.Location` of the input

        Returns
        -------
        `np.ndarray`
            Amplitude and phase of waves
        """
        cterm = 1j * self._wave_nums[l.index] * l.depth_within

        if l.wave_field == WaveField.within:
            return (self._waves_a[l.index] * np.exp(cterm) +
                    self._waves_b[l.index] * np.exp(-cterm))
        elif l.wave_field == WaveField.outcrop:
            return 2 * self._waves_a[l.index] * np.exp(cterm)
        elif l.wave_field == WaveField.incoming_only:
            return self._waves_a[l.index] * np.exp(cterm)
        else:
            raise NotImplementedError

    def calc_accel_tf(self, lin, lout):
        """Compute the acceleration transfer function.

        Parameters
        ----------
        lin : :class:`~site.Location`
            Location of input
        lout : :class:`~site.Location`
            Location of output. Note that this would typically be midheight
            of the layer.

        """
        tf = self.wave_at_location(lout) / self.wave_at_location(lin)
        return tf

    def calc_stress_tf(self, lin, lout, damped):
        """Compute the stress transfer function.

        Parameters
        ----------
        lin : :class:`~site.Location`
            Location of input
        lout : :class:`~site.Location`
            Location of output. Note that this would typically be midheight
            of the layer.

        """
        tf = self.calc_strain_tf(lin, lout)
        if damped:
            # Scale by complex shear modulus to include the influence of
            # damping
            tf *= lout.layer.comp_shear_mod
        else:
            tf *= lout.layer.shear_mod

        return tf

    def calc_strain_tf(self, lin, lout):
        """Compute the strain transfer function from `lout` to
        `location_in`.

        The strain transfer function from the acceleration at layer `n`
        (outcrop) to the mid-height of layer `m` (within) is defined as

        Parameters
        ----------
        lin : :class:`~site.Location`
            Location of input
        lout : :class:`~site.Location`
            Location of output. Note that this would typically be midheight
            of the layer.

        Returns
        -------
        strain_tf : :class:`numpy.ndarray`
            Transfer function to be applied to an acceleration FAS.
        """
        # FIXME: Correct discussion for using acceleration FAS
        # Strain(angFreq, z=h_m/2)
        # ------------------------ =
        #    accel_n(angFreq)
        #
        #          i k*_m [ A_m exp(i k*_m h_m / 2) - B_m exp(-i k*_m h_m / 2)]
        #          ------------------------------------------------------------
        #                         -angFreq^2 (2 * A_n)
        #
        assert lout.wave_field == WaveField.within

        ang_freqs = self.motion.angular_freqs
        # The numerator cannot be computed using wave_at_location() because
        # it is A - B.
        cterm = 1j * self._wave_nums[lout.index, :] * lout.depth_within
        numer = (1j * self._wave_nums[lout.index, :] *
                 (self._waves_a[lout.index, :] * np.exp(cterm) -
                  self._waves_b[lout.index, :] * np.exp(-cterm)))
        denom = -ang_freqs ** 2 * self.wave_at_location(lin)

        # Only compute transfer function for non-zero frequencies
        mask = ~np.isclose(ang_freqs, 0)
        tf = np.zeros_like(mask, dtype=np.complex)
        # Scale into units from gravity
        tf[mask] = GRAVITY * numer[mask] / denom[mask]

        return tf

class EquivalentLinearCalculator(LinearElasticCalculator):
    """Class for performing equivalent-linear elastic site response."""

    def __init__(self, strain_ratio=0.65, tolerance=0.01, max_iterations=15):
        """Initialize the class.

        Parameters
        ----------
        strain_ratio: float, default=0.65
            Ratio between the maximum strain and effective strain used to
            compute strain compatible properties.

        tolerance: float, default=0.01
            Tolerance in the iterative properties, which would cause the
            iterative process to terminate.

        max_iterations: int, default=15
            Maximum number of iterations to perform.
        """
        super().__init__()
        self._strain_ratio = strain_ratio
        self._tolerance = tolerance
        self._max_iterations = max_iterations

    def __call__(self, motion, profile, loc_input):
        """Perform the wave propagation.

        Parameters
        ----------
        motion: :class:`~.base.motion.Motion`
            Input motion.

        profile: :class:`~.base.site.Profile`
            Site profile.

        loc_input: :class:`~.base.site.Location`
            Location of the input motion.
        """
        super().__call__(motion, profile, loc_input)

        self._estimate_strains()

        iteration = 0
        while iteration < self.max_iterations:
            self._calc_waves(motion.angular_freqs, profile)
            for index, layer in enumerate(profile[:-1]):
                loc_layer = Location(index, layer, 'within',
                                     layer.thickness / 2)
                # Compute the representative strain(s) within the layer. FDM
                #  will provide a vector of strains.
                layer.strain = self._calc_strain(loc_input, loc_layer, motion)
            # Maximum error (damping and shear modulus) over all layers
            max_error = max(l.max_error for l in profile)
            if max_error < self.tolerance:
                break
            iteration += 1

        # Compute the maximum strain within the profile.
        for index, layer in enumerate(profile[:-1]):
            loc_layer = Location(index, layer, 'within', layer.thickness / 2)
            layer.strain_max = self._calc_strain_max(loc_input, loc_layer,
                                                     motion)

    def _estimate_strains(self):
        """Compute an estimate of the strains."""
        # Estimate the strain based on the PGV and shear-wave velocity
        for l in self._profile:
            l.reset()
            l.strain = self._motion.pgv / l.initial_shear_vel

    @property
    def strain_ratio(self):
        return self._strain_ratio

    @property
    def tolerance(self):
        return self._tolerance

    @property
    def max_iterations(self):
        return self._max_iterations

    @classmethod
    def calc_strain_ratio(cls, mag):
        """Compute the effective strain ratio using Idriss and Sun (1992).

        Parameters
        ----------
        mag: float
            Magnitude of the input motion.

        Returns
        -------
        strain_ratio : float
            Effective strain ratio

        References
        ----------
        .. [1] Idriss, I. M., & Sun, J. I. (1992). SHAKE91: A computer program
            for conducting equivalent linear seismic response analyses of
            horizontally layered soil deposits. Center for Geotechnical
            Modeling, Department of Civil and Environmental Engineering,
            University of California, Davis, CA.
        """
        return (mag - 1) / 10

    def _calc_strain(self, loc_input, loc_layer, motion, *args):
        """Compute the strain used for iterations of material properties."""
        strain_max = self._calc_strain_max(loc_input, loc_layer, motion, *args)
        return self.strain_ratio * strain_max

    def _calc_strain_max(self, loc_input, loc_layer, motion, *args):
        """Compute the effective strain at the center of a layer."""
        return motion.calc_peak(
            self.calc_strain_tf(loc_input, loc_layer))

class FrequencyDependentEqlCalculator(EquivalentLinearCalculator):
    """Class for performing equivalent-linear elastic site response with
    frequency-dependent modulii and damping.

    Parameters
    ----------
    use_smooth_spectrum: bool, default=False
        Use the Kausel & Assimaki (2002) smooth spectrum for the strain.
        Otherwise, the complete Fourier amplitude spectrum is used.
    strain_ratio: float, default=1.00
        ratio between the maximum strain and effective strain used to compute
        strain compatible properties. There is not clear guidance the use of
        the effective strain ratio. However, given the nature of the method,
        it would make sense not to include the an effective strain ratio.
    tolerance: float, default=0.01
        tolerance in the iterative properties, which would cause the iterative
        process to terminate.
    max_iterations: int, default=15
        maximum number of iterations to perform.

    References
    ----------
    .. [1] Kausel, E., & Assimaki, D. (2002). Seismic simulation of inelastic
        soils via frequency-dependent moduli and damping. Journal of
        Engineering Mechanics, 128(1), 34-47.
    """

    def __init__(self,
                 use_smooth_spectrum=False,
                 strain_ratio=1.0,
                 tolerance=0.01,
                 max_iterations=15):
        """Initialize the class."""
        super().__init__(strain_ratio, tolerance, max_iterations)

        # Use the smooth strain spectrum as proposed by Kausel and Assimaki
        self._use_smooth_spectrum = use_smooth_spectrum

    def _estimate_strains(self):
        """Estimate the strains by running an EQL site response.

        This step was recommended in Section 8.3.1 of Zalachoris (2014).
        """
        eql = EquivalentLinearCalculator()
        eql(self._motion, self._profile, self._loc_input)

    def _calc_strain(self, loc_input, loc_layer, motion, *args):
        freqs = np.array(motion.freqs)
        strain_tf = self.calc_strain_tf(loc_input, loc_layer)
        strain_fas = np.abs(strain_tf * motion.fourier_amps)
        # Maximum strain in the time domain modified by the effective strain
        # ratio
        strain_eff = self.strain_ratio * motion.calc_peak(strain_tf)

        if self._use_smooth_spectrum:
            # Equation (8)
            freq_avg = (np.trapz(freqs * strain_fas, x=freqs) /
                        np.trapz(strain_fas, x=freqs))

            # Find the average strain at frequencies less than the average
            # frequency
            # Equation (8)
            mask = (freqs < freq_avg)
            strain_avg = np.trapz(strain_fas[mask], x=freqs[mask]) / freq_avg

            # Normalize the frequency and strain by the average values
            freqs /= freq_avg
            strain_fas /= strain_avg

            # Fit the smoothed model at frequencies greater than the average
            A = np.c_[-freqs[~mask], -np.log(freqs[~mask])]
            a, b = np.linalg.lstsq(A, np.log(strain_fas[~mask]))[0]
            # This is a modification of the published method that ensures a
            # smooth transition in the strain
            shape = np.minimum(1, np.exp(-a * freqs) / np.power(freqs, b))
            strains = strain_eff * shape
        else:
            strains = strain_eff * strain_fas / np.max(strain_fas)

        return strains


#!/usr/bin/env python
# coding: utf-8
# This is a demo, given a LINEMOD image 0002.png and the mask-RCNN result(mask and 2D boundingbox), do 6D pose estimation for the lamp:
# First based on RGB, then use depth map to refine the z-direction, finally refine both rotation and translation.
#Prerequisite before generating the codebook:
#ckpt: under workspace_path/experiments/<experiment_name>/checkpoints_lambda250/checkpoints/ckpt-<num_iterations>-1
#Rendered imgs and edgemaps under reference rotations \bar_R(Generated by render_codebook.py) under the path: path_embedding_data
#image,depth map, mask
#mesh.ply under the path: path_model

import cv2
import os
import numpy as np
import tensorflow as tf
import sonnet as snt
import open3d as o3d
from pysixd_stuff.pysixd import inout
from est_utils import est_tra_w_tz,rectify_rot,depth_refinement,rotation_error_icp
obj_id=14
num_iterations=30000
LATENT_SPACE_SIZE = 128
NUM_FILTER = [128, 256, 512, 512]
KERNEL_SIZE_ENCODER = 5
STRIDES =[2, 2, 2, 2]
BATCH_NORM = False
image_size=128
embedding_dim = 128
num_embeddings = 92232
K_train =np.array([572.41140, 0, 325.26110, 0, 573.57043, 242.04899, 0, 0, 1]).reshape((3, 3))#Should be consistent with \bar_R images
Radius_render_train = 700

K_test = np.array([572.4114, 0.0, 325.2611, 0.0, 573.57043, 242.04899, 0.0, 0.0, 1.0]).reshape(3,3)
experiment_name='linemod_{:02d}_softmax_edge'.format(obj_id)
path_workspath='./ws/'
path_embedding_data='./embedding92232s/{:02d}'.format(obj_id)
path_model='./ws/meshes/obj_{:02d}.ply'.format(obj_id)
model_ply = inout.load_ply(path_model)
model_o3d = o3d.io.read_point_cloud(path_model)

#### Step 0: Load pose estimation network
class Encoder(snt.AbstractModule):
    def __init__(self, latent_space_size, num_filters, kernel_size, strides, batch_norm, name='encoder'):
        super(Encoder, self).__init__(name=name)
        self._latent_space_size = latent_space_size
        self._num_filters = num_filters
        self._kernel_size = kernel_size
        self._strides = strides
        self._batch_normalization = batch_norm

    @property
    def latent_space_size(self):
        return self._latent_space_size

    @property
    def encoder_layers(self):
        layers = []
        x = self._input
        layers.append(x)
        for filters, stride in zip(self._num_filters, self._strides):
            padding = 'same'
            x = tf.layers.conv2d(
                inputs=x,
                filters=filters,
                kernel_size=self._kernel_size,
                strides=stride,
                padding=padding,
                kernel_initializer=tf.contrib.layers.xavier_initializer_conv2d(),
                activation=tf.nn.relu,
            )
            if self._batch_normalization:
                x = tf.layers.batch_normalization(x, training=self._is_training)
            layers.append(x)
        return layers

    @property
    def encoder_out(self):
        x = self.encoder_layers[-1]
        encoder_out = tf.contrib.layers.flatten(x)

        return encoder_out

    @property
    def z(self):
        x = self.encoder_out
        # construct data
        z = tf.layers.dense(
            x,
            self._latent_space_size,
            activation=None,
            kernel_initializer=tf.contrib.layers.xavier_initializer(),
            name=None
        )
        return z

    def _build(self, x, is_training=False):
        self._input = x
        self._is_training = is_training
        return self.z

class VectorQuantizer(snt.AbstractModule):
    def __init__(self, embedding_dim, num_embeddings, name='vq_center'):
        super(VectorQuantizer, self).__init__(name=name)
        self._embedding_dim = embedding_dim
        self._num_embeddings = num_embeddings

        with self._enter_variable_scope():
            initializer = tf.uniform_unit_scaling_initializer()
            self._w = tf.get_variable('embedding', [embedding_dim, num_embeddings], initializer=initializer, trainable=True)

    def _build(self, inputs):
        input_shape = tf.shape(inputs)
        with tf.control_dependencies([
            tf.Assert(tf.equal(input_shape[-1], self._embedding_dim),[input_shape])]):
            flat_inputs = tf.reshape(inputs, [-1, self._embedding_dim])
            w = self.embeddings.read_value()

        distances = -tf.matmul(tf.nn.l2_normalize(flat_inputs, axis=1), tf.nn.l2_normalize(w, axis=0))
        encoding_indices = tf.argmax(- distances, 1)
        encoding_indices = tf.reshape(encoding_indices, tf.shape(inputs)[:-1])
        return {'encoding_indices': encoding_indices, }

    @property
    def embeddings(self):
        return self._w

# Build modules.
graph_estpose=tf.Graph()
with graph_estpose.as_default():
    with tf.variable_scope(experiment_name):  # .split('_')[0]+'_'+experiment_name.split('_')[1]):
        I_x = tf.placeholder(tf.float32, shape=(None, image_size, image_size, 4))
        with tf.variable_scope('encoder'):
            encoder = Encoder(latent_space_size=LATENT_SPACE_SIZE,num_filters=NUM_FILTER,kernel_size=KERNEL_SIZE_ENCODER,strides=STRIDES,batch_norm=BATCH_NORM)
        z = encoder(I_x)
        network_vars = tf.trainable_variables()
        print(network_vars)
        vq_codebook = VectorQuantizer(embedding_dim=embedding_dim,num_embeddings=num_embeddings)

        # For evaluation, make sure is_training=False!
        with tf.variable_scope('validation'):
            nn_item = vq_codebook(z)

        # Bounding box informations for foreground model in pose template repository
        codebook_obj_bbs = np.load(os.path.join(path_embedding_data,'obj_bbs.npy'))
        codebook_rotations = np.load(os.path.join(path_embedding_data,'rot_infos.npz'))['rots']

    saver = tf.train.Saver(network_vars, save_relative_paths=False)
    embedding = tf.placeholder(tf.float32, shape=[embedding_dim, num_embeddings])
    embedding_assign_op = tf.assign(vq_codebook.embeddings, embedding)

gpu_options = tf.GPUOptions(allow_growth=True, per_process_gpu_memory_fraction=0.9)
config = tf.ConfigProto(gpu_options=gpu_options)
sess_estpose=tf.Session(graph=graph_estpose,config=config)
print('Step 0, Load Rotation Estimation Net')
with sess_estpose.as_default():
    with graph_estpose.as_default():
        saver.restore(sess_estpose, '{:s}/experiments/{:s}/checkpoints_lambda250/checkpoints/chkpt-{:d}'.format(path_workspath,experiment_name,num_iterations-1))
        arr_codebook = np.load(os.path.join(path_embedding_data,'edgeLambda250_codebook.npy'))
        sess_estpose.run(embedding_assign_op, {embedding: arr_codebook.T})

print('Step 0, Load RGB image')
image = cv2.imread('./demo_data/0002.png')
img_depth = inout.load_depth2('./demo_data/0002_depth.png')
depth_scale=1.0
img_depth=img_depth*depth_scale

print('Step 1, Load 2D detection and mask')
img_masks = np.load('./demo_data/0002_mask.npy')
for cc in range(3):
    image[:, :, cc] = np.where(img_masks[:, :, 0] == 1,image[:, :, cc], 0)
img_depth = np.where(img_masks[:, :, 0] == 1, img_depth, 0)

obj_bb=np.array([326, 54, 103, 151])
#obj_bb is the 2D bounding box on the image, with 4 elements: x,y,w,h; where (x,y) is the location of the left-top corner
#Thus center 2D location is (x+w/2, h+w/2)

verbose=False
if verbose:
    image_vis=image.copy()
    cv2.rectangle(image_vis,(obj_bb[0],obj_bb[1]),(obj_bb[0]+obj_bb[2],obj_bb[1]+obj_bb[3]),(255,0,0),2)
    cv2.imshow('img',image_vis)
    cv2.waitKey()

print('Step 2, Pose estimation')
with sess_estpose.as_default():
    with graph_estpose.as_default():
        img_bgr=image.copy()

        x,y,w,h=obj_bb
        H,W,_=img_bgr.shape
        size = int(np.maximum(h, w) * 1.2)
        left = int(np.max([x + w / 2 - size / 2, 0]))
        right = int(np.min([x + w / 2 + size / 2, W]))
        top = int(np.max([y + h / 2 - size / 2, 0]))
        bottom = int(np.min([y + h / 2 + size / 2, H]))

        crop = img_bgr[top:bottom, left:right].copy()
        crop_depth=img_depth[top:bottom,left:right].copy()

        query_bgr = cv2.resize(crop, (image_size,image_size))
        query_edge = np.expand_dims(cv2.Canny(query_bgr, 50, 150),2)
        query = np.expand_dims((np.concatenate((query_bgr, query_edge), axis=-1) /255.),0)

        idx=sess_estpose.run([nn_item],feed_dict={I_x:query})
        idx=idx[0]['encoding_indices'][0]
        est_rot_cb= codebook_rotations[idx] #This should be a 3x3 matrix, which indicates an initial model-to-camera rotation estimation

        if verbose:
            #path='../Edge-Network/embedding92232s/{:02d}/imgs/{:05d}.png'.format(obj_id,idx)
            #est_bgr=cv2.imread(path)
            cv2.imshow('bbox',query_bgr)
            cv2.imshow('bbox_depth',crop_depth)
            #cv2.imshow('estimated rotation',est_bgr)
            cv2.waitKey()

        K00_ratio = K_test[0, 0] / K_train[0, 0]
        K11_ratio = K_test[1, 1] / K_train[1, 1]

        mean_K_ratio = np.mean([K00_ratio, K11_ratio])

        render_bb = codebook_obj_bbs[idx].squeeze()
        est_bb = obj_bb.copy() #Same as obj_bb, thus est_bb=[x,y,w,h], is the 2D bounding box detected, where (x,y) is the 2D location of left-top corner
        #Center of the 2D bbox is (x_c,y_c)=(est_bb[0]+est_bb[2]/2,est_bb[1]+est_bb[3]/2)
        diag_bb_ratio = np.linalg.norm(np.float32(render_bb[2:])) / np.linalg.norm(np.float32(est_bb[2:]))

        mm_tz = diag_bb_ratio * mean_K_ratio * Radius_render_train

        center_obj_x_train = render_bb[0] + render_bb[2] / 2. - K_train[0, 2]
        center_obj_y_train = render_bb[1] + render_bb[3] / 2. - K_train[1, 2]

        center_obj_x_test = est_bb[0] + est_bb[2] / 2 - K_test[0, 2]
        center_obj_y_test = est_bb[1] + est_bb[3] / 2 - K_test[1, 2]

        est_tra = est_tra_w_tz(mm_tz,Radius_render_train,K_test,center_obj_x_test,center_obj_y_test,K_train,center_obj_x_train,center_obj_y_train)
        est_rot=rectify_rot(est_rot_cb,est_tra)

        max_mean_dist_factor=2.0
        #Refine the z-direction only.
        mm_tz,max_mean_dist=depth_refinement(crop_depth, model_ply,est_rot.astype(np.float32),est_tra.flatten(), K_test, (W, H), max_mean_dist_factor=max_mean_dist_factor)#5.0)

        est_tra = est_tra_w_tz(mm_tz,Radius_render_train,K_test,center_obj_x_test,center_obj_y_test, K_train,center_obj_x_train,center_obj_y_train)
        est_rot=rectify_rot(est_rot_cb,est_tra)

        #Further icp refinement for rotation and translation
        if True:
            est_rot,est_tra,_= rotation_error_icp(img_depth, model_o3d,  obj_bb, est_rot, est_tra.flatten(),K_test.copy(),
                                                width=W, height=H, max_mean_dist=max_mean_dist,
                                                max_mean_dist_factor=max_mean_dist_factor,
                                                regist_error_threshold=3.0, fitness_threshold=0.7)

print('Step 3: Write result')
print('Estimated rotation - model2cam, 3x3 rotation matrix:') ###3x3 rotation matrix
print(est_rot.astype(np.float32).reshape((3,3)))
print('Estimated translation - model2cam, in mm:') ###translation vector, with unit mm
print(est_tra.astype(np.float32).reshape((1,3)))
print('Detected 2D bounding box - on 2D image, (x_topleft,y_topleft,bbox_width,bbox_height:')
print(np.array(obj_bb).astype(np.int32).reshape((1,4))) ### 2D bounding box

path_result='./result.txt'
with open(path_result, 'w') as f:
    line_tpl=', '.join(['{:.8f}'] * 9) + '\n' + ', '.join(['{:.8f}'] * 3)
    Rt = est_rot.astype(np.float32).flatten().tolist() + est_tra.flatten().tolist()
    txt=line_tpl.format(*Rt)
    f.write(txt)

'''
Expected result:
if refine for both rotation and translation:
est_rot: [-0.20446137, 0.94286263, -0.26306966, 0.85182256, 0.03896938, -0.52237886, -0.48227987, -0.33089498, -0.81111938],
est_tra: [90.49808699, -137.58854774, 849.52665827]
if refine only z-direction
est_rot: [-0.15618885, 0.96754020, -0.19867308, 0.86759037, 0.03824597, -0.49580660, -0.47211438, -0.24980631, -0.84540218],
est_tra: [85.44441029, -146.03545470, 840.30747686]
'''


"""Load transformations from URDF files.

See :doc:`transform_manager` for more information.
"""
import os
import numpy as np
from bs4 import BeautifulSoup
from .transform_manager import TransformManager
from .transformations import transform_from, concat
from .rotations import (
    active_matrix_from_extrinsic_roll_pitch_yaw, matrix_from_axis_angle,
    norm_vector)

class UrdfTransformManager(TransformManager):
    """Transformation manager that can load URDF files.

    URDF is the `Unified Robot Description Format <http://wiki.ros.org/urdf>`_.
    URDF allows to define joints between links that can be rotated about one
    axis. This transformation manager allows to set the joint angles after
    joints have been added or loaded from an URDF.

    .. warning::

        Note that this module requires the Python package beautifulsoup4.

    .. note::

        Joint angles must be given in radians.

    Parameters
    ----------
    strict_check : bool, optional (default: True)
        Raise a ValueError if the transformation matrix is not numerically
        close enough to a real transformation matrix. Otherwise we print a
        warning.

    check : bool, optional (default: True)
        Check if transformation matrices are valid and requested nodes exist,
        which might significantly slow down some operations.
    """
    def __init__(self, strict_check=True, check=True):
        super(UrdfTransformManager, self).__init__(strict_check, check)
        self._joints = {}
        self.collision_objects = []
        self.visuals = []

    def add_joint(self, joint_name, from_frame, to_frame, child2parent, axis,
                  limits=(float("-inf"), float("inf")), joint_type="revolute"):
        """Add joint.

        Parameters
        ----------
        joint_name : str
            Name of the joint

        from_frame : Hashable
            Child link of the joint

        to_frame : Hashable
            Parent link of the joint

        child2parent : array-like, shape (4, 4)
            Transformation from child to parent

        axis : array-like, shape (3,)
            Rotation axis of the joint (defined in the child frame)

        limits : pair of float, optional (default: (-inf, inf))
            Lower and upper joint angle limit

        joint_type : str, optional (default: 'revolute')
            Joint type: revolute or prismatic (continuous is the same as
            revolute)
        """
        self.add_transform(from_frame, to_frame, child2parent)
        self._joints[joint_name] = (
            from_frame, to_frame, child2parent, norm_vector(axis), limits,
            joint_type)

    def set_joint(self, joint_name, value):
        """Set joint position.

        Note that joint values are clipped to their limits.

        Parameters
        ----------
        joint_name : str
            Name of the joint

        value : float
            Joint angle in radians in case of revolute joints or position
            in case of prismatic joint.

        Raises
        ------
        KeyError
            If joint_name is unknown
        """
        if joint_name not in self._joints:
            raise KeyError("Joint '%s' is not known" % joint_name)
        from_frame, to_frame, child2parent, axis, limits, joint_type = \
            self._joints[joint_name]
        # this is way faster than np.clip:
        value = min(max(value, limits[0]), limits[1])
        if joint_type == "revolute":
            joint_rotation = matrix_from_axis_angle(
                np.hstack((axis, (value,))))
            joint2A = transform_from(
                joint_rotation, np.zeros(3), strict_check=self.strict_check)
        else:
            assert joint_type == "prismatic"
            joint_offset = value * axis
            joint2A = transform_from(
                np.eye(3), joint_offset, strict_check=self.strict_check)
        self.add_transform(from_frame, to_frame, concat(
            joint2A, child2parent, strict_check=self.strict_check,
            check=self.check))

    def get_joint_limits(self, joint_name):
        """Get limits of a joint.

        Parameters
        ----------
        joint_name : str
            Name of the joint

        Returns
        -------
        limits : pair of float
            Lower and upper joint angle limit

        Raises
        ------
        KeyError
            If joint_name is unknown
        """
        if joint_name not in self._joints:
            raise KeyError("Joint '%s' is not known" % joint_name)
        return self._joints[joint_name][4]

    def load_urdf(self, urdf_xml, mesh_path=None, package_dir=None):
        """Load URDF file into transformation manager.

        Parameters
        ----------
        urdf_xml : str
            Robot definition in URDF

        mesh_path : str, optional (default: None)
            Path in which we search for meshes that are defined in the URDF.
            Meshes will be ignored if it is set to None and no 'package_dir'
            is given.

        package_dir : str, optional (default: None)
            Some URDFs start file names with 'package://' to refer to the ROS
            package in which these files (textures, meshes) are located. This
            variable defines to which path this prefix will be resolved.
        """
        robot_name, links, joints = parse_urdf(
            urdf_xml, mesh_path, package_dir, self.strict_check)
        initialize_urdf_transform_manager(self, robot_name, links, joints)

    def plot_visuals(self, frame, ax=None, ax_s=1, wireframe=False,
                     convex_hull_of_mesh=True, alpha=0.3):  # pragma: no cover
        """Plot all visuals in a given reference frame.

        Visuals can be boxes, spheres, cylinders, or meshes. Note that visuals
        that cannot be connected to the reference frame are omitted.

        Parameters
        ----------
        frame : Hashable
            Reference frame

        ax : Matplotlib 3d axis, optional (default: None)
            If the axis is None, a new 3d axis will be created

        ax_s : float, optional (default: 1)
            Scaling of the new matplotlib 3d axis

        wireframe : bool, optional (default: False)
            Plot wireframe (surface otherwise)

        convex_hull_of_mesh : bool, optional (default: True)
            Displays convex hull of meshes instead of the original mesh. This
            makes plotting a lot faster with complex meshes.

        alpha : float, optional (default: 0.3)
            Alpha value of the surface / wireframe that will be plotted

        Returns
        -------
        ax : Matplotlib 3d axis
            New or old axis
        """
        return self._plot_objects(
            self.visuals, frame, ax, ax_s, wireframe, convex_hull_of_mesh,
            alpha)

    def plot_collision_objects(
            self, frame, ax=None, ax_s=1, wireframe=True,
            convex_hull_of_mesh=True, alpha=1.0):  # pragma: no cover
        """Plot all collision objects in a given reference frame.

        Collision objects can be boxes, spheres, cylinders, or meshes. Note
        that collision objects that cannot be connected to the reference frame
        are omitted.

        Parameters
        ----------
        frame : Hashable
            Reference frame

        ax : Matplotlib 3d axis, optional (default: None)
            If the axis is None, a new 3d axis will be created

        ax_s : float, optional (default: 1)
            Scaling of the new matplotlib 3d axis

        wireframe : bool, optional (default: True)
            Plot wireframe (surface otherwise)

        convex_hull_of_mesh : bool, optional (default: True)
            Displays convex hull of meshes instead of the original mesh. This
            makes plotting a lot faster with complex meshes.

        alpha : float, optional (default: 1)
            Alpha value of the surface / wireframe that will be plotted

        Returns
        -------
        ax : Matplotlib 3d axis
            New or old axis
        """
        return self._plot_objects(
            self.collision_objects, frame, ax, ax_s, wireframe,
            convex_hull_of_mesh, alpha)

    def _plot_objects(self, objects, frame, ax=None, ax_s=1, wireframe=True,
                      convex_hull_of_mesh=True, alpha=1.0):  # pragma: no cover
        """Plot all objects in a given reference frame.

        Objects can be boxes, spheres, cylinders, or meshes. Note that objects
        that cannot be connected to the reference frame are omitted.

        Parameters
        ----------
        objects : list
            Objects that will be plotted

        frame : Hashable
            Reference frame

        ax : Matplotlib 3d axis, optional (default: None)
            If the axis is None, a new 3d axis will be created

        ax_s : float, optional (default: 1)
            Scaling of the new matplotlib 3d axis

        wireframe : bool, optional (default: True)
            Plot wireframe (surface otherwise)

        convex_hull_of_mesh : bool, optional (default: True)
            Displays convex hull of meshes instead of the original mesh. This
            makes plotting a lot faster with complex meshes.

        alpha : float, optional (default: 1)
            Alpha value of the surface / wireframe that will be plotted

        Returns
        -------
        ax : Matplotlib 3d axis
            New or old axis
        """
        if ax is None:
            from .plot_utils import make_3d_axis
            ax = make_3d_axis(ax_s)
        for obj in objects:
            ax = obj.plot(
                self, frame, ax, wireframe=wireframe,
                convex_hull=convex_hull_of_mesh, alpha=alpha)
        return ax

def parse_urdf(urdf_xml, mesh_path=None, package_dir=None, strict_check=True):
    """Parse information from URDF file.

    Parameters
    ----------
    urdf_xml : str
        Robot definition in URDF

    mesh_path : str, optional (default: None)
        Path in which we search for meshes that are defined in the URDF.
        Meshes will be ignored if it is set to None and no 'package_dir'
        is given.

    package_dir : str, optional (default: None)
        Some URDFs start file names with 'package://' to refer to the ROS
        package in which these files (textures, meshes) are located. This
        variable defines to which path this prefix will be resolved.

    strict_check : bool, optional (default: True)
        Raise a ValueError if the transformation matrix is not numerically
        close enough to a real transformation matrix. Otherwise we print a
        warning.

    Returns
    -------
    robot_name : str
        Name of the robot

    links : list of Link
        Links of the robot

    joints : list of Joint
        Joints of the robot

    Raises
    ------
    UrdfException
        If URDF is not valid
    """
    urdf = BeautifulSoup(urdf_xml, "xml")

    # URDF XML schema:
    # https://github.com/ros/urdfdom/blob/master/xsd/urdf.xsd

    robot = urdf.find("robot")
    if robot is None:
        raise UrdfException("Robot tag is missing.")

    if not robot.has_attr("name"):
        raise UrdfException("Attribute 'name' is missing in robot tag.")

    robot_name = robot["name"]

    materials = dict([
        _parse_material(material)
        for material in robot.findAll("material", recursive=False)])

    links = [_parse_link(link, materials, mesh_path, package_dir, strict_check)
             for link in robot.findAll("link", recursive=False)]

    link_names = [link.name for link in links]
    joints = [_parse_joint(joint, link_names, strict_check)
              for joint in robot.findAll("joint", recursive=False)]

    return robot_name, links, joints

def initialize_urdf_transform_manager(tm, robot_name, links, joints):
    """Initializes transform manager from previously parsed URDF data.

    Parameters
    ----------
    tm : UrdfTransformManager
        Transform manager

    robot_name : str
        Name of the robot

    links : list of Link
        Links of the robot

    joints : list of Joint
        Joints of the robot
    """
    tm.add_transform(links[0].name, robot_name, np.eye(4))
    _add_links(tm, links)
    _add_joints(tm, joints)

def _parse_material(material):
    """Parse material."""
    if not material.has_attr("name"):
        raise UrdfException("Material name is missing.")
    colors = material.findAll("color")
    if len(colors) not in [0, 1]:
        raise UrdfException("More than one color is not allowed.")
    if len(colors) == 1:
        color = _parse_color(colors[0])
    else:
        color = None
    # TODO texture is currently ignored
    return material["name"], color

def _parse_color(color):
    """Parse color."""
    if not color.has_attr("rgba"):
        raise UrdfException("Attribute 'rgba' of color tag is missing.")
    return np.fromstring(color["rgba"], sep=" ")

def _parse_link(link, materials, mesh_path, package_dir, strict_check):
    """Create link."""
    if not link.has_attr("name"):
        raise UrdfException("Link name is missing.")

    result = Link()
    result.name = link["name"]

    visuals, visual_transforms = _parse_link_children(
        link, "visual", materials, mesh_path, package_dir, strict_check)
    result.visuals = visuals
    result.transforms.extend(visual_transforms)

    collision_objects, collision_object_transforms = _parse_link_children(
        link, "collision", dict(), mesh_path, package_dir, strict_check)
    result.collision_objects = collision_objects
    result.transforms.extend(collision_object_transforms)

    inertial = link.find("inertial")
    if inertial is not None:
        result.inertial_frame[:, :] = _parse_origin(inertial, strict_check)
        result.mass = _parse_mass(inertial)
        result.inertia[:, :] = _parse_inertia(inertial)
        result.transforms.append(
            ("inertial_frame:%s" % result.name, result.name,
             result.inertial_frame))

    return result

def _parse_link_children(link, child_type, materials, mesh_path, package_dir,
                         strict_check):
    """Parse collision objects or visuals."""
    children = link.findAll(child_type)
    shape_objects = []
    transforms = []
    for i, child in enumerate(children):
        if child.has_attr("name"):
            name = "%s:%s/%s" % (child_type, link["name"], child["name"])
        else:
            name = "%s:%s/%s" % (child_type, link["name"], i)

        color = None
        if child_type == "visual":
            material = child.find("material")
            if material is not None:
                material_name, color = _parse_material(material)
                if color is None and material_name in materials:
                    color = materials[material_name]

        child2link = _parse_origin(child, strict_check)
        transforms.append((name, link["name"], child2link))

        shape_objects.extend(_parse_geometry(
            child, name, color, mesh_path, package_dir))
    return shape_objects, transforms

def _parse_geometry(child, name, color, mesh_path, package_dir):
    """Parse geometric primitives (box, cylinder, sphere) or meshes."""
    geometry = child.find("geometry")
    if geometry is None:
        raise UrdfException("Missing geometry tag in link '%s'" % name)
    result = []
    for shape_type in ["box", "cylinder", "sphere", "mesh"]:
        shapes = geometry.findAll(shape_type)
        Cls = shape_classes[shape_type]
        for shape in shapes:
            shape_object = Cls(
                name, mesh_path=mesh_path, package_dir=package_dir,
                color=color)
            shape_object.parse(shape)
            result.append(shape_object)
    return result

def _parse_origin(entry, strict_check):
    """Parse transformation."""
    origin = entry.find("origin")
    translation = np.zeros(3)
    rotation = np.eye(3)
    if origin is not None:
        if origin.has_attr("xyz"):
            translation = np.fromstring(origin["xyz"], sep=" ")
        if origin.has_attr("rpy"):
            roll_pitch_yaw = np.fromstring(origin["rpy"], sep=" ")
            # URDF and KDL use the active convention for rotation matrices.
            # For more details on how the URDF parser handles the
            # conversion from Euler angles, see this blog post:
            # https://orbitalstation.wordpress.com/tag/quaternion/
            rotation = active_matrix_from_extrinsic_roll_pitch_yaw(
                roll_pitch_yaw)
    return transform_from(
        rotation, translation, strict_check=strict_check)

def _parse_mass(inertial):
    """Parse link mass."""
    mass = inertial.find("mass")
    if mass is not None and mass.has_attr("value"):
        result = float(mass["value"])
    else:
        result = 0.0
    return result

def _parse_inertia(inertial):
    """Parse inertia matrix."""
    inertia = inertial.find("inertia")

    result = np.zeros((3, 3))
    if inertia is None:
        return result

    if inertia.has_attr("ixx"):
        result[0, 0] = float(inertia["ixx"])
    if inertia.has_attr("ixy"):
        ixy = float(inertia["ixy"])
        result[0, 1] = ixy
        result[1, 0] = ixy
    if inertia.has_attr("ixz"):
        ixz = float(inertia["ixz"])
        result[0, 2] = ixz
        result[2, 0] = ixz
    if inertia.has_attr("iyy"):
        result[1, 1] = float(inertia["iyy"])
    if inertia.has_attr("iyz"):
        iyz = float(inertia["iyz"])
        result[1, 2] = iyz
        result[2, 1] = iyz
    if inertia.has_attr("izz"):
        result[2, 2] = float(inertia["izz"])
    return result

def _parse_joint(joint, link_names, strict_check):
    """Create joint object."""
    j = Joint()

    if not joint.has_attr("name"):
        raise UrdfException("Joint name is missing.")
    j.joint_name = joint["name"]

    if not joint.has_attr("type"):
        raise UrdfException("Joint type is missing in joint '%s'."
                            % j.joint_name)

    parent = joint.find("parent")
    if parent is None:
        raise UrdfException("No parent specified in joint '%s'"
                            % j.joint_name)
    if not parent.has_attr("link"):
        raise UrdfException("No parent link name given in joint '%s'."
                            % j.joint_name)
    j.parent = parent["link"]
    if j.parent not in link_names:
        raise UrdfException("Parent link '%s' of joint '%s' is not "
                            "defined." % (j.parent, j.joint_name))

    child = joint.find("child")
    if child is None:
        raise UrdfException("No child specified in joint '%s'"
                            % j.joint_name)
    if not child.has_attr("link"):
        raise UrdfException("No child link name given in joint '%s'."
                            % j.joint_name)
    j.child = child["link"]
    if j.child not in link_names:
        raise UrdfException("Child link '%s' of joint '%s' is not "
                            "defined." % (j.child, j.joint_name))

    j.joint_type = joint["type"]

    if j.joint_type in ["planar", "floating"]:
        raise UrdfException("Unsupported joint type '%s'" % j.joint_type)
    if j.joint_type not in ["revolute", "continuous", "prismatic",
                            "fixed"]:
        raise UrdfException("Joint type '%s' is not allowed in a URDF "
                            "document." % j.joint_type)

    j.child2parent = _parse_origin(joint, strict_check)

    j.joint_axis = np.array([1, 0, 0])
    if j.joint_type in ["revolute", "continuous", "prismatic"]:
        axis = joint.find("axis")
        if axis is not None and axis.has_attr("xyz"):
            j.joint_axis = np.fromstring(axis["xyz"], sep=" ")

    j.limits = _parse_limits(joint)
    return j

def _parse_limits(joint):
    """Parse joint limits."""
    limit = joint.find("limit")
    lower, upper = float("-inf"), float("inf")
    if limit is not None:
        if limit.has_attr("lower"):
            lower = float(limit["lower"])
        if limit.has_attr("upper"):
            upper = float(limit["upper"])
    return lower, upper

def _add_links(tm, links):
    """Add previously parsed links.

    Parameters
    ----------
    tm : UrdfTransformManager
        Transform manager

    links : list of Link
        Joint information from URDF
    """
    for link in links:
        tm.visuals.extend(link.visuals)
        tm.collision_objects.extend(link.collision_objects)

        for from_frame, to_frame, transform in link.transforms:
            tm.add_transform(from_frame, to_frame, transform)

def _add_joints(tm, joints):
    """Add previously parsed joints.

    Parameters
    ----------
    tm : UrdfTransformManager
        Transform manager

    joints : list of Joint
        Joint information from URDF
    """
    for joint in joints:
        if joint.joint_type in ["revolute", "continuous"]:
            tm.add_joint(
                joint.joint_name, joint.child, joint.parent,
                joint.child2parent, joint.joint_axis, joint.limits,
                "revolute")
        elif joint.joint_type == "prismatic":
            tm.add_joint(
                joint.joint_name, joint.child, joint.parent,
                joint.child2parent, joint.joint_axis, joint.limits,
                "prismatic")
        else:
            assert joint.joint_type == "fixed"
            tm.add_transform(joint.child, joint.parent, joint.child2parent)

class Link(object):
    """Link from URDF file.

    This class is only required temporarily while we parse the URDF.

    Attributes
    ----------
    name : str
        Link name

    visuals : list of Geometry
        Visual geometries

    collision_objects : list of Geometry
        Geometries for collision calculation

    transforms : list
        Transformations given as tuples: name of frame A, name of frame B,
        transform A2B

    inertial_frame : array, shape (4, 4)
        Pose of inertial frame with respect to the link

    mass : float
        Mass of the link

    inertia : array, shape (3, 3)
        Inertia matrix
    """
    def __init__(self):
        self.name = None
        self.visuals = []
        self.collision_objects = []
        self.transforms = []
        self.inertial_frame = np.eye(4)
        self.mass = 0.0
        self.inertia = np.zeros((3, 3))

class Joint(object):
    """Joint from URDF file.

    This class is only required temporarily while we parse the URDF.

    Attributes
    ----------
    child : str
        Name of the child

    parent : str
        Name of the parent frame

    child2parent : array-like, shape (4, 4)
        Transformation from child to parent

    joint_name : str
        Name of the joint that defines the transformation

    joint_axis : array-like, shape (3,)
        Rotation axis of the joint (defined in the child frame)

    joint_type : str
        Either 'fixed' or 'revolute'

    limits : pair of float
        Lower and upper joint angle limit
    """
    def __init__(self):
        self.child = None
        self.parent = None
        self.child2parent = np.eye(4)
        self.joint_name = None
        self.joint_axis = None
        self.joint_type = "fixed"
        self.limits = float("-inf"), float("inf")

class Geometry(object):
    """Geometrical object."""
    def __init__(self, frame, mesh_path, package_dir, color):
        self.frame = frame
        self.mesh_path = mesh_path
        self.package_dir = package_dir
        self.color = color

    def parse(self, xml):
        """Parse parameters of geometry."""

    def plot(self, tm, frame, ax=None, alpha=0.3, wireframe=True,
             convex_hull=True):
        """Plot geometry."""

class Box(Geometry):
    """Geometrical object: box."""
    def __init__(self, frame, mesh_path, package_dir, color):
        super(Box, self).__init__(frame, mesh_path, package_dir, color)
        self.size = np.zeros(3)

    def parse(self, xml):
        """Parse box size."""
        if xml.has_attr("size"):
            self.size[:] = np.fromstring(xml["size"], sep=" ")

    def plot(self, tm, frame, ax=None, alpha=0.3, wireframe=True,
             convex_hull=True):  # pragma: no cover
        """Plot box."""
        A2B = tm.get_transform(self.frame, frame)
        color = self.color if self.color is not None else "k"
        from .plot_utils import plot_box
        return plot_box(
            ax, self.size, A2B, wireframe=wireframe, alpha=alpha, color=color)

class Sphere(Geometry):
    """Geometrical object: sphere."""
    def __init__(self, frame, mesh_path, package_dir, color):
        super(Sphere, self).__init__(frame, mesh_path, package_dir, color)
        self.radius = 0.0

    def parse(self, xml):
        """Parse sphere radius."""
        if not xml.has_attr("radius"):
            raise UrdfException("Sphere has no radius.")
        self.radius = float(xml["radius"])

    def plot(self, tm, frame, ax=None, alpha=0.3, wireframe=True,
             convex_hull=True):  # pragma: no cover
        """Plot sphere."""
        center = tm.get_transform(self.frame, frame)[:3, 3]
        color = self.color if self.color is not None else "k"
        from .plot_utils import plot_sphere
        return plot_sphere(
            ax, self.radius, center, wireframe=wireframe, alpha=alpha,
            color=color)

class Cylinder(Geometry):
    """Geometrical object: cylinder."""
    def __init__(self, frame, mesh_path, package_dir, color):
        super(Cylinder, self).__init__(frame, mesh_path, package_dir, color)
        self.radius = 0.0
        self.length = 0.0

    def parse(self, xml):
        """Parse cylinder radius and length."""
        if not xml.has_attr("radius"):
            raise UrdfException("Cylinder has no radius.")
        self.radius = float(xml["radius"])
        if not xml.has_attr("length"):
            raise UrdfException("Cylinder has no length.")
        self.length = float(xml["length"])

    def plot(self, tm, frame, ax=None, alpha=0.3, wireframe=True,
             convex_hull=True):  # pragma: no cover
        """Plot cylinder."""
        A2B = tm.get_transform(self.frame, frame)
        color = self.color if self.color is not None else "k"
        from .plot_utils import plot_cylinder
        return plot_cylinder(
            ax, self.length, self.radius, 0.0, A2B, wireframe=wireframe,
            alpha=alpha, color=color)

class Mesh(Geometry):
    """Geometrical object: mesh."""
    def __init__(self, frame, mesh_path, package_dir, color):
        super(Mesh, self).__init__(frame, mesh_path, package_dir, color)
        self.filename = None
        self.scale = np.ones(3)

    def parse(self, xml):
        """Parse mesh filename and scale."""
        if self.mesh_path is None and self.package_dir is None:
            self.filename = None
        else:
            if not xml.has_attr("filename"):
                raise UrdfException("Mesh has no filename.")
            if self.mesh_path is not None:
                self.filename = os.path.join(self.mesh_path, xml["filename"])
            else:
                assert self.package_dir is not None
                self.filename = xml["filename"].replace(
                    "package://", self.package_dir)
            if xml.has_attr("scale"):
                self.scale = np.fromstring(xml["scale"], sep=" ")

    def plot(self, tm, frame, ax=None, alpha=0.3, wireframe=True,
             convex_hull=True):  # pragma: no cover
        """Plot mesh."""
        from .plot_utils import plot_mesh
        A2B = tm.get_transform(self.frame, frame)
        color = self.color if self.color is not None else "k"
        return plot_mesh(
            ax, self.filename, A2B, self.scale, wireframe=wireframe,
            convex_hull=convex_hull, alpha=alpha, color=color)

shape_classes = {"box": Box,
                 "sphere": Sphere,
                 "cylinder": Cylinder,
                 "mesh": Mesh}

class UrdfException(Exception):
    """Exception while parsing URDF files."""


import os

import numpy as np

import torch
from torch import nn
from boxlist import cat_boxlist, boxlist_iou

import math

INF = 100000000

class SigmoidFocalLoss(nn.Module):
    def __init__(self, gamma, alpha):
        super().__init__()

        self.gamma = gamma
        self.alpha = alpha
        self.eps = 1e-6

    def forward(self, out, target):
        n_class = out.shape[1]
        class_ids = torch.arange(
            1, n_class + 1, dtype=target.dtype, device=target.device
        ).unsqueeze(0)

        t = target.unsqueeze(1)
        p = torch.sigmoid(out)
        p = torch.clamp(p, self.eps, 1-self.eps) # for numerical stability

        gamma = self.gamma
        alpha = self.alpha

        term1 = (1 - p) ** gamma * torch.log(p)
        term2 = p ** gamma * torch.log(1 - p)

        # print(term1.sum(), term2.sum())

        loss = (
            -(t == class_ids).float() * alpha * term1
            - ((t != class_ids) * (t >= 0)).float() * (1 - alpha) * term2
        )

        return loss.sum()

def get_num_gpus():
    return int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1

def reduce_sum(tensor):
    if get_num_gpus() <= 1:
        return tensor
    import torch.distributed as dist
    tensor = tensor.clone()
    dist.all_reduce(tensor, op=dist.reduce_op.SUM)
    return tensor

def cat(tensors, dim=0):
    """
    Efficient version of torch.cat that avoids a copy if there is only a single element in a list
    """
    assert isinstance(tensors, (list, tuple))
    if len(tensors) == 1:
        return tensors[0]
    return torch.cat(tensors, dim)

def permute_and_flatten(layer, N, A, C, H, W):
    layer = layer.view(N, -1, C, H, W)
    layer = layer.permute(0, 3, 4, 1, 2)
    layer = layer.reshape(N, -1, C)
    return layer

def concat_box_prediction_layers(pred_cls, pred_reg):
    pred_cls_flattened = []
    pred_reg_flattened = []
    # for each feature level, permute the outputs to make them be in the
    # same format as the labels. Note that the labels are computed for
    # all feature levels concatenated, so we keep the same representation
    # for the objectness and the pred_reg
    for pred_cls_per_level, pred_reg_per_level in zip(pred_cls, pred_reg):
        N, AxC, H, W = pred_cls_per_level.shape
        Cx16 = pred_reg_per_level.shape[1]
        C = Cx16 // 16
        A = 1
        pred_cls_per_level = permute_and_flatten(
            pred_cls_per_level, N, A, C, H, W
        )
        pred_cls_flattened.append(pred_cls_per_level)

        pred_reg_per_level = permute_and_flatten(
            pred_reg_per_level, N, A, (C*16), H, W
        )
        pred_reg_flattened.append(pred_reg_per_level)
    # concatenate on the first dimension (representing the feature levels), to
    # take into account the way the labels were generated (with all feature maps
    # being concatenated as well)
    pred_cls = cat(pred_cls_flattened, dim=1).reshape(-1, C)
    pred_reg = cat(pred_reg_flattened, dim=1).reshape(-1, C*16)
    return pred_cls, pred_reg

class PoseLoss(object):
    def __init__(self, gamma, alpha, anchor_sizes, anchor_strides, positive_num, positive_lambda,
                    loss_weight_cls, loss_weight_reg, internal_K, diameters, target_coder):
        self.cls_loss_func = SigmoidFocalLoss(gamma, alpha)
        # self.centerness_loss_func = nn.BCEWithLogitsLoss(reduction="sum")
        # self.matcher = Matcher(fg_iou_threshold, bg_iou_threshold, True)
        self.anchor_sizes = anchor_sizes
        self.anchor_strides = anchor_strides
        self.positive_num = positive_num
        self.positive_lambda = positive_lambda
        self.loss_weight_cls = loss_weight_cls
        self.loss_weight_reg = loss_weight_reg
        self.internal_K = internal_K
        self.target_coder = target_coder
        self.diameters = diameters

    def ObjectSpaceLoss(self, pred, target_3D_in_camera_frame, cls_labels, anchors, weight=None):
        if not isinstance(self.diameters, torch.Tensor):
            self.diameters = torch.FloatTensor(self.diameters).to(device=pred.device).view(-1)

        diameter_ext = self.diameters[cls_labels.view(-1,1).repeat(1, 8*3).view(-1, 3, 1)]

        cellNum = pred.shape[0]
        pred_filtered = pred.view(cellNum, -1, 16)[torch.arange(cellNum), cls_labels]

        pred_xy = self.target_coder.decode(pred_filtered, anchors)
        # target_xy = self.target_coder.decode(target, anchors)

        pred_xy = pred_xy.view(-1,2,8).transpose(1,2).contiguous().view(-1,2)

        # construct normalized 2d
        B = torch.inverse(self.internal_K).mm(torch.cat((pred_xy.t(), torch.ones_like(pred_xy[:,0]).view(1,-1)), dim=0)).t()
        # compute projection matrices
        P = torch.bmm(B.view(-1, 3, 1), B.view(-1, 1, 3)) / torch.bmm(B.view(-1, 1, 3), B.view(-1, 3, 1))

        target_3D_in_camera_frame = target_3D_in_camera_frame.view(-1, 3, 1)
        px = torch.bmm(P, target_3D_in_camera_frame)

        target_3D_in_camera_frame = target_3D_in_camera_frame / diameter_ext
        px = px / diameter_ext
        scaling_factor = 50 # 0.02d
        losses = nn.SmoothL1Loss(reduction='none')(scaling_factor * px, scaling_factor * target_3D_in_camera_frame).view(cellNum, -1).mean(dim=1)
        losses = losses / scaling_factor

        if weight is not None and weight.sum() > 0:
            return (losses * weight).sum()
        else:
            assert losses.numel() != 0
            return losses.sum()

    def prepare_targets(self, targets, anchors):
        cls_labels = []
        reg_targets = []
        aux_raw_boxes = []
        aux_3D_in_camera_frame = []
        level_cnt = len(anchors[0])
        for im_i in range(len(targets)):
            pose_targets_per_im = targets[im_i]
            bbox_targets_per_im = pose_targets_per_im.to_object_boxlist()
            assert bbox_targets_per_im.mode == "xyxy"
            bboxes_per_im = bbox_targets_per_im.bbox
            labels_per_im = pose_targets_per_im.class_ids + 1
            anchors_per_im = cat_boxlist(anchors[im_i])
            num_gt = bboxes_per_im.shape[0]
            assert(level_cnt == len(anchors[im_i]))
            #
            rotations_per_im = pose_targets_per_im.rotations
            translations_per_im = pose_targets_per_im.translations
            mask_per_im = pose_targets_per_im.mask

            #
            anchor_sizes_per_level_interest = self.anchor_sizes[:level_cnt]
            anchor_strides_per_level_interst = self.anchor_strides[:level_cnt]
            gt_object_sizes = bbox_targets_per_im.box_span()

            num_anchors_per_level = [len(anchors_per_level.bbox) for anchors_per_level in anchors[im_i]]

            anchors_cx_per_im = (anchors_per_im.bbox[:, 2] + anchors_per_im.bbox[:, 0]) / 2.0
            anchors_cy_per_im = (anchors_per_im.bbox[:, 3] + anchors_per_im.bbox[:, 1]) / 2.0
            anchors_cx_per_im = torch.clamp(anchors_cx_per_im, min = 0, max = mask_per_im.shape[1] - 1).long()
            anchors_cy_per_im = torch.clamp(anchors_cy_per_im, min = 0, max = mask_per_im.shape[0] - 1).long()

            mask_at_anchors = mask_per_im[anchors_cy_per_im, anchors_cx_per_im]
            mask_labels = []
            for gt_i in range(num_gt):
                valid_mask = (mask_at_anchors == (gt_i+1))
                mask_labels.append(valid_mask)
            mask_labels = torch.stack(mask_labels).t()
            mask_labels = mask_labels.long()

            # random selecting candidates from each level first
            candidate_idxs = [[] for i in range(num_gt)]
            start_idx = 0
            gt_sz = gt_object_sizes.view(1,-1).repeat(level_cnt,1)
            lv_sz = torch.FloatTensor(anchor_sizes_per_level_interest).type_as(gt_sz)
            lv_sz = lv_sz.view(-1,1).repeat(1,num_gt)
            dk = torch.log2(gt_sz/lv_sz).abs()
            nk = torch.exp(-self.positive_lambda * (dk * dk))
            nk = self.positive_num * nk / nk.sum(0, keepdim=True)
            nk = (nk + 0.5).int()
            for level in range(level_cnt):
                end_idx = start_idx + num_anchors_per_level[level]
                is_in_mask_per_level = mask_labels[start_idx:end_idx, :]
                #
                for gt_i in range(num_gt):
                    posi_num = nk[level][gt_i]

                    valid_pos = is_in_mask_per_level[:, gt_i].nonzero().view(-1)
                    posi_num = min(posi_num, len(valid_pos))
                    # rand_idx = torch.randint(0, len(valid_pos), (int(posi_num),)) # randoms with replacement
                    rand_idx = torch.randperm(len(valid_pos))[:posi_num] # randoms without replacement
                    candi_pos = valid_pos[rand_idx] + start_idx
                    candidate_idxs[gt_i].append(candi_pos)
                #
                start_idx = end_idx

            # flagging selected positions
            roi = torch.full_like(mask_labels, -INF)
            for gt_i in range(num_gt):
                tmp_idx = torch.cat(candidate_idxs[gt_i], dim=0)
                roi[tmp_idx, gt_i] = 1

            anchors_to_gt_values, anchors_to_gt_indexs = roi.max(dim=1)
            cls_labels_per_im = labels_per_im[anchors_to_gt_indexs]
            cls_labels_per_im[anchors_to_gt_values == -INF] = 0 # background setting

            mask_visibilities, _ = mask_labels.max(dim=1)
            # logical_and, introduced only after pytorch 1.5
            # ignored_indexs = torch.logical_and(mask_visibilities==1, cls_labels_per_im==0)
            ignored_indexs = (mask_visibilities == 1) * (cls_labels_per_im == 0)
            cls_labels_per_im[ignored_indexs] = -1 # positions within mask but not selected will not be touched

            #
            matched_boxes = bboxes_per_im[anchors_to_gt_indexs]
            matched_classes = (labels_per_im - 1)[anchors_to_gt_indexs]
            matched_rotations = rotations_per_im[anchors_to_gt_indexs]
            matched_translations = translations_per_im[anchors_to_gt_indexs]

            #
            matched_3Ds = pose_targets_per_im.keypoints_3d[matched_classes]
            # matched_Ks = pose_targets_per_im.K.repeat(matched_classes.shape[0], 1, 1)
            # TODO
            # assert equals self K
            if not isinstance(self.internal_K, torch.Tensor):
                self.internal_K = torch.FloatTensor(self.internal_K).to(device=matched_3Ds.device).view(3, 3)
            reg_targets_per_im = self.target_coder.encode(
                self.internal_K, matched_3Ds,
                matched_rotations, matched_translations,
                anchors_per_im.bbox)
            cls_labels.append(cls_labels_per_im)
            reg_targets.append(reg_targets_per_im)
            aux_raw_boxes.append(matched_boxes)
            matched_3D_in_camera_frame = torch.bmm(matched_rotations, matched_3Ds.transpose(1, 2)) + matched_translations
            aux_3D_in_camera_frame.append(matched_3D_in_camera_frame.transpose(1, 2))

        return cls_labels, reg_targets, aux_raw_boxes, aux_3D_in_camera_frame

    def __call__(self, pred_cls, pred_reg, targets, anchors):
        labels, reg_targets, aux_raw_boxes, aux_3D_in_camera_frame = self.prepare_targets(targets, anchors)

        N = len(labels)
        pred_cls_flatten, pred_reg_flatten = concat_box_prediction_layers(pred_cls, pred_reg)

        labels_flatten = torch.cat(labels, dim=0)
        reg_targets_flatten = torch.cat(reg_targets, dim=0)
        aux_raw_boxes_flatten = torch.cat(aux_raw_boxes, dim=0)
        aux_3D_in_camera_frame_flatten = torch.cat(aux_3D_in_camera_frame, dim=0)
        anchors_flatten = torch.cat([cat_boxlist(anchors_per_image).bbox for anchors_per_image in anchors], dim=0)

        pos_inds = torch.nonzero(labels_flatten > 0).squeeze(1)

        valid_cls_inds = torch.nonzero(labels_flatten >= 0).squeeze(1)
        cls_loss = self.cls_loss_func(pred_cls_flatten[valid_cls_inds], labels_flatten[valid_cls_inds])

        if pos_inds.numel() > 0:
            pred_reg_flatten = pred_reg_flatten[pos_inds]
            cls_label_flatten = labels_flatten[pos_inds] - 1 # start from class 0
            reg_targets_flatten = reg_targets_flatten[pos_inds]
            aux_raw_boxes_flatten = aux_raw_boxes_flatten[pos_inds]
            aux_3D_in_camera_frame_flatten = aux_3D_in_camera_frame_flatten[pos_inds]
            anchors_flatten = anchors_flatten[pos_inds]

            reg_loss = self.ObjectSpaceLoss(
                pred_reg_flatten, aux_3D_in_camera_frame_flatten,
                cls_label_flatten, anchors_flatten
                )
        else:
            reg_loss = pred_reg_flatten.sum()

        return cls_loss * self.loss_weight_cls, reg_loss * self.loss_weight_reg


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

"""
A module for parsing information from various files.
"""

from __future__ import (absolute_import, division, print_function, unicode_literals)
import numpy as np
import os

from arkane.statmech import determine_qm_software
from arkane.qchem import QChemLog
from arkane.gaussian import GaussianLog
from arkane.molpro import MolproLog

from arc.common import get_logger
from arc.species.converter import get_xyz_string, standardize_xyz_string
from arc.arc_exceptions import InputError, ParserError

##################################################################

logger = get_logger()

def parse_frequencies(path, software):
    """
    Parse the frequencies from a freq job output file.
    """
    lines = _get_lines_from_file(path)
    freqs = np.array([], np.float64)
    if software.lower() == 'qchem':
        for line in lines:
            if ' Frequency:' in line:
                items = line.split()
                for i, item in enumerate(items):
                    if i:
                        freqs = np.append(freqs, [(float(item))])
    elif software.lower() == 'gaussian':
        with open(path, 'r') as f:
            line = f.readline()
            while line != '':
                if 'Frequencies --' in line:
                    freqs = np.append(freqs, [float(frq) for frq in line.split()[2:]])
                line = f.readline()
    elif software.lower() == 'molpro':
        read = False
        for line in lines:
            if 'Nr' in line and '[1/cm]' in line:
                continue
            if read:
                if line == os.linesep:
                    read = False
                    continue
                freqs = np.append(freqs, [float(line.split()[-1])])
            if 'Low' not in line and 'Vibration' in line and 'Wavenumber' in line:
                read = True
    else:
        raise ParserError('parse_frequencies() can currently only parse Molpro, QChem and Gaussian files,'
                          ' got {0}'.format(software))
    logger.debug('Using parser.parse_frequencies. Determined frequencies are: {0}'.format(freqs))
    return freqs

def parse_t1(path):
    """
    Parse the T1 parameter from a Molpro coupled cluster calculation.
    """
    lines = _get_lines_from_file(path)
    t1 = None
    for line in lines:
        if 'T1 diagnostic:' in line:
            t1 = float(line.split()[-1])
    return t1

def parse_e_elect(path, zpe_scale_factor=1.):
    """
    Parse the zero K energy, E0, from an sp job output file.
    """
    if not os.path.isfile(path):
        raise InputError('Could not find file {0}'.format(path))
    log = determine_qm_software(fullpath=path)
    try:
        e_elect = log.loadEnergy(zpe_scale_factor) * 0.001  # convert to kJ/mol
    except Exception:
        logger.warning('Could not read e_elect from {0}'.format(path))
        e_elect = None
    return e_elect

def parse_xyz_from_file(path):
    """
    Parse xyz coordinated from:
    .xyz - XYZ file
    .gjf - Gaussian input file
    .out or .log - ESS output file (Gaussian, QChem, Molpro)
    other - Molpro or QChem input file
    """
    lines = _get_lines_from_file(path)
    file_extension = os.path.splitext(path)[1]

    xyz = None
    relevant_lines = list()

    if file_extension == '.xyz':
        relevant_lines = lines[2:]
    elif file_extension == '.gjf':
        start_parsing = False
        for line in lines:
            if start_parsing and line and line != '\n' and line != '\r\n':
                relevant_lines.append(line)
            elif start_parsing:
                break
            else:
                splits = line.split()
                if len(splits) == 2 and all([s.isdigit() for s in splits]):
                    start_parsing = True
    elif 'out' in file_extension or 'log' in file_extension:
        log = determine_qm_software(fullpath=path)
        coords, number, _ = log.loadGeometry()
        xyz = get_xyz_string(coords=coords, numbers=number)
    else:
        record = False
        for line in lines:
            if '$end' in line or '}' in line:
                break
            if record and len(line.split()) == 4:
                relevant_lines.append(line)
            elif '$molecule' in line:
                record = True
            elif 'geometry={' in line:
                record = True
        if not relevant_lines:
            raise ParserError('Could not parse xyz coordinates from file {0}'.format(path))
    if xyz is None and relevant_lines:
        xyz = ''.join([line for line in relevant_lines if line])
    return standardize_xyz_string(xyz)

def parse_dipole_moment(path):
    """
    Parse the dipole moment in Debye from an opt job output file.
    """
    lines = _get_lines_from_file(path)
    log = determine_qm_software(path)
    dipole_moment = None
    if isinstance(log, GaussianLog):
        # example:
        # Dipole moment (field-independent basis, Debye):
        # X=             -0.0000    Y=             -0.0000    Z=             -1.8320  Tot=              1.8320
        read = False
        for line in lines:
            if 'dipole moment' in line.lower() and 'debye' in line.lower():
                read = True
            elif read:
                dipole_moment = float(line.split()[-1])
                read = False
    elif isinstance(log, QChemLog):
        # example:
        #     Dipole Moment (Debye)
        #          X       0.0000      Y       0.0000      Z       2.0726
        #        Tot       2.0726
        skip = False
        read = False
        for line in lines:
            if 'dipole moment' in line.lower() and 'debye' in line.lower():
                skip = True
            elif skip:
                skip = False
                read = True
            elif read:
                dipole_moment = float(line.split()[-1])
                read = False
    elif isinstance(log, MolproLog):
        # example:
        #  Dipole moment /Debye                   2.96069859     0.00000000     0.00000000
        for line in lines:
            if 'dipole moment' in line.lower() and '/debye' in line.lower():
                splits = line.split()
                dm_x, dm_y, dm_z = float(splits[-3]), float(splits[-2]), float(splits[-1])
                dipole_moment = (dm_x ** 2 + dm_y ** 2 + dm_z ** 2) ** 0.5
    else:
        raise ParserError('Currently dipole moments can only be parsed from either Gaussian, Molpro, or QChem '
                          'optimization output files')
    if dipole_moment is None:
        raise ParserError('Could not parse the dipole moment')
    return dipole_moment

def parse_polarizability(path):
    """
    Parse the polarizability from a freq job output file, returns the value in Angstrom^3.
    """
    lines = _get_lines_from_file(path)
    polarizability = None
    for line in lines:
        if 'Isotropic polarizability for W' in line:
            # example:  Isotropic polarizability for W=    0.000000       11.49 Bohr**3.
            # 1 Bohr = 0.529177 Angstrom
            polarizability = float(line.split()[-2]) * 0.529177 ** 3
    return polarizability

def _get_lines_from_file(path):
    """
    A helper function for getting a list of lines from the file at `path`.
    """
    if os.path.isfile(path):
        with open(path, 'r') as f:
            lines = f.readlines()
    else:
        raise InputError('Could not find file {0}'.format(path))
    return lines

def process_conformers_file(conformers_path):
    """
    Parse coordinates and energies from an ARC conformers file of either species or TSs.

    Args:
        conformers_path (str): The path to an ARC conformers file
                               (either a "conformers_before_optimization" or
                               a "conformers_after_optimization" file).

    Returns:
        xyz (list): Entries are optimized xyz's in a string format.
    Returns:
        energies (list): Entries float numbers representing the energies in kJ/mol.
    """
    if not os.path.isfile(conformers_path):
        raise ValueError('Conformers file {0} could not be found'.format(conformers_path))
    with open(conformers_path, 'r') as f:
        lines = f.readlines()
    xyzs, energies = list(), list()
    line_index = 0
    while line_index < len(lines):
        if 'conformer' in lines[line_index] and ':' in lines[line_index] and lines[line_index].strip()[-2].isdigit():
            xyz, energy = '', None
            line_index += 1
            print(lines[line_index].strip())
            while line_index < len(lines) and lines[line_index].strip() and 'SMILES' not in lines[line_index]\
                    and 'energy' not in lines[line_index].lower() and 'guess method' not in lines[line_index].lower():
                xyz += lines[line_index]
                line_index += 1
            while line_index < len(lines) and 'conformer' not in lines[line_index]:
                if 'relative energy:' in lines[line_index].lower():
                    energy = float(lines[line_index].split()[2])
                line_index += 1
            xyzs.append(xyz)
            energies.append(energy)
        else:
            line_index += 1
    return xyzs, energies


#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import ccsd_lambda

einsum = lib.einsum

def kernel(mycc, eris=None, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO):
    if eris is None: eris = mycc.ao2mo()
    return ccsd_lambda.kernel(mycc, eris, t1, t2, l1, l2, max_cycle, tol,
                              verbose, make_intermediates, update_lambda)

# l2, t2 as ijab
def make_intermediates(mycc, t1, t2, eris):
    nocc, nvir = t1.shape
    foo = eris.fock[:nocc,:nocc]
    fov = eris.fock[:nocc,nocc:]
    fvo = eris.fock[nocc:,:nocc]
    fvv = eris.fock[nocc:,nocc:]

    tau = t2 + einsum('ia,jb->ijab', t1, t1) * 2

    v1 = fvv - einsum('ja,jb->ba', fov, t1)
    v1-= numpy.einsum('jbac,jc->ba', eris.ovvv, t1)
    v1+= einsum('jkca,jkbc->ba', eris.oovv, tau) * .5

    v2 = foo + einsum('ib,jb->ij', fov, t1)
    v2-= numpy.einsum('kijb,kb->ij', eris.ooov, t1)
    v2+= einsum('ikbc,jkbc->ij', eris.oovv, tau) * .5

    v3 = einsum('ijcd,klcd->ijkl', eris.oovv, tau)
    v4 = einsum('ljdb,klcd->jcbk', eris.oovv, t2)
    v4+= numpy.asarray(eris.ovvo)

    v5 = fvo + numpy.einsum('kc,jkbc->bj', fov, t2)
    tmp = fov - numpy.einsum('kldc,ld->kc', eris.oovv, t1)
    v5+= numpy.einsum('kc,kb,jc->bj', tmp, t1, t1)
    v5-= einsum('kljc,klbc->bj', eris.ooov, t2) * .5
    v5+= einsum('kbdc,jkcd->bj', eris.ovvv, t2) * .5

    w3 = v5 + numpy.einsum('jcbk,jb->ck', v4, t1)
    w3 += numpy.einsum('cb,jb->cj', v1, t1)
    w3 -= numpy.einsum('jk,jb->bk', v2, t1)

    woooo = numpy.asarray(eris.oooo) * .5
    woooo+= v3 * .25
    woooo+= einsum('jilc,kc->jilk', eris.ooov, t1)

    wovvo = v4 - numpy.einsum('ljdb,lc,kd->jcbk', eris.oovv, t1, t1)
    wovvo-= einsum('ljkb,lc->jcbk', eris.ooov, t1)
    wovvo+= einsum('jcbd,kd->jcbk', eris.ovvv, t1)

    wovoo = einsum('icdb,jkdb->icjk', eris.ovvv, tau) * .25
    wovoo+= numpy.einsum('jkic->icjk', numpy.asarray(eris.ooov).conj()) * .5
    wovoo+= einsum('icbk,jb->icjk', v4, t1)
    wovoo-= einsum('lijb,klcb->icjk', eris.ooov, t2)

    wvvvo = einsum('jcak,jb->bcak', v4, t1)
    wvvvo+= einsum('jlka,jlbc->bcak', eris.ooov, tau) * .25
    wvvvo-= numpy.einsum('jacb->bcaj', numpy.asarray(eris.ovvv).conj()) * .5
    wvvvo+= einsum('kbad,jkcd->bcaj', eris.ovvv, t2)

    class _IMDS: pass
    imds = _IMDS()
    imds.ftmp = lib.H5TmpFile()
    dtype = numpy.result_type(t2, eris.vvvv).char
    imds.woooo = imds.ftmp.create_dataset('woooo', (nocc,nocc,nocc,nocc), dtype)
    imds.wovvo = imds.ftmp.create_dataset('wovvo', (nocc,nvir,nvir,nocc), dtype)
    imds.wovoo = imds.ftmp.create_dataset('wovoo', (nocc,nvir,nocc,nocc), dtype)
    imds.wvvvo = imds.ftmp.create_dataset('wvvvo', (nvir,nvir,nvir,nocc), dtype)
    imds.woooo[:] = woooo
    imds.wovvo[:] = wovvo
    imds.wovoo[:] = wovoo
    imds.wvvvo[:] = wvvvo
    imds.v1 = v1
    imds.v2 = v2
    imds.w3 = w3
    imds.ftmp.flush()
    return imds

# update L1, L2
def update_lambda(mycc, t1, t2, l1, l2, eris, imds):
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    fov = eris.fock[:nocc,nocc:]
    mo_e_o = eris.mo_energy[:nocc]
    mo_e_v = eris.mo_energy[nocc:] + mycc.level_shift
    v1 = imds.v1 - numpy.diag(mo_e_v)
    v2 = imds.v2 - numpy.diag(mo_e_o)

    l1new = numpy.zeros_like(l1)
    l2new = numpy.zeros_like(l2)

    mba = einsum('klca,klcb->ba', l2, t2) * .5
    mij = einsum('kicd,kjcd->ij', l2, t2) * .5
    m3 = einsum('klab,ijkl->ijab', l2, numpy.asarray(imds.woooo))
    tau = t2 + einsum('ia,jb->ijab', t1, t1) * 2
    tmp = einsum('ijcd,klcd->ijkl', l2, tau)
    oovv = numpy.asarray(eris.oovv)
    m3 += einsum('klab,ijkl->ijab', oovv, tmp) * .25
    tmp = einsum('ijcd,kd->ijck', l2, t1)
    m3 -= einsum('kcba,ijck->ijab', eris.ovvv, tmp)
    m3 += einsum('ijcd,cdab->ijab', l2, eris.vvvv) * .5

    l2new += oovv
    l2new += m3
    fov1 = fov + einsum('kjcb,kc->jb', oovv, t1)
    tmp = einsum('ia,jb->ijab', l1, fov1)
    tmp+= einsum('kica,jcbk->ijab', l2, numpy.asarray(imds.wovvo))
    tmp = tmp - tmp.transpose(1,0,2,3)
    l2new += tmp - tmp.transpose(0,1,3,2)
    tmp = einsum('ka,ijkb->ijab', l1, eris.ooov)
    tmp+= einsum('ijca,cb->ijab', l2, v1)
    tmp1vv = mba + einsum('ka,kb->ba', l1, t1)
    tmp+= einsum('ca,ijcb->ijab', tmp1vv, oovv)
    l2new -= tmp - tmp.transpose(0,1,3,2)
    tmp = einsum('ic,jcba->jiba', l1, eris.ovvv)
    tmp+= einsum('kiab,jk->ijab', l2, v2)
    tmp1oo = mij + einsum('ic,kc->ik', l1, t1)
    tmp-= einsum('ik,kjab->ijab', tmp1oo, oovv)
    l2new += tmp - tmp.transpose(1,0,2,3)

    l1new += fov
    l1new += einsum('jb,ibaj->ia', l1, eris.ovvo)
    l1new += einsum('ib,ba->ia', l1, v1)
    l1new -= einsum('ja,ij->ia', l1, v2)
    l1new -= einsum('kjca,icjk->ia', l2, imds.wovoo)
    l1new -= einsum('ikbc,bcak->ia', l2, imds.wvvvo)
    l1new += einsum('ijab,jb->ia', m3, t1)
    l1new += einsum('jiba,bj->ia', l2, imds.w3)
    tmp =(t1 + einsum('kc,kjcb->jb', l1, t2)
          - einsum('bd,jd->jb', tmp1vv, t1)
          - einsum('lj,lb->jb', mij, t1))
    l1new += numpy.einsum('jiba,jb->ia', oovv, tmp)
    l1new += numpy.einsum('icab,bc->ia', eris.ovvv, tmp1vv)
    l1new -= numpy.einsum('jika,kj->ia', eris.ooov, tmp1oo)
    tmp = fov - einsum('kjba,jb->ka', oovv, t1)
    l1new -= numpy.einsum('ik,ka->ia', mij, tmp)
    l1new -= numpy.einsum('ca,ic->ia', mba, tmp)

    eia = lib.direct_sum('i-j->ij', mo_e_o, mo_e_v)
    l1new /= eia
    l2new /= lib.direct_sum('ia+jb->ijab', eia, eia)

    time0 = log.timer_debug1('update l1 l2', *time0)
    return l1new, l2new

if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf.cc import gccsd

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.spin = 0
    mol.build()
    mf = scf.RHF(mol).run()
    mf0 = mf
    mf = scf.addons.convert_to_ghf(mf)
    mycc = gccsd.GCCSD(mf)
    eris = mycc.ao2mo()
    mycc.kernel(eris=eris)
    l1, l2 = mycc.solve_lambda(mycc.t1, mycc.t2, eris=eris)
    l1 = mycc.spin2spatial(l1, mycc.mo_coeff.orbspin)
    l2 = mycc.spin2spatial(l2, mycc.mo_coeff.orbspin)
    print(lib.finger(l1[0]) --0.0030030170069977758)
    print(lib.finger(l1[1]) --0.0030030170069977758)
    print(lib.finger(l2[0]) --0.041444910588788492 )
    print(lib.finger(l2[1]) - 0.1077575086912813   )
    print(lib.finger(l2[2]) --0.041444910588788492 )
    print(abs(l2[1]-l2[1].transpose(1,0,2,3)-l2[0]).max())
    print(abs(l2[1]-l2[1].transpose(0,1,3,2)-l2[0]).max())

    from pyscf.cc import ccsd
    mycc0 = ccsd.CCSD(mf0)
    eris0 = mycc0.ao2mo()
    mycc0.kernel(eris=eris0)
    t1 = mycc0.t1
    t2 = mycc0.t2
    imds = ccsd_lambda.make_intermediates(mycc0, t1, t2, eris0)
    l1, l2 = ccsd_lambda.update_lambda(mycc0, t1, t2, t1, t2, eris0, imds)
    l1ref, l2ref = ccsd_lambda.update_lambda(mycc0, t1, t2, l1, l2, eris0, imds)
    t1 = mycc.spatial2spin(t1, mycc.mo_coeff.orbspin)
    t2 = mycc.spatial2spin(t2, mycc.mo_coeff.orbspin)
    l1 = mycc.spatial2spin(l1, mycc.mo_coeff.orbspin)
    l2 = mycc.spatial2spin(l2, mycc.mo_coeff.orbspin)
    imds = make_intermediates(mycc, t1, t2, eris)
    l1, l2 = update_lambda(mycc, t1, t2, l1, l2, eris, imds)
    l1 = mycc.spin2spatial(l1, mycc.mo_coeff.orbspin)
    l2 = mycc.spin2spatial(l2, mycc.mo_coeff.orbspin)
    print(abs(l1[0]-l1ref).max())
    print(abs(l2[1]-l2ref).max())