"""Benchmark for Storm.

Note that the original test omitted the 'step 7 multiview' test which was incomplete, 
so its omitted here as well.

"""

import datetime
import os
thisdir = os.path.dirname(__file__)
import sys
import time

import storm
from storm.locals import *


class Zoo(object):
    __storm_table__ = 'Zoo'
    ID = Int(primary=True, default=storm.expr.Sequence('Zoo_ID_seq'))
    Name = Unicode()
    Founded = Date()
    Opens = Time()
    LastEscape = DateTime()
    Admission = Float()
    
    def __init__(self, **kwargs):
        for k, v in kwargs.iteritems():
            setattr(self, k, v)


class Animal(object):
    __storm_table__ = 'Animal'
    ID = Int(primary=True, default=storm.expr.Sequence('Animal_ID_seq'))
    ZooID = Int()
    zoo = Reference(ZooID, Zoo.ID)
    Name = Unicode()
    Species = Unicode()
    Legs = Int(default=4)
    LastEscape = DateTime()
    Lifespan = Float()
    MotherID = Int()
##    mother = Reference(MotherID, Animal.ID)
    PreferredFoodID = Int()
    AlternateFoodID = Int()
    
    def __init__(self, **kwargs):
        for k, v in kwargs.iteritems():
            setattr(self, k, v)


