"""
Tests for Downsampled Filters/Factors/Classifiers
"""
from functools import partial

import pandas as pd
from pandas.util.testing import assert_frame_equal

from zipline.errors import NoFurtherDataError
from zipline.pipeline import (
    Pipeline,
    CustomFactor,
    CustomFilter,
    CustomClassifier,
    SimplePipelineEngine,
)
from zipline.pipeline.data.testing import TestingDataSet
from zipline.pipeline.domain import (
    CA_EQUITIES,
    EquitySessionDomain,
    GB_EQUITIES,
    US_EQUITIES,
)
from zipline.pipeline.factors import SimpleMovingAverage
from zipline.pipeline.filters.smoothing import All
from zipline.testing import ZiplineTestCase, parameter_space, ExplodingObject
from zipline.testing.fixtures import (
    WithTradingSessions,
    WithSeededRandomPipelineEngine,
    WithAssetFinder,
)
from zipline.utils.classproperty import classproperty
from zipline.utils.input_validation import _qualified_name
from zipline.utils.numpy_utils import int64_dtype


class NDaysAgoFactor(CustomFactor):
    inputs = [TestingDataSet.float_col]

    def compute(self, today, assets, out, floats):
        out[:] = floats[0]


class NDaysAgoFilter(CustomFilter):
    inputs = [TestingDataSet.bool_col]

    def compute(self, today, assets, out, bools):
        out[:] = bools[0]


class NDaysAgoClassifier(CustomClassifier):
    inputs = [TestingDataSet.categorical_col]
    dtype = TestingDataSet.categorical_col.dtype

    def compute(self, today, assets, out, cats):
        out[:] = cats[0]


class ComputeExtraRowsTestCase(WithTradingSessions, ZiplineTestCase):

    DATA_MIN_DAY = pd.Timestamp('2012-06', tz='UTC')
    DATA_MAX_DAY = pd.Timestamp('2015', tz='UTC')
    TRADING_CALENDAR_STRS = ('NYSE', 'LSE', 'TSX')

    # Test with different window_lengths to ensure that window length is not
    # used when calculating exra rows for the top-level term.
    factor1 = TestingDataSet.float_col.latest
    factor11 = NDaysAgoFactor(window_length=11)
    factor91 = NDaysAgoFactor(window_length=91)

    filter1 = TestingDataSet.bool_col.latest
    filter11 = NDaysAgoFilter(window_length=11)
    filter91 = NDaysAgoFilter(window_length=91)

    classifier1 = TestingDataSet.categorical_col.latest
    classifier11 = NDaysAgoClassifier(window_length=11)
    classifier91 = NDaysAgoClassifier(window_length=91)

    all_terms = [
        factor1,
        factor11,
        factor91,
        filter1,
        filter11,
        filter91,
        classifier1,
        classifier11,
        classifier91,
    ]

    @parameter_space(
        calendar_name=TRADING_CALENDAR_STRS,
        base_terms=[
            (factor1, factor11, factor91),
            (filter1, filter11, filter91),
            (classifier1, classifier11, classifier91),
        ],
        __fail_fast=True
    )
    def test_yearly(self, base_terms, calendar_name):
        downsampled_terms = tuple(
            t.downsample('year_start') for t in base_terms
        )
        all_terms = base_terms + downsampled_terms

        all_sessions = self.trading_sessions[calendar_name]
        end_session = all_sessions[-1]

        years = all_sessions.year
        sessions_in_2012 = all_sessions[years == 2012]
        sessions_in_2013 = all_sessions[years == 2013]
        sessions_in_2014 = all_sessions[years == 2014]

        # Simulate requesting computation where the unaltered lookback would
        # land exactly on the first date in 2014.  We shouldn't request any
        # additional rows for the regular terms or the downsampled terms.
        for i in range(0, 30, 5):
            start_session = sessions_in_2014[i]
            self.check_extra_row_calculations(