from IPython.display import Image
Image(url='http://i.imgur.com/l9Yzfcs.gif')

from math import pi as PI
from math import sin, cos
import random

from plotly.widgets import GraphWidget
from plotly.graph_objs import *
import plotly.plotly as py
import plotly.tools as tls

from IPython.html import widgets 
from IPython.display import display, clear_output

root = 12

# color the tree with a gradient from root_col to tip_col
# interpolate linearly to get color at a given position in the gradient
def get_col(root_col, tip_col, iterat):
    r = ((iterat*1.0/root)*(root_col[0]-tip_col[0]))+tip_col[0]
    g = ((iterat*1.0/root)*(root_col[1]-tip_col[1]))+tip_col[1]
    b = ((iterat*1.0/root)*(root_col[2]-tip_col[2]))+tip_col[2]
    return '#%02x%02x%02x' % (r,g,b)
    
def tree_graph(iterat=12, branch_angle=18.0):
    # angle to radian factor
    ang2rad = PI/180.0
    # experiment with trunk length (try 120)
    t = 120
    # experiment with factor to contract the trunk each iteratation (try 0.7)
    r = 0.7
    # starting orientation (initial 90 deg)
    theta = 90.0 * ang2rad
    # experiment with angle of the branch (try 18 deg)
    dtheta = branch_angle * ang2rad
    # experiment with gradient color choices
    root_col = (40,40,40)
    tip_col = (250,250,250)
    # experiment with factor to increase random angle variation as child branches get smaller
    iterscale = 6.0
    # center of bottom
    origin = (250, 500)
    root=iterat
    # make the tree
    
    def fractal_tree(lines, iterat, origin, t, r, theta, dtheta, root_col, tip_col, randomize=False):
        """
        draws branches
        iterat:     iteratation number, stop when iterat == 0
        origin:   x,y coordinates of the start of this branch
        t:        current trunk length
        r:        factor to contract the trunk each iteratation
        theta:    starting orientation
        dtheta:   angle of the branch
        """
        if iterat == 0:
            return lines
        # render the branch
        x0, y0 = origin

        # randomize the length
        randt = random.random()*t
        if randomize:
            x, y = x0 + randt * cos(theta), y0 - randt * sin(theta)
        else:
            x, y = x0 + cos(theta), y0 - sin(theta)
        # color the branch according to its position in the tree
        col = get_col(root_col, tip_col, iterat)
        # add to traces
        lines.append(Scatter(x=[x0, x], y=[y0, y], mode='lines', line=Line(color=col, width=1)))
        # recursive calls
        if randomize:
            fractal_tree(lines, iterat-1, (x,y), t * r, r, theta + (random.random())*(iterscale/(iterat+1))*dtheta, dtheta, root_col, tip_col, randomize)
            fractal_tree(lines, iterat-1, (x,y), t * r, r, theta - (random.random())*(iterscale/(iterat+1))*dtheta, dtheta, root_col, tip_col, randomize) 
        else: 
            fractal_tree(lines, iterat-1, (x,y), t * r, r, theta + dtheta, dtheta, root_col, tip_col, randomize)
            fractal_tree(lines, iterat-1, (x,y), t * r, r, theta - dtheta, dtheta, root_col, tip_col, randomize)
    
    lines = []
    fractal_tree(lines, iterat, origin, t, r, theta, dtheta, root_col, tip_col, True)
    
    # group the lines by similar color
    branches = {}
    for line in lines:
        color = line['line']['color']
        if color not in branches:
            branches[color] = line
        else:
            branches[color]['x'].extend(line['x'])
            branches[color]['y'].extend(line['y'])
            branches[color]['x'].append(None)
            branches[color]['y'].append(None)
            
    
    branch_data = [branches[c] for c in branches]
    
    return branch_data

g = GraphWidget('https://plot.ly/~chris/4103')

iterations = widgets.IntSliderWidget()
iterations.min= 2
iterations.max= 14
iterations.value = 12
iterations.description = 'iterations'

branch_angle = widgets.IntSliderWidget()
branch_angle.min = 5
branch_angle.max = 90
branch_angle.value = 18
branch_angle.description = 'branch angle'

seed = widgets.ButtonWidget()
seed.description = 'Grow a new tree'

def regraph_tree(_):
    branch_data = tree_graph(iterations.value, branch_angle.value)

    g.restyle({'x': [[]], 'y': [[]]})
    g.add_traces(branch_data)

iterations.on_trait_change(regraph_tree, 'value')
branch_angle.on_trait_change(regraph_tree, 'value')
seed.on_click(regraph_tree)
    
display(iterations)
display(branch_angle)
display(seed)
display(g)

layout = Layout(yaxis=YAxis(autorange='reversed'), width=500, showlegend=False)

g.relayout(layout)

# CSS styling within IPython notebook - feel free to re-use
from IPython.core.display import HTML
import urllib2

HTML(urllib2.urlopen('http://bit.ly/1Bf5Hft').read())