# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import numpy as np
import sys
import os
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
ROOT_DIR = os.path.dirname(BASE_DIR)
sys.path.append(os.path.join(ROOT_DIR, 'utils'))
from box_util import get_3d_box

class ScannetDatasetConfig(object):
    def __init__(self):
        self.dataset = 'scannet'
        self.num_class = 18
        self.num_heading_bin = 24 # angle: -pi/2~pi/2, so divide 0~2*pi into 24 bin
        self.num_size_cluster = 18

        self.type2class = {'cabinet':0, 'bed':1, 'chair':2, 'sofa':3, 'table':4, 'door':5, 'window':6,'bookshelf':7,'picture':8, 'counter':9, 'desk':10, 'curtain':11, 'refrigerator':12, 'showercurtrain':13, 'toilet':14, 'sink':15, 'bathtub':16, 'garbagebin':17}
        #self.type2class = {'wall':0, 'floor':1, 'cabinet':2, 'bed':3, 'chair':4, 'sofa':5, 'table':6, 'door':7,'window':8,'bookshelf':9,'picture':10, 'counter':11, 'blinds':12, 'desk':13, 'shelves':14, 'curtain':15, 'dresser':16, 'pillow':17, 'mirror':18, 'floormat':19, 'clothes':20, 'ceiling':21, 'books':22, 'refrigerator':23, 'television':24, 'paper':25, 'towel':26, 'showercurtrain':27, 'box':28, 'whiteboard':29, 'person':30, 'nightstand':31, 'toilet':32, 'sink':33, 'lamp':34, 'bathtub':35, 'bag':36}
        self.type2class_room = {'other':0, 'wall':1, 'floor':2}
        self.class2type = {self.type2class[t]:t for t in self.type2class}
        self.class2type_room = {self.type2class_room[t]:t for t in self.type2class_room}
        self.nyu40ids = np.array([3,4,5,6,7,8,9,10,11,12,14,16,24,28,33,34,36,39])
        #self.nyu40ids = np.array([0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36])
        self.nyu40ids_room = np.array([1,2])

        self.nyu40id2class = {nyu40id: i for i,nyu40id in enumerate(list(self.nyu40ids))}
        self.nyu40id2class_sem = {nyu40id: i for i,nyu40id in enumerate(list(self.nyu40ids))}
        self.mean_size_arr = np.load(os.path.join(ROOT_DIR,'scannet/meta_data/scannet_means.npz'))['arr_0']
        #self.mean_size_arr = np.load(os.path.join(ROOT_DIR,'scannet/meta_data/scannet_means_v2.npz.npy'))[:self.num_class,:]
        self.type_mean_size = {}
        for i in range(self.num_size_cluster):
            self.type_mean_size[self.class2type[i]] = self.mean_size_arr[i,:]

    def class2angle(self, pred_cls, residual, to_label_format=True):
        return 0

    '''
    def angle2class(self, angle):
        # assert(False)
        num_class = self.num_heading_bin
        angle = angle%(2*np.pi)
        assert(angle>=0 and angle<=2*np.pi)
        angle_per_class = 2*np.pi/float(num_class)
        shifted_angle = (angle+angle_per_class/2)%(2*np.pi)
        class_id = int(shifted_angle/angle_per_class)
        residual_angle = shifted_angle - (class_id*angle_per_class+angle_per_class/2)
        return class_id, residual_angle
    def class2angle(self, pred_cls, residual, to_label_format=True):
        num_class = self.num_heading_bin
        angle_per_class = 2*np.pi/float(num_class)
        angle_center = pred_cls * angle_per_class
        angle = angle_center + residual
        if to_label_format and angle>np.pi:
            angle = angle - 2*np.pi
        return angle
    '''
    def angle2class2(self, angle):
        ''' modify according to sunrgbd
            scannet_angle: angle: -pi/2 ~ pi/2
            1: angle += pi/2 -> 0~pi
            2: class*(2pi/N) + number = angle + pi/2
        '''
        class_id, residual_angle = self.angle2class(angle + np.pi / 2)
        return class_id, residual_angle

    def class2angle2(self, pred_cls, residual, to_label_format=True):
        angle = self.class2angle(pred_cls, residual)
        angle = angle - np.pi / 2
        return angle

    def size2class(self, size, type_name):
        ''' Convert 3D box size (l,w,h) to size class and size residual '''
        size_class = self.type2class[type_name]
        size_residual = size - self.type_mean_size[type_name]
        return size_class, size_residual

    def class2size(self, pred_cls, residual):
        ''' Inverse function to size2class '''
        return self.mean_size_arr[pred_cls, :] + residual

    def param2obb(self, center, heading_class, heading_residual, size_class, size_residual):
        heading_angle = self.class2angle(heading_class, heading_residual)
        box_size = self.class2size(int(size_class), size_residual)
        obb = np.zeros((7,))
        obb[0:3] = center
        obb[3:6] = box_size
        obb[6] = heading_angle*-1
        return obb

    def param2obb2(self, center, heading_class, heading_residual, size_class, size_residual):
        heading_angle = self.class2angle(heading_class, heading_residual)
        box_size = self.class2size(int(size_class), size_residual)
        obb = np.zeros((7,))
        obb[0:3] = center
        obb[3:6] = box_size
        obb[6] = heading_angle
        return obb

def rotate_aligned_boxes(input_boxes, rot_mat):
    centers, lengths = input_boxes[:,0:3], input_boxes[:,3:6]
    new_centers = np.dot(centers, np.transpose(rot_mat))

    dx, dy = lengths[:,0]/2.0, lengths[:,1]/2.0
    new_x = np.zeros((dx.shape[0], 4))
    new_y = np.zeros((dx.shape[0], 4))

    for i, crnr in enumerate([(-1,-1), (1, -1), (1, 1), (-1, 1)]):
        crnrs = np.zeros((dx.shape[0], 3))
        crnrs[:,0] = crnr[0]*dx
        crnrs[:,1] = crnr[1]*dy
        crnrs = np.dot(crnrs, np.transpose(rot_mat))
        new_x[:,i] = crnrs[:,0]
        new_y[:,i] = crnrs[:,1]

    new_dx = 2.0*np.max(new_x, 1)
    new_dy = 2.0*np.max(new_y, 1)
    new_lengths = np.stack((new_dx, new_dy, lengths[:,2]), axis=1)

    return np.concatenate([new_centers, new_lengths], axis=1)


# 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


import torch
import numpy as np
from Torch_rl.agent.core_policy import Agent_policy_based
import torch.nn as nn
from torch.autograd import Variable as V
from copy import deepcopy
from Torch_rl.common.distribution import *
from torch.optim import Adam
from torch.autograd import Variable
from gym import spaces
from Torch_rl.common.util import csv_record
from Torch_rl.common.util import get_gae

class gpu_foward(nn.Module):
    def __init__(self, model):
        super(gpu_foward, self).__init__()
        model.to_gpu()
        self.model = model
    def forward(self,obs):
        obs = obs.cuda()
        out = self.model(obs)
        return out

