# coding: utf-8

# In[5]:

import numpy as np
np.seterr(divide='ignore', invalid='ignore') # 0割warning消し
from numpy import sin,cos
import math 
import matplotlib.pyplot as plt
#from itertools import chain

get_ipython().magic('matplotlib inline')
import glob
import yaml
import os
import cv2
from collections import OrderedDict

class Poses:
    def __init__(self):
        self.yamlPaths = []
        self.allOriginalPoses=[]
        self.fromWhichYaml=[]
        self.fromWhichImage=[]
        self.normalizedPoses=[]
        self.normalizedSizeRatio=[]
        self.normalizedCenter=[]
        self.normalizedAngle=[]
    
    def loadSingleYaml(self, yamlPath, imgPath):
        f=open(yamlPath,'r')
        f.readline() # 最初の２行は読み飛ばさないとエラーになる
        f.readline() # 最初の２行は読み飛ばさないとエラーになる
        aYaml = yaml.load(f)
        f.close()    
        # 人数☓１８☓３がベタで並んでいる数値を、poses[人数][18][3] に変換する
        n=aYaml['sizes'][0] # 人数
        data=np.array([float(i) for i in aYaml['data'] ])
        points = np.array([data[i:i+3] for i in range(0, n*18*3,3)])
        poses = [points[i:i+18] for i in range(0, n*18,18)]
        #　allOriginalPoses　nに結合
        self.allOriginalPoses.extend( poses )
        # どのymlに含まれるか、どの画像に含まれるかを保持する
        yamlFileName, ext = os.path.splitext( os.path.basename(yamlPath) )
        for i in range(n):
            self.fromWhichYaml.append( yamlPath )        
            self.fromWhichImage.append( imgPath+yamlFileName.replace('_pose', '') +".jpeg")
        
    # load pose yamls from directory
    def loadPosesFromYamlDirectory(self, directoryPath, imgPath):
        for filenameWithPathAsString in glob.glob(directoryPath+"*.yml"):
            self.yamlPaths.append(filenameWithPathAsString)
        # 各姿勢yamlファイルを読み込む
        for yamlPath in self.yamlPaths:
            self.loadSingleYaml(yamlPath, imgPath)

    def normalizeAllPoses(self):
        for i, pose in enumerate( self.allOriginalPoses ):
            # rlhipを中心とした位置移動
            center=np.array([(pose[8][0]+pose[11][0])/2.0,(pose[8][1]+pose[11][1])/2.0])
            self.normalizedCenter.append(center) # rlhip中心位置の保持
            normalizedPose0=[]              # rlhip中心位置補正した点を格納する
            for point in pose:
                normalizedPose0.append(np.array([point[0],point[1]])-center)
            # 胴体長を基準としたサイズ移動
            size=np.linalg.norm( np.array([pose[1][0],pose[1][1]])-center )   
            self.normalizedSizeRatio.append(size)
            normalizedPose1=[]              # rlhip中心位置補正し、サイズ調整した点を格納する
            for point in normalizedPose0:
                normalizedPose1.append(point/size)
            # 胴体向きを基準とした回転
            angle=math.atan2(normalizedPose1[1][1],normalizedPose1[1][0])
            self.normalizedAngle.append(angle)
            normalizedPose2=[]              # rlhip中心位置補正し、サイズ調整し、回転補正した点を格納する
            for point in normalizedPose1:
                rotationMatrix = np.matrix( (
                    ( cos(np.pi/2+angle), sin(np.pi/2+angle)),
                    (-sin(np.pi/2+angle), cos(np.pi/2+angle))
                ) )
                normalizedPose2.append( np.ravel( np.dot(rotationMatrix, point.T) ) )
            self.normalizedPoses.append(normalizedPose2)

    def visualizePose(self, poseId):
        pose = self.normalizedPoses[poseId]
        x = []
        y = []
        for j in range(14):
            x.append(pose[j][0])
            y.append(pose[j][1])
        fig = plt.figure()
        plt.xlim([-2,2])
        plt.ylim([-2,2])
        ax = fig.add_subplot(1,1,1)
        ax.scatter(x,y)
        ax.set_title('plot')
        ax.set_xlabel('x')
        ax.set_ylabel('y')
        plt.plot([x[1], x[8]], [y[1], y[8]], color='r', linestyle='-', linewidth=2)
        plt.plot([x[1], x[11]], [y[1], y[11]], color='g', linestyle='-', linewidth=2)

    def visualizeAllPoses(self):
        for i, pose in enumerate( self.normalizedPoses ):
            self.visualizePose(i)
    
    def findClosetPair(self, sourceID):
        myNormalizedPose = self.normalizedPoses[sourceID]
        closetID = 0
        closetNormSum = 10000000000000
        for i, pose in enumerate(self.normalizedPoses):
            normSum = 0
            for j in [1,2,3,5,6,8,11]: #range(12):         
                normSum=normSum+np.linalg.norm(pose[j]-myNormalizedPose[j])
            if(normSum<closetNormSum and i!=sourceID):
                closetID = i
                closetNormSum = normSum
        return closetID
    
    def visuaizeImage(self, imageId):
        img = cv2.imread(self.fromWhichImage[imageId], cv2.IMREAD_COLOR)
        plt.figure(figsize=(10,10))
        plt.imshow(np.array(img)) 
        
    def normalizeImageAndSave(self, imageId, imgPath):
        img = cv2.imread(self.fromWhichImage[imageId], cv2.IMREAD_COLOR) 
        center = tuple(self.normalizedCenter[imageId])
        size = tuple(np.array([img.shape[1], img.shape[0]]))
        angle = self.normalizedAngle[imageId]*180.0/np.pi+90.0
        scale = 400.0/self.normalizedSizeRatio[imageId]
        rotationMatrix = cv2.getRotationMatrix2D(center, angle, scale)
        rotationImg = cv2.warpAffine(img, rotationMatrix, size, flags=cv2.INTER_CUBIC)
        [x,y]=self.normalizedCenter[imageId]
        croppedImg = rotationImg[y-500:y+0, x-150:x+150]        
        plt.figure(figsize=(10,10))
        plt.imshow(croppedImg) 
        cv2.imwrite(imgPath,croppedImg)

    def composeTwoImages(self, srcImageId, replaceImageId):
        srcImg = cv2.imread(self.fromWhichImage[srcImageId], cv2.IMREAD_COLOR) 
        replaceImg = cv2.imread(self.fromWhichImage[replaceImageId], cv2.IMREAD_COLOR) 

        plt.figure(figsize=(6,6))
        plt.imshow(cv2.cvtColor(srcImg, cv2.COLOR_BGR2RGB)) 

        [x,y]=self.normalizedCenter[replaceImageId]
        s=self.normalizedSizeRatio[replaceImageId]
        print(s)
        croppedImg = replaceImg[y-s:y, x-int(s/2):x+int(s/2)]

        maskImg = cv2.imread("/Users/jun/Desktop/new/mask.png", cv2.IMREAD_COLOR) 
        maskImg = cv2.resize(maskImg, (croppedImg.shape[0], croppedImg.shape[1]) )

        replaceImgCenter = tuple(np.array([0,int(s/2)]))
        size = tuple(np.array([croppedImg.shape[1], croppedImg.shape[0]]))
        angle = self.normalizedAngle[replaceImageId]*180.0/np.pi+90.0-(self.normalizedAngle[srcImageId]*180.0/np.pi+90.0)
        scale = 1
        rotationMatrix = cv2.getRotationMatrix2D(replaceImgCenter, angle, scale)
        rotationImg = cv2.warpAffine(croppedImg, rotationMatrix, size, flags=cv2.INTER_CUBIC)
        
        resizeRotationImg = np.zeros((srcImg.shape[0], srcImg.shape[1], 3), np.uint8)

        h, w = rotationImg.shape[:2]
        s2=self.normalizedSizeRatio[srcImageId]
        [x,y] = self.normalizedCenter[srcImageId]
        
        pts1 = np.float32( [[0,0],[w,0],[w,h],[0,h]] )
        pts2 = np.float32( [[x-s2/2,y-s2],
                            [x+s2/2,   y-s2],
                            [x+s2/2,   y],
                            [x-s2/2,y]])
        M = cv2.getPerspectiveTransform(pts1,pts2)
        overlayImg = cv2.warpPerspective(rotationImg, M,(srcImg.shape[1], srcImg.shape[0]))
        h, w = maskImg.shape[:2]
        pts1 = np.float32( [[0,0],[w,0],[w,h],[0,h]] )
        pts2 = np.float32( [[x-150,y-400],
                            [x+150, y-400],
                            [x+150,   y-100],
                            [x-150,y-100]])
        M = cv2.getPerspectiveTransform(pts1,pts2)

        maskImg = cv2.warpPerspective(maskImg, M,(srcImg.shape[1], srcImg.shape[0]))
        mask = cv2.threshold(maskImg, 128, 255, cv2.THRESH_BINARY_INV)[1]
        mask = cv2.bitwise_not(mask)
        overlayImgWithMask = cv2.bitwise_and(overlayImg, mask)
        mask = cv2.bitwise_not(mask)
        srcImg = cv2.bitwise_and(srcImg, mask)
        result = cv2.bitwise_or(srcImg, overlayImgWithMask)
        
        plt.figure(figsize=(6,6))
        plt.imshow(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) 
        


# In[6]:

aPose = Poses()
aPose.loadPosesFromYamlDirectory("/Users/jun/Desktop/new/naked_standing_women_poses/",
                                "/Users/jun/Desktop/new/naked_standing_women/")
aPose.normalizeAllPoses()
#aPose.normalizeImageAndSave(0,'/Users/jun/Download/target0.jpg')
aPose.findClosetPair(0)
aPose.composeTwoImages(0,85)


# In[ ]: