from time import time as current_time
import sys, itertools
import numpy
from scipy.interpolate import interp1d
from scipy.integrate import odeint
from scipy.optimize import fmin
from matplotlib import pyplot

global time0
time0 = current_time()

tau = 2*numpy.pi
def taylor_sin(a):
    x = a % tau
    r = (x
        -x**3/6
        +x**5/120
        -x**7/5040
        +x**9/362880
        -x**11/39916800
        +x**13/6227020800
        -x**15/1307674368000
        +x**17/355687428096000
        -x**19/121645100408832000
        +x**21/51090942171709440000
        -x**23/25852016738884976640000
        +x**25/15511210043330985984000000
        -x**27/10888869450418352160768000000
        #+x**29/8841761993739701954543616000000
        #-x**31/8222838654177922817725562880000000
        #+x**33/8683317618811886495518194401280000000
        #-x**35/10333147966386144929666651337523200000000
        #+x**37/13763753091226345046315979581580902400000000
        #-x**39/20397882081197443358640281739902897356800000000
    )
    return r
def taylor_cos(a):
    x = a % tau
    r = (1
        -x**2/2
        +x**4/24
        -x**6/720
        +x**8/40320
        -x**10/3628800
        +x**12/479001600
        -x**14/87178291200
        +x**16/20922789888000
        -x**18/6402373705728000
        +x**20/2432902008176640000
        -x**22/1124000727777607680000
        +x**24/620448401733239439360000
        -x**26/403291461126605635584000000
        #+x**28/304888344611713860501504000000
        #-x**30/265252859812191058636308480000000
        #+x**32/263130836933693530167218012160000000
        #-x**34/295232799039604140847618609643520000000
        #+x**36/371993326789901217467999448150835200000000
        #-x**38/523022617466601111760007224100074291200000000
    )
    return r

### trig approximations
cos = taylor_cos
sin = taylor_sin

### radii:
global r0, r_roll, length
r0 = 1.151467 #meters -- the length of cable in the first full-circle rotation
r_anchor = .2 #meters -- the radius of the anchor point
r_roll  = .2045575 #meters -- the max radius of the roll of cable
length = 578.63682 #meters -- the length of the cable
phi_max = 2 * (length - r0) / r_roll #rad
r = lambda phi: r0 + r_roll*(1 - phi/(2*phi_max))*phi # cable length let out
r_dot_co = lambda phi: r_roll*(1 - phi/phi_max) #*phi1
r_sat = lambda phi: numpy.sqrt(r(phi)**2 + r_anchor) # radius of satellite

### gravitational acceleration:
geopotential = 3.987 * 10**14 #cubic meters per second^2
earth_radius = 6.3675 * 10**6 #meters
balloon_height = 40*10**3
height = lambda phi: balloon_height + r(phi) * cos(phi)
g = lambda phi: geopotential / (earth_radius + height(phi))**2

### quadratic friction:
density = .003 #kg/m^3
kappa = .25 # for a sphere
radius = .1 #meters -- radius of the satellite cage
Area = numpy.pi * radius **2 # for a sphere
quad_co = kappa * density * Area
# fric_quad = quad_co * (r * phi1 * vcm_co)**2

### mass:
global m, Mass
m = 39.7 #kilograms -- mass of the satellite cage (overestimate)
l = 0.16 #kilograms/meter -- density of Super Max cable, 16mm: 269.8kN
l = 0.41 #kilograms/meter -- density of Super Max cable, 26mm: 647.4kN
max_tension = 647.4 #kN, given by choice of cable above.
Mass = 670 #kilograms -- mass of /everything/
M = lambda phi: Mass - m - l*r(phi)
r_cm = lambda phi: (m*r(phi) + .5*l*r(phi)**2) / (m + l*r(phi) + M(phi))
v_cm_co = lambda phi: 1 - r_cm(phi)/r(phi)

