#!/bin/bash
"exec" "blender" "--python" "$0"

# PHYSICS DEMO - NEWTON'S CRADLE
#
# This Blender program simulates a Newton's cradle with 5 balls.
#

import bpy
import math
import mathutils

scene = bpy.context.scene

# REMOVE ALL EXISTING OBJECTS
#
# i.e, the cube, the lamp, the camera if we're running this script for
# the first time, or a previous copy of this model if we've modified
# this script and are re-running it, or anything else that happens to
# be loaded into Blender

# This causes bpy.ops.rigidbody.object_add() to fail with incorrect context
#bpy.ops.wm.read_homefile()

bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete()

# Remove any old handlers.  This prevents the crash triggered in
# 5a5d50101cb88d20c10ff032b57ee45057ce54ea, which was apparently due
# to an old frame_change_pre handler accessing deleted objects.

del bpy.app.handlers.frame_change_pre[:]
del bpy.app.handlers.frame_change_post[:]
del bpy.app.handlers.load_pre[:]
del bpy.app.handlers.load_post[:]
del bpy.app.handlers.render_pre[:]
del bpy.app.handlers.render_post[:]
del bpy.app.handlers.save_pre[:]
del bpy.app.handlers.save_post[:]
del bpy.app.handlers.scene_update_pre[:]
del bpy.app.handlers.scene_update_post[:]

# RENDER SETTINGS

#scene.render.image_settings.file_format = 'FFMPEG'
scene.render.resolution_x = 166
scene.render.resolution_y = 132
scene.render.resolution_percentage = 100
#scene.render.engine = 'CYCLES'

# MATERIALS
#
# For some reason, the import picks up an object and not just a
# texture, so we delete it, as the scene should be empty at this
# point.  Also, if the link_append doesn't work, it will print a
# warning message but won't raise an exception.  That doesn't happen
# until we try to access the material.
#
# Trying to import a second material without first deleting the object
# from the first one crashes Blender 2.69

bpy.ops.wm.link_append(directory="../Blender Texture Disc/material/caststeel.blend/Material/",
                       filename="caststeel", autoselect=True)
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete()

try:
    caststeel = bpy.data.materials['caststeel']
except KeyError:
    caststeel = None

bpy.ops.wm.link_append(directory="../Blender Texture Disc/material/walnut.blend/Material/",
                       filename="walnut", autoselect=True)
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete()

try:
    walnut = bpy.data.materials['walnut']
except KeyError:
    walnut = None

# THE CAMERA

bpy.ops.object.camera_add()
camera = bpy.context.active_object

camera.location = (15.3416, -4.5762, 12.8521)
camera.rotation_euler = (1.1087, 0.0133, 1.1483)

# THE LIGHTING

bpy.ops.object.lamp_add(type='SUN')
bpy.context.active_object.location = (10,10,25)

bpy.data.worlds['World'].zenith_color = mathutils.Color((0.9, 0.9, 0.8))
bpy.data.worlds['World'].light_settings.use_environment_light = True
bpy.data.worlds['World'].light_settings.environment_color = 'SKY_COLOR'

# THE 1x1x5 TOP THAT THE BALLS HANG FROM

bpy.ops.mesh.primitive_cube_add()
bpy.ops.rigidbody.object_add()
top = bpy.context.active_object
top.scale = (0.5,2.5,0.5)
top.location = (0,2,5)
top.active_material = walnut
top.rigid_body.type = 'PASSIVE'

# THE BALLS AND RODS

balls = []
rods = []

# The rigid body simulation works a lot better if the balls aren't
# actually touching, which produces a series of two body collisions,
# instead of a more complicated five body collision, which would
# require considering the elasticity of the balls to achieve an
# accurate result.
#
# Balls are centered at integer coordinates along the y axis.  In an
# earlier version of this program, the balls had a fixed radius of 0.5
# (i.e, they were just touching).  Now I allow the radius to adjusted
# by the user, with the point of contact between the balls and the
# rods still at (0,y,0.5)

ball_radius = 0.4

