__all__ = ["Generator"]

from common.settings import make_video_path
from wdb import create_database, delete_database, StateTransitionDatabase
from wgenerator.settings import IMAGE_WIDTH, IMAGE_HEIGHT
from wgui.wlayermanager import wlayermanager
import numpy
import os
import pickle

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!")