def throw(phi, time):
    """The diff eqs for the windup."""
    phi, phi1 = phi
    prime_vars = [
      # phi' = phi'
        phi1,
      # phi'' =
        ( (r_sat(phi)*r_anchor/r(phi) - r_dot_co(phi)) * phi1**2
         + g(phi) * ( sin(phi) - cos(phi)*r_anchor/r(phi) )
         - quad_co * ( v_cm_co(phi) * r(phi) * phi1 )**2 / m
         ) / r_sat(phi)
      ]
    return prime_vars

#good final times: 20 185 1874
def solve_fun(func, init_vars, timespace):
    time0 = current_time()
    time = numpy.linspace(*timespace)
    soln = odeint(func, init_vars, time)
    print "  ::  took %s seconds" % (current_time() - time0)
    return time, soln

TimespaceTooSmall = Warning(" *** Requires larger timespace!!! *** ")
def interpret(soln, max_tension):
    phi  = soln[:, 0]
    phi1 = soln[:, 1]
    v_tan = v_cm_co(phi) * phi1 * r(phi) / 1000 #kilometers/second
    tension = (m*r_sat(phi)/r(phi))*(phi1**2 * r(phi) - g(tau/4))/1000 #kN
    cable_moment = l*((r(phi) - r_cm(phi))**3 - r_cm(phi)**3)/3 #kg m^2
    moment = m*(r(phi) - r_cm(phi))**2 + cable_moment - M(phi)*r_cm(phi)**2
    momentum = moment*phi1**2 #kg m^2/s^2
    
    try:
        maxindex = numpy.argmax(tension >= max_tension)
        max_arg = numpy.argmax(v_tan[:maxindex])
    except ValueError: max_arg = numpy.argmax(v_tan)
    
    if phi[max_arg] > abs(phi[-1] - tau):
        raise TimespaceTooSmall
    
    return phi, phi1, v_tan, tension, moment, momentum, max_arg

def solver(func, init_vars, timespace=(0,600,10**5), max_tension=max_tension,
    returns=['time', 'variables', 'slices', 'maxes']):
    time0 = current_time()
    for n in itertools.count(1):
        time, soln = solve_fun(func, init_vars, timespace)
        
        try: variables, max_arg = interpret(soln, max_tension)
        except TimespaceTooSmall:
            timespace = map(lambda x: x*.2*1.5**n, timespace)
            continue
        break
    
    getslice = lambda x: numpy.ndarray.__getslice__(x, 0, max_arg+1)
    getmax   = lambda x: numpy.ndarray.__getitem__( x,    max_arg   )
    slices = map(getslice, (time,) + variables)
    maxes  = map(getmax,   (time,) + variables)
    
    print ": took %s seconds to solve" % (current_time() - time0)
    
    out = []
    for item in returns: out.append(vars()[item])
    return tuple(out)