for y in range(5):

    bpy.ops.mesh.primitive_uv_sphere_add()
    bpy.ops.rigidbody.object_add()
    ball = bpy.context.active_object
    ball.location = (0,y,0.5 - ball_radius)
    ball.scale = (ball_radius, ball_radius, ball_radius)
    ball.active_material = caststeel
    balls.append(ball)

    # Stock cylinders have radius 1, length 2, and are aligned along
    # the z-axis.  Ours are z-scaled by 2, so they have length 4 and
    # connect the tops of the balls (at z=0.5) to the bottom of the
    # top (at z=4.5)

    bpy.ops.mesh.primitive_cylinder_add()
    bpy.ops.rigidbody.object_add()
    rod = bpy.context.active_object
    rod.scale = (.05,.05,2)
    rod.location = (0,y,2.5)
    rod.active_material = caststeel
    rods.append(rod)

    # The newly created rod is already the active object, so just
    # select the ball and parent it to the rod.

    ball.select = True
    bpy.ops.object.parent_set()

    # Now I mimic bpy.ops.rigidbody.connect(), which is a script in
    # bl_operators/rigidbody.py that first creates an empty object...
    #
    # The location of the hinge constraint is where the rod meets the
    # top, and we rotate 90 degrees about the y-axis to put the
    # constraint's local z-axis (its axis of motion) along our x-axis

    ob = bpy.data.objects.new("Constraint", object_data=None)
    ob.location = rod.location
    ob.location[2] += 2
    ob.rotation_euler = (0, math.radians(90), 0)
    scene.objects.link(ob)
    scene.objects.active = ob
    ob.select = True

    # ...then creates a rigid body constraint...

    bpy.ops.rigidbody.constraint_add()
    con_obj = bpy.context.active_object
    con_obj.empty_draw_type = 'ARROWS'
    con = con_obj.rigid_body_constraint
    con.type = 'HINGE'

    # ...and assigns the two connected objects to it.

    con.object1 = top
    con.object2 = rod

    # Now add a fixed constraint where the rod meets the ball

    ob = bpy.data.objects.new("Constraint", object_data=None)
    ob.location = rod.location
    ob.location[2] -= 2
    scene.objects.link(ob)
    scene.objects.active = ob
    ob.select = True

    bpy.ops.rigidbody.constraint_add()
    con_obj = bpy.context.active_object
    con_obj.empty_draw_type = 'SPHERE'
    con_obj.empty_draw_size = 0.1
    con = con_obj.rigid_body_constraint
    con.type = 'FIXED'

    con.object1 = rod
    con.object2 = ball

    # Parent the constraint to the ball.
    #
    # Leaving the constraint unparented seems to affect only the GUI,
    # not the simulation.

    bpy.ops.object.select_all(action='DESELECT')
    scene.objects.active = ball
    ob.select = True
    bpy.ops.object.parent_set()

    # Set physics parameters - restitution is called "bounciness" in GUI
    ball.rigid_body.friction = 0
    ball.rigid_body.restitution = 1
    ball.rigid_body.angular_damping = 0
    ball.rigid_body.linear_damping = 0

    rod.rigid_body.friction = 0
    rod.rigid_body.restitution = 1
    rod.rigid_body.angular_damping = 0
    rod.rigid_body.linear_damping = 0

    # Setting the rods to have very small mass produces bizarre effects!
    #rod.rigid_body.mass = 0.001

# Deselect everything

bpy.ops.object.select_all(action='DESELECT')
scene.objects.active = None

# At 24 fps, the default step rate of 60 requires interpolation on
# alternate frames, degrading accuracy noticeably in the kinetic
# energy calculation (try it).  Changing to 72 steps per second gives
# us exactly 3 steps per frame and nice kinetic energy values.

scene.rigidbody_world.steps_per_second = 72

# Store our objects in a convenient place for Python Console

bpy.balls = balls
bpy.rods = rods

# rotate_object(obj, angle, direction, point)
#
# obj - a Blender object
# angle - angle in radians
# direction - axis to rotate around (a vector from the origin)
# point - point to rotate around (a vector)

def rotate_object(obj, angle, direction, point):
    R = mathutils.Matrix.Rotation(angle, 4, direction)
    T = mathutils.Matrix.Translation(point)
    M = T * R * T.inverted()
    obj.location = M * obj.location
    obj.rotation_euler.rotate(M)

# updateInitialPositions - adjust the initial angle of the first ball
# and the radii of the balls, then update the scene to recompute the
# world matrices so we can stash away the initial locations of the
# balls.

def extract_location(matrix_world):
    return (matrix_world[0][3], matrix_world[1][3], matrix_world[2][3])

def updateInitialPositions(self, context):
    rods[0].location = (0,0,2.5)
    rods[0].rotation_euler = (0,0,0)
    rotate_object(rods[0], math.radians(-scene.InitialAngle), (1,0,0), (0,0,4.5))

    ball_radius = scene.BallRadius
    for i in range(5):
        balls[i].location = (0,i,0.5-ball_radius)
        balls[i].scale = (ball_radius, ball_radius, ball_radius)

    scene.update()

    for i in range(5):
        balls[i].data["InitialLocation"] = extract_location(balls[i].matrix_world)
        rods[i].data["InitialLocation"] = extract_location(rods[i].matrix_world)

    return None

