import bpy
import os
import mathutils
import sys

# Take a series of screenshots of the UI where the 3D view comes from an auxiliary screenshot camera,
# but make it look like the UI is looping the 3D animation.
# This makes things a little confusing since we want 400 frames,
# but the animation is only 20 frames long.  Gymnastics ensue.


def viewMatrixFromCamera(cam):
    """
    :param cam: the camera whose point of view is of interest
    :return: the view matrix suitable for use in region_3d.view_matrix
    """
    return cam.matrix_world.inverted()

def viewFrom(screen):
    """
    :param screen:
    :return: the "first" area in the specified screen that is a 3D view.
    """
    for obj in screen.areas:
        if (obj.type == 'VIEW_3D'):
            return obj
    return None

def filenameFor(idx, scene=bpy.context.scene):
    return os.path.join(scene.render.filepath,
        "ss%04d.png"%idx)
        
        
def exp1(region_3d):
    return region_3d.perspective_matrix * region_3d.view_matrix

def cameraFor(idx):
    """
    :param idx:
    :return: which camera's viewpoint will we use for the 3D view when it's time to take screenshot number idx.
    """
    if (True or idx<=200):
        return bpy.data.objects['Camera.002']
    # I was using two cameras for two halves, but then I changed my mind and used the same camera for both.
    return bpy.data.objects['Camera.003']

def screenShotAt(idx):
    """
    take a screenshot of the current window state and store it in the file for idx
    """

    fn = filenameFor(idx)

    if not os.path.isfile(fn):
        print ("snap")
        bpy.ops.screen.screenshot(filepath=fn)
    else:
        print("%s already exists"%fn)

def update3dView(idx, v3d, scn=bpy.context.scene):
    """
    Adjust the window for screenshot number idx:
    this includes the view matrix in the 3D view (v3d) as well as shading mode and active object.
    """

    cam = cameraFor(idx)
    
    scn.frame_current = idx
    scn.update()
    
    v3d.spaces[0].region_3d.view_matrix = viewMatrixFromCamera(cam)
    
    if (idx > 200):
        v3d.spaces[0].viewport_shade = 'TEXTURED'
        wantAO = bpy.data.objects["hat"]
    else:
        v3d.spaces[0].viewport_shade = 'WIREFRAME'
        wantAO = bpy.data.objects["kiwi Armature"]
        
    scn.objects.active = wantAO
    for obj in scn.objects:
        obj.select = (obj==wantAO)


scn = bpy.context.scene
screen = bpy.context.screen

# we stash our screenshot index as a custom property on this camera object
cam = bpy.data.objects["Camera.002"]
idx = cam.data["cursor"]

i2 = (idx - scn.frame_start) % (scn.frame_end - scn.frame_start +1) + scn.frame_start

if (scn.frame_current == idx and scn.frame_current != i2):
    # we can get the view matrix from the axuiliary camera,
    # but can't take the screenshot yet because
    # we have to set the frame_current to something inside the animation time bounds.
    print ("adjust camera")
    v3d = viewFrom(screen)
    update3dView(idx, v3d)

    scn.frame_current = i2
else:
    print ("screenshot %d"%idx)
    screenShotAt(idx)

    end = scn.frame_end
    end = 400

    for q in range(end):
        idx = idx+1

        if (idx > end):
            idx = 1
        if not (os.path.isfile(filenameFor(idx))):
            break
        # skip screenshots we have already captured

    cam.data["cursor"] = idx

    # let's adjust the 3D view to match the auxiliary camera
    v3d = viewFrom(screen)

    update3dView(idx, v3d)