import time,psycopg2,random

NROWS = 1000000
NATTRS = 32             # number of attributes per row

attr_columns = ["a%02d" % n for n in range(NATTRS) ]

def query(s,*args,**kwargs):
    print(s)
    return c.execute( s,*args,**kwargs )

def bench( f, *args, **kwargs ):
    t = time.time()
    try:
        return f( *args, **kwargs )
    finally:
        dt = time.time()-t
        c = db.cursor()
        query( "SELECT pg_relation_size( 'testins' )")
        size, = c.fetchone()
        query( "SELECT count(*) FROM testins")
        rows, = c.fetchone()
        print("%s: %.03fs %d rows %10.03f kB" % (f.__name__, dt, rows, size/1024))

def drop():
    query( "DROP TABLE profiles_btree" );
    query( "DROP TABLE profiles_gist" );
    query( "DROP TABLE profiles_bloom" );
    query( "DROP TABLE profiles_gin" );

def build():
    query( "CREATE UNLOGGED TABLE profiles_btree ( id INTEGER NOT NULL,"
        + ",".join( "%s INTEGER NULL" % c for c in attr_columns )
        + " );" )

    query( "CREATE UNLOGGED TABLE profiles_bloom ( id INTEGER NOT NULL,"
        + ",".join( "%s INTEGER NULL" % c for c in attr_columns )
        + " );" )

    query( "CREATE UNLOGGED TABLE profiles_gist ( id INTEGER NOT NULL, attrs INTEGER[] NOT NULL )" );
    query( "CREATE UNLOGGED TABLE profiles_gin  ( id INTEGER NOT NULL, attrs INTEGER[] NOT NULL )" );

def fill():
    print("Generate...")
    insert_attrs = []
    insert_arrays = []
    for nrow in range( NROWS ):
        attrs = [ None ] * NATTRS
        array = []

        # Make 5 columns with 50% True, 50% NULL
        for pos in range( 5 ):
            if random.randint(0,1):
                attrs[pos] = 1
                array.append( pos )

        # Put 5 True values at random in the other columns
        for n in range( 5 ):
            pos = random.randint( 5, NATTRS-1 )
            attrs[ pos ] = 1
            array.append( pos )

        insert_attrs.append( [nrow] + attrs )
        insert_arrays.append( [nrow, array] )

    print("INSERT")
    c.executemany( "INSERT INTO profiles_btree VALUES ("+",".join( "%s" for n in range(NATTRS+1))+")", insert_attrs )
    query( "INSERT INTO profiles_bloom SELECT * FROM profiles_btree")

    print("INSERT")
    c.executemany( "INSERT INTO profiles_gist VALUES (%s,%s)", insert_arrays )
    query( "INSERT INTO profiles_gin   SELECT * FROM profiles_gist")

    for table in "btree","gist","bloom","gin":
        query( "ALTER TABLE profiles_%s ADD PRIMARY KEY(id)" % table )


db = psycopg2.connect("dbname=test")
c = db.cursor()
if 1:
    try:
        drop()
        query("COMMIT")
    except Exception as e:
        print(e)
        query("ROLLBACK")

    try:
        build()
        query("COMMIT")
    except Exception as e:
        print(e)
        query("ROLLBACK")

    fill()
    query("COMMIT")

    query("CREATE INDEX ON profiles_bloom USING bloom(%s)" % ",".join(attr_columns) )
    query("CREATE INDEX ON profiles_gist  USING GIST (attrs gist__int_ops)" )
    query("CREATE INDEX ON profiles_gist  USING GIN  (attrs gin__int_ops)" )
    for col in attr_columns:
        query("CREATE INDEX ON profiles_btree(%s) WHERE %s IS NOT NULL"%(col,col))
    query("COMMIT")
    query("VACUUM ANALYZE")