# CUSTOM PROPERTIES
#
# I can't create a property on a single object with a callback, to do
# that I need to assign it to a class.  Might as well use Scene.

# InitialAngle - rotation of the first ball
# BallRadius - size of the balls (limited to 0.5 so they can't overlap)

bpy.types.Scene.InitialAngle = bpy.props.IntProperty(default=90,
                                                     update=updateInitialPositions)

bpy.types.Scene.BallRadius = bpy.props.FloatProperty(default=ball_radius,
                                                     min=0.1, max=0.5,
                                                     update=updateInitialPositions)

updateInitialPositions(None, None)

# SimulationLength, LoopSimulation - adjust length and behavior of simulation
#
# Need to set both animation length and physics cache length.

def updateSimulationLengthFunc(self, context):
    scene.frame_end = scene.SimulationLength
    scene.rigidbody_world.point_cache.frame_end = scene.SimulationLength
    return None

bpy.types.Scene.SimulationLength = bpy.props.IntProperty(default=250, update=updateSimulationLengthFunc)
bpy.types.Scene.LoopSimulation = bpy.props.BoolProperty(default=True)

# Calculate and print kinetic and potential energies to a text view

text = bpy.data.texts.new(name='Energy')

def calculate_and_print_energies():

    # Calculate kinetic energy - sum of 1/2 m v^2
    # Pythagorean theorem: delta^2 = (delta x)^2 + (delta y)^2 + (delta z)^2
    # speed = delta/time
    # mass is 1; time is 1/fps of a second
    # KE = 1/2 * delta^2 * fps^2
    #
    # omega = (delta radians)/time = fps * (delta radians)
    # Moment of inertia of rod rotating around end: I = m L^2 / 3 = 16/3 (m is 1; L is 4)
    # KE = 1/2 I omega^2 = 8/3 omega^2

    kinetic_energy = 0
    potential_energy = 0

    for ball in balls:
        delta_squared = 0
        for i in range(3):
            delta_squared += (ball.matrix_world[i][3] - ball.data["PreviousLocation"][i])**2
            ball.data["PreviousLocation"][i] = ball.matrix_world[i][3]

        kinetic_energy += delta_squared * scene.render.fps**2 / 2
        potential_energy += 9.81 * ball.matrix_world[2][3]

    for rod in rods:
        delta_squared = 0
        for i in range(3):

            delta_squared += (rod.matrix_world[i][3] - rod.data["PreviousLocation"][i])**2
            rod.data["PreviousLocation"][i] = rod.matrix_world[i][3]

        kinetic_energy += delta_squared * scene.render.fps**2 / 2

        rotation = rod.matrix_world.to_euler().x
        omega = scene.render.fps * (rotation - rod.data["PreviousRotation"])
        kinetic_energy += (8/3) * omega**2
        rod.data["PreviousRotation"] = rotation

        potential_energy += 9.81 * rod.matrix_world[2][3]

    text.write("Frame " + str(scene.frame_current) + " KE " + str(kinetic_energy) + " PE " + str(potential_energy) + " E " + str(kinetic_energy + potential_energy) + "\n")

# Handler to stop simulation at end (if requested) and print energy
# calculation at each frame.

def frame_change_handler(scene):

    # At start of simulation, use initial location of balls as previous location

    if scene.frame_current == 2:
        for i in range(5):
            balls[i].data["PreviousLocation"] = balls[i].data["InitialLocation"]
            rods[i].data["PreviousLocation"] = rods[i].data["InitialLocation"]
            rods[i].data["PreviousRotation"] = rods[i].rotation_euler[0]

    if scene.frame_current != 1:
        calculate_and_print_energies()

    if scene.frame_current == scene.SimulationLength and not scene.LoopSimulation:
        bpy.ops.screen.animation_cancel(restore_frame=False)

bpy.app.handlers.frame_change_pre.append(frame_change_handler)

# Add a Panel to 3D View Properties with simulation controls

class VIEW3D_Newtons_Cradle(bpy.types.Panel):
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'TOOL_PROPS'
    bl_label = "Newton's Cradle"

    def draw(self, context):
        layout = self.layout

        row = layout.row()
        row.prop(scene, 'InitialAngle', text="Initial Angle")

        row = layout.row()
        row.prop(scene, 'BallRadius', text="Ball Radius")

        row = layout.row()
        row.prop(scene, 'SimulationLength', text="Simulation Length")

        row = layout.row()
        row.prop(scene, 'LoopSimulation', text="Loop Simulation")

bpy.utils.register_class(VIEW3D_Newtons_Cradle)

# Local Variables:
# mode: python
# End: