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

# Form implementation generated from reading ui file 'tennis_GUI.ui'
#
# Created by: PyQt5 UI code generator 5.15.0
#
# WARNING: Any manual changes made to this file will be lost when pyuic5 is
# run again.  Do not edit this file unless you know what you are doing.


from PyQt5 import QtCore, QtGui, QtWidgets
import sys
import math
import numpy as np
current_image = None
import cv2

red = (0,0,255)
blue = (255,0,0)
green = (0,255,0)

class Ui_MainWindow(object):
    def setupUi(self, MainWindow):
        global current_image
        MainWindow.setObjectName("MainWindow")
        MainWindow.resize(1122, 851)
        self.centralwidget = QtWidgets.QWidget(MainWindow)
        self.centralwidget.setObjectName("centralwidget")
        self.image_box = QtWidgets.QLabel(self.centralwidget)
        self.image_box.setGeometry(QtCore.QRect(10, 0, 1101, 461))
        self.image_box.setText("")
        current_image = "tennis_court2.png"
        self.image_box.setPixmap(QtGui.QPixmap(current_image))
        self.image_box.setObjectName("image_box")
        self.label = QtWidgets.QLabel(self.centralwidget)
        self.label.setGeometry(QtCore.QRect(50, 580, 101, 16))
        self.label.setObjectName("label")
        self.label_2 = QtWidgets.QLabel(self.centralwidget)
        self.label_2.setGeometry(QtCore.QRect(60, 520, 81, 16))
        self.label_2.setObjectName("label_2")
        self.label_3 = QtWidgets.QLabel(self.centralwidget)
        self.label_3.setGeometry(QtCore.QRect(30, 550, 101, 16))
        self.label_3.setObjectName("label_3")
        self.label_4 = QtWidgets.QLabel(self.centralwidget)
        self.label_4.setGeometry(QtCore.QRect(20, 610, 121, 16))
        self.label_4.setObjectName("label_4")
        self.height_box = QtWidgets.QLineEdit(self.centralwidget)
        self.height_box.setGeometry(QtCore.QRect(160, 520, 113, 22))
        self.height_box.setObjectName("height_box")
        self.theta_box = QtWidgets.QLineEdit(self.centralwidget)
        self.theta_box.setGeometry(QtCore.QRect(160, 550, 113, 22))
        self.theta_box.setObjectName("theta_box")
        self.speed_box = QtWidgets.QLineEdit(self.centralwidget)
        self.speed_box.setGeometry(QtCore.QRect(160, 580, 113, 22))
        self.speed_box.setObjectName("speed_box")
        self.rev_per_sec_box = QtWidgets.QLineEdit(self.centralwidget)
        self.rev_per_sec_box.setGeometry(QtCore.QRect(160, 610, 113, 22))
        self.rev_per_sec_box.setObjectName("rev_per_sec_box")
        self.hit_button = QtWidgets.QPushButton(self.centralwidget)
        self.hit_button.setGeometry(QtCore.QRect(750, 550, 121, 51))
        font = QtGui.QFont()
        font.setFamily("Lucida Bright")
        font.setPointSize(18)
        font.setBold(True)
        font.setWeight(75)
        self.hit_button.setFont(font)
        self.hit_button.setObjectName("hit_button")
        self.hit_button.clicked.connect(self.show_trajectory)
        self.groupBox = QtWidgets.QGroupBox(self.centralwidget)
        self.groupBox.setGeometry(QtCore.QRect(510, 520, 181, 111))
        self.groupBox.setObjectName("groupBox")
        self.red_clr = QtWidgets.QRadioButton(self.groupBox)
        self.red_clr.setGeometry(QtCore.QRect(20, 20, 95, 21))
        self.red_clr.setObjectName("red_clr")
        self.greeen_clr = QtWidgets.QRadioButton(self.groupBox)
        self.greeen_clr.setGeometry(QtCore.QRect(20, 50, 95, 21))
        self.greeen_clr.setObjectName("greeen_clr")
        self.blue_clr = QtWidgets.QRadioButton(self.groupBox)
        self.blue_clr.setGeometry(QtCore.QRect(20, 80, 95, 21))
        self.blue_clr.setObjectName("blue_clr")
        self.clear_all = QtWidgets.QPushButton(self.centralwidget)
        self.clear_all.setGeometry(QtCore.QRect(890, 560, 93, 28))
        self.clear_all.setObjectName("clear_all")
        self.groupBox_2 = QtWidgets.QGroupBox(self.centralwidget)
        self.groupBox_2.setGeometry(QtCore.QRect(330, 520, 131, 111))
        self.groupBox_2.setObjectName("groupBox_2")
        self.enable_drag_button = QtWidgets.QRadioButton(self.groupBox_2)
        self.enable_drag_button.setGeometry(QtCore.QRect(20, 30, 95, 21))
        self.enable_drag_button.setObjectName("enable_drag_button")
        self.disable_drag_button = QtWidgets.QRadioButton(self.groupBox_2)
        self.disable_drag_button.setGeometry(QtCore.QRect(20, 60, 95, 21))
        self.disable_drag_button.setObjectName("disable_drag_button")
        self.label_5 = QtWidgets.QLabel(self.centralwidget)
        self.label_5.setGeometry(QtCore.QRect(20, 490, 111, 16))
        self.label_5.setObjectName("label_5")
        self.shot_position = QtWidgets.QLineEdit(self.centralwidget)
        self.shot_position.setGeometry(QtCore.QRect(160, 490, 113, 22))
        self.shot_position.setObjectName("shot_position")
        MainWindow.setCentralWidget(self.centralwidget)
        self.menubar = QtWidgets.QMenuBar(MainWindow)
        self.menubar.setGeometry(QtCore.QRect(0, 0, 1122, 25))
        self.menubar.setObjectName("menubar")
        MainWindow.setMenuBar(self.menubar)
        self.statusbar = QtWidgets.QStatusBar(MainWindow)
        self.statusbar.setObjectName("statusbar")
        MainWindow.setStatusBar(self.statusbar)
        self.clear_all.clicked.connect(self.clear_canvas)
        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

    def retranslateUi(self, MainWindow):
        _translate = QtCore.QCoreApplication.translate
        MainWindow.setWindowTitle(_translate("MainWindow", "Tennis-Simulator"))
        self.label.setText(_translate("MainWindow", "Speed (m/s):"))
        self.label_2.setText(_translate("MainWindow", "Height (m):"))
        self.label_3.setText(_translate("MainWindow", "Angle (degrees):"))
        self.label_4.setText(_translate("MainWindow", "Spin speed (rev/s):"))
        self.hit_button.setText(_translate("MainWindow", "Hit"))
        self.hit_button.setShortcut(_translate("MainWindow", "Return"))
        self.groupBox.setTitle(_translate("MainWindow", "Color:"))
        self.red_clr.setText(_translate("MainWindow", "Red"))
        self.greeen_clr.setText(_translate("MainWindow", "Green"))
        self.blue_clr.setText(_translate("MainWindow", "Blue"))
        self.clear_all.setText(_translate("MainWindow", "Clear all"))
        self.clear_all.setShortcut(_translate("MainWindow", "Esc"))
        self.groupBox_2.setTitle(_translate("MainWindow", "Air drag:"))
        self.enable_drag_button.setText(_translate("MainWindow", "Enabled"))
        self.disable_drag_button.setText(_translate("MainWindow", "Disabled"))
        self.label_5.setText(_translate("MainWindow", "Shot position (m):"))
        self.shot_position.setText(_translate("MainWindow", "0"))
        self.height_box.setText(_translate("MainWindow", "0"))
        self.rev_per_sec_box.setText(_translate("MainWindow", "0"))
        self.theta_box.setText(_translate("MainWindow", "0"))
        self.speed_box.setText(_translate("MainWindow", "0"))

    def show_trajectory(self):
        global current_image
        px_per_m = 46.31
        gnd_y = 396


        speed = float(self.speed_box.text())
        rps = float(self.rev_per_sec_box.text())
        angle = float(self.theta_box.text())
        height = float(self.height_box.text())
        T = 3
        points = 200
        radius = 0.0342
        x_offset = float(self.shot_position.text())

        #x_offset_px = px_per_m*x_offset

        color = green

        if(self.red_clr.isChecked()):
            color = red
        if(self.greeen_clr.isChecked()):
            color = green
        if(self.blue_clr.isChecked()):
            color = blue



        if (self.enable_drag_button.isChecked()):
            drag = True
        else:
            drag = False

        image = cv2.imread(current_image)
        (Sx,Sy) = calculate_trajectory(height,x_offset,speed,rps,angle,T,points,drag)

        for i in range(Sx.shape[0]):
            X_m = Sx[i]
            X_px = int(X_m * px_per_m)


            Y_m = Sy[i]
            Y_px = int(Y_m * px_per_m)
            Y_px = gnd_y - Y_px

            cv2.circle(image, (X_px, Y_px), 3, color, -1)

        #cv2.imshow("Trajectory:", image)
        cv2.imwrite('temp_traj.png', image)
        current_image = 'temp_traj.png'

        self.image_box.setPixmap(QtGui.QPixmap(current_image))
        #cv2.waitKey(0)


    def clear_canvas(self):
        global current_image
        current_image = "tennis_court2.png"
        self.image_box.setPixmap(QtGui.QPixmap(current_image))