class PPO_Agent(Agent_policy_based):
    def __init__(self, env, policy_model, value_model,
                 lr=5e-4, ent_coef=0.01, vf_coef=0.5,
                 ## hyper-parawmeter
                 gamma=0.90, lam=0.95, cliprange=0.2, batch_size=64, value_train_round=10,
                 running_step=2048, running_ep=20, value_regular=0.01,
                 ## decay
                 decay=False, decay_rate=0.9,
                 ##
                 path=None):
        self.gpu = False
        self.env = env
        self.gamma = gamma
        self.lam = lam
        self.ent_coef = ent_coef
        self.vf_coef = vf_coef
        self.cliprange = cliprange

        self.value_train_step = value_train_round

        self.sample_rollout = running_step
        self.sample_ep = running_ep
        self.batch_size = batch_size

        self.loss_cal = torch.nn.SmoothL1Loss()

        self.policy = policy_model
        if value_model == "shared":
            self.value = policy_model
        elif value_model == "copy":
            self.value = deepcopy(policy_model)
        else:
            self.value = value_model

        self.dist = make_pdtype(env.action_space, policy_model)

        self.policy_model_optim = Adam(self.policy.parameters(), lr=lr)
        self.value_model_optim = Adam(self.value.parameters(), lr=lr, weight_decay=value_regular)
        if decay:
            self.policy_model_decay_optim = torch.optim.lr_scheduler.ExponentialLR(self.policy_model_optim, decay_rate,
                                                                            last_epoch=-1)
            self.value_model_decay_optim = torch.optim.lr_scheduler.ExponentialLR(self.value_model_optim, decay_rate,
                                                                             last_epoch=-1)

        #torch.nn.utils.clip_grad_norm_(self.policy.parameters(), 1, norm_type=2)
        #torch.nn.utils.clip_grad_norm_(self.value.parameters(), 1, norm_type=2)

        super(PPO_Agent, self).__init__(path)
        example_input = Variable(torch.rand((100,)+self.env.observation_space.shape))
        self.writer.add_graph(self.policy, input_to_model=example_input)

        self.backward_step_show_list = ["pg_loss", "entropy", "vf_loss"]
        self.backward_ep_show_list = ["pg_loss", "entropy", "vf_loss"]

        self.training_round = 0
        self.running_step = 0
        self.record_sample = None
        self.training_step = 0

    def update(self, sample):
        step_len = len(sample["s"])

        time_round = np.ceil(step_len/self.batch_size)
        time_left = time_round*self.batch_size-step_len
        array = list(range(step_len)) +list(range(int(time_left)))
        loss_re, pgloss_re, enloss_re, vfloss_re = [], [], [], []

        for key in sample.keys():
            temp = torch.stack(list(sample[key])).squeeze()
            if self.gpu:
                sample[key] = temp.cuda()
            else:
                sample[key] = temp

        for train_time in range(int(time_round)):
            index = array[train_time*self.batch_size: (train_time+1)*self.batch_size]
        # for index in range(step_len):
            training_s = sample["s"][index].detach()
            training_a = sample["a"][index].detach()
            training_r = sample["r"][index].detach()
            R = sample["return"][index].detach()
            old_value = sample["value"][index].detach()
            old_neglogp = sample["logp"][index].detach()
            advs = sample["advs"][index].detach()

            " CALCULATE THE LOSS"
            " Total loss = Policy gradient loss - entropy * entropy coefficient + Value coefficient * value loss"

            " the value loss"
            value_now = self.value.forward(training_s).squeeze()
            # value loss
            value_clip = old_value + torch.clamp(old_value - value_now, min=-self.cliprange,
                                                 max=self.cliprange)  # Clipped value
            vf_loss1 = self.loss_cal(value_now, R)  # Unclipped loss
            vf_loss2 = self.loss_cal(value_clip, R)  # clipped loss
            vf_loss = .5 * torch.max(vf_loss1, vf_loss2)

            #generate Policy gradient loss
            outcome = self.policy.forward(training_s)
            # new_neg_lop = torch.empty(size=(self.batch_size,))
            # for time in range(self.batch_size):
            #     new_policy = self.dist(outcome[time])
            #     new_neg_lop[time] = new_policy.log_prob(training_a[time])
            new_policy = self.dist(outcome)
            new_neg_lop = new_policy.log_prob(training_a)
            ratio = torch.exp(torch.diag(new_neg_lop) - old_neglogp)
            pg_loss1 = -advs * ratio
            pg_loss2 = -advs * torch.clamp(ratio, 1.0 - self.cliprange, 1.0 + self.cliprange)
            pg_loss = .5 * torch.max(pg_loss1, pg_loss2).mean()

            # entropy
            entropy = new_policy.entropy().mean()
            loss = pg_loss - entropy * self.ent_coef + vf_loss * self.vf_coef

            self.policy_model_optim.zero_grad()
            pg_loss.backward()
            self.policy_model_optim.step()
            for _ in range(self.value_train_step):
                value_now = self.value.forward(training_s).squeeze()
                # value loss
                value_clip = old_value + torch.clamp(old_value - value_now, min=-self.cliprange,
                                                     max=self.cliprange)  # Clipped value
                vf_loss1 = self.loss_cal(value_now, R)  # Unclipped loss
                vf_loss2 = self.loss_cal(value_clip, R)  # clipped loss
                vf_loss = .5 * torch.max(vf_loss1, vf_loss2)
                self.value_model_optim.zero_grad()
                vf_loss1.backward()
                self.value_model_optim.step()
            # approxkl = self.loss_cal(neg_log_pac, self.record_sample["neglogp"])
            # self.cliprange = torch.gt(torch.abs(ratio - 1.0).mean(), self.cliprange)
            loss_re = loss.detach().numpy()
            pgloss_re.append(pg_loss.detach().numpy())
            enloss_re.append(entropy.detach().numpy())
            vfloss_re.append(vf_loss1.detach().numpy())
        return np.sum(loss_re), {"pg_loss": np.sum(pgloss_re),
                                   "entropy": np.sum(enloss_re),
                                   "vf_loss": np.sum(vfloss_re)}

    def load_weights(self, filepath):
        model = torch.load(filepath+"ppo.pkl")
        self.policy.load_state_dict(model["policy"].state_dict())
        self.value.load_state_dict(model["value"].state_dict())

    def save_weights(self, filepath, overwrite=False):
        torch.save({"policy": self.policy,"value": self.value}, filepath + "PPO.pkl")

    def behavior_clone(self, buffer):
        sample_ = buffer.sample(32)
        action_label = sample_["a"].squeeze()
        action_predict = self.policy(sample_["s"].cuda())
        loss_bc = self.loss_cal(action_label.cuda(), action_predict)
        loss = loss_bc
        self.policy_model_optim.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(self.policy.parameters(), 1, norm_type=2)
        self.policy_model_optim.step()
        return loss.cpu().detach().numpy()

    def value_pretrain(self, record_sample, new_sample_len):
        train_times = np.floor(new_sample_len/128)
        round_loss = 0
        for io in range(train_times-1):
            index = list(range(128 * io, 128 * (io + 1)))
            if self.gpu:
                predict = torch.from_numpy(np.array(record_sample["s"])[index]).cuda()
                lable = torch.from_numpy(np.array(record_sample["return"]))[index].cuda()
            else:
                perdict = torch.from_numpy(np.array(record_sample["s"])[index])
                lable = torch.from_numpy(np.array(record_sample["return"]))[index]
            value_now = self.value.forward(perdict)
            # value loss
            vf_loss = self.loss_cal(value_now,lable)  # Unclipped loss
            self.value_model_optim.zero_grad()
            vf_loss.backward()
            self.value_model_optim.step()
            round_loss += vf_loss.cpu().detach().numpy()
        return round_loss

    def cuda(self):
        self.policy = gpu_foward(self.policy)
        self.value = gpu_foward(self.value)
        self.loss_cal = self.loss_cal.cuda()
        self.gpu = True


from __future__ import print_function
import sys
import torch
import torch.backends.cudnn as cudnn
import numpy as np
from data import BaseTransform, VOC_300, VOC_512
import cv2
from layers.functions import Detect, PriorBox
import matplotlib.patches as patches
from collections import OrderedDict
from PyQt5 import QtWidgets, QtCore, QtGui
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from models.RFB_Net_vgg import build_net
import time

