import unittest
import warnings
import nose
from nose.tools import assert_equal
from datetime import datetime

import numpy as np
import pandas as pd
from pandas import DataFrame
from pandas.io import data as web
from pandas.io.data import DataReader, SymbolWarning
from pandas.util.testing import (assert_series_equal, assert_produces_warning,
                                 network, assert_frame_equal)
from numpy.testing import assert_array_equal


def assert_n_failed_equals_n_null_columns(wngs, obj, cls=SymbolWarning):
    all_nan_cols = pd.Series(dict((k, pd.isnull(v).all()) for k, v in
                                  obj.iteritems()))
    n_all_nan_cols = all_nan_cols.sum()
    valid_warnings = pd.Series([wng for wng in wngs if isinstance(wng, cls)])
    assert_equal(len(valid_warnings), n_all_nan_cols)
    failed_symbols = all_nan_cols[all_nan_cols].index
    msgs = valid_warnings.map(lambda x: x.message)
    assert msgs.str.contains('|'.join(failed_symbols)).all()


class TestGoogle(unittest.TestCase):

    @network
    def test_google(self):
        # asserts that google is minimally working and that it throws
        # an exception when DataReader can't get a 200 response from
        # google
        start = datetime(2010, 1, 1)
        end = datetime(2013, 01, 27)

        self.assertEquals(
            web.DataReader("F", 'google', start, end)['Close'][-1],
            13.68)

        self.assertRaises(Exception, web.DataReader, "NON EXISTENT TICKER",
                          'google', start, end)

    @network
    def test_get_quote_fails(self):
        self.assertRaises(NotImplementedError, web.get_quote_google,
                          pd.Series(['GOOG', 'AAPL', 'GOOG']))

    @network
    def test_get_goog_volume(self):
        df = web.get_data_google('GOOG')
        self.assertEqual(df.Volume.ix['OCT-08-2010'], 2863473)

    @network
    def test_get_multi1(self):
        sl = ['AAPL', 'AMZN', 'GOOG']
        pan = web.get_data_google(sl, '2012')

        def testit():
            ts = pan.Close.GOOG.index[pan.Close.AAPL > pan.Close.GOOG]
            self.assertEquals(ts[0].dayofyear, 96)

        if (hasattr(pan, 'Close') and hasattr(pan.Close, 'GOOG') and
            hasattr(pan.Close, 'AAPL')):
            testit()
        else:
            self.assertRaises(AttributeError, testit)

    @network
    def test_get_multi2(self):
        with warnings.catch_warnings(record=True) as w:
            pan = web.get_data_google(['GE', 'MSFT', 'INTC'], 'JAN-01-12',
                                      'JAN-31-12')
            result = pan.Close.ix['01-18-12']
            assert_n_failed_equals_n_null_columns(w, result)

            # sanity checking

            assert np.issubdtype(result.dtype, np.floating)
            result = pan.Open.ix['Jan-15-12':'Jan-20-12']
            self.assertEqual((4, 3), result.shape)
            assert_n_failed_equals_n_null_columns(w, result)


