#!/usr/bin/env python3
""" Aligner for faceswap.py """

#import logging
from threading import Lock

import cv2
import numpy as np

#logger = logging.get#logger(__name__)  # pylint: disable=invalid-name

_MEAN_FACE = np.array([[0.010086, 0.106454], [0.085135, 0.038915], [0.191003, 0.018748],
                       [0.300643, 0.034489], [0.403270, 0.077391], [0.596729, 0.077391],
                       [0.699356, 0.034489], [0.808997, 0.018748], [0.914864, 0.038915],
                       [0.989913, 0.106454], [0.500000, 0.203352], [0.500000, 0.307009],
                       [0.500000, 0.409805], [0.500000, 0.515625], [0.376753, 0.587326],
                       [0.435909, 0.609345], [0.500000, 0.628106], [0.564090, 0.609345],
                       [0.623246, 0.587326], [0.131610, 0.216423], [0.196995, 0.178758],
                       [0.275698, 0.179852], [0.344479, 0.231733], [0.270791, 0.245099],
                       [0.192616, 0.244077], [0.655520, 0.231733], [0.724301, 0.179852],
                       [0.803005, 0.178758], [0.868389, 0.216423], [0.807383, 0.244077],
                       [0.729208, 0.245099], [0.264022, 0.780233], [0.350858, 0.745405],
                       [0.438731, 0.727388], [0.500000, 0.742578], [0.561268, 0.727388],
                       [0.649141, 0.745405], [0.735977, 0.780233], [0.652032, 0.864805],
                       [0.566594, 0.902192], [0.500000, 0.909281], [0.433405, 0.902192],
                       [0.347967, 0.864805], [0.300252, 0.784792], [0.437969, 0.778746],
                       [0.500000, 0.785343], [0.562030, 0.778746], [0.699747, 0.784792],
                       [0.563237, 0.824182], [0.500000, 0.831803], [0.436763, 0.824182]])

_MEAN_FACE_3D = np.array([[4.056931, -11.432347, 1.636229],   # 8 chin LL
                          [1.833492, -12.542305, 4.061275],   # 7 chin L
                          [0.0, -12.901019, 4.070434],        # 6 chin C
                          [-1.833492, -12.542305, 4.061275],  # 5 chin R
                          [-4.056931, -11.432347, 1.636229],  # 4 chin RR
                          [6.825897, 1.275284, 4.402142],     # 33 L eyebrow L
                          [1.330353, 1.636816, 6.903745],     # 29 L eyebrow R
                          [-1.330353, 1.636816, 6.903745],    # 34 R eyebrow L
                          [-6.825897, 1.275284, 4.402142],    # 38 R eyebrow R
                          [1.930245, -5.060977, 5.914376],    # 54 nose LL
                          [0.746313, -5.136947, 6.263227],    # 53 nose L
                          [0.0, -5.485328, 6.76343],          # 52 nose C
                          [-0.746313, -5.136947, 6.263227],   # 51 nose R
                          [-1.930245, -5.060977, 5.914376],   # 50 nose RR
                          [5.311432, 0.0, 3.987654],          # 13 L eye L
                          [1.78993, -0.091703, 4.413414],     # 17 L eye R
                          [-1.78993, -0.091703, 4.413414],    # 25 R eye L
                          [-5.311432, 0.0, 3.987654],         # 21 R eye R
                          [2.774015, -7.566103, 5.048531],    # 43 mouth L
                          [0.509714, -7.056507, 6.566167],    # 42 mouth top L
                          [0.0, -7.131772, 6.704956],         # 41 mouth top C
                          [-0.509714, -7.056507, 6.566167],   # 40 mouth top R
                          [-2.774015, -7.566103, 5.048531],   # 39 mouth R
                          [-0.589441, -8.443925, 6.109526],   # 46 mouth bottom R
                          [0.0, -8.601736, 6.097667],         # 45 mouth bottom C
                          [0.589441, -8.443925, 6.109526]])   # 44 mouth bottom L

_EXTRACT_RATIOS = dict(legacy=0.375, face=0.5, head=0.625)

def get_matrix_scaling(matrix):
    """ Given a matrix, return the cv2 Interpolation method and inverse interpolation method for
    applying the matrix on an image.

    Parameters
    ----------
    matrix: :class:`numpy.ndarray`
        The transform matrix to return the interpolator for

    Returns
    -------
    tuple
        The interpolator and inverse interpolator for the given matrix. This will be (Cubic, Area)
        for an upscale matrix and (Area, Cubic) for a downscale matrix
    """
    x_scale = np.sqrt(matrix[0, 0] * matrix[0, 0] + matrix[0, 1] * matrix[0, 1])
    y_scale = (matrix[0, 0] * matrix[1, 1] - matrix[0, 1] * matrix[1, 0]) / x_scale
    avg_scale = (x_scale + y_scale) * 0.5
    if avg_scale >= 1.:
        interpolators = cv2.INTER_CUBIC, cv2.INTER_AREA
    else:
        interpolators = cv2.INTER_AREA, cv2.INTER_CUBIC
    #logger.trace("interpolator: %s, inverse interpolator: %s", interpolators[0], interpolators[1])
    return interpolators

def transform_image(image, matrix, size, padding=0):
    """ Perform transformation on an image, applying the given size and padding to the matrix.

    Parameters
    ----------
    image: :class:`numpy.ndarray`
        The image to transform
    matrix: :class:`numpy.ndarray`
        The transformation matrix to apply to the image
    size: int
        The final size of the transformed image
    padding: int, optional
        The amount of padding to apply to the final image. Default: `0`

    Returns
    -------
    :class:`numpy.ndarray`
        The transformed image
    """
    '''logger.trace("image shape: %s, matrix: %s, size: %s. padding: %s",
                 image.shape, matrix, size, padding)'''
    # transform the matrix for size and padding
    mat = matrix * (size - 2 * padding)
    mat[:, 2] += padding

    # transform image
    interpolators = get_matrix_scaling(mat)
    retval = cv2.warpAffine(image, mat, (size, size), flags=interpolators[0])
    #logger.trace("transformed matrix: %s, final image shape: %s", mat, image.shape)
    return retval

