from time import time as current_time
import sys
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
r0 = 1. #meters -- the length of cable in the first full-circle rotation
r_anchor = .2 #meters -- the radius of the anchor point
r_winch  = .2 #meters -- the max radius of the roll of cable
length = 1000 #meters -- the length of the cable
phi_max = 2 * (length - r0) / r_winch #rad
r = lambda phi: r0 + r_winch*(1 - phi/(2*phi_max))*phi # cable let out
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

### cable data:
global dia, l, max_tension
#diameter in m: [l in kg/m, breaking point in kN]
cables = {.006: [.023,   41.2], .008: [.039,   65.7],
          .009: [.042,   82.4], .010: [.059,  105.9],
          .011: [.063,  133.4], .012: [.095,  161.9],
          .014: [.128,  215.8], .016: [.160,  269.8],
          .018: [.208,  343.3], .020: [.255,  407.1],
          .022: [.305,  490.5], .024: [.358,  569.0],
          .026: [.410,  647.4], .028: [.465,  725.9],
          .030: [.520,  799.5], .032: [.570,  868.1],
          .034: [.625,  941.2], .036: [.680, 1020.2],
          .038: [.740, 1098.7], .040: [.840, 1245.8],
          .042: [.930, 1373.4], .044: [1.02, 1491.1],
          .046: [1.11, 1618.6], .048: [1.21, 1755.9],
          .050: [1.31, 1893.3], .052: [1.41, 2020.8],
          .056: [1.63, 2315.1], .060: [1.75, 2472.0],
          .064: [2.00, 2766.3], .068: [2.26, 3099.9],
          .072: [2.54, 3413.8], .080: [3.13, 4139.7],
          .088: [3.79, 4934.3], .096: [4.51, 5768.1]}
dia = .026
cable_choices = numpy.array(cables.keys())
nearest = lambda dia: cable_choices[numpy.abs(dia - cable_choices).argmin()]
l, max_tension = cables[nearest(dia)]

### masses:
global m, Mass
m = 39.7 #kilograms -- mass of the satellite cage (overestimate)
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)

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

#good final times: 20 185 1874
def solve_fun(func, init_vars, space=(0, 1874, 10**5)):
    time0 = current_time()
    time = numpy.linspace(*space)
    soln = odeint(func, init_vars, time)
    print "took %s seconds to solve" % (current_time() - time0)
    return time, soln
    
def interpret(soln, max_tension=max_tension):
    phi  = soln[:, 0]
    phi1 = soln[:, 1]
    v_tan = v_cm_co(phi) * phi1 * r(phi) / 1000 #kilometers/second
    cos_t = lambda phi: cos(numpy.arctan(tan(phi))) # cos(theta)
    tension = (m/cos_t(phi))*(phi1**2 * r(phi) - g(tau/4))/1000 #kiloNewtons
    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)
    
    return phi, phi1, v_tan, cos_t, tension, moment, momentum, max_arg

def plot(time, soln, graphs=['all'], i=7):
    phi, phi1, v_tan, cos_t, tension, moment, momentum, max_arg = \
        interpret(soln)
    phi_max, phi1_max, v_tan_max, tension_max, moment_max, momentum_max = \
        map(lambda x: numpy.ndarray.__getitem__(x, max_arg),
            [phi, phi1, v_tan, tension, moment, momentum])
    
    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, m, Mass, l, dia, max_tension
    r0, m, Mass, dia = guesses
    l, max_tension = cables[nearest(dia)]
    
    time, soln = solve_fun(func, init_vars)
    
    phi, phi1, v_tan, cos_t, tension, moment, momentum, max_arg = \
        interpret(soln)
    phi_max, phi1_max, v_tan_max, tension_max, moment_max, momentum_max = \
        map(lambda x: numpy.ndarray.__getitem__(x, max_arg),
            [phi, phi1, v_tan, tension, moment, momentum])
    
    f.write("%s, " % r0)
    f.write("%s, " % m)
    f.write("%s, " % Mass)
    f.write("%s, " % phi_max)
    f.write("%s, " % phi1_max)
    f.write("%s, " % tension_max)
    f.write("%s\n" % v_tan_max)

    return 2/v_tan_max + Mass/600 + phi1.max() / tau + cable*50

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

#outfile = 'C:\\Users\\dave\\Desktop\\out'
outfile = '../out'
def optimize():
    fil = file(outfile, 'w')
    fil.write('r0, sat. mass, Total Mass, cable diameter, '
            + 'phi, omega, Tension, Max Vel.\n')
    guess = (r0, 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():
    time, soln = solve_fun(throw, inits)
    plot(time, soln, ['all', '!gravity'])
    pyplot.show()

def compare():
    global m, Mass
    m = 39.7
    Mass = 650
    time, soln = solve_fun(throw, inits)
    plot(time, soln, ['tension'], 8)
    pyplot.title("Mass = %i" % Mass)
    
    m = 39.7
    Mass = 670
    time, soln = solve_fun(throw, inits)
    plot(time, soln, ['tension'], 7)
    pyplot.title("Mass = %i" % Mass)
    
    m = 39.7
    Mass = 690
    time, soln = solve_fun(throw, inits)
    plot(time, soln, ['tension'], 6)
    pyplot.title("Mass = %i" % Mass)
    
    m = 41.81320863
    Mass = 4172
    time, soln = solve_fun(throw, inits)
    plot(time, soln, ['tension'], 5) 
    pyplot.title("Mass = %i" % Mass)   
    
    pyplot.show()

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