"""

CHANGES:
Added abiity to load arbitrarily shaped craniotomies (not just circlar)
Added ability to do multiple craniotomies in one pass
Integrated with CranioDesigner, a craniotomy design package, developed by Flavio JK da Silva (4/2/2020)
"""
import pygame
from RESOURCES.GUI_elements_by_flav import *
from RESOURCES.Calculate_Plane_and_plot import *
from RobotCommands import *
import thorlabs_apt as apt
import thorlabs_apt.core as core
import threading
from drill_craniotomy2 import *
from drill_holes2 import *
import queue
import subprocess

from tkinter import filedialog
from tkinter import *
import os

import CranioHelperFunctions

red         = (255,0,0)
green       = (0,255,0)
blue        = (0,0,255)
gray        = (150,150,150)
darkgray    = (50,50,50)
lightgray   = (200,200,200)
black       = (0,0,0)
white       = (255,255,255)
yellow      = (255,255,0)
pink        = (255, 192, 203)
pink1= (255, 181, 197, 255),
pink2= (238, 169, 184, 255),
pink3= (205, 145, 158, 255),
pink4= (139, 99, 108, 255),

lightpurple  = (160,12,75)
darkpurple  = (51,5,25)

roboXs = []
roboYs = []
roboZs = []

ScrXYZ = (0,0,0) # in pixels

BREGMAscrn = [217,440,580] # in pixels (BREGMA)
BREGMArobo = [10.0, 10.0, 15.0] # in ROBOT Coordinates (BREGMA) mm
BREGMArat = [0.0, 0.0, 0.0] # in ROBOT Coordinates (BREGMA) mm
#graphic_scaleXY = 7.3333   #66 pix for 9mm dist from bregma to lambda = 7.333 pix/mm
graphic_scaleXY = 10.0  #66 pix for 9mm dist from bregma to lambda = 7.333 pix/mm
graphic_scaleZ = 5.0
Robo000scrn = [0,0,0]

Robo000scrn[0] = BREGMAscrn[0]  + BREGMArobo[0] * graphic_scaleXY
Robo000scrn[1] = BREGMAscrn[1]  - BREGMArobo[1] * graphic_scaleXY
Robo000scrn[2] = BREGMAscrn[2]  - BREGMArobo[2] * graphic_scaleZ

###########################################################################################################
# screw holes
hole_idx = 0
screw_holes_rat = [(-2.67, -5.625), (2.67,-5.625), (5, 7.68), (2.65, 9), (-2.65, 9), (-5, 7.68), (-1,12) ]
screw_holes_scrn = []

craniotomies = []
craniotomies_scrn = []
craniotomies_rat = []
###########################################################################################################
# Craniotomy trajectory is a circle of r radius, drawn on best fit plane generated by 5 j=chosen ooints on the skull
###########################################################################################################

drill_vel = 2.0
#######################################
def CranioBot_GUI_ELEMENTS(myscreen):
    global drill_vel, hole_idx, craniotomies_rat, screw_holes_rat
    # BOXES
    boxes = []
    #def __init__(self, surface, x, y, w, h, fill_color, line_color=black):

    # BUTTONS
    buttons = []
    #(self, surface, index, x, y, w, h,
    #                text,fsize = 18,
    #                face_color = (150,150,150),
    #                text_color = (0,0,0)):
    buttons.append(MyButton(myscreen,0, 500,20, 100,30,"HOME ROBOT",14))
    buttons.append(MyButton(myscreen,0, 500,60, 100,30,"TO BREGMA",14))

    buttons.append(MyButton(myscreen,1, 350,215, 120,30,"SELECT BREGMA",14))
##    buttons.append(MyButton(myscreen,11, 350,215, 120,30,"CLR POINT",14))
##    buttons.append(MyButton(myscreen,30, 350,255, 120,30,"BEST FIT PLANE",14))

    buttons.append(MyButton(myscreen,40, 60,290, 150,25,"Select Existing Craniotomy",14))
    buttons.append(MyButton(myscreen,41, 220,290, 150,25,"Create New Craniotomy Path",14))

    buttons.append(MyButton(myscreen,2, 480,520, 20,18,"0.1",12))          # Y - Up
    buttons.append(MyButton(myscreen,12, 480,500, 20,18,"1",12))          # Y - Up
    buttons.append(MyButton(myscreen,22, 480,480, 20,18,"5",12))          # Y - Up

    buttons.append(MyButton(myscreen,3, 480,580, 20,18,"0.1",12))         # Y - DN
    buttons.append(MyButton(myscreen,13, 480,600, 20,18,"1",12))
    buttons.append(MyButton(myscreen,23, 480,620, 20,18,"5",12))

    buttons.append(MyButton(myscreen,4, 450,550, 20,20,"0.1",12))         # X - L
    buttons.append(MyButton(myscreen,14, 430,550, 20,20,"1",12))         # X - L
    buttons.append(MyButton(myscreen,24, 410,550, 20,20,"5",12))         # X - L

    buttons.append(MyButton(myscreen,5, 510,550, 20,20,"0.1",12))         # X - R
    buttons.append(MyButton(myscreen,15, 530,550, 20,20,"1",12))         # X - L
    buttons.append(MyButton(myscreen,25, 550,550, 20,20,"5",12))         # X - L

    buttons.append(MyButton(myscreen,6, 600,520, 20,18,"0.1",12))         # Z - up
    buttons.append(MyButton(myscreen,16, 600,500, 20,18,"1",12))         # Z - up
    buttons.append(MyButton(myscreen,26, 600,480, 20,18,"5",12))         # Z - up

    buttons.append(MyButton(myscreen,7, 600,580, 20,18,"0.1",12))         # Z - Dn
    buttons.append(MyButton(myscreen,17, 600,600, 20,18,"1",12))         # Z - Dn
    buttons.append(MyButton(myscreen,27, 600,620, 20,18,"5",12))         # Z - Dn
    buttons.append(MyButton(myscreen,8, 100,640, 100,60,"Next Hole",22))
    buttons.append(MyButton(myscreen,9, 230,640, 100,60,"Craniotomy",22))
    #buttons.append(MyButton(myscreen,10,280,720, 100,100,"DRILL",16))

    buttons.append(MyButton(myscreen,11, 540,660, 140,60,"PAUSE",22))
    buttons.append(MyButton(myscreen,11, 20,730, 660,60,"STOP NOW!!",26,(255,0,0)))

    buttons.append(MyButton(myscreen,11, 485,670, 20,20,"+",22))
    buttons.append(MyButton(myscreen,11, 485,690, 20,20,"-",26))
    # CIRCLES (IF ANY)
    circles = []

    # LEDS (LIGHTS AND NOSE POKES)
    LEDs = []
    #LEDs.append(MyLED(myscreen,0, 15,15,  30,"OFF", red, darkgray)) # L LIGHTS

    # LABELS
    labels = []
    # def __init__(surface,x, y, label_name,label_pos, text ,fsize = 12):
##    labels.append(MyLabel(myscreen,50,20,   "1.) Align robot to steriotaxic equipment, then press 'HOME ROBOT'",16))
##    labels.append(MyLabel(myscreen,50,40,   "2.) Use screen buttons to move drill to bregma. ",16))
##    labels.append(MyLabel(myscreen,50,60,   "3.) Then press 'SELECT BREGMA'",16))
##    labels.append(MyLabel(myscreen,50,80,   "4.) Now use GUI to move Robot to set at least 4 more",16))
##    labels.append(MyLabel(myscreen,50,100,  "     points on skull. Press 'SELECT POINT'",16))
##    labels.append(MyLabel(myscreen,50,120,  "5.) Select best-fit plane through 5 points.",16))
    labels.append(MyLabel(myscreen,50,20,   "1.) Align robot to steriotaxic equipment, then press 'HOME ROBOT'",16))
    labels.append(MyLabel(myscreen,50,40,   "2.) Select species (rats or mice). ",16))
    labels.append(MyLabel(myscreen,50,60,   "3.) Use screen buttons to move drill to bregma. ",16))
    labels.append(MyLabel(myscreen,50,80,   "4.) Then press 'SELECT BREGMA'",16))
    labels.append(MyLabel(myscreen,100,200, "           Robo                      Rat",16))

    labels.append(MyLabel(myscreen,370,645, "    Craniotomy",16))
    labels.append(MyLabel(myscreen,80,710, "Modify using UP \ DN keys                        Button or Space Bar",14))

    labels.append(MyLabel(myscreen,410,530, "X",14))
    labels.append(MyLabel(myscreen,468,480, "Y",14))
    labels.append(MyLabel(myscreen,590,480, "Z",14))
    labels.append(MyLabel(myscreen,630,585, "Skull Surface",14))
    labels.append(MyLabel(myscreen,640,480, "Home Z",14)) # Z axis
    labels.append(MyLabel(myscreen,345,330, "Home XY",14)) # XY axis

    labels.append(MyLabel(myscreen,50,125,   "Rats",16))
    labels.append(MyLabel(myscreen,50,150,   "Mice",16))

    # RADIO BUTTONS
    radio_buttons = []
    #(self,screen,index, x,y,radius, ONOFF, on_color)
    radio_buttons.append(radio_button(myscreen,0,   110,130, 7, "ON", (255,200,200))) # RATS
    radio_buttons.append(radio_button(myscreen,1,   110,155, 7, "OFF", (255,200,200))) # Mice

    """AtlasGuide: Software for stereotaxic guidance using 3D CT/MRI hybrid atlases of developing mouse brains
    Xin Li,a Manisha Aggarwal,a Johnny Hsu,a Hangyi Jiang,a and Susumu Moria,b,*
    J Neurosci Methods. Author manuscript; available in PMC 2014 Oct 30.

    bregma-lambda distance on mouse ~= 4.2 mm

    """

    # INFO BOXES
    info_boxes = []
    #def __init__(self, surface,x, y, w, h, label,label_pos, text ,fsize = 14,surface_color = (255,255,255),BLINKING = False):

    # ROBO COORDS
    # Point 1
    info_boxes.append(InfoBox( myscreen,100,220,120,20,"Bregma",'LEFT','0, 0, 0'))
    # Point 2
##    info_boxes.append(InfoBox( myscreen,100,200,120,20,"Other",'LEFT','0, 0, 0'))
##    # Point 3
##    info_boxes.append(InfoBox( myscreen,100,220,120,20,"Other",'LEFT','0, 0, 0'))
##    # Point 4
##    info_boxes.append(InfoBox( myscreen,100,240,120,20,"Other",'LEFT','0, 0, 0'))
##    # Point 5
##    info_boxes.append(InfoBox( myscreen,100,260,120,20,"Other",'LEFT','0, 0, 0'))

    #RAT COORD
    # Point 1
    info_boxes.append(InfoBox( myscreen,220,220,120,20,"",'LEFT','0, 0, 0'))
    # Point 2
##    info_boxes.append(InfoBox( myscreen,220,200,120,20,"",'LEFT','0, 0, 0'))
##    # Point 3
##    info_boxes.append(InfoBox( myscreen,220,220,120,20,"",'LEFT','0, 0, 0'))
##    # Point 4
##    info_boxes.append(InfoBox( myscreen,220,240,120,20,"",'LEFT','0, 0, 0'))
##    # Point 5
##    info_boxes.append(InfoBox( myscreen,220,260,120,20,"",'LEFT','0, 0, 0'))

    # USER INPUT BOXES
    user_inputs = []
    #def __init__(self, surface,x, y, w, h, label_name,label_pos, text ,fsize = 12):

