#!/usr/bin/env python

import iixr

# Test basic data types.

numbers = [12345678, 0, 1, 127, 128, 255, 256]

f = open("test", "wb")

w = iixr.FileWriter(f)

for number in numbers:

    w.write_number(number)

w.close()

f = open("test", "rb")

r = iixr.FileReader(f)

for number in numbers:

    n = r.read_number()

    print number == n, number, n

r.close()

# Test positions.

all_doc_positions = [

    [

        (123, [1, 3, 5, 15, 25]),

        (124, [0, 100])

    ],

    [

        (78, [9]),

        (196, [10, 11])

    ]

    ]

f = open("test", "wb")

w = iixr.PositionWriter(f)

for doc_positions in all_doc_positions:

    for docnum, positions in doc_positions:

        w.write_positions(docnum, positions)

    w.reset()

w.close()

f = open("test", "rb")

r = iixr.PositionReader(f)

for doc_positions in all_doc_positions:

    for docnum, positions in doc_positions:

        d, p = r.read_positions()

        print docnum == d, docnum, d

        print positions == p, positions, p

    r.reset()

r.close()

f = open("test", "wb")

w = iixr.PositionWriter(f)

offsets = []

for doc_positions in all_doc_positions:

    offsets.append(

        w.write_all_positions(doc_positions)

        )

w.close()

f = open("test", "rb")

r = iixr.PositionReader(f)

offsets.reverse()

all_doc_positions.reverse()

for offset, doc_positions in zip(offsets, all_doc_positions):

    dp = r.read_all_positions(offset)

    print doc_positions == dp, doc_positions, dp

r.close()

# Test fields.

doc_fields = [

    (123, ["testing", "fields", "stored", "compressed"]),

    (456, ["fields", "for a second", "document"]),

    (789, ["field value"]),

    (1234, []),

    (2345, ["abc", "def"]),

    (3456, ["apple", "banana", "cherry"]),

    (4567, ["drue", "eple"])

    ]

f = open("testF", "wb")

w = iixr.FieldWriter(f)

for docnum, fields in doc_fields:

    w.write_fields(docnum, fields)

w.close()

f = open("testF", "rb")

r = iixr.FieldReader(f)

for docnum, fields in doc_fields:

    dn, df = r.read_fields()

    print docnum == dn, docnum, dn

    print fields == df, fields, df

r.close()

# Test field index files.

indexed_docs = [

    (123, 100000987),

    (456, 100004321),

    (789, 100008765)

    ]

f = open("testFI", "wb")

w = iixr.FieldIndexWriter(f)

for docnum, offset in indexed_docs:

    w.write_document(docnum, offset)

w.close()

f = open("testFI", "rb")

r = iixr.FieldIndexReader(f)

for docnum, offset in indexed_docs:

    dn, o = r.read_document()

    print docnum == dn, docnum, dn

    print offset == o, offset, o

r.close()

# Test field dictionaries.

f = open("testF", "wb")

w = iixr.FieldWriter(f)

f2 = open("testFI", "wb")

w2 = iixr.FieldIndexWriter(f2)

wd = iixr.FieldDictionaryWriter(w, w2, 3)

for docnum, fields in doc_fields:

    wd.write_fields(docnum, fields)

wd.close()

f = open("testF", "rb")

r = iixr.FieldReader(f)

f2 = open("testFI", "rb")

r2 = iixr.FieldIndexReader(f2)

rd = iixr.FieldDictionaryReader(r, r2)

doc_fields_reversed = doc_fields[:]

doc_fields_reversed.reverse()

for docnum, fields in doc_fields_reversed:

    df = rd.read_fields(docnum)

    print fields == df, fields, df

for docnum in (13579, 246810):

    df = rd.read_fields(docnum)

    print df is None, df

rd.close()

# Test terms.

terms = [

    ("aardvark",  100000123),

    ("anteater",  100000456),

    ("badger",    100000789),

    ("bull",     1000001234),

    ("bulldog",  1000002345),

    ("cat",      1000003456)

    ]

f = open("test", "wb")

w = iixr.TermWriter(f)

for term, offset in terms:

    w.write_term(term, offset)

w.close()

f = open("test", "rb")

r = iixr.TermReader(f)

for term, offset in terms:

    t, o = r.read_term()

    print term == t, term, t

    print offset == o, offset, o