class ZooMark(object):
    
    def setup(self, uri, iterations):
        global db, store, ITERATIONS
        
        db = create_database('postgres://scott:tiger@127.0.0.1/test')
        
        store = Store(db)
        ITERATIONS = iterations
    
    def teardown(self):
        store.close()
        conn = db.connect()
        conn._raw_connection.set_isolation_level(0)
        conn.raw_execute("DROP TABLE Animal", None)
        conn.raw_execute("DROP TABLE Zoo", None)
        conn.raw_execute("DROP SEQUENCE Animal_ID_seq", None)
        conn.raw_execute("DROP SEQUENCE Zoo_ID_seq", None)
        conn.close()
        
    def step_1_create_tables(self):
        # Storm doesn't do CREATE TABLE, etc for you
        conn = db.connect()
        conn._raw_connection.set_isolation_level(0)
        conn.raw_execute('CREATE SEQUENCE Animal_ID_seq START 1;', None)
        conn.raw_execute('CREATE TABLE Animal (MotherID INT4, LastEscape TIMESTAMP, '
                      'Name VARCHAR(100), ZooID INT4, Age INTERVAL(53), '
                      'PreferredFoodID INT4, Lifespan FLOAT4, '
                      """ID INT4 DEFAULT nextval('Animal_ID_seq'), """
                      'AlternateFoodID INT4, PreviousZoos TEXT, '
                      'Legs INT4 DEFAULT 4, Species VARCHAR(100), '
                      'PRIMARY KEY (ID));', None)
        conn.raw_execute('CREATE INDEX iAnimalZooID ON Animal (ZooID);', None)
        conn.raw_execute('CREATE SEQUENCE Zoo_ID_seq START 1;', None)
        conn.raw_execute('CREATE TABLE Zoo (Name VARCHAR(255), Admission DECIMAL(6, 2), '
                      'Founded DATE, LastEscape TIMESTAMP, '
                      '''ID INT4 DEFAULT nextval('Zoo_ID_seq'), '''
                      'Opens TIME, PRIMARY KEY (ID));', None)
        conn.raw_execute('CREATE INDEX iZooID ON Zoo (ID);', None)
        conn.close()
    
    def step_1a_populate(self):
        wap = Zoo(Name=u'Wild Animal Park',
                  Founded=datetime.date(2000, 1, 1),
                  # 59 can give rounding errors with divmod, which
                  # AdapterFromADO needs to correct.
                  Opens=datetime.time(8, 15, 59),
                  LastEscape=datetime.datetime(2004, 7, 29, 5, 6, 7),
                  Admission=4.95)
        store.add(wap)
        
        sdz = Zoo(Name = u'San Diego Zoo',
                  # This early date should play havoc with a number
                  # of implementations.
                  Founded = datetime.date(1835, 9, 13),
                  Opens = datetime.time(9, 0, 0),
                  Admission = 0,
                  )
        store.add(sdz)
        
        bio = Zoo(Name = u'Montr\xe9al Biod\xf4me',
                  Founded = datetime.date(1992, 6, 19),
                  Opens = datetime.time(9, 0, 0),
                  Admission = 11.75,
                  )
        store.add(bio)
        
        seaworld = Zoo(Name = u'Sea_World', Admission = 60)
        store.add(seaworld)
        
        # Let's add a crazy futuristic Zoo to test large date values.
        lp = Zoo(Name = u'Luna Park',
                 Founded = datetime.date(2072, 7, 17),
                 Opens = datetime.time(0, 0, 0),
                 Admission = 134.95,
                 )
        store.add(lp)
        store.flush()
        
        # Animals
        leopard = Animal(Species=u'Leopard', Lifespan=73.5,)
        store.add(leopard)
        leopard.ZooID = wap.ID
        leopard.LastEscape = datetime.datetime(2004, 12, 21, 8, 15, 0, 999907)
        
        store.add(Animal(Species=u'Lion', ZooID=wap.ID))
        store.add(Animal(Species=u'Slug', Legs=1, Lifespan=.75))
        store.add(Animal(Species=u'Tiger', ZooID=sdz.ID))
        
        # Override Legs.default with itself just to make sure it works.
        store.add(Animal(Species=u'Bear', Legs=4))
        store.add(Animal(Species=u'Ostrich', Legs=2, Lifespan=103.2))
        store.add(Animal(Species=u'Centipede', Legs=100))
        
        store.add(Animal(Species=u'Emperor Penguin', Legs=2, ZooID=seaworld.ID))
        store.add(Animal(Species=u'Adelie Penguin', Legs=2, ZooID=seaworld.ID))
        
        store.add(Animal(Species=u'Millipede', Legs=1000000, ZooID=sdz.ID))
        
        # Add a mother and child to test relationships
        bai_yun = Animal(Species=u'Ape', Nameu='Bai Yun', Legs=2)
        store.add(bai_yun)
        store.add(Animal(Species=u'Ape', Name=u'Hua Mei', Legs=2,
                         MotherID=bai_yun.ID))
        store.flush()
        store.commit()
    
    def step_2_insert(self):
        for x in xrange(ITERATIONS):
            store.add(Animal(Species=u'Tick', Name=u'Tick %d' % x, Legs=8))
            
        # add flush() so that SQL is actually emitted.  otherwise, 
        # no SQL is emitted in this step (it occurs at the start of step 3).
        # adding store.flush() raises the time for 100 iterations from .02 seconds
        # to .11 seconds
        store.flush()
        
    def step_3_Properties(self):
        # use list(store.find(...)) so that the SQL is actually emitted.
        # otherwise, no SQL is emitted in this step.
        # adding list() raises the time for 100 iterations from .02 seconds (assuming flush() in step 2)
        # to 1.5 seconds
        for x in xrange(ITERATIONS):
            # Zoos
            WAP = list(store.find(Zoo, Zoo.Name==u'Wild Animal Park'))
            SDZ = list(store.find(Zoo, Zoo.Founded==datetime.date(1835, 9, 13)))
            Biodome = list(store.find(Zoo, Zoo.Name==u'Montr\xe9al Biod\xf4me'))
            seaworld = list(store.find(Zoo, Zoo.Admission == float(60)))
            
            # Animals
            leopard = list(store.find(Animal, Animal.Species == u'Leopard'))
            ostrich = list(store.find(Animal, Animal.Species==u'Ostrich'))
            millipede = list(store.find(Animal, Animal.Legs==1000000))
            ticks = list(store.find(Animal, Animal.Species==u'Tick'))

    def step_4_Expressions(self):
        for x in xrange(ITERATIONS):
            assert len(list(store.find(Zoo))) == 5
            assert len(list(store.find(Animal))) == ITERATIONS + 12
            assert len(list(store.find(Animal, Animal.Legs==4))) == 4
            assert len(list(store.find(Animal, Animal.Legs == 2))) == 5
            assert len(list(store.find(Animal, Animal.Legs >= 2, Animal.Legs < 20))) == ITERATIONS + 9
            assert len(list(store.find(Animal, Animal.Legs > 10))) == 2
            assert len(list(store.find(Animal, Animal.Lifespan > 70))) == 2
            assert len(list(store.find(Animal, Animal.Species.like(u'L%')))) == 2
            assert len(list(store.find(Animal, Animal.Species.like(u'%pede')))) == 2
            
            assert len(list(store.find(Animal, Animal.LastEscape != None))) == 1
            assert len(list(store.find(Animal, Animal.LastEscape == None))) == ITERATIONS + 11
            
            # In operator (containedby)
            assert len(list(store.find(Animal, Animal.Species.like(u'%pede%')))) == 2
            assert len(list(store.find(Animal,
                    Animal.Species.is_in((u'Lion', u'Tiger', u'Bear'))))) == 3
            
            # Try In with cell references
            class thing(object): pass
            pet, pet2 = thing(), thing()
            pet.Name, pet2.Name = u'Slug', u'Ostrich'
            assert len(list(store.find(Animal, Animal.Species.is_in((pet.Name, pet2.Name))))) == 2
            
            # logic and other functions
            class LENGTH(storm.expr.NamedFunc):
                name = "LENGTH"
            name = 'Lion'
            assert len(list(store.find(Animal, LENGTH(Animal.Species) == len(name)))) == ITERATIONS + 3
            
            assert len(list(store.find(Animal, Animal.Species.like(u'%i%')))) == ITERATIONS + 7
            
            # Test now(), today(), year(), month(), day()
            class NOW(storm.expr.NamedFunc):
                name = "now"
            assert len(list(store.find(Zoo, Zoo.Founded != None, Zoo.Founded < NOW()))) == 3
            assert len(list(store.find(Animal, Animal.LastEscape == NOW()))) == 0
            class date_part(storm.expr.NamedFunc):
                name = "date_part"
            assert len(list(store.find(Animal, date_part('year', Animal.LastEscape) == 2004))) == 1
            assert len(list(store.find(Animal, date_part('month', Animal.LastEscape) == 12))) == 1
            assert len(list(store.find(Animal, date_part('day', Animal.LastEscape) == 21))) == 1
    
    def step_5_Aggregates(self):
        def getall(*args, **kwargs):
            return store.execute(Select(*args, **kwargs)).get_all()
        
        for x in xrange(ITERATIONS):
            # views
            legs = getall(Animal.Legs)
            legs.sort()
            
            expected = {u'Leopard': 73.5,
                        u'Slug': .75,
                        u'Tiger': None,
                        u'Lion': None,
                        u'Bear': None,
                        u'Ostrich': 103.2,
                        u'Centipede': None,
                        u'Emperor Penguin': None,
                        u'Adelie Penguin': None,
                        u'Millipede': None,
                        u'Ape': None,
                        u'Tick': None,
                        }
            for species, lifespan in getall([Animal.Species, Animal.Lifespan]):
                assert lifespan == expected[species]
            
            expected = [u'Montr\xe9al Biod\xf4me', u'Wild Animal Park']
            # changed "datetime.now()" to be "date.today()" to provide type
            # consistency (other platforms require it)
            e = getall([Zoo.Name], storm.expr.And(Zoo.Founded != None,
                                    Zoo.Founded <= datetime.date.today(),
                                    Zoo.Founded >= datetime.date(1990, 1, 1)))
            values = [val[0] for val in e]
            assert set(values) == set(expected)
            
            # distinct
            legs = [x[0] for x in getall([Animal.Legs], distinct=True)]
            legs.sort()
    
    def step_6_Editing(self):
        for x in xrange(ITERATIONS):
            # Edit
            SDZ = store.find(Zoo, Zoo.Name==u'San Diego Zoo')[0]
            store.find(Zoo, Zoo.ID==SDZ.ID).set(
                     Name=u'The San Diego Zoo',
                     Founded = datetime.date(1900, 1, 1),
                     Opens = datetime.time(7, 30, 0),
                     Admission = 35.00)
            
            # Test edits
            SDZ = store.find(Zoo, Zoo.Name==u'The San Diego Zoo')[0]
            assert SDZ.Founded == datetime.date(1900, 1, 1), SDZ.Founded
            
            # Change it back
            store.find(Zoo, Zoo.ID==SDZ.ID).set(
                     Name = u'San Diego Zoo',
                     Founded = datetime.date(1835, 9, 13),
                     Opens = datetime.time(9, 0, 0),
                     Admission = 0)
            
            # Test re-edits
            SDZ = store.find(Zoo, Zoo.Name==u'San Diego Zoo')[0]
            assert SDZ.Founded == datetime.date(1835, 9, 13), SDZ.Founded