def plot(solved, graphs=['all'], i=7):
    time, variables, slices, maxes = solved
    phi, phi1, v_tan, tension, moment, momentum = variables
    (phi_max, phi1_max, v_tan_max,
         tension_max, moment_max, momentum_max) = maxes
    
    def test(*args):
        b = 0
        for string in args:
            b += string in graphs or \
                ('all' in graphs and "!" + string not in graphs)
        return bool(b)
    
    i -= 1
    if test('gravity'):
        pyplot.figure(i)
        pyplot.plot(time, g(phi))
        pyplot.grid(True)
        pyplot.ylabel("gravitational acceleration (m/s^2)")
        pyplot.xlabel("time (seconds)")
    
    i -= 1
    if test('phi', 'r', 'radius'):
        pyplot.figure(i)
        pyplot.plot(time, phi / tau, label="phi (cycles)")
        pyplot.plot(time, r(phi), label="r (m)")
        pyplot.grid(True)
        pyplot.legend(loc="upper left")
        pyplot.title("phi and radius vs. time")
        pyplot.ylabel("phi (cycles) and radius (meters)")
        pyplot.xlabel("time (seconds)")
    
    i -= 1
    if test("phi'", "omega", "phi1"):
        pyplot.figure(i)
        pyplot.plot(time, phi1 / tau)
        pyplot.grid(True)
        pyplot.ylabel("phi' (Hertz)")
        pyplot.xlabel("time (seconds)")

    i -= 1
    if test('v_tan', 'tension'):
        pyplot.figure(i)
        rect = [.1, .1, .8, .8]
        
        ax1 = pyplot.axes(rect)
        ax1.yaxis.tick_left()
        ax1.xaxis.grid(True)
        pyplot.plot(phi / tau, v_tan, marker='+', markevery=len(phi)/50)
        pyplot.ylabel("+: tangential velocity (km/s)")
        pyplot.xlabel("phi (cycles)")
        pyplot.ylim(ymin=0)
        
        ax2 = pyplot.axes(rect, frameon=False, sharex=ax1)
        ax2.yaxis.tick_right()
        ax2.yaxis.set_label_position('right')
        pyplot.setp(ax2.get_xticklabels(), visible=False)
        pyplot.plot(phi / tau, tension, marker='x', markevery=len(phi)/50)
        pyplot.ylabel("x: Tension (kN)")
        pyplot.ylim(ymin=0)
        
        pyplot.axvline(x=phi_max/tau)
    
    i -= 1
    if test('moment'):
        pyplot.figure(i)
        pyplot.plot(phi / tau, moment)
        pyplot.grid(True)
        pyplot.ylabel("Moment of Inertia (kg m^2)")
        pyplot.xlabel("phi (cycles)")
    
    i -= 1
    if test('momentum'):
        pyplot.figure(i)
        pyplot.plot(phi / tau, momentum)
        pyplot.grid(True)
        pyplot.ylabel("Angular Momentum (kg m^2 s^-2)")
        pyplot.xlabel("phi (cycles)")
    
    #import pdb; pdb.set_trace()

def opt_fun(guesses, f, func, init_vars):
    """Solve, plot, and compare the diff eq."""
    global r0, r_roll, length, m, Mass
    r0, r_roll, length, m, Mass = guesses
    
    time, soln = solve_fun(func, init_vars)
    
    phi, phi1, v_tan, tension, moment, momentum, max_arg = \
        interpret(soln)
    phi, phi1, v_tan, tension, moment, momentum, time = \
        map(lambda x: numpy.ndarray.__getslice__(x, 0, max_arg+1),
            [phi, phi1, v_tan, tension, moment, momentum, time])
    
    time, (phi, phi1, v_tan, tension, moment, momentum, max_arg) = \
        solver(func, init_vars, returns=['time', 'slices'])
    
    # Parameter values.
    f.write("%s, " % r_roll)
    f.write("%s, " % length)
    f.write("%s, " % r0)
    f.write("%s, " % m)
    f.write("%s, " % Mass)
    
    # Values at the max velocity point.
    f.write("%s, " %     phi[-1])
    f.write("%s, " %    phi1[-1])
    f.write("%s, " % tension[-1])
    f.write("%s\n" %   v_tan[-1])

    # If you want strict limits, you can change the return to:
    #return 1/v_tan + 1000*(Mass > Mass_limit) + 1000*(phi1.max() > phi1_limit)
    return 2/v_tan[-1] + Mass/600 + phi1[-1] / tau

v0 = numpy.sqrt(g(0) * r(0))
w0 = v0/r0
inits = [0, w0]

#outfile = 'C:\\Users\\dave\\Desktop\\out'
outfile = '/home/jandew/Documents/Impossible Challenges/microsatellite/out'
def optimize():
    fil = file(outfile, 'w')
    fil.write(
    'r0, r_roll, length of cable, m, Mass, phi, omega, Tension, Max Vel.\n')
    guess = (r0, r_roll, length, m, Mass)
    #try: print opt_fun(guess, fil, throw, inits)
    try: print fmin(opt_fun, guess, args=(fil, throw, inits))
    finally: fil.close()

def solve():
    solved = solver(throw, inits)
    plot(solved, ['all', '!gravity'])
    pyplot.show()

if len(sys.argv) == 2:
    func_name = sys.argv[1]
    if func_name in dir():
        exec(func_name + '()')