import unittest

from whoosh import analysis, fields, formats, index, qparser, query, searching, scoring
from whoosh.filedb.filestore import RamStorage
from whoosh.query import *
from whoosh.searching import Searcher
from whoosh.scoring import FieldSorter

class TestSearching(unittest.TestCase):
    def make_index(self):
        s = fields.Schema(key = fields.ID(stored = True),
                          name = fields.TEXT,
                          value = fields.TEXT)
        st = RamStorage()
        ix = st.create_index(s)
        
        w = ix.writer()
        w.add_document(key = u"A", name = u"Yellow brown", value = u"Blue red green render purple?")
        w.add_document(key = u"B", name = u"Alpha beta", value = u"Gamma delta epsilon omega.")
        w.add_document(key = u"C", name = u"One two", value = u"Three rendered four five.")
        w.add_document(key = u"D", name = u"Quick went", value = u"Every red town.")
        w.add_document(key = u"E", name = u"Yellow uptown", value = u"Interest rendering outer photo!")
        w.commit()
        
        return ix
    
    def _get_keys(self, stored_fields):
        return sorted([d.get("key") for d in stored_fields])
    
    def _docs(self, q, s):
        return self._get_keys([s.stored_fields(docnum) for docnum
                               in q.docs(s)])
    
    def _doc_scores(self, q, s, w):
        return self._get_keys([s.stored_fields(docnum) for docnum, score
                               in q.doc_scores(s, weighting = w)])
    
    def test_empty_index(self):
        schema = fields.Schema(key = fields.ID(stored=True), value = fields.TEXT)
        st = RamStorage()
        self.assertRaises(index.EmptyIndexError, st.open_index, schema)
    
    def test_docs_method(self):
        ix = self.make_index()
        s = ix.searcher()
        
        self.assertEqual(self._get_keys(s.documents(name = "yellow")), [u"A", u"E"])
        self.assertEqual(self._get_keys(s.documents(value = "red")), [u"A", u"D"])
    
    def _run_query(self, q, result):
        ix = self.make_index()
        s = ix.searcher()
        self.assertEqual(self._docs(q, s), result)
        
    def test_term(self):
        self._run_query(Term("name", u"yellow"), [u"A", u"E"])
        self._run_query(Term("value", u"zeta"), [])
        self._run_query(Term("value", u"red"), [u"A", u"D"])
        
    def test_require(self):
        self._run_query(Require(Term("value", u"red"), Term("name", u"yellow")),
                        [u"A"])
        
    def test_and(self):
        self._run_query(And([Term("value", u"red"), Term("name", u"yellow")]),
                        [u"A"])
        
    def test_or(self):
        self._run_query(Or([Term("value", u"red"), Term("name", u"yellow")]),
                        [u"A", u"D", u"E"])
    
    def test_or_minmatch(self):
        schema = fields.Schema(k=fields.STORED, v=fields.TEXT)
        st = RamStorage()
        ix = st.create_index(schema)
        
        w = ix.writer()
        w.add_document(k=1, v=u"alfa bravo charlie delta echo")
        w.add_document(k=2, v=u"bravo charlie delta echo foxtrot")
        w.add_document(k=3, v=u"charlie delta echo foxtrot golf")
        w.add_document(k=4, v=u"delta echo foxtrot golf hotel")
        w.add_document(k=5, v=u"echo foxtrot golf hotel india")
        w.add_document(k=6, v=u"foxtrot golf hotel india juliet")
        w.commit()
        
        s = ix.searcher()
        q = Or([Term("v", "echo"), Term("v", "foxtrot")], minmatch=2)
        r = s.search(q)
        self.assertEqual(sorted(d["k"] for d in r), [2, 3, 4, 5])
    
    def test_not(self):
        self._run_query(Or([Term("value", u"red"), Term("name", u"yellow"), Not(Term("name", u"quick"))]),
                        [u"A", u"E"])
    
    def test_not2(self):
        schema = fields.Schema(name=fields.ID(stored=True), value=fields.TEXT)
        storage = RamStorage()
        ix = storage.create_index(schema)
        writer = ix.writer()
        writer.add_document(name=u"a", value=u"alfa bravo charlie delta echo")
        writer.add_document(name=u"b", value=u"bravo charlie delta echo foxtrot")
        writer.add_document(name=u"c", value=u"charlie delta echo foxtrot golf")
        writer.add_document(name=u"d", value=u"delta echo golf hotel india")
        writer.add_document(name=u"e", value=u"echo golf hotel india juliet")
        writer.commit()
        
        searcher = ix.searcher()
        p = qparser.QueryParser("value")
        results = searcher.search(p.parse("echo NOT golf"))
        self.assertEqual(sorted([d["name"] for d in results]), ["a", "b"])
        
        results = searcher.search(p.parse("echo NOT bravo"))
        self.assertEqual(sorted([d["name"] for d in results]), ["c", "d", "e"])
        searcher.close()
        
        ix.delete_by_term("value", u"bravo")
        ix.commit()
        
        searcher = ix.searcher()
        results = searcher.search(p.parse("echo NOT charlie"))
        self.assertEqual(sorted([d["name"] for d in results]), ["d", "e"])
        searcher.close()
    
    def test_andnot(self):
        self._run_query(AndNot(Term("name", u"yellow"), Term("value", u"purple")),
                        [u"E"])
    
    def test_variations(self):
        self._run_query(Variations("value", u"render"), [u"A", u"C", u"E"])
    
    def test_topnot(self):
        self._run_query(Not(Term("name", "yellow")), [u"B", u"C", u"D"])
    
    def test_wildcard(self):
        self._run_query(Or([Wildcard('value', u'*red*'), Wildcard('name', u'*yellow*')]),
                        [u"A", u"C", u"D", u"E"])
    
