"""
Ordinary Least Squares (OLS) to generate multiple linear regression
taken from SciPy CookBook [http://www.scipy.org/Cookbook/OLS] is used
by the Solar Wind processor to compute Bx linear fit
coefficients.
"""
from __future__ import division
from scipy import stats
from numpy import c_, ones, dot, diff
from scipy.linalg import inv, solve, det
from numpy import log, pi, sqrt, square, diagonal
from numpy.random import randn, seed
import time

class ols:
    """
    Author: Vincent Nijs (+ ?)
    Email: v-nijs at kellogg.northwestern.edu
    Last Modified: Mon Jan 15 17:56:17 CST 2007
    Dependencies: See import statement at the top of this file

    Class for multi-variate regression using OLS.

    Args:
        y (array-like): Dependent variable.
        y_varnm (str): Variable label for y.
        x (array-like): Independent variables. Note that a constant is added by default.
        x_varnm (str or list): Variable labels for the independent variables.

    Attributes:
        b (array-like): Coefficients.
        p (array-like): Coefficient p-values.

    Methods:
        estimate(): Estimate coefficients and basic stats.
        dw(): Calculate the Durbin-Waston statistic.
        omni(): Omnibus test for normality.
        JB(): Calculate residual skewness, kurtosis, and perform the JB test for normality.
        ll(): Calculate model log-likelihood and two information criteria.
        summary(): Print model output to screen.

    Example:
        # Create an OLS instance
        m = ols(y, x, y_varnm='y', x_varnm=['x1', 'x2', 'x3', 'x4'])

        # Print the coefficients
        print(m.b)

        # Print the coefficient p-values
        print(m.p)
    """

    def __init__(self,y,x,y_varnm = 'y',x_varnm = ''):
        """
        Initializes the ols class.

        Args:
            y (array-like): Dependent variable.
            x (array-like): Independent variables. Note that a constant is added by default.
            y_varnm (str): Variable label for y.
            x_varnm (str or list): Variable labels for the independent variables.
        """
        self.y = y
        self.x = c_[ones(x.shape[0]),x]
        self.y_varnm = y_varnm
        if not isinstance(x_varnm,list): 
            self.x_varnm = ['const'] + list(x_varnm)
        else:
            self.x_varnm = ['const'] + x_varnm

        # Estimate model using OLS
        self.estimate()

    def estimate(self):
        """
        Estimates the coefficients and basic statistics for the OLS model.

        This method calculates the coefficients, residuals, standard errors, t-statistics,
        p-values, R-squared, adjusted R-squared, F-statistic, and F-statistic p-value.

        Returns:
            None
        """
        # estimating coefficients, and basic stats
        self.inv_xx = inv(dot(self.x.T,self.x))
        xy = dot(self.x.T,self.y)
        self.b = dot(self.inv_xx,xy)                    # estimate coefficients

        self.nobs = self.y.shape[0]                     # number of observations
        self.ncoef = self.x.shape[1]                    # number of coef.
        self.df_e = self.nobs - self.ncoef              # degrees of freedom, error 
        self.df_r = self.ncoef - 1                      # degrees of freedom, regression 

        self.e = self.y - dot(self.x,self.b)            # residuals
        self.sse = dot(self.e,self.e)/self.df_e         # SSE
        self.se = sqrt(diagonal(self.sse*self.inv_xx))  # coef. standard errors
        self.t = self.b / self.se                       # coef. t-statistics
        self.p = (1-stats.t.cdf(abs(self.t), self.df_e)) * 2    # coef. p-values

        self.R2 = 1 - self.e.var()/self.y.var()         # model R-squared
        self.R2adj = 1-(1-self.R2)*((self.nobs-1)/(self.nobs-self.ncoef))   # adjusted R-square

        self.F = (self.R2/self.df_r) / ((1-self.R2)/self.df_e)  # model F-statistic
        self.Fpv = 1-stats.f.cdf(self.F, self.df_r, self.df_e)  # F-statistic p-value

    def dw(self):
        """
        Calculates the Durbin-Waston statistic.

        Returns:
            float: The Durbin-Waston statistic.
        """
        de = diff(self.e,1)
        dw = dot(de,de) / dot(self.e,self.e);

        return dw

    def omni(self):
        """
        Calculates the Omnibus test for normality.

        Returns:
            tuple: A tuple containing the chi-square test statistic and the p-value.
        """
        return stats.normaltest(self.e) 
    
    def JB(self):
        """
        Calculate residual skewness, kurtosis, and perform the JB test for normality.

        Returns:
            tuple: A tuple containing the JB test statistic, JB test p-value, skewness, and kurtosis.
        """

        # Calculate residual skewness and kurtosis
        skew = stats.skew(self.e) 
        kurtosis = 3 + stats.kurtosis(self.e) 
        
        # Calculate the Jarque-Bera test for normality
        JB = (self.nobs/6) * (square(skew) + (1/4)*square(kurtosis-3))
        JBpv = 1-stats.chi2.cdf(JB,2);

        return JB, JBpv, skew, kurtosis

    def ll(self):
        """
        Calculate model log-likelihood, AIC, and BIC criterion values.

        Returns:
            tuple: A tuple containing the model log-likelihood, AIC, and BIC.
        """
        
        # Model log-likelihood, AIC, and BIC criterion values 
        ll = -(self.nobs*1/2)*(1+log(2*pi)) - (self.nobs/2)*log(dot(self.e,self.e)/self.nobs)
        aic = -2*ll/self.nobs + (2*self.ncoef/self.nobs)
        bic = -2*ll/self.nobs + (self.ncoef*log(self.nobs))/self.nobs

        return ll, aic, bic
    
    def summary(self):
        """
        Prints model output to the screen.

        This method prints the model output, including the dependent variable, method used, date and time, number of observations, number of variables, coefficients, standard errors, t-statistics, p-values, model statistics, and residual statistics.

        Returns:
            None
        """

        # local time & date
        t = time.localtime()

        # extra stats
        ll, aic, bic = self.ll()
        JB, JBpv, skew, kurtosis = self.JB()
        omni, omnipv = self.omni()

        # printing output to screen
        print('\n==============================================================================')
        print("Dependent Variable: " + self.y_varnm)
        print("Method: Least Squares")
        print("Date: ", time.strftime("%a, %d %b %Y",t))
        print("Time: ", time.strftime("%H:%M:%S",t))
        print('# obs:               %5.0f' % self.nobs)
        print('# variables:     %5.0f' % self.ncoef) 
        print('==============================================================================')
        print('variable     coefficient     std. Error      t-statistic     prob.')
        print('==============================================================================')
        for i in range(len(self.x_varnm)):
            print('''% -5s          % -5.6f     % -5.6f     % -5.6f     % -5.6f''' % tuple([self.x_varnm[i],self.b[i],self.se[i],self.t[i],self.p[i]]) )
        print('==============================================================================')
        print('Models stats                         Residual stats')
        print('==============================================================================')
        print('R-squared            % -5.6f         Durbin-Watson stat  % -5.6f' % tuple([self.R2, self.dw()]))
        print('Adjusted R-squared   % -5.6f         Omnibus stat        % -5.6f' % tuple([self.R2adj, omni]))
        print('F-statistic          % -5.6f         Prob(Omnibus stat)  % -5.6f' % tuple([self.F, omnipv]))
        print('Prob (F-statistic)   % -5.6f			JB stat             % -5.6f' % tuple([self.Fpv, JB]))
        print('Log likelihood       % -5.6f			Prob(JB)            % -5.6f' % tuple([ll, JBpv]))
        print('AIC criterion        % -5.6f         Skew                % -5.6f' % tuple([aic, skew]))
        print('BIC criterion        % -5.6f         Kurtosis            % -5.6f' % tuple([bic, kurtosis]))
        print('==============================================================================')

