from cozmo_fsm import *
import pandas as pd
import cv2
from matplotlib import pyplot as plt
import numpy as np 
import os
from BuildData import AddTrainingSetImage
from BuildData import MakeTrainingSetFromFolders
from kNN_Cozmo import Get_KNN_IsFloor
from kNN_Cozmo import Train_KNN

#Used for storing images in a folder marked with your name
folder_name = "JD_Table"
#folder it will store images in

#Starting number for pics
picNum = 0
#IMPORTANT id value for surface
id_num = 7
#Number of images you want to take
num_images = 60
#Boolean update dataset with new images
Update_Training_Set = True
#Boolean create new data set
Create_New_Set = True

#Path to data from collected images
data_path = 'data/occupancy_v2.csv'

#JD User  Table #7 wood tabble
#9 paper
#4 blue yoga mat

#Megan User  Table #3

class CollectImages(StateMachineProgram):
    #plt.switch_backend('Agg')
    def __init__(self):
        super().__init__(viewer_crosshairs=True, cam_viewer=True)

    def start(self):
        super().start()
        #self.KNN = Train_KNN()
        global picNum, num_images, data_path, Create_New_Set
        self.max_img_num = picNum+num_images
        robot.camera.color_image_enabled = True
        if(not os.path.isdir(folder_name)): os.mkdir(folder_name)

        if(os.path.isfile(data_path) and Create_New_Set): os.remove(data_path)

    class ProjectToGround(StateNode):
        def start(self,event=None):
            super().start(event)
            camera_center = (320/2, 240/2)
            point = self.robot.kine.project_to_ground(*camera_center)
            print('Camera center point is at %4.1f, %4.1f' % (point[0], point[1]))

    class GrabPatch(StateNode):
        def start(self,event=None):
            global picNum, id_num, user_name, Update_Training_Set, num_images, data_path, Create_New_Set
            print("grabbing ", picNum)
            super().start(event)
            img = np.array(self.robot.world.latest_image.raw_image)

            # Boost green to compensate for Cozmo camera idiosyncracies
            img[:,:,1] = np.minimum(231,img[:,:,1]) * 1.10

            patch = img[105:135, 145:175, :]
            patch2 = cv2.cvtColor(patch, cv2.COLOR_RGB2BGR)
            if Update_Training_Set: AddTrainingSetImage(patch2,id_num, data_path)
            #Get_KNN_IsFloor(patch2,self.parent.KNN)
            cv2.imwrite('%s/doorpic%d.jpg' % (folder_name, picNum), patch2)
            picNum += 1

            if(picNum < self.parent.max_img_num):
                self.post_success()
            else:
                print("Done Collecting Images\n")
                if(Create_New_Set):
                    print("Adding Yoga Mat Images\n")
                    MakeTrainingSetFromFolders("JD_YogaMat/",4, data_path)
                print("Dataset Additions Complete")
                self.post_failure()

    def setup(self):
        """
    
            project: SetHeadAngle(-25) =T(1)=> forward
    
            forward: self.GrabPatch() =S=> Forward(20) =T(2)=> turn
            turn: self.GrabPatch() =S=> Turn(55) =T(2)=> forward
        """
        
        # Code generated by genfsm on Fri May  8 13:52:30 2020:
        
        project = SetHeadAngle(-25) .set_name("project") .set_parent(self)
        forward = self.GrabPatch() .set_name("forward") .set_parent(self)
        forward1 = Forward(20) .set_name("forward1") .set_parent(self)
        turn = self.GrabPatch() .set_name("turn") .set_parent(self)
        turn1 = Turn(55) .set_name("turn1") .set_parent(self)
        
        timertrans1 = TimerTrans(1) .set_name("timertrans1")
        timertrans1 .add_sources(project) .add_destinations(forward)
        
        successtrans1 = SuccessTrans() .set_name("successtrans1")
        successtrans1 .add_sources(forward) .add_destinations(forward1)
        
        timertrans2 = TimerTrans(2) .set_name("timertrans2")
        timertrans2 .add_sources(forward1) .add_destinations(turn)
        
        successtrans2 = SuccessTrans() .set_name("successtrans2")
        successtrans2 .add_sources(turn) .add_destinations(turn1)
        
        timertrans3 = TimerTrans(2) .set_name("timertrans3")
        timertrans3 .add_sources(turn1) .add_destinations(forward)
        
        return self


  # now we are thinking of doing k-nearest neighbors, or k means clustering
    # want to get the hue pixels of our images (eventually, filter responses) and send into KMeans
    # to do: get the hue pixels, create the histogram