def render_online():
    """
    Renders a spinning square moving on a straight line.
    """
    movie_name = "square_online"
    save_dir = make_video_path(movie_name+".pngvin")
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    
    print "Rendering pngvin movie", movie_name
    
    delete_database(movie_name)
    create_database(movie_name, [2])
    db = StateTransitionDatabase(movie_name)
    
    surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, IMAGE_WIDTH, IMAGE_HEIGHT)
    ctx = cairo.Context(surface)
    
    num_frames=32
    
    x0 = 0.2
    y0 = 0.2
    theta0 = 0.0
    
    dx=0.02
    dy=0.02
    dtheta=2.0*math.pi/num_frames
    width=0.1
    
    prev_state = numpy.array([IMAGE_WIDTH*x0, IMAGE_WIDTH*y0, theta0])

    for i in xrange(num_frames):
        clear(ctx)

        ctx.set_source_rgb(0,0,0)
        ctx.rectangle(0, 0, 1, 1)
        ctx.fill()

        ctx.set_source_rgb(1, 1, 1)
        ctx.translate(x0+dx*i, y0+dy*i)
        ctx.rotate(theta0+dtheta*i)
        ctx.rectangle(-width/2, -width/2, width, width)
        ctx.fill()

        # Render Obstacle
        clear(ctx)

        surface.write_to_png(make_video_path(save_dir, "frame-"+left_align_videoformat(i)+".png"))

        if i>0:
            # Input the transition into the database
            new_state = numpy.array([IMAGE_WIDTH*(x0+dx*i), IMAGE_WIDTH*(y0+dy*i), theta0+dtheta*i])
            db.add_transitions(prev_state[:2], new_state[:2])
            prev_state = new_state
        
        print "Rendered frame", i
    
    print "Done."
    def __init__(self, dataset, number_of_objects=1, number_of_transitions=1000, number_of_frames=64, number_of_movies=4, debug=False, dt=1, **other_kwargs):
        if isinstance(debug, str):
            debug = (debug == "True")

        delete_database(dataset)
        create_database(dataset, self.PARAMETER_GROUPS)
        self.dataset = dataset
        self.db = StateTransitionDatabase(dataset)
        self.number_of_objects = int(number_of_objects)
        self.number_of_transitions = int(number_of_transitions)
        self.number_of_frames = int(number_of_frames)
        self.number_of_movies = int(number_of_movies)
        self.dt = float(dt)
def render_bounce(movie_id=0, start_state=None, gravity=numpy.array([0, 1]), bounce_factor=0.75, input_to_database=True):
    """
    Renders a falling and bouncing square.
    """
    
    dataset = "square_bounce"
    num_frames = 128
    square_side = 50.
    half_square_side = square_side/2.0
    
    dt = 0.5
    
    X_LIMITS = [half_square_side, IMAGE_WIDTH-half_square_side]
    Y_LIMITS = [half_square_side, IMAGE_HEIGHT-half_square_side]
    
    # State: [x, y, x', y']
    if start_state is None:
        delete_database(dataset)
        create_database(dataset, [2, 2])
        render_bounce(0, numpy.array([IMAGE_WIDTH/2., 0., Y_LIMITS[0], 0.]))
        render_bounce(1, numpy.array([IMAGE_WIDTH/2., 2., Y_LIMITS[0], 0.]))
        render_bounce(2, numpy.array([IMAGE_WIDTH*2./3, -1, IMAGE_HEIGHT/2., -0.4]))
        render_bounce(3, numpy.array([IMAGE_WIDTH/4., 0., Y_LIMITS[1], 0.]))
        render_bounce(4, numpy.array([IMAGE_WIDTH*3./4, 0., Y_LIMITS[1], 2]))
        render_bounce(5, numpy.array([IMAGE_WIDTH/2., 4, Y_LIMITS[1], 0.6]))
        render_bounce(6, numpy.array([IMAGE_WIDTH/2., 4., IMAGE_HEIGHT/2., 0.]))
        render_bounce(7, numpy.array([X_LIMITS[0], 5., Y_LIMITS[0], 0.]), input_to_database=False)
        render_bounce(8, numpy.array([IMAGE_WIDTH/2., -1, IMAGE_HEIGHT*3./4, 0.6]), input_to_database=False)
        render_bounce(9, numpy.array([IMAGE_WIDTH/5., 2, IMAGE_HEIGHT/3., 0.4]), input_to_database=False)
        return
    
    movie_name = dataset + "_" + str(movie_id)
    save_dir = make_video_path(movie_name + ".pngvin")
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    
    print "Rendering pngvin movie", movie_name
    
    db = StateTransitionDatabase(dataset)
    
    surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, IMAGE_WIDTH, IMAGE_HEIGHT)
    ctx = cairo.Context(surface)
    
    state = start_state
