import bpy
import sys
import numpy as np
from mathutils import Vector
DEG_TO_RAD = np.pi / 180.


def _to_rad(data):
    return (DEG_TO_RAD * data[0], DEG_TO_RAD * data[1], DEG_TO_RAD * data[2])


def setup_scene():
    for block in bpy.data.meshes:
        if block.users == 0:
            bpy.data.meshes.remove(block)
    bpy.ops.object.mode_set(mode='OBJECT')
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete()
    bpy.context.scene.render.resolution_x = 720
    bpy.context.scene.render.resolution_y = 480
    bpy.context.scene.frame_start = 0
    bpy.context.scene.render.image_settings.color_mode = 'RGB'
    bpy.context.scene.render.image_settings.file_format = 'PNG'


def add_camera_light():
    bpy.ops.object.camera_add(
        location=(0.0, -3.5, 2.2),
        rotation=_to_rad((60.8, 0.0, 0.0)),
    )
    bpy.context.scene.camera = bpy.context.object
    bpy.ops.object.light_add(
        type='SUN',
        radius=1,
        location=(33.2, -23.6, 17.4),
        rotation=_to_rad((208, 140, 181)),
    )
    bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[0].default_value = (1, 1, 1, 1)


def make_plane():
    """Underlying plane. Use collision so that it can support cloth on top."""
    bpy.ops.mesh.primitive_plane_add(location=(0, 0, 0))
    bpy.ops.transform.resize(value=(6.0, 6.0, 6.0))
    bpy.ops.object.modifier_add(type='COLLISION')


def make_cloth():
    """Make a plane at height 0.2m (it then settles on the plane) and sub-divide."""
    bpy.ops.mesh.primitive_plane_add(location=(0, 0, 0.2))
    bpy.ops.object.editmode_toggle()
    bpy.ops.mesh.subdivide(number_cuts=20)
    bpy.ops.object.editmode_toggle()
    bpy.ops.object.modifier_add(type='CLOTH')

    # Colors
    mat = bpy.data.materials.new(name="CustomColor")
    mat.use_nodes = False
    mat.diffuse_color = (0,1,0,0)
    bpy.context.object.data.materials.append(mat)

    return bpy.context.object


def reset(cloth):
    """
    Remove vertex groups. The second for loop over actions removes keyframes,
    which seems to avoid keyframes spilling over from prior actions.
    """
    for key in cloth.vertex_groups.keys():
        cloth.vertex_groups.remove(cloth.vertex_groups[key])
    cloth.modifiers["Cloth"].settings.vertex_group_mass = ''
    for a in bpy.data.actions:
        bpy.data.actions.remove(a)


def define_action(cloth, vertex_idxs, frames_per_act):
    """Define an action: pick up cloth at these vertex indices. Uses keyframing.

    We add an armature and make it the parent of a pinned vertex group. Parent set:
    https://docs.blender.org/api/current/bpy.ops.object.html#bpy.ops.object.parent_set
    """
    if len(vertex_idxs) == 0:
        print('A no-op action, just letting the scene settle.')
        return
    bpy.ops.object.armature_add(enter_editmode=False, location=(0, 0, 0))

    # Select the cloth _in_addition_ to the armature. The active one is the parent.
    arm = bpy.context.object
    cloth.select_set(True)
    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.parent_set(type='ARMATURE_NAME')

    # Set cloth as the active object, and pin the designated vertices.
    bpy.ops.object.select_all(action='DESELECT')
    bpy.context.view_layer.objects.active = cloth
    gripped_group = bpy.context.object.vertex_groups.new(name='Pinned')
    gripped_group.add(vertex_idxs, 1.0, 'ADD')
    cloth.modifiers["Cloth"].settings.vertex_group_mass = 'Pinned'

    # The arm modifier must be at the TOP of the modifier stack (over cloth)!
    for _ in range(5):
        bpy.ops.object.modifier_move_up(modifier="Armature")

    # Also need pinned set for the armature. The arm's default name is 'Bone'.
    b_group = bpy.context.object.vertex_groups["Bone"]
    b_group.add(vertex_idxs, 1.0, 'ADD')

    # Now animate by assigning the arm to different positions at certain keyframes.
    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.posemode_toggle()
    start_frame = bpy.context.scene.frame_current
    frames = [start_frame + 1, start_frame + frames_per_act]
    print('\nkeyframing at: {}'.format(frames))

    # Keyframe inserts need to use `b.[...]` attributes, like `b.location`.
    for b in arm.pose.bones:
        b.keyframe_insert("location", frame=frames[0])
    c = arm.pose.bones["Bone"]

    # Create keyframe to move armatures upwards by 0.3m, and also moves cloth.
    c.location += Vector((0, 0.3, 0))
    for b in arm.pose.bones:
        b.keyframe_insert("location", frame=frames[1])
    bpy.ops.object.posemode_toggle()


def advance_simulator(t, frames_per_act):
    """Actually take steps in the physics simulator based on current keyframes.

    Whereas `define_action` method above defines the keyframes for the action,
    this explicitly take steps in the physics simulator by setting the frames.
    Images are saved AFTER we set the frame, with time and frame indices.
    """
    scene = bpy.context.scene
    start = scene.frame_current + 1  # we already set at this frame earlier
    print('Now stepping, going from {} to {}'.format(start, start + frames_per_act))
    filename = 'images/act_{}_img_%06d.png'.format(str(t).zfill(2))
    for f in range(start, start + frames_per_act):
        scene.frame_set(f)  # Advance the simulated physics!
        scene.render.filepath = filename % f
        bpy.ops.render.render(write_still=True)

        # ---- debugging position AFTER setting frame ---- #
        depsgraph = bpy.context.evaluated_depsgraph_get()
        cloth_upd = cloth.evaluated_get(depsgraph)
        v_c  = [cloth.matrix_world @ v.co for v in list(cloth.data.vertices)]
        v_cd = [cloth_upd.matrix_world @ v.co for v in list(cloth_upd.data.vertices)]
        print('after frame {}, vertex idx 0: {}, {}'.format(f, v_c[0], v_cd[0]))


if __name__ == '__main__':
    setup_scene()
    add_camera_light()
    make_plane()
    cloth = make_cloth()
    frames_per_act = 50
    bpy.context.scene.frame_set(0)

    # First action: no-op (lets the scene settle). Equivalent to no items in list.
    t = 0
    vertex_idxs = []
    reset(cloth)
    define_action(cloth, vertex_idxs, frames_per_act)
    advance_simulator(t, frames_per_act)

    # Second action: grasp at corner and pull upwards. TODO: physics issue??
    t = 1
    vertex_idxs = [0]
    reset(cloth)
    define_action(cloth, vertex_idxs, frames_per_act)
    advance_simulator(t, frames_per_act)