class AlignedFace():
    """ Class to align a face.

    Holds the aligned landmarks and face image, as well as associated matrices and information
    about an aligned face.

    Parameters
    ----------
    landmarks: :class:`numpy.ndarray`
        The original 68 point landmarks that pertain to the given image for this face
    image: :class:`numpy.ndarray`, optional
        The original frame that contains the face that is to be aligned. Pass `None` if the aligned
        face is not to be generated, and just the co-ordinates should be calculated.
    centering: ["legacy", "face", "head"], optional
        The type of extracted face that should be loaded. "legacy" places the nose in the center of
        the image (the original method for aligning). "face" aligns for the nose to be in the
        center of the face (top to bottom) but the center of the skull for left to right. "head"
        aligns for the center of the skull (in 3D space) being the center of the extracted image,
        with the crop holding the full head. Default: `"face"`
    size: int, optional
        The size in pixels, of each edge of the final aligned face. Default: `64`
    coverage_ratio: float, optional
        The amount of the aligned image to return. A ratio of 1.0 will return the full contents of
        the aligned image. A ratio of 0.5 will return an image of the given size, but will crop to
        the central 50%% of the image.
    dtype: str, optional
        Set a data type for the final face to be returned as. Passing ``None`` will return a face
        with the same data type as the original :attr:`image`. Default: ``None``
    is_aligned_face: bool, optional
        Indicates that the :attr:`image` is an aligned face rather than a frame.
        Default: ``False``
    """
    def __init__(self, landmarks, image=None, centering="face", size=64, coverage_ratio=1.0,
                 dtype=None, is_aligned=False):
        '''logger.trace("Initializing: %s (image shape: %s, centering: '%s', size: %s, "
                     "coverage_ratio: %s, dtype: %s, is_aligned: %s)", self.__class__.__name__,
                     image if image is None else image.shape, centering, size, coverage_ratio,
                     dtype, is_aligned)'''
        self._frame_landmarks = landmarks
        self._centering = centering
        self._size = size
        self._dtype = dtype
        self._is_aligned = is_aligned
        self._matrices = dict(legacy=_umeyama(landmarks[17:], _MEAN_FACE, True)[0:2],
                              face=None,
                              head=None)
        self._padding = self._padding_from_coverage(size, coverage_ratio)

        self._cache = self._set_cache()

        self._face = self._extract_face(image)
        '''logger.trace("Initialized: %s (matrix: %s, padding: %s, face shape: %s)",
                     self.__class__.__name__, self._matrices["legacy"], self._padding,
                     self._face if self._face is None else self._face.shape)'''

    @property
    def size(self):
        """ int: The size (in pixels) of one side of the square extracted face image. """
        return self._size

    @property
    def padding(self):
        """ int: The amount of padding (in pixels) that is applied to each side of the
        extracted face image for the selected extract type. """
        return self._padding[self._centering]

    @property
    def matrix(self):
        """ :class:`numpy.ndarray`: The 3x2 transformation matrix for extracting and aligning the
        core face area out of the original frame, with no padding or sizing applied. The returned
        matrix is offset for the given :attr:`centering`. """
        if self._matrices[self._centering] is None:
            matrix = self._matrices["legacy"].copy()
            matrix[:, 2] -= self.pose.offset[self._centering]
            self._matrices[self._centering] = matrix
            #logger.trace("original matrix: %s, new matrix: %s", self._matrices["legacy"], matrix)
        return self._matrices[self._centering]

    @property
    def pose(self):
        """ :class:`lib.align.PoseEstimate`: The estimated pose in 3D space. """
        with self._cache["pose"][1]:
            if self._cache["pose"][0] is None:
                lms = cv2.transform(np.expand_dims(self._frame_landmarks, axis=1),
                                    self._matrices["legacy"]).squeeze()
                self._cache["pose"][0] = PoseEstimate(lms)
        return self._cache["pose"][0]

    @property
    def adjusted_matrix(self):
        """ :class:`numpy.ndarray`: The 3x2 transformation matrix for extracting and aligning the
        core face area out of the original frame with padding and sizing applied. """
        with self._cache["adjusted_matrix"][1]:
            if self._cache["adjusted_matrix"][0] is None:
                matrix = self.matrix.copy()
                mat = matrix * (self._size - 2 * self.padding)
                mat[:, 2] += self.padding
                #logger.trace("adjusted_matrix: %s", mat)
                self._cache["adjusted_matrix"][0] = mat
        return self._cache["adjusted_matrix"][0]

    @property
    def face(self):
        """ :class:`numpy.ndarray`: The aligned face at the given :attr:`size` at the specified
        :attr:`coverage` in the given :attr:`dtype`. If an :attr:`image` has not been provided
        then an the attribute will return ``None``. """
        return self._face

    @property
    def original_roi(self):
        """ :class:`numpy.ndarray`: The location of the extracted face box within the original
        frame. """
        with self._cache["original_roi"][1]:
            if self._cache["original_roi"][0] is None:
                roi = np.array([[0, 0],
                                [0, self._size - 1],
                                [self._size - 1, self._size - 1],
                                [self._size - 1, 0]])
                roi = np.rint(self.transform_points(roi, invert=True)).astype("int32")
                #logger.trace("original roi: %s", roi)
                self._cache["original_roi"][0] = roi
        return self._cache["original_roi"][0]

    @property
    def landmarks(self):
        """ :class:`numpy.ndarray`: The 68 point facial landmarks aligned to the extracted face
        box. """
        with self._cache["landmarks"][1]:
            if self._cache["landmarks"][0] is None:
                lms = self.transform_points(self._frame_landmarks)
                #logger.trace("aligned landmarks: %s", lms)
                self._cache["landmarks"][0] = lms
        return self._cache["landmarks"][0]

    @property
    def interpolators(self):
        """ tuple: (`interpolator` and `reverse interpolator`) for the :attr:`adjusted matrix`. """
        with self._cache["interpolators"][1]:
            if self._cache["interpolators"][0] is None:
                interpolators = get_matrix_scaling(self.adjusted_matrix)
                #logger.trace("interpolators: %s", interpolators)
                self._cache["interpolators"][0] = interpolators
        return self._cache["interpolators"][0]

    @classmethod
    def _set_cache(cls):
        """ Set the cache items.

        Items are cached so that they are only created the first time they are called.
        Each item includes a threading lock to make cache creation thread safe.

        Returns
        -------
        dict
            The Aligned Face cache
        """
        return dict(pose=[None, Lock()],
                    original_roi=[None, Lock()],
                    landmarks=[None, Lock()],
                    adjusted_matrix=[None, Lock()],
                    interpolators=[None, Lock()],
                    cropped_roi=[dict(), Lock()],
                    cropped_size=[dict(), Lock()],
                    cropped_slices=[dict(), Lock()])

    def transform_points(self, points, invert=False):
        """ Perform transformation on a series of (x, y) co-ordinates in world space into
        aligned face space.

        Parameters
        ----------
        points: :class:`numpy.ndarray`
            The points to transform
        invert: bool, optional
            ``True`` to reverse the transformation (i.e. transform the points into world space from
            aligned face space). Default: ``False``

        Returns
        -------
        :class:`numpy.ndarray`
            The transformed points
        """
        retval = np.expand_dims(points, axis=1)
        mat = cv2.invertAffineTransform(self.adjusted_matrix) if invert else self.adjusted_matrix
        retval = cv2.transform(retval, mat, retval.shape).squeeze()
        '''logger.trace("invert: %s, Original points: %s, transformed points: %s",
                     invert, points, retval)'''
        return retval

    def _extract_face(self, image):
        """ Extract the face from a source image and populate :attr:`face`. If an image is not
        provided then ``None`` is returned.

        Parameters
        ----------
        image: :class:`numpy.ndarray` or ``None``
            The original frame to extract the face from. ``None`` if the face should not be
            extracted

        Returns
        -------
        :class:`numpy.ndarray` or ``None``
            The extracted face at the given size, with the given coverage of the given dtype or
            ``None`` if no image has been provided.
        """
        if image is None:
            #logger.debug("_extract_face called without a loaded image. Returning empty face.")
            if self._is_aligned:
                raise ValueError("An aligned face must be provided if calling with "
                                 "'is_aligned=True'")
            return None
        if self._is_aligned and self._centering != "head":  # Crop out the sub face from full head
            image = self._convert_centering(image)

        if self._is_aligned and image.shape[0] != self._size:  # Resize the given aligned face
            interp = cv2.INTER_CUBIC if image.shape[0] < self._size else cv2.INTER_AREA
            retval = cv2.resize(image, (self._size, self._size), interpolation=interp)
        elif self._is_aligned:
            retval = image
        else:
            retval = transform_image(image, self.matrix, self._size, self.padding)
        retval = retval if self._dtype is None else retval.astype(self._dtype)
        return retval

    def _convert_centering(self, image):
        """ When the face being loaded is pre-aligned, the loaded image will have 'head' centering
        so it needs to be cropped out to the appropriate centering.

        This function temporarily converts this object to a full head aligned face, extracts the
        sub-cropped face to the correct centering, revers the sub crop and returns the cropped
        face.

        Parameters
        ----------
        image: :class:`numpy.ndarray`
            The original head-centered aligned image

        Returns
        -------
        :class:`numpy.ndarray`
            The aligned image with the correct centering
        """
        # Input image is sized up because of integer rounding
        src_size = self.size - (self._size * _EXTRACT_RATIOS[self._centering])
        head_size = 2 * int(np.rint(src_size / (1 - _EXTRACT_RATIOS["head"]) / 2))
        if head_size != image.shape[0]:
            interp = cv2.INTER_CUBIC if image.shape[0] < head_size else cv2.INTER_AREA
            image = cv2.resize(image, (head_size, head_size), interpolation=interp)

        # store requested size + centering whilst temporary converting to full head extract
        old_centering = self._centering
        old_size = self.size
        self._centering = "head"
        self._size = image.shape[0]

        # crop the requested centering from image
        size = self.get_cropped_size(old_centering)
        out = np.zeros((size, size, image.shape[-1]), dtype=image.dtype)
        slices = self.get_cropped_slices(old_centering)
        out[slices["out"][0], slices["out"][1], :] = image[slices["in"][0], slices["in"][1], :]

        # Revert back to the correct centering and size and reset the cache
        self._centering = old_centering
        self._size = old_size
        self._cache = self._set_cache()
        '''logger.trace("Cropped from aligned extract: (centering: %s, in shape: %s, out shape: %s)",
                     old_centering, image.shape, out.shape)'''
        return out

    @classmethod
    def _padding_from_coverage(cls, size, coverage_ratio):
        """ Return the image padding for a face from coverage_ratio set against a
            pre-padded training image.

        Parameters
        ----------
        size: int
            The final size of the aligned image in pixels
        coverage_ratio: float
            The ratio of the final image to pad to

        Returns
        -------
        dict
            The padding required, in pixels for 'head', 'face' and 'legacy' face types
        """
        retval = {_type: round((size * (coverage_ratio - (1 - _EXTRACT_RATIOS[_type]))) / 2)
                  for _type in ("legacy", "face", "head")}
        #logger.trace(retval)
        return retval

    def get_cropped_roi(self, centering):
        """ Obtain the region of interest within an aligned face set to centered coverage for
        an alternative centering

        Parameters
        ----------
        centering: ["legacy", "face"]
            The type of centering to obtain the region of interest for. "legacy" places the nose
            in the center of the image (the original method for aligning). "face" aligns for the
            nose to be in the center of the face (top to bottom) but the center of the skull for
            left to right.

        Returns
        -------
        :class:`numpy.ndarray`
            The (`left`, `top`, `right`, `bottom` location of the region of interest within an
            aligned face centered on the head for the given centering
        """
        if self._centering != "head":
            raise ValueError("Sub ROI can only be obtained from an aligned face with 'head' "
                             "centering")
        with self._cache["cropped_roi"][1]:
            if centering not in self._cache["cropped_roi"][0]:
                offset = self.pose.offset.get(centering, np.float32((0, 0)))  # legacy = 0,0
                offset -= self.pose.offset["head"]
                offset *= ((self._size - self._padding["head"]) / 2)

                center = np.rint(offset + self._size / 2).astype("int32")
                padding = self.get_cropped_size(centering) // 2
                roi = np.array([center - padding, center + padding]).ravel()
                '''logger.trace("centering: '%s', center: %s, padding: %s, sub roi: %s",
                             centering, center, padding, roi)'''
                self._cache["cropped_roi"][0][centering] = roi
        return self._cache["cropped_roi"][0][centering]

    def get_cropped_size(self, centering):
        """ Obtain the size of a cropped face from a full head centered image.

        Parameters
        ----------
        centering: ["legacy", "face"]
            The type of centering to obtain the region of interest for. "legacy" places the nose
            in the center of the image (the original method for aligning). "face" aligns for the
            nose to be in the center of the face (top to bottom) but the center of the skull for
            left to right.

        Returns
        -------
        int
           The pixel size of a sub-crop image from a full head aligned image

        Notes
        -----
        The ROI in relation to the source image is calculated by rounding the padding of one side
        to the nearest integer then applying this padding to the center of the crop, so the size
        is calculated in the same way.
        """
        if self._centering != "head":
            raise ValueError("Sub ROI can only be obtained from an aligned face with 'head' "
                             "centering")
        with self._cache["cropped_size"][1]:
            if not self._cache["cropped_size"][0].get(centering):
                src_size = self.size - (self._size * _EXTRACT_RATIOS["head"])
                size = 2 * int(np.rint(src_size / (1 - _EXTRACT_RATIOS[centering]) / 2))
                #logger.trace("centering: %s, size: %s, crop_size: %s", centering, self._size, size)
                self._cache["cropped_size"][0][centering] = size
        return self._cache["cropped_size"][0][centering]

    def get_cropped_slices(self, centering):
        """ Obtain the slices to turn a full head extract into an alternatively centered extract.

        Parameters
        ----------
        centering: ["legacy", "face"]
            The type of centering to obtain the region of interest for. "legacy" places the nose
            in the center of the image (the original method for aligning). "face" aligns for the
            nose to be in the center of the face (top to bottom) but the center of the skull for
            left to right.

        Returns
        -------
        dict
            The slices for an input full head image and output cropped image
        """
        if self._centering != "head":
            raise ValueError("Cropped slices can only be obtained from an aligned face with "
                             "'head' centering")
        with self._cache["cropped_slices"][1]:
            if not self._cache["cropped_slices"][0].get(centering):
                size = self.get_cropped_size(centering)
                roi = self.get_cropped_roi(centering)
                slice_in = [slice(max(roi[1], 0), roi[3]), slice(max(roi[0], 0), roi[2])]
                slice_out = [slice(max(roi[1] * -1, 0), size - max(0, roi[3] - self.size)),
                             slice(max(roi[0] * -1, 0), size - max(0, roi[2] - self.size))]
                self._cache["cropped_slices"][0][centering] = {"in": slice_in, "out": slice_out}
                '''logger.trace("centering: %s, cropped_slices: %s",
                             centering, self._cache["cropped_slices"][0][centering])'''
        return self._cache["cropped_slices"][0][centering]