#    state[:2] += half_square_side
    
    print "Initial values:", state
    
    def timestep(state, fixpoint_iterations=10):
        next_state = state
        d0 = numpy.array((state[1], gravity[0], state[3], gravity[1]))
        for i in xrange(fixpoint_iterations):
            d = numpy.array((next_state[1], gravity[0], next_state[3], gravity[1]))
            next_state = state + 0.5 * (d0 + d) * dt
        return next_state
    
    for i in xrange(num_frames):
        
        if i>0:
            next_state = timestep(state)

            # Collisions with edges
            for axis, limits in ((0, X_LIMITS), (2, Y_LIMITS)):
                if next_state[axis] < limits[0]:
                    next_state[axis+1] *= -bounce_factor
                    next_state[axis] = 2*limits[0] - next_state[axis]
                elif next_state[axis] > limits[1]:
                    next_state[axis+1] *= -bounce_factor
                    next_state[axis] = 2*limits[1] - next_state[axis]
            
            if input_to_database:
                # Input the transition into the database
                normalizer = numpy.array((state[0], 0, state[2], 0))
                db.add_transitions(state - normalizer, next_state - normalizer)
        
            state = next_state
        
        x = (state[0]-half_square_side)/IMAGE_WIDTH
        y = (state[2]-half_square_side)/IMAGE_HEIGHT
        
        clear(ctx)
        ctx.rectangle(0, 0, 1, 1)
        ctx.set_source_rgb(0,0,0)
        ctx.fill()
        
        ctx.set_source_rgb(1, 1, 1)
        ctx.rectangle(x, y, square_side/IMAGE_WIDTH, square_side/IMAGE_WIDTH)
        ctx.fill()
        
        surface.write_to_png(make_video_path(save_dir, "frame-"+left_align_videoformat(i)+".png"))

    print "Completed rendering", movie_name
def generate_pendulum(number_of_transitions, l, g=9.81, dt=0.001, radius=24):
    """
    Generates number_of_transitions random movements of a pendulum with start
    angle theta_0 from the y axis
    """
    
    amplitude = 40
    dataset = "pendulum"
    
    # State: [x, y, x', y']
    delete_database(dataset)
    create_database(dataset, [1])
    
    print "Generating %i random accelerating movements"%(number_of_transitions)
    db = StateTransitionDatabase(dataset)
    
    omega = math.sqrt(float(g)/l)
    T = 2*math.pi*1.0/omega
    
    theta_0 = numpy.random.uniform(math.radians(-amplitude), math.radians(amplitude), (number_of_transitions, 1))
    t = numpy.random.uniform(0, T, (number_of_transitions, 1))
    
    from_phi = theta_0 * numpy.cos(omega*t)
    to_phi = theta_0 * numpy.cos(omega*(t+dt))
    
    db.add_transitions(from_phi, to_phi)
    
    print "Done."
    
    print "Generating movement sequences for testing..."
    num_movies = 4
    num_frames = 64
    
    for movie_id in xrange(num_movies):
        movie_name = dataset + "_" + str(movie_id)
        save_dir = make_video_path(movie_name + ".pngvin")
        import shutil
        if os.path.exists(save_dir):
            shutil.rmtree(save_dir)
        os.makedirs(save_dir)
        
        print "Rendering pngvin movie to %s"%(save_dir)
        
        surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, IMAGE_WIDTH, IMAGE_HEIGHT)
        ctx = cairo.Context(surface)
        
        theta_0 = numpy.random.uniform(math.radians(-amplitude),math.radians(amplitude))
        t = numpy.atleast_2d(numpy.linspace(0, 3*T, num_frames)).T
        
        phi = theta_0 * numpy.cos(omega*dt*t)
        
        for i in xrange(num_frames):
            x = (IMAGE_WIDTH/2 + l*numpy.sin(phi[i]))/IMAGE_WIDTH
            y = l*numpy.cos(phi[i])/IMAGE_HEIGHT
            
            clear(ctx)
            ctx.rectangle(0, 0, 1, 1)
            ctx.set_source_rgb(0,0,0)
            ctx.fill()
            
            ctx.set_source_rgb(1, 1, 1)
