# Path: drain/model.py
# def y_subset(y, query=None, aux=None, subset=None, dropna=False, outcome='true',
#              k=None, p=None, ascending=False, score='score', p_of='notnull'):
#     """
#     Subset a model "y" dataframe
#     Args:
#         query: operates on y, or aux if present
#         subset: takes a dataframe or index thereof and subsets to that
#         dropna: means drop missing outcomes
#         return: top k (count) or p (proportion) if specified
#         p_of: specifies what the proportion is relative to
#             'notnull' means proportion is relative to labeled count
#             'true' means proportion is relative to positive count
#             'all' means proportion is relative to total count
#     """
#     if query is not None:
#         if aux is None:
#             y = y.query(query)
#         else:
#             s = aux.ix[y.index]
#             if len(s) != len(y):
#                 logging.warning('y not a subset of aux')
#             y = y.ix[s.query(query).index]
# 
#     if subset is not None:
#         if hasattr(subset, 'index'):
#             subset = subset.index
#         y = y.ix[y.index.intersection(subset)]
# 
#     if dropna:
#         y = y.dropna(subset=[outcome])
# 
#     if k is not None and p is not None:
#         raise ValueError("Cannot specify both k and p")
#     elif k is not None:
#         k = k
#     elif p is not None:
#         if p_of == 'notnull':
#             k = int(p*y[outcome].notnull().sum())
#         elif p_of == 'true':
#             k = int(p*y[outcome].sum())
#         elif p_of == 'all':
#             k = int(p*len(y))
#         else:
#             raise ValueError('Invalid value for p_of: %s' % p_of)
#     else:
#         k = None
# 
#     if k is not None:
#         y = y.sort_values(score, ascending=ascending).head(k)
# 
#     return y


import pandas as pd
import numpy as np
from drain.model import y_subset


y = pd.DataFrame({
    'score':[.1,1,.2,.3,0],
    'true':[True, False, np.nan,False,True],
    'test':[True, True, True, False, False],
    'attr':[False, True, False, False, False],
})

def test_subset_query():
    assert set(y_subset(y, query='attr').index) == set([1])

def test_subset_dropna():
   assert set(y_subset(y, dropna=True).index) == set([0,1,3,4])

def test_subset_k():

# Path: drain/model.py
# def y_subset(y, query=None, aux=None, subset=None, dropna=False, outcome='true',
#              k=None, p=None, ascending=False, score='score', p_of='notnull'):
#     """
#     Subset a model "y" dataframe
#     Args:
#         query: operates on y, or aux if present
#         subset: takes a dataframe or index thereof and subsets to that
#         dropna: means drop missing outcomes
#         return: top k (count) or p (proportion) if specified
#         p_of: specifies what the proportion is relative to
#             'notnull' means proportion is relative to labeled count
#             'true' means proportion is relative to positive count
#             'all' means proportion is relative to total count
#     """
#     if query is not None:
#         if aux is None:
#             y = y.query(query)
#         else:
#             s = aux.ix[y.index]
#             if len(s) != len(y):
#                 logging.warning('y not a subset of aux')
#             y = y.ix[s.query(query).index]
# 
#     if subset is not None:
#         if hasattr(subset, 'index'):
#             subset = subset.index
#         y = y.ix[y.index.intersection(subset)]
# 
#     if dropna:
#         y = y.dropna(subset=[outcome])
# 
#     if k is not None and p is not None:
#         raise ValueError("Cannot specify both k and p")
#     elif k is not None:
#         k = k
#     elif p is not None:
#         if p_of == 'notnull':
#             k = int(p*y[outcome].notnull().sum())
#         elif p_of == 'true':
#             k = int(p*y[outcome].sum())
#         elif p_of == 'all':
#             k = int(p*len(y))
#         else:
#             raise ValueError('Invalid value for p_of: %s' % p_of)
#     else:
#         k = None
# 
#     if k is not None:
#         y = y.sort_values(score, ascending=ascending).head(k)
# 
#     return y


# Path: tests/test_model.py
import pandas as pd
import numpy as np
from drain.model import y_subset



y = pd.DataFrame({
    'score':[.1,1,.2,.3,0],
    'true':[True, False, np.nan,False,True],
    'test':[True, True, True, False, False],
    'attr':[False, True, False, False, False],
})

def test_subset_query():
    assert set(y_subset(y, query='attr').index) == set([1])

def test_subset_dropna():
   assert set(y_subset(y, dropna=True).index) == set([0,1,3,4])

def test_subset_k():

   assert set(y_subset(y, k=2).index) == set([1,3])