##    user_inputs.append(get_user_input( myscreen,373,300,50,17,"5.) Enter crainiotomy center X offset from Bregma (mm) ",'LEFT','0'))
##    user_inputs.append(get_user_input( myscreen,373,320,50,17,"6.) Enter crainiotomy center Y offset from Bregma (mm) ",'LEFT','-1.7'))
##    user_inputs.append(get_user_input( myscreen,373,340,50,17,"7.) Enter radius(remember to subtract drill diam)(mm)     ",'LEFT','5.0'))

    user_inputs.append(get_user_input(myscreen,450,370,50,17,"Scrn X",'TOP','435'))
    user_inputs.append(get_user_input(myscreen,530,370,50,17,"Scrn Y",'TOP','360'))
    user_inputs.append(get_user_input(myscreen,610,370,50,17,"Scrn Z",'TOP','360'))

    user_inputs.append(get_user_input(myscreen,450,410,50,17,"Drill X",'TOP','10'))
    user_inputs.append(get_user_input(myscreen,530,410,50,17,"Drill Y",'TOP','10'))
    user_inputs.append(get_user_input(myscreen,610,410,50,17,"Drill Z",'TOP','10'))

    user_inputs.append(get_user_input(myscreen,450,450,50,17,"Rat X",'TOP','10'))
    user_inputs.append(get_user_input(myscreen,530,450,50,17,"Rat Y",'TOP','10'))
    user_inputs.append(get_user_input(myscreen,610,450,50,17,"Rat Z",'TOP','10'))

    user_inputs.append(get_user_input(myscreen,360,680,50,17,"Depth (mm)",'TOP','0.0'))
    user_inputs.append(get_user_input(myscreen,430,680,50,17,"Vel",'TOP',str(drill_vel)))

    user_inputs.append(get_user_input(myscreen,20,640,50,17,"Hole #",'TOP',str(hole_idx)))
    user_inputs.append(get_user_input(myscreen,20,680,50,17," of ",'TOP',str(len(screw_holes_rat))))
    # NOTE:  Spaces are required otherwise this will also be put into info box labeled "offset.."

    # TOGGLES
    toggles = []

    return buttons, radio_buttons, boxes, LEDs, toggles, info_boxes, user_inputs, labels

class MyDrill:
    """
    MyDrill class
    button_state = "UP" or "DN"
    """
    global BREGMAscrn
    def __init__(self, surface, x, y):
        self.myscreen = surface
        self.rect = pygame.Rect(x-10,y-10,20,20)
        self.x = x
        self.y = y
        self.z = 10
        self.color = (255,0,0)

    def draw_drillXY(self, x, y):
        myscreen = self.myscreen
        col = self.color
        pygame.draw.circle(myscreen,col,(x, y),10,2)# Circle
        pygame.draw.circle(myscreen,col,(x, y),3,0) # Filled Circle
        pygame.draw.line(myscreen,col,(x-20, y),(x+20, y) ,2)
        pygame.draw.line(myscreen,col,(x, y-20),(x, y+20) ,2)
        self.rect = pygame.Rect(x-10,y-10,20,20)

    def draw_drillZ(self, x, y):
        #line(surface, color, start_pos, end_pos, width) -> Rect
        myscreen = self.myscreen
        col = self.color
        #black = (0,0,0)
        pygame.draw.line(myscreen,col,(x-20, y),(x+20, y) ,2) # ____
        pygame.draw.line(myscreen,col,(x-5, y-10),(x, y) ,2)  #  \
        pygame.draw.line(myscreen,col,(x+5, y-10),(x, y) ,2)  #   /

        # SKULL SURFACE
        pygame.draw.line(myscreen,black,(630, BREGMAscrn[2]+2),(690, BREGMAscrn[2]+2) ,5) # SKULL TOP

def draw_circle_craniotomy(myscreen,user_inputs): # User inputs are offset from bregma and Raduis in mm
    global graphic_scaleXY,BREGMAscrn

    x = BREGMAscrn[0]
    z = BREGMAscrn[2]

    for usr_input in user_inputs:
        if "X offset" in usr_input.label:
            #x = int(BregmaScrn[0] - float(usr_input.text)*graphic_scaleXY )
            x = int(x - float(usr_input.text)*graphic_scaleXY )
        if "Y offset" in usr_input.label:
            y = int(BREGMAscrn[1] - float(usr_input.text)*graphic_scaleXY )
        if "radius" in usr_input.label:
            r = float(usr_input.text)*graphic_scaleXY #pix/mm
    radius = int(r)
    pygame.draw.circle(myscreen,(0,0,255),(x, y),radius,1)

def draw_screw_holes(myscreen):
    global screw_holes_rat, BREGMArobo, BREGMAscrn,graphic_scaleXY
    for hole in screw_holes_rat:
        location_robo = (int(hole[0]+ BREGMArobo[0]),
                         int(hole[1] + BREGMArobo[1]))# in Robot coords
        location_scrn = (int(hole[0]*graphic_scaleXY + BREGMAscrn[0]),
                         int(hole[1] *graphic_scaleXY + BREGMAscrn[1]))
        pygame.draw.circle(myscreen,(0,255,0),location_scrn,2,0)

def draw_screw_hole(myscreen, hole, color):
    global BREGMArobo, BREGMAscrn,graphic_scaleXY

##    location_robo = (int(hole[0]+ BREGMArobo[0]),
##                     int(hole[1] + BREGMArobo[1]))# in Robot coords
    location_scrn = (int(BREGMAscrn[0] - hole[0]*graphic_scaleXY ),
                     int(BREGMAscrn[1] + hole[1] *graphic_scaleXY))
    pygame.draw.circle(myscreen,color,location_scrn,2,0)

#################################
#  MAIN PROGRAM

def setupGUI():
    # FLAGS
    BOX_SELECTED = False
    LEFT_MOUSE_DOWN = False
    RIGHT_MOUSE_DOWN = False
    BUTTON_SELECTED = False
    MOVING_DRILL =  False
    ROBOT_HOMED = False
    ROBOT_CONNECTED = False
    BREGMA_SELECTED = False
    DRILLING_HOLE = False
    SETTING_UP = True
    UP_KEY = False
    DN_KEY = False
    MOVED_UP = False
    MOVED_DN = False
    MOVED_XY = False
    MOVING_UP = None
    MOVING_XY = None
    MOVING_DN = None
    MOVING = None
    RATS = True
    MICE = False
    CIRC_CRANIOTOMY = False
    circle = None
    queStep = queue.Queue()
    step_msg = None
    queHoles = queue.Queue()
    queCran = queue.Queue()
    depth = 0.0
    moveThreadX = None
    moveThreadY = None
    moveThreadZ = None
    # GLOBALS
    global BREGMAscrn, BREGMArobo, graphic_scaleXY, graphic_scaleZ
    global screw_holes_rat, screw_holes_scrn, craniotomies, craniotomies_rat, craniotomies_scrn,drill_vel
    global Robo000scrn, skull_Zsurface_scrn,hole_idx

    pygame.init()
    os.environ['SDL_VIDEO_WINDOW_POS'] = "%d,%d" % (10,30) # Place window on computer screen
    UMNlogo = pygame.image.load(r'.\RESOURCES\UMNlogo.PNG')
    pygame.display.set_icon(UMNlogo)
    TNElogo = pygame.image.load(r'.\RESOURCES\TNE logo.jpg')
    TNElogo = pygame.transform.scale(TNElogo, (70, 55))
    rats_skull_graphic = pygame.image.load(r'.\RESOURCES\skullgraphicReal.png')
    mice_skull_graphic = pygame.image.load(r'.\RESOURCES\mouse.png')
    pygame.display.set_caption('CraniobotGUI.py 2.3 by F. da Silva & Mark Schatza Nov. 20, 2019 (updated 4/6/2020)') # Enter your window caption here

    screw_holes_robo = []

    # At BREGMA IF NO ROBOT IS CONNECTED
    DRILLscrn = [213, 486, 549]  # BREGMA
    DRILLrat =  [0.0, 0.0, 0.0]  #
    DRILLrobo =  [10.0, 10.0, 15.0] # ROBOT 000

    ########################
    # Prepare Screen
    #########################
    myscreen = pygame.display.set_mode((700,900),pygame.RESIZABLE,32)
    myscreen.fill(pink1)

    #########################
    #  Create GUI elements
    #########################
    buttons, radio_buttons, boxes,LEDs,toggles,info_boxes,user_inputs,labels = CranioBot_GUI_ELEMENTS( myscreen )
    drill = MyDrill(myscreen,215, 473)
    drill_static_graphic = MyDrill(myscreen,510, 540)
    # WARNING MSGS
    warning = WARNINGBox( myscreen,500,150,160,160,"",'LEFT','WARNING!!',22)
    clk_time_start = time.perf_counter()
    START_TIME = time.perf_counter()
    cur_time =  time.perf_counter()

    START_EXPT = False
    PAUSE_STARTED = False

    # FLAGS
    SETTING_UP = True
    RUNNING = False

    SKULL_SURFACE = False
    BRAIN_SURFACE = False
    ###############################
    # Connect to motors
    Mx,My,Mz = connect_to_3motors()
    print("Mx,My,Mz: ",Mx,My,Mz)

    if Mx is not None and My is not None and Mz is not None:
            Mx.disable()
            My.disable()
            Mz.disable()
            time.sleep (0.1) #in sec
            Mx.enable()
            My.enable()
            Mz.enable()
            time.sleep (0.1) #in sec
            DRILLrobo = [Mx.position, My.position, Mz.position]
            ROBOT_CONNECTED = True ## REMOVE ME IF USING OTHER COMP
    #############################
    #  MAIN LOOP
    #############################

    while SETTING_UP:
        '''
         SYSTEM EVENTS
        '''

        for button in buttons: # Check for collision with EXISTING buttons
            if ( button.index == 2 or button.index == 12 or
                button.index == 22 or button.index == 3 or
                button.index == 13 or button.index == 23 ):
                if moveThreadY is not None and moveThreadY.is_alive():
                    button.face_color = lightgray
                else:
                    button.face_color = gray
            elif ( button.index == 4 or button.index == 14 or
                 button.index == 24 or button.index == 5 or
                 button.index == 15 or button.index == 25 ):
                 if moveThreadX is not None and moveThreadX.is_alive():
                     button.face_color = lightgray
                 else:
                     button.face_color = gray
            elif ( button.index == 6 or button.index == 16 or
                 button.index == 26 or button.index == 7 or
                 button.index == 17 or button.index == 27 ):
                     if moveThreadZ is not None and moveThreadZ.is_alive():
                         button.face_color = lightgray
                     else:
                         button.face_color = gray

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                exit_game()

            ################################
            ############################################################################
            #  HANDLE GUI EVENTS
            ###########################################################################################################

            # STOP ON SPACE PRESS
            elif (event.type == pygame.KEYDOWN):
                if event.key == pygame.K_SPACE:
                    for button in buttons:
                        if button.text == "STOP":
                            button.UP_DN = "DN"
                            if Mx is not None and My is not None and Mz is not None:
                                Mx.disable()
                                My.disable()
                                Mz.disable()
                                queHoles.put('STOP')
                                queCran.put('STOP')
                            button.text = 'GO'
                            button.face_color = (0,255,0)
                        elif button.text == 'GO':
                            button.UP_DN = "DN"
                            if Mx is not None and My is not None and Mz is not None:
                                Mx.enable()
                                My.enable()
                                Mz.enable()
                                queHoles.put('GO')
                                queCran.put('GO')
                            button.text = 'STOP'
                            button.face_color = (255,0,0)

                # ARROW KEYS TO GO UP/DN WHEN DRILLING HOLES
                elif event.key == pygame.K_DOWN:
                   if moveThreadZ is None or not moveThreadZ.is_alive():
                      if DRILLING_HOLE: # CHANGE ME for new coord system
                          moveThreadZ = threading.Thread(target =  move_step, args =(0.1,Mz))
                          moveThreadZ.start()

                elif event.key == pygame.K_UP:
                   if moveThreadZ is None or not moveThreadZ.is_alive():
                      if DRILLING_HOLE: # CHANGE ME for new coord system
                          moveThreadZ = threading.Thread(target =  move_step, args =(-0.1,Mz))
                          moveThreadZ.start()
            #-----------------------------------------
            # MOUSE MOVE
            elif (event.type == pygame.MOUSEMOTION):#
                cur_x,cur_y = pygame.mouse.get_pos()

                if MOVING_DRILL:

                    DRILLscrn[0] = cur_x #
                    DRILLscrn[1] = cur_y #
                    #print("DRILLscrn",DRILLscrn[0],DRILLscrn[1],"BREGMAscrn: ",BREGMAscrn[0],BREGMAscrn[1])
                    DRILLrobo[0] = -(DRILLscrn[0] - Robo000scrn[0])/graphic_scaleXY
                    DRILLrobo[1] = (DRILLscrn[1] - Robo000scrn[1])/graphic_scaleXY
                    DRILLrobo[2] = (DRILLscrn[2] - Robo000scrn[2])/graphic_scaleZ
                    if DRILLrobo[0] < 0.0 or DRILLrobo[0] > 25.0:
                        print("1.ROBO X OUT OF BOUNDS!!!!!")
                    if DRILLrobo[1] < 0.0 or DRILLrobo[1] > 25.0:
                        print("1.ROBO Y OUT OF BOUNDS!!!!!")
                    if DRILLrobo[2] < 0.0 or DRILLrobo[2] > 25.0:
                        print("1.ROBO Z OUT OF BOUNDS!!!!!")
                    DRILLrat[0] = DRILLrobo[0] - BREGMArobo[0]
                    DRILLrat[1] = DRILLrobo[1] - BREGMArobo[1]
                    DRILLrat[2] = DRILLrobo[2] - BREGMArobo[2]

            # MOUSE DOWN
            elif (event.type == pygame.MOUSEBUTTONDOWN ):#Mouse Clicked
                #print("MOUSE CLICKED")
                LEFT_MOUSE_DOWN = False
                RIGHT_MOUSE_DOWN = False

                cur_x,cur_y = pygame.mouse.get_pos()
                print("cur_x: ",cur_x,"cur_y: ",cur_y)
                if event.button == 1:
                    LEFT_MOUSE_DOWN = True
                elif event.button == 3:
                    RIGHT_MOUSE_DOWN = True

                elif event.button == 4:  #Wheel roll UP
                    MOUSE_WHEEL_SCROLL_UP = True

                elif event.button == 5: #Wheel roll Down
                    MOUSE_WHEEL_SCROLL_DN = True

                stepSize = 1 # Change with buttons
                # BUTTONS
                if LEFT_MOUSE_DOWN:
                    # RADIO BUTTONS
                    rad_button_idx = 0
                    rad_button_on = 0
                    for rad_button in radio_buttons: # Check for collision with EXISTING buttons
                            if button.rect.collidepoint(cur_x,cur_y):
                                        rad_button.ONOFF = "ON"
                                        if  rad_button.index == 0:
                                            MICE = False
                                            RATS = True
                                            radio_buttons[1].ONOFF = "OFF"
                                        if  rad_button.index == 1:
                                            MICE = True
                                            RATS = False
                                            radio_buttons[0].ONOFF = "OFF"

                    # BUTTONS
                    for button in buttons: # Check for collision with EXISTING buttons
                        if button.rect.collidepoint(cur_x,cur_y):
                               #########
                               if "Select Existing" in button.text:
                                    button.UP_DN = "DN"
                                    print("existing selected")
                                    cwd = os.getcwd()
                                    RODENT, screw_holes_scrn, screw_holes_rat, craniotomies_scrn, craniotomies_rat = CranioHelperFunctions.load_design(cwd,graphic_scaleXY ,BREGMAscrn)
                                    print("screw_holes_scrn: ",screw_holes_scrn)
                                    print("craniotomies_rat: ",craniotomies_rat)
                                    if "RAT" in RODENT:
                                        RATS = True
                                        MICE = False
                                    if "M" in RODENT:
                                        RATS = False
                                        MICE = True
                               elif "Create" in button.text:
                                    print("Create New Craniotomy")
                                    button.UP_DN = "DN"
                                    subprocess.call("CraniotomyDesigner1.3.py", shell=True)

                               # DRILL Y
                               elif button.index == 2 or button.index == 12 or button.index == 22: # y"^" Move drill to UP (+y)
                                   if button.index == 2: stepSize = 0.1
                                   elif button.index == 12: stepSize = 1
                                   elif button.index == 22: stepSize = 5

                                   if moveThreadY is None or not moveThreadY.is_alive():
                                      button.UP_DN = "DN"
                                      if DRILLrobo[1] - stepSize < 0.0:
                                            DRILLrobo[1] = 0.0
                                            print("1.ROBO Y OUT OF BOUNDS IF MOVED!!!!!")
                                      else: DRILLrobo[1] -= stepSize # mm
                                      if BREGMA_SELECTED: # CHANGE ME for new coord system
                                            DRILLscrn[1] = DRILLrobo[1]*graphic_scaleXY + Robo000scrn[1]# pixels
                                            DRILLrat[1] = DRILLrobo[1] - BREGMArobo[1]# mm

                                      moveThreadY = threading.Thread(target =  move_step, args =(-1.0 * stepSize,My))
                                      moveThreadY.start()

                               if button.index == 3 or button.index == 13 or button.index == 23: # y"v" Move drill to UP (-y)
                                   if button.index == 3: stepSize = 0.1
                                   elif button.index == 13: stepSize = 1
                                   elif button.index == 23: stepSize = 5

                                   if moveThreadY is None or not moveThreadY.is_alive():
                                      button.UP_DN = "DN"
                                      if DRILLrobo[1] + stepSize > 25.0:
                                            DRILLrobo[1] = 25.0
                                            print("2.ROBO Y OUT OF BOUNDS IF MOVED!!!!!")

                                      else: DRILLrobo[1] += stepSize # mm
                                      if BREGMA_SELECTED: # CHANGE ME for new coord system
                                            DRILLscrn[1] = DRILLrobo[1]*graphic_scaleXY + Robo000scrn[1]# pixels
                                            DRILLrat[1] = DRILLrobo[1] - BREGMArobo[1]# mm

                                      moveThreadY = threading.Thread(target =  move_step, args =(stepSize,My))
                                      moveThreadY.start()

                               #########
                               # DRILL X
                               elif button.index == 4 or button.index == 14 or button.index == 24: #x"<" Move drill to left (+x in robo, -x in scrn)
                                   if button.index == 4: stepSize = 0.1
                                   elif button.index == 14: stepSize = 1
                                   elif button.index == 24: stepSize = 5

                                   if moveThreadX is None or not moveThreadX.is_alive():
                                      button.UP_DN = "DN"
                                      if DRILLrobo[0] + stepSize > 25.0:
                                            DRILLrobo[0] = 25.0
                                            print("3.ROBO X OUT OF BOUNDS IF MOVED!!!!!")

                                      else: DRILLrobo[0] += stepSize # mm

                                      if BREGMA_SELECTED: # CHANGE ME for new coord system
                                            DRILLscrn[0] = -DRILLrobo[0]*graphic_scaleXY + Robo000scrn[0]# pixels
                                            DRILLrat[0] = DRILLrobo[0] - BREGMArobo[0]# mm

                                      moveThreadX = threading.Thread(target =  move_step, args =(stepSize,Mx))
                                      moveThreadX.start()

                               elif button.index == 5 or button.index == 15 or button.index == 25: #x">" Move drill to right (-x in robo)
                                   if button.index == 5: stepSize = 0.1
                                   elif button.index == 15: stepSize = 1
                                   elif button.index == 25: stepSize = 5

                                   if moveThreadX is None or not moveThreadX.is_alive():
                                      button.UP_DN = "DN"
                                      if DRILLrobo[0] - stepSize < 0.0:
                                            DRILLrobo[0] = 0.0
                                            print("4.ROBO X OUT OF BOUNDS IF MOVED!!!!!")
                                      else: DRILLrobo[0] -= stepSize # mm
                                      if BREGMA_SELECTED: # CHANGE ME for new coord system
                                            DRILLscrn[0] = -DRILLrobo[0]*graphic_scaleXY + Robo000scrn[0]# pixels
                                            DRILLrat[0] = DRILLrobo[0] - BREGMArobo[0]# mm

                                      moveThreadX = threading.Thread(target =  move_step, args =(-1.0 * stepSize,Mx))
                                      moveThreadX.start()

                               #########
                               # DRILL Z
                               elif button.index == 6 or button.index == 16 or button.index == 26: #z"^" Move drill to UP -z robo -y screen)
                                   if button.index == 6: stepSize = 0.1
                                   elif button.index == 16: stepSize = 1
                                   elif button.index == 26: stepSize = 5

                                   if moveThreadZ is None or not moveThreadZ.is_alive():
                                      button.UP_DN = "DN"
                                      if DRILLrobo[2] - stepSize < 0.0:
                                            DRILLrobo[2] = 0.0
                                            print("5.ROBO Z OUT OF BOUNDS IF MOVED!!!!!")
                                      else: DRILLrobo[2] -= stepSize # mm
                                      if BREGMA_SELECTED: # CHANGE ME for new coord system
                                            DRILLscrn[2] = DRILLrobo[2]*graphic_scaleZ + Robo000scrn[2]# pixels
                                            DRILLrat[2] = DRILLrobo[2] - BREGMArobo[2]# mm

                                      moveThreadZ = threading.Thread(target =  move_step, args =(-1.0 * stepSize,Mz))
                                      moveThreadZ.start()

                               elif button.index == 7 or button.index == 17 or button.index == 27: #z"v" Move drill to DN (+z robo, +y screen)
                                   if button.index == 7: stepSize = 0.1
                                   elif button.index == 17: stepSize = 1
                                   elif button.index == 27: stepSize = 5

                                   if moveThreadZ is None or not moveThreadZ.is_alive():
                                      button.UP_DN = "DN"

                                      if BREGMA_SELECTED: # CHANGE ME for new coord system
                                            if DRILLrobo[2] + stepSize > 25.0:
                                                DRILLrobo[2] = 25.0
                                                print("2.ROBO Z OUT OF BOUNDS IF MOVED!!!!!")
                                            else: DRILLrobo[2] += stepSize # mm
                                            # Only allow small steps if at brain surface
                                            if DRILLrobo[2] - BREGMArobo[2] >=1.2 and not (BRAIN_SURFACE == True and stepSize == 0.1):
                                                BRAIN_SURFACE = True
                                                DRILLrobo[2] -= stepSize
                                            else:
                                                BRAIN_SURFACE = False
                                                DRILLscrn[2] = DRILLrobo[2]*graphic_scaleZ + Robo000scrn[2]# pixels
                                                DRILLrat[2] = DRILLrobo[2] - BREGMArobo[2]# mm
                                                if DRILLrat[2] >=0:
                                                    SKULL_SURFACE = True
                                                else: SKULL_SURFACE = False
                                                moveThreadZ = threading.Thread(target =  move_step, args =(stepSize,Mz))
                                                moveThreadZ.start()

                                      else:
                                          moveThreadZ = threading.Thread(target =  move_step, args =(stepSize,Mz))
                                          moveThreadZ.start()
                                          if DRILLrobo[2] + stepSize > 25.0:
                                              DRILLrobo[2] = 25.0
                                              print("2.ROBO Z OUT OF BOUNDS IF MOVED!!!!!")
                                          else: DRILLrobo[2] += stepSize # mm

                               ##################
                               # HOME ROBOT HERE
                               elif "HOME ROBOT" in button.text :  # This is Robot 000 in screen coordinates
                                    button.UP_DN = "DN"

                                    for box in info_boxes: # Clear info boxes
                                        box.text = "0, 0, 0"

                                    if Mx is not None and My is not None and Mz is not None:
                                        ROBOT_CONNECTED = True
                                        home_all(Mx,My,Mz)
                                        x = Mx.position
                                        y = My.position
                                        z = Mz.position
                                        print("Cur Motor positions: ",x,y,z)
                                        print("ROBOT HOMED")
                                        DRILLrobo = [x, y, z]
                                    else:
                                        ROBOT_CONNECTED = False
                                        print("3.ROBOT NOT CONNECTED")

                                    if ROBOT_CONNECTED:
                                        DRILLrobo[0] = 0.0# IN ROBOT COORDINATES
                                        DRILLrobo[1] = 0.0# IN ROBOT COORDINATES
                                        DRILLrobo[2] = 0.0# IN ROBOT COORDINATES

                                        DRILLrat[0] = -10# IN RAT COORDINATES
                                        DRILLrat[1] = 10# IN RAT COORDINATES
                                        DRILLrat[2] = 15.0# IN RAT COORDINATES
                                    else:
                                        DRILLrobo[0] = 0.0#BREGMAscrn[0]+12.5*graphic_scaleXY# IN ROBOT COORDINATES
                                        DRILLrobo[1] = 0.0#BREGMAscrn[1]-12.5*graphic_scaleXY# IN ROBOT COORDINATES
                                        DRILLrobo[2] = 0.0#BREGMAscrn[2]-25.0*graphic_scaleZ # IN ROBOT COORDINATES

                                    #DRILLscrn = [215, 473, 549]
                                    DRILLscrn = [Robo000scrn[0],Robo000scrn[1],Robo000scrn[2]]

                                    ROBOT_HOMED = True

                               ####################################
                               # BACK TO BREGMA
                               elif "TO BREGMA" in button.text :  #
                                    if BREGMA_SELECTED:
                                        button.UP_DN = "DN"

                                        try:
                                            if MOVING is None or not MOVING.is_alive(): #is_alive is a Tread function
                                                print("MOVING TO BREGMA: ", BREGMArobo[0],BREGMArobo[1],BREGMArobo[2])
                                                MOVING = threading.Thread(target=move_to_XYZ, args = (BREGMArobo[0],
                                                                                BREGMArobo[1],
                                                                                BREGMArobo[2]-5,
                                                                                Mx,My,Mz,""))
                                                MOVING.start()
                                        except:
                                            print('cleanup')
                                            print('movement error')

                                        MOVING = None
                                    else: print("BREGMA NOT SELECTED YET")

                               ####################################
                               # SELECT BEST-FIT PLANE POINTS
                               elif "SELECT POINT" in button.text or "SELECT BREGMA" in button.text:
                                    button.UP_DN = "DN"
                                    #n = len(roboXs)+1

                                    ##########################
                                    # POINT 1: BREGMA
                                    n=1
                                    if n == 1: # Bregma in screen coord
                                        DRILLrat = [0.0, 0.0, 0.0] # in RAT COORDINATES

                                        # LOCATION OF ROBOT 000 on SCREEN
                                        if ROBOT_CONNECTED: #
                                            BREGMArobo[0] = Mx.position
                                            BREGMArobo[1] = My.position
                                            BREGMArobo[2] = Mz.position
                                            BREGMArobo = DRILLrobo.copy()
                                            Robo000scrn[0] = BREGMAscrn[0]+BREGMArobo[0]*graphic_scaleXY # ROBOT x00 on screen
                                            Robo000scrn[1] = BREGMAscrn[1]-BREGMArobo[1]*graphic_scaleXY # ROBOT 0y0 on screen
                                            Robo000scrn[2] = BREGMAscrn[2]-BREGMArobo[2]*graphic_scaleZ  # ROBOT 00z on screen

                                        else:
                                            print("ROBOT NOT CONNECTED")
                                            #BREGMArobo = [213,486,10]
                                            BREGMArobo[0] = -(BREGMAscrn[0] - Robo000scrn[0])/graphic_scaleXY
                                            BREGMArobo[1] = (BREGMAscrn[1] - Robo000scrn[1])/graphic_scaleXY
                                            BREGMArobo[2] = (BREGMAscrn[2] - Robo000scrn[2])/graphic_scaleZ
                                            if BREGMArobo[0] <= 0.0 or BREGMArobo[0] >= 25.0:
                                                print("5.ROBO X OUT OF BOUNDS!!!!!")
                                            if BREGMArobo[1] <= 0.0 or BREGMArobo[1] >= 25.0:
                                                print("5.ROBO Y OUT OF BOUNDS!!!!!")
                                            if BREGMArobo[2] <= 0.0 or BREGMArobo[2] >= 25.0:
                                                print("5.ROBO Z OUT OF BOUNDS!!!!!")

                                        #DRILLscrn = BREGMAscrn !!!!! Why doesn't this work????
                                        DRILLscrn = BREGMAscrn.copy()

                                        print("BREGMA: DRILLscrn",DRILLscrn[0],DRILLscrn[1],\
                                              "BREGMAscrn: ",BREGMAscrn[0],BREGMAscrn[1])
                                        info_boxes[0].text = str(round(BREGMArobo[0],2))+', ' + \
                                                             str(round(BREGMArobo[1],2)) +', '+ \
                                                             str(round(BREGMArobo[2],2))
                                        #info_boxes[5].text = str(BREGMArat[0]) +', '+ \
                                        info_boxes[1].text = str(BREGMArat[0]) +', '+ \
                                                             str(BREGMArat[1]) +', '+ \
                                                             str(BREGMArat[2])
                                        BREGMA_SELECTED = True
                                        #button.text = "SELECT POINT"
                                    # ROBOT COORDINATES
                                    if ROBOT_CONNECTED:
                                        DRILLrobo = [Mx.position, My.position, Mz.position]

                                    else: # For screen graphics when ROBOT is NOT connected
                                        DRILLrobo[0] = -(DRILLscrn[0] - Robo000scrn[0])/graphic_scaleXY #in mm (graphic_scaleXY in pix/mm)
                                        DRILLrobo[1] = (DRILLscrn[1] - Robo000scrn[1])/graphic_scaleXY
                                        DRILLrobo[2] = (DRILLscrn[2] - Robo000scrn[2])/graphic_scaleZ
                                        if DRILLrobo[0] <= 0.0 or DRILLrobo[0] >= 25.0:
                                            print("4.ROBO X OUT OF BOUNDS!!!!!")
                                        if DRILLrobo[1] <= 0.0 or DRILLrobo[1] >= 25.0:
                                            print("4.ROBO Y OUT OF BOUNDS!!!!!")
                                        if DRILLrobo[2] <= 0.0 or DRILLrobo[2] >= 25.0:
                                            print("4.ROBO Z OUT OF BOUNDS!!!!!")

                                    print("DrillRobo: ",DRILLrobo, "\nRobo000scrn: ",Robo000scrn)

                                    # RAT COORDINATES
                                    DRILLrat[0] = DRILLrobo[0] - BREGMArobo[0]
                                    DRILLrat[1] = DRILLrobo[1] - BREGMArobo[1]
                                    DRILLrat[2] = DRILLrobo[2] - BREGMArobo[2]

                                    # SCREEN COORDINATES
                                    DRILLscrn[0] = Robo000scrn[0] - DRILLrobo[0] * graphic_scaleXY
                                    #DRILLscrn[0] = (DRILLrobo[0] - BREGMArobo[0])*graphic_scaleXY + BREGMAscrn[0]  # IN SCREEN COORDINATES
                                    DRILLscrn[1] = (DRILLrobo[1] - BREGMArobo[1])*graphic_scaleXY + BREGMAscrn[1]
                                    DRILLscrn[2] = (DRILLrobo[2] - BREGMArobo[2])*graphic_scaleZ  + BREGMAscrn[2]

                                    # POINT LOCATIONS
                                    roboXs.append(DRILLrobo[0]) # For calculating best fit plane
                                    roboYs.append(DRILLrobo[1])
                                    roboZs.append(DRILLrobo[2])
                                    print("Point Xs: ",roboXs)
                                    print("Point Ys: ",roboYs)
                                    print("Point Zs: ",roboZs)

                                    # POINT 2
                                    if n == 2: # in screen coord

                                        info_boxes[1].text = str(round(DRILLrobo[0],2)) +', ' + \
                                                             str(round(DRILLrobo[1],2)) +', ' + \
                                                             str(round(DRILLrobo[2],2))

                                        info_boxes[6].text = str(round(DRILLrat[0],2))  + ', '+\
                                                             str(round(DRILLrat[1],2))  + ', '+\
                                                             str(round(DRILLrat[2],2))                                        # POINT 3
                                    # POINT 3
                                    if n == 3:
                                        info_boxes[2].text = str(round(DRILLrobo[0],2)) +', ' + \
                                                             str(round(DRILLrobo[1],2)) +', ' + \
                                                             str(round(DRILLrobo[2],2))
                                        info_boxes[7].text = str(round(DRILLrat[0],2))  + ', '+\
                                                             str(round(DRILLrat[1],2))  + ', '+\
                                                             str(round(DRILLrat[2],2))
                                    # POINT 4
                                    if n == 4:
                                        info_boxes[3].text = str(round(DRILLrobo[0],2)) +', ' + \
                                                             str(round(DRILLrobo[1],2)) +', ' + \
                                                             str(round(DRILLrobo[2],2))
                                        info_boxes[8].text = str(round(DRILLrat[0],2))  + ', '+\
                                                             str(round(DRILLrat[1],2))  + ', '+\
                                                             str(round(DRILLrat[2],2))
                                    # POINT 5
                                    if n == 5:
                                        info_boxes[4].text = str(round(DRILLrobo[0],2)) +', ' + \
                                                             str(round(DRILLrobo[1],2)) +', ' + \
                                                             str(round(DRILLrobo[2],2))
                                        info_boxes[9].text = str(round(DRILLrat[0],2))  + ', '+\
                                                             str(round(DRILLrat[1],2))  + ', '+\
                                                             str(round(DRILLrat[2],2))

                               ############################################
                               elif "CLR" in button.text:
                                    button.UP_DN = "DN"
                                    if len(roboXs) >0:
                                        roboXs.pop()
                                        roboYs.pop()
                                        roboZs.pop()
                                        n = len(roboXs)
                                        info_boxes[n].text = '(0, 0, 0)'
                                        info_boxes[n+5].text = '(0, 0, 0)'
