import numpy as N
import numpy.linalg as L
import scipy.linalg
from scipy.sandbox.models.model import LikelihoodModel, LikelihoodModelResults
from scipy.sandbox.models import utils

class ols_model(LikelihoodModel):
    
    """
    A simple ordinary least squares model.
    """

    def logL(self, b, Y):
        return -scipy.linalg.norm(self.whiten(Y) - N.dot(self.wdesign, b))**2 / 2.

    def __init__(self, design):
        LikelihoodModel.__init__(self)
        self.initialize(design)

    def initialize(self, design):
        self.design = design
        self.wdesign = self.whiten(design)
        self.calc_beta = L.pinv(self.wdesign)
        self.normalized_cov_beta = N.dot(self.calc_beta,
                                         N.transpose(self.calc_beta))
        self.df_resid = self.wdesign.shape[0] - utils.rank(self.design)

    def whiten(self, Y):
        return Y
    
    def est_coef(self, Y):
        """
        Estimate coefficients using lstsq, returning fitted values, Y
        and coefficients, but initialize is not called so no
        psuedo-inverse is calculated.
        """
            
        Z = self.whiten(Y)

        lfit = Results(L.lstsq(self.wdesign, Z)[0], Y)
        lfit.predict = N.dot(self.design, lfit.beta)


    def fit(self, Y):
        """
        Full \'fit\' of the model including estimate of covariance matrix, (whitened)
        residuals and scale. 

        """
    
        Z = self.whiten(Y)

        lfit = Results(N.dot(self.calc_beta, Z), Y,
                       normalized_cov_beta=self.normalized_cov_beta)

        lfit.df_resid = self.df_resid
        lfit.predict = N.dot(self.design, lfit.beta)
        lfit.resid = Z - N.dot(self.wdesign, lfit.beta)
        lfit.scale = N.add.reduce(lfit.resid**2) / lfit.df_resid

        lfit.Z = Z # just in case
        
        return lfit

class ar_model(ols_model):
    """
    A regression model with an AR(1) covariance structure.

    Eventually, this will be AR(p) -- all that is needed is to
    determine the self.whiten method from AR(p) parameters.
    """

    def __init__(self, design, rho=0):
        self.rho = rho
        ols_model.__init__(self, design)


    def whiten(self, X):
        factor = 1. / N.sqrt(1 - self.rho**2)
        return N.concatenate([[X[0]], (X[1:] - self.rho * X[0:-1]) * factor])

class wls_model(ols_model):
    """

    A regression model with diagonal but non-identity covariance
    structure. The weights are proportional to the inverse of the
    variance of the observations.

    """

    def __init__(self, design, weights=1):
        self.weights = weights
        ols_model.__init__(self, design)


    def whiten(self, X):
        if X.ndim == 1:
            return X * N.sqrt(self.weights)
        elif X.ndim == 2:
            c = N.sqrt(self.weights)
            v = N.zeros(X.shape, N.float64)
            for i in range(X.shape[1]):
                v[:,i] = X[:,i] * c
            return v
    

class Results(LikelihoodModelResults):
    """
    This class summarizes the fit of a linear regression model.
    It handles the output of contrasts, estimates of covariance, etc.
    """

    def __init__(self, beta, Y, normalized_cov_beta=None, scale=1.):
        LikelihoodModelResults.__init__(self, beta, normalized_cov_beta, scale)
        self.Y = Y

    def norm_resid(self):
        """
        Residuals, normalized to have unit length.

        Note: residuals are whitened residuals.
        """

        if not hasattr(self, 'resid'):
            raise ValueError, 'need normalized residuals to estimate standard deviation'

        sdd = utils.recipr(self.sd) / N.sqrt(self.df)
        return  self.resid * N.multiply.outer(N.ones(self.Y.shape[0]), sdd)


    def predict(self, design):
        """
        Return fitted values from a design matrix.
        """
        return N.dot(design, self.beta)

    def Rsq(self, adjusted=False):
        """
        Return the R^2 value for each row of the response Y.
        """
        self.Ssq = N.std(self.Z,axis=0)**2
        ratio = self.scale / self.Ssq
        if not adjusted: ratio *= ((self.Y.shape[0] - 1) / self.df_resid)
        return 1 - ratio


def isestimable(C, D):
    """
    From an q x p contrast matrix C and an n x p design matrix D, checks
    if the contrast C is estimable by looking at the rank of hstack([C,D]) and
    verifying it is the same as the rank of D.
    """

    if C.ndim == 1:
        C.shape = (C.shape[0], 1)
    new = N.hstack([C, D])
    if utils.rank(new) != utils.rank(D):
        return False
    return True