# -*- coding: utf-8 -*-
"""Convenience Wrappers

Created on Sat Oct 30 14:56:35 2010

Author: josef-pktd
License: BSD
"""

import numpy as np
import scikits.statsmodels.api as sm
from scikits.statsmodels import GLS, WLS, OLS

def remove_nanrows(y, x):
    '''remove common rows in [y,x] that contain at least one nan

    TODO: this should be made more flexible,
     arbitrary number of arrays and 1d or 2d arrays

    duplicate: Skipper added sm.tools.drop_missing

    '''
    mask = ~np.isnan(y)
    mask *= ~(np.isnan(x).any(-1))  #* or &
    y = y[mask]
    x = x[mask]
    return y, x


def linmod(y, x, weights=None, sigma=None, add_const=True, filter_missing=True,
           **kwds):
    '''get linear model with extra options for entry

    dispatches to regular model class and does not wrap the output

    If several options are exclusive, for example sigma and weights, then the
    chosen class depends on the implementation sequence.
    '''

    if filter_missing:
        y, x = remove_nanrows(y, x)
        #do the same for masked arrays

    if add_const:
        x = sm.add_constant(x, prepend=True)

    if not sigma is None:
        return GLS(y, x, sigma=sigma, **kwds)
    elif not weights is None:
        return WLS(y, x, weights=weights, **kwds)
    else:
        return OLS(y, x, **kwds)