import argparse 
from pythonosc import dispatcher 
from pythonosc import osc_server 

import threading 
import time
import math
import numpy as np
import curses

import rtde_receive
import rtde_control


from helpfunctions import *
from ui import UI
from oscsend import OSCSend





rtde_r=None
rtde_c=None
server=None

#Looptimings
frequency_loop=500
frequency_rtde=None 


#Robot Control Parameters
# Parameters
velocity = 0.1
acceleration = 0.1
lookahead_time = 0.1
gain = 300
joint_q_aim=[0,0,0,0,0,0] #joint rotations aiming for
joint_q = [0,0,0,0,0,0] #joint rotation to be send
velocitylimit=[1,1,1,3,3,3] #joint velocity limits


#osc receive statistics
oscreceiveintervals = np.zeros(100)
oscreceiveintervals_pos=0
last_oscreceive=0

def handle_received_osc_joint(address, *args): 
    oscreceiveintervals[oscreceiveintervals_pos]=time.time()-last_oscreceive
    last_oscreceive=time.time()
    oscreceiveintervals_pos+=1
    oscreceiveintervals_pos%=np.size(oscreceiveintervals)
    
    print(f"Received OSC message on {address}") 
    for iarg, arg in enumerate(args): 
        #print(f"Data: {arg}") 
        joint_q_aim[iarg]=float(arg)

 
def setupOSC(ip, port): 
    # Set up the dispatcher for the server 
    disp = dispatcher.Dispatcher() 
    disp.map("/joint", handle_received_osc_joint)  # You can change '/filter' to any address pattern you expect 
 
    # Set up OSC server 
    server = osc_server.ThreadingOSCUDPServer((ip, port), disp) 
    print(f"Serving on {server.server_address}") 
    server_thread = threading.Thread(target=server.serve_forever) 
    server_thread.start() 
 
    return server,server_thread



def setupRTDE(robotip,frequency):
    dt = 1 / frequency
    rtde_r = rtde_receive.RTDEReceiveInterface(robotip, frequency)

    rtde_c = rtde_control.RTDEControlInterface(robotip)

    return rtde_r
 
 
joint_min=[toRad(-360),toRad(-360),toRad(-360),toRad(-360),toRad(-360),toRad(0)]
joint_max=[toRad(360),toRad(360),toRad(360),toRad(360),toRad(360),toRad(360)]


ui = UI(joint_min,joint_max)
oscsend = OSCSend()



def main(stdscr):

    ui.begin(stdscr)

    
    loop_dt = 1 / frequency_loop
    
    
    
    
    loopcounter = 0
    loopdurations = np.zeros(100)
    loopdurations_pos=0
    last_receivertde=0
    last_sendrtde=0
    

    q=[0,0,0,0,0,0]

    
    try:
        while ui.getKey() != ord('q'):
            looptime = time.time()
            
            start = time.time()
            

            #Get robot joint angled
            if looptime-last_receivertde>=1.0/frequency_receivertde:
                last_receivertde=looptime

                if rtde_r is not None:
                    q=rtde_r.getActualQ()
                else:
                    _offsets=[12.2,5.3,0,6,12.4,612.91]
                    _factors=[1.2,1.14,1.214,0.9921,1,1.092]
                    
                    for _i,_offset in enumerate(_offsets):
                        q[_i]=(math.sin(time.time()*_factors[_i]+_offset)/2.0+0.5)*(joint_max[_i]-joint_min[_i])+joint_min[_i]


            #send data to all osc receivers
            oscsend.send(looptime,q)

            if looptime-last_sendrtde>=1.0/frequency_sendrtde:
                last_sendrtde=looptime

                for i in range(6): #update joint_q with max speeds
                    joint_q[i] += constrain(joint_q_aim[i]-joint_q[i], -velocitylimit[i]/frequency_sendrtde,velocitylimit[i]/frequency_sendrtde) #constrain joint speeds

                if rtde_c is not None:
                    rtde_c.servoJ(joint_q, velocity, acceleration, 1.0/frequency_sendrtde, lookahead_time, gain)


            #Display
            ui.update(q,looptime)
                
            #Looptiming
            end = time.time()
            loopduration = end - start
            loopdurations[loopdurations_pos]=loopduration
            loopdurations_pos+=1
            loopdurations_pos%=np.size(loopdurations)

            ui.setDebugtext(0,"loopcounter="+str(loopcounter))
            ui.setDebugtext(1,"loop usage="+str(round(np.mean(loopdurations)/loop_dt,3)))
            ui.setDebugtext(2,"osc receive frequency="+str(round(1.0/np.mean(oscreceiveintervals),0))+" \t pos="+str(oscreceiveintervals_pos))

            
            

            if loopduration < loop_dt:
                time.sleep(loop_dt - loopduration)
            loopcounter += 1

    except KeyboardInterrupt:
        
        print("\nKeyboardInterrupt")
    curses.endwin()
        

if __name__ == "__main__": 
    parser = argparse.ArgumentParser() 
    parser.add_argument("--listenip",
        default="127.0.0.1", help="The ip to listen on")
    parser.add_argument("--listenport",
        type=int, default=5005, help="The port to listen on")

    parser.add_argument("--clientip",
        default="127.0.0.1", help="The ip to connect to")
    parser.add_argument("--port",
        type=int, default=5005, help="The port to connect to")
    

    parser.add_argument("--robotip",
        default="127.0.0.1", help="The ip to connect to")
    
    parser.add_argument("--test", action='store_true', default=False, help="Disable network and robot communication")
    
    
    parser.add_argument(
        "-fs",
        "--frequencysend",
        dest="frequencysend",
        help="the frequency at which the robot control data is send",
        type=float,
        default=500.0,
        metavar="<frequency>")

    parser.add_argument(
        "-fr",
        "--frequencyreceive",
        dest="frequencyreceive",
        help="the frequency at which the robot actualq data is received",
        type=float,
        default=500.0,
        metavar="<frequency>")
    
    args = parser.parse_args()

    frequency_receivertde=args.frequencyreceive
    frequency_sendrtde=args.frequencysend

 
    
    if args.test is not True:
        server,server_thread=setupOSC(args.listenip, args.listenport) 

        oscsend.setupReceivers([args.clientip],args.port)

        rtde_r=setupRTDE(args.robotip,frequency_rtde)

    
    curses.wrapper(main)

    if server is not None:
        input("Ended. Press Enter to shutdown")
        server.shutdown() 
        server_thread.join() 
    if rtde_c is not None:
        rtde_c.servoStop()
        rtde_c.stopScript()