#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
# This file is part of pandas-learn
# https://github.com/RichLewis42/pandas-learn
#
# Licensed under the MIT license:
# http://www.opensource.org/licenses/MIT
# Copyright (c) 2015, Rich Lewis <rl403@cam.ac.uk>

"""
tests.test_adaptor.classifier
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Tests for the classifier adaptor module of pdlearn.
"""

import pytest

from ..test_utils import DATA_FRAME, SERIES, ONE_D_ARRAY, TWO_D_ARRAY
from pdlearn.utils import is_frame, is_series, CompatabilityWarning

import pdlearn.adaptor.classifier
import pdlearn.adaptor.model
import numpy as np
import pandas as pd

TARGETS = pd.DataFrame(TWO_D_ARRAY, columns=['T', 'G'])

class MockClassifier(object):

    def __init__(self, multitask=False):
        if multitask:
            self.classes_ = [np.array([0, 1]), np.array([0, 1])]
        else:
            self.classes_ = np.array([0, 1])

        self._num_tasks = 2 if multitask else 1

    def fit(self, X, y=None):
        pass

    def predict(self, X):
        if self._num_tasks == 1:
            return np.array([1] * len(X))
        else:
            return [np.array([1] * len(X))] * self._num_tasks

    def predict_proba(self, X):

        if self._num_tasks == 1:
            return np.array([[0.1, 0.9]] * len(X))
        else:
            return [np.array([[0.1, 0.9]] * len(X))] * self._num_tasks

    def predict_log_proba(self, X):
        if self._num_tasks == 1:
            return np.array([[-1, -0.04]] * len(X))
        else:
            return [np.array([[-1, -0.04]] * len(X))] * self._num_tasks


@pdlearn.adaptor.classifier
class ChildClassifier(MockClassifier):

    def __init__(self, features=None, targets=None, multitask=False):

        if targets:
            self.target_names_ = pd.Index(targets)
            multitask = multitask or len(self.target_names_) > 1
        if features:
            self.feature_names_ = pd.Index(features)

        super(ChildClassifier, self).__init__(multitask)

class TestClassifier(object):

    """ Tests for pdlearn.adaptor.classifier """

    def test_not_multitask(self):

        assert not ChildClassifier(targets=['a']).multitask_

    def test_multitask(self):

        assert ChildClassifier(targets=['a', 'b']).multitask_


    def test_fit_with_arr(self):

        child = ChildClassifier().fit(TWO_D_ARRAY)
        assert not child.pandas_mode_

        child.fit(TWO_D_ARRAY, ONE_D_ARRAY)
        assert not child.pandas_mode_

    def test_with_hybrid(self):

        child = ChildClassifier().fit(TWO_D_ARRAY, SERIES)
        assert child.pandas_mode_
        assert np.array_equal(child.target_names_, ['T'])

        child.fit(DATA_FRAME, ONE_D_ARRAY)
        assert child.pandas_mode_
        assert np.array_equal(child.feature_names_, ['x', 'y'])

    def test_fit_with_df(self):

        child = ChildClassifier().fit(DATA_FRAME)
        assert child.pandas_mode_
        assert np.array_equal(child.feature_names_, ['x', 'y'])

    def test_fit_with_df_ser(self):