r.close()

# Test terms in index files.

indexed_terms = [

    ("aardvark",  100000123, 200000321),

    ("anteater",  100000456, 200000654),

    ("badger",    100000789, 200000987),

    ("bull",     1000001234, 200004321),

    ("bulldog",  1000002345, 200005432),

    ("cat",      1000003456, 200006543)

    ]

f = open("test", "wb")

w = iixr.TermIndexWriter(f)

for term, offset, info_offset in indexed_terms:

    w.write_term(term, offset, info_offset)

w.close()

f = open("test", "rb")

r = iixr.TermIndexReader(f)

for term, offset, info_offset in indexed_terms:

    t, o, i = r.read_term()

    print term == t, term, t

    print offset == o, offset, o

    print info_offset == i, info_offset, i

r.close()

# Test dictionaries with only term data.

f = open("test", "wb")

w = iixr.TermWriter(f)

f2 = open("testI", "wb")

w2 = iixr.TermIndexWriter(f2)

f3 = open("testP", "wb")

w3 = iixr.PositionWriter(f3)

wd = iixr.TermDictionaryWriter(w, w2, w3, 3)

for term, offset in terms:

    wd._write_term(term, offset)

wd.close()

f = open("test", "rb")

r = iixr.TermReader(f)

f2 = open("testI", "rb")

r2 = iixr.TermIndexReader(f2)

f3 = open("testP", "rb")

r3 = iixr.PositionReader(f3)

rd = iixr.TermDictionaryReader(r, r2, r3)

terms_reversed = terms[:]

terms_reversed.reverse()

for term, offset in terms_reversed:

    o = rd._find_term(term)

    print offset == o, offset, o

for term in ("dog", "dingo"):

    o = rd._find_term(term)

    print o is None, o

rd.close()

# Test dictionaries with term and position data.

terms_with_positions = [

    ("aardvark",  [(1, [2, 45, 96]), (20, [13])]),

    ("anteater",  [(1, [43, 44])]),

    ("badger",    [(7, [2, 22, 196]), (19, [55, 1333]), (21, [0])]),

    ("bull",      [(6, [128]), (16, [12])]),

    ("bulldog",   [(43, [17, 19, 256, 512])]),

    ("cat",       [(123, [12, 145, 196]), (1200, [113])])

    ]

f = open("test", "wb")

w = iixr.TermWriter(f)

f2 = open("testI", "wb")

w2 = iixr.TermIndexWriter(f2)

f3 = open("testP", "wb")

w3 = iixr.PositionWriter(f3)

wd = iixr.TermDictionaryWriter(w, w2, w3, 3)

for term, doc_positions in terms_with_positions:

    wd.write_term_positions(term, doc_positions)

wd.close()

f = open("test", "rb")

r = iixr.TermReader(f)

f2 = open("testI", "rb")

r2 = iixr.TermIndexReader(f2)

f3 = open("testP", "rb")

r3 = iixr.PositionReader(f3)

rd = iixr.TermDictionaryReader(r, r2, r3)

terms_reversed = terms_with_positions[:]

terms_reversed.reverse()

for term, doc_positions in terms_reversed:

    dp = rd.find_positions(term)

    print doc_positions == dp, doc_positions, dp

for term in ("dog", "dingo"):

    dp = rd.find_positions(term)

    print dp is None, dp

rd.close()

# Test high-level index operations.

docs = [

    (1, "The cat sat on the mat"),

    (2, "Every good boy deserves football"),

    (13, "One good turn deserves another"),

    (14, "Every man for himself"),

    (25, "Red sky at night shepherd's delight"),

    (36, "She sells sea shells on the sea shore")

    ]

doc_tests = [

    ("Every", [(2, [0]), (14, [0])]),

    ("good", [(2, [1]), (13, [1])]),

    ("deserves", [(2, [3]), (13, [3])]),

    ("sea", [(36, [2, 6])])

    ]

index = iixr.Index("test_index")

wi = index.get_writer(3)

for docnum, text in docs:

    for position, term in enumerate(text.split()):

        wi.add_position(term, docnum, position)

wi.close()

rd = index.get_reader()

for term, doc_positions in doc_tests:

    dp = rd.find_positions(term)

    print doc_positions == dp, doc_positions, dp

index.close()

# vim: tabstop=4 expandtab shiftwidth=4