#            ctx.rectangle(x, y, 10, 10)
            ctx.arc(x, y, float(radius)/IMAGE_WIDTH, 0., 2 * math.pi)
            ctx.fill()
            
            surface.write_to_png(os.path.join(save_dir, "frame-"+left_align_videoformat(i)+".png"))
        
        numpy.save(os.path.join(save_dir, "state_sequence.npy"), phi)
    
    print "Done."
class Generator:
    def __init__(self, dataset, number_of_objects=1, number_of_transitions=1000, number_of_frames=64, number_of_movies=4, debug=False, dt=1, **other_kwargs):
        if isinstance(debug, str):
            debug = (debug == "True")

        delete_database(dataset)
        create_database(dataset, self.PARAMETER_GROUPS)
        self.dataset = dataset
        self.db = StateTransitionDatabase(dataset)
        self.number_of_objects = int(number_of_objects)
        self.number_of_transitions = int(number_of_transitions)
        self.number_of_frames = int(number_of_frames)
        self.number_of_movies = int(number_of_movies)
        self.dt = float(dt)

    def generate(self):
        print
        print "Launching generator %s"%(self.__class__.__name__)
        print
        
        print "Generating training database with %i transitions..."%(self.number_of_transitions)
        self.db.add_transitions(*self.generate_training_transitions())
        print "Training database generated."
        print

        print "Generating testing movies with %i objects..."%(self.number_of_objects)
        print

        for movie_i in xrange(self.number_of_movies):
            print "Generating movie %i of %i..."%(movie_i+1, self.number_of_movies)
            save_dir = make_video_path("%s_%i.pngvin"%(self.dataset, movie_i))
            
            import shutil
            if os.path.exists(save_dir):
                print "Directory already exists:", save_dir
                shutil.rmtree(save_dir)
                print "Deleted directory:", save_dir
            os.makedirs(save_dir)
            
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            
            wlm = wlayermanager()
            
            print "Calculating state sequences..."
            states = [self.generate_testing_sequence(self.number_of_frames) for i in xrange(self.number_of_objects)]
            map(numpy.save, (os.path.join(save_dir, "state_sequence_%i.npy"%i) for i in xrange(len(states))), states)
            map(pickle.dump, map(self.generate_metadata, xrange(len(states))), (open(os.path.join(save_dir, "metadata_%i.pickle"%i), "w") for i in xrange(len(states))))
            print "State sequences calculated."
            
            print "Generating movie..."
            wlm.add_layer(self.generate_testing_movie(states))
            print "Movie generated."
            
            print "Saving movie to %s"%(save_dir)
            wlm.exportPNGVIN(save_dir, width=IMAGE_WIDTH, height=IMAGE_HEIGHT)
            print "Movie saved."
            
            print "Finished generating movie %i of %i"%(movie_i+1, self.number_of_movies)
            print

        print "Finished generating testing movies."
    
    def generate_training_transitions(self):
        """Implementations should return a length 2 sequence of equal size
        numpy arrays.
        """
        raise NotImplementedError("This class is abstract!")
    
    def generate_testing_sequence(self, num_frames):
        """Implementations should return a numpy array representing a sequence
        of states to be tracked.
        """
        raise NotImplementedError("This class is abstract!")
    
    def generate_testing_movie(self, states):
        """Implementations should return a wvideo.
        @param states: A sequence of states generated by
                        generate_testing_sequence
        @see generate_testing_sequence
        """
        raise NotImplementedError("This class is abstract!")

    def generate_metadata(self, obj_i):
        """Implementations should return a dictionary.
        @param obj_i: Which object to generate metadata for.
        """
        raise NotImplementedError("This class is abstract!")