classes = ['aeroplane', 'ship', 'storage_tank', 'baseball_diamond', 'tennis_court', 'basketball_court',
           'ground_track_field', 'harbor', 'bridge', 'vehicle']

class RFB_GUI(QtWidgets.QMainWindow):
    def __init__(self):
        super(RFB_GUI, self).__init__()

        MyMessageBox(self)
        self.setWindowTitle("RFB-GUI Demo Program")
        self.resize(1280, 900)
        self.setFocus()

        self.file_item = QtWidgets.QAction('Open image', self)
        self.file_item.setShortcut('Ctrl+O')
        self.file_item.triggered.connect(self.select_file)

        self.label = DragLabel("Please drag image here\nor\nPress Ctrl+O to select", self)
        self.label.addAction(self.file_item)
        self.setCentralWidget(self.label)

        self.priorbox = PriorBox(self.cfg)
        self.cuda = True
        self.numclass = 21
        self.net = build_net('test', self.input_size, self.numclass)  # initialize detector
        state_dict = torch.load(self.trained_model)

        new_state_dict = OrderedDict()
        for k, v in state_dict.items():
            head = k[:7]
            if head == 'module.':
                name = k[7:]
            else:
                name = k
            new_state_dict[name] = v
        self.net.load_state_dict(new_state_dict)
        self.net.eval()
        if self.cuda:
            self.net = self.net.cuda()
            cudnn.benchmark = True
        else:
            self.net = self.net.cpu()
        print('Finished loading model!')

    def select_file(self):
        file_path = QtWidgets.QFileDialog.getOpenFileName(self, 'Select image',
                                                          r'/home/zcy/data/NWPU_VHR-10_dataset/positive_image_set',
                                                          "Image files(*.bmp *.jpg *.pbm *.pgm *.png *.ppm *.xbm *.xpm)"
                                                          ";;All files (*.*)")
        # try:
        self.detect(file_path[0], self)
        # except Exception as e:
        #     QtWidgets.QMessageBox.information(self, "Alert", str(e))

    def detect(self, file_name, object):
        print(file_name)
        start_time = time.time()
        img = cv2.imread(file_name.strip())
        if img is None:
            QtWidgets.QMessageBox.information(self, "Alert", "Please select images")
            return
        scale = torch.Tensor([img.shape[1], img.shape[0],
                              img.shape[1], img.shape[0]])
        detector = Detect(object.numclass, 0, object.cfg)
        transform = BaseTransform(object.net.size, (123, 117, 104), (2, 0, 1))
        with torch.no_grad():
            x = transform(img).unsqueeze(0)
            if object.cuda:
                x = x.cuda()
                scale = scale.cuda()
        out = object.net(x)
        with torch.no_grad():
            priors = object.priorbox.forward()
            if object.cuda:
                priors = priors.cuda()
        boxes, scores = detector.forward(out, priors)
        boxes = boxes[0]
        scores = scores[0]
        boxes *= scale
        boxes = boxes.cpu().numpy()
        scores = scores.cpu().numpy()

        result_set = []
        for j in range(1, object.numclass):
            max_ = max(scores[:, j])
            inds = np.where(scores[:, j] > 0.2)[0]  # conf > 0.6
            if inds is None:
                continue
            c_bboxes = boxes[inds]
            c_scores = scores[inds, j]
            c_dets = np.hstack((c_bboxes, c_scores[:, np.newaxis])).astype(
                np.float32, copy=False)
            keep = object.nms_py(c_dets, 0.6)
            c_dets = c_dets[keep, :]
            c_bboxes = c_dets[:, :4]
            for bbox in c_bboxes:
                # Create a Rectangle patch
                rect = patches.Rectangle((int(bbox[0]), int(bbox[1])), int(bbox[2]) - int(bbox[0]) + 1,
                                         int(bbox[3]) - int(bbox[1]) + 1, linewidth=1, edgecolor='r')
                result_set.append(str(rect))
                cv2.rectangle(img, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (255, 0, 0), 2)
                cv2.imwrite("my_test.png", img)

        end_time = time.time()
        print(end_time - start_time)
        img_data = QtGui.QPixmap("my_test.png")
        height = object.height()
        width = object.height() / img_data.height() * img_data.width()
        img_data = img_data.scaled(width, height)
        object.label.resize(width, height)
        object.label.setPixmap(img_data)
        self.setFocus()

    def nms_py(self, dets, thresh):
        x1 = dets[:, 0]
        y1 = dets[:, 1]
        x2 = dets[:, 2]
        y2 = dets[:, 3]
        scores = dets[:, 4]

        areas = (x2 - x1 + 1) * (y2 - y1 + 1)
        order = scores.argsort()[::-1]
        ndets = dets.shape[0]
        suppressed = np.zeros((ndets), dtype=np.int)
        keep = []
        for _i in range(ndets):
            i = order[_i]
            if suppressed[i] == 1:
                continue
            keep.append(i)
            ix1 = x1[i]
            iy1 = y1[i]
            ix2 = x2[i]
            iy2 = y2[i]
            iarea = areas[i]
            for _j in range(_i + 1, ndets):
                j = order[_j]
                if suppressed[j] == 1:
                    continue
                xx1 = max(ix1, x1[j])
                yy1 = max(iy1, y1[j])
                xx2 = min(ix2, x2[j])
                yy2 = min(iy2, y2[j])
                w = max(0.0, xx2 - xx1 + 1)
                h = max(0.0, yy2 - yy1 + 1)
                inter = w * h
                ovr = inter / (iarea + areas[j] - inter)
                if ovr >= thresh:
                    suppressed[j] = 1
        return keep

class MyMessageBox(QMessageBox):
    def __init__(self, parent):
        super().__init__()
        self.setWindowTitle('Input Size')
        self.setText("Please choose the input image size ")
        self.setFont(QtGui.QFont("Ubuntu Mono", 14))

        _300_button = self.addButton(self.tr('300 x 300'), QMessageBox.ActionRole)
        _512_button = self.addButton(self.tr('512 x 512'), QMessageBox.ActionRole)
        cancel_button = self.addButton(' Cancel ', QMessageBox.ActionRole)
        self.exec_()
        button = self.clickedButton()
        if button == _300_button:
            parent.input_size = 300
            parent.cfg = VOC_300
            parent.trained_model = 'weights/RFB_vgg_NWPU_300.pth'
        elif button == _512_button:
            parent.input_size = 512
            parent.cfg = VOC_512
            parent.trained_model = 'weights/RFB_vgg_NWPU_512.pth'
        elif button == cancel_button:
            sys.exit()

class DragLabel(QLabel):
    def __init__(self, text, parent):
        super().__init__(text, parent)
        self.parent = parent
        self.setAcceptDrops(True)
        self.setFont(QtGui.QFont("Ubuntu Mono", 30))
        self.setAlignment(Qt.AlignCenter)

    def dragEnterEvent(self, QDragEnterEvent):
        QDragEnterEvent.accept()

    def dropEvent(self, QDropEvent):
        self.parent.detect(QDropEvent.mimeData().text()[7:], self.parent)

if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    rfb_gui = RFB_GUI()
    rfb_gui.show()
    sys.exit(app.exec_())


from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import os
from typing import Dict, List, Tuple

import numpy as np
import tensorflow as tf
from collections import namedtuple

from rb.processings.gec.transformer.transformer import Transformer
from rb.processings.gec.transformer.transformer_scheduler import CustomSchedule
from rb.processings.gec.transformer.utils import create_masks
from rb.processings.gec import beam_search
from rb.utils.downloader import check_version, download_model
from rb.core.lang import Lang
import tensorflow_datasets as tfds

