#!/usr/bin/env python3


#*******************************
# dokme black tanha baraye halate tasvir ghabele estefade ast na baraye halate mogheyat.
# dar halate mogheyat darsurat eshtebah dar rasm bayad kole safhe pak shavad
#*******************************

from os import environ
environ['PYGAME_HIDE_SUPPORT_PROMPT'] = '1'

import pygame
from pygame.locals import *
import time
import os
import numpy as np
import random
import cv2
from keras.models import load_model
import math

class Graphic:

    Button_state={"Color": "White"}
    x_y_of_number=[]
    data=[]

    def __init__(self):

        pygame.init()
        #infoObject = pygame.display.Info()
        #self.screen = pygame.display.set_mode((infoObject.current_w, infoObject.current_h),FULLSCREEN)
        self.screen = pygame.display.set_mode((960, 540))
        self.clock = pygame.time.Clock()
        pygame.display.set_caption('BlackBoard')

        self.image = pygame.image.load("screen_bar.jpg")
        self.black_mask=pygame.Surface((960,466))
        self.black_mask.fill((0,0,0))
        #self.image = pygame.transform.scale(self.image, (720, 480)) 
        self.screen.blit(self.image, (0, 466))

        self.radius = 10
        self.ball_surface = pygame.Surface((self.radius*2,self.radius*2))
        self.ball_surface.set_colorkey((0, 0, 0))
        self.ball_surface = self.ball_surface.convert_alpha()
        pygame.draw.circle(self.ball_surface, (0,255,0), (self.radius,self.radius), self.radius, 2)

        self.Color_ball_surface = pygame.Surface((self.radius*2,self.radius*2))
        self.Color_ball_surface.set_colorkey((0, 0, 0))
        self.Color_ball_surface = self.ball_surface.convert_alpha()
        pygame.draw.circle(self.Color_ball_surface, (255,255,255), (self.radius,self.radius), self.radius, self.radius)

        # self.myfont = pygame.font.SysFont("Arial", 22)
        # self.Line_txt_surface = self.myfont.render("1", False, (225, 255, 255))

        self.model = load_model('MNIST_CNN.h5')
        #self.model = load_model('MNIST_LSTM.h5') #RNN

        self.brush=None
        self.Hold=False
        self.Button="white"
        self.last_brush=(0,0)
        self.img_counter=0
        self.folder_name=time.time()

    def mouse_rectangle(self):

        buttons_rectangle = {}
        buttons_rectangle["Black_pos"]=[(317,502),(331,518)]
        buttons_rectangle["White_pos"]=[(334,502),(349,518)]
        buttons_rectangle["Eraser_pos"]=[(354,502),(396,518)]
        buttons_rectangle["Save_pos"]=[(227,485),(257,515)]
        
        for event in pygame.event.get():
            if(event.type==pygame.MOUSEBUTTONDOWN):

                for P in buttons_rectangle:
                    if(buttons_rectangle[P][0][0]<event.pos[0]<buttons_rectangle[P][1][0] and buttons_rectangle[P][0][1]<event.pos[1]<buttons_rectangle[P][1][1]):
                        self.Button=P.replace("_pos",'')

                if(event.button == 1): 
                    self.brush=event.pos

            elif(event.type == pygame.MOUSEBUTTONUP):
                if(event.button == 1):
                    self.brush = None

            elif(event.type==pygame.MOUSEMOTION):
                if self.brush:
                    self.brush = event.pos
        #print(self.brush)

    def Buttons_check(self):

        if(self.Button=="White"):
            Graphic.Button_state["Color"]="White"
        elif(self.Button=="Black"):
            Graphic.Button_state["Color"]="Black"
        elif(self.Button=="Eraser"):
            Graphic.Button_state["Color"]="Eraser"
        elif(self.Button=="Save"):
            Graphic.Button_state["Color"]="Save"
        self.Button="" #reset


    def draw(self):

        #print(Graphic.Button_state)

        B_P=Graphic.Balls_pos()
        self.screen.blit(self.ball_surface, B_P[Graphic.Button_state["Color"]])

        if(self.brush):
            self.screen.blit(self.Color_ball_surface,self.brush)

        if(Graphic.Button_state["Color"]=="Eraser"):
            self.screen.blit(self.black_mask,(0,0))
            Graphic.Button_state["Color"]="White"
            Graphic.x_y_of_number.clear()
            Graphic.data.clear()

        elif(Graphic.Button_state["Color"]=="White"):
            pygame.draw.circle(self.Color_ball_surface, (255,255,255), (self.radius,self.radius), self.radius, self.radius)

        elif(Graphic.Button_state["Color"]=="Black"):
            pygame.draw.circle(self.Color_ball_surface, (0,0,0), (self.radius,self.radius), self.radius, self.radius)

    def pygame_to_opencv(self):

        view_ROI = pygame.surfarray.array3d(self.screen)        
        view_ROI = view_ROI.transpose([1, 0, 2]) #convert from (width, height, channel) to (height, width, channel)
        view_ROI=view_ROI[0:466,0:]

        self.img_bgr = cv2.cvtColor(view_ROI, cv2.COLOR_RGB2BGR)
        self.img_gray= cv2.cvtColor(view_ROI,cv2.COLOR_RGB2GRAY)

        #cv2.imshow("img_gray",self.img_gray)

    def hand_writing_detection_using_pictures_CNN(self):

        contours,hierarchy = cv2.findContours(self.img_gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)#_NONE

        if len(contours)>0:
            #cv2.drawContours(self.img_bgr, contours, -1, (0,255,0), 5)
            # find the biggest countour (c) by the area
            #c = max(contours, key = cv2.contourArea)
            num_rec= 0
            for cnt in contours:
                x,y,w,h = cv2.boundingRect(cnt)
                if (w>28 and h>28):
                    num_rec+=1
                    #print(num_rec)
                    roi= self.img_gray[y:y+h,x:x+w]
                    #roi=np.pad(roi, 5, mode='constant', constant_values=0)
                    roi=roi/255 #normalization
                    #cv2.imshow('ROI',roi)
                    gray_small = cv2.resize(roi, (28, 28))
                    gray_small = gray_small.reshape(1,28,28,1)
                    pred=self.model.predict(gray_small) 
                    max_pred=np.amax(pred)
                    pred_number=np.argmax(pred,axis=1)

                    if(pred_number!=10):# and max_pred>0.90):
                        LABEL_COLOR = (0,255,0)
                        LABEL_TEXT = str(pred_number)+"="+str(max_pred)
                        cv2.rectangle(self.img_bgr,(x,y),(x+w,y+h),(255,0,0),1)
                        cv2.putText(self.img_bgr, LABEL_TEXT, (x, y-8), cv2.FONT_HERSHEY_SIMPLEX, 1/2, LABEL_COLOR, 2)

    def hand_writing_detection_using_pictures_RNN(self):

        contours,hierarchy = cv2.findContours(self.img_gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)#_NONE

        if len(contours)>0:
            #cv2.drawContours(self.img_bgr, contours, -1, (0,255,0), 5)
            # find the biggest countour (c) by the area
            #c = max(contours, key = cv2.contourArea)
            num_rec= 0
            for cnt in contours:
                x,y,w,h = cv2.boundingRect(cnt)
                if (w>28 and h>28):
                    num_rec+=1
                    #print(num_rec)
                    roi= self.img_gray[y:y+h,x:x+w]
                    #roi=np.pad(roi, 5, mode='constant', constant_values=0)
                    roi=roi/255 #normalization
                    #cv2.imshow('ROI',roi)
                    gray_small = cv2.resize(roi, (28, 28))
                    gray_small=np.expand_dims(gray_small, axis=0)
                    pred=self.model.predict(gray_small)
                    max_pred=np.amax(pred)
                    pred_number=np.argmax(pred,axis=1)

                    if(pred_number!=10):#and max_pred>0.90):
                        LABEL_COLOR = (0,255,0)
                        LABEL_TEXT = str(pred_number)+"="+str(max_pred)
                        cv2.rectangle(self.img_bgr,(x,y),(x+w,y+h),(255,0,0),1)
                        cv2.putText(self.img_bgr, LABEL_TEXT, (x, y-8), cv2.FONT_HERSHEY_SIMPLEX, 1/2, LABEL_COLOR, 2)   

    def get_hand_writing_coordinates(self):

        if(self.brush):
            if(self.brush[1]<466):
                if(Graphic.x_y_of_number):
                    if(Graphic.x_y_of_number[-1]!=self.brush):
                        self.Hold=True
                        Graphic.x_y_of_number.append(self.brush)
                else:
                    Graphic.x_y_of_number.append(self.brush)
        else:
            if(self.Hold==True):
                self.Hold=False
                Graphic.data.append(Graphic.x_y_of_number)
                Graphic.x_y_of_number=[]

        #print(len(Graphic.data),Graphic.x_y_of_number)

    def flatten(t):
        return [item for sublist in t for item in sublist]

    def save(self):
        if(Graphic.Button_state["Color"]=="Save"):
            #print(Graphic.data)
            #print(Graphic.flatten(Graphic.data))
            np.savez("dataset/normal_data"+str(time.time()),np.array(Graphic.data,dtype=object)) 
            np.savez("dataset/flatten_data"+str(time.time()),np.array(Graphic.flatten(Graphic.data),dtype=object))
            Graphic.Button_state["Color"]="Eraser"

    def Balls_pos():

        Ball_pos = {
        "White":(332,500),
        "Black":(313,500),
        "Eraser":(352,500),
        "Save":(232,490)
        }
        return Ball_pos

