"""Tests for common behaviors shared by all ComputableTerms.
"""
import numpy as np

from zipline.lib.labelarray import LabelArray
from zipline.pipeline import Classifier, Factor, Filter
from zipline.testing import parameter_space
from zipline.utils.numpy_utils import (
    categorical_dtype,
    datetime64ns_dtype,
    float64_dtype,
    int64_dtype,
    NaTns,
)

from .base import BaseUSEquityPipelineTestCase


class Floats(Factor):
    inputs = ()
    window_length = 0
    dtype = float64_dtype


class AltFloats(Factor):
    inputs = ()
    window_length = 0
    dtype = float64_dtype


class Dates(Factor):
    inputs = ()
    window_length = 0
    dtype = datetime64ns_dtype


class AltDates(Factor):
    inputs = ()
    window_length = 0
    dtype = datetime64ns_dtype


class Bools(Filter):
    inputs = ()
    window_length = 0


class AltBools(Filter):
    inputs = ()
    window_length = 0


class Strs(Classifier):
    inputs = ()
    window_length = 0
    dtype = categorical_dtype
    missing_value = None


class AltStrs(Classifier):
    inputs = ()
    window_length = 0
    dtype = categorical_dtype
    missing_value = None


class Ints(Classifier):
    inputs = ()
    window_length = 0
    dtype = int64_dtype
    missing_value = -1


class AltInts(Classifier):
    inputs = ()
    window_length = 0
    dtype = int64_dtype
    missing_value = -1