class Beam(namedtuple("Beam", ["log_prob", "ids", "sentence", "length"])):
  """A finished beam

  Args:
    probs: Log probability of them beam
    ids: List of ids of the the beam
    sentence: Sentence in string
    lengths: Length of the beam
  """
  pass

class GecCorrector():

    def __init__(self, d_model: str=64, beam: int=8,
                lm_path: str=None, normalize: bool=False, weight_lm: float=1., check_updates = True):
        self.d_model = d_model
        self.beam = beam
        self.lm_path = lm_path
        self.normalize = normalize
        self.weight_lm = weight_lm

        if d_model == 768:
            name = 'transformer_768'
        elif d_model == 64:
            name = 'transformer_64'

        self.checkpoint_path = os.path.join('resources', 'ro', 'models', 'gec', name)
        self.vocabulary_path = os.path.join(self.checkpoint_path, 'tokenizer_ro')
        if check_updates and check_version(Lang.RO, ['models', 'gec', name]):
            download_model(Lang.RO, ['models', 'gec', name])
        self.transformer, self.optimizer = self.__load_model()
        self.lm_model = self.__load_lm()
        self.tokenizer_ro = self.__load_tokenizer()

    def __load_tokenizer(self):
        tokenizer_ro = tfds.features.text.SubwordTextEncoder.load_from_file(self.vocabulary_path)
        return tokenizer_ro

    def __load_lm(self):
        if self.lm_path:
            # install kenlm from https://github.com/kpu/kenlm, but it works without it (no reranking for beams)
            import kenlm
            lm_model = kenlm.Model(lm_path)
        else:
            lm_model = None
        return lm_model

    def __load_model(self):

        self.__set_params_model(self.d_model)
        self.vocab_size = self.dict_size + 2

        learning_rate = CustomSchedule(self.d_model)
        optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98,
                                        epsilon=1e-9)

        transformer = Transformer(self.num_layers, self.d_model, self.num_heads, self.dff,
                            self.vocab_size, self.vocab_size,
                            pe_input=self.vocab_size,
                            pe_target=self.vocab_size,
                            rate=self.dropout)

        ckpt = tf.train.Checkpoint(transformer=transformer, optimizer=optimizer)

        ckpt_manager = tf.train.CheckpointManager(ckpt, self.checkpoint_path, max_to_keep=5)
        if ckpt_manager.latest_checkpoint:
            ckpt.restore(ckpt_manager.latest_checkpoint)
        else:
            tf.compat.v1.logging.error('no checkpoints for transformers, aborting')
            return None

        return transformer, optimizer

    def __set_params_model(self, d_model: int=64):
        if d_model == 64:
            self.num_layers = 2
            self.seq_length = 64
            self.dff = 64
            self.num_heads = 2
            self.dropout = 0.1
            self.dict_size = 1024

            self.num_layers = 2
        elif d_model == 768:
            self.num_layers = 6
            self.seq_length = 512
            self.dff = 2048
            self.num_heads = 8
            self.dropout = 0.1
            self.dict_size = (2**15)

    def __init_beam(self, vocab_size, end_token_id, beam_width=1):

        length_penalty = 0.6 if self.normalize else 0.0
        config = beam_search.BeamSearchConfig(
            beam_width=beam_width,
            vocab_size=vocab_size,
            eos_token=end_token_id,
            length_penalty_weight=length_penalty,
            choose_successors_fn=beam_search.choose_top_k)

        beam_state = beam_search.BeamSearchState(
            log_probs=tf.nn.log_softmax(tf.ones(config.beam_width)),
            lengths=tf.constant(
                1, shape=[config.beam_width], dtype=tf.int32),
            finished=tf.zeros(
                [config.beam_width], dtype=tf.bool))
        return config, beam_state

    def correct_sentence(self, inp_sentence: str):
        start_token, end_token = [self.tokenizer_ro.vocab_size], [self.tokenizer_ro.vocab_size + 1]
        inp_sentence = inp_sentence.strip()
        in_sentence = inp_sentence
        inp_sentence = start_token + self.tokenizer_ro.encode(inp_sentence) + end_token
        if len(inp_sentence) > self.d_model:
            print('Sentence is too long (more than 766 subwords), cannot correct it.')
            return in_sentence

        start_token_id, end_token_id = self.tokenizer_ro.vocab_size, self.tokenizer_ro.vocab_size + 1

        # duplicate x beam_width == batch size
        encoder_input = tf.expand_dims(inp_sentence, 0)
        encoder_input = tf.tile(encoder_input, [self.beam, 1])

        decoder_input = [start_token_id] * self.beam
        output = tf.expand_dims(decoder_input, 1) # for batch size == beam_wisth

        # beam search init
        config, beam_state = self.__init_beam(vocab_size=(self.dict_size + 2),
                                                    end_token_id=end_token_id,
                                                    beam_width=self.beam)
        beam_values = tf.constant(start_token_id, shape=(1, self.beam))
        beam_parents = tf.zeros((2, self.beam), dtype=tf.int32)

        for i in range(self.d_model + 2):
            enc_padding_mask, combined_mask, dec_padding_mask = create_masks(
                encoder_input, output)

            predictions, attention_weights = self.transformer(encoder_input, output,
                                                                False, enc_padding_mask,
                                                                combined_mask, dec_padding_mask)
            # !predictions.shape == (batch_size, i, vocab_size) (predicts a softmax for each existing word)
            beam_pred = tf.squeeze(predictions[: ,-1:, :], 1)  # (batch_size, 1, vocab_size), select only the last word
            bs_output, beam_state = beam_search.beam_search_step(time_=i, logits=beam_pred,
                                                                beam_state=beam_state, config=config)

            # add new predictions to the beams decoder
            bs_output_predicted_ids = tf.expand_dims(bs_output.predicted_ids, axis=0)
            beam_values = tf.concat([beam_values, bs_output_predicted_ids], axis=0)
            res = tf.cast(beam_search.gather_tree_py(beam_values.numpy(), beam_parents.numpy()), dtype=tf.int32)
            output = tf.transpose(res)

            bs_output_beam_parent_ids = tf.expand_dims(bs_output.beam_parent_ids, axis=0)
            beam_parents = tf.concat([beam_parents, bs_output_beam_parent_ids], axis=0)

            all_finished = tf.reduce_all(beam_state.finished) # and
            if all_finished:    break

        beams = []
        for i, out in enumerate(output):
            b = Beam(log_prob=beam_state.log_probs[i].numpy(), sentence="", ids=out.numpy(), length=len(out.numpy()))
            beams.append(b)

        sentences = self.__decode_sentence(beams)
        # for sent in sentences:
        #     print(sent.log_prob, sent.sentence)
        sentences = sorted(sentences, key=lambda beam: beam.log_prob, reverse=True)
        return sentences[0].sentence

    def __decode_sentence(self, beams):
        sentences = []

        for beam in beams:
            sentence_ids = []

            for i in beam.ids:
                if i < self.tokenizer_ro.vocab_size:
                    sentence_ids.append(i)
                if i == self.tokenizer_ro.vocab_size + 1:
                    break

            predicted_sentence = self.tokenizer_ro.decode(sentence_ids)

            if self.lm_model:
                lm_prob = lm_model.score(predicted_sentence, bos=True, eos=True)

                if self.normalize:
                    cand_prob = beam.log_prob + self.weight_lm * lm_prob * (1.0/beam.length)
                else:
                    cand_prob = beam.log_prob + self.weight_lm * lm_prob
            else:
                cand_prob = beam.log_prob

            sentences.append(Beam(log_prob=cand_prob, sentence=predicted_sentence,
                                 ids=beam.ids, length=len(beam.ids)))
        return sentences