class PoseEstimate():
    """ Estimates pose from a generic 3D head model for the given 2D face landmarks.

    Parameters
    ----------
    landmarks: :class:`numpy.ndarry`
        The original 68 point landmarks aligned to 0.0 - 1.0 range

    References
    ----------
    Head Pose Estimation using OpenCV and Dlib - https://www.learnopencv.com/tag/solvepnp/
    3D Model points - http://aifi.isr.uc.pt/Downloads/OpenGL/glAnthropometric3DModel.cpp
    """
    def __init__(self, landmarks):
        self._distortion_coefficients = np.zeros((4, 1))  # Assuming no lens distortion
        self._xyz_2d = None

        self._camera_matrix = self._get_camera_matrix()
        self._rotation, self._translation = self._solve_pnp(landmarks)
        self._offset = self._get_offset()

    @property
    def xyz_2d(self):
        """ :class:`numpy.ndarray` projected (x, y) coordinates for each x, y, z point at a
        constant distance from adjusted center of the skull (0.5, 0.5) in the 2D space. """
        if self._xyz_2d is None:
            xyz = cv2.projectPoints(np.float32([[6, 0, -2.3], [0, 6, -2.3], [0, 0, 3.7]]),
                                    self._rotation,
                                    self._translation,
                                    self._camera_matrix,
                                    self._distortion_coefficients)[0].squeeze()
            self._xyz_2d = xyz - self._offset["head"]
        return self._xyz_2d

    @property
    def offset(self):
        """ dict: The amount to offset a standard 0.0 - 1.0 umeyama transformation matrix for a
        from the center of the face (between the eyes) or center of the head (middle of skull)
        rather than the nose area. """
        return self._offset

    @classmethod
    def _get_camera_matrix(cls):
        """ Obtain an estimate of the camera matrix based off the original frame dimensions.

        Returns
        -------
        :class:`numpy.ndarray`
            An estimated camera matrix
        """
        focal_length = 4
        camera_matrix = np.array([[focal_length, 0, 0.5],
                                  [0, focal_length, 0.5],
                                  [0, 0, 1]], dtype="double")
        #logger.trace("camera_matrix: %s", camera_matrix)
        return camera_matrix

    def _solve_pnp(self, landmarks):
        """ Solve the Perspective-n-Point for the given landmarks.

        Takes 2D landmarks in world space and estimates the rotation and translation vectors
        in 3D space.

        Parameters
        ----------
        landmarks: :class:`numpy.ndarry`
            The original 68 point landmark co-ordinates relating to the original frame

        Returns
        -------
        rotation: :class:`numpy.ndarray`
            The solved rotation vector
        translation: :class:`numpy.ndarray`
            The solved translation vector
        """
        points = landmarks[[6, 7, 8, 9, 10, 17, 21, 22, 26, 31, 32, 33, 34,
                            35, 36, 39, 42, 45, 48, 50, 51, 52, 54, 56, 57, 58]]
        _, rotation, translation = cv2.solvePnP(_MEAN_FACE_3D,
                                                points,
                                                self._camera_matrix,
                                                self._distortion_coefficients,
                                                flags=cv2.SOLVEPNP_ITERATIVE)
        #logger.trace("points: %s, rotation: %s, translation: %s", points, rotation, translation)
        return rotation, translation

    def _get_offset(self):
        """ Obtain the offset between the original center of the extracted face to the new center
        of the head in 2D space.

        Returns
        -------
        :class:`numpy.ndarray`
            The x, y offset of the new center from the old center.
        """
        points = dict(head=(0, 0, -2.3), face=(0, -1.5, 4.2))
        offset = dict()
        for key, pnts in points.items():
            center = cv2.projectPoints(np.float32([pnts]),
                                       self._rotation,
                                       self._translation,
                                       self._camera_matrix,
                                       self._distortion_coefficients)[0].squeeze()
            #logger.trace("center %s: %s", key, center)
            offset[key] = center - (0.5, 0.5)
        #logger.trace("offset: %s", offset)
        return offset

def _umeyama(source, destination, estimate_scale):
    """Estimate N-D similarity transformation with or without scaling.

    Imported, and slightly adapted, directly from:
    https://github.com/scikit-image/scikit-image/blob/master/skimage/transform/_geometric.py

    Parameters
    ----------
    source: :class:`numpy.ndarray`
        (M, N) array source coordinates.
    destination: :class:`numpy.ndarray`
        (M, N) array destination coordinates.
    estimate_scale: bool
        Whether to estimate scaling factor.

    Returns
    -------
    :class:`numpy.ndarray`
        (N + 1, N + 1) The homogeneous similarity transformation matrix. The matrix contains
        NaN values only if the problem is not well-conditioned.

    References
    ----------
    .. [1] "Least-squares estimation of transformation parameters between two
            point patterns", Shinji Umeyama, PAMI 1991, :DOI:`10.1109/34.88573`
    """
    # pylint:disable=invalid-name,too-many-locals
    num = source.shape[0]
    dim = source.shape[1]

    # Compute mean of source and destination.
    src_mean = source.mean(axis=0)
    dst_mean = destination.mean(axis=0)

    # Subtract mean from source and destination.
    src_demean = source - src_mean
    dst_demean = destination - dst_mean

    # Eq. (38).
    A = dst_demean.T @ src_demean / num

    # Eq. (39).
    d = np.ones((dim,), dtype=np.double)
    if np.linalg.det(A) < 0:
        d[dim - 1] = -1

    T = np.eye(dim + 1, dtype=np.double)

    U, S, V = np.linalg.svd(A)

    # Eq. (40) and (43).
    rank = np.linalg.matrix_rank(A)
    if rank == 0:
        return np.nan * T
    if rank == dim - 1:
        if np.linalg.det(U) * np.linalg.det(V) > 0:
            T[:dim, :dim] = U @ V
        else:
            s = d[dim - 1]
            d[dim - 1] = -1
            T[:dim, :dim] = U @ np.diag(d) @ V
            d[dim - 1] = s
    else:
        T[:dim, :dim] = U @ np.diag(d) @ V

    if estimate_scale:
        # Eq. (41) and (42).
        scale = 1.0 / src_demean.var(axis=0).sum() * (S @ d)
    else:
        scale = 1.0

    T[:dim, dim] = dst_mean - scale * (T[:dim, :dim] @ src_mean.T)
    T[:dim, :dim] *= scale

    return T


import keras
import sys
import os
import numpy as np
from keras.models import load_model
from PIL import Image

imsize = (64, 64)

# モデルファイルの指定、判別するファイルは実行時の引数から取得
name = sys.argv[1]
testpic = name
keras_param = "./cnn.h5"

def load_image(path):
    # 画像サイズを都合の良いサイズに変換し、numpy用に変換する
    img = Image.open(path)
    img = img.convert('RGB')
    img = img.resize(imsize)
    img = np.asarray(img)
    img = img / 255.0
    return img

model = load_model(keras_param)
img = load_image(testpic)
prd = model.predict(np.array([img]))
# input()
print(prd)  # 精度を表示
prelabel = np.argmax(prd, axis=1)
if prelabel == 0:
    print(">>> 萌実ちゃん")
elif prelabel == 1:
    print(">>> ヨメミちゃん")


import numpy as np
import cv2
import mmcv
# from mmcv.runner import load_checkpoint
# from mmdet.models import build_detector
# from mmdet.core import multiclass_nms
from mmdet.core import get_classes
from mmdet.apis import init_detector, inference_detector

def detector_prepare(cfg_path,
                     ckpt_path,
                     is_pretrained=False):
    # read config file
    cfg = mmcv.Config.fromfile(cfg_path)
    if not is_pretrained:
        cfg.model.pretrained = None
    model = init_detector(cfg_path, ckpt_path, device='cuda:0')
    return cfg, model

def det(img, detector_model, resize_width, resize_height, score_thres=0.7):
    if isinstance(img, str):
        img = cv2.imread(img)
    elif not isinstance(img, np.ndarray):
        raise 'img type error.'

    raw_height, raw_width = img.shape[:2]
    resizeimg = cv2.resize(img, (resize_width, resize_height), interpolation=cv2.INTER_CUBIC)

    height_scale = resize_height / raw_height
    width_scale  = resize_width  / raw_width

    result = inference_detector(detector_model, resizeimg)
    labels = np.concatenate([
        np.full(bbox.shape[0], i, dtype=np.int32)
        for i, bbox in enumerate(result)
    ])
    bboxes = np.vstack(result)

    labels_list = []; bboxes_list = []
    for i, bbox in enumerate(bboxes):
        score = bbox[4]
        if score > score_thres:
            x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3]
            if x1 > 0 and y1 > 0 and x2 < resize_width - 1 and y2 < resize_height - 1:
                bboxes_list.append([x1 / width_scale, y1 / height_scale,
                                    x2 / width_scale, y2 / height_scale,
                                    score])
                labels_list.append(labels[i])
    return bboxes_list, labels_list