##                               ##########################################
##
##                               elif button.index == 40: # Circular Craniotomy
##                                     CIRC_CRANIOTOMY = True
##
##                               elif button.index == 41: # Draw Craniotomy Path
##                                     CIRC_CRANIOTOMY = False
##                               ############################################
                               elif button.index == 9: # "Craniotomy":
                                    button.UP_DN = "DN"
                                    cranios_combined =[]
                                    c3d = []
                                    print("craniotomies: ",craniotomies_rat)
                                    print("BREGMArobo :",BREGMArobo)
                                    # Turn 2D craniotomies into 3D at BREGMA depth
                                    for c in craniotomies_rat:  # NOTE:  At this point, crniotomies are only list of (x,Y), no z
                                        for p in c:
                                            c3d.append((p[0]+BREGMArobo[0],p[1]+BREGMArobo[1],BREGMArobo[2]))

                                        c3d.append((p[0]+BREGMArobo[0],p[1]+BREGMArobo[1],BREGMArobo[2]-10)) # create new point 10 mm above last point
                                        print ("c3d ",c3d)

##                                    for c in c3d:
##                                        cranios_combined.append(c)

                                   # print("cranios_combined: ",c3d)

                                    # XXXX RUN CRANIOTOMIES THROUGH BEST FIT PLANE FIRST XXXXXX
                                    if ROBOT_CONNECTED:
                                        #for cranio in craniotomies:
                                        cranThread = threading.Thread(target = drill_carniotomy,
                                                                      args = (BREGMArobo,Mx,My,Mz, c3d, depth, queCran))
                                        cranThread.start()

                                    else:
                                        print('ROBOT IS NOT CONNECTED!')

##                                    if ROBOT_CONNECTED:
##                                        print("CRANIOTOMY STARTED....")
##                                        #for cranio in craniotomies:
##                                    if circle is not None: # This is the projected circle onto the best fit plabe
##
##
##                                        circle = best_fit_plot(roboXs,roboYs,roboZs,5.5,BREGMArobo,0,-1.7)
##                                        cranThread = threading.Thread(target = drill_circ_carniotomy,
##                                                                      args = (BREGMArobo,Mx,My,Mz,roboXs,roboYs,roboZs,
##                                                                              circle, depth, queCran ))
##                                        cranThread.start()
##
##                                    else:
##                                        print('ROBOT IS NOT CONNECTED!')

                                    # Increase Depth for next pass
                                    depth = depth+0.1

                                    #else:
                                    #    print('Define Circle! Hint: press best fit plane')
                               ############################################
                               elif button.text == "Next Hole":
                                    DRILLING_HOLE = True
                                    button.UP_DN = "DN"
                                    if BREGMA_SELECTED and ROBOT_CONNECTED: #convert to ROBO coordinates
                                        if hole_idx >=len(screw_holes_rat):
                                            DRILLING_HOLE = False
                                            print (hole_idx, " holes drilled")
                                            hole_idx = 0
                                        else:
##                                            holeX = round(screw_holes_rat[hole_idx][0] + BREGMArobo[0], 1)
##                                            holeY = round(screw_holes_rat[hole_idx][1] + BREGMArobo[1], 1)
                                            hole = (screw_holes_rat[hole_idx][0],screw_holes_rat[hole_idx][1])
##                                            move_to_hole = threading.Thread(target =  move_step, args =(stepSize,Mz))
##                                            move_to_hole.start()

                                            holesThread = threading.Thread(target = drill_holes, args = (hole,BREGMArobo,Mx,My,Mz, queHoles))
                                            holesThread.start()
                                            print("Drilling Hole ", hole_idx, " at: ",hole)
                                            hole_idx += 1
                                            queHoles.put('')

                                    else:
                                        print("Bregma MUST be selected first!!!")
                               ############################################
                               elif button.text == "PAUSE":
                                    button.UP_DN = "DN"
                                    queHoles.put('PAUSE')
                                    queCran.put('PAUSE')
                                    button.text = 'RESUME'
                               elif button.text == 'RESUME':
                                    button.UP_DN = "DN"
                                    queHoles.put('RESUME')
                                    queCran.put('RESUME')
                                    button.text = 'PAUSE'
                               ############################################
                               elif button.text == "STOP":
                                    button.UP_DN = "DN"
                                    if Mx is not None and My is not None and Mz is not None:
                                        Mx.disable()
                                        My.disable()
                                        Mz.disable()
                                        queHoles.put('STOP')
                                        queCran.put('STOP')
                                    button.text = 'GO'
                                    button.face_color = (0,255,0)

                               elif button.text == 'GO':
                                    button.UP_DN = "DN"
                                    if Mx is not None and My is not None and Mz is not None:
                                        Mx.enable()
                                        My.enable()
                                        Mz.enable()
                                        queHoles.put('GO')
                                        queCran.put('GO')
                                    button.text = 'STOP'
                                    button.face_color = (255,0,0)

                               #############################################
                               elif "BEST FIT" in button.text:
                                    button.UP_DN = "DN"
                                    for usr_input in user_inputs:
                                        if 'X offset' in usr_input.label:
                                            X_Offset = float(usr_input.text)
                                        if 'Y offset' in usr_input.label:
                                            Y_Offset = float(usr_input.text)

                                        if 'radius' in usr_input.label:
                                            Radius = float(usr_input.text)
                                    ##    # BEST FIT CIRCLE
                                    circle = best_fit_plot(roboXs,roboYs,roboZs,5.5,BREGMArobo,0,-1.7)
                                    ##    points = []

                               #############################################

                               elif "+" in button.text:
                                    button.UP_DN = "DN"
                                    queHoles.put('vel_up')
                                    queCran.put('vel_up')
                                    drill_vel += 0.25  # This is for diplay only!
                                                      # Not directly tied to drill vel in craniotomy
                                    if drill_vel >= 2.25: drill_vel = 2.25

                               elif "-" in button.text:
                                    button.UP_DN = "DN"
                                    queHoles.put('vel_dn')
                                    queCran.put('vel_dn')
                                    drill_vel -= 0.25
                                    if drill_vel <= 0.25: drill_vel = 0.25# This is for diplay only!
                                                      # Not directly tied to drill vel in craniotomy

                    #############################################

                                    # This is for diplay only!
                                                      # Not directly tied to drill vel in craniotomy
                    #############################################

                    if drill.rect.collidepoint(cur_x,cur_y):
                        MOVING_DRILL = True
                        drill.color = (255,255,255)
                        print("moving drill: ",MOVING_DRILL)

            elif (event.type == pygame.MOUSEBUTTONUP ):#Mouse UN-Clicked
                    for button in buttons: # Check for collision with EXISTING buttons
                        if button.rect.collidepoint(cur_x,cur_y):
                               button.UP_DN = "UP"
                    LEFT_MOUSE_DOWN = False
                    RIGHT_MOUSE_DOWN = False
                    MOVING_DRILL = False
                    drill.color = (255,0,0)
        ########################
        # REDRAW GUI ELEMENTS
        ########################
        myscreen.fill(pink4)
        myscreen.blit(TNElogo,(625,5))
        if RATS:  myscreen.blit(rats_skull_graphic,(94,350))
        if MICE:  myscreen.blit(mice_skull_graphic,(124,361))
        #########################
        #  Create GUI elements
        #########################
        for button in buttons:
            button.draw()

        for button in radio_buttons:
            button.draw()

        for LED in LEDs:
            LED.draw()

        for lbl in labels: # Last go on top
            lbl.draw()

        for info in info_boxes: # Last go on top
            info.draw()

        # USER INPUTS
        for user_input in user_inputs:
            if CIRC_CRANIOTOMY:
                 user_input.draw()
            else: # Draw Path
                if not "5.)"  in user_input.label and not "6.)" in user_input.label and  not "7.)" in user_input.label:
                    user_input.draw()

        # UPDATE DRILL INFO and GRAPHICS
        for user_input in user_inputs:
            if "Scrn X" in user_input.label: user_input.text = str(round(DRILLscrn[0],2))
            if "Scrn Y" in user_input.label: user_input.text = str(round(DRILLscrn[1],2))

            if "Drill X" in user_input.label: user_input.text = str(round(DRILLrobo[0],2))
            if "Drill Y" in user_input.label: user_input.text = str(round(DRILLrobo[1],2))
            if "Drill Z" in user_input.label: user_input.text = str(round(DRILLrobo[2],2))

            if "Rat X" in user_input.label: user_input.text = str(round(DRILLrat[0],2))
            if "Rat Y" in user_input.label: user_input.text = str(round(DRILLrat[1],2))
            if "Rat Z" in user_input.label: user_input.text = str(round(DRILLrat[2],2))
            if "Depth" in user_input.label: user_input.text = str(round(depth,2))
            if "Vel" in user_input.label: user_input.text = str(round(drill_vel,1))

            if "Hole #" in user_input.label: user_input.text = str(hole_idx)
            if " of " in user_input.label: user_input.text = str(len(screw_holes_rat))
            # NOTE:  Spaces are required otherwise this will also be put into info box labeled "offset.."

        # WARNING MSG
        if BRAIN_SURFACE:
            warning.text = "DANGER!!\nPOTENTIAL BRAIN\nCONTACT!!!\nClick again only if\nyou're 100% sure"
            warning.draw()
        elif SKULL_SURFACE:
            warning.text = "    MAKE SURE\n    DRILL IS ON!\n         and in \n          FWD \n       direction!"
            warning.draw()

        if DRILLrat[2] < 0:
            SKULL_SURFACE = False
            #elif DRILLrat[2] < -4:
            BRAIN_SURFACE = False

        # DRAW ROBOT 000 ON SCREEN IN GREEN
        pygame.draw.line(myscreen,(0,255,0),(Robo000scrn[0]-20, Robo000scrn[1]),(Robo000scrn[0]+20, Robo000scrn[1]) ,2)
        pygame.draw.line(myscreen,(0,255,0),(Robo000scrn[0], Robo000scrn[1]-20),(Robo000scrn[0], Robo000scrn[1]+20) ,2)
        pygame.draw.line(myscreen,(0,255,0),(630, Robo000scrn[2]),(690, Robo000scrn[2]) ,2)

        # DRAW STATIC DRILL Graphic
        drill_static_graphic.draw_drillXY(490, 560)
        if not BREGMA_SELECTED:
            drill.draw_drillZ(660, 505)

        # DRAW DRILL ACTUAL LOCATION
        else: # BREGMA_SELECTED:
            drill.draw_drillXY(int(DRILLscrn[0]),int(DRILLscrn[1]))
            drill.draw_drillZ(660, int(DRILLscrn[2]))

            if len(craniotomies_scrn)  > 0:
                 #print("len(craniotomies_scrn: ",len(craniotomies_scrn))
                 for c in craniotomies_scrn:
                         p = 0
                         for point in c:
                                 hcolor = (0,0,255)
                                 cranotomies_robo = (int(point[0]+ BREGMArobo[0]),
                                                 int(point[1] + BREGMArobo[1]))# in Robot coords

        ##                         craniotomies_scrn = ((int(point[0] - BREGMAscrn[0])/graphic_scaleXY),
        ##                                         (int(point[1]-BREGMAscrn[1])/graphic_scaleXY))

                                 pygame.draw.circle(myscreen, (0,0,255), point, 3,2) # BLUE
                                 if p > 0:
                                        # Draw line between points
                                        pygame.draw.line(myscreen,hcolor,(prev_point),(point),1)

                                 prev_point = point
                                 p+=1

                         # Draw line between ist and last point
                         pygame.draw.line(myscreen,hcolor,(c[0]),(c[len(c)-1]),1)

            if len(screw_holes_scrn) > 0:
                for hole in screw_holes_scrn:
                    hcolor = (150,255,255)
                    location_robo = (int(hole[0]+ BREGMArobo[0]),
                                     int(hole[1] + BREGMArobo[1]))# in Robot coords

                    location_scrn = ((int(hole[0] - BREGMAscrn[0])/graphic_scaleXY),
                                     (int(hole[1]-BREGMAscrn[1])/graphic_scaleXY))

                    #draw_screw_hole(myscreen,location_scrn, hcolor)
                    pygame.draw.circle(myscreen, (0,0,255), hole, 3,2) # BLUE
                    pygame.draw.line(myscreen,(0,0,255),(hole[0],hole[1] - 10),(hole[0], hole[1] + 10),1)
                    pygame.draw.line(myscreen,(0,0,255),(hole[0] -10 ,hole[1]),(hole[0]+10, hole[1]),1)

        pygame.display.flip()

