diff --git a/frc2024.png b/frc2024.png new file mode 100644 index 0000000..4a384db Binary files /dev/null and b/frc2024.png differ diff --git a/main.py b/main.py new file mode 100644 index 0000000..1c22339 --- /dev/null +++ b/main.py @@ -0,0 +1,141 @@ +import math +import pygame as pg +from pygame.locals import * +from sys import exit +import numpy as np + +doubleClickDuration = 200 + +nodeColor = (255, 255, 255) +nodeRadius = 15 + +lineApproximationLineColor = (127, 127, 127, 0.5) +lineApproximationLineWidth = 3 + +curvePointCount = 300 +curvePointColor = (255, 255, 0) +curvePointRadius = 2 + +curveEditPointColor = (0, 255, 255) +curveEditPointRadius = 5 + +pg.init() + +# fieldImg = Background("frc2024.png", (0, 0)) + +fieldImg = pg.image.load("frc2024.png") + +screen_width = fieldImg.get_width() +screen_height = fieldImg.get_height() + +screen = pg.display.set_mode((screen_width, screen_height)) +fieldImg = pg.image.load("frc2024.png").convert_alpha() +pg.display.set_caption("Auto Planner") + + +curveEditPoints = [] +nodes = [] + +def addNode(pos): + nodes.append(pos) + if len(nodes) > 1: + index = len(nodes)-1 + # Middle point between current point and previous point + editPos = (nodes[index-1][0]+pos[0])/2,(nodes[index-1][1]+pos[1])/2 + curveEditPoints.append(editPos) + +def bezier(p0, p1, p2): + #for p in [p0, p1, p2]: + # pg.draw.circle(screen, (255, 255, 255), p, 5) + for t in np.arange(0, 1, 1/curvePointCount): + px = p0[0]*(1-t)**2 + 2*(1-t)*t*p1[0] + p2[0]*t**2 + py = p0[1]*(1-t)**2 + 2*(1-t)*t*p1[1] + p2[1]*t**2 + pg.draw.circle(screen, curvePointColor, (px, py), curvePointRadius) + +def refresh(): + pg.draw.rect(screen, (0, 0, 0), (0, 0, screen_width, screen_height)) + screen.blit(fieldImg, screen.get_rect()) + for i in range(0,len(curveEditPoints)): + pg.draw.line(screen, lineApproximationLineColor, nodes[i], curveEditPoints[i], lineApproximationLineWidth) + pg.draw.line(screen, lineApproximationLineColor, curveEditPoints[i], nodes[i+1], lineApproximationLineWidth) + bezier(nodes[i], curveEditPoints[i], nodes[i+1]) + pg.draw.circle(screen, curveEditPointColor, curveEditPoints[i], curveEditPointRadius) + for pos in nodes: + pg.draw.circle(screen, nodeColor, pos, nodeRadius) + pg.display.update() + +def getElemAt(pos): + for i in range(0,len(curveEditPoints)): + if getDist(pos, curveEditPoints[i], curveEditPointRadius): + return 1, i + for i in range(0,len(nodes)): + if getDist(pos, nodes[i], nodeRadius): + return 0, i + return -1, -1 + +def getDist(pos1, pos2, dist): + return math.sqrt(math.pow(pos1[0]-pos2[0], 2) + math.pow(pos1[1]-pos2[1], 2)) <= dist + +def singleClick(): + pos = pg.mouse.get_pos() + addNode(pos) + refresh() + +def doubleClick(): + curveEditPoints = [] + nodes = [] + refresh() + +refresh() + +running = True +last_click = -1 + +clickType = -1 +clickIndex = -1 + +while running: + for event in pg.event.get(): + + if event.type == pg.MOUSEBUTTONDOWN: + clickType, clickIndex = getElemAt(pg.mouse.get_pos()) + + now = pg.time.get_ticks() + if now - last_click <= doubleClickDuration: + if clickType == -1: + pass + if clickType == 0: + print(nodes) + if clickIndex > 0: + if clickIndex < len(nodes)-1: + newPos = (nodes[clickIndex-1][0]+nodes[clickIndex][0])/2,(nodes[clickIndex-1][1]+nodes[clickIndex][1])/2 + curveEditPoints[clickIndex] = newPos + curveEditPoints.pop(clickIndex-1) + elif clickIndex == 0 and len(nodes) > 1: + curveEditPoints.pop(clickIndex) + # if len(nodes) != 1: + nodes.pop(clickIndex) + # else: nodes = [] + refresh() + + + else: + if clickType == -1: + singleClick() + last_click = pg.time.get_ticks() + + if event.type == pg.MOUSEMOTION and clickType != -1: + if clickType == 0: + nodes[clickIndex] = pg.mouse.get_pos() + if clickType == 1: + curveEditPoints[clickIndex] = pg.mouse.get_pos() + refresh() + + if event.type == pg.MOUSEBUTTONUP: + if clickType != -1: + clickType = -1 + clickIndex = -1 + + if event.type == pg.QUIT: + running = False +pg.quit()