# Do the necessary imports
import argparse
import shutil
import base64
from datetime import datetime
import os
import cv2
import numpy as np
import socketio
import eventlet
import eventlet.wsgi
from PIL import Image
from flask import Flask
from io import BytesIO, StringIO
import json
import pickle
import matplotlib.image as mpimg
import time

# Import functions for perception and decision making
from perception import perception_step
from manuever import rock_manuever
# from manuever import rover_stuck
from decision import decision_step
from perception import rover_stuck
# from perception import is_rover_stuck
from supporting_functions import update_rover, create_output_images
# Initialize socketio server and Flask application 
# (learn more at: https://python-socketio.readthedocs.io/en/latest/)
sio = socketio.Server()
app = Flask(__name__)

# Read in ground truth map and create 3-channel green version for overplotting
# NOTE: images are read in by default with the origin (0, 0) in the upper left
# and y-axis increasing downward.
ground_truth = mpimg.imread('map_bw.png')
# This next line creates arrays of zeros in the red and blue channels
# and puts the map into the green channel.  This is why the underlying 
# map output looks green in the display image
ground_truth_3d = np.dstack((ground_truth*0, ground_truth*255, ground_truth*0)).astype(np.float)

# Define RoverState() class to retain rover state parameters
class RoverState():
    def __init__(self):
        self.start_time = None # To record the start time of navigation
        self.total_time = None # To record total duration of naviagation
        self.start_xpos = 0 # record the x coordinates of the starting postion
        self.start_ypos = 0 # record the y coordinates of the starting postion
        self.start_initpos_count = 0 # variable used to only record the initial positions of the rover
        self.img = None # Current camera image
        self.pos = None # Current position (x, y)
        self.yaw = None # Current yaw angle
        self.pitch = None # Current pitch angle
        self.roll = None # Current roll angle
        self.vel = None # Current velocity
        self.steer = 0 # Current steering angle
        self.throttle = 0 # Current throttle value
        self.brake = 0 # Current brake value
        self.nav_angles = None # Angles of navigable terrain pixels
        self.nav_dists = None # Distances of navigable terrain pixels
        self.ground_truth = ground_truth_3d # Ground truth worldmap
        self.mode = 'forward' # Current mode (can be forward or stop)
        self.throttle_set = 0.5 # Throttle setting when accelerating
        self.brake_set = 5 # Brake setting when braking
        # The stop_forward and go_forward fields below represent total count
        # of navigable terrain pixels.  This is a very crude form of knowing
        # when you can keep going and when you should stop.  Feel free to
        # get creative in adding new fields or modifying these!
        self.stop_forward = 500 # Threshold to initiate stopping
        self.go_forward = 500 # Threshold to go forward again
        self.max_vel = 1# Maximum velocity (meters/second)
        # Image output from perception step
        # Update this image to display your intermediate analysis steps
        # on screen in autonomous mode
        self.vision_image = np.zeros((160, 320, 3), dtype=np.float) 
        # Worldmap
        # Update this image with the positions of navigable terrain
        # obstacles and rock samples
        self.worldmap = np.zeros((200, 200, 3), dtype=np.float) 
        self.samples_pos = None # To store the actual sample positions
        self.samples_to_find = 0 # To store the initial count of samples
        self.samples_located = 0 # To store number of samples located on map
        self.samples_collected = 0 # To count the number of samples collected
        self.near_sample = 0 # Will be set to telemetry value data["near_sample"]
        self.picking_up = 0 # Will be set to telemetry value data["picking_up"]
        self.send_pickup = False # Set to True to trigger rock pickup
        self.rock_dist = 0 #will be set to distance from rock in perception.py
        self.rock_ang = 0 #will be set to angle from rock in perception.py
        self.rock_stop_forward = 70 #stopping distance to stop in front of rock for pickup
        self.a = 0 #variable to test 
        self.c = 0 #variable to limit looping in worldmap update
        self.rock_x_world = 0 #variable to save the xpos rock coordinates
        self.rock_y_world = 0 #variable to save the ypos rock coordinates
        self.rover_x_world = 0 #variable to save the xpos of the rover
        self.rover_y_world = 0 #variable to save the ypos of the rover
        self.rover_x_world_rev = 0 #variable to save the xpos of the rover
        self.rover_y_world_rev = 0 #variable to save the ypos of the rover
        # self.rover_x_pos = np.ones(5, dtype=np.float) #variable to save an array of x coordinates and then perform std deviation to determine if rover is stuck
        # self.rover_y_pos = np.ones(5, dtype=np.float) #variable to save an array of x coordinates and then perform std deviation to determine if rover is stuck
        self.stuck_time = 0 #variable to gather data samples to determine if rover is stuck
        #self.straight_reverse = 'False' #flat that enables the rover to back out straight
        self.reverse = 'False' #flag that enables the rover to turn around
        self.manuever_flag = 'False'    #flag that is used to prevent exessive switch of Rover drive parameters when it is near the rock
        self.turn180 = 'False'
        self.steering_bias = 5
# Initialize our rover 
Rover = RoverState()
# Variables to track frames per second (FPS)
# Intitialize frame counter
frame_counter = 0
# Initalize second counter
second_counter = time.time()
fps = None


