# Alec Acosta
# Blender Particle Baking
# 
# V 1.1
# Updates:
#  1.1 -
#   Added option for visibility keyframing
#   Fixed issue with particles that start later having the wrong starting position
#   Removed some redundant code
#
# How to use:
#  1. Create particle system static or animated wtih physics
#  2. Create an instance geometry such as a cube or vertice if you want to attach nulls
#  3. Select particle system then select instance geometry
#  4. Run script and look for a new set of objects named "particle.001" etc.
#
#
import bpy

# Key inputs. Change occordingly.
KEYFRAME_LOCATION = True
KEYFRAME_ROTATION = True
KEYFRAME_SCALE = True
KEYFRAME_VISIBILITY = False  # Viewport and render visibility.

def create_objects_for_particles(ps, obj):
    # Duplicate the given object for every particle and return the duplicates.
    # Use instances instead of full copies.
    obj_list = []
    mesh = obj.data
    for i, _ in enumerate(ps.particles):
        dupli = bpy.data.objects.new(
                    # Still hard coded change 3 to number of trailing numbers in particle names
                    name="particle.{:03d}".format(i),
                    object_data=mesh)
        # 2.80: link to objects in collection
        bpy.context.scene.collection.objects.link(dupli)
        obj_list.append(dupli)
    return obj_list

def match_and_keyframe_objects(ps, obj_list, start_frame, end_frame):
    # Match and keyframe the objects to the particles for every frame in the given range.
    for frame in range(start_frame, end_frame + 1):
        bpy.context.scene.frame_set(frame)
        for p, obj in zip(ps.particles, obj_list):
            match_object_to_particle(p, obj)
            keyframe_obj(obj)

def match_object_to_particle(p, obj):
    # Match the location, rotation, scale and visibility of the object to the particle.
    loc = p.location
    rot = p.rotation
    size = p.size
    if p.alive_state == 'ALIVE':
        vis = True
    else:
        vis = False
    obj.location = loc
    # Set rotation mode to quaternion to match particle rotation.
    obj.rotation_mode = 'QUATERNION'
    obj.rotation_quaternion = rot
    obj.scale = (size, size, size)
    # 2.80: hide_viewport instead of hide
    obj.hide_viewport = not(vis)
    obj.hide_render = not(vis)

def keyframe_obj(obj):
    # Keyframe location, rotation, scale and visibility if specified.
    if KEYFRAME_LOCATION:
        obj.keyframe_insert("location")
    if KEYFRAME_ROTATION:
        obj.keyframe_insert("rotation_quaternion")
    if KEYFRAME_SCALE:
        obj.keyframe_insert("scale")
    if KEYFRAME_VISIBILITY:
        # 2.80: hide_viewport instead of hide
        obj.keyframe_insert("hide_viewport")
        obj.keyframe_insert("hide_render")

def main():
    # Assume only 2 objects are selected.
    # The active object should be the one with the particle system.
    ps_obj = bpy.context.object
    obj = [obj for obj in bpy.context.selected_objects if obj != ps_obj][0]
    # 2.80: Get the dependency graph
    dg = bpy.context.evaluated_depsgraph_get()
    # 2.80: Get an evaluated object from it
    eval_obj = ps_obj.evaluated_get(dg)
    ps = eval_obj.particle_systems[0]  # Assume only 1 particle system is present.
    start_frame = bpy.context.scene.frame_start
    end_frame = bpy.context.scene.frame_end
    obj_list = create_objects_for_particles(ps, obj)
    match_and_keyframe_objects(ps, obj_list, start_frame, end_frame)

if __name__ == '__main__':
    main()