import re
import unittest
from ir_datasets.datasets.codec import CodecQuery, CodecDoc
from ir_datasets.formats import TrecQrel
from .base import DatasetIntegrationTest


class TestCodec(DatasetIntegrationTest):
    def test_docs(self):
        self._test_docs('codec', count=729824, items={
            0: CodecDoc('5f8fddcc5a1a25d5fce9ec5e331ba8dd', "Apple's Tim Cook calls for calm heads on China, US trade", re.compile("^Apple's Tim Cook calls for calm heads on China, US trade\nApple CEO Tim Cook\\.Getty Images\nApple Chief.{2099}d last year when he called for China to increase trade and continue opening itself up to the world\\.\n$", flags=48), 'https://www.cnbc.com/2018/03/24/apples-tim-cook-calls-for-calm-heads-on-china-us-trade.html'),
            9: CodecDoc('ac2c854a75425587461307c506c48330', 'Buying the yen to protect against a trade war is ‘pretty peculiar,’ wealth manager says', re.compile('^Buying the yen to protect against a trade war is ‘pretty peculiar,’ wealth manager says\nVIDEO3:5803:.{2303}d that it would look to be exempted from the metal tariffs, though they came into effect on Friday\\.\n$', flags=48), 'https://www.cnbc.com/2018/03/26/yen-to-protect-against-a-trade-war-is-wrong.html?__source=twitter%7Cmain'),
            729823: CodecDoc('2dc726c7c5de68b9766e35336fa214a0', 'Damped vibration', re.compile('^Damped vibration\nThe simple harmonic oscillations discussed above continue forever, at constant ampl.{103}tors behave somewhat differently, however\\. Harmonic oscillations tend to die away as time goes on\\.…\n$', flags=48), 'https://www.britannica.com/science/damped-vibration'),
        })

    def test_queries(self):
        self._test_queries('codec', count=42, items={
            0: CodecQuery('economics-1', "How has the UK's Open Banking Regulation benefited challenger banks?", 'finance', 'UK’s Open Banking regulation, which has parallels to the EU’s second payment service directive (PSD2), went live in January 2018. This piece of legislation “will require banks to open their payments infrastructure and customer data assets to third parties”. As a result, banks no longer have a monopoly on user data if clients grant permission. \n\nChallenger banks are small, recently created retail banks that compete directly with the longer-established banks in the UK. Specifically, seeking market share from the "big four" UK retail banks (Barclays, HSBC, Lloyds Banking Group, and NatWest Group). The banks distinguish themselves from the historic banks by modern financial technology practices, such as online-only operations, that avoid the costs and complexities of traditional banking. The largest UK-operating challenger banks include Atom Bank, Revolut, Starling, N26, and Tide.\n\nRelevant documents and entities will discuss how challenger banks have used open banking to develop new products or capture market share from traditional retail banks in the UK.'),
            9: CodecQuery('economics-18', 'Was the crash that followed the dot-com bubble an overreaction considering the ultimate success of the internet?', 'finance', 'The dot-com buddle from 1995-2000 saw incredible growth in stocks that were considered to have anything to do with the internet. NASDAQ rose 400% only to fall 78% from its peak by 2002. Several high-profile failures included Pets.com, Boo.com, Woldcom, and Global Crossing. However, many internet-based companies survived and went on to thrive, i.e. eBay, Amazon, Qualcomm, Cisco Systems. Venture Capital pulled back from the internet space for a period.\n\nHowever, with FANNG stocks and incredible software/internet-based companies performance from the mid-2000s until the present day, it could be argued that the direction of the dot-com bubble was absolutely correct. Nonetheless, the magnitude of equity growth and capital allocation to companies with limited commercial value was irrational and required a market correction.'),
            41: CodecQuery('politics-23', 'Is the rise of European populism a threat to the European Union?', 'politics', '"Populist" is a broad term that describes, typically a politician, who targets people who feel that established elite groups disregard their concerns. Some critics highlight negative connotations, including criticising foreign migration or minorities.\n\nIn recent years, populism has been stronger in Eastern Europe, i.e. Bulgaria, Hungary, Austrian, and Poland, etc. have seen the rise of populist politicians. Right-wing populist movements have also gained momentum in France, Spain, the United Kingdom and other parts of Europe. In Hungary and Poland, some critics argue this has led to an erosion of the rule of law, increased persecution, and authoritarianism.\n\nEuropean Union is a political union of democratic nations. However, radical right-wing politics reject what the EU stands for and how it works, i.e. against European supranational integration and push for national policies. Populists also criticise the EU\'s perceived bureaucracy and failures - common arguments during the Brexit Leave campaign. The EU is founded based on shared democratic values that countries need to be relatively ideologically aligned to function within a political union. There is also the threat that disillusioned Eastern European countries will turn away from the EU and toward Russia.'),
        })
        self._test_queries('codec/economics', count=14, items={
            0: CodecQuery('economics-1', "How has the UK's Open Banking Regulation benefited challenger banks?", 'finance', 'UK’s Open Banking regulation, which has parallels to the EU’s second payment service directive (PSD2), went live in January 2018. This piece of legislation “will require banks to open their payments infrastructure and customer data assets to third parties”. As a result, banks no longer have a monopoly on user data if clients grant permission. \n\nChallenger banks are small, recently created retail banks that compete directly with the longer-established banks in the UK. Specifically, seeking market share from the "big four" UK retail banks (Barclays, HSBC, Lloyds Banking Group, and NatWest Group). The banks distinguish themselves from the historic banks by modern financial technology practices, such as online-only operations, that avoid the costs and complexities of traditional banking. The largest UK-operating challenger banks include Atom Bank, Revolut, Starling, N26, and Tide.\n\nRelevant documents and entities will discuss how challenger banks have used open banking to develop new products or capture market share from traditional retail banks in the UK.'),
            9: CodecQuery('economics-18', 'Was the crash that followed the dot-com bubble an overreaction considering the ultimate success of the internet?', 'finance', 'The dot-com buddle from 1995-2000 saw incredible growth in stocks that were considered to have anything to do with the internet. NASDAQ rose 400% only to fall 78% from its peak by 2002. Several high-profile failures included Pets.com, Boo.com, Woldcom, and Global Crossing. However, many internet-based companies survived and went on to thrive, i.e. eBay, Amazon, Qualcomm, Cisco Systems. Venture Capital pulled back from the internet space for a period.\n\nHowever, with FANNG stocks and incredible software/internet-based companies performance from the mid-2000s until the present day, it could be argued that the direction of the dot-com bubble was absolutely correct. Nonetheless, the magnitude of equity growth and capital allocation to companies with limited commercial value was irrational and required a market correction.'),
            13: CodecQuery('economics-23', 'Offering non-accounting services arguably creates a conflict of interest for the Big Four. Is this the reason for their inability to uncover recent financial scandals?', 'finance', 'The Big Four are the four largest global accounting firms that dominate corporate accounting, i.e. Deloitte, Ernst & Young (EY), PricewaterhouseCoopers (PwC), and Klynveld Peat Marwick Goerdeler (KPMG). As well as offering accounting services, these firms also over other many higher-margin services such as tax, consultancy, and technology services (~80% revenue). \n\nInternal controls (firewalls, departments, etc.) are meant to prohibit any cooperation between audit and non-audit services in winning customer contracts or favour. However, critics would argue that this does not go far enough and can lead to less rigorous audits, calling for the Big Four to be broken up. Critics highlight poor practices in recent accounting scandals, including  Wirecard, Carillion, Satyam Computer Services, South Africa examples, etc. \n\nHowever, it should be highlighted that even if the Big Four were separated, this would not necessarily lead to no accounting scandals. Some fraudulent practices may be hard to catch even if objective and following best practices. Also, there are other areas of possible conflict of interest that non-accounting work, including how the corporate hires the auditor.'),
        })
        self._test_queries('codec/history', count=14, items={
            0: CodecQuery('history-1', 'Would the United Kingdom have been ready for WWII without the time gained through Appeasement?', 'history', "Many argue Britain's army was depleted in the early 1930s and stretched across the globe. UK defence spending had fallen significantly during the 1920s, from over £700 million in 1919 to 100 million in 1931.\n\nBetween 1934 and 1939, the UK launched a substantial programme of re-arming, recognising that war with Hitler was becoming increasingly likely. Although Appeasement was also motivated by Chamberlain's desire to end war, some argue this meant that the UK was more prepared in 1939 when war eventually broke out.  \n\nDespite these efforts, Germany was still better prepared for war under Hilter's single-minded preparation since he came to power in 1933. However, without Appeasement, the differential might have been much worse."),
            9: CodecQuery('history-19', 'How close did the world come to nuclear war during the Cuban Missile Crisis?', 'history', 'During the Cuban Missile Crisis, leaders of the United States and the Soviet Union engaged in a 13-day political and military standoff in October 1962 over the installation of nuclear-armed Soviet missiles on Cuba. This was the peak of the Cold War and a high-stakes political and military situation, given the potential devastation of nuclear weapons. However, exactly how close we came to nuclear armageddon is still debated. \n\nThe political leaders within this crisis were JFK for the United States, Nikita Khrushchev of the Soviet Union, and Fidel Casto of Cuba. Many highlights that JFK and Khrushchev were measured in their leadership styles and understood nuclear war meant mutual destruction. Some highlight exchanging of letters and other communications to prevent a nuclear war. Nonetheless, many highlight the mistrust and fear between both sides and how a single false move could have led to a disaster. Both sides were actively preparing for a nuclear war, and some within each camp through nuclear strikes was likely.'),
            13: CodecQuery('history-25', 'How responsible was Rasputin for the fall of the Romanov dynasty?', 'history', "The Russian Tsar, Nicholas II, abdicated from power in 1917, bringing the 300-year-old Romanov dynasty to an end. Some historians suggest that Grigori Rasputin's scandalous reputation helped discredit the Tsarist government and helped to lay the foundation for the Russian Revolution. \n\nHowever, many historians argue that although Rasputin was a useful propaganda tool against the Tsar, much larger factors were at play. For example, Nicholas II was viewed as a weak leader, the Russo-Japanese War, Bloody Sunday, Tsarina unpopularity, and WWI. There was also significant economic issues, including inflation and food shortages."),
        })
        self._test_queries('codec/politics', count=14, items={
            0: CodecQuery('politics-1', 'Is Scottish Independence inevitable?', 'politics', "This questions focuses on the long-term political, economic and social reasoning behind whether Scotland will likely become independent. Short-term facts and opinions are less central to this question.\n\nSome argue that Scottish independence is inevitable given the surge of support towards SNP in recent decades. Labour has lost political weight in Scotland since the Blair era, and Conservatives historically struggle to penetrate a more left-wing Scottish demographic. Brexit further exacerbated this political unease, i.e. right-leaning Britain and left-leaning Scotland. While the younger demographic is more likely to be independence supporters and the older demographic is more likely to be pro-union.\n\nThere are several arguments that Scotland independence is not inevitable. Historically there has been a union for 300+ years covering a full range of circumstances (world wars, successes). Economically, Britain funds lots of Scottish spending through Barnett-based funding, Scotland's oil is less valuable than previously, and there is much economic uncertainly around currency and debt. Some also argue that Scotland has large political independence due to devolution; thus, independence is unnecessary."),
            9: CodecQuery('politics-16', 'Why did Hilary Clinton lose the 2016 US presidential election?', 'politics', 'Hillary Clinton lost the 2016 US presidential election to Donald Trump in 2016. Political commentators highlight many reasons for Clinton\'s loss. \n\nFor example, Donald Trump managed to craft a strong populist message that resonated with many voters who were disenfranchised with current politics, particularly with the "political elite" who some felt Clinton represented. The Democratats were also divided, specifically far-left factions led by Bernie Sanders. Political gridlock under an Obama Administration that, rightly or wrongly, some felt the public wanting change. Hilary also had relatively low personal approval ratings, which were not helped by the FBI investigation into her use of email. Some also point to external factors, including Russian interference.'),
            13: CodecQuery('politics-23', 'Is the rise of European populism a threat to the European Union?', 'politics', '"Populist" is a broad term that describes, typically a politician, who targets people who feel that established elite groups disregard their concerns. Some critics highlight negative connotations, including criticising foreign migration or minorities.\n\nIn recent years, populism has been stronger in Eastern Europe, i.e. Bulgaria, Hungary, Austrian, and Poland, etc. have seen the rise of populist politicians. Right-wing populist movements have also gained momentum in France, Spain, the United Kingdom and other parts of Europe. In Hungary and Poland, some critics argue this has led to an erosion of the rule of law, increased persecution, and authoritarianism.\n\nEuropean Union is a political union of democratic nations. However, radical right-wing politics reject what the EU stands for and how it works, i.e. against European supranational integration and push for national policies. Populists also criticise the EU\'s perceived bureaucracy and failures - common arguments during the Brexit Leave campaign. The EU is founded based on shared democratic values that countries need to be relatively ideologically aligned to function within a political union. There is also the threat that disillusioned Eastern European countries will turn away from the EU and toward Russia.'),
        })

    def test_qrels(self):
        self._test_qrels('codec', count=6186, items={
            0: TrecQrel('economics-8', '089a22846f6ba15fb4ef4cca0a884dd4', 2, 'Q0'),
            9: TrecQrel('economics-8', '24de81ea95c7df32941e8bd200d3528a', 2, 'Q0'),
            6185: TrecQrel('politics-23', 'f7975664f7fc45416c7256d12d06fbdf', 1, 'Q0'),
        })
        self._test_qrels('codec/economics', count=1970, items={
            0: TrecQrel('economics-8', '089a22846f6ba15fb4ef4cca0a884dd4', 2, 'Q0'),
            9: TrecQrel('economics-8', '24de81ea95c7df32941e8bd200d3528a', 2, 'Q0'),
            1969: TrecQrel('economics-23', 'ebd36155b9808933bbbfc26af6d18dec', 1, 'Q0'),
        })
        self._test_qrels('codec/history', count=2024, items={
            0: TrecQrel('history-20', '00aa648a657bdf73369bcb093030cc41', 0, 'Q0'),
            9: TrecQrel('history-20', '0d239d8dea605cd079d1a0144aa6ed46', 1, 'Q0'),
            2023: TrecQrel('history-25', 'dba9433750c2505f2a6a69661d4eb2fd', 0, 'Q0'),
        })
        self._test_qrels('codec/politics', count=2192, items={
            0: TrecQrel('politics-12', '02d8176599da0bfc15caaf7e0b3bba6b', 3, 'Q0'),
            9: TrecQrel('politics-12', '0e23e45de303e5f440a5a17ccd7974ec', 1, 'Q0'),
            2191: TrecQrel('politics-23', 'f7975664f7fc45416c7256d12d06fbdf', 1, 'Q0'),
        })


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