if __name__ == "__main__":
    corrector = GecCorrector(checkpoint_path='checkpoints/10m_transformer_768_2',
                vocabulary_path='checkpoints/10m_transformer_768_2/tokenizer_ro', d_model=768)

    decoded_senntence = corrector.correct_sentence('Am mers la magazi sa cuumper leegume si fructe de mancaat.')
    print(decoded_senntence)


from __future__ import print_function, division

from torch.autograd import Variable
import torch
from PIL import Image
import torch.nn as nn
import pandas as pd
from skimage import io, transform
import torch.optim as optim
from torch.optim import lr_scheduler
import numpy as np
from torch.utils.data import Dataset, DataLoader
import torchvision
from torchvision import datasets, models, transforms, utils
import matplotlib.pyplot as plt
import time
import os
import copy

import serial
import warnings
import serial.tools.list_ports
import cv2 as OpenCV

def Resnet18(pesos,clases = 5):
    ResnetFacial = models.resnet18(pretrained=True)
    num_ftrs = ResnetFacial.fc.in_features
    ResnetFacial.fc = nn.Linear(num_ftrs, clases)
    ResnetFacial = ResnetFacial.to(device)
    ResnetFacial.load_state_dict(torch.load(pesos))
    return ResnetFacial

def Normalizar(x):
    normalized = (x-np.min(x))/(np.max(x)-np.min(x))
    return normalized

def Conversion(imagen):
    pil_image = OpenCV.cvtColor(imagen, OpenCV.COLOR_BGR2RGB)
    image = Image.fromarray(pil_image)
    trans = transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
    image = trans(image)
    image = Variable(image, requires_grad=True)
    image = image.unsqueeze(0)
    return image.cuda()

def Prueba(ResnetFacial, imagen):
    with torch.no_grad():
        ResnetFacial.eval()
        outputs = ResnetFacial(Conversion(imagen))
        _, preds = torch.max(outputs,1)
        Resultados = Normalizar(outputs.data.cpu().numpy())
    return ''.join(str(clases[preds]))

def BuscaArduino():
    arduino_ports = [
        p.device
        for p in serial.tools.list_ports.comports()
        if 'Arduino' in p.description
    ]
    if not arduino_ports:
        raise IOError("No Arduino")
    if len(arduino_ports) > 1:
        warnings.warn('Usar primer arduino encontrado')
    arduino = serial.Serial(arduino_ports[0],9600)
    return arduino

def ControlPorGestos(inicio = True):
    elapsed = 0;
    font = OpenCV.FONT_HERSHEY_SIMPLEX
    cam = OpenCV.VideoCapture(0)
    ret_val, img = cam.read()
    prueba = Prueba(Modelo,img)
    inicio = time.time()
    while True:
        ret_val, img = cam.read()
        if elapsed >= 0.1:
            prueba = Prueba(Modelo,img)
            try:
                Accion(prueba)
            except serial.serialutil.SerialException:
                prueba = 'No hay arduino'
            inicio = fin
        ImgDone = OpenCV.putText(img,prueba,(50,400), font, 3,(255,255,255),2,OpenCV.LINE_AA)
        OpenCV.imshow('Reconocimiento facial', ImgDone)
        if OpenCV.waitKey(1) == 27:
            break  # esc to quit
        fin = time.time()
        elapsed = fin - inicio
    OpenCV.destroyAllWindows()
    cam.release();
    arduino.close();

def Accion(Gesto):
    if Gesto == clases[0]:
        arduino.write(b'0')
    if Gesto == clases[1]:
        arduino.write(b'1')
    if Gesto == clases[2]:
        arduino.write(b'2')
    if Gesto == clases[3]:
        arduino.write(b'3')
    if Gesto == clases[4]:
        arduino.write(b'4')

clases = ['Abajo','Arriba','Derecha','Izquierda','Neutro']
arduino = BuscaArduino();
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
pesos = 'Resnet18Facial.pth'
Modelo = Resnet18(pesos)
Lectura = OpenCV.VideoCapture(0)

valido, foto = Lectura.read()
if valido == True:
    print('Foto lista')
else:
    print('Error al iniciar cámara')
Lectura.release()

ControlPorGestos();
print(arduino_ports[0])


import numpy as np
import pylab as plt
import h5py
import os
import pandas

import radiograph

class GlobalOffsets():
    """Singleton object to hold the global position information computed by Chris W."""

    def __init__(self):
        self.readdata = False
        try:
            p = pandas.read_csv("20211027_global_corrections_sr.csv")
            self.offsetdata = p.set_index([" Moving HDF5"])
            self.readdata = True
        except Exception:
            pass

    def position_offsets_mm(self, filename):
        if not self.readdata:
            print("Warning! Could not read offsets for file {}".format(filename))
            return 0.0, 0.0
        try:
            row = self.offsetdata.loc[filename]
            return row[" Shift[0] (um)"]/1000, row[" Shift[1] (um)"]/1000
        except Exception:
            print("Warning! Exception computing offsets for file {}".format(filename))
            return 0.0, 0.0

_globalOffsets = GlobalOffsets()

def position_offsets_mm(filename):
    """Use the singleton GlobalOffsets object to compute this file's offsets in mm. (Add to HDF5 values.)"""
    return _globalOffsets.position_offsets_mm(filename)