# Define telemetry function for what to do with incoming data
@sio.on('telemetry')
def telemetry(sid, data):

    global frame_counter, second_counter, fps
    frame_counter+=1
    # Do a rough calculation of frames per second (FPS)
    if (time.time() - second_counter) > 1:
        fps = frame_counter
        frame_counter = 0
        second_counter = time.time()
    print("Current FPS: {}".format(fps))

    if data:
        global Rover
        # Initialize / update Rover with current telemetry
        Rover, image = update_rover(Rover, data)

        if np.isfinite(Rover.vel):

            # Execute the perception and decision steps to update the Rover's state
            #when a rock is found execute perception and rock_manuever picking up the rock
            if Rover.samples_located > Rover.samples_collected:             
                Rover = perception_step(Rover)
                Rover = rock_manuever(Rover)
                if Rover.vel == 0:
                    Rover = rover_stuck(Rover)
            if Rover.samples_located == Rover.samples_collected and Rover.reverse == 'False':  #if Rover.reverse flag is false and no rocks to pickup operate normally
                Rover.manuever_flag == 'False'
                Rover.mode = 'forward'
                Rover = perception_step(Rover)
                Rover = decision_step(Rover)
            #flag that is true when rover is stuck and needs to reverse straight back
            if Rover.reverse == 'True':
                Rover.steer = 0
                Rover.brake = 0
                Rover.throttle = -1
            #flag that is set to false when the rover velocity reaches -1m/s
            if Rover.reverse == 'True' and Rover.vel < -1:
                Rover.reverse = 'False'
                Rover.steer = 0
            #Flag that is true when the rover needs to perform a 180 degree turn
            if Rover.turn180 == 'True'and Rover.manuever_flag == 'False':
                Rover.steer = 15
                Rover.throttle = 0
                Rover.brake = 0
            #Set the turn180 flag back to false when there is navigable terrain in front of Rover    
            if len(Rover.nav_angles) > Rover.stop_forward:
                Rover.turn180 = 'False'
            if Rover.vel<0:
                Rover.steer = 0
                
            # Create output images to send to server
            out_image_string1, out_image_string2 = create_output_images(Rover)

            # The action step!  Send commands to the rover!
 
            # Don't send both of these, they both trigger the simulator
            # to send back new telemetry so we must only send one
            # back in respose to the current telemetry data.

            # If in a state where want to pickup a rock send pickup command
            print('##rover pickup status##')
            print(Rover.send_pickup)
            if Rover.send_pickup and not Rover.picking_up:
                print('##met conditions to send pickup##')
                print(Rover.send_pickup)
                send_pickup()
                print('##sent the pickup command##')
            else:
                # Reset Rover flags
				# Send commands to the rover!
                Rover.send_pickup = False
                commands = (Rover.throttle, Rover.brake, Rover.steer)
                send_control(commands, out_image_string1, out_image_string2)

        # In case of invalid telemetry, send null commands
        else:

            # Send zeros for throttle, brake and steer and empty images
            send_control((0, 0, 0), '', '')

        # If you want to save camera images from autonomous driving specify a path
        # Example: $ python drive_rover.py image_folder_path
        # Conditional to save image frame if folder was specified
        if args.image_folder != '':
            timestamp = datetime.utcnow().strftime('%Y_%m_%d_%H_%M_%S_%f')[:-3]
            image_filename = os.path.join(args.image_folder, timestamp)
            image.save('{}.jpg'.format(image_filename))

    else:
        sio.emit('manual', data={}, skip_sid=True)

@sio.on('connect')
def connect(sid, environ):
    print("connect ", sid)
    send_control((0, 0, 0), '', '')
    sample_data = {}
    sio.emit(
        "get_samples",
        sample_data,
        skip_sid=True)

def send_control(commands, image_string1, image_string2):
    # Define commands to be sent to the rover
    data={
        'throttle': commands[0].__str__(),
        'brake': commands[1].__str__(),
        'steering_angle': commands[2].__str__(),
        'inset_image1': image_string1,
        'inset_image2': image_string2,
        }
    # Send commands via socketIO server
    sio.emit(
        "data",
        data,
        skip_sid=True)
    eventlet.sleep(0)
# Define a function to send the "pickup" command 
def send_pickup():
    print("Picking up")
    pickup = {}
    sio.emit(
        "pickup",
        pickup,
        skip_sid=True)
    eventlet.sleep(0)
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Remote Driving')
    parser.add_argument(
        'image_folder',
        type=str,
        nargs='?',
        default='',
        help='Path to image folder. This is where the images from the run will be saved.'
    )
    args = parser.parse_args()
    
    #os.system('rm -rf IMG_stream/*')
    if args.image_folder != '':
        print("Creating image folder at {}".format(args.image_folder))
        if not os.path.exists(args.image_folder):
            os.makedirs(args.image_folder)
        else:
            shutil.rmtree(args.image_folder)
            os.makedirs(args.image_folder)
        print("Recording this run ...")
    else:
        print("NOT recording this run ...")
    
    # wrap Flask application with socketio's middleware
    app = socketio.Middleware(sio, app)

    # deploy as an eventlet WSGI server
    eventlet.wsgi.server(eventlet.listen(('', 4567)), app)