if __name__ == '__main__':

    ##########################
    ### testing the ols class
    ##########################

    # creating simulated data and variable labels
    seed(1)
    data =  randn(100,5)            # the data array

    # intercept is added, by default
    m = ols(data[:,0],data[:,1:],y_varnm = 'y',x_varnm = ['x1','x2','x3','x4'])
    m.summary()

    # if you have rpy installed, use it to test the results
    have_rpy =  False
    try:
        print("\n")
        print("="*30)
        print("Validating OLS results in R")
        print("="*30)

        import rpy
        have_rpy = True
    except ImportError:
        print("\n")
        print("="*30)
        print("Validating OLS-class results in R")
        print("="*30)
        print("rpy is not installed")
        print("="*30)

    if have_rpy:
        y = data[:,0]
        x1 = data[:,1]
        x2 = data[:,2]
        x3 = data[:,3]
        x4 = data[:,4]
        rpy.set_default_mode(rpy.NO_CONVERSION)
        linear_model = rpy.r.lm(rpy.r("y ~ x1 + x2 + x3 + x4"), data = rpy.r.data_frame(x1=x1,x2=x2,x3=x3,x4=x4,y=y))
        rpy.set_default_mode(rpy.BASIC_CONVERSION)
        print(linear_model.as_py()['coefficients'])
        summary = rpy.r.summary(linear_model)
        print(summary)