#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

try:
    setupGUI()
finally:
    print('cleaned')
    core._cleanup()


"Calculates how a participant's estimated k influences learning speed."

import pickle
import os
import statsmodels.api as sm
from scipy.special import logit
import numpy as np
from bfit_samples import get_samples, get_subject_meank
from bdata import bdata, N, NTRIALS
from mpl_meanresp_curve import KMAXCURVE, KSPEED_CURVES_FN

def main():
    "Calculates the mean k and mean responses for each participant."
    with open(KSPEED_CURVES_FN, 'rb') as inpf:
        mpl_kcurves = [pickle.load(inpf) for k in range(KMAXCURVE + 1)]
    # Determine interval with maximum difference in mean response
    ini, end = None, None
    dif = 0
    for i in range(NTRIALS - 100):
        j = i + 100
        this_dif = 0
        for trial in range(i, j):
            this_dif += mpl_kcurves[0][trial] - mpl_kcurves[2][trial]
        this_dif /= 100
        if this_dif > dif:
            dif = this_dif
            ini, end = i, j
    print(dif, ini, end)
    if not os.path.exists('mean_k.pickle'):
        samples = get_samples()
        samples = samples.sample(10000)
        mean_k = [get_subject_meank(samples, i) for i in range(N)]
        with open('mean_k.pickle', 'wb') as outf:
            pickle.dump((ini, end), outf)
            pickle.dump(mean_k, outf)
    else:
        with open('mean_k.pickle', 'rb') as inpf:
            ini, end = pickle.load(inpf)
            mean_k = pickle.load(inpf)
    # ini, end = 200, 300
    mean_resp = [np.mean(y[ini:end]) for x, y in bdata]
    mod = sm.OLS(logit(mean_resp), [(1, k) for k in mean_k])
    res = mod.fit()
    print(res.summary())

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import models, migrations
import django.db.models.deletion
import smart_selects.db_fields

class Migration(migrations.Migration):

    dependencies = [
    ]

    operations = [
        migrations.CreateModel(
            name='BCategory',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('category', models.CharField(max_length=255, unique=True, null=True)),
                ('background_color', models.CharField(help_text=b'Ex: #FF0000', max_length=255, null=True, verbose_name=b'Background and Title Color')),
                ('category_icon', models.CharField(help_text=b'Ex: fa fa-car', max_length=255, unique=True, null=True)),
                ('created', models.DateTimeField(auto_now_add=True, null=True)),
                ('updated', models.DateTimeField(auto_now=True, null=True)),
            ],
            options={
                'verbose_name': 'Category',
                'verbose_name_plural': 'Categories',
            },
        ),
        migrations.CreateModel(
            name='BSubCategory',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('sub_category', models.CharField(max_length=255, null=True, verbose_name=b'Sub Category')),
                ('created', models.DateTimeField(auto_now_add=True, null=True)),
                ('updated', models.DateTimeField(auto_now=True, null=True)),
                ('category', models.ForeignKey(on_delete=django.db.models.deletion.SET_NULL, to='business_categories.BCategory', null=True)),
            ],
            options={
                'verbose_name': 'Sub Category',
                'verbose_name_plural': 'Sub Categories',
            },
        ),
        migrations.CreateModel(
            name='BSubCategoryType',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('subcat_type', models.CharField(max_length=255, null=True, verbose_name=b'Sub Category Type')),
                ('created', models.DateTimeField(auto_now_add=True, null=True)),
                ('updated', models.DateTimeField(auto_now=True, null=True)),
                ('category', models.ForeignKey(on_delete=django.db.models.deletion.SET_NULL, to='business_categories.BCategory', null=True)),
                ('sub_category', smart_selects.db_fields.ChainedForeignKey(chained_model_field=b'category', to='business_categories.BSubCategory', chained_field=b'category', null=True, auto_choose=True)),
            ],
            options={
                'verbose_name': 'Sub Category Type',
                'verbose_name_plural': 'Sub Category Types',
            },
        ),
        migrations.AlterUniqueTogether(
            name='bsubcategorytype',
            unique_together=set([('category', 'sub_category', 'subcat_type')]),
        ),
        migrations.AlterUniqueTogether(
            name='bsubcategory',
            unique_together=set([('category', 'sub_category')]),
        ),
    ]


# -*- coding: utf-8 -*-

from flask import Blueprint
from flask import redirect, request
from flask.helpers import url_for, flash
from flask_login import login_required
from models import Client, db
from flask.templating import render_template
from flask.views import MethodView
from messages import ClientMessages, CssMessages

bp = Blueprint('client', __name__)

class ClientAPIForm(MethodView):

    @login_required
    def post(self):

        client = Client()

        client.nome =  request.form.get('nome')
        client.razao_social =  request.form.get('razao_social')
        client.rua =  request.form.get('rua')
        client.cep =  request.form.get('cep')
        client.bairro =  request.form.get('bairro')
        client.cidade =  request.form.get('cidade')
        client.email = request.form.get('email')
        client.observacao =  request.form.get('observacao')

        client.save()

        flash(ClientMessages.CLIENT_ADD_SUCCESS, CssMessages.SUCCESS)
        return redirect(url_for('.ClientAPIForm'))

    @login_required
    def get(self):
        return render_template('client_add.html')

bp.add_url_rule('/client', view_func=ClientAPIForm.as_view('ClientAPIForm'))

class ClientAPI(MethodView):

    @login_required
    def get(self, id=None):

        client = Client.query.get(id)

        if client:
            return render_template('client_update.html', client=client)
        else:
            return render_template('page_404.html')

    @login_required
    def post(self, id=None):

        client = Client.query.get(id)

        client.nome =  request.form.get('nome')
        client.razao_social =  request.form.get('razao_social')
        client.rua =  request.form.get('rua')
        client.cep =  request.form.get('cep')
        client.bairro =  request.form.get('bairro')
        client.cidade =  request.form.get('cidade')
        client.email = request.form.get('email')
        client.observacao =  request.form.get('observacao')

        client.save()

        flash(ClientMessages.CLIENT_UPDATE_SUCCESS, 'success')

        return redirect(url_for('.ClientAPI', id=id))

bp.add_url_rule('/client/<id>', view_func=ClientAPI.as_view('ClientAPI'), methods=['GET', 'POST'])

class ClientAPIList(MethodView):

    @login_required
    def get(self, id=None):

        if id:
            clients = Client.query.filter_by(id=id)
        else:
            clients = Client.query.all()

        return render_template('client_list.html', clients=clients)

bp.add_url_rule('/clients', view_func=ClientAPIList.as_view('ClientAPIList'))

# @bp.before_request
# def restrict_bp_to_admins():
#     if not users.is_current_user_admin():
#         return redirect(users.create_login_url(request.url))


"""
Created on Sat Jun 13 15:31:32 2020

@author: david
"""

from pastml.acr import pastml_pipeline
from Bio import SeqIO
from ete3 import Tree
import pandas as pd
import numpy as np

def fasta2csv(fasta_file,csv_file):

    "Get sequences from fasta"
    seq_dic = {}
    for record in SeqIO.parse(fasta_file, "fasta"):
        seq_dic[record.id] = [i for i in record.seq] #str(record.seq)
    sites = len(seq_dic[next(iter(seq_dic))])
    traits = ["site" + str(i) for i in range(sites)]
    seqs_df = pd.DataFrame.from_dict(seq_dic, orient='index', columns=traits)
    seqs_df.to_csv(csv_file)

    return traits

def reconstruct(tree_file,align_file,csv_file="temp-align.csv"):

    if align_file.endswith('.fasta'):
        "Convert fasta to csv"
        traits = fasta2csv(align_file,csv_file) # convert fasta to csv
        data = csv_file # Path to the table containing tip/node annotations, in csv or tab format
    else:
        df = pd.read_csv(align_file, index_col=0)
        traits = [column for column in df.columns]
        data = align_file

    tree = tree_file # Path to the tree in newick format

    # Columns present in the annotation table,
    # for which we want to reconstruct ancestral states
    columns = traits #['Country']

    # Path to the output compressed map visualisation
    html_compressed = "tree-000_map.html"

    # (Optional) path to the output tree visualisation
    html = "tree-000_tree.html"

    pastml_pipeline(data=data, data_sep=',', columns=columns, tree=tree, verbose=True)

    #pastml_pipeline(data=data, data_sep=',', columns=columns, name_column=traits[0], tree=tree,
    #                html_compressed=html_compressed, html=html, verbose=True)

def label_internal_nodes(tree):

    internal_cntr = 0 # counter for internal nodes encountered
    for node in tree.traverse("preorder"):
        if node.is_root():
            node.name = 'root'
        else:
            if not node.is_leaf():
                node.name = 'n' + str(internal_cntr)
                internal_cntr += 1
    return tree

