class TestChangeOfSids(TestCase):
    def setUp(self):
        self.sids = range(90)
        self.env = TradingEnvironment()
        self.env.write_data(equities_identifiers=self.sids)

        self.sim_params = factory.create_simulation_parameters(
            start=datetime(1990, 1, 1, tzinfo=pytz.utc),
            end=datetime(1990, 1, 8, tzinfo=pytz.utc),
            env=self.env,
        )

    def test_all_sids_passed(self):
        algo = BatchTransformAlgorithmSetSid(
            sim_params=self.sim_params,
            env=self.env,
        )
        source = DifferentSidSource()
        algo.run(source)
        for i, (df, date) in enumerate(zip(algo.history, source.trading_days)):
            self.assertEqual(df.index[-1], date, "Newest event doesn't \
                             match.")

            for sid in self.sids[:i]:
                self.assertIn(sid, df.columns)

            self.assertEqual(df.iloc[-1].iloc[-1], i)

    def test_compute_lifetimes(self):
        num_assets = 4
        env = TradingEnvironment()
        trading_day = env.trading_day
        first_start = pd.Timestamp('2015-04-01', tz='UTC')

        frame = make_rotating_asset_info(
            num_assets=num_assets,
            first_start=first_start,
            frequency=env.trading_day,
            periods_between_starts=3,
            asset_lifetime=5
        )

        env.write_data(equities_df=frame)
        finder = env.asset_finder

        all_dates = pd.date_range(
            start=first_start,
            end=frame.end_date.max(),
            freq=trading_day,
        )

        for dates in all_subindices(all_dates):
            expected_with_start_raw = full(
                shape=(len(dates), num_assets),
                fill_value=False,
                dtype=bool,
            )
            expected_no_start_raw = full(
                shape=(len(dates), num_assets),
                fill_value=False,
                dtype=bool,
            )

            for i, date in enumerate(dates):
                it = frame[['start_date', 'end_date']].itertuples()
                for j, start, end in it:
                    # This way of doing the checks is redundant, but very
                    # clear.
                    if start <= date <= end:
                        expected_with_start_raw[i, j] = True
                        if start < date:
                            expected_no_start_raw[i, j] = True

            expected_with_start = pd.DataFrame(
                data=expected_with_start_raw,
                index=dates,
                columns=frame.index.values,
            )
            result = finder.lifetimes(dates, include_start_date=True)
            assert_frame_equal(result, expected_with_start)

            expected_no_start = pd.DataFrame(
                data=expected_no_start_raw,
                index=dates,
                columns=frame.index.values,
            )
            result = finder.lifetimes(dates, include_start_date=False)
            assert_frame_equal(result, expected_no_start)

def create_test_df_source(sim_params=None, env=None, bars='daily'):
    if bars == 'daily':
        freq = pd.datetools.BDay()
    elif bars == 'minute':
        freq = pd.datetools.Minute()
    else:
        raise ValueError('%s bars not understood.' % bars)

    if sim_params and bars == 'daily':
        index = sim_params.trading_days
    else:
        if env is None:
            env = TradingEnvironment()

        start = pd.datetime(1990, 1, 3, 0, 0, 0, 0, pytz.utc)
        end = pd.datetime(1990, 1, 8, 0, 0, 0, 0, pytz.utc)

        days = env.days_in_range(start, end)

        if bars == 'daily':
            index = days
        if bars == 'minute':
            index = pd.DatetimeIndex([], freq=freq)

            for day in days:
                day_index = env.market_minutes_for_day(day)
                index = index.append(day_index)

    x = np.arange(1, len(index) + 1)

    df = pd.DataFrame(x, index=index, columns=[0])

    return DataFrameSource(df), df

    def setUpClass(cls):
        metadata = {
            0: {
                'symbol': 'CLG06',
                'root_symbol': 'CL',
                'start_date': pd.Timestamp('2005-12-01', tz='UTC'),
                'notice_date': pd.Timestamp('2005-12-20', tz='UTC'),
                'expiration_date': pd.Timestamp('2006-01-20', tz='UTC')},
            1: {
                'root_symbol': 'CL',
                'symbol': 'CLK06',
                'start_date': pd.Timestamp('2005-12-01', tz='UTC'),
                'notice_date': pd.Timestamp('2006-03-20', tz='UTC'),
                'expiration_date': pd.Timestamp('2006-04-20', tz='UTC')},
            2: {
                'symbol': 'CLQ06',
                'root_symbol': 'CL',
                'start_date': pd.Timestamp('2005-12-01', tz='UTC'),
                'notice_date': pd.Timestamp('2006-06-20', tz='UTC'),
                'expiration_date': pd.Timestamp('2006-07-20', tz='UTC')},
            3: {
                'symbol': 'CLX06',
                'root_symbol': 'CL',
                'start_date': pd.Timestamp('2006-02-01', tz='UTC'),
                'notice_date': pd.Timestamp('2006-09-20', tz='UTC'),
                'expiration_date': pd.Timestamp('2006-10-20', tz='UTC')}
        }

        env = TradingEnvironment(load=noop_load)
        env.write_data(futures_data=metadata)
        cls.asset_finder = env.asset_finder

