import numpy as np
import pandas as pd

class OLS_Simu:
    """Generate an OLS object with randomized data
       
    Parameters
    ----------
    n, k, beta: int, int, list
        n is the sample size, k is the number of parameters including constant, 
        beta is a list of predetermined OLS parameters. 
    """
    def __init__(self, n, k, beta):

        self.n = n
        self.k = k
        
        # generate X matrix    
        self.const = np.ones(self.n)
        self.const = self.const[np.newaxis, :]
        self.X_inde = np.random.randn(self.k-1, self.n)
        self.X = np.concatenate((self.const.T, self.X_inde.T), axis=1)
        
        # set beta a vertical vector
        self.beta = beta
        if len(self.beta) == self.k:
            self.beta_array = np.array(self.beta)
            self.beta_array = self.beta_array[np.newaxis, :].T
        else:
            raise Exception('length of beta is not equal to k!')
            
        # generate set disturbance term as a vertical vector
        self.u = np.random.randn(self.n)
        self.u = self.u[np.newaxis, :].T
        
        # generate y
        self.y = self.X@self.beta_array + self.u
        
        # estimates and residuals
        self.beta_hat = np.linalg.inv(self.X.T@self.X)@self.X.T@self.y
        self.resid =  self.y - self.X@self.beta_hat
    
#################################################
    def ols_estimates(self):
        """Estimate the model and return the beta_hat and residuals."""
        self.beta_hat = np.linalg.inv(self.X.T@self.X)@self.X.T@self.y
        self.resid =  self.y - self.X@self.beta_hat
        return self.beta_hat, self.resid
    
#################################################
    def anova(self):
        """ANOVA decomposition"""
        self.TSS = np.linalg.norm(self.y)
        self.ESS = np.linalg.norm(self.X@self.beta_hat)
        self.RSS = np.linalg.norm(self.y)-np.linalg.norm(self.X@self.beta_hat)
        return self.TSS, self.ESS, self.RSS

#################################################
    def proj_mat_P(self):
        """Projection matrix P"""
        self.proj_matP = self.X@np.linalg.inv(self.X.T@self.X)@self.X.T
        return self.proj_matP
#################################################
    def proj_mat_M(self):
        """Projection matrix M"""
        self.proj_matM =  np.eye(self.X.shape[0]) - self.X@np.linalg.inv(self.X.T@self.X)@self.X.T
        return self.proj_matM
#################################################
    def s_sqr(self):
        """Unbiased estimator of sigma square."""
        return np.sum(self.resid**2)/(len(self.resid)-self.k)
#################################################
    def cov_beta_hat(self):
        """Covariance matrix of estimated coefficients beta, standar"""
        self.cov_beta_hat = np.sum(self.resid**2)/(len(self.resid)-self.k)*np.linalg.inv(self.X.T@self.X)
        self.cov_beta_hat_prin_diag = self.cov_beta_hat.diagonal()
        return (self.cov_beta_hat, self.cov_beta_hat_prin_diag)


##################################################
##################################################
##################################################
##################################################
def gen_X(n, k):
    """
    Parameters
    ----------
    n, k: int
        generate a random explanatory matrxi with constant term of size n by k, n is the sample size,
        k is the number of parameters.
    Returns
    ----------
    X: ndarray
        matrix styled n by k random array generated by np.random.randn().
    """
    const = np.ones(n)
    const = const[np.newaxis, :]
    X_inde = np.random.randn(k-1, n)
    X = np.concatenate((const.T, X_inde.T), axis=1)
    return X

#################################################
def gen_beta(params = [2, 3, 4, 5]):
    """
    Parameters
    ----------
    params: list
        input a list of beta parameters
    
    Returns
    ----------
    beta_array: ndnarry
        an array of beta coefficients, no nothing input, default value
        params = [2, 3, 4, 5]
    """
    beta_array = np.array(params)
    beta_array = beta_array[np.newaxis, :].T
    return beta_array

#################################################
def gen_u(n):
    """
    Parameters
    ----------
    params: int
        input a integer for the size of disturbance term
    
    Returns
    ----------
    u: ndnarry
        an array of disturbance term
    """
    u = np.random.randn(n)
    u = u[np.newaxis, :].T
    return u

#################################################
def ols(y, X):
    """
    Parameters
    ----------
    y, x: ndarray
        input ndarrays produced by gen_X
    
    Returns
    ----------
    beta_hat: ndnarry
        an array of estimated coefficients
    """
    beta_hat = np.linalg.inv(X.T@X)@X.T@y
    return beta_hat

    # TSS = np.linalg.norm(y)
    # ESS = np.linalg.norm(X@beta_hat)
    # RSS = np.linalg.norm(y-X@beta_hat)
#################################################    
def proj_mat_P(X):
    """computing projection matrix P. Input X represents regressors.
    Parameters
    ----------
    X: ndarray
        input ndarrays produced by gen_X
    
    Returns
    ----------
    proj_matP: ndnarry
        Projection matrix P
    """
    proj_matP = X@np.linalg.inv(X.T@X)@X.T
    return proj_matP
#################################################
def proj_mat_M(X):
    """computing projection matrix M. Input X represents regressors.
    Parameters
    ----------
    X: ndarray
        input ndarrays produced by gen_X
    
    Returns
    ----------
    proj_matM: ndnarry
        Projection matrix M
    """
    proj_matM =  np.eye(X.shape[0]) - X@np.linalg.inv(X.T@X)@X.T
    return proj_matM
#################################################
def s_qr(resid, k):
    """
    Unbiased estimator of sigma square.
    Parameters
    ----------
    params: ndarray
        input an matrix-form ndarray of residuals
    
    Returns
    ----------
    beta_array: float
        return a float number
    """
    return np.sum(resid**2)/(len(resid)-k)

#################################################
def cov_beta_hat(resid, k, X):
    """
    Covariance matrix of estimated coefficients beta, standard 
    Parameters
    ----------
    resid, k, X: ndarray, integer, ndarray

    
    Returns
    ----------
    cov_beta_hat, prin_diag: ndarray, ndarray

    """
    cov_beta_hat = np.sum(resid**2)/(len(resid)-k)*np.linalg.inv(X.T@X)
    cov_beta_hat_prin_diag = cov_beta_hat.diagonal()
    return (cov_beta_hat, cov_beta_hat_prin_diag)