"""
    Reconstruct ancestral states using Sankoff's max parsimony algorithm (see Felsenstein p15-18)
    Parsimony scores are computed assuming all transitions have a cost of one.
    Tip/ancestral features are input and output as dictionaries
    Internal nodes are labeled as n<X> where X is an int determined by the position of the node in a pre-order traversal
"""
def reconstruct_MP(tree,feature_dic):

    state_set = set([feature_dic[node.name] for node in tree.traverse() if node.is_leaf()])
    states = len(state_set)

    "Add a state2int dict map so this works with non-integer data types"
    state2int = {state:index for index, state in enumerate(state_set)}
    int2state = {index:state for index, state in enumerate(state_set)}

    "Post-order traversal to compute costs in terms of required state transitions"
    for node in tree.traverse("postorder"):
        if node.is_leaf():
            costs = [np.inf]*states
            tip_state = state2int[feature_dic[node.name]] # feature_dic[node.name]
            costs[tip_state] = 0
            node.add_features(costs = costs)
        else:
            costs = [0]*states
            for i in range(states):
                child_costs = []
                for child in node.children:
                    temp_costs = [0]*states
                    for j in range(states):
                        temp_costs[j] = child.costs[j]
                        if i != j:
                            temp_costs[j] = temp_costs[j] + 1 # add cost for transitioning between i and j
                    child_costs.append(temp_costs)
                costs[i] = sum([min(c) for c in child_costs])
            node.add_features(costs = costs)

    "Pre-order traversal to select anc states based on least cost parsimony score"
    anc_dic = {}
    internal_cntr = 0 # counter for internal nodes encountered
    for node in tree.traverse("preorder"):
        costs = node.costs
        if node.is_root():
            root_state = costs.index(min(costs)) # or np.argmin(node.costs)
            node.add_features(state = root_state)
            anc_dic['root'] = root_state
        else:
            parent_state = node.up.state
            least_cost = min(costs)
            least_cost_state = costs.index(least_cost)
            if parent_state == least_cost_state:
                anc_state = parent_state
            else:
                parent_st_cost = costs[parent_state]
                if parent_st_cost < (least_cost+1):
                    anc_state = parent_state # if parent state costs less than transitioning to least cost state
                else:
                    anc_state = least_cost_state
            node.add_features(state = anc_state)
            if node.is_leaf():
                anc_dic[node.name] = anc_state
            else:
                name = 'n' + str(internal_cntr)
                anc_dic[name] = anc_state
                internal_cntr += 1

    "Covert from integers back to original states"
    for k,v in anc_dic.items():
        anc_dic[k] = int2state[v]

    "Should return tree and anc_dic"
    return tree, anc_dic

if __name__ == '__main__':

    import time

    "Old ML test data"
    #path = './test-sets/testTF_timeVaryingSiteEffects_june2020/'
    #tree_file = path + 'tree-000.tre'
    #fasta_file = path + 'tree-000.fasta'
    #csv_file = path + 'tree-000.csv'
    #reconstruct(tree_file,fasta_file,csv_file)

    "New covid test data for testing MP reconstructions"
    tree_file = 'covid_ancestral_D614G_testMP.tre'
    csv_file = 'covid_features_D614G_testMP.csv'

    tree = Tree(tree_file, format=1)
    tree = label_internal_nodes(tree)
    #for node in tree.traverse("preorder"): print(node.name)

    df = pd.read_csv(csv_file,index_col='node')
    feature_dic = {}
    for index, row in df.iterrows():
        feature_dic[index] = row['nsp12_P323L+S_D614G']
    tic = time.perf_counter()
    tree, anc_dic = reconstruct_MP(tree,feature_dic)
    toc = time.perf_counter()
    elapsed = toc - tic
    print(f"Elapsed time: {elapsed:0.4f} seconds")

    "Plot"
    import balticmod as bt
    import matplotlib as mpl
    from matplotlib import pyplot as plt
    import TreeUtils
    import seaborn as sns

    "Write tree with fit vals to multi-type Newick file"
    absolute_time = 2020.67
    mtt_file = 'covid_ancestral_D614G_testMP_mtt.tre'
    fig_file = 'covid_ancestral_D614G_testMP_MP.png'
    TreeUtils.write_MTT_newick(tree,mtt_file)

    sns.set(style="darkgrid")
    myTree=bt.loadNewick(mtt_file,absoluteTime=False)
    myTree.traverse_tree() ## required to set heights
    myTree.setAbsoluteTime(absolute_time) ## set absolute time of all branches by specifying date of most recent tip
    myTree.treeStats() ## report stats about tree

    cmap = mpl.cm.get_cmap('tab10', 10)

    fig,ax = plt.subplots(figsize=(20,20),facecolor='w')

    x_attr=lambda k: k.absoluteTime ## x coordinate of branches will be absoluteTime attribute
    c_func=lambda k: 'darkorange' if k.traits['type']=='1' else 'steelblue' ## colour of branches
    s_func=lambda k: 50-30*k.height/myTree.treeHeight ## size of tips
    z_func=lambda k: 100

    cu_func=lambda k: 'k' ## for plotting a black outline of tip circles
    su_func=lambda k: 2*(50-30*k.height/myTree.treeHeight) ## black outline in twice as big as tip circle
    zu_func=lambda k: 99
    myTree.plotTree(ax,x_attr=x_attr,colour_function=c_func) ## plot branches
    myTree.plotPoints(ax,x_attr=x_attr,size_function=s_func,colour_function=c_func,zorder_function=z_func) ## plot circles at tips
    myTree.plotPoints(ax,x_attr=x_attr,size_function=su_func,colour_function=cu_func,zorder_function=zu_func) ## plot circles under tips (to give an outline)

    "Add legend the hard way"
    import matplotlib.patches as mpatches
    blue_patch = mpatches.Patch(color='steelblue', label = 'Spike 614D')
    red_patch = mpatches.Patch(color='darkorange', label = 'Spike 614G')
    handles = [blue_patch,red_patch]

    "For Regions"
    ax.legend(handles=handles,prop={'size': 24}) #loc='upper left'

    "Add month labels as xticks"
    step_freq = 1/12
    xticks = np.arange(2020,absolute_time,step_freq)
    ax.set_xticks(xticks)
    labels = ['Jan','Feb','Mar','Apr','May','June','July','Aug','Sep']
    ax.set_xticklabels(labels, fontsize=24) #rotation='vertical'
    ax.set_xlabel('Time', fontsize=24)

    ax.set_ylim(-5,myTree.ySpan+5)
    plt.savefig(fig_file, dpi=300)


'''
Author: James Hibben
A simple python implementation of an RSA cryptographic algorithm, for
MSU CSC 590-002 in the Spring Semester of 2016.

This file was written in conjunction with Zach Gilmer for an assignment in
implementing both symmetric-key and public-key (asymmetric-key) algorithms.

This program is based on avalonalex's RSA.py project, found
at https://gist.github.com/avalonalex/2122098. Some parts are removed or edited
due to alternate library imports and alternate ways to deal with certain
functions. Additionally, some functions were restructured to run as inner
functions in cases where they were only called from one function.
'''

from fractions import gcd
from random import randint
from math import log
from copy import copy

def coPrime(a, b):
    """Determine if a and b are coprime; return false if gcd(a,b) is anything
       other than 1"""
    if gcd(a,b) != 1:
        return False
    return True

def modExp(a, d, n):
    '''returns (a ** d) mod n'''
    def int2baseTwo(x):
        '''
        Convert x from positive integer to base-2.
        Returns list of bits in reverse order.
        '''
        assert x >= 0
        bitInverse = []
        while x != 0:
            bitInverse.append(x & 1)
            x >>= 1
        return bitInverse

    assert d >= 0
    assert n >= 0
    base2D = int2baseTwo(d)
    base2DLength = len(base2D)
    modArray = []
    result = 1
    for i in range(1, base2DLength + 1):
        if i == 1:
            modArray.append(a % n)
        else:
            modArray.append((modArray[i - 2] ** 2) % n)
    for i in range(0, base2DLength):
        if base2D[i] == 1:
            result *= base2D[i] * modArray[i]
    return result % n

def millerRabinTest(n, k):
    '''
    Miller-Rabin pseudo-prime test.
    Test is a fast approximation for a prime with good accuracy.
    Returning True: most likely a prime
    Returning False: definitely a composite
    k -> used to determine accuracy
    '''
    # go ahead and define function(s) here; the actual Miller-Rabin test is
    #   after the functions
    def extractTwos(m):
        '''
        Strict: m is non-negative.
        Counts the number of 0-bits there are at the end of bin(m), which can
        represented as (2 ** n) - 1.
        Returns a tuple (s, d) such that m = (2 ** i) * d.
        Only called within millerRabinTest
        '''
        assert m >= 0
        i = 0
        while m & (2 ** i) == 0:
            i += 1
        return i, m >> i

    def isComposite(a):
        x = modExp(a, d, n)
        if x == 1 or x == (n - 1):
            return None
        else:
            for i in range(1, s):
                x = modExp(x, 2, n)
                if x == 1:
                    return False
                elif x == (n - 1):
                    return None
            return False

    # millerRabinTest body
    assert n >= 1  # ensure that n is larger than 1
    assert k > 0  # ensure k is positive

    # 2 is a prime; show this
    if n == 2:
        return True

    # all even integers other than 2 are composites; show this
    if n % 2 == 0:
        return False

    extracted = extractTwos(n - 1)
    s = extracted[0]
    d = extracted[1]
    assert 2 ** s * d == n - 1

    for i in range(0, k):
        a = randint(2, n - 2)
        if isComposite(a) == False:
            return False
    return True  # reasonably certain that n is prime

