#!/usr/bin/python
# This file is part of the Coral project.
# Please read the copyright notice in the LICENSE file.
#
# -*- coding: UTF-8 -*-
#
# arbitrary_model_generator - generate models of arbitrary size/type/etc
#
# SYNOPSIS
#
# Explained in usage()
#
# DESCRIPTION
#
# Creates arbitrary models, meaning
# - which metamodels used
# - amount of files used
# - amount of root elements/file used
# - amount of model elements used
#
# This program has been used to stresstest Coral's model framework
# and I/O routines
#
# AUTHORS
#
# Marcus Alanen <maalanen@abo.fi> 2005
#

import random
import sys
import string

class Memoize:
    """Extremely simple memoizer class which caches previous function
    calls with same parameters forever."""
    def __init__(self, theobj, thename):
        self.values = {}
        func = getattr(theobj, thename)
        v = self.values

        def replacing_func(*args):
            if v.has_key(args):
                return v[args]
            else:
                ret = func(*args)
                v[args] = ret
                return ret
        setattr(theobj, thename, replacing_func)


def get_lower_and_upper(lower, upper):
    """Returns a tuple (l, u) where l = @lower and u is somewhere
    between @lower an @upper+1. If @upper == -1 this is taken to
    mean infinity, in which case we return a random number (greater
    than @lower)."""
    if upper == -1:
        upper = lower + random.randrange(1, lower*10 + 10)
    else:
        upper += 1

    return lower, random.randrange(lower, upper)

# This (inverted) list from http://www.w3.org/TR/REC-xml/#charsets for XML 1.0
# Although XML 1.1 would be a bit better, there are lot's of incompatible
# readers out there...
xml_non_printables =   "\000\001\002\003\004\005\006\007" \
                     + "\010\013\014\016\017\020\021\022" \
                     + "\023\024\025\026\027\030\031\032" \
                     + "\033\034\035\036\037"