def create_test_panel_source(sim_params=None, env=None, source_type=None):
    start = sim_params.first_open \
        if sim_params else pd.datetime(1990, 1, 3, 0, 0, 0, 0, pytz.utc)

    end = sim_params.last_close \
        if sim_params else pd.datetime(1990, 1, 8, 0, 0, 0, 0, pytz.utc)

    if env is None:
        env = TradingEnvironment()

    index = env.days_in_range(start, end)

    price = np.arange(0, len(index))
    volume = np.ones(len(index)) * 1000

    arbitrary = np.ones(len(index))

    df = pd.DataFrame({'price': price,
                       'volume': volume,
                       'arbitrary': arbitrary},
                      index=index)
    if source_type:
        source_types = np.full(len(index), source_type)
        df['type'] = source_types

    panel = pd.Panel.from_dict({0: df})

    return DataPanelSource(panel), panel

    def setUp(self):
        self.constants = {
            # Every day, assume every stock starts at 2, goes down to 1,
            # goes up to 4, and finishes at 3.
            USEquityPricing.low: 1,
            USEquityPricing.open: 2,
            USEquityPricing.close: 3,
            USEquityPricing.high: 4,
        }
        self.asset_ids = [1, 2, 3, 4]
        self.dates = date_range('2014-01', '2014-03', freq='D', tz='UTC')
        self.loader = PrecomputedLoader(
            constants=self.constants,
            dates=self.dates,
            sids=self.asset_ids,
        )

        self.asset_info = make_simple_equity_info(
            self.asset_ids,
            start_date=self.dates[0],
            end_date=self.dates[-1],
        )
        environment = TradingEnvironment()
        environment.write_data(equities_df=self.asset_info)
        self.asset_finder = environment.asset_finder
        self.assets = self.asset_finder.retrieve_all(self.asset_ids)

    def setUpClass(cls):
        metadata = {
            0: {
                "symbol": "CLG06",
                "root_symbol": "CL",
                "start_date": pd.Timestamp("2005-12-01", tz="UTC"),
                "notice_date": pd.Timestamp("2005-12-20", tz="UTC"),
                "expiration_date": pd.Timestamp("2006-01-20", tz="UTC"),
            },
            1: {
                "root_symbol": "CL",
                "symbol": "CLK06",
                "start_date": pd.Timestamp("2005-12-01", tz="UTC"),
                "notice_date": pd.Timestamp("2006-03-20", tz="UTC"),
                "expiration_date": pd.Timestamp("2006-04-20", tz="UTC"),
            },
            2: {
                "symbol": "CLQ06",
                "root_symbol": "CL",
                "start_date": pd.Timestamp("2005-12-01", tz="UTC"),
                "notice_date": pd.Timestamp("2006-06-20", tz="UTC"),
                "expiration_date": pd.Timestamp("2006-07-20", tz="UTC"),
            },
            3: {
                "symbol": "CLX06",
                "root_symbol": "CL",
                "start_date": pd.Timestamp("2006-02-01", tz="UTC"),
                "notice_date": pd.Timestamp("2006-09-20", tz="UTC"),
                "expiration_date": pd.Timestamp("2006-10-20", tz="UTC"),
            },
        }

        env = TradingEnvironment(load=noop_load)
        env.write_data(futures_data=metadata)
        cls.asset_finder = env.asset_finder

    def test_yahoo_bars_to_panel_source(self):
        env = TradingEnvironment()
        finder = AssetFinder(env.engine)
        stocks = ['AAPL', 'GE']
        env.write_data(equities_identifiers=stocks)
        start = pd.datetime(1993, 1, 1, 0, 0, 0, 0, pytz.utc)
        end = pd.datetime(2002, 1, 1, 0, 0, 0, 0, pytz.utc)
        data = factory.load_bars_from_yahoo(stocks=stocks,
                                            indexes={},
                                            start=start,
                                            end=end)
        check_fields = ['sid', 'open', 'high', 'low', 'close',
                        'volume', 'price']

        copy_panel = data.copy()
        sids = finder.map_identifier_index_to_sids(
            data.items, data.major_axis[0]
        )
        copy_panel.items = sids
        source = DataPanelSource(copy_panel)
        for event in source:
            for check_field in check_fields:
                self.assertIn(check_field, event)
            self.assertTrue(isinstance(event['volume'], (integer_types)))
            self.assertTrue(event['sid'] in sids)

    def setUp(self):
        self.assets = [1, 2, 3]
        self.dates = date_range("2014-01", "2014-03", freq="D", tz="UTC")

        asset_info = make_simple_asset_info(self.assets, start_date=self.dates[0], end_date=self.dates[-1])
        env = TradingEnvironment()
        env.write_data(equities_df=asset_info)
        self.asset_finder = env.asset_finder

 def test_sids(self):
     # Ensure that the sids property of the AssetFinder is functioning
     env = TradingEnvironment()
     env.write_data(equities_identifiers=[1, 2, 3])
     sids = env.asset_finder.sids
     self.assertEqual(3, len(sids))
     self.assertTrue(1 in sids)
     self.assertTrue(2 in sids)
     self.assertTrue(3 in sids)

    def setUp(self):
        self.assets = [1, 2, 3]
        self.dates = date_range('2014-01-01', '2014-02-01', freq='D', tz='UTC')

        asset_info = make_simple_asset_info(
            self.assets,
            start_date=self.dates[0],
            end_date=self.dates[-1],
        )
        env = TradingEnvironment()
        env.write_data(equities_df=asset_info)
        self.asset_finder = env.asset_finder

    def setUpClass(cls):
        cls.test_data_dir = TempDirectory()
        cls.db_path = cls.test_data_dir.getpath("adjustments.db")
        all_days = TradingEnvironment().trading_days
        cls.calendar_days = all_days[all_days.slice_indexer(TEST_CALENDAR_START, TEST_CALENDAR_STOP)]
        daily_bar_reader = MockDailyBarSpotReader()
        writer = SQLiteAdjustmentWriter(cls.db_path, cls.calendar_days, daily_bar_reader)
        writer.write(SPLITS, MERGERS, DIVIDENDS)

        cls.assets = TEST_QUERY_ASSETS
        cls.asset_info = EQUITY_INFO
        cls.bcolz_writer = SyntheticDailyBarWriter(cls.asset_info, cls.calendar_days)
        cls.bcolz_path = cls.test_data_dir.getpath("equity_pricing.bcolz")
        cls.bcolz_writer.write(cls.bcolz_path, cls.calendar_days, cls.assets)

    def setUp(self):
        self.__calendar = date_range('2014', '2015', freq=trading_day)
        self.__assets = assets = Int64Index(arange(1, 20))

        # Set up env for test
        env = TradingEnvironment()
        env.write_data(
            equities_df=make_simple_asset_info(
                assets,
                self.__calendar[0],
                self.__calendar[-1],
            ))
        self.__finder = env.asset_finder
        self.__mask = self.__finder.lifetimes(self.__calendar[-10:])

    def test_trading_environment(self):
        benchmark_returns, treasury_curves = \
            factory.load_market_data()

        env = TradingEnvironment(
            benchmark_returns,
            treasury_curves,
            period_start=datetime(2008, 1, 1, tzinfo=pytz.utc),
            period_end=datetime(2008, 12, 31, tzinfo=pytz.utc),
            capital_base=100000,
        )
        #holidays taken from: http://www.nyse.com/press/1191407641943.html
        new_years = datetime(2008, 1, 1, tzinfo=pytz.utc)
        mlk_day = datetime(2008, 1, 21, tzinfo=pytz.utc)
        presidents = datetime(2008, 2, 18, tzinfo=pytz.utc)
        good_friday = datetime(2008, 3, 21, tzinfo=pytz.utc)
        memorial_day = datetime(2008, 5, 26, tzinfo=pytz.utc)
        july_4th = datetime(2008, 7, 4, tzinfo=pytz.utc)
        labor_day = datetime(2008, 9, 1, tzinfo=pytz.utc)
        tgiving = datetime(2008, 11, 27, tzinfo=pytz.utc)
        christmas = datetime(2008, 5, 25, tzinfo=pytz.utc)
        a_saturday = datetime(2008, 8, 2, tzinfo=pytz.utc)
        a_sunday = datetime(2008, 10, 12, tzinfo=pytz.utc)
        holidays = [
            new_years,
            mlk_day,
            presidents,
            good_friday,
            memorial_day,
            july_4th,
            labor_day,
            tgiving,
            christmas,
            a_saturday,
            a_sunday
        ]

        for holiday in holidays:
            self.assertTrue(not env.is_trading_day(holiday))

        first_trading_day = datetime(2008, 1, 2, tzinfo=pytz.utc)
        last_trading_day = datetime(2008, 12, 31, tzinfo=pytz.utc)
        workdays = [first_trading_day, last_trading_day]

        for workday in workdays:
            self.assertTrue(env.is_trading_day(workday))

        self.assertTrue(env.last_close.month == 12)
        self.assertTrue(env.last_close.day == 31)

    def test_compute_lifetimes(self):
        num_assets = 4
        env = TradingEnvironment()
        trading_day = env.trading_day
        first_start = pd.Timestamp('2015-04-01', tz='UTC')

        frame = make_rotating_asset_info(
            num_assets=num_assets,
            first_start=first_start,
            frequency=env.trading_day,
            periods_between_starts=3,
            asset_lifetime=5
        )

        env.write_data(equities_df=frame)
        finder = env.asset_finder

        all_dates = pd.date_range(
            start=first_start,
            end=frame.end_date.max(),
            freq=trading_day,
        )

        for dates in all_subindices(all_dates):
            expected_mask = full(
                shape=(len(dates), num_assets),
                fill_value=False,
                dtype=bool,
            )

            for i, date in enumerate(dates):
                it = frame[['start_date', 'end_date']].itertuples()
                for j, start, end in it:
                    if start <= date <= end:
                        expected_mask[i, j] = True

            # Filter out columns with all-empty columns.
            expected_result = pd.DataFrame(
                data=expected_mask,
                index=dates,
                columns=frame.index.values,
            )

            actual_result = finder.lifetimes(dates)
            assert_frame_equal(actual_result, expected_result)

    def _handle_market_close(self, completed_date):

        # increment the day counter before we move markers forward.
        self.day_count += 1.0

        # Get the next trading day and, if it is past the bounds of this
        # simulation, return the daily perf packet
        next_trading_day = TradingEnvironment.instance().\
            next_trading_day(completed_date)

        # Check if any assets need to be auto-closed before generating today's
        # perf period
        if next_trading_day:
            self.check_asset_auto_closes(next_trading_day=next_trading_day)

        # Take a snapshot of our current performance to return to the
        # browser.
        daily_update = self.to_dict(emission_type='daily')

        # On the last day of the test, don't create tomorrow's performance
        # period.  We may not be able to find the next trading day if we're at
        # the end of our historical data
        if self.market_close >= self.last_close:
            return daily_update

        # move the market day markers forward
        env = TradingEnvironment.instance()
        self.market_open, self.market_close = \
            env.next_open_and_close(self.day)
        self.day = env.next_trading_day(self.day)

        # Roll over positions to current day.
        self.todays_performance.rollover()
        self.todays_performance.period_open = self.market_open
        self.todays_performance.period_close = self.market_close

        # If the next trading day is irrelevant, then return the daily packet
        if (next_trading_day is None) or (next_trading_day >= self.last_close):
            return daily_update

        # Check for any dividends and auto-closes, then return the daily perf
        # packet
        self.check_upcoming_dividends(next_trading_day=next_trading_day)
        return daily_update

    def setUp(self):
        self.sids = range(90)
        self.env = TradingEnvironment()
        self.env.write_data(equities_identifiers=self.sids)

        self.sim_params = factory.create_simulation_parameters(
            start=datetime(1990, 1, 1, tzinfo=pytz.utc),
            end=datetime(1990, 1, 8, tzinfo=pytz.utc),
            env=self.env,
        )

    def setUp(self):
        self.__calendar = date_range('2014', '2015', freq=trading_day)
        self.__assets = assets = Int64Index(arange(1, 20))

        # Set up env for test
        env = TradingEnvironment()
        env.write_data(
            equities_df=make_simple_asset_info(
                assets,
                self.__calendar[0],
                self.__calendar[-1],
            ),
        )
        self.__finder = env.asset_finder

        # Use a 30-day period at the end of the year by default.
        self.__mask = self.__finder.lifetimes(
            self.__calendar[-30:],
            include_start_date=False,
        )

def minutes_for_days():
    """
    500 randomly selected days.
    This is used to make sure our test coverage is unbaised towards any rules.
    We use a random sample because testing on all the trading days took
    around 180 seconds on my laptop, which is far too much for normal unit
    testing.

    We manually set the seed so that this will be deterministic.
    Results of multiple runs were compared to make sure that this is actually
    true.

    This returns a generator of tuples each wrapping a single generator.
    Iterating over this yeilds a single day, iterating over the day yields
    the minutes for that day.
    """
    env = TradingEnvironment()
    random.seed('deterministic')
    return ((env.market_minutes_for_day(random.choice(env.trading_days)),)
            for _ in range(500))

def mixed_frequency_expected_index(count, frequency):
    """
    Helper for enumerating expected indices for test_mixed_frequency.
    """
    env = TradingEnvironment.instance()
    minute = MIXED_FREQUENCY_MINUTES[count]

    if frequency == '1d':
        return [env.previous_open_and_close(minute)[1], minute]
    elif frequency == '1m':
        return [env.previous_market_minute(minute), minute]