def newKey(a, b, k):
    '''
    Try to find two pseudo primes ~ between a, b to generate the keys.
    Raises a ValueError if it cannot find one.
    '''
    # define inner functions
    def findPrime(a, b, k):
        '''
        Returns a pseudo-prime roughly between a, b (may be greater than b).
        Limit the number of attempts to find a pseudo-prime to
        10*ln(x)+3 (most likely thousands of attempts, so it shouldn't fail).
        If it can't find a prime, throw a ValueError to halt the program.
        '''
        x = randint(a, b)
        for i in range(0, int(10 * log(x) + 3)):
            if millerRabinTest(x, k):
                return x
            else:
                x += 1
        raise ValueError

    def extendedGCD(a, b):
        if a == 0:
            return b, 0, 1
        else:
            g, y, x = extendedGCD(b % a, a)
            return g, x - (b // a) * y, y

    def modMultInverse(e, totient):
        '''Modular Multiplicative Inverse function'''
        if coPrime(e, totient):
            combined = extendedGCD(e, totient)
            return combined[1] % totient
        else:
            return 0

    # newKey function body
    try:
        p = findPrime(a, b, k)
        while True:
            q = findPrime(a, b, k)
            if q != p:
                break
    except:
        raise ValueError

    n = p * q
    m = (p - 1) * (q - 1)

    while True:
        e = randint(1, m)
        if coPrime(e, m):
            break

    d = modMultInverse(e, m)
    return (n, e, d)

def encrypt(message, modN, e, blockSize):
    '''
    given a string message, public keys (e, n), and blockSize,
    encrypt using RSA algorithms
    '''
    # inner functions
    def string2numList(strn):
        '''convert an ASCII string to a list of integers'''
        return [ ord(char) for char in strn ]

    def numList2blocks(lst, blockSize):
        '''
        Takes a list of integers (range: 0-127) and combines them to be
        blockSize using base 256. If the length is not blockSize, then
        random integers will be generated to fill it.
        '''
        returnList = []
        toProcess = copy(lst)
        if len(toProcess) % blockSize != 0:
            for i in range(0, blockSize - len(toProcess) % blockSize):
                toProcess.append(randint(32, 126))
        for i in range(0, len(toProcess), blockSize):
            block = 0
            for j in range(0, blockSize):
                block += toProcess[i + j] << (8 * (blockSize - j - 1))
            returnList.append(block)
        return returnList

    numList = string2numList(message)
    numBlocks = numList2blocks(numList, blockSize)
    return [ modExp(blocks, e, modN) for blocks in numBlocks ]

def decrypt(cipher, modN, d, blockSize):
    '''reverses encryption'''
    # inner functions
    def numList2string(lst):
        '''convert a list of integers to a string based on their ASCII values'''
        return ''.join(map(chr, lst))

    def blocks2numList(blocks, blockSize):
        '''inverses numList2blocks'''
        toProcess = copy(blocks)
        returnList = []
        for numBlock in toProcess:
            inner = []
            for i in range(0, blockSize):
                inner.append(numBlock % 256)
                numBlock >>= 8
            inner.reverse()
            returnList.extend(inner)
        return returnList

    numBlocks = [ modExp(blocks, d, modN) for blocks in cipher ]
    numList = blocks2numList(numBlocks, blockSize)
    return numList2string(numList)


#!/usr/bin/env python3
#
# For SVN->Git conversions w/ SVN externals. Assumes a linear history (because SVN).
#
# SVN does not store the external revision in parent repo revisions, it just gets the latest. So we assume the newest commit on the external at the time of the parent commit is the correct commit
#
# This ended up quite a mess. I'm not going to put any effort into improving this since it got the job done.
#

import argparse
import os
import shlex
from subprocess import CalledProcessError, check_call, check_output

from colorama import Fore

debug = True

def green(*args, **kwargs):
    print(Fore.GREEN, *args, Fore.RESET, **kwargs)

def _parse_args():
    parser = argparse.ArgumentParser(description='Link a git repository to a submodule, preserving history')
    parser.add_argument('--parent-branch', default=None, help='Parent branch to operate on (default: default branch after clone)')
    parser.add_argument('--submodule-name', default=None, help='Submodule name (default: last element of repo path (git default))')
    parser.add_argument('parent_repo', help='Parent repository working directory')
    parser.add_argument('submodule_repo', help='Submodule repository URI')
    return parser.parse_args()

def _git_run(cmd_str, return_output=False, env=None, **kwargs):
    check_fn = check_call
    if return_output:
        check_fn = check_output
    if debug:
        green(cmd_str.format(**kwargs))
    return check_fn(shlex.split(cmd_str.format(**kwargs)), env=env)

def _git_out(cmd_str, **kwargs):
    out = _git_run(cmd_str, return_output=True, **kwargs).strip()
    if debug:
        green('  = %s' % out)
    return out.decode('utf-8')

def inspect_repo(repo, ref=None):
    # set ref to a sha to include commits up to but not including ref
    rval = []
    for line in check_output(shlex.split("git -C {repo} log '--format=format:%H %at'".format(repo=repo))).splitlines():
        sha, stamp = line.decode('utf-8').strip().split()
        if sha == ref:
            break
        rval.append((sha, stamp))
    return list(reversed(rval))

def rev_parse(clone, abbrev_ref=False):
    return _git_out('git -C {clone} rev-parse {abbrev_ref}HEAD',
                    clone=clone,
                    abbrev_ref='--abbrev-ref ' if abbrev_ref else '')

def clone(repo, clone, branch=None):
    if not os.path.exists(clone):
        green('Cloning {repo} to {clone}'.format(repo=repo, clone=clone))
        _git_run('git clone {branch}{repo} {clone}',
                 branch='-b %s ' % branch if branch else '',
                 repo=repo,
                 clone=clone)
    elif branch:
        checkout(clone, branch)
    return rev_parse(clone=clone, abbrev_ref=True)

def checkout(clone, ref):
    _git_run('git -C {clone} checkout {ref}',
             clone=clone,
             ref=ref)

def find_submodule_commit(parent_commit, submodule_clone, submodule_ref=None):
    last_submodule_commit = ('none', -1)
    for submodule_commit in inspect_repo(submodule_clone, ref=submodule_ref):
        if submodule_commit[1] > parent_commit[1]:
            break
        last_submodule_commit = submodule_commit
    return last_submodule_commit

def find_commit_pair(parent_clone, submodule_clone, parent_ref=None, submodule_ref=None):
    # locate the oldest commit on the parent that is not older than the oldest commit on the submodule.
    oldest_submodule_commit = ('none', -1)
    for parent_commit in inspect_repo(parent_clone, ref=parent_ref):
        submodule_commit = find_submodule_commit(parent_commit, submodule_clone, submodule_ref=submodule_ref)
        if submodule_commit[1] != -1:
            return (parent_commit, submodule_commit)
    else:
        return (None, None)

def submodule_add(parent_clone, submodule_clone, submodule_repo, submodule_name):
    parent_commit, first_submodule_commit = find_commit_pair(parent_clone, submodule_clone)
    green('Adding initial submodule {name} at commit {submodule_commit} on parent commit {parent_commit}'
        .format(
            name=submodule_name,
            submodule_commit=first_submodule_commit[0],
            parent_commit=parent_commit[0]))
    # remove in case switching from a branch that has it
    #_git_run('rm -rf {clone}/{name}', clone=parent_clone, name=submodule_name)
    checkout(parent_clone, parent_commit[0])
    # --force is needed when adding it on a second branch
    _git_run('git -C {clone} submodule add --force {repo} {name}'
        .format(
            clone=parent_clone,
            repo=submodule_repo,
            name=submodule_name))
    return first_submodule_commit[0]

def submodule_update(parent_clone, submodule_clone, submodule_repo, parent_branch, submodule_name, parent_commit):
    if parent_commit:
        parent_ref = parent_commit[0]
        checkout(parent_clone, parent_ref)
        submodule_ref = find_submodule_commit(parent_commit, submodule_clone)[0]
    else:
        submodule_ref = submodule_add(parent_clone, submodule_clone, submodule_repo, submodule_name)
        parent_ref = _git_out('git -C {clone} rev-parse HEAD', clone=parent_clone)
    green('Checking out submodule {name} at commit {submodule_commit} and amending to parent commit {parent_commit}'
        .format(
            name=submodule_name,
            submodule_commit=submodule_ref,
            parent_commit=parent_ref if parent_ref else 'initial'))
    checkout('%s/%s' % (parent_clone, submodule_name), submodule_ref)
    _git_run('git -C {clone} add -- {name}',
             clone=parent_clone,
             name=submodule_name)
    _git_run('git -C {clone} commit --no-edit --amend --allow-empty-message', clone=parent_clone)
    new_ref = _git_out('git -C {clone} rev-parse HEAD', clone=parent_clone)
    new_ref_time = _git_out("git -C {clone} log '--format=format:%at' -1 {ref}", clone=parent_clone, ref=new_ref)
    checkout(parent_clone, parent_branch)
    _git_run('git -C {clone} rebase {new_commit}', clone=parent_clone, new_commit=new_ref)
    return ((new_ref, new_ref_time), submodule_ref)

if __name__ == '__main__':
    args = _parse_args()
    parent_clone = os.path.basename(args.parent_repo.rstrip('/'))
    submodule_clone = os.path.basename(args.submodule_repo.rstrip('/'))
    parent_branch = clone(args.parent_repo, parent_clone, args.parent_branch)
    clone(args.submodule_repo, submodule_clone)
    submodule_name = args.submodule_name
    if not submodule_name:
        submodule_name = os.path.basename(args.submodule_repo)
    new_commit = True
    while new_commit:
        new_commit, submodule_ref = submodule_update(parent_clone, submodule_clone, args.submodule_repo, parent_branch, submodule_name, new_commit if new_commit is not True else None)
        new_commit = find_commit_pair(parent_clone, submodule_clone, parent_ref=new_commit[0], submodule_ref=submodule_ref)[0]
    initial = _git_out('git -C {clone} rev-list --max-parents=0 HEAD', clone=parent_clone)
    _git_run("""git -C {clone} filter-branch -f --commit-filter 'export GIT_COMMITTER_NAME="$GIT_AUTHOR_NAME"; export GIT_COMMITTER_EMAIL="$GIT_AUTHOR_EMAIL"; export GIT_COMMITTER_DATE="$GIT_AUTHOR_DATE"; git commit-tree "$@"' -- {initial}..HEAD""", clone=parent_clone, initial=initial)


import os
import time

from urllib.parse import urlparse

from django.core.management.base import BaseCommand, CommandError

import requests

from images.models import Image, ImageSourceResult

def get_any_value(d: dict, *keys):
    """
    Returns the first valid key's value.
    """

    for key in keys:
        if key in d:
            return d.get(key)
    return None

class Command(BaseCommand):
    def handle(self, *args, **options):
        """
        Handle image source fetching.
        """

        for image in Image.objects.filter(
            source_queries=None, source_url=None
        ).exclude(file=None).exclude(file='')[:100]:
            # Handle the 100 daily images.

            file_url = image.file.url

            with requests.get(file_url) as f:
                r = requests.post(
                    "https://saucenao.com/search.php",
                    params={
                        "db": 999,
                        "output_type": 2,
                        "testmode": 1,
                        "numres": 16,
                        "api_key": os.getenv("BACKEND_SAUCENAO_TOKEN"),
                    },
                    files=dict(file=f.content),
                )

                try:
                    results = r.json()["results"]
                except:
                    self.stderr.write(self.style.ERROR(f"ERROR - Key error in response - {r.status_code}"))
                    continue
                possible_sources = []

                for result in results:
                    if float(result["header"]["similarity"]) >= 85:
                        possible_sources.append(
                            {
                                "title": get_any_value(results, "title", "index_name"),
                                "similarity": float(result["header"]["similarity"]),
                                "ext_urls": result["data"].get("ext_urls"),
                                "source": result["data"].get("source"),
                                "artist": {
                                    "name": get_any_value(
                                        result["data"],
                                        "member_name",
                                        "creator",
                                        "author_name",
                                        "artist",
                                    ),
                                    "url": get_any_value(result["data"], "author_url"),
                                },
                            }
                        )

                if len(possible_sources) == 0:
                    ImageSourceResult.objects.create(
                        image=image, status=0, similarity=0
                    )
                    self.stderr.write(
                        f"QUERIED: ERROR - Image ID: {image.id} - Status code: {r.status_code}"
                    )
                    self.stdout.write("Sleeping 10s...\n")
                    time.sleep(10)
                    continue

                possible_sources = sorted(
                    possible_sources, key=lambda d: d["similarity"], reverse=True
                )

                source = possible_sources[0]

                ImageSourceResult.objects.create(
                    image=image,
                    status=r.status_code,
                    title=source["title"],
                    similarity=source["similarity"],
                    ext_urls=source["ext_urls"],
                    source=source["source"],
                    artist_name=source["artist"]["name"],
                    artist_url=source["artist"]["url"],
                )

                image.source_url = (
                    source["ext_urls"][0]
                    if source["ext_urls"] is not None and len(source["ext_urls"])
                    else None
                )

                source_hostname = urlparse(image.source_url).netloc

                source_names = {
                    "www.pixiv.net": "Pixiv",
                    "danbooru.donmai.us": "Danbooru",
                    "gelbooru.com": "Gelbooru",
                    "deviantart.com": "DevianArt",
                    "www.mangaupdates.com": "Manga Updates",
                    "mangadex.org": "Mangadex",
                    "yande.re": "yande.re",
                }

                image.source_name = (
                    source_names[source_hostname]
                    if source_hostname in source_names
                    else None
                )
                image.save()

                self.stdout.write(
                    self.style.SUCCESS(
                        f"QUERIED - Image ID: {image.id} - Results: {len(possible_sources)} - Status code: {r.status_code}"
                    )
                )
                self.stdout.write("Sleeping 10s...\n")
                time.sleep(10)