#! /usr/bin/env python

# golf_angle_optimization.py

import numpy
import pylab, matplotlib
import scipy.optimize

import cgolf

cgolf.h = 0.01

Vs = numpy.arange(45.,82,2)
rpms = numpy.arange(2000.,4100,100)
angles = numpy.zeros([len(Vs),len(rpms)])


# Optimize the angle for all of the different launch speeds
for i,V in enumerate(Vs):
    for j,rpm in enumerate(rpms):
        angles[i,j] = scipy.optimize.brent(
            lambda theta: -cgolf.get_total_distance(
                *cgolf.get_trajectory(V,theta,rpm)), brack=[1,15,40],
            tol=0.01)

# Contour plotting
pylab.figure(figsize=(6.,5))
pylab.contourf(rpms,Vs,angles,50)

colorbar = pylab.colorbar(format='%.0f',orientation='horizontal',
                          ticks=matplotlib.ticker.MaxNLocator(5),
                          shrink=0.8)
colorbar.set_label('Optimum angle (deg)')

pylab.xlim(2000,4000)
pylab.ylim(45,80)

pylab.xlabel('Spin rate (rpm)')
pylab.ylabel('Launch speed (m/s)')

# Display the plot
pylab.show()