class TestYahoo(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        try:
            import lxml
        except ImportError:
            raise nose.SkipTest

    @network
    def test_yahoo(self):
        # asserts that yahoo is minimally working and that it throws
        # an exception when DataReader can't get a 200 response from
        # yahoo
        start = datetime(2010, 1, 1)
        end = datetime(2013, 01, 27)

        self.assertEquals( web.DataReader("F", 'yahoo', start,
                                          end)['Close'][-1], 13.68)

    @network
    def test_yahoo_fails(self):
        start = datetime(2010, 1, 1)
        end = datetime(2013, 01, 27)
        self.assertRaises(Exception, web.DataReader, "NON EXISTENT TICKER",
                          'yahoo', start, end)

    @network
    def test_get_quote_series(self):
        df = web.get_quote_yahoo(pd.Series(['GOOG', 'AAPL', 'GOOG']))
        assert_series_equal(df.ix[0], df.ix[2])

    @network
    def test_get_quote_string(self):
        df = web.get_quote_yahoo('GOOG')

    @network
    def test_get_quote_stringlist(self):
        df = web.get_quote_yahoo(['GOOG', 'AAPL', 'GOOG'])
        assert_series_equal(df.ix[0], df.ix[2])

    @network
    def test_get_components_dow_jones(self):
        df = web.get_components_yahoo('^DJI') #Dow Jones
        assert isinstance(df, pd.DataFrame)
        self.assertEqual(len(df), 30)

    @network
    def test_get_components_dax(self):
        df = web.get_components_yahoo('^GDAXI') #DAX
        assert isinstance(df, pd.DataFrame)
        self.assertEqual(len(df), 30)
        self.assertEqual(df[df.name.str.contains('adidas', case=False)].index,
                         'ADS.DE')

    @network
    def test_get_components_nasdaq_100(self):
        """as of 7/12/13 the conditional will test false because the link is
        invalid"""
        df = web.get_components_yahoo('^NDX') #NASDAQ-100
        assert isinstance(df, pd.DataFrame)

        if len(df) > 1:
            # Usual culprits, should be around for a while
            assert 'AAPL' in df.index
            assert 'GOOG' in df.index
            assert 'AMZN' in df.index
        else:
            expected = DataFrame({'exchange': 'N/A', 'name': '@^NDX'},
                                 index=['@^NDX'])
            assert_frame_equal(df, expected)

    @network
    def test_get_data_single_symbol(self):
        #single symbol
        #http://finance.yahoo.com/q/hp?s=GOOG&a=09&b=08&c=2010&d=09&e=10&f=2010&g=d
        df = web.get_data_yahoo('GOOG')
        self.assertEqual(df.Volume.ix['OCT-08-2010'], 2859200)

    @network
    def test_get_data_multiple_symbols(self):
        sl = ['AAPL', 'AMZN', 'GOOG']
        pan = web.get_data_yahoo(sl, '2012')

        def testit():
            ts = pan.Close.GOOG.index[pan.Close.AAPL > pan.Close.GOOG]
            self.assertEquals(ts[0].dayofyear, 96)

        if hasattr(pan.Close, 'GOOG') and hasattr(pan.Close, 'AAPL'):
            testit()
        else:
            self.assertRaises(AttributeError, testit)

    @network
    def test_get_data_multiple_symbols_two_dates(self):
        pan = web.get_data_yahoo(['GE', 'MSFT', 'INTC'], 'JAN-01-12', 'JAN-31-12')
        result = pan.Close.ix['01-18-12']
        self.assertEqual(len(result), 3)

        # sanity checking
        assert np.issubdtype(result.dtype, np.floating)

        expected = np.array([[ 18.99,  28.4 ,  25.18],
                             [ 18.58,  28.31,  25.13],
                             [ 19.03,  28.16,  25.52],
                             [ 18.81,  28.82,  25.87]])
        result = pan.Open.ix['Jan-15-12':'Jan-20-12']
        self.assertEqual(expected.shape, result.shape)

    @network
    def test_get_date_ret_index(self):
        pan = web.get_data_yahoo(['GE', 'INTC', 'IBM'], '1977', '1987',
                                 ret_index=True)
        self.assert_(hasattr(pan, 'Ret_Index'))
        if hasattr(pan, 'Ret_Index') and hasattr(pan.Ret_Index, 'INTC'):
            tstamp = pan.Ret_Index.INTC.first_valid_index()
            result = pan.Ret_Index.ix[tstamp]['INTC']
            self.assertEqual(result, 1.0)

        # sanity checking
        assert np.issubdtype(pan.values.dtype, np.floating)


class TestYahooOptions(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        try:
            import lxml
        except ImportError:
            raise nose.SkipTest

        # aapl has monthlies
        cls.aapl = web.Options('aapl', 'yahoo')
        today = datetime.today()
        year = today.year
        month = today.month + 1
        if month > 12:
            year = year + 1
            month = 1
        cls.expiry = datetime(year, month, 1)

    @classmethod
    def tearDownClass(cls):
        del cls.aapl, cls.expiry

    @network
    def test_get_options_data(self):
        try:
            calls, puts = self.aapl.get_options_data(expiry=self.expiry)
        except IndexError:
            warnings.warn("IndexError thrown no tables found")
        else:
            assert len(calls)>1
            assert len(puts)>1

    @network
    def test_get_near_stock_price(self):
        try:
            calls, puts = self.aapl.get_near_stock_price(call=True, put=True,
                                                        expiry=self.expiry)
        except IndexError:
            warnings.warn("IndexError thrown no tables found")
        else:
            self.assertEqual(len(calls), 5)
            self.assertEqual(len(puts), 5)

    @network
    def test_get_call_data(self):
        try:
            calls = self.aapl.get_call_data(expiry=self.expiry)
        except IndexError:
            warnings.warn("IndexError thrown no tables found")
        else:
            assert len(calls)>1

    @network
    def test_get_put_data(self):
        try:
            puts = self.aapl.get_put_data(expiry=self.expiry)
        except IndexError:
            warnings.warn("IndexError thrown no tables found")
        else:
            assert len(puts)>1


class TestOptionsWarnings(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        try:
            import lxml
        except ImportError:
            raise nose.SkipTest

        with assert_produces_warning():
            cls.aapl = web.Options('aapl')

        today = datetime.today()
        cls.year = today.year
        cls.month = today.month + 1
        if cls.month > 12:
            cls.year += 1
            cls.month = 1

    @classmethod
    def tearDownClass(cls):
        del cls.aapl, cls.year, cls.month

    @network
    def test_get_options_data_warning(self):
        with assert_produces_warning():
            print('month: {0}, year: {1}'.format(self.month, self.year))
            try:
                self.aapl.get_options_data(month=self.month, year=self.year)
            except IndexError:
                warnings.warn("IndexError thrown no tables found")

    @network
    def test_get_near_stock_price_warning(self):
        with assert_produces_warning():
            print('month: {0}, year: {1}'.format(self.month, self.year))
            try:
                calls_near, puts_near = self.aapl.get_near_stock_price(call=True,
                                                                    put=True,
                                                                    month=self.month,
                                                                    year=self.year)
            except IndexError:
                warnings.warn("IndexError thrown no tables found")

    @network
    def test_get_call_data_warning(self):
        with assert_produces_warning():
            print('month: {0}, year: {1}'.format(self.month, self.year))
            try:
                self.aapl.get_call_data(month=self.month, year=self.year)
            except IndexError:
                warnings.warn("IndexError thrown no tables found")

    @network
    def test_get_put_data_warning(self):
        with assert_produces_warning():
            print('month: {0}, year: {1}'.format(self.month, self.year))
            try:
                self.aapl.get_put_data(month=self.month, year=self.year)
            except IndexError:
                warnings.warn("IndexError thrown no tables found")


class TestDataReader(unittest.TestCase):
    @network
    def test_read_yahoo(self):
        gs = DataReader("GS", "yahoo")
        assert isinstance(gs, DataFrame)

    @network
    def test_read_google(self):
        gs = DataReader("GS", "google")
        assert isinstance(gs, DataFrame)

    @network
    def test_read_fred(self):
        vix = DataReader("VIXCLS", "fred")
        assert isinstance(vix, DataFrame)

    @network
    def test_read_famafrench(self):
        for name in ("F-F_Research_Data_Factors",
                     "F-F_Research_Data_Factors_weekly", "6_Portfolios_2x3",
                     "F-F_ST_Reversal_Factor"):
            ff = DataReader(name, "famafrench")
            assert ff
            assert isinstance(ff, dict)


class TestFred(unittest.TestCase):
    @network
    def test_fred(self):
        """
        Throws an exception when DataReader can't get a 200 response from
        FRED.
        """
        start = datetime(2010, 1, 1)
        end = datetime(2013, 01, 27)

        self.assertEquals(
            web.DataReader("GDP", "fred", start, end)['GDP'].tail(1),
            15984.1)

        self.assertRaises(Exception, web.DataReader, "NON EXISTENT SERIES",
                          'fred', start, end)

    @network
    def test_fred_nan(self):
        start = datetime(2010, 1, 1)
        end = datetime(2013, 01, 27)
        df = web.DataReader("DFII5", "fred", start, end)
        assert pd.isnull(df.ix['2010-01-01'])

    @network
    def test_fred_parts(self):
        start = datetime(2010, 1, 1)
        end = datetime(2013, 01, 27)
        df = web.get_data_fred("CPIAUCSL", start, end)
        self.assertEqual(df.ix['2010-05-01'], 217.23)

        t = df.CPIAUCSL.values
        assert np.issubdtype(t.dtype, np.floating)
        self.assertEqual(t.shape, (37,))

    @network
    def test_fred_part2(self):
        expected = [[576.7],
                    [962.9],
                    [684.7],
                    [848.3],
                    [933.3]]
        result = web.get_data_fred("A09024USA144NNBR", start="1915").ix[:5]
        assert_array_equal(result.values, np.array(expected))

    @network
    def test_invalid_series(self):
        name = "NOT A REAL SERIES"
        self.assertRaises(Exception, web.get_data_fred, name)


if __name__ == '__main__':
    nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],
                   exit=False)
