import plotly
import time
import json
import math
from numpy import sin, cos, pi, array
import numpy as np
import scipy.integrate as integrate

# Fill in the config.json file in this directory with your plotly username, 
# plotly API key, and your generated plotly streaming tokens
# Sign up to plotly here: https://plot.ly/ssu
# View your API key and streaming tokens here: https://plot.ly/settings

with open('./config.json') as config_file:
    plotly_user_config = json.load(config_file)

username = "Streaming-Demos" # plotly_user_config['plotly_username'] 
api_key = plotly_user_config['plotly_api_key']
stream_tokens = plotly_user_config['plotly_streaming_tokens']

# Initialize your plotly object
p = plotly.plotly(username, api_key)

# Define the Double Pendulum state equations and parameters
# Code adapted from http://matplotlib.org/examples/animation/double_pendulum_animated.html (Thx!)

G =  9.8 # acceleration due to gravity, in m/s^2
L1 = 1.0 # length of pendulum 1 in m
L2 = 1.0 # length of pendulum 2 in m
M1 = 1.0 # mass of pendulum 1 in kg
M2 = 1.0 # mass of pendulum 2 in kg


def derivs(state, t):

    dydx = np.zeros_like(state)
    dydx[0] = state[1]

    del_ = state[2]-state[0]
    den1 = (M1+M2)*L1 - M2*L1*cos(del_)*cos(del_)
    dydx[1] = (M2*L1*state[1]*state[1]*sin(del_)*cos(del_)
               + M2*G*sin(state[2])*cos(del_) + M2*L2*state[3]*state[3]*sin(del_)
               - (M1+M2)*G*sin(state[0]))/den1

    dydx[2] = state[3]

    den2 = (L2/L1)*den1
    dydx[3] = (-M2*L2*state[3]*state[3]*sin(del_)*cos(del_)
               + (M1+M2)*G*sin(state[0])*cos(del_)
               - (M1+M2)*L1*state[1]*state[1]*sin(del_)
               - (M1+M2)*G*sin(state[2]))/den2

    return dydx

# th1 and th2 are the initial angles (degrees)
# w10 and w20 are the initial angular velocities (degrees per second)
th1 = 120.0
w1 = 0.0
th2 = -10.0
w2 = 0.0

rad = pi/180

# Initialize your plotly real-time streaming graph with a REST API call
# Embed your stream token in one of the traces of a plotly-data object - one token per trace
# Also embed 'maxpoints', the number of points that you want plotted at a time

# The `iplot` command will embed our plotly graph as an iframe in this notebook
# Each plotly graph has a unique url that you can share 
# Anyone can view your streaming graph in real-time - 
# viewers see the same data at the same time 
# (try it! Open up this notebook up in two different browser windows and observer
# that the graphs are plotting identical data!)

# The unique URL for this graph is https://plot.ly/~streaming-demos/4

def axes(range):
    return {
          "autotick" : True,
          "showgrid" : False,
          "showline" : False,
          "zeroline" : False,
          "ticks": '',  
          "range": range
    }
p.iplot([
        {
            'x': [], 
            'y': [], 
            'type': 'scatter', 
            'mode': 'lines+markers', 
            'line': {
                'opacity': 0.8
            },
            'marker': {
                'size': 12
            },
            'stream': {
                'token': stream_tokens[4], 
            }
        },
        {
            'x': [], 
            'y': [], 
            'type': 'scatter', 
            'mode': 'lines', 
            'line': {
                'color': 'rgba(31, 119, 180, 0.15)'
            },
            'stream': {
                'token': stream_tokens[5], 
                'maxpoints': 100
            }
        }],
        layout = {
            'xaxis': axes([-2, 2]), 
            'yaxis': axes([-2, 0.5]), 
            'title': 'Double Pendulum Simulation',
            'showlegend': False
        },
        filename='streaming double pendulum', fileopt='overwrite')

# Now stream data to this plot!
# Write to a plotly stream with plotly-formatted data
# In this case, we'll continually integrate the double pendulum equations 20 points a time, 
# and write the solution, one point at a time

# Our plotly-object is formatted as: {'x': x_data, 'y': y_data}

s0 = plotly.stream(stream_tokens[4])
s1 = plotly.stream(stream_tokens[5])
state = np.array([th1, w1, th2, w2])*pi/180.
while True:
    # create a time array from 0..100 sampled at 0.1 second steps
    dt = 0.05
    t = np.arange(0.0, 20, dt)
    # solve!
    y = integrate.odeint(derivs, state, t)
    x1 = L1*sin(y[:,0])
    y1 = -L1*cos(y[:,0])
    x2 = L2*sin(y[:,2]) + x1
    y2 = -L2*cos(y[:,2]) + y1
    # Write the solution to plotly's servers, one point at a time
    for (x1i, y1i, x2i, y2i) in zip(x1, y1, x2, y2):
        s0.write({'x': [0, x1i, x2i], 'y': [0, y1i, y2i]})
        s1.write({'x': x2i, 'y': y2i})
        time.sleep(0.08) # plot points 80 ms at a time

    # Set the new initial state
    state = np.array([y[-1,0], y[-1, 1], y[-1, 2], y[-1, 3]])

# When you're done, close your stream!
s.close()

# CSS styling within IPython notebook
from IPython.core.display import HTML
import urllib2
def css_styling():
    url = 'https://raw.githubusercontent.com/plotly/python-user-guide/master/custom.css'
    styles = urllib2.urlopen(url).read()
    return HTML(styles)

css_styling()