def nms(bboxes_list,
        labels_list,
        is_pretrained=False,
        class_names=get_classes('coco'),
        is_rcnn=True):
    '''
    non maximum suppress on detection result boxes
    input bboxes list and labels list, return ndarray of nms result
    result format: det_bboxes: [[x1, y1, x2, y2, score],...]  det_labels: [0 0 0 0 1 1 1 2 2 ...]
    '''
    # read config file
    if is_rcnn:
        cfg_path = FASTER_RCNN_CONFIG
    cfg = mmcv.Config.fromfile(cfg_path)
    if not is_pretrained:
        cfg.model.pretrained = None

    # NMS
    multi_bboxes = []
    multi_scores = []
    for i, bbox in enumerate(bboxes_list):
        # only show vehicles
        # if 2 <= (labels_list[i] + 1) <= 8:  # vehicles
        if 0 <= labels_list[
                i] <= 7:  # choose what to keep, now keep person and all vehicles
            multi_bboxes.append(bbox[0:4])
            # temp = [0 for _ in range(len(class_names))]
            temp = [0 for _ in range(3)]
            temp[labels_list[i] + 1] = bbox[4]
            multi_scores.append(temp)

    # if result is null
    if not multi_scores:
        return np.array([]), np.array([])

    if is_rcnn:
        det_bboxes, det_labels = multiclass_nms(
            torch.from_numpy(np.array(multi_bboxes).astype(np.float32)),
            torch.from_numpy(np.array(multi_scores).astype(np.float32)),
            cfg.model.test_cfg.rcnn.score_thr, cfg.model.test_cfg.rcnn.nms,
            cfg.model.test_cfg.rcnn.max_per_img)
    else:
        det_bboxes, det_labels = multiclass_nms(
            torch.from_numpy(np.array(multi_bboxes).astype(np.float32)),
            torch.from_numpy(np.array(multi_scores).astype(np.float32)),
            cfg.test_cfg.score_thr, cfg.test_cfg.nms, cfg.test_cfg.max_per_img)

    return det_bboxes.numpy(), det_labels.numpy()

def show(full_img,
        det_bboxes,
        det_labels,
        target_class,
        save_dir,
        show_scale=0.05,
        score_thres=0.):
    # show the detection result and save it
    # load full image
    import copy
    if isinstance(det_bboxes, list):
        det_bboxes = np.array(det_bboxes)
    if isinstance(det_labels, list):
        det_labels = np.array(det_labels)
    det_bboxes = copy.deepcopy(det_bboxes)
    det_labels = copy.deepcopy(det_labels)

    full_height, full_width = full_img.shape[:2]
    full_img = mmcv.imresize(full_img, (int(full_width * show_scale), int(full_height * show_scale)))

    # transfer scale of detection results
    det_bboxes[:, 0:4] *= show_scale

    # save result after NMS
    mmcv.imshow_det_bboxes(
        full_img.copy(),
        det_bboxes,
        det_labels,
        class_names=['unsure', target_class],
        score_thr=score_thres,
        out_file=save_dir,
        show=False,
        wait_time=0,
    )
    return None


from utils.core import *
import recipe_planner.utils as recipe
import recipe_planner.recipe
import numpy as np
import random

from queue import PriorityQueue

StringToGridSquare = {
        "Tomato"   : Counter,
        "Lettuce"  : Counter,
        "Onion"    : Counter,
        "Plate"    : Counter,
        "Zorg"     : Counter,
        "Cutboard" : Cutboard,
        "Delivery" : Delivery,
        }

StringToObject = {
        "Tomato"  : Tomato,
        "Lettuce" : Lettuce,
        "Onion"   : Onion,
        "Zorg"    : Zorg,
        }

class MinPriorityQueue(PriorityQueue):
    """Used for min priority queue in BRTDP algorithm."""

    def __init__(self):
        PriorityQueue.__init__(self)
        self.counter = 0
    def put(self, item, priority):
        PriorityQueue.put(self, (priority, self.counter, item))
        self.counter += 1
    def get(self, *args, **kwargs):
        _, _, item = PriorityQueue.get(self, args, kwargs)
        return item

class Stack:
    def __init__(self):
        self.items = []
    def push(self, item):
        self.items.append(item)
    def pop(self):
        return self.items.pop()
    def empty(self):
        return len(self) == 0
    def __len__(self):
        return len(self.items)

def is_smaller(p_, p):
    if type(p) is not tuple:
        return p_[0] < p
    else: return  p_ < p

def get_single_actions(env, agent):
    actions = []

    agent_locs = list(map(lambda a: a.location, env.sim_agents))

    # Check valid movement actions
    for t in [(0, 1), (0, -1), (-1, 0), (1, 0)]:
        new_loc = env.world.inbounds(tuple(np.asarray(agent.location) + np.asarray(t)))
        # Check to make sure not at boundary
        if new_loc not in agent_locs:
            gs = env.world.get_gridsquare_at(new_loc)
            # Can move into floors
            if not gs.collidable:
                actions.append(t)
            # Can interact with deliveries
            elif isinstance(gs, Delivery):
                actions.append(t)
            # Can interact with others if at least one of me or gs is holding something, or mergeable
            elif gs.holding is None and agent.holding is not None:
                actions.append(t)
            elif gs.holding is not None and isinstance(gs.holding, Object) and agent.holding is None:
                actions.append(t)
            elif gs.holding is not None and isinstance(gs.holding, Object) and\
                agent.holding is not None and mergeable(agent.holding, gs.holding):
                actions.append(t)
    # doing nothing is always possible
    actions.append((0, 0))

    return actions

def euclidean_dist(A, B):
    return np.linalg.norm(B - A)

def manhattan_dist(A, B):
    A_x, A_y = A
    B_x, B_y = B
    return float(abs(A_x - B_x) + abs(A_y - B_y))

def get_closest(obj_locations, source_loc):
    assert obj_locations, "Obj locations is probably empty? {}".format(obj_locations)

    closest_dist = np.inf
    for loc in obj_locations:
        dist = manhattan_dist(loc, source_loc)
        if dist < closest_dist:
            closest_dist = dist
    return closest_dist

def get_min_dist_between(A_locations, B_locations):
    min_dist = np.inf

    for A in A_locations:
        A_min = get_closest(B_locations, A)
        if A_min < min_dist: min_dist = A_min
    return min_dist

def get_obj(obj_string, type_, state, location=(None, None)):
    # Core.Supply objects
    if type_ == "is_supply":
        if obj_string not in StringToGridSquare:
            raise NotImplementedError("{} is not recognized.".format(obj_string))
        return StringToGridSquare[obj_string](location)

    # Core.Object(Food / Plate) objects
    elif type_ == "is_object":
        if "-" in obj_string:
            obj_strs = obj_string.split("-")
            # just getting objects
            objects = [get_obj(obj_string=s,
                type_="is_object", state=FoodState.FRESH) for s in obj_strs]
            # getting into right food env
            for i, s in enumerate(obj_strs):
                if s == "Plate":
                    continue
                objects[i] = get_obj(obj_string=s,
                        type_="is_object",
                        state=objects[i].contents[0].state_seq[-1])
            o = Object(location, objects[0].contents[0])
            for obj in objects[1:]:
                o.merge(obj.contents[0])
            return o
        elif obj_string == "Plate":
            return Object(location, Plate())
        elif obj_string in StringToObject:
            obj = StringToObject[obj_string]()
            obj.set_state(state)
            return Object(location, obj)
        else:
            raise NotImplementedError("Type {} is not recognized".format(type_))

def get_subtask_action_obj(subtask):
    if isinstance(subtask, recipe.Get):
        obj = get_obj(obj_string=subtask.args[0], type_="is_supply", state=None)
    elif isinstance(subtask, recipe.Chop):
        obj = get_obj(obj_string="Cutboard", type_="is_supply", state=None)
    elif isinstance(subtask, recipe.Deliver):
        obj = get_obj(obj_string="Delivery", type_="is_supply", state=None)
    elif isinstance(subtask, recipe.Merge):
        obj = None
    elif subtask is None:
        obj = None
    else:
        raise ValueError("Did not recognize subtask {} so could not find the appropriate subtask location".format(subtask))
    return obj

def get_subtask_obj(subtask):
    if isinstance(subtask, recipe.Chop):
        # start off raw, get chopped
        start_obj = get_obj(obj_string=subtask.args[0],
                type_="is_object", state=FoodState.FRESH)
        goal_obj = get_obj(obj_string=subtask.args[0],
                type_="is_object", state=FoodState.CHOPPED)

    elif isinstance(subtask, recipe.Merge):
        # only need in last state
        obj1 = get_obj(obj_string=subtask.args[0],
                type_="is_object", state=FoodState.FRESH)
        obj2 = get_obj(obj_string=subtask.args[1],
                type_="is_object", state=FoodState.FRESH)
        object_list = [obj1, obj2]
        start_obj = [] # initial state before merging
        # expected objects in their last food state

        for i, o in enumerate(object_list):
            if isinstance(o.contents[0], Plate):
                start_obj.append(copy.copy(o))
                continue

            object_list[i] = get_obj(obj_string=subtask.args[i],
                    type_="is_object", state=o.contents[0].state_seq[-1])
            # Must be in the last state before merging
            start_obj.append(get_obj(obj_string=subtask.args[i],
                type_="is_object", state=o.contents[0].state_seq[-1]))

        # Merging objects
        object_list[0].merge(object_list[1])
        goal_obj = object_list[0]

    elif isinstance(subtask, recipe.Deliver):
        start_obj = get_obj(obj_string=subtask.args[0],
                type_="is_object", state=FoodState.FRESH)
        # Correct the state.
        state = start_obj.contents[0].state_seq[-1] if not isinstance(start_obj.contents[0], Plate) else start_obj.contents[1].state_seq[-1]
        start_obj = get_obj(obj_string=subtask.args[0],
                type_="is_object", state=state)
        goal_obj = copy.copy(start_obj)

    elif subtask is None:
        return None, None

    else:
        raise NotImplementedError("{} was not recognized".format(subtask))

    return start_obj, goal_obj


import os

import cv2
import numpy as np
import torch
import torch.cuda.amp as amp
import torch.nn.functional as func
import torch.nn.utils
import torchvision.ops as cv_ops
import tqdm

import datasets.tools as tools
import utils.bbox_ops as bbox_ops