class OneScan():
    """
    Represent one raster scan on the target.
    """

    def __init__(self, filename, verbose=True, globaloffset=True):
        f = h5py.File(filename, "r")
        self.verbose = verbose
        self.h5 = f
        self.filename = "/".join(filename.split(os.path.sep)[-3:])
        self.basename = os.path.basename(filename)
        self.angle = float(self.basename.split("_")[3])
        try:
            self.nsplits = f["numSplits"][()]
        except KeyError:
            self.nsplits = 1

        self.good = f["detInfo/goodChans"][:]
        try:
            f["eds"]
            self.eds = f["eds/rate_target_line"][:]
            self.eds_unc = f["eds/rate_unc_target_line"][:]
        except KeyError:
            print("No EDS data in {}. Yikes".format(filename))

        global_offsets = np.zeros(2, dtype=float)
        if globaloffset:
            global_offsets = np.array(position_offsets_mm(self.basename))
        self.target = f["steps/corrected_sampleFixed"][:, :2] + global_offsets
        if "endPos_sampleFixed" in f["steps"]:
            self.targetend = f["steps/endPos_sampleFixed"][:, :2] + global_offsets
        else:
            # If endPos isn't there, assume these are nanodots.
            self.targetend = self.target
        self.command = f["steps/command_sampleFixed"][:, :2] + global_offsets
        self.global_offsets = global_offsets
        self.offset = self.target-self.command
        self.drift = self.targetend-self.target
        if self.nsplits == 1:
            self.position = self.target + 0.5*self.drift
        else:
            self.position = f["steps/subDwellPos_sampleFixed"][:, :2] + global_offsets

        self.duration = f["times/duration"][:]
        self.start = f["times/start/ns"][:]*1e-9
        self.medianrates = {}
        self.xdet = {}
        self.ydet = {}
        for k, tesgrp in f["tes"].items():
            if not k.startswith("chan"):
                continue
            if "sig" not in tesgrp:
                continue
            cnum = int(k.replace("chan", ""))
            self.medianrates[cnum] = np.median(tesgrp["sig"][:, 0]/self.duration)
            try:
                self.xdet[cnum], self.ydet[cnum] = tesgrp["xy_microns"][:]
            except KeyError:  # Assume the old way
                self.xdet[cnum] = tesgrp["xdet"][0, 0]
                self.ydet[cnum] = tesgrp["ydet"][0, 0]
        self.medianrate = np.median([x for x in self.medianrates.values()])

        self.Nx, self.Xcoords = findgrid(self.target[:, 0])
        self.Ny, self.Ycoords = findgrid(self.target[:, 1])
        self._compute_good_bad_list()

    def _compute_good_bad_list(self):
        self.goodtes = []
        self.badtes = []
        for cnum, mr in self.medianrates.items():
            if abs(mr/self.medianrate-1.0) > 0.3:
                if self.verbose:
                    print("Channel {} is bad because median Pt Lα rate={:.2f} standard={:.2f}".format(
                        cnum, mr, self.medianrate))
                self.badtes.append(cnum)
            else:
                self.goodtes.append(cnum)

    def tes_count_array(self):
        "Return array `counts` where `counts[i, j]` is TES counts for good tes `i` and dwell `j`."
        return np.vstack([self.h5["tes/chan{:03d}/sig".format(c)][:, 0] for c in self.goodtes])

    def tes_rate_array(self):
        return self.tes_count_array()/self.duration

    def compute_eds_scale_factor(self):
        tes_rates = self.tes_rate_array()
        median_tes = np.median(tes_rates, axis=0)  # median over TESs
        raw_factor = np.median(median_tes/self.eds)  # median over dwell steps
        # assume absorption is due to 50-50 SiO2 and Cu, total T=3.30 µm thick
        mean_mu = (.005956 + .221777)/2  # loss per µm
        T = 3.30  # µm of total thickness
        sample_absorption = np.exp(-T/np.cos(self.angle/180*np.pi)*mean_mu)
        return raw_factor/sample_absorption

    def plot_tes_rate(self):
        plt.clf()
        plt.subplot(111, aspect="equal")
        plt.xlabel("TES X position (µm)")
        plt.ylabel("TES Y position (µm)")
        plt.title("Pt L$\\alpha$ rate (cps) {}".format(self.basename[:-5]))
        cm = plt.cm.inferno
        for k, v in self.medianrates.items():
            plt.plot(self.xdet[k], self.ydet[k], "s", ms=14, color=cm(.5*v/self.medianrate))
            if v > self.medianrate*1.5:
                txtcolor = "k"
            else:
                txtcolor = "w"
            plt.text(self.xdet[k], self.ydet[k], "{:.1f}".format(
                v), color=txtcolor, ha="center", va="center", fontsize="x-small")

    def plot_eds_summary(self):
        "Plot a summary of the EDS rate vs time for this scan."
        plt.clf()
        ax = plt.subplot(211)
        idx = np.arange(1, 1+len(self.eds))
        scale = np.median(self.eds)
        e = self.eds/scale
        de = self.eds_unc/scale

        j = 0
        y = self.target[:, 1]
        plt.plot(idx[[0, -1]], [1, 1], "k")
        while True:
            use = np.abs(y-y[j]) < 1e-4
            # plt.errorbar(idx[use], e[use], de[use], fmt=".")
            plt.plot(idx[use], e[use], ".")
            plt.plot(idx[use], 1+de[use], idx[use], 1-de[use], color="k", lw=1, alpha=0.5)

            # Fix `use` for next time
            use[:j] = True
            if np.all(use):
                break
            j = (~use).argmax()

        plt.xlabel("Scan step #")
        plt.ylabel("EDS rate relative to median")
        plt.title("SEM-side energy-dispersive sensor (EDS) rate")
        msg = "Horizontal lines show median EDS rate ± 1σ: {:.1f}±{:.1f} cps".format(
            scale, np.median(de)*scale)
        plt.text(.1, .1, msg, transform=ax.transAxes)

        plt.subplot(212, aspect="equal")
        x, y = self.target.T
        cm = plt.cm.inferno
        rate = self.eds
        vmin, vmax = np.percentile(rate, [5, 95])
        plt.scatter(x*1000, y*1000, s=1400/self.Nx, marker="s",
                    c=rate, vmin=vmin, vmax=vmax, cmap=cm)
        plt.xlabel("X location (µm)")
        plt.ylabel("Y location (µm)")
        plt.title("EDS rate vs scan position")
        plt.colorbar(location="bottom", label="EDS counts per second")

    def plot_targets(self, color=None, axis=None, showEnd=True):
        x, y = self.target.T
        if axis is None:
            plt.clf()
            axis = plt.subplot(111, aspect="equal")
            plt.xlim([x.min()-.0007, x.max()+.0007])
            plt.ylim([y.min()-.0007, y.max()+.0007])
            plt.title("Scan target locations: {}".format(self.basename))
            plt.xlabel("X location (mm)")
            plt.ylabel("Y location (mm)")

        if color is None:
            if "outer_1" in self.basename:
                color = "r"
            elif "outer_2" in self.basename:
                color = "orange"
            elif "inner_1" in self.basename:
                color = "b"
            elif "inner_2" in self.basename:
                color = "c"
            elif "inner_3" in self.basename:
                color = "g"
            else:
                color = "k"

        plt.plot(x, y, "s", color=color, ms=10, fillstyle="none")
        if showEnd:
            x, y = self.targetend.T
            plt.plot(x, y, "o", color=color, ms=4)

    def plot_timewt_drift(self, axis=None):
        "Plot time-weighted drift"
        alldrifts = []
        for d in self.drift:
            dist = 1e6*((d**2).sum()**0.5)
            dsec = np.linspace(-dist*.5, dist*.5, 65)
            alldrifts.extend(dsec)

        if axis is None:
            plt.clf()
            axis = plt.subplot(111)
            plt.grid(True)

        Nbins = 100
        plt.hist(np.abs(alldrifts), Nbins, [0, 200], histtype="step", color="green", lw=2)
        plt.xlim([0, 200])
        plt.xlabel("Position error (nm)")
        plt.ylabel("Time at error (sec per 2 nm bin)")
        plt.title("Time-weighted drift {}".format(self.basename))
        drift_quantiles = np.percentile(np.abs(alldrifts), [50, 90])
        print("Drift median: {:.2f} nm, 90%ile: {:.2f} nm".format(*drift_quantiles))
        plt.text(.5, .9, "Drift < {:.2f} nm 50% of time".format(
            drift_quantiles[0]), transform=axis.transAxes, color="r", bbox=dict(facecolor="w", alpha=0.9))
        plt.text(.5, .8, "Drift < {:.2f} nm 90% of time".format(
            drift_quantiles[1]), transform=axis.transAxes, color="r", bbox=dict(facecolor="w", alpha=0.9))

