from visual import *
import math
import random
import numpy
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import matplotlib.lines as mlines

# TODO: have max speed allowable for making putt

#------------------------------------------------global constants-------------------------------------------
M = .045            # golf ball about .045 kg
G = 9.8             # acceleration of gravity
MU = .04            # rough estimate for coefficient of friction
HOLE_RADIUS = .054  # radius of hole in m

#-----------------------------------------------classes for models-------------------------------------------
"""A 3D model of a surface. The user can pick from several predefined
models, defined by the Type enum."""
class Green_Map(object):
    """Enum for selecting partial derivatives."""
    class Partial(object):
        X = 0
        Y = 1

    """Enum for the type of surface to generate."""
    class Type(object):
        PLANE = 1
        STEEP_PLANE = 2
        CENTRAL_HILL = 3
        BOWL = 4
        RIDGE = 5
    
    STEP = .01 # each point is .01 m apart
    z_vals = numpy.empty((900, 900)) # 2-D array for 9m x 9m surface

    """Generate the specified surface.
    Param:
        choice - the model to generate
    """
    def __init__(self, choice):
        max_x = self.z_vals.shape[0]
        max_y = self.z_vals.shape[1]
        if choice == self.Type.PLANE: # plane z = y/9
            for i in range(max_x):
                for j in range(max_y):
                    self.z_vals[i][j] = j/float(max_x)
        elif choice == self.Type.STEEP_PLANE: # plane z = 4y/9
            for i in range(max_x):
                for j in range(max_y):
                    self.z_vals[i][j] = j/(float(max_x)/4.0)
        elif choice == self.Type.CENTRAL_HILL: # z = 1/(1+r^2)
            center = vector(max_x/2.0*self.STEP, max_y/2.0*self.STEP, 0.0)
            for i in range(max_x):
                for j in range(max_y):
                    current = vector(i*self.STEP, j*self.STEP, 0)
                    self.z_vals[i][j] = 1.0/(1.0 + mag2(center-current))
        elif choice == self.Type.BOWL: # z = .025r^2
            center = vector(max_x/2.0*self.STEP, max_y/2.0*self.STEP, 0.0)
            for i in range(max_x):
                for j in range(max_y):
                    current = vector(i*self.STEP, j*self.STEP, 0)
                    self.z_vals[i][j] = .025*mag2(center-current)
        elif choice == self.Type.RIDGE: # z = y^(1/3)
            max_delta = .1
            for i in range(max_x):
                for j in range(max_y):
                    y = (j-max_x/2.0)/float(max_x/2.0)
                    if abs(y) < .001:
                        y = .001
                    self.z_vals[i][j] = abs(y)**(1/3.0)*abs(y)/y
        else:
            raise ValueError("Invalid choice.")

    """Returns a random number within delta of the given value."""
    def randrange(self, num, delta):
        min_val = num-delta
        return 2*delta*random.random() + min_val

    """Adds a scatter plot of the surface to the given axes.
    Param:
        ax - the Axes3D object to plot the graph on
        graph_step - how many points to skip while graphing
    """
    def graph(self, ax, graph_step):        
        size = self.z_vals.size/(int(graph_step**2 + .5))
        X = numpy.empty(size)
        for i in range(size):
            X[i] = i/(self.z_vals.shape[0]/graph_step) * graph_step * self.STEP

        Y = numpy.empty(size)
        for i in range(size):
            Y[i] = i%(self.z_vals.shape[1]/graph_step) * graph_step * self.STEP

        Z = numpy.empty(size)
        for i in range(0, self.z_vals.shape[0], graph_step):
            for j in range(0, self.z_vals.shape[1], graph_step):
                Z[i + j/graph_step] = self.z_vals[i][j]

        ax.scatter(X, Y, Z, color='g', marker='o')

    """Return the value of the partial derivative with respect to the given variable
    at the given position (x,y).
    Param:
        var - the variable to take the partial with respect to
        pos - a vector <x, y>
    Return:
        df/dx or df/dy
    """
    def partial(self, var, pos):
        retval = 0
        if var == self.Partial.X:
            x = pos[0]
            z1 = self.get(pos)
            z2 = self.get((x+self.STEP, pos[1]))
            retval = (z2-z1)/self.STEP
        elif var == self.Partial.Y:
            y = pos[1]
            z1 = self.get(pos)
            z2 = self.get((pos[0], y+self.STEP))
            retval = (z2-z1)/self.STEP
        return retval
    
    """Get the z-value at the given position (x,y).
    Param:
        pos - a vector <x, y>
    Return:
        f(x, y)
    """
    def get(self, pos):
        x = int(pos[0]*1/self.STEP +.5)
        y = int(pos[1]*1/self.STEP +.5)
        return self.z_vals[x][y]

    """Get a vector <x, y, f(x, y)> at the given position (x,y)."""
    def get_vec(self, pos):
        return vector(pos[0], pos[1], self.get(pos))

    """Gets an upward normal vector to the surface at (x,y). The z component
    is guaranteed to be 1.
    Param:
        pos - a vector <x, y>
    Return:
        the normal vector <-df/dx, -df/dy, 1>
    """
    def up_normal(self, pos):
        return vector(-self.partial(Green_Map.Partial.X, pos), -self.partial(Green_Map.Partial.Y, pos), 1)

"""A class that represents the correct path to the hole."""
class Path(object):
    X = []
    Y = []
    Z = []

    """Initializes a new path object."""
    def __init__(self):
        self.X = []
        self.Y = []
        self.Z = []

    """Adds a point to this path.
    Param:
        r - a vector <x, y, z>
    """
    def add_point(self, r):
        self.X.append(r[0])
        self.Y.append(r[1])
        self.Z.append(r[2])

    """Adds a line graph of this path to the given axes.
    Param:
        ax - the Axes3D object to plot the path on
    """
    def graph(self, ax):
        ax.plot(self.X, self.Y, self.Z, color='k')
        
#------------------------------------------------simulation-------------------------------------------
"""The simulation!"""
def main():
    surf_options = ("Exit", "Plane", "Steep plane", "Central hill", "Bowl", "Ridge")
    debug_options = ("None", "Show error (recommended)", "Show plots", "Everything... grab some popcorn")
    choice = display_menu("\nSelect a surface type:", surf_options)
    while choice != 0:
        # create surface
        print "\nGenerating 9m x 9m surface..."
        green_map = Green_Map(choice)
        
        # get start and end positions
        print "\nChoose start and end positions on interval (0, 9)."
        start_x = get_input("x coordinate of start: ", 0, 9)
        start_y = get_input("y coordinate of start: ", 0, 9)
        end_x = get_input("x coordinate of end: ", 0, 9)
        end_y = get_input("y coordinate of end: ", 0, 9)

        # get debug level
        debug = display_menu("\nChoose level of debug info:", debug_options)

        # run simulation
        print "\nRunning simulation..."
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d') # init graph
        speed, direction = find_path(green_map, ax, vector(start_x, start_y, 0), vector(end_x, end_y, 0), debug)

        # display results
        print "\nRequired speed (m/s):", speed
        print "Required direction (degrees CCW from straight line to hole):", direction
        print "Close graph to continue..."
        plt.show()

        # again?
        choice = display_menu("\nSelect a surface type:", surf_options)

"""Displays a menu and returns the user's choice."""
def display_menu(message, options):
    print message
    i = 0
    for s in options:
        print '\t' + str(i) + " - " + s
        i += 1
    choice = int(raw_input("Choice: "))
    return choice

"""Gets user input, which must be between the given values."""
def get_input(message, min_val, max_val):
    x = float(raw_input(message))
    while x < min_val or x > max_val:
        print x, "is not between", min_val, "and", max_val
        print "Shame on you."
        x = float(raw_input(message))
    return x

