"""
证件照

# lzq #
"""

from tkinter import *
from PIL import Image,ImageTk
import cv2
import numpy as np

import UI
from VariableFrame import VariableFrame

#
class CertificateFrame(VariableFrame):
    start={}        #鼠标移动的起始坐标
    area = {}
    res_bgr = (0,0,255)
    iterCount = 1
    var1 = None
    var2 = None
    def __init__(self, canv, master=None):
        VariableFrame.__init__(self, canv, master)
        
        bcolor = UI.bg
        fcolor = UI.fg
        w = 27
        h = 1
        f = UI.font
        fy = 5
        dy = 33
        fx = 0
        sx = 80
        tx = 160
        
        i=0
        sizeLb = Label(self, text="尺寸", 
                        bg=bcolor, fg=fcolor, 
                        width=w, height=h, 
                        font=f)
        sizeLb.place(x=fx, y=fy+i*dy)
        self.var1 = StringVar()
        self.var1.set("0")
        
        i+=1
        oneRbtn = Radiobutton(self, text="一寸", 
                        bg=bcolor, fg=fcolor, 
                        selectcolor=UI.radioSelectColor,
                        width=8, height=h, 
                        font=f,
                        variable=self.var1,
                        value="one",
                        command=lambda: self.setSize(295, 413))
        oneRbtn.place(x=fx, y=fy+i*dy)
        #i+=1
        twoRbtn = Radiobutton(self, text="二寸", 
                        bg=bcolor, fg=fcolor, 
                        selectcolor=UI.radioSelectColor,
                        width=8, height=h, 
                        font=f,
                        variable=self.var1,
                        value="two",
                        command=lambda: self.setSize(413, 578))
        twoRbtn.place(x=sx, y=fy+i*dy)
        
        i+=2
        colorLb = Label(self, text="背景色", 
                        bg=bcolor, fg=fcolor, 
                        width=w, height=h, 
                        font=f)
        colorLb.place(x=fx, y=fy+i*dy)
        self.var2 = StringVar()
        self.var2.set("red")
        i+=1
        redRbtn = Radiobutton(self, text="红色",
                        bg=bcolor, fg=fcolor, selectcolor=UI.radioSelectColor,
                        width=8, height=h, 
                        font=f, 
                        variable=self.var2,
                        value="red",
                        command=lambda: self.setColor((0,0,255)))
        redRbtn.place(x=fx, y=fy+i*dy)
        
        #i+=1
        blueRbtn = Radiobutton(self, text="蓝色", 
                        bg=bcolor, fg=fcolor, selectcolor=UI.radioSelectColor,
                        width=8, height=h, 
                        font=f, 
                        variable=self.var2,
                        value="blue",
                        command=lambda: self.setColor((255,0,0)))
        blueRbtn.place(x=sx, y=fy+i*dy)
        
        #i+=1
        whiteRbtn = Radiobutton(self, text="白色", 
                        bg=bcolor, fg=fcolor, selectcolor=UI.radioSelectColor,
                        width=8, height=h, 
                        font=f, 
                        variable=self.var2,
                        value="white",
                        command=lambda: self.setColor((255,255,255)))
        whiteRbtn.place(x=tx, y=fy+i*dy)
        
        i+=2
        countScale = Scale(self,
                            from_=0,  # 设置最小值
                            to=50,  # 设置最大值
                            resolution=1,  # 设置步距值
                            orient=HORIZONTAL,  # 设置水平方向
                            label="迭代次数",
                            bg=bcolor, fg=fcolor, 
                            font=f, 
                            command=self.setCount)
        countScale.set(5)
        countScale.place(x=fx, y=fy+i*dy)
        
        i+=3
        confirmBtn = Button(self, text="确认", 
                        bg=bcolor, fg=fcolor, 
                        width=w, height=h, 
                        font=f, 
                        command=lambda: self.confirm())
        confirmBtn.place(x=fx, y=fy+i*dy)
        
        i+=4
        offerMg = Message(self,
                        text="原图背景为纯色，光线均匀，生成结果更准确",
                        width=self["width"],
                        bg=bcolor, fg=fcolor)
        offerMg.place(x=fx, y=fy+i*dy)
        
        canv.tag_bind("border", "<ButtonPress-1>", self.recordStart)
        canv.tag_bind("border", "<B1-Motion>", self.move)
    
    def setSize(self, w, h):
        """
            设置证件照尺寸
        """
        self.area["w"] = w
        self.area["h"] = h
        x = (self.canv.width-w)/2
        y = (self.canv.height-h)/2
        x,y = self.canv.realCoord(x, y)
        self.canv.delete("border")
        self.canv.create_rectangle(x, y, x+w, y+h, 
                        width=3, 
                        outline="#ee88aa",
                        fill="#ffffff",
                        stipple="gray12",
                        tags=("border"))
        self.area["x"] = x
        self.area["y"] = y
    
    def setColor(self, bgr):
        """
            设置证件照背景色
        """
        self.res_bgr = bgr

    def recordStart(self,event):
        """
            鼠标按下，记录鼠标位置
        """
        ex,ey = self.canv.realCoord(event.x, event.y)
        self.start["x"] = ex
        self.start["y"] = ey
        



    def move(self,event):
        """
            鼠标移动选择区域
        """
        ex,ey = self.canv.realCoord(event.x, event.y)
        x0 = ex-self.start["x"]
        y0 = ey-self.start["y"]
        self.canv.move("border", x0, y0)
        self.area["x"] += x0
        self.area["y"] += y0
        self.start["x"] = ex
        self.start["y"] = ey
    

    def setCount(self, scaleVar):
        """
            Scale回调函数，设置迭代次数
        """
        self.iterCount = int(scaleVar)

    def confirm(self):
        """
            确认设置完成，开始生成证件照
        """
        xx, yy = self.canv.getCoord()
        x1 = int(self.area["x"] - xx)
        y1 = int(self.area["y"] - yy)
        x2 = int(x1+self.area["w"])
        y2 = int(y1+self.area["h"])
        
        w,h = self.canv.img.size
        w = int(w*self.canv.scale)
        h = int(h*self.canv.scale)
        img = self.canv.img.resize((w, h))
        
        if x1>=0 and y1>=0 and x2<=w and y2<=h :
            if self.area["w"] == 295:
                cut_h = y1+int((y2-y1)*0.66)
                reduce = 0
                ext = int(self.area["w"]/2)
            else:
                cut_h = y1+int((y2-y1)*0.6)
                reduce = int((x2-x1)/10)
                ext = int(self.area["w"]/9)
            
            origin1 = img.crop((0, 0, w, cut_h))
            origin2 = img   #.crop((0, cut_h, w, h))
            
            origin1_cv2 = self.PIL_to_cv2(origin1)
            rect1 = (x1+reduce, y1, x2-x1-2*reduce, cut_h)
            res1_cv2 = self.create(origin1_cv2, rect1, self.res_bgr, self.iterCount)
            res1 = self.cv2_to_PIL(res1_cv2)
            
            origin2_cv2 = self.PIL_to_cv2(origin2)
            if x1-ext > 0:
                #s1 = 20
                #s1 = x1-20
                s1 = x1-ext
            else:
                s1 = int(x1/2)
                #s1 = 10
            if x2+ext < w:
                #s2 = w-20
                #s2 = x2+20
                s2 = x2+ext
            else:
                s2 = x2+int((w-x2)/2)
                #s2 = w-10
            rect2 = (s1, y1, s2-s1, h)  #(s1, 0, s2-s1, h-cut_h)
            res2_cv2 = self.create(origin2_cv2, rect2, self.res_bgr, self.iterCount)
            
            #img.paste(res1, (0,0))
            #img.paste(res2, (0,cut_h))
            #res2.paste(res1, (0,0))
            
            rows,cols = res1_cv2.shape[:2]
            for i in range(y1,rows):
                for j in range(x1,x2):
                    if tuple(res2_cv2[i,j]) != self.res_bgr and tuple(res1_cv2[i,j]) == self.res_bgr:
                        res2_cv2[i,j] = self.res_bgr
            
            res2 = self.cv2_to_PIL(res2_cv2)
            res = res2.crop((x1, y1, x2, y2))
            self.canv.delete("border")
            self.canv.refresh(res)
    
    def create(self, img, rect, bgr=(0,0,255), iterCount=1):
        """
            生成证件照。
            img为cv2.Image
        """
        mask = np.zeros(img.shape[:2], np.uint8)
        bgModel = np.zeros((1,65), np.float64)
        fgModel = np.zeros((1,65), np.float64)
        cv2.grabCut(img, mask, rect, bgModel, fgModel, iterCount, cv2.GC_INIT_WITH_RECT)
        mask2 = np.where((mask == 2) | (mask == 0), 0, 1).astype(np.uint8)
        out = img * mask2[:, :, np.newaxis]
        
        rows,cols = out.shape[:2]
        hsv = cv2.cvtColor(out,cv2.COLOR_BGR2HSV)
        #遍历替换
        for i in range(rows):
            for j in range(cols):
                if hsv[i,j].any() == 0:
                    if (i-10>=0 and hsv[i-10,j].any()==0) or (j-10>=0 and hsv[i,j-10].any()==0) or (i+10<rows and hsv[i+10,j].any()==0) or (j+10<cols and hsv[i,j+10].any()==0):
                        out[i,j] = bgr      #此处替换颜色，为BGR通道
        return out
    
    
    def cv2_to_PIL(self, img):
        """
            cv2图像转换为PIL.Image，
            img为cv2图像
        """
        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
        image = Image.fromarray(img)
        return image
        
        
    def PIL_to_cv2(self, img):
        """
            PIL.Image转换为cv2图像，
            img为PIL.Image
        """
        img_np = np.asarray(img)
        image = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
        return image