class OneSplitScan(OneScan):
    """
    Represent one raster scan on the target, split N ways
    """

    def __init__(self, filename):
        f = h5py.File(filename, "r")
        self.h5 = f
        self.filename = "/".join(filename.split(os.path.sep)[-3:])
        self.basename = os.path.basename(filename)
        self.angle = float(self.basename.split("_")[3])

        self.good = f["detInfo/goodChans"][:]

        self.Nsteps, self.Nsplits = f["eds/sig"].shape
        self.eds = f["eds/sig"][:].ravel()
        self.eds_unc = f["eds/unc"][:].ravel()

        x = f["steps/subDwellPos_sampleFixed_x"][:].ravel()
        y = f["steps/subDwellPos_sampleFixed_y"][:].ravel()
        self.target = np.vstack((x, y)).T
        commandx = np.repeat(f["steps/command_sampleFixed"][:, 0], self.Nsplits)
        commandy = np.repeat(f["steps/command_sampleFixed"][:, 1], self.Nsplits)
        self.command = np.vstack((commandx, commandy)).T
        self.offset = self.target-self.command

        # Drift has to be hacked: no sub-dwell start/end info!
        self.drift = np.roll(self.target, -1, axis=0)-self.target
        self.drift[-1, :] = 0.0

        self.duration = np.repeat(f["times/duration"][:self.Nsteps], self.Nsplits) / self.Nsplits

        self.medianrates = {}
        self.xdet = {}
        self.ydet = {}
        for k in f["tes"].keys():
            if not k.startswith("chan"):
                continue
            cnum = int(k.replace("chan", ""))
            g = f["tes"][k]
            self.medianrates[cnum] = np.median(g["sig"][:].ravel()/self.duration)
            self.xdet[cnum] = g["xdet"][0, 0]
            self.ydet[cnum] = g["ydet"][0, 0]
        self.medianrate = np.median([x for x in self.medianrates.values()])
        self.Nx, self.Xcoords = findgrid(self.target[:, 0])
        self.Ny, self.Ycoords = findgrid(self.target[:, 1])
        self._compute_good_bad_list()

def findgrid(x, mindist=2e-4):
    """
    Given a vector `x`
    """
    vmin = np.median(x[np.abs(x-x.min()) < mindist])
    v = [vmin]
    while True:
        use = x > v[-1]+mindist
        if use.sum() == 0:
            break
        v2 = (x[use]).min()
        use = np.abs(x-v2) < mindist
        if use.sum() == 0:
            break
        v2 = np.median(x[use])
        v.append(v2)
    N = len(v)
    v = np.array(v)
    return N, v

def summarize_scan(scan):
    plt.clf()
    plt.gcf().suptitle("File: {}".format(scan.filename))
    ax = plt.subplot(411)
    ax.set_aspect("equal")
    plt.title("Scan targets and commands needed to hit them")
    plt.plot(scan.target[:, 0], scan.target[:, 1], "xk", label="Targets")
    plt.plot(scan.command[:, 0], scan.command[:, 1], "o-", color="r", ms=4, label="Commands")
    # plt.plot(scan.command[:,0], scan.command[:,1], "-", color="r")
    plt.xlabel("X location (mm)")
    plt.ylabel("Y location (mm)")
    plt.legend()

    dx, dy = (scan.drift.T)*1e6
    d = np.sqrt(dx**2+dy**2)
    ax = plt.subplot(412)

    fx = np.poly1d(np.polyfit(scan.Xcoords, np.arange(scan.Nx)+0.5, 1))
    fy = np.poly1d(np.polyfit(scan.Ycoords, np.arange(scan.Ny)+0.5, 1))
    ix = np.asarray(fx(scan.target[:, 0]), dtype=int)
    iy = np.asarray(fy(scan.target[:, 1]), dtype=int)
    map = np.zeros((scan.Ny, scan.Nx), dtype=float)
    for i, (x, y) in enumerate(zip(dx, dy)):
        map[iy[i], ix[i]] = (x**2+y**2)**.5
    xstep = (scan.Xcoords[-1]-scan.Xcoords[0])/scan.Nx
    ystep = (scan.Ycoords[-1]-scan.Ycoords[0])/scan.Ny
    lft = scan.Xcoords[0]-.5*xstep
    rgt = scan.Xcoords[-1]+.5*xstep
    bot = scan.Ycoords[0]-.5*ystep
    top = scan.Ycoords[-1]+.5*ystep
    plt.subplot(412)
    plt.title("Drift during each step")
    plt.imshow(map, extent=(lft, rgt, bot, top), origin="lower", aspect="equal", cmap=plt.cm.magma)
    for i, (x, y) in enumerate(zip(dx, dy)):
        tx, ty = scan.target[i, :2]
        plt.plot([tx, tx+x/1e6], [ty, ty+y/1e6], color="white")
        plt.plot(tx, ty, "o", color="w", ms=2)

    plt.subplot(413)
    plt.title("Drift (change in offset) with each step")
    plt.plot(dx, "-o", ms=3, label="x")
    plt.plot(dy, "-o", ms=3, label="y")
    plt.plot(np.zeros_like(dx), "k")
    plt.ylabel("Δoffset (nm)")
    plt.legend()

    dist = (dx**2+dy**2)**.5
    titles = ("X drift (nm)", "Y drift (nm)", "Total drift (nm)")
    for i, d in enumerate((dx, dy, dist)):
        ax = plt.subplot(4, 4, 13+i)
        if d.min() >= 0:
            hrange = [0, 500]
        else:
            hrange = [-250, 250]
        plt.hist(d, 100, hrange, histtype="step", color="C{}".format(i))
        rms = (d**2).mean()**.5
        mad = np.median(np.abs(d))
        text = "RMS: {:.2f} nm".format(rms)
        plt.text(.9, .9, text, ha="right", transform=ax.transAxes)
        text = "MAD: {:.2f} nm".format(mad)
        plt.text(.9, .8, text, ha="right", transform=ax.transAxes)
        plt.title(titles[i])
        plt.xlim(hrange)

    ax = plt.subplot(4, 4, 16)
    plt.title("Drift")
    ax.set_aspect("equal")
    good = dist < 100
    xrange = dx.max()-dx.min()
    yrange = dy.max()-dy.min()
    if yrange > xrange and False:
        plt.plot(dy[good], dx[good], "ok", ms=3)
        plt.plot(dy[~good], dx[~good], "or", ms=3)
        plt.plot(.1, .4, marker=r'$\circlearrowleft$', ms=20, transform=ax.transAxes, color="g")
        plt.text(.03, .2, "Rotate for xy", color="g", transform=ax.transAxes)
        plt.xlabel("ΔY (nm)")
        plt.ylabel("ΔC (nm)")
    else:
        plt.plot(dx[good], dy[good], "ok", ms=3)
        plt.plot(dx[~good], dy[~good], "or", ms=3)
        plt.xlabel("ΔX (nm)")
        plt.ylabel("ΔY (nm)")
    plt.tight_layout()

eds_scale_factor = {
    -45.0: .149195,
    -37.5: .003347,
    -30.0: .003230,
    -22.5: .003117,
    -15.0: .002998,
    -7.5: .003098,
    0.0: .003039,
    7.5: .003043,
    15.0: .003028,
    22.5: .003139,
    30.0: .003102,
    37.5: .003224,
    45.0: .003182,
}

def radiograph_edges_std(angle_deg):
    xshift = np.tan(angle_deg*np.pi/180.)*.01
    x0 = -.067+xshift
    x9 = -.033+xshift
    x0 = -.07+xshift
    x9 = -.03+xshift
    y0, y9 = -.009, -.002
    return x0, x9, y0, y9

def radiograph_edges_full(xnom, ynom, magX, magY):
    # Pixelize a region that's (6.5 mm/mag) around the nominal positions
    border = 6.5
    x0 = xnom.min()-border/magX
    x9 = xnom.max()+border/magX
    y0 = ynom.min()-border/magY
    y9 = ynom.max()+border/magY
    return x0, x9, y0, y9