def train_one_epoch(model, criterion, optimizer, lr_scheduler, data_loader, epoch_idx, dist_logger):
    model.train()
    scaler = amp.GradScaler()
    processor = tqdm.tqdm(data_loader, disable=not dist_logger.is_master_rank)

    epoch_losses = []
    epoch_losses_class = []
    epoch_losses_bbox = []
    epoch_losses_centerness = []
    epoch_losses_instance_mask = []

    for img, points, targets in processor:
        img = img.cuda(non_blocking=True)
        points = points.cuda(non_blocking=True)
        targets = {key: val.cuda(non_blocking=True) for key, val in targets.items()}

        with amp.autocast():
            pred = model(img, points, targets['distance'])
            losses = criterion(points, pred, targets)

            loss_class = losses['class']
            loss_bbox = losses['bbox']
            loss_centerness = losses['centerness']
            # loss_instance_mask = losses['instance_mask']

            loss = loss_class + loss_bbox + loss_centerness  # + loss_instance_mask

        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 3.)
        scaler.step(optimizer)
        scaler.update()
        lr_scheduler.step()

        epoch_losses.append(loss.clone().detach())
        epoch_losses_class.append(loss_class.clone().detach())
        epoch_losses_bbox.append(loss_bbox.clone().detach())
        epoch_losses_centerness.append(loss_centerness.clone().detach())
        # epoch_losses_instance_mask.append(loss_instance_mask.clone().detach())

        cur_loss = dist_logger.reduce_loss(loss)
        avg_loss = dist_logger.reduce_epoch_loss(epoch_losses)
        processor.set_description(f'Epoch: {epoch_idx + 1}, cur_loss: {cur_loss:.3f}, avg_loss: {avg_loss:.3f}')

    dist_logger.save_model(model)
    # dist_logger.update_tensorboard(super_tag='loss', tag_scaler_dict={
    #     'loss': dist_logger.reduce_epoch_loss(epoch_losses),
    #     'class': dist_logger.reduce_epoch_loss(epoch_losses_class),
    #     'bbox': dist_logger.reduce_epoch_loss(epoch_losses_bbox),
    #     'centerness': dist_logger.reduce_epoch_loss(epoch_losses_centerness),
    #     # 'instance_mask': dist_logger.reduce_epoch_loss(epoch_losses_instance_mask)
    # }, idx=epoch_idx)

@torch.no_grad()
def val_one_epoch(model, data_loader, coco_gt, dist_logger, epoch_idx, nms_cfg):
    pred_instances = []
    nms_pre, cls_score_thr, iou_thr = nms_cfg['nms_pre'], nms_cfg['cls_score_thr'], nms_cfg['iou_thr']
    _, _, label_to_cat_map = tools.get_cat_label_map(coco_gt, tools.COCO_CLASSES)
    # print(label_to_cat_map)

    model.eval()
    processor = tqdm.tqdm(data_loader, disable=not dist_logger.is_master_rank)
    for img, points, img_ids in processor:
        img = img.cuda(non_blocking=True)
        points = points.cuda(non_blocking=True)
        img_info_list = coco_gt.loadImgs(img_ids.numpy())

        pred = model(img, points)
        class_pred = pred['class'].sigmoid()  # [B, num_points, num_classes]
        centerness_pred = pred['centerness'].sigmoid()  # [B, num_points]
        bbox_pred = bbox_ops.convert_distance_to_bbox(points, pred['distance'])  # [B, num_points, 4]
        # instance_mask_pred = pred['instance_mask'].sigmoid()  # [B, num_points, pooler_size, pooler_size]

        # print(class_pred.shape, centerness_pred.shape, bbox_pred.shape, instance_mask_pred.shape)
        # exit(-1)

        cls_pred_scores, cls_pred_indexes = class_pred.max(dim=-1)  # [B, num_points]

        batch_size, _, num_classes = class_pred.shape
        _, _, ih, iw = img.shape

        for batch_idx in range(batch_size):
            b_cls_pred_scores = cls_pred_scores[batch_idx]
            b_cls_pred_indexes = cls_pred_indexes[batch_idx]
            b_centerness_pred = centerness_pred[batch_idx]
            b_bbox_pred = bbox_pred[batch_idx, :]  # [num_points, 4]

            _, top_idx = (b_cls_pred_scores * b_centerness_pred).topk(nms_pre)
            top_class_pred_scores = b_cls_pred_scores[top_idx]
            top_class_pred_indexes = b_cls_pred_indexes[top_idx]
            top_centerness_pred = b_centerness_pred[top_idx]
            top_bbox_pred = b_bbox_pred[top_idx, :]  # [topk, 4]

            nms_scores = top_class_pred_scores * top_centerness_pred
            top_bbox_pred = cv_ops.clip_boxes_to_image(top_bbox_pred, size=(ih, iw))

            valid_mask = top_class_pred_scores > cls_score_thr
            valid_class_pred_scores = top_class_pred_scores[valid_mask]
            valid_class_pred_indexes = top_class_pred_indexes[valid_mask]
            valid_nms_scores = nms_scores[valid_mask]
            valid_bbox_pred = top_bbox_pred[valid_mask, :]

            keep_idx = cv_ops.batched_nms(valid_bbox_pred, valid_nms_scores, valid_class_pred_indexes, iou_thr)
            keep_class_pred_scores = valid_class_pred_scores[keep_idx]
            keep_class_pred_indexes = valid_class_pred_indexes[keep_idx]
            keep_bbox_pred = valid_bbox_pred[keep_idx, :]

            oh, ow = img_info_list[batch_idx]['height'], img_info_list[batch_idx]['width']
            keep_bbox_pred = bbox_ops.recover_bboxes(keep_bbox_pred, oh, ow, ih, iw)
            keep_bbox_pred = cv_ops.box_convert(keep_bbox_pred, in_fmt='xyxy', out_fmt='xywh')

            for cls_score, cls_idx, bbox in zip(keep_class_pred_scores, keep_class_pred_indexes, keep_bbox_pred):
                # poly = coco_mask.frPyObjects(poly.permute(1, 0).reshape(1, -1).detach().cpu().double().numpy(), oh, ow)
                # rle = coco_mask.merge(poly)
                # rle['counts'] = rle['counts'].decode('utf-8')

                pred_instances.append({
                    'image_id': int(img_ids[batch_idx]),
                    'category_id': label_to_cat_map[int(cls_idx) + 1],
                    'bbox': [float(str('%.1f' % coord)) for coord in bbox.tolist()],
                    # 'segmentation': rle,
                    'score': float(str('%.1f' % cls_score))
                })

    dist_logger.save_pred_instances_local_rank(pred_instances)
    dist_logger.save_val_file()
    dist_logger.evaluate(coco_gt)


import numpy as np
from MyDQN.logger import Logger
from MyDQN import vrep
import time
import random
import cv2 as cv

image_pix = 84  # 输入图像的维度 image_pix * image_pix 灰度图

class EnvGrasp(object):
    def __init__(self):
        self.total_success = 0
        self.total_try = 0
        self.logger = Logger('./logs_grasp')
        self.work_space = np.asarray([[-0.7, -0.3], [-0.2, 0.2], [0.1, 0.4]])

        self.object_work_space = np.asarray([[-0.6, -0.4], [-0.1, 0.1], [0.01, 0.31]])

        vrep.simxFinish(-1)
        self.sim_client = vrep.simxStart('127.0.0.1', 19997, True, True, 5000, 5)  # Connect to V-REP on port 19997

        self.current_state = []     # 2d

        self.target = []    # 3d

        self.current_handle = 0
        self.current_position = [0, 0, 0]   # 3d

        self.pre_distance = 0

        self.correct_count = 0

        if self.sim_client == -1:
            print('Failed to connect to simulation (V-REP remote API server). Exiting.')
            exit()
        else:
            print('Connected to simulation.')

    def get_state(self):
        sim_ret, self.current_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target',
                                                                vrep.simx_opmode_blocking)
        sim_ret, self.current_position = vrep.simxGetObjectPosition(self.sim_client, self.current_handle, -1,
                                                                    vrep.simx_opmode_blocking)

        self.current_position = np.asarray(self.current_position)
        # print('current position : ', self.current_position, 'target position : ', self.target)
        self.current_state = self.current_position - self.target
        frame = self.get_sensor_data()
        return self.current_state[0:2], frame, self.current_position

    def get_reward(self):
        done = 0
        if self.current_position[0] < self.work_space[0][0] or self.current_position[0] > self.work_space[0][1] \
                        or self.current_position[1] < self.work_space[1][0] or self.current_position[1] > self.work_space[1][1] \
                        or self.current_position[2] < self.work_space[2][0] or self.current_position[2] > self.work_space[2][1]:
            print('arm is out of workspace!')
            done = 1
            self.correct_count = 0
            return -1, done
        distance = np.linalg.norm(self.current_state)
        # print('distance : ', distance, ' predistance : ', self.pre_distance)
        reward = -distance
        # print(distance)
        if distance < self.pre_distance:
            reward += 0.3
        else:
            pass
        if distance < 0.015:
            self.correct_count += 1
            print('reached !')
            reward += 0.5
        else:
            self.correct_count = 0
        if self.correct_count >= 5:
            # grasp
            success = self.grasp()
            if success:
                print('grasp success ! ')
                reward = 1
            else:
                print('grasp failed !')
                reward = -1
            self.correct_count = 0
            done = 1
        self.pre_distance = distance
        return reward, done

    def reset(self):
        sim_ret, self.UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target',
                                                                   vrep.simx_opmode_blocking)
        vrep.simxStopSimulation(self.sim_client, vrep.simx_opmode_blocking)
        vrep.simxStartSimulation(self.sim_client, vrep.simx_opmode_blocking)
        time.sleep(1)
        sim_ret, self.RG2_tip_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_tip', vrep.simx_opmode_blocking)
        sim_ret, gripper_position = vrep.simxGetObjectPosition(self.sim_client, self.RG2_tip_handle, -1,
                                                               vrep.simx_opmode_blocking)
        vrep.simxSetObjectPosition(self.sim_client, self.UR5_target_handle, -1, (-0.47, 0, 0.3),
                                   vrep.simx_opmode_blocking)
        vrep.simxSetObjectOrientation(self.sim_client, self.UR5_target_handle, -1, (np.pi / 2, 0, np.pi / 2),
                                      vrep.simx_opmode_blocking)
        while gripper_position[2] > 0.4:  # V-REP bug requiring multiple starts and stops to restart
            vrep.simxStopSimulation(self.sim_client, vrep.simx_opmode_blocking)
            vrep.simxStartSimulation(self.sim_client, vrep.simx_opmode_blocking)
            time.sleep(1)
            sim_ret, gripper_position = vrep.simxGetObjectPosition(self.sim_client, self.RG2_tip_handle, -1,
                                                                   vrep.simx_opmode_blocking)
            vrep.simxSetObjectPosition(self.sim_client, self.UR5_target_handle, -1, (-0.47, 0, 0.3),
                                       vrep.simx_opmode_blocking)
            vrep.simxSetObjectOrientation(self.sim_client, self.UR5_target_handle, -1, (np.pi / 2, 0, np.pi / 2),
                                          vrep.simx_opmode_blocking)
        # print('started!')

        target_x = random.random() * (self.object_work_space[0][1] - self.object_work_space[0][0]) + self.object_work_space[0][0]
        target_y = random.random() * (self.object_work_space[1][1] - self.object_work_space[1][0]) + self.object_work_space[1][0]
        # target_z = random.random() * (self.work_space[2][1] - self.work_space[2][0]) + self.work_space[2][0]
        target_z = 0.075

        self.target = [target_x, target_y, target_z]
        # print()
        # vrep.simxPauseCommunication(self.sim_client, False)

        # orientation_y = -random.random() * np.pi + np.pi / 2

        # print()
        # vrep.simxPauseCommunication(self.sim_client, False)
        # object_orientation = [-np.pi / 2, orientation_y, -np.pi / 2]

        object_orientation = [0, 0, 0]
        curr_mesh_file = '/home/chen/stl-model/cup.obj'
        curr_shape_name = 'shape001'
        object_color = [0, 0.6, 1]
        ret_resp, ret_ints, ret_floats, ret_strings, ret_buffer = vrep.simxCallScriptFunction(self.sim_client,
                                                                                'remoteApiCommandServer',
                                                                                vrep.sim_scripttype_childscript,
                                                                                'importShape',
                                                                                [0, 0, 255, 0],
                                                                                self.target + object_orientation + object_color,
                                                                                [curr_mesh_file,
                                                                                curr_shape_name],
                                                                                bytearray(),
                                                                                vrep.simx_opmode_blocking)
        self.target = np.asarray(self.target)
        sim_ret, self.current_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target',
                                                                vrep.simx_opmode_blocking)
        sim_ret, self.current_position = vrep.simxGetObjectPosition(self.sim_client, self.current_handle, -1,
                                                                    vrep.simx_opmode_blocking)
        self.pre_distance = np.linalg.norm(self.target[0:2] - self.current_position[0:2])
        self.current_state, frame, pos = self.get_state()
        return self.current_state, frame, pos

    def open_griper(self):
        gripper_motor_velocity = 0.5
        gripper_motor_force = 20
        sim_ret, RG2_gripper_handle = vrep.simxGetObjectHandle(self.sim_client, 'RG2_openCloseJoint',
                                                               vrep.simx_opmode_blocking)
        sim_ret, gripper_joint_position = vrep.simxGetJointPosition(self.sim_client, RG2_gripper_handle,
                                                                    vrep.simx_opmode_blocking)
        vrep.simxSetJointForce(self.sim_client, RG2_gripper_handle, gripper_motor_force, vrep.simx_opmode_blocking)
        vrep.simxSetJointTargetVelocity(self.sim_client, RG2_gripper_handle, gripper_motor_velocity,
                                        vrep.simx_opmode_blocking)
        gripper_fully_opened = False
        while gripper_joint_position < 0.0536:  # Block until gripper is fully open
            sim_ret, new_gripper_joint_position = vrep.simxGetJointPosition(self.sim_client, RG2_gripper_handle,
                                                                        vrep.simx_opmode_blocking)
            if new_gripper_joint_position <= gripper_joint_position:
                return gripper_fully_opened
            gripper_joint_position = new_gripper_joint_position
        gripper_fully_opened = True
        return gripper_fully_opened

    def move_to(self, tool_position):
        sim_ret, UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target', vrep.simx_opmode_blocking)
        # UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client,'UR5_target',vrep.simx_opmode_blocking)
        sim_ret, UR5_target_position = vrep.simxGetObjectPosition(self.sim_client, UR5_target_handle, -1,
                                                                  vrep.simx_opmode_blocking)
        # print(UR5_target_position)
        move_direction = np.asarray(tool_position)
        move_magnitude = np.linalg.norm(move_direction)
        if move_magnitude == 0 or not move_magnitude == move_magnitude:
            move_step = [0, 0, 0]
            num_move_steps = 0
            print('magnitude error~!', move_magnitude, tool_position)
        else:
            move_step = 0.005 * move_direction / move_magnitude
            num_move_steps = int(np.floor(move_magnitude / 0.005))

        # print('move direction : ', move_direction, 'move magnitude : ', move_magnitude, 'move step : ', move_step, 'num : ', num_move_steps)

        for step_iter in range(num_move_steps):
            vrep.simxSetObjectPosition(self.sim_client, UR5_target_handle, -1, (UR5_target_position[0] +
                                                                                move_step[0] * min(step_iter,
                                                                                num_move_steps),
                                                                                UR5_target_position[1] +
                                                                                move_step[1] * min(step_iter,
                                                                                num_move_steps),
                                                                                UR5_target_position[2]),
                                       vrep.simx_opmode_blocking)

    def move_down(self, direction=1):
        sim_ret, UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target', vrep.simx_opmode_blocking)
        # UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client,'UR5_target',vrep.simx_opmode_blocking)
        sim_ret, UR5_target_position = vrep.simxGetObjectPosition(self.sim_client, UR5_target_handle, -1,
                                                                  vrep.simx_opmode_blocking)
        # print(UR5_target_position)
        # time.sleep(200)
        if direction == 1:
            move_direction = np.asarray([0, 0, 0.039 - UR5_target_position[2]])
            move_magnitude = UR5_target_position[2] - 0.039
        else:
            move_direction = np.asarray([0, 0, 0.301 - UR5_target_position[2]])
            move_magnitude = 0.301 - UR5_target_position[2]
        if move_magnitude == 0 or not move_magnitude == move_magnitude:
            move_step = [0, 0, 0]
            num_move_steps = 0
            print('magnitude error~!', move_magnitude)
        else:
            move_step = 0.02 * move_direction / move_magnitude
            num_move_steps = int(np.floor(move_magnitude / 0.02))
            # print('move_direction : ', move_direction, 'move_magnitude : ', move_magnitude)
            # print('move step : ', move_step, 'num : ', num_move_steps)

        for step_iter in range(num_move_steps):
            vrep.simxSetObjectPosition(self.sim_client, UR5_target_handle, -1, (UR5_target_position[0],
                                                                                UR5_target_position[1],
                                                                                UR5_target_position[2] +
                                                                                move_step[2] * (step_iter+1)),
                                                                                vrep.simx_opmode_blocking)

    def close_gripper(self):
        gripper_motor_velocity = -0.5
        gripper_motor_force = 100
        sim_ret, RG2_gripper_handle = vrep.simxGetObjectHandle(self.sim_client, 'RG2_openCloseJoint',
                                                               vrep.simx_opmode_blocking)
        sim_ret, gripper_joint_position = vrep.simxGetJointPosition(self.sim_client, RG2_gripper_handle,
                                                                    vrep.simx_opmode_blocking)
        vrep.simxSetJointForce(self.sim_client, RG2_gripper_handle, gripper_motor_force, vrep.simx_opmode_blocking)
        vrep.simxSetJointTargetVelocity(self.sim_client, RG2_gripper_handle, gripper_motor_velocity,
                                        vrep.simx_opmode_blocking)
        gripper_fully_closed = False
        while gripper_joint_position > -0.047:  # Block until gripper is fully closed
            sim_ret, new_gripper_joint_position = vrep.simxGetJointPosition(self.sim_client, RG2_gripper_handle,
                                                                            vrep.simx_opmode_blocking)
            # print(gripper_joint_position)
            if new_gripper_joint_position >= gripper_joint_position:
                return gripper_fully_closed
            gripper_joint_position = new_gripper_joint_position
        gripper_fully_closed = True
        return gripper_fully_closed

    def grasp(self):
        self.total_try += 1
        self.open_griper()
        self.move_down()
        self.close_gripper()
        self.move_down(-1)  # move up
        time.sleep(1)
        sim_ret, RG2_gripper_handle = vrep.simxGetObjectHandle(self.sim_client, 'RG2_openCloseJoint',
                                                               vrep.simx_opmode_blocking)
        sim_ret, gripper_joint_position = vrep.simxGetJointPosition(self.sim_client, RG2_gripper_handle,
                                                                    vrep.simx_opmode_blocking)
        if 0.055 >= gripper_joint_position >= -0.04:
            success = True
        else:
            success = False
        if success:
            self.total_success += 1

        accuracy = float(self.total_success) / self.total_try
        info = {'grasp success rate': accuracy}
        for tag, value in info.items():
            self.logger.scalar_summary(tag, value, step=self.total_try)

        return success

    def get_sensor_data(self, is_save=False):
        sim_ret, cam_handle = vrep.simxGetObjectHandle(self.sim_client, 'Vision_sensor_hand', vrep.simx_opmode_blocking)

        # Get color image from simulation
        sim_ret, resolution, raw_image = vrep.simxGetVisionSensorImage(self.sim_client, cam_handle, 0,
                                                                       vrep.simx_opmode_blocking)
        color_img = np.asarray(raw_image)
        # print(color_img.shape)
        color_img.shape = (resolution[1], resolution[0], 3)
        # print(color_img)
        color_img = color_img.astype(np.float) / 255

        color_img[color_img < 0] += 1

        # color_img *= 255
        color_img = np.fliplr(color_img)

        test_show_image = color_img * 255
        test_show_image = test_show_image.astype(np.uint8)

        test_show_image = cv.cvtColor(test_show_image, cv.COLOR_BGR2GRAY)

        if is_save:
            # print(test_show_image.shape)
            # test_show_image = Image.fromarray(test_show_image)
            # test_show_image.show()
            cv.imwrite('/home/chen/PycharmProjects/Reinforcement/VisualGrasp/image/' + 'grasp_success' + str(time.asctime(time.localtime(time.time()))) + '.png', test_show_image)

        resize_color = cv.resize(test_show_image, (image_pix, image_pix))

        resize_color = np.asarray(resize_color)
        # cv.imwrite('t.png', resize_color)
        resize_color = resize_color.astype(np.float) / 255.
        # np.set_printoptions(threshold=np.inf, suppress=True)
        # f = open('test.txt', 'w')
        # print(resize_color, file=f)
        # f.close()
        # print(resize_color.shape)
        return resize_color

    def step(self, action):
        # action : [x_ratio, y_ratio] ~ [-1, 1]
        move_magnitude = 0.05
        action = np.append(action, 0)
        action = action * move_magnitude
        self.move_to(action)
        self.current_state, frame, pos = self.get_state()
        reward, done = self.get_reward()
        return reward, self.current_state, done, frame, pos