class Main:

    def __init__(self):
        Graphic.__init__(self)

    def run(self):
        x=0
        y=0
        c=1
        #n=0

        while True:
            
            self.screen.blit(self.image, (0, 466)) #clean_screen

            #pos=pygame.mouse.get_pos()
            keyState = pygame.key.get_pressed()
            
            Graphic.mouse_rectangle(self)
            Graphic.Buttons_check(self)
            Graphic.draw(self)
            Graphic.pygame_to_opencv(self)
            Graphic.get_hand_writing_coordinates(self)
            Graphic.save(self)

            # if (keyState[pygame.K_r]): #press R to detect
            #     Graphic.hand_writing_detection_using_coordinates(self)

            if (keyState[pygame.K_c]): #press C to detect 
                Graphic.hand_writing_detection_using_pictures_CNN(self)
                #Graphic.hand_writing_detection_using_pictures_RNN(self)

            if (keyState[pygame.K_ESCAPE] or cv2.waitKey(1) & 0xFF == ord('q')):
                pygame.quit()
                break
                #quit()

            # if (keyState[pygame.K_RIGHT]):
            #     x+=1
            # if (keyState[pygame.K_LEFT]):
            #     x-=1
            # if (keyState[pygame.K_UP]):
            #     y-=1
            # if (keyState[pygame.K_DOWN]):
            #     y+=1
            # self.screen.blit(self.ball_surface,(x,y))
            # print((x,y))

            #self.img_bgr = cv2.resize(self.img_bgr, (480, 270))
            cv2.imshow("out_put",self.img_bgr)
            milliseconds=self.clock.tick(30) # fps=30
            pygame.display.flip()
            
# Execute:
Main().run()
cv2.destroyAllWindows()