import re
import unittest
import ir_datasets
from ir_datasets.datasets.pmc import PmcDoc, TrecCdsQuery, TrecCds2016Query
from ir_datasets.formats import TrecQrel
from .base import DatasetIntegrationTest


_logger = ir_datasets.log.easy()


class TestPmc(DatasetIntegrationTest):
    def test_pmc_docs(self):
        self._test_docs('pmc/v1', count=733111, items={
            0: PmcDoc('1469887', 'Environmental Health Perspectives', 'Neoplastic transformation of cultured mammalian cells by estrogens and estrogenlike chemicals.', re.compile('^Estrogens are clearly carcinogenic in humans and rodents but the mechanisms by which these hormones .{1529}tiple effects of estrogens acting together cause genetic alterations leading to cell transformation\\.$', flags=48), ''),
            9: PmcDoc('1500872', 'Nucleic Acids Research', 'Metabolic regulation of \nApoB\n mRNA editing is associated with phosphorylation of APOBEC-1 complementation factor', re.compile('^Apolipoprotein B \\(\napoB\n\\) mRNA editing is a nuclear event that minimally requires the RNA substrate,.{1451}tion of ACF with APOBEC\\-1 and thereby increasing the probability of editosome assembly and activity\\.$', flags=48), re.compile('^INTRODUCTION\nApoB\n mRNA editing involves the site\\-specific deamination of cytidine 6666 to uridine w.{33462}modulate \napoB\n mRNA editing in the context of current models of editosome composition and assembly\\.$', flags=48)),
            733110: PmcDoc('3808453', 'PLoS Computational Biology', 'The Influence of Synaptic Weight Distribution on Neuronal Population Dynamics', re.compile('^The manner in which different distributions of synaptic weights onto cortical neurons shape their sp.{1093}s enable the network to respond faster and with more stability in the face of external fluctuations\\.$', flags=48), re.compile('^Introduction\nExperiments analyzing the distribution of synaptic weights impinging onto neurons typic.{74264}ns obtained from different synaptic weight distributions\\.\n\\(PDF\\)\nClick here for additional data file\\.$', flags=48)),
        })
        self._test_docs('pmc/v2', count=1255260, items={
            0: PmcDoc('4504085', 'Burns', 'Preventing childhood scalds within the home: Overview of systematic reviews and a systematic review of primary studies', re.compile('^Highlights\n•\nWe performed an overview of published systematic reviews and a systematic review of pri.{388}long with thermometers or thermostatic mixing valves is effective in reducing hot water temperature\\.$', flags=48), re.compile('^1\nIntroduction\nChildren are at particular risk of thermal injuries\\. Globally, thermal injuries are t.{26998}, engineering and educational approaches to reduce scalds risk\\.\nConflict of interest statement\nNone\\.$', flags=48)),
            9: PmcDoc('4210750', 'Bundesgesundheitsblatt, Gesundheitsforschung, Gesundheitsschutz', 'Web-based questionnaires to capture acute infections in long-term cohorts', re.compile('^Background\nIncidence of acute respiratory infections \\(ARI\\) and gastrointestinal infections \\(GII\\) are.{2411}election bias\n is possible and must be kept in mind when discussing generalizability of the results\\.$', flags=48), re.compile('^Acute respiratory \\(ARI\\) and gastrointestinal infections \\(GII\\) have an important share of the overall.{18063}er forms of data collection as long as its limitations—specifically selection bias—are kept in mind\\.$', flags=48)),
            1255259: PmcDoc('4486443', 'OMICS : a Journal of Integrative Biology', 'Towards Development of Clustering Applications for Large-Scale Comparative Genotyping and Kinship Analysis Using Y-Short Tandem Repeats', re.compile('^Abstract\nY\\-chromosome short tandem repeats \\(Y\\-STRs\\) are genetic markers with practical applications .{1186}ential for further development towards fully automatic clustering of any large\\-scale genotypic data\\.$', flags=48), re.compile('^Introduction\nY\\-\nchromosome short tandem repeats\n \\(Y\\-STRs\\) are a class of genetic markers found only .{32192}s, implementing the algorithm within the aforementioned clustering tools is experimentally feasible\\.$', flags=48)),
        })

    def test_pmc_queries(self):
        self._test_queries('pmc/v1/trec-cds-2014', count=30, items={
            0: TrecCdsQuery('1', 'diagnosis', 'A 58-year-old African-American woman presents to the ER with episodic pressing/burning anterior chest pain that began two days earlier for the first time in her life. The pain started while she was walking, radiates to the back, and is accompanied by nausea, diaphoresis and mild dyspnea, but is not increased on inspiration. The latest episode of pain ended half an hour prior to her arrival. She is known to have hypertension and obesity. She denies smoking, diabetes, hypercholesterolemia, or a family history of heart disease. She currently takes no medications. Physical examination is normal. The EKG shows nonspecific changes.', '58-year-old woman with hypertension and obesity presents with exercise-related episodic chest pain radiating to the back.'),
            9: TrecCdsQuery('10', 'diagnosis', 'A physician is called to see a 67-year-old woman who underwent cardiac catheterization via the right femoral artery earlier in the morning. She is now complaining of a cool right foot. Upon examination she has a pulsatile mass in her right groin with loss of distal pulses, and auscultation reveals a bruit over the point at which the right femoral artery was entered.', '67-year-old woman status post cardiac catheterization via right femoral artery, now with a cool, pulseless right foot and right femoral bruit.'),
            29: TrecCdsQuery('30', 'treatment', 'A 72-year-old man complains of increasing calf pain when walking uphill. The symptoms have gradually increased over the past 3 months. The patient had an uncomplicated myocardial infarction 2 years earlier and a transient ischemic attack 6 months ago. Over the past month, his blood pressure has worsened despite previous control with diltiazem, hydrochlorothiazide, and propranolol. His is currently taking isosorbide dinitrate, hydrochlorothiazide, and aspirin. On physical examination, his blood pressure is 151/91 mm Hg, and his pulse is 67/min. There is a right carotid bruit. His lower extremities are slightly cool to the touch and have diminished pulses at the dorsalis pedis.', '72-year-old man with calf pain while walking uphill. History of ischemic heart disease and worsening hypertension despite medication compliance. On physical exam he has a right carotid bruit and his lower extremities are cool, with diminished dorsalis pedis pulses.'),
        })
        self._test_queries('pmc/v1/trec-cds-2015', count=30, items={
            0: TrecCdsQuery('1', 'diagnosis', 'A 44 yo male is brought to the emergency room after multiple bouts of vomiting that has a "coffee ground" appearance. His heart rate is 135 bpm and blood pressure is 70/40 mmHg. Physical exam findings include decreased mental status and cool extremities. He receives a rapid infusion of crystalloid solution followed by packed red blood cell transfusion and is admitted to the ICU for further care.', 'A 44-year-old man with coffee-ground emesis, tachycardia, hypoxia,  hypotension and  cool, clammy extremities.'),
            9: TrecCdsQuery('10', 'diagnosis', 'A 38 year old woman complains of severe premenstrual and menstrual pelvic pain, heavy, irregular periods and occasional spotting between periods. Past medical history remarkable for  two years of infertility treatment and an ectopic pregnancy at age 26.', 'A 38 year old woman with severe dysmenorrhea, menorrhagia, and menometrorrhagia. PMH of infertility treatment and ectopic pregnancy'),
            29: TrecCdsQuery('30', 'treatment', 'A 47 year old male who fell on his outstretched left arm presents with pain and bruising on the inside and outside of the elbow, swelling, and inability to bend the arm. On the x-ray, the ulna has dislocated posteriorly from the trochlea of the humerus. The radius has dislocated from the capitulum of the humerus.', 'A 47 year old male who fell on his outstretched left arm presents with pain, swelling, and inability to bend the arm. The x-ray, shows posterior elbow dislocation.'),
        })
        self._test_queries('pmc/v2/trec-cds-2016', count=30, items={
            0: TrecCds2016Query('1', 'diagnosis', '\n   78 M w/ pmh of CABG in early [**Month (only) 3**] at [**Hospital6 4406**]\n   (transferred to nursing home for rehab on [**12-8**] after several falls out\n   of bed.) He was then readmitted to [**Hospital6 1749**] on\n   [**3120-12-11**] after developing acute pulmonary edema/CHF/unresponsiveness?.\n   There was a question whether he had a small MI; he reportedly had a\n   small NQWMI. He improved with diuresis and was not intubated. \n   .\n   Yesterday, he was noted to have a melanotic stool earlier this evening\n   and then approximately 9 loose BM w/ some melena and some frank blood\n   just prior to transfer, unclear quantity.\n    ', '78 M transferred to nursing home for rehab after CABG. Reportedly readmitted with a small NQWMI. Yesterday, he was noted to have a melanotic stool and then today he had approximately 9 loose BM w/ some melena and some frank blood just prior to transfer, unclear quantity.', 'A 78 year old male presents with frequent stools and melena.'),
            9: TrecCds2016Query('10', 'diagnosis', '\n   The patient is a 55-year-old woman with hepatic sarcoidosis and\n   regenerative hyperplasia s/p TIPS [**10/3245**] placed [**1-27**] variceal bleeding\n   and portal hypertensive gastropathy s/p TIPS re-do with angioplasty and\n   portal vein embolectomy, who was brought to the ED by her husband for\n   evaluation after he noted worsening asterixis. While in the waiting room \n   the pt became more combative and then unresponsive. \n\n   In the ED: VS - Temp 97.9F, HR 115, BP 122/80, R 18, O2-sat 98% 2L NC.\n   She was unresponsive but able to protect her airway and so not\n   intubated. She vomited x1 and received Zofran as well as 1.5 L NS. Labs\n   were significant for K 5.5, BUN 46, Cr 2.2 (up from baseline of 0.8),\n   and ammonia of 280. Stool was Guaiac negative. A urinalysis and CXR\n   were done and are pending, and a FAST revealed\n   hepatosplenomegaly but no intraperitoneal fluid. \n   \n   On arrival to the ICU the pt had another episode of emesis. NGT was\n   placed to suction and 1.5L bilious material was drained.\n   Allergies:\n   Cipro (Oral) (Ciprofloxacin Hcl)\n   Hives;\n   Doxycycline\n   Hives; hallucin\n   Paxil (Oral) (Paroxetine Hcl)\n   hair loss;\n   Quinine\n   Rash;\n   Compazine (Injection) (Prochlorperazine Edisylate)\n   muscle spasm;\n   Levaquin (Oral) (Levofloxacin)\n   tendinitis of t\n   Lithium\n   Hives;\n    ', 'A 55y/o F with sarcoidosis, COPD, idiopathic cardiomyopathy with EF 40% and diastolic dysfunction, varices s/p TIPS and hypothyroidism presenting today with confusion. She was brought to the ED by her husband for evaluation after he noted worsening asterixis.  While in the waiting room the pt became more combative and then unresponsive.  In the ED: VS - Temp 97.9F, HR 115, BP 122/80, R 18, O2-sat 98% 2L NC.  She was unresponsive but able to protect her airway and so not intubated. She vomited x1 and received Zofran as well as 1.5 L NS. Labs were significant for K 5.5, BUN 46, Cr 2.2 (up from baseline of 0.8), and ammonia of 280. Stool was Guaiac negative. A urinalysis and CXR were done and are pending, and a FAST revealed hepatosplenomegaly but no intraperitoneal fluid.', 'A 55-year-old woman with sarcoidosis, presenting today with confusion and worsening asterixis.   In the waiting room, the pt became more combative and then unresponsive. Ammonia level 280 on admission.'),
            29: TrecCds2016Query('30', 'treatment', '\n   85 y/o F with PMHx of HTN, HL, h/o breast CA and 3cm renal pelvis\n   transitional cell tumor who presented for nephrectomy on [**2575-8-15**].  Her\n   post op course was complicated by agitation thought due to narcotics.\n   Today, she was restarted on her home meds and while on telemetry, pt\n   was noted to be bradycardic to 40s.  Pt was triggered for SBP of 70 and\n   HR of 40 during which she remained asymptomatic.  She was given 1L IVF\n   and her HR/BP trended back up to baseline.  However, there was a second\n   event an hour later when she sat up and became bradycardic in the 30s\n   with associated hypotension. Second episode occurred with position change\n   and again, pt developped junctional rhythm in 30s. \n   \n   home meds:\n   Verapamil 240mg daily\n   Lisinopril 5mg\n   Rosuvastatin 10mg\n   Meclizine 25 TID PRN\n   Imipramine 25 QHS\n   Colace 100mg \n   Loratidine 10mg daily\n\n   Physical Examination\n   T: 98 BP: 111/47 P: 74 R: 16 O2: 98% on 2L NC\n   General: oriented to person only, NAD, comfortable\n   HEENT: Sclera anicteric, dry MM, oropharynx clear\n   Neck: supple, unable to appreciate JVP due to habitus\n   Lungs: poor effort but [**Month (only) 199**] BS at bases and some audible airway\n   secretion in upper airways\n   CV: Regular rate and rhythm, no m/r/g, diff to auscult [**2-13**] habitus\n   Abdomen:  diffusely tender, bowel sounds present, multiple surgical\n   incisions, clean dry and intact, abd binder in place\n   GU: foley in place\n   Ext: cool, no edema, 1+ pulses, pneumoboots in place\n    ', '85 y/o F with PMHx of HTN, HL, h/o breast CA and 3cm renal pelvis transitional cell tumor who presented for nephrectomy.  Her post op course was complicated by agitation thought due to narcotics.  Today, she was restarted on her home meds and while on telemetry, pt was noted to be bradycardic to 40s.  Pt was triggered for SBP of 70 and HR of 40 during which she remained asymptomatic.  She was given 1L IVF and her HR/BP trended back up to baseline.  However, there was a second event an hour later when she sat up and became bradycardic in the 30s with associated hypotension. Second episode occurred with position change and again, pt developped junctional rhythm in 30s.', 'An 85 year-old woman on  verapamil  presents with junctional heart rhythm in 30s with associated hypotension.'),
        })

    def test_pmc_qrels(self):
        self._test_qrels('pmc/v1/trec-cds-2014', count=37949, items={
            0: TrecQrel('1', '1033658', 0, '0'),
            9: TrecQrel('1', '1037001', 0, '0'),
            37948: TrecQrel('30', '80153', 0, '0'),
        })
        self._test_qrels('pmc/v1/trec-cds-2015', count=37807, items={
            0: TrecQrel('1', '1065003', 1, '0'),
            9: TrecQrel('1', '117132', 0, '0'),
            37806: TrecQrel('30', '64646', 2, '0'),
        })
        self._test_qrels('pmc/v2/trec-cds-2016', count=37707, items={
            0: TrecQrel('1', '1036067', 0, '0'),
            9: TrecQrel('1', '1160569', 0, '0'),
            37706: TrecQrel('30', '65042', 0, '0'),
        })


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