import theano
import theano.sparse
from theano import tensor
from pylearn2.models.autoencoder import DenoisingAutoencoder
from pylearn2.space import VectorSpace
from theano.sparse.sandbox.sp2 import sampling_dot

from pylearn2.expr.basic import theano_norms

class SparseDenoisingAutoencoder(DenoisingAutoencoder):
    """
    Denoising autoencoder working with only sparse inputs and efficient
    reconstruction sampling

    Parameters
    ----------
    corruptor : WRITEME
    nvis : WRITEME
    nhid : WRITEME
    act_enc : WRITEME
    act_dec : WRITEME
    tied_weights : WRITEME
    irange : WRITEME
    rng : WRITEME

    References
    ----------
    Y. Dauphin, X. Glorot, Y. Bengio. Large-Scale Learning of Embeddings with
    Reconstruction Sampling. In Proceedings of the 28th International
    Conference on Machine Learning (ICML 2011).
    """

    def __init__(self):
        raise NotImplementedError(
        'This class has been deprecated since 2012.'\
        'In Feb, 2015, all historical codes are hence removed.')


from styx_msgs.msg import TrafficLight
import tensorflow as tf
import numpy
import os
import cv2
import rospy

MAX_IMAGE_WIDTH = 300
MAX_IMAGE_HEIGHT = 300
RECORD_IMAGES = False

class TLClassifier(object):
    def __init__(self):
        self.ModelGraph = None
        self.TFSession = None
        self.TrafficLights = {1: TrafficLight.RED,
                              2: TrafficLight.YELLOW,
                              3: TrafficLight.GREEN,
                              4: TrafficLight.UNKNOWN}

        PathToTheModel = os.path.dirname(os.path.realpath(__file__)) + "/frozen_graph.pb"

        NetWorkConfig = tf.ConfigProto()
        NetWorkConfig.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1

        self.ModelGraph = tf.Graph()
        with tf.Session(graph=self.ModelGraph, config=NetWorkConfig) as sess:
            self.TFsession = sess
            GraphDef = tf.GraphDef()
            with tf.gfile.GFile(PathToTheModel, 'rb') as fid:
                SerializedGraph = fid.read()
                GraphDef.ParseFromString(SerializedGraph)
                tf.import_graph_def(GraphDef, name='')

    def get_classification(self, image):
        """Determines the color of the traffic light in the image
        Args:
            image (cv::Mat): image containing the traffic light
        Returns:
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)
        """
        TrafficLightIndex, Probability = self.Predict(image)

        if TrafficLightIndex is not None:
            rospy.logdebug("class: %d, probability: %f", TrafficLightIndex, Probability)

        return TrafficLightIndex

    def Predict(self, CapturedImage, MinScore=0.5):
        image_tensor = self.ModelGraph.get_tensor_by_name('image_tensor:0')
        detection_boxes = self.ModelGraph.get_tensor_by_name('detection_boxes:0')
        detection_scores = self.ModelGraph.get_tensor_by_name('detection_scores:0')
        detection_classes = self.ModelGraph.get_tensor_by_name('detection_classes:0')

        CapturedImage = cv2.resize(CapturedImage, (MAX_IMAGE_WIDTH, MAX_IMAGE_HEIGHT))
        CapturedImage = cv2.cvtColor(CapturedImage, cv2.COLOR_BGR2RGB)

        (Boxes, Scores, Classes) = self.TFSession.run(
            [detection_boxes, detection_scores, detection_classes],
            feed_dict={image_tensor: numpy.expand_dims(CapturedImage, axis=0)})

        Scores = numpy.squeeze(Scores)
        Classes = numpy.squeeze(Classes)
        Boxes = numpy.squeeze(Boxes)

        for i, box in enumerate(Boxes):
            if Scores[i] > MinScore:
                light_class = self.TrafficLights[Classes[i]]
                rospy.logdebug("Traffic Light Class detected: %d", light_class)
                return light_class, Scores[i]

        return None, None


import numpy as np
import os

#data_dir = '/root/ssd_data/luna_segment_attribute/'
data_dir = './data/'

# attribute data index
#0 name
#1 pos_x
#2 pos_y
#3 pos_z
#4 size : < 5 samll, 5 < and 10 >, medium, 10 > big
#5 mal: 2.5 < benign, 2.5 < and 3.5 < neutral, 3.5 < malignant, range 1 ~ 5
#6 sphericiy: linear <-> ovoid <-> round, range 1 ~5
#7 margin: poorly defined <-> sharp, range 1 ~5
#8 spiculation: marked <-> None, range 1~5
#9 texture: nonsolid <-> part-solid <-> solid, range 1~5
#10 calcification: popcorn, laminated, solid, non-central, central, absent, 6 class
#11 internal_structure: soft tissue, fluid, fat, air, 4 class
#12 lobulation:marked <-> none, range 1 ~ 5
#13 subtlety: extreamely subtle <-> obvious,range 1 ~ 5
#14 hit_count

# check mal is benign, neutral, malignant
# 0 benign, 1 neutral, 2 malignant

def check_mal(mal):
    res = ''
    if (mal <= 2.5):
        res = 0
    elif (2.5 < mal <= 3.5):
        res = 1
    else:
        res = 2
    return res

#check rule and make sentneces
def rule(label, prob_list, size_rule_prob, rule_count_name, size_rule_count_name, attr_dict):
    size = label[4]
    mal = label[5]
    sphericity = label[6]
    margin = label[7]
    spiculation = label[8]
    texture = label[9]
    calcification = label[10]
    internal_structure = label[11]
    lobulation = label[12]
    subtlety = label[13]

    print ("#############################################")
    mal_res = check_mal(mal)

    #### check attribute has rule
    rule_idx = []
    size_idx = 0

    #benign rule check
    if (mal_res == 0):
        print ("mal is benign")
        if (lobulation <= 3):
            rule_idx.append(1)
        if (margin >= 3):
            rule_idx.append(2)
        if (sphericity >= 3):
            rule_idx.append(3)
        if (spiculation <= 2):
            rule_idx.append(4)
        if (subtlety <= 3):
            rule_idx.append(5)
        if (texture <= 3):
            rule_idx.append(6)
    #neutral rule check
    elif (mal_res == 1):
        print("mal is neutral")
        if (3 < lobulation and lobulation <= 4):
            rule_idx.append(1)
        if (3 > margin and lobulation >= 2):
            rule_idx.append(2)
        if (3 > sphericity and lobulation >= 2):
            rule_idx.append(3)
        if (2 < spiculation and lobulation <= 3):
            rule_idx.append(4)
        if (3 < subtlety and lobulation <= 4):
            rule_idx.append(5)
        if (3 < texture and lobulation <= 4):
            rule_idx.append(6)
    #malignant rule check
    elif (mal_res == 2):
        print("mal is malignant")
        if (calcification == 6):
            rule_idx.append(0)
        if (lobulation > 4):
            rule_idx.append(1)
        if (margin < 2):
            rule_idx.append(2)
        if (sphericity < 2):
            rule_idx.append(3)
        if (spiculation > 3):
            rule_idx.append(4)
        if (subtlety > 4):
            rule_idx.append(5)
        if (texture > 4):
            rule_idx.append(6)

    #check size rule
    if (size < 5):
        size_idx = 0
    if ( 5 <= size and size < 10 ):
        size_idx = 1
    if ( 10 < size):
        size_idx = 2

    prob_dict = {}

    #get size rule prob
    prob_dict[size_rule_count_name[size_idx]] = size_rule_prob[size_idx][mal_res]

    #get rule prob
    for i in range(len(rule_idx)):
        prob_dict[rule_count_name[rule_idx[i]]] = prob_list[mal_res][rule_idx[i]][mal_res]

    #sorting rule prob
    prob_dict_sort = sorted(prob_dict.items(), key=lambda kv: kv[1], reverse=True)

    #print top 3 rule sentence
    top_count = 0
    print('because', end=' ')
    for key, value in prob_dict_sort:
        top_count = top_count + 1
        print (attr_dict[mal_res][key], end=' ')

        if (key == 'medium' or key == 'small' or key == 'big'):
            print ( '('+str(size)+'mm)', end=' ')
        if (top_count >= 3):
            break
        else :
            print ('and', end=' ')
    print()

#count nodule per malignancy, size rule at dataset
#ex) size < 5 nodule all count, maliginant count, neutral count, benign count

def size_count_rule(label, size_rule_count_list):
    size = label[4]
    mal = label[5]

    mal_result = check_mal(mal)

    if (size < 5):
        size_rule_count_list[0][mal_result] = size_rule_count_list[0][mal_result] + 1
    if ( 5 <= size and size < 10 ):
        size_rule_count_list[1][mal_result] = size_rule_count_list[1][mal_result] + 1
    if ( 10 < size):
        size_rule_count_list[2][mal_result] = size_rule_count_list[2][mal_result] + 1

#count benign rule at dataset
def benign_count_rule(label, rule_count_list):
    size = label[4]
    mal = label[5]
    sphericity = label[6]
    margin = label[7]
    spiculation = label[8]
    texture = label[9]
    calcification = label[10]
    internal_structure = label[11]
    lobulation = label[12]
    subtlety = label[13]

    mal_result = check_mal(mal)

    if (lobulation <= 3):
        rule_count_list[1][mal_result] = rule_count_list[1][mal_result] + 1
    if (margin >= 3):
        rule_count_list[2][mal_result] = rule_count_list[2][mal_result] + 1
    if (sphericity >= 3):
        rule_count_list[3][mal_result] = rule_count_list[3][mal_result] + 1
    if (spiculation <= 2):
        rule_count_list[4][mal_result] = rule_count_list[4][mal_result] + 1
    if (subtlety <= 3):
        rule_count_list[5][mal_result] = rule_count_list[5][mal_result] + 1
    if (texture <= 3):
        rule_count_list[6][mal_result] = rule_count_list[6][mal_result] + 1