#        for wcls in dir(scoring):
#            if wcls is scoring.Weighting: continue
#            if isinstance(wcls, scoring.Weighting):
#                for query, result in tests:
#                    self.assertEqual(self._doc_scores(query, s, wcls), result)
    
    def test_range(self):
        schema = fields.Schema(id=fields.ID(stored=True), content=fields.TEXT)
        st = RamStorage()
        ix = st.create_index(schema)
        
        w = ix.writer()
        w.add_document(id=u"A", content=u"alfa bravo charlie delta echo")
        w.add_document(id=u"B", content=u"bravo charlie delta echo foxtrot")
        w.add_document(id=u"C", content=u"charlie delta echo foxtrot golf")
        w.add_document(id=u"D", content=u"delta echo foxtrot golf hotel")
        w.add_document(id=u"E", content=u"echo foxtrot golf hotel india")
        w.commit()
        s = ix.searcher()
        qp = qparser.QueryParser("content", schema=schema)
        
        q = qp.parse(u"charlie [delta TO foxtrot]")
        self.assertEqual(q.__class__.__name__, "And")
        self.assertEqual(q[0].__class__.__name__, "Term")
        self.assertEqual(q[1].__class__.__name__, "TermRange")
        self.assertEqual(q[1].start, "delta")
        self.assertEqual(q[1].end, "foxtrot")
        self.assertEqual(q[1].startexcl, False)
        self.assertEqual(q[1].endexcl, False)
        ids = sorted([d['id'] for d in s.search(q)])
        self.assertEqual(ids, [u'A', u'B', u'C'])
        
        q = qp.parse(u"foxtrot {echo TO hotel]")
        self.assertEqual(q.__class__.__name__, "And")
        self.assertEqual(q[0].__class__.__name__, "Term")
        self.assertEqual(q[1].__class__.__name__, "TermRange")
        self.assertEqual(q[1].start, "echo")
        self.assertEqual(q[1].end, "hotel")
        self.assertEqual(q[1].startexcl, True)
        self.assertEqual(q[1].endexcl, False)
        ids = sorted([d['id'] for d in s.search(q)])
        self.assertEqual(ids, [u'B', u'C', u'D', u'E'])
        
        q = qp.parse(u"{bravo TO delta}")
        self.assertEqual(q.__class__.__name__, "TermRange")
        self.assertEqual(q.start, "bravo")
        self.assertEqual(q.end, "delta")
        self.assertEqual(q.startexcl, True)
        self.assertEqual(q.endexcl, True)
        ids = sorted([d['id'] for d in s.search(q)])
        self.assertEqual(ids, [u'A', u'B', u'C'])
    
    def test_range_clusiveness(self):
        schema = fields.Schema(id=fields.ID(stored=True))
        st = RamStorage()
        ix = st.create_index(schema)
        w = ix.writer()
        for letter in u"abcdefg":
            w.add_document(id=letter)
        w.commit()
        s = ix.searcher()
        
        def do(startexcl, endexcl, string):
            q = TermRange("id", "b", "f", startexcl, endexcl)
            r = "".join(sorted(d['id'] for d in s.search(q)))
            self.assertEqual(r, string)
            
        do(False, False, "bcdef")
        do(True, False, "cdef")
        do(True, True, "cde")
        do(False, True, "bcde")
        
    def test_open_ranges(self):
        schema = fields.Schema(id=fields.ID(stored=True))
        st = RamStorage()
        ix = st.create_index(schema)
        w = ix.writer()
        for letter in u"abcdefg":
            w.add_document(id=letter)
        w.commit()
        s = ix.searcher()
        
        qp = qparser.QueryParser("id", schema=schema)
        def do(qstring, result):
            q = qp.parse(qstring)
            r = "".join(sorted([d['id'] for d in s.search(q)]))
            self.assertEqual(r, result)
            
        do(u"[b TO]", "bcdefg")
        do(u"[TO e]", "abcde")
        do(u"[b TO d]", "bcd")
        do(u"{b TO]", "cdefg")
        do(u"[TO e}", "abcd")
        do(u"{b TO d}", "c")
    
    def test_keyword_or(self):
        schema = fields.Schema(a=fields.ID(stored=True), b=fields.KEYWORD)
        st = RamStorage()
        ix = st.create_index(schema)
        
        w = ix.writer()
        w.add_document(a=u"First", b=u"ccc ddd")
        w.add_document(a=u"Second", b=u"aaa ddd")
        w.add_document(a=u"Third", b=u"ccc eee")
        w.commit()
        
        qp = qparser.QueryParser("b", schema=schema)
        searcher = ix.searcher()
        qr = qp.parse(u"b:ccc OR b:eee")
        self.assertEqual(qr.__class__, query.Or)
        r = searcher.search(qr)
        self.assertEqual(len(r), 2)
        self.assertEqual(r[0]["a"], "Third")
        self.assertEqual(r[1]["a"], "First")

    def test_merged(self):
        sc = fields.Schema(id=fields.ID(stored=True), content=fields.TEXT)
        st = RamStorage()
        ix = st.create_index(sc)
        w = ix.writer()
        w.add_document(id=u"alfa", content=u"alfa")
        w.add_document(id=u"bravo", content=u"bravo")
        w.add_document(id=u"charlie", content=u"charlie")
        w.add_document(id=u"delta", content=u"delta")
        w.commit()
        
        s = ix.searcher()
        r = s.search(query.Term("content", u"bravo"))
        self.assertEqual(len(r), 1)
        self.assertEqual(r[0]["id"], "bravo")
        
        w = ix.writer()
        w.add_document(id=u"echo", content=u"echo")
        w.commit()
        self.assertEqual(len(ix.segments), 1)
        
        s = ix.searcher()
        r = s.search(query.Term("content", u"bravo"))
        self.assertEqual(len(r), 1)
        self.assertEqual(r[0]["id"], "bravo")
        
    def test_multireader(self):
        sc = fields.Schema(id=fields.ID(stored=True), content=fields.TEXT)