class ArbitraryModelGenerator:
    """The arbitrary model generator."""
    def __init__(self):
        import coral
        self.XMI = coral.lang.XMI12
        self.CORE = coral.lang.CORE01
        
    def get_all_types(self, metamodels):
        """Returns a tuple with 1) a hash of all metaclasses from the
        list of @metamodels. This visits Coral Package objects transitively.
        2) a tuple of all Coral Packages used."""
        import coral
        def fill(mm, h, packages):
            for mc in m.__elements__():
                if mc.isDatatypeDefinition():
                    continue
                if mc.isAbstract():
                    continue
                if isinstance(mc, coral.Package):
                    if not packages.has_key(mc):
                        packages[mc] = True
                        fill(mc, h, packages)
                    continue
                h[mc] = True
            return h
    
        h = {}
        packages = {}
        for m in metamodels:
            packages[m] = True
            fill(m, h, packages)
    
        return h.keys(), tuple(packages.keys())

    def select_of_type(self, thelist, thetype):
        """Return a tuple of elements from @thelist which are an instance
        of @thetype. (or a child type of @thetype)."""
        return tuple([f for f in thelist if f.oclIsKindOf(thetype)])

    def choose_one_of_type(self, thelist, thetype):
        """Returns an element from @thelist which is an instance
        of @thetype. If no such element exists, None is returned."""
        f = self.select_of_type(thelist, thetype)
        if not f:
            return None
        return self.choose_one(f)
    
    def choose_one(self, thelist):
        """Returns a value from @thelist which must be nonempty."""
        e = thelist[random.randrange(0, len(thelist))]
        return e
    
    def collect_concrete_subclasses(self, mc):
        """Returns all transitive subclasses of @mc (including @mc)
        which are concrete (i.e. isAbstract() returns false)."""
        def helper(thetype, h):
            if not thetype.isAbstract():
                h[thetype] = True
            for i in thetype.__subclasses__():
                helper(i, h)
    
        ret = {}
        helper(mc, ret)
        return ret.keys()
    
    
    def choose_suitable_property(self, element, valid_packages):
        """Returns a list of (property, subclass) items for
        @element. The properties are Coral Property objects,
        which are valid for @element, and each subclass is a valid
        concrete target metaclass _and_ belong to the list of @valid_packages.

        @valid_packages are used to restrict the returned metaclasses
        to belong to certain metamodels only.

        E.g.

        For a UML 1.4 Model, we might return
        (ownedElement, Package)
        (ownedElement, Class)
        ...
        etc.

        Then you just pick one from the list, create a new element
        and put in the corresponding slot of @element

        """
        mc = element.__type__()
        candidates = []
        for prop in mc.__properties__():
            if prop.isUnserializable():
                continue
            if prop.isAnonymous():
                continue
            if not prop.isComposition():
                continue
            target_mc = prop.getTarget()
            if target_mc.isDatatypeDefinition():
                continue
            subs = self.collect_concrete_subclasses(target_mc)
            for s in subs:
                if s.__package__() in valid_packages:
                    candidates.append((prop, s))
        # Don't only return the XMI extension thingy,
        # otherwise we get too many of them...
        if len(candidates) == 1 and candidates[0][1] == self.XMI.Extension:
            return []
        
        return candidates

    def __set_a_primitive(self, atom, unique_strings_hash):
        """Sets @atom (a Coral Atom) to some random (valid) value.
        Supports Coral Strings, Integers, Doubles, Enumerations."""
        target = atom.__type__()

        if target == self.CORE.String or target == self.CORE.Name:
            while True:
                s = ""
                for i in xrange(0, random.randrange(0, 20)):
                    while True:
                        c = chr(random.randrange(0, 256))
                        # The '/' makes eclipse names go bonkers
                        if unique_strings_hash:
                            if c in string.ascii_letters:
                                break
                        else:
                            if c not in xml_non_printables:
                                break
                    s += c
                if not unique_strings_hash:
                    break
                if not unique_strings_hash.has_key(s):
                    unique_strings_hash[s] = True
                    break
            atom.__setString__(s)
        elif target == self.CORE.Integer:
            atom.__setIValue__(random.randrange(-sys.maxint, sys.maxint))
        elif target == self.CORE.Double:
            atom.__setDValue__(random.random())
        elif target.isEnumerationDefinition():
            atom.__setEValue__(random.randrange(0, target.size()))
        else:
            pass

    def __fill_element(self, e, unique_strings_hash):
        target_mc = e.__type__()
        #
        # Set the various primitive datatypes of the newly
        # created element
        #
        for p in target_mc.__properties__():
            target = p.getTarget()
            if not target.isDatatypeDefinition():
                continue
            # Unserializable properties won't show up in XMI
            # So don't change them
            if p.isUnserializable():
                continue
            if p.isAnonymous():
                continue
            assert p.isComposition()
            if not unique_strings_hash and random.random() < 0.5:
                continue
            if p.isMultSingle():
                self.__set_a_primitive(getattr(e, p.getName()), unique_strings_hash)
                continue
            
            lower, upper = get_lower_and_upper(p.getMultLower(),
                                               p.getMultUpper())

            slot = e.__getSlotByName__(p.getName())
            for i in xrange(lower, upper):
                atom = target.newAtom()
                self.__set_a_primitive(atom, unique_strings_hash)
                if p.isAtMostOne():
                    setattr(e, p.getName(), atom)
                else:
                    slot.append(atom)
    
        
    def create_models(self, outputdir, n_files, n_roots_per_file,
                      metamodel_strings, n_elements,
                      compare_after_saving, unique_strings):
        """The main method. Creates models totalling @n_elements
        amount of elements and saves them into @n_files. The files
        contain @n_roots_per_file root elements.  Valid metamodels are
        given as a list of strings in @metamodel_strings, such as
        ["UML14", "XMIDI20", ...]. The @outputdir is a prefix where to
        place the files.

        The created models are not only arbitrarily deep, but also
        arbitrarily connect to each other, even to elements in other
        files.

        Finally, the files are loaded back, compared with the original ones
        using coral.isSameModel.
        """
        
        import coral
        elements = []
        
        metamodels = []
        
        for m in metamodel_strings:
            L = getattr(coral.lang, m)
            metamodels.append(L)
        
        all_types, packages = self.get_all_types(metamodels)

        #
        # Only allow the roots of the various packages as
        # top-level elements. This makes it easier to create models as
        # well since e.g. we don't choose an element which can't
        # have any children. (We assume that the root elements can have
        # children...)
        #
        root_candidates = []
        for i in metamodels:
            # XMI.Extension elements are not allowed to be top-level in Coral
            r = i.getRoot()
            if r == self.XMI.Extension:
                continue
            root_candidates.append(r)
        
        unique_strings_hash = None
        if unique_strings:
            unique_strings_hash = {}
            unique_strings_hash[""] = True
    
        projects = []
    
        #
        # Create all the projects
        #
        for f in xrange(n_files):
            p = coral.Project()
            p.setFilename(outputdir + str(f) + ".xmi")
            projects.append(p) 
            for i in xrange(n_roots_per_file):
                e = self.choose_one(root_candidates).newAtom()
                e.__setProject__(p)
                elements.append(e)
                self.__fill_element(e, unique_strings_hash)

        #
        # Populate root elements with lot's of (transitive) children
        #
        for i in xrange(n_files * n_roots_per_file, n_elements+1):
            if i == n_elements or i % 100 == 0:
                print >>sys.stderr, "(Created", i, "elements of", n_elements, "...)"
                if i == n_elements:
                    break
            parent = None
            while True:
                parent = self.choose_one(elements)
                mc = parent.__type__()
                prop_and_mc_list = self.choose_suitable_property(parent,
                                                                 packages)
                if prop_and_mc_list:
                    assert parent != None
                    prop_and_mc = self.choose_one(prop_and_mc_list)
                    prop, target_mc = prop_and_mc
                    assert prop.isComposition()
                    assert not target_mc.isDatatypeDefinition()
                    slot = parent.__getSlotByName__(prop.getName())
                    if prop.getMultUpper() != -1 \
                           and len(slot) >= prop.getMultUpper():
                        # This slot is already full!
                        continue
                    break

            e = target_mc.newAtom()
            if prop.isAtMostOne():
                setattr(parent, prop.getName(), e)
            else:
                slot.add(e)
            elements.append(e)
            assert e.__parent__() == parent

            self.__fill_element(e, unique_strings_hash)

        #
        # Add associations between elements (and thus files)
        #
        elements = tuple(elements) 
        for idx, e in enumerate(elements):
            if idx % 100 == 0:
                print >>sys.stderr, "(Setting associations", idx, "elements of", \
                      len(elements), "...)"
            for p in e.__type__().__properties__():
                if p.isUnserializable():
                    continue
                if p.isAnonymous():
                    continue
                if not p.isAssociation():
                    continue
                oppo = p.getOppositeProperty()
                if oppo and oppo.isComposition():
                    continue
                lower, upper = get_lower_and_upper(p.getMultLower(),
                                                   p.getMultUpper())
                target = p.getTarget()
                if target.isDatatypeDefinition():
                    continue

                s = e.__getSlotByName__(p.getName())
                if len(s) + upper > p.getMultUpper():
                    upper = p.getMultUpper() - len(s)

                for i in xrange(lower, upper):
                    dest = self.choose_one_of_type(elements, target)
                    if dest is None:
                        break
                    if p.isAtMostOne():
                        setattr(e, p.getName(), dest)
                    else:
                        s.append(dest)
                
        #
        # Save results
        #
        print >>sys.stderr, "(Saving...)"
        for p in projects:
            coral.saveDocument(p)
    
        # Should we compare what we generated?
        if not compare_after_saving:
            return
        
        #
        # Rename projects so we can load the "new" ones
        # instead.
        #
        print >>sys.stderr, "(Renaming...)"
        filenames = [ p.getFilename() for p in projects ]
        for p in projects:
            p.setFilename(p.getFilename() + "_renamed")
    
        #
        # Try to load them back
        #
        print >>sys.stderr, "(Loading...)"
        new_projects = []
        for f in filenames:
            p = coral.loadDocument(f)
            new_projects.append(p)
        #
        # Compare
        #
        print >>sys.stderr, "(Comparing...)"
        for oldp, newp in map(None, projects, new_projects):
            assert len(oldp.contents()) == len(newp.contents())
            print >>sys.stderr, "(Comparing", newp.getFilename(), "...)"
            for oldroot, newroot in map(None, oldp.contents(),newp.contents()):
                b = coral.isSameModel(oldroot, newroot)
                assert b
                    
        print >>sys.stderr, "Finished! Everything OK!"
        