"""Tries various initial speeds and starting directions until it finds a
path that gets the ball within the radius of the hole. It only considers
the path in 2D, since the z-value at every (x, y) position has to
be a point on the surface.
Param:
    green_map - 3-D model of the surface
    ax - the Axes3D object to use for plotting
    r0 - the starting position, a vecotr <x0, y0, 0>
    rf - the position of the hole, a vector <xf, yf, 0>
    debug - int inidicating how much information to print
Return:
    v0 - the required initial speed
    direction - the required starting direction (in degrees CCW from a straight-line
                path to the hole
"""
def find_path(green_map, ax, r0, rf, debug):
    # get rough initial estimate for speed
    v0 = estimate_speed(green_map, r0, rf)
    debug_init(green_map, v0, r0, rf, debug)
    num_resets = 0 # used if simulation has trouble finding path

    # start simulation
    direction = 0.0 # first try direction toward hole
    error, was_long, missed_right, path, reset = try_path(green_map, r0, rf, v0, direction, debug)
    last_was_long = was_long
    last_missed_right = missed_right
    speed_step = v0/20.0 # init increment for changing speed
    if was_long: # make speed increment negative if putt went too far
        speed_step *= -1
    direction_step = 1.0 # init increment for changing direction (1 degree CCW)
    if not missed_right: # make direction increment negative (CW) if putt missed to left
        direction_step *= -1
    debug_find_path(green_map, r0, rf, path, error, was_long, missed_right, speed_step, direction_step, debug)
    
    # refine estimate until ball would go in hole
    while error > HOLE_RADIUS:
        # check if the ball missed in a different way, if so switch direction and decrease step
        if was_long != last_was_long: # reverse speed increment
            speed_step /= -10.0
        if missed_right != last_missed_right: # reverse direction increment
            direction_step /= -10.0
        # increment speed, direction
        v0 += speed_step
        direction += direction_step
        # try a new path
        last_was_long = was_long
        last_missed_right = missed_right
        error, was_long, missed_right, path, reset = try_path(green_map, r0, rf, v0, direction, debug)
        debug_find_path(green_map, r0, rf, path, error, was_long, missed_right, speed_step, direction_step, debug)
        
        # check if simulation needs to be reset
        if reset or abs(speed_step) < 1.0e-3 or abs(direction_step) < 1.0e-3: # simulation failed, try again
            num_resets += 1
            v0 = estimate_speed(green_map, r0, rf)*(2**(num_resets)) # try again with greater v0
            direction = 0.0 # start by trying direction toward hole
            error, was_long, missed_right, path, reset = try_path(green_map, r0, rf, v0, direction, debug)
            last_was_long = was_long
            last_missed_right = missed_right
            speed_step = v0/20.0 # init increment for changing speed
            if was_long: # make speed increment negative if putt went too far
                speed_step *= -1
            direction_step = 1.0 # init increment for changing direction (1 degree)
            if not missed_right: # make direction increment negative (CW) if putt missed to left
                direction_step *= -1
            if debug > 0:
                print "reset..."
            debug_find_path(green_map, r0, rf, path, error, was_long, missed_right, speed_step, direction_step, debug)

    # update graph
    plot_hole(ax, green_map, r0, rf) # add start and stop positions to graph
    green_map.graph(ax, 30) # add the surface to the graph
    path.graph(ax) # add the path the ball took to the graph
    ax.set_zlim3d(-4.5, 4.5)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    # done!
    return v0, direction

"""Estimates the initial speed needed to reach the hole from the given
starting position.
Param:
    green_map - 3D model of surface
    r0 - start position vector
    rf - position vector of hole
Return:
    an initial speed estimate
"""
def estimate_speed(green_map, r0, rf):
    # E0 + Wfric = Ef
    # K0 + Wfric = Ugf
    # .5*M*v0^2 = MG(zf-z0) + MU*M*G*D
    # .5*v0^2 = G(zf-z0) + MU*G*D
    # v0 = sqrt(2G(zf-zo) + 2MU*G*D)
    # where D is distance between r0 and rf
    distance = mag(rf - r0)
    z0 = green_map.get(r0)
    zf = green_map.get(rf)
    v0squared = 2*G*(zf-z0) + 2*MU*G*distance
    if (v0squared < 0):
        return .1
    else:
        return math.sqrt(v0squared)

"""Tries a path and runs until the ball starts moving away from the hole
(after which it has no chance of going in).
Param:
    green_map - 3D model of surface
    r0 - start position vector
    rf - position vector of hole
    v0 - initial speed
    direction - initial direction, in degrees CCW from straight-line path
    debug - int describing level of debug information to print
Return:
    error - the minimum distance the ball passed to the hole
    was_long - boolean indicating if the ball went past the hole
    missed_right - boolean indicating if the ball missed on the right side
    path - the Path object representing the path the ball travelled
    reset - boolean indicating whether the simulation needs to be reset (true
            if direction doesn't point toward the hole)
"""
def try_path(green_map, r0, rf, v0, direction, debug):
    # init path vars
    r = vector(r0[0], r0[1], r0[2]) # current displacement
    rnet = rf - r0 # displacement vector from start to hole
    path = Path()
    path.add_point(green_map.get_vec(r))
    # init velocity vector
    dir_vec = norm(rnet) # unit vector toward hole
    dir_vec = rotate(dir_vec, radians(direction), green_map.up_normal(r0)) # rotate
    v = v0 * dir_vec # scale direction vector to get velocity
    # init distance vars
    current_dist_to_hole = mag(rnet) # current distance from hole
    min_dist_to_hole = current_dist_to_hole # closest distance ball passed to hole
    dist_to_hole = mag(rf-r0) # distance from start to hole
    total_dist = 0 # displacement of ball from start
    # init time vars
    dt = .001 # small time increment

    # update r, v until ball starts going wrong way
    moving_toward = dot((rnet), v) > 0 # boolean to track if ball is actually moving toward hole
    if not moving_toward: # init velocity pointing away from hole, tell simulation to reset
        return min_dist_to_hole, False, False, path, True
    while moving_toward:
        try:
            v = update_v(green_map, r, v, dt, debug)
        except IndexError:
            return min_dist_to_hole, False, False, path, True
        r = update_r(r, v, dt)
        current_dist_to_hole = mag(rf - r)
        total_dist = mag(r-r0)
        if current_dist_to_hole < min_dist_to_hole: # update minimum dist, if necessary
            min_dist_to_hole = current_dist_to_hole
        moving_toward = dot((rf-r), v) > 0 # v still has component toward hole if dot > 0
        path.add_point(green_map.get_vec(r))
    debug_try_path(direction, v0, v, r, r0, rnet, debug)

    # get return values
    was_long = total_dist > dist_to_hole*.99 # check if end displacement is greater than dist to hole
    missed_right = cross(r-r0, rnet)[2] > 0 # final pos is to right of hole if z-comp of cross product points up
                               
    return min_dist_to_hole, was_long, missed_right, path, False

"""Updates the given velocity vector by calculating the net force and then
applying an acceleration to the velocity over the time increment.
Param:
    green_map - 3D model of the surface
    r - the current position vector
    v - the current velocity vector
    dt - a small time increment
    debug - int indicating amount of information to print
Return:
    an updated velocity vector
"""
def update_v(green_map, r, v, dt, debug):
    # get direction vector for gravity
    fx = green_map.partial(Green_Map.Partial.X, (r[0], r[1])) # df/dx
    fy = green_map.partial(Green_Map.Partial.Y, (r[0], r[1])) # df/dy
    grav_dir = vector(-fx, -fy, 0)
    grav_dir = norm(grav_dir)
    # get direction vector for friction
    fric_dir = -norm(v) # acts opposite to direction of travel
    # get angle of inclination
    up_normal = green_map.up_normal((r[0], r[1])) # updward unit normal vector
    theta = math.acos(1/mag(up_normal))   
    # get magnitude of friction force
    fric_mag = MU*M*G*math.cos(theta)
    # get magnitude of gravity force
    grav_mag = M*G*math.sin(theta)
    # get net force
    fric_force = fric_mag * fric_dir
    grav_force = grav_mag * grav_dir
    fnet = fric_force + grav_force
    # return new velocity
    accel = fnet / M
    v = v + accel*dt
    v[2] = 0
    debug_update_v(green_map, r, v, grav_dir, fric_dir, theta, debug)
    return v

"""Updates the given position vector using the time increment and
the given instantaneous veolocity.
Param:
    r - current position vector
    v - current velocity vector
    dt - small time increment
Return:
    an updated position vector
"""
def update_r(r, v, dt):
    return r + v*dt

"""Plots the start location and end location on the graph.
Param:
    ax - the Axes3D object to plot the hole on
    green_map - 3D model of the surface
    r0 - start position vector
    rf - end position vector
"""
def plot_hole(ax, green_map, r0, rf):
    start_vec = green_map.get_vec(r0) # add z-value to position vectors (z=0 initially)
    end_vec = green_map.get_vec(rf)
    X1 = [start_vec[0]]
    Y1 = [start_vec[1]]
    Z1 = [start_vec[2]]
    X2 = [end_vec[0]]
    Y2 = [end_vec[1]]
    Z2 = [end_vec[2]]
    ax.scatter(X1, Y1, Z1, color='b', marker='^', s=160)
    ax.scatter(X2, Y2, Z2, color='r', marker='^', s=160)
    proxy_artist1 = mlines.Line2D([], [], color='b', marker='^')
    proxy_artist2 = mlines.Line2D([], [], color='r', marker='^')
    ax.legend([proxy_artist1, proxy_artist2], ["Start", "End"])

#-----------------------------------------------debug functions-------------------------------------------
"""Function for debugging the init state of the putt."""
def debug_init(green_map, v0, r0, rf, debug):
    if debug > 0:
        print "INIT"
        print "Start pos", green_map.get_vec(r0)
        print "End pos", green_map.get_vec(rf)
        print "Init speed estimate", v0

"""Function for debugging find_path."""
def debug_find_path(green_map, r0, rf, path, error, was_long, missed_right, speed_step, direction_step, debug):
    if debug > 1:
        print "FIND PATH"
        print "Error", error
        print "Was long", was_long
        print "Missed right", missed_right
        print "V step", speed_step
        print "Dir step", direction_step

        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        plot_hole(ax, green_map, r0, rf)
        green_map.graph(ax, 30)
        ax.set_xlabel('x')
        ax.set_ylabel('y')
        ax.set_zlabel('z')
        path.graph(ax)
        ax.set_zlim3d(-4.5, 4.5)
        plt.show()
    elif debug > 0:
        print "Error", error

"""Function for debugging try_path."""
def debug_try_path(direction, v0, v, r, r0, rnet, debug):
    if debug > 2:
        print "TRY PATH"
        print "Init direction", direction
        print "Init speed", v0
        print "Final v", v
        print "Final speed", mag(v)
        print "Final r", r
        print "Change in r", r-r0
        print "Required change in r", rnet
        raw_input("Press enter...")

"""Function for debugging update_v."""
def debug_update_v(green_map, r, v, grav_dir, fric_dir, theta, debug):
    if debug > 2:
        print "UPDATE_V"
        print "r", r
        print "v", v
        print "fx", green_map.partial(Green_Map.Partial.X, (r[0], r[1]))
        print "fy", green_map.partial(Green_Map.Partial.Y, (r[0], r[1]))
        print "grav dir", grav_dir
        print "fric dir", fric_dir
        print "angle of incl", theta
        raw_input("Press enter...")

main()