import re
import unittest
from ir_datasets.formats import TrecQrel, TrecDoc, TrecQuery
from .base import DatasetIntegrationTest


class TestTrecArabic(DatasetIntegrationTest):
    def test_trec_arabic_docs(self):
        self._test_docs('trec-arabic', count=383872, items={
            0: TrecDoc('19940513_AFP_ARB.0001', re.compile('^\n\n \\&HT;    العلم الفلسطيني لم يُرفع فوق كنيس اريحا \n\n\n\n اريحا \\(الضفة الغربية\\) 31\\-5 \\(اف ب\\)\\- يقوم احد .{783}ارا عن استيائهم وغضبهم من ذلك ونظموا تظاهرات عدة الا ان الشرطة الاسرائيلية قامت كل مرة بتفريقهم\\. \n\n\n$', flags=48), re.compile('^<HEADLINE>\n \\&HT;    العلم الفلسطيني لم يُرفع فوق كنيس اريحا \n</HEADLINE>\n<TEXT>\n<P>\n اريحا \\(الضفة ال.{858}ائهم وغضبهم من ذلك ونظموا تظاهرات عدة الا ان الشرطة الاسرائيلية قامت كل مرة بتفريقهم\\. \n</P>\n</TEXT>\n$', flags=48)),
            9: TrecDoc('19940513_AFP_ARB.0010', re.compile('^\n\n \\&HT;    وفد الجامعة \\.\\. \n\n\n\n صنعاء \\- وكانت الجامعة العربية استبعدت فكرة ارسال قوة فصل خلال اجتماعه.{1334}ب في حين صرح قائد البحرية الجنوبية العقيد علي قاسم طالب بانه تم صد الهجوم الشمالي في هذا القطاع\\. \n\n\n$', flags=48), re.compile('^<HEADLINE>\n \\&HT;    وفد الجامعة \\.\\. \n</HEADLINE>\n<TEXT>\n<P>\n صنعاء \\- وكانت الجامعة العربية استبعدت فك.{1423}ح قائد البحرية الجنوبية العقيد علي قاسم طالب بانه تم صد الهجوم الشمالي في هذا القطاع\\. \n</P>\n</TEXT>\n$', flags=48)),
            383871: TrecDoc('20001220_AFP_ARB.0148', re.compile('^\n\nبوش يختار مهاجرا كوبيا وزيرا للاسكان\n\n\n\n اوستن \\(الولايات المتحدة\\) 02\\-21 \\(اف ب\\)\\- اعلن الرئيس المنتخ.{642}بير في قضية الطفل الكوبي اللاجىء أليان غونزاليس من خلال مشاركته في جهود منع عودة الطفل الى كوبا\\. \n\n\n$', flags=48), re.compile('^<HEADLINE>\nبوش يختار مهاجرا كوبيا وزيرا للاسكان\n</HEADLINE>\n<TEXT>\n<P>\n اوستن \\(الولايات المتحدة\\) 02\\-.{710} الطفل الكوبي اللاجىء أليان غونزاليس من خلال مشاركته في جهود منع عودة الطفل الى كوبا\\. \n</P>\n</TEXT>\n$', flags=48)),
        })

    def test_trec_arabic_queries(self):
        self._test_queries('trec-arabic/ar2001', count=25, items={
            0: TrecQuery(query_id='1', title='فنون العرض و المؤسسات الاسلامية في العالم العربي', description='ما هو اثر المؤسسات الاسلامية على فنون العرض مثل الرقص\n و الموسيقى في العالم العربي؟', narrative='المقالات المتعلقة  بالفنون الرياضية او التشكيلية\n او بفنون العرض خارج العالم العربي              \nاو بالسلوكيات الدينية خارج اطار فنون العرض  \nاو بالديون و القروض المالية  لا علاقة لها بالموضوع'),
            9: TrecQuery(query_id='10', title='القضاء على شلل الأطفال في الشرق الأوسط', description='هل تحقق أي تقدم  في القضاء على شلل الأطفال \n في الشرق الأوسط؟', narrative='ان المقالات  التي تتحدث عن القضاء على شلل الاطفال عالميا\n  او في مناطق اخرى  في العالم  من غير ذكر الشرق الاوسط\n ليست لها علاقة بالموضوع\n و يشمل الشرق الاوسط هنا الصومال و افغانستان و باكستان'),
            24: TrecQuery(query_id='25', title='الدور الاوروبى والامريكى فى عملية السلام فى الشرق الأوسط', description='ما هي أدوار الدول الأوربية وامريكا في عملية السلام فى الشرق\n الأسط؟', narrative='يتعلق  بهذا الموضوع كل مقال يخص التدخل الاوروبى والامريكى في\n  القرارات العربية لتوجيه عملية السلام في الشرق الاوسط\n وما لا يرتبط بهذا الموضوع هو التدخل الاوروبى و الامريكى فى \nالشؤون الداخلية لدول الشرق الاوسط')
        })
        self._test_queries('trec-arabic/ar2002', count=50, items={
            0: TrecQuery(query_id='26', title='مجلس المقاومة الوطني الكردستاني', description='كيف ينظر مجلس المقاومة الوطنية الى الإستقلال المحتمل للاكراد؟', narrative='الموضوع يتضمن نصوص متعلقة بتحركات مجلس المقاومة الوطنية ، مقالات تتحدث عن قيادة اوجلان ضمن جهود الاكراد للاستقلال .'),
            9: TrecQuery(query_id='35', title='ميناء سلمان وام قصر', description='ما أهمية العلاقة التجارية البحرية بين مينء سلمان في البحرين\n  وأم قصر في العراق؟', narrative='المقالات التي اعتبرت مناسبة تناقش المصالح الاقتصادية في وجود\n خط تجاري بحري بين ميناء سلمان في البحرين و ام قصر في العراق'),
            49: TrecQuery(query_id='75', title='فيروسات الكمبيوتر في الوطن العربي', description='ما هو تأثير فيروسات الكمبيوتر على القطاع الصناعي والتجاري في الدول العربية؟', narrative='المقالات التي اعتبرت جزءا من موضوع البحث تناولت المشاكل الناجمة\n عن فيروسات الكمبيوتر في الدول العربية. المقالات التي تناولت الآثار\n الناجمة عن الفيروسات في دول اخرى اعتبرت خارج نطاق البحث')
        })

    def test_trec_arabic_qrels(self):
        self._test_qrels('trec-arabic/ar2001', count=22744, items={
            0: TrecQrel(query_id='1', doc_id='19940515_AFP_ARB.0095', relevance=0, iteration='0'),
            9: TrecQrel(query_id='1', doc_id='19940606_AFP_ARB.0009', relevance=0, iteration='0'),
            22743: TrecQrel(query_id='25', doc_id='20001220_AFP_ARB.0044', relevance=0, iteration='0')
        })
        self._test_qrels('trec-arabic/ar2002', count=38432, items={
            0: TrecQrel(query_id='26', doc_id='19940515_AFP_ARB.0115', relevance=0, iteration='0'),
            9: TrecQrel(query_id='26', doc_id='19940519_AFP_ARB.0114', relevance=0, iteration='0'),
            38431: TrecQrel(query_id='75', doc_id='20001211_AFP_ARB.0042', relevance=0, iteration='0')
        })


if __name__ == '__main__':
    unittest.main()