def compute_radiograph(scan, parity=False, rotation=0.0, mag=6340, voxsize_nm=50.0, ncombine=1, std_extent=True):
    """
    Convert a scan to its 2d radiograph
    """
    this_eds_scale_factor = eds_scale_factor[scan.angle]
    cosang = np.cos(scan.angle*np.pi/180.)
    magX = mag * cosang**2
    magY = mag * cosang
    print("Magnifications: {:.2f}, {:.2f}. EDS factor: {:.6f}".format(
        magX, magY, this_eds_scale_factor))
    voxsize_mm = voxsize_nm/1e6
    xnom = np.asarray(scan.position[:, 0])
    ynom = np.asarray(scan.position[:, 1])
    Nstep = len(xnom)
    assert Nstep == len(ynom)

    if std_extent:
        x0, x9, y0, y9 = radiograph_edges_std(scan.angle)
    else:
        x0, x9, y0, y9 = radiograph_edges_full(xnom, ynom, magX, magY)
    nx = int((x9-x0)/voxsize_mm+0.5)
    ny = int((y9-y0)/voxsize_mm+0.5)
    dx = nx*voxsize_mm
    dy = ny*voxsize_mm
    x9 = x0+dx
    y9 = y0+dy
    print("Making a ({:}x{:}) map with dx={:.2f} µm dy={:.2f} µm".format(nx, ny, dx*1e3, dy*1e3))
    if nx > 2000 or ny > 1000:
        raise ValueError("too many bins: {}x{}".format(nx, ny))

    Wt = np.zeros((nx, ny), dtype=float)
    Vt = np.zeros((nx, ny), dtype=float)
    Ct = np.zeros((nx, ny), dtype=float)
    xidx, yidx = [a for a in np.indices((nx+1, ny+1))]
    xedges = voxsize_mm*xidx+x0  # In mm
    yedges = voxsize_mm*yidx+y0  # In mm
    xpix = voxsize_mm*0.5+xedges[:-1]
    ypix = voxsize_mm*0.5+yedges[:-1]
    xpixctr = voxsize_mm*np.linspace(.5, nx-.5, nx)+x0
    ypixctr = voxsize_mm*np.linspace(.5, ny-.5, ny)+y0

    f = scan.h5
    sintheta = np.sin(rotation/57.2958)
    costheta = np.cos(rotation/57.2958)
    sig = 50
    sig2 = sig**2
    plt.figure(2)
    plt.clf()
    ax = plt.subplot(111)
    ax.set_aspect("equal")
    plt.title("Pixel positions on sample (µm) $\\theta={:.1f}^\\circ$".format(scan.angle))
    total_photons = 0.0
    for channum in scan.goodtes:
        k = "chan{:03d}".format(channum)
        g = f["tes/{}".format(k)]
        print(k)
        try:
            X = scan.xdet[channum]/1e3  # In mm
            Y = scan.ydet[channum]/1e3
            X, Y = costheta*X-sintheta*Y, sintheta*X+costheta*Y
            if parity:
                X = -X
        except IndexError:
            continue
        plt.plot(X/magX*1000, Y/magY*1000, "or")
        xtes = xnom + X/magX  # In mm
        ytes = ynom + Y/magY
        for idx in range(0, Nstep, ncombine):
            eds_total = (scan.eds[idx:idx+ncombine]*scan.duration[idx:idx+ncombine]).sum()
            expectedsig = eds_total*this_eds_scale_factor
            this_signal = g["sig"][idx:idx+ncombine].sum()/expectedsig
            sig_unc_combined = (g["sig_unc"][idx:idx+ncombine]**-2).sum()**-0.5
            this_unc = sig_unc_combined/expectedsig
            thisX = xtes[idx]
            thisY = ytes[idx]
            if this_signal <= 0:
                # print(k, idx, this_signal)
                continue

            ix = int((thisX-x0)/voxsize_mm)
            i0 = ix-2
            i9 = ix+3
            if i9 < 0:
                continue
            if i0 < 0:
                i0 = 0
            if i9 >= nx+1:
                i9 = nx
            iy = int((thisY-y0)/voxsize_mm)
            j0 = iy-2
            j9 = iy+3
            if j9 < 0:
                continue
            if j0 < 0:
                j0 = 0
            if j9 >= ny+1:
                j9 = ny
            # print(xtes[idx], xpix[ix, iy], ytes[idx], ypix[ix, iy])
            # print(xtes[idx]-xpix[ix, iy], ytes[idx]-ypix[ix, iy])
            # print((xtes[idx]-x0)/voxsize_mm, (ytes[idx]-y0)/voxsize_mm)
            # print("i,j: ", ix, i0, i9, iy, j0, j9)
            pixd2 = (thisX-xpix[i0:i9, j0:j9])**2+(thisY-ypix[i0:i9, j0:j9])**2
            wt_spot = np.exp(-0.5*(pixd2*1e12/sig2))
            wt_sn = this_unc**-2
            weight = wt_spot*wt_sn
            Vt[i0:i9, j0:j9] += weight*this_unc
            Wt[i0:i9, j0:j9] += weight
            Ct[i0:i9, j0:j9] += weight*this_signal
            total_photons += this_signal
    Wt[Wt <= 0] = Wt.min()/10
    rg = radiograph.Radiograph(Ct/Wt, Vt/Wt, Wt, xpixctr, ypixctr,
                               (xedges, yedges), scan.basename, scan.angle)
    rg.total_photons = total_photons
    return rg

def compute_8_radiographs(scan, voxsize_nm=50.0, std_extent=True):
    """
    Compute radiographs from `scan` 8 ways: w/o and w/ x-flip parity,
    and with the TES array coordinates rotated by (0, 90, 180, 270) degrees.
    Return a list of all 8, and (try to) plot them, too.
    """
    rg1 = [compute_radiograph(scan, parity=False, rotation=r, voxsize_nm=voxsize_nm,
                              std_extent=std_extent) for r in (0, 90, 180, 270)]
    rg2 = [compute_radiograph(scan, parity=True, rotation=r, voxsize_nm=voxsize_nm,
                              std_extent=std_extent) for r in (0, 90, 180, 270)]
    rg1.extend(rg2)
    try:
        plot_8_radiographs(rg1)
    except Exception:
        pass
    return rg1

def plot_8_radiographs(rgs):
    "See compute_8_radiographs()"
    midpct = 98
    plo, phi = np.nanpercentile(rgs[0].rg.ravel(), [50-midpct/2, 50+midpct/2])
    # dp = phi-plo
    vmin = plo  # -dp*5/midpct
    vmax = phi  # +dp*5/midpct
    plt.clf()
    for i, rg in enumerate(rgs):
        j = 1+2*i
        if i > 3:
            j = 2*i-6
        plt.subplot(4, 2, j, aspect="equal")
        xpix, ypix = rg.xyedges
        z = rg.rg
        plt.pcolor(xpix, -ypix, z, vmin=vmin, vmax=vmax, cmap=plt.cm.inferno)
        rotation = [0, 90, 180, 270][i % 4]
        if i > 3:
            title = "x$\\rightarrow -$x parity flip, TES plane rotation {}$^\\circ$".format(
                rotation)
            plt.yticks([])
        else:
            title = "No parity flip, TES plane rotation {}$^\\circ$".format(rotation)
        plt.title(title)
    plt.tight_layout()


import cv2
import numpy as np
from PIL import Image

facerecog = cv2.face.LBPHFaceRecognizer_create()
detector= cv2.CascadeClassifier("haarcascade_frontalface_default.xml")
facerecog.read('train2.yml')
#img = ("test_images/t1.jpg").copy()

#loading the image and converting it to gray scale
pilImage=Image.open("test_images/t1.jpg").convert('L')
#Now we are converting the PIL image into numpy array
imageNp=np.array(pilImage,'uint8')
# extract the face from the training image sample
faces=detector.detectMultiScale(imageNp)
#If a face is there then append that in the list as well as Id of it
for (x,y,w,h) in faces:
    #cv2.rectangle(img,(x-50, y-50),(x+w+50,y+h+50),(255,0,0),2)
    Id, conf = facerecog.predict(pilImage[y:y+h,x:x+w])
    print("Id = " + Id)
    break


#!/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]
'''