#count neutral rule at dataset
def neutral_count_rule(label, rule_count_list):
    size = label[4]
    mal = label[5]
    sphericity = label[6]
    margin = label[7]
    spiculation = label[8]
    texture = label[9]
    calcification = label[10]
    internal_structure = label[11]
    lobulation = label[12]
    subtlety = label[13]

    mal_result = check_mal(mal)

    if (3 < lobulation and lobulation <= 4):
        rule_count_list[1][mal_result] = rule_count_list[1][mal_result] + 1
    if ( 3 > margin and lobulation >= 2):
        rule_count_list[2][mal_result] = rule_count_list[2][mal_result] + 1
    if ( 3 > sphericity and lobulation >= 2):
        rule_count_list[3][mal_result] = rule_count_list[3][mal_result] + 1
    if (2 < spiculation and lobulation <= 3):
        rule_count_list[4][mal_result] = rule_count_list[4][mal_result] + 1
    if (3 < subtlety and lobulation <= 4):
        rule_count_list[5][mal_result] = rule_count_list[5][mal_result] + 1
    if (3 < texture and lobulation<= 4):
        rule_count_list[6][mal_result] = rule_count_list[6][mal_result] + 1

#count mal_count_rule at dataset
def mal_count_rule(label, rule_count_list):
    size = label[4]
    mal = label[5]
    sphericity = label[6]
    margin = label[7]
    spiculation = label[8]
    texture = label[9]
    calcification = label[10]
    internal_structure = label[11]
    lobulation = label[12]
    subtlety = label[13]

    mal_result = check_mal(mal)

    if (calcification == 6):
        rule_count_list[0][mal_result] = rule_count_list[0][mal_result] + 1
    if (lobulation > 4):
        rule_count_list[1][mal_result] = rule_count_list[1][mal_result] + 1
    if (margin < 2):
        rule_count_list[2][mal_result] = rule_count_list[2][mal_result] + 1
    if (sphericity < 2):
        rule_count_list[3][mal_result] = rule_count_list[3][mal_result] + 1
    if (spiculation > 3):
        rule_count_list[4][mal_result] = rule_count_list[4][mal_result] + 1
    if (subtlety > 4):
        rule_count_list[5][mal_result] = rule_count_list[5][mal_result] + 1
    if (texture > 4):
        rule_count_list[6][mal_result] = rule_count_list[6][mal_result] + 1
    # print (reason)

def print_label(label):
    print ("=========================================================================================================================================================")
    print ("name ", label[0], "size", label[4], "mal", label[5], "sphericity", label[6], "margin", label[7])
    print ("spiculation ", label[8], "texture", label[9], "calcification", label[10], "internal_structure", label[11], "lobulation", label[12], "subtlety", label[13])

#calc and print rule prob.
def calc_prob(mal_rule_count_list, neutral_rule_count_list, benign_rule_count_list, rule_count_name,
               mal_rule_prob, neutral_rule_prob, benign_rule_prob,
              size_rule_count_list, size_rule_count_name, size_rule_prob):
    print("==================================================")
    print("mal rule prob")
    for i in range(7):
        if (np.sum(mal_rule_count_list[i]) != 0):
            print(rule_count_name[i], " rule prob", end=' ')
            print("benign", mal_rule_count_list[i][0] / np.sum(mal_rule_count_list[i]), end=' ')
            print("neutral", mal_rule_count_list[i][1] / np.sum(mal_rule_count_list[i]), end=' ')
            print("mal", mal_rule_count_list[i][2] / np.sum(mal_rule_count_list[i]), end=' ')
            print("all count", np.sum(mal_rule_count_list[i]))
            mal_rule_prob[i] = mal_rule_count_list[i] / np.sum(mal_rule_count_list[i])
        else:
            print(rule_count_name[i], " does not have count")

    print("==================================================")
    print("neutral rule prob")
    for i in range(7):
        if (np.sum(neutral_rule_count_list[i]) != 0):
            print(rule_count_name[i], " rule prob", end=' ')
            print("benign", neutral_rule_count_list[i][0] / np.sum(neutral_rule_count_list[i]), end=' ')
            print("neutral", neutral_rule_count_list[i][1] / np.sum(neutral_rule_count_list[i]), end=' ')
            print("mal", neutral_rule_count_list[i][2] / np.sum(neutral_rule_count_list[i]), end=' ')
            print("all count", np.sum(neutral_rule_count_list[i]))
            neutral_rule_prob[i] = neutral_rule_count_list[i] / np.sum(neutral_rule_count_list[i])
        else:
            print(rule_count_name[i], " does not have count")

    print("==================================================")
    print("benign rule prob")
    for i in range(7):
        if (np.sum(benign_rule_count_list[i]) != 0):
            print(rule_count_name[i], " rule prob", end=' ')
            print("benign", benign_rule_count_list[i][0] / np.sum(benign_rule_count_list[i]), end=' ')
            print("neutral", benign_rule_count_list[i][1] / np.sum(benign_rule_count_list[i]), end=' ')
            print("mal", benign_rule_count_list[i][2] / np.sum(benign_rule_count_list[i]), end=' ')
            print("all count", np.sum(benign_rule_count_list[i]))
            benign_rule_prob[i] = benign_rule_count_list[i] / np.sum(benign_rule_count_list[i])
        else:
            print(rule_count_name[i], " does not have count")
    print("==================================================")
    print("size rule prob")
    for i in range(3):
        if (np.sum(size_rule_count_list[i]) != 0):
            print(size_rule_count_name[i], " rule prob", end=' ')
            print("benign", size_rule_count_list[i][0] / np.sum(size_rule_count_list[i]), end=' ')
            print("neutral", size_rule_count_list[i][1] / np.sum(size_rule_count_list[i]), end=' ')
            print("mal", size_rule_count_list[i][2] / np.sum(size_rule_count_list[i]), end=' ')
            print("all count", np.sum(size_rule_count_list[i]))
            size_rule_prob[i] = size_rule_count_list[i] / np.sum(size_rule_count_list[i])
        else:
            print(rule_count_name[i], " does not have count")

labels = []
nodule_labels = []

#list for count,
# 7attribute(calcification, lobulation, margin, sphericiy, spiculation, subtlety, texture)
# 3 malignamcy(benign, neutral, malignant)
mal_rule_count_list = np.zeros((7,3))
neutral_rule_count_list = np.zeros((7,3))
benign_rule_count_list = np.zeros((7,3))

mal_rule_prob = np.zeros((7,3))
neutral_rule_prob = np.zeros((7,3))
benign_rule_prob = np.zeros((7,3))

#rule count attribute key
rule_count_name = ['calcification', 'lobulation', 'margin', 'sphericity',
                   'spiculation', 'subtlety', 'texture']

#attribute rule explain word
mal_attr_dict = {'calcification': 'calcification is absent', 'lobulation': 'lobulation is none', 'margin': 'margin is Poorly defined',
                 'sphericity': 'sphericity is Linear', 'spiculation': 'spiculation is marked', 'subtlety': 'subtlety is obvious', 'texture': 'texture is solid',
                 'medium': 'size is medium', 'small': 'size is small', 'big': 'size is big'}
neutral_attr_dict = {'calcification': '', 'lobulation': 'lobulation is meidum', 'margin': 'margin is meidum',
                     'sphericity': 'sphericity is ovoid', 'spiculation': 'spiculation is meidum', 'subtlety': 'subtlety is subtle', 'texture': 'texture is part solid',
                     'medium': 'size is medium', 'small': 'size is small', 'big': 'size is big'}
benign_attr_dict = {'calcification': '', 'lobulation': 'lobulation is marked', 'margin': 'margin is poorly sharp',
                    'sphericity': 'sphericity is Round', 'spiculation': 'spiculation is marked', 'subtlety': 'subtlety is subtle', 'texture': 'texture is nonsolid',
                    'medium': 'size is medium', 'small': 'size is small', 'big': 'size is big'}

attr_dict = [benign_attr_dict, neutral_attr_dict, mal_attr_dict]

#list for size count
# 3 rule (small, medium, bing),
# 3 malignamcy(benign, neutral, malignant)a
size_rule_count_list = np.zeros((3,3))
#list for size prob
size_rule_prob = np.zeros((3,3))

size_rule_count_name = ['small', 'medium', 'big']

#load attribute
for idx in range(887):
    # l = np.load(os.path.join(data_dir, '%s_attribute.npy' % idx))

    num = str(idx)
    for i in range(3 - len(str(idx))):
        num = str('0') + num

    path = data_dir + num + '_attribute.npy'
    l = np.load(path)

    if np.all(l == 0):
        l = np.array([])
    labels.append(l)

for i, l in enumerate(labels):
    if len(l) > 0:
        for label in l:
            nodule_labels.append([np.concatenate([[i], label[1:]])])

nodule_labels = np.concatenate(nodule_labels, axis=0)

print (np.shape(nodule_labels))

#attribute score rounding
for i in range(np.shape(nodule_labels)[0]):
    nodule_labels[i][4] = round(nodule_labels[i][4], 2)
    nodule_labels[i][5] = round(nodule_labels[i][5], 2)
    nodule_labels[i][6] = round(nodule_labels[i][6], 2)
    nodule_labels[i][7] = round(nodule_labels[i][7], 2)
    nodule_labels[i][8] = round(nodule_labels[i][8], 2)
    nodule_labels[i][9] = round(nodule_labels[i][9], 2)
    nodule_labels[i][10] = round(nodule_labels[i][10], 2)
    nodule_labels[i][11] = round(nodule_labels[i][11], 2)
    nodule_labels[i][12] = round(nodule_labels[i][12], 2)
    nodule_labels[i][13] = round(nodule_labels[i][13], 2)

#count benign, neutral, malignant per rule
for i in range(np.shape(nodule_labels)[0]):
    #print_label(nodule_labels[i])
    mal_count_rule(nodule_labels[i], mal_rule_count_list)
    neutral_count_rule(nodule_labels[i], neutral_rule_count_list)
    benign_count_rule(nodule_labels[i], benign_rule_count_list)
    size_count_rule(nodule_labels[i], size_rule_count_list)

#calculate benign, neutral, malignancy prob
calc_prob(mal_rule_count_list, neutral_rule_count_list, benign_rule_count_list, rule_count_name,
            mal_rule_prob, neutral_rule_prob, benign_rule_prob,
            size_rule_count_list, size_rule_count_name, size_rule_prob)

prob_list = [benign_rule_prob, neutral_rule_prob, mal_rule_prob]
print (size_rule_prob)
print (benign_rule_prob)
print (neutral_rule_prob)
print (mal_rule_prob)

#print all noudle sentence
for i in range(np.shape(nodule_labels)[0]):
    rule(nodule_labels[i], prob_list, size_rule_prob, rule_count_name, size_rule_count_name, attr_dict)