def calculate_trajectory(height,x_offset,speed,rps,angle_degrees,T,points,drag):
    points = points

    V = speed
    rps = rps
    angle_degrees = angle_degrees  # degrees

    theta = (math.pi / 180) * angle_degrees
    Vx_in = V * math.cos(theta)
    Vy_in = V * math.sin(theta)

    T = T
    t = np.linspace(0, T, points)

    Vx = np.zeros(t.shape[0])
    Vy = np.zeros(t.shape[0])
    Sx = np.zeros(t.shape[0])
    Sy = np.zeros(t.shape[0])


    Vx[0] = Vx_in
    Vy[0] = Vy_in

    Sx[0] = x_offset
    Sy[0] = height



    density_air = 1.21
    m = 0.059
    g = 9.8

    radius = 0.0342
    g = 9.8

    V_spin = radius * rps

    for i in range(1, t.shape[0]):

        cross_area = math.pi * (radius ** 2)
        delta_t = t[i] - t[i - 1]

        theta = math.atan(Vy[i - 1] / Vx[i - 1])

        V_net = (Vx[i - 1] ** 2 + Vy[i - 1] ** 2) ** 0.5

        '''Cd = 0.55 + 1/ (22.5 + 4.2 * (V_net / V_spin)**2.5)**0.4
        Fd = Cd *cross_area * density_air * (V_net**2) /2'''
        if(drag):
            Fd = 0.0014 * V_net ** 2
        else:
            Fd = 0

        Cl = 1 / (2 + (V_net / V_spin))
        Fl = Cl * cross_area * density_air * (V_net ** 2) / 2

        Fx = -Fd * math.cos(theta) - Fl * math.sin(theta)
        Fy = Fl * math.cos(theta) - Fd * math.sin(theta) - m * g

        Vx[i] = Vx[i - 1] + (1 / m) * (Fx) * delta_t
        Vy[i] = Vy[i - 1] + (1 / m) * (Fy) * delta_t

        Sx[i] = Sx[i - 1] + Vx[i - 1] * delta_t + 0.5 * (Fx / m) * delta_t ** 2
        Sy[i] = Sy[i - 1] + Vy[i - 1] * delta_t + 0.5 * (Fy / m) * delta_t ** 2

        if (Sy[i] <= 0):
            # hitting ground....
            #print("Bounce at: T= ", t[i], " secs")
            #print("Distance: ", Sx[i])
            Vy[i] = -Vy[i-1]*0.95
            Vx[i] = Vx[i-1]*0.9

        if (abs(Sx[i] - 11.885) < 0.3) and (Sy[i] <= 0.914):


            Vx[i] = -Vx[i]*0.1
            Vy[i] = Vy[i]*0.7

    ##Sx = Sx[0:clash]
    # Sy = Sy[0:clash]

    return (Sx,Sy)


def my_excepthook(type, value, tback):
    # log the exception here

    # then call the default handler
    sys.__excepthook__(type, value, tback)


if __name__ == "__main__":


    sys.excepthook = my_excepthook
    app = QtWidgets.QApplication(sys.argv)
    MainWindow = QtWidgets.QMainWindow()
    ui = Ui_MainWindow()
    ui.setupUi(MainWindow)
    MainWindow.show()
    sys.exit(app.exec())