class FillNATestCase(BaseUSEquityPipelineTestCase):

    @parameter_space(
        null_locs=[
            # No NaNs.
            np.zeros((4, 4), dtype=bool),
            # All NaNs.
            np.ones((4, 4), dtype=bool),
            # NaNs on Diagonal
            np.eye(4, dtype=bool),
            # Nans every third element.
            (np.arange(16).reshape(4, 4) % 3) == 0,
        ]
    )
    def test_fillna_with_scalar(self, null_locs):
        shape = (4, 4)
        num_cells = shape[0] * shape[1]

        floats = np.arange(num_cells, dtype=float).reshape(shape)
        floats[null_locs] = np.nan
        float_fillval = 999.0
        float_expected = np.where(null_locs, float_fillval, floats)
        float_expected_zero = np.where(null_locs, 0.0, floats)

        dates = (np.arange(num_cells, dtype='i8')
                 .view('M8[D]')
                 .astype('M8[ns]')
                 .reshape(shape))
        dates[null_locs] = NaTns
        date_fillval = np.datetime64('2014-01-02', 'ns')
        date_expected = np.where(null_locs, date_fillval, dates)

        strs = np.arange(num_cells).astype(str).astype(object).reshape(shape)
        strs[null_locs] = None
        str_fillval = "filled"
        str_expected = np.where(null_locs, str_fillval, strs)

        ints = np.arange(num_cells, dtype='i8').reshape(shape)
        ints[null_locs] = -1
        int_fillval = 777
        int_expected = np.where(null_locs, int_fillval, ints)

        terms = {
            'floats': Floats().fillna(float_fillval),
            # Make sure we accept integer as a fill value on float-dtype
            # factors.
            'floats_fill_zero': Floats().fillna(0),
            'dates': Dates().fillna(date_fillval),
            'strs': Strs().fillna(str_fillval),
            'ints': Ints().fillna(int_fillval),
        }

        expected = {
            'floats': float_expected,
            'floats_fill_zero': float_expected_zero,
            'dates': date_expected,
            'strs': self.make_labelarray(str_expected),
            'ints': int_expected,
        }

        self.check_terms(
            terms,
            expected,
            initial_workspace={
                Floats(): floats,
                Dates(): dates,
                Strs(): self.make_labelarray(strs),
                Ints(): ints,
            },
            mask=self.build_mask(self.ones_mask(shape=(4, 4))),
        )

    @parameter_space(
        null_locs=[
            # No NaNs.
            np.zeros((4, 4), dtype=bool),
            # # All NaNs.
            np.ones((4, 4), dtype=bool),
            # NaNs on Diagonal
            np.eye(4, dtype=bool),
            # Nans every third element.
            (np.arange(16).reshape((4, 4)) % 3) == 0,
        ]
    )
    def test_fillna_with_expression(self, null_locs):
        shape = (4, 4)
        mask = self.build_mask(self.ones_mask(shape=(4, 4)))
        state = np.random.RandomState(4)
        assets = self.asset_finder.retrieve_all(mask.columns)

        def rand_vals(dtype):
            return state.randint(1, 100, shape).astype(dtype)

        floats = np.arange(16, dtype=float).reshape(shape)
        floats[null_locs] = np.nan
        float_fillval = rand_vals(float)
        float_expected = np.where(null_locs, float_fillval, floats)
        float_expected_1d = np.where(null_locs, float_fillval[:, [0]], floats)

        dates = (np.arange(16, dtype='i8')
                 .view('M8[D]')
                 .astype('M8[ns]').
                 reshape(shape))
        dates[null_locs] = NaTns
        date_fillval = rand_vals('M8[D]').astype('M8[ns]')
        date_expected = np.where(null_locs, date_fillval, dates)
        date_expected_1d = np.where(null_locs, date_fillval[:, [1]], dates)

        strs = np.arange(16).astype(str).astype(object).reshape(shape)
        strs[null_locs] = None
        str_fillval = rand_vals(str)
        str_expected = np.where(null_locs, str_fillval, strs)
        str_expected_1d = np.where(null_locs, str_fillval[:, [2]], strs)

        ints = np.arange(16).reshape(shape)
        ints[null_locs] = -1
        int_fillval = rand_vals(int64_dtype)
        int_expected = np.where(null_locs, int_fillval, ints)
        int_expected_1d = np.where(null_locs, int_fillval[:, [3]], ints)

        terms = {
            'floats': Floats().fillna(AltFloats()),
            'floats_1d': Floats().fillna(AltFloats()[assets[0]]),

            'dates': Dates().fillna(AltDates()),
            'dates_1d': Dates().fillna(AltDates()[assets[1]]),

            'strs': Strs().fillna(AltStrs()),
            'strs_1d': Strs().fillna(AltStrs()[assets[2]]),

            'ints': Ints().fillna(AltInts()),
            'ints_1d': Ints().fillna(AltInts()[assets[3]]),
        }

        expected = {
            'floats': float_expected,
            'floats_1d': float_expected_1d,
            'dates': date_expected,
            'dates_1d': date_expected_1d,
            'strs': self.make_labelarray(str_expected),
            'strs_1d': self.make_labelarray(str_expected_1d),
            'ints': int_expected,
            'ints_1d': int_expected_1d,
        }

        self.check_terms(
            terms,
            expected,
            initial_workspace={
                Floats(): floats,
                Dates(): dates,
                Strs(): self.make_labelarray(strs),
                Ints(): ints,

                AltFloats(): float_fillval,
                AltDates(): date_fillval,
                AltStrs(): self.make_labelarray(str_fillval),
                AltInts(): int_fillval,
            },
            mask=mask,
        )

    def should_error(self, f, exc_type, expected_message):
        with self.assertRaises(exc_type) as e:
            f()

        message = str(e.exception)
        self.assertIn(expected_message, message)

    def test_bad_inputs(self):
        def dtype_for(o):
            return np.array([o]).dtype

        self.should_error(
            lambda: Floats().fillna('3.0'),
            TypeError,
            " from {!r} to {!r} according to the rule 'same_kind'"
            .format(dtype_for('3.0'), np.dtype(float))
        )

        self.should_error(
            lambda: Dates().fillna('2014-01-02'),
            TypeError,
            "from {!r} to {!r} according to the rule 'same_kind'"
            .format(dtype_for('2014-01-02'), np.dtype('M8[ns]'))
        )

        self.should_error(
            lambda: Ints().fillna('300'),
            TypeError,
            "from {!r} to {!r} according to the rule 'same_kind'"
            .format(dtype_for('300'), np.dtype('i8')),
        )

        self.should_error(
            lambda: Strs().fillna(10.0),
            TypeError,
            "Fill value 10.0 is not a valid choice for term Strs with dtype"
            " object.\n\n"
            "Coercion attempt failed with: "
            "String-dtype classifiers can only produce strings or None."

        )

    def make_labelarray(self, strs):
        return LabelArray(strs, missing_value=None)