def main(arguments):
    def usage():
        print """
Usage: generate_arbitrary_models [ --help | --version ] [-f files] [ -c ]
                                 [ -n elements ] [ -r roots ]
                                 [ -m metamodel ] [ -p prefix ]
Options:
  --help         Display this information
  --version      Display version information and exit
  -c             Compare models saved with models in memory (crash-me)
  -f files       Amount of files to generate (default 1)
  -n elements    Total amount of elements to create (default 10000)
  -r roots       Roots per file (default 5)
  -m metamodel   Metamodel to use. Can be issued multiple times
                 (default UML14)
  -p prefix      Output prefix (default ./m)
  --unique-strings Creates unique strings for all 1..1 strings

A set of files are created (given with -f) and populated with
elements (given by -n). The element types are taken from the metamodels
specified by -m.
        """

    def version():
        import coral
        print "This is Coral version", coral.version, \
              " (arbitrary_model_generator)"
        
    import sys
    import getopt
    import os

    try:
        opts, args = getopt.getopt(arguments, "chp:r:f:n:m:",
                                   ["help", "version", "unique-strings" ])
    except getopt.GetoptError:
        usage()
        sys.exit(2)

    n_files = 1
    n_roots_per_file = 5
    metamodel_strings = [ "UML14" ]
    _set_metamodels = False
    n_elements = 10000
    prefix = "./m"
    crashme = False
    unique_strings = False

    for o,a in opts:
        if o in ("-h", "--help"):
            usage()
            sys.exit(0)
        elif o in ("-c"):
            crashme = True
        elif o in ("--version"):
            version()
            sys.exit(0)
        elif o in ("-p"):
            prefix = a
        elif o in ("-f"):
            n_files = int(a)
        elif o in ("-r"):
            n_roots_per_file = int(a)
        elif o in ("-n"):
            n_elements = int(a)
        elif o in ("-m"):
            if not _set_metamodels:
                _set_metamodels = True
                metamodel_strings = []
            metamodel_strings.append(a)
        elif o in ("--unique-strings"):
            unique_strings = True
        else:
            usage()
            sys.exit(2)

    a = ArbitraryModelGenerator()
    # Much faster execution for these pure functions
    m1 = Memoize(a, "collect_concrete_subclasses")
    m2 = Memoize(a, "choose_suitable_property")

    a.create_models(prefix, n_files, n_roots_per_file,
                    metamodel_strings, n_elements,
                    crashme, unique_strings)

if __name__ == "__main__":
    main(sys.argv[1:])
