# -*- 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 mpmath import *
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(1293, 851)
        MainWindow.setDocumentMode(False)
        self.centralwidget = QtWidgets.QWidget(MainWindow)
        self.centralwidget.setObjectName("centralwidget")
        self.image_box = QtWidgets.QLabel(self.centralwidget)
        self.image_box.setGeometry(QtCore.QRect(10, 10, 1281, 461))
        self.image_box.setText("")
        current_image = "tennis_court3.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(220, 580, 101, 16))
        self.label.setObjectName("label")
        self.label_2 = QtWidgets.QLabel(self.centralwidget)
        self.label_2.setGeometry(QtCore.QRect(230, 520, 81, 16))
        self.label_2.setObjectName("label_2")
        self.label_3 = QtWidgets.QLabel(self.centralwidget)
        self.label_3.setGeometry(QtCore.QRect(200, 550, 101, 16))
        self.label_3.setObjectName("label_3")
        self.label_4 = QtWidgets.QLabel(self.centralwidget)
        self.label_4.setGeometry(QtCore.QRect(190, 610, 121, 16))
        self.label_4.setObjectName("label_4")
        self.height_box = QtWidgets.QLineEdit(self.centralwidget)
        self.height_box.setGeometry(QtCore.QRect(330, 520, 113, 22))
        self.height_box.setObjectName("height_box")
        self.theta_box = QtWidgets.QLineEdit(self.centralwidget)
        self.theta_box.setGeometry(QtCore.QRect(330, 550, 113, 22))
        self.theta_box.setObjectName("theta_box")
        self.speed_box = QtWidgets.QLineEdit(self.centralwidget)
        self.speed_box.setGeometry(QtCore.QRect(330, 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(330, 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(920, 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.groupBox = QtWidgets.QGroupBox(self.centralwidget)
        self.groupBox.setGeometry(QtCore.QRect(680, 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.setChecked(True)
        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(1060, 560, 93, 28))
        self.clear_all.setObjectName("clear_all")
        self.groupBox_2 = QtWidgets.QGroupBox(self.centralwidget)
        self.groupBox_2.setGeometry(QtCore.QRect(500, 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.setChecked(True)
        self.disable_drag_button.setObjectName("disable_drag_button")
        self.label_5 = QtWidgets.QLabel(self.centralwidget)
        self.label_5.setGeometry(QtCore.QRect(190, 490, 111, 16))
        self.label_5.setObjectName("label_5")
        self.shot_position = QtWidgets.QLineEdit(self.centralwidget)
        self.shot_position.setGeometry(QtCore.QRect(330, 490, 113, 22))
        self.shot_position.setObjectName("shot_position")
        self.groupBox_3 = QtWidgets.QGroupBox(self.centralwidget)
        self.groupBox_3.setGeometry(QtCore.QRect(330, 650, 131, 121))
        self.groupBox_3.setObjectName("groupBox_3")
        self.grass_court = QtWidgets.QRadioButton(self.groupBox_3)
        self.grass_court.setGeometry(QtCore.QRect(10, 20, 95, 21))
        self.grass_court.setChecked(True)
        self.grass_court.setObjectName("grass_court")
        self.hard_court = QtWidgets.QRadioButton(self.groupBox_3)
        self.hard_court.setGeometry(QtCore.QRect(10, 50, 95, 21))
        self.hard_court.setObjectName("hard_court")
        self.clay_court = QtWidgets.QRadioButton(self.groupBox_3)
        self.clay_court.setGeometry(QtCore.QRect(10, 80, 95, 21))
        self.clay_court.setObjectName("clay_court")
        self.show_animation_button = QtWidgets.QPushButton(self.centralwidget)
        self.show_animation_button.setGeometry(QtCore.QRect(930, 610, 111, 41))
        self.show_animation_button.setObjectName("show_animation_button")
        self.label_6 = QtWidgets.QLabel(self.centralwidget)
        self.label_6.setGeometry(QtCore.QRect(930, 670, 55, 16))
        self.label_6.setObjectName("label_6")
        self.animation_speed = QtWidgets.QSlider(self.centralwidget)
        self.animation_speed.setGeometry(QtCore.QRect(1050, 670, 121, 20))
        self.animation_speed.setMinimum(1)
        self.animation_speed.setMaximum(5)
        self.animation_speed.setSingleStep(1)
        self.animation_speed.setProperty("value", 1)
        self.animation_speed.setSliderPosition(1)
        self.animation_speed.setOrientation(QtCore.Qt.Horizontal)
        self.animation_speed.setInvertedAppearance(False)
        self.animation_speed.setInvertedControls(False)
        self.animation_speed.setObjectName("animation_speed")
        self.label_7 = QtWidgets.QLabel(self.centralwidget)
        self.label_7.setGeometry(QtCore.QRect(10, 490, 71, 16))
        self.label_7.setObjectName("label_7")
        self.sim_time = QtWidgets.QLineEdit(self.centralwidget)
        self.sim_time.setGeometry(QtCore.QRect(80, 490, 61, 22))
        self.sim_time.setObjectName("sim_time")
        self.label_8 = QtWidgets.QLabel(self.centralwidget)
        self.label_8.setGeometry(QtCore.QRect(10, 530, 55, 16))
        self.label_8.setObjectName("label_8")
        self.data_points = QtWidgets.QLineEdit(self.centralwidget)
        self.data_points.setGeometry(QtCore.QRect(80, 530, 61, 22))
        self.data_points.setObjectName("data_points")
        MainWindow.setCentralWidget(self.centralwidget)
        self.menubar = QtWidgets.QMenuBar(MainWindow)
        self.menubar.setGeometry(QtCore.QRect(0, 0, 1293, 25))
        self.menubar.setObjectName("menubar")
        MainWindow.setMenuBar(self.menubar)
        self.statusbar = QtWidgets.QStatusBar(MainWindow)
        self.statusbar.setObjectName("statusbar")
        MainWindow.setStatusBar(self.statusbar)



        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

        self.hit_button.clicked.connect(self.show_trajectory)
        self.clear_all.clicked.connect(self.clear_canvas)
    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.groupBox_3.setTitle(_translate("MainWindow", "Court type:"))
        self.grass_court.setText(_translate("MainWindow", "Grass"))
        self.hard_court.setText(_translate("MainWindow", "Hard-court"))
        self.clay_court.setText(_translate("MainWindow", "Clay-court"))
        self.show_animation_button.setText(_translate("MainWindow", "Show animation"))
        self.show_animation_button.setShortcut(_translate("MainWindow", "Shift+Return"))
        self.label_6.setText(_translate("MainWindow", "Speed:"))
        self.label_7.setText(_translate("MainWindow", "Sim. Time:"))
        self.sim_time.setText(_translate("MainWindow", "2"))
        self.label_8.setText(_translate("MainWindow", "Data pts:"))
        self.data_points.setText(_translate("MainWindow", "200"))

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


        speed = float(self.speed_box.text())
        rpm = float(self.rev_per_sec_box.text())
        angle = float(self.theta_box.text())
        height = float(self.height_box.text())
        T = 3
        points = 150
        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.clay_court.isChecked()):
            court = "clay"
        if (self.hard_court.isChecked()):
            court = "hard"
        if (self.grass_court.isChecked()):
            court = "grass"


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

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

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


            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)
        np.set_printoptions(suppress=True)

        #print(Sx)
        #print(Sy)


        #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_court3.png"
        self.image_box.setPixmap(QtGui.QPixmap(current_image))

def get_kl(omega,v,V_net,radius,density_air,vx,vy):
    S = abs(omega*radius)/V_net
    kl = (1/2)*S*(math.pi*pow(radius,2))*density_air*(pow(V_net,2))/abs(omega*V_net)

    if ( (v>0 and omega>0) or (v<0 and omega<0)):
        return kl
    else:
        return -kl

def get_sign(val):
    if (val>0):
        return 1
    elif (val<0):
        return  -1
    else:
        return 0
def calculate_trajectory(height,x_offset,speed,rps,angle_degrees,T,points,drag,court):
    points = points

    V = speed
    rps = rps*2*3.14/60
    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
    #making time array...
    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.2
    m = 0.059
    g = 9.8

    radius = 0.0342
    g = 9.8

    V_spin = radius * rps

    k = 0.0744
    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])


        #print("---------------------")
        V_net = pow(pow(Vx[i-1],2) + pow(Vy[i-1],2),0.5)
        #print("Velocity: ",V_net,Vx[i-1],Vy[i-1],"Theta: ",theta)

        #This was working good enough...
        if(rps != 0):
        
            klx = get_kl(rps,Vy[i-1],V_net,radius,density_air,Vx[i-1],Vy[i-1])
            kly = get_kl(rps,Vx[i-1],V_net,radius,density_air,Vx[i-1],Vy[i-1])
            #print("KLX,KLY: ",klx,kly)
        else:
            klx = 0
            kly = 0



        fdx = -get_sign(Vx[i - 1]) * (0.5) * density_air * (3.14 * (radius ** 2)) * 0.6 * (Vx[i - 1] ** 2)
        fdy = -get_sign(Vy[i - 1]) * (0.5) * density_air * (3.14 * (radius ** 2)) * 0.6 * (Vy[i - 1] ** 2)

        fx = fdx - klx*(abs(rps*Vy[i-1]))
        fy = fdy + kly*(abs(rps*Vx[i-1]))-m*g


        #print("Force-x: ", fx, "  Force-Y: ", fy)
        ax = fx/m
        ay = fy/m
        #print("aX,aY: ", ax, ay)


        Vx[i] = Vx[i-1] +   ax*delta_t
        Vy[i] = Vy[i-1] +   ay*delta_t



        Sx[i] = Sx[i - 1] + Vx[i - 1] * delta_t + 0.5 * ax * delta_t ** 2
        Sy[i] = Sy[i - 1] + Vy[i - 1] * delta_t + 0.5 * ay * delta_t ** 2

        if (Sy[i] <= 0.00 and Vy[i] < 0):
            # hitting ground....
            #print("Bounce at: T= ", t[i], " secs")
            #print("Distance: ", Sx[i])
            #print("Collision at: ",Sx[i])
            rest_coffs_x = {"clay": 0.7, "grass": 0.9, "hard": 0.78}
            rest_coffs_y = { "clay" : 0.85 ,"grass":0.6 , "hard" : 0.83}

            #print("V before:",Vx[i],Vy[i])


            Vy[i] = -Vy[i]*rest_coffs_y[court]
            Vx[i] = Vx[i]*rest_coffs_x[court]
            #print("V AFTER:", Vx[i], Vy[i])

        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]
    print(Sx)
    print(Sy)
    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())
