import numpy as np
import pysal
import scipy.sparse as SP
import itertools as iter
from scipy.stats import f, chisqprob
import numpy.linalg as la
from utils import spbroadcast

"""
Tools for different regimes procedure estimations
"""

__author__ = "Luc Anselin luc.anselin@asu.edu, \
        Daniel Arribas-Bel darribas@asu.edu, \
        Pedro V. Amaral pedro.amaral@asu.edu"


class Chow:

    '''
    Chow test of coefficient stability across regimes. The test is a
    particular case of the Wald statistic in which the constraint are setup
    according to the spatial or other type of regime structure

    ...

    Parameters
    ==========
    reg     : regression object
              Regression object from PySAL.spreg which is assumed to have the
              following attributes:

                    * betas : coefficient estimates
                    * vm    : variance covariance matrix of betas
                    * kr    : Number of variables varying across regimes
                    * kryd  : Number of endogenous variables varying across regimes
                    * kf    : Number of variables fixed (global) across regimes
                    * nr    : Number of regimes

    Attributes
    ==========
    joint   : tuple
              Pair of Wald statistic and p-value for the setup of global
              regime stability, that is all betas are the same across
              regimes.
    regi    : array
              kr x 2 array with Wald statistic (col 0) and its p-value (col 1)
              for each beta that varies across regimes. The restrictions
              are setup to test for the global stability (all regimes have the
              same parameter) of the beta.

    Examples
    ========
    >>> import numpy as np
    >>> import pysal
    >>> from ols_regimes import OLS_Regimes
    >>> db = pysal.open(pysal.examples.get_path('columbus.dbf'),'r')
    >>> y_var = 'CRIME'
    >>> y = np.array([db.by_col(y_var)]).reshape(49,1)
    >>> x_var = ['INC','HOVAL']
    >>> x = np.array([db.by_col(name) for name in x_var]).T
    >>> r_var = 'NSA'
    >>> regimes = db.by_col(r_var)
    >>> olsr = OLS_Regimes(y, x, regimes, constant_regi='many', nonspat_diag=False, spat_diag=False, name_y=y_var, name_x=x_var, name_ds='columbus', name_regimes=r_var, regime_err_sep=False)
    >>> print olsr.name_x_r #x_var
    ['CONSTANT', 'INC', 'HOVAL']
    >>> print olsr.chow.regi
    [[ 0.01020844  0.91952121]
     [ 0.46024939  0.49750745]
     [ 0.55477371  0.45637369]]
    >>> print 'Joint test:'
    Joint test:
    >>> print olsr.chow.joint
    (0.6339319928978806, 0.8886223520178802)
    '''

    def __init__(self, reg):
        kr, kf, kryd, nr, betas, vm = reg.kr, reg.kf, reg.kryd, reg.nr, reg.betas, reg.vm
        if betas.shape[0] != vm.shape[0]:
            if kf > 0:
                betas = betas[0:vm.shape[0], :]
                kf = kf - 1
            else:
                brange = []
                for i in range(nr):
                    brange.extend(range(i * (kr + 1), i * (kr + 1) + kr))
                betas = betas[brange, :]
        r_global = []
        regi = np.zeros((reg.kr, 2))
        for vari in np.arange(kr):
            r_vari = buildR1var(vari, kr, kf, kryd, nr)
            r_global.append(r_vari)
            q = np.zeros((r_vari.shape[0], 1))
            regi[vari, :] = wald_test(betas, r_vari, q, vm)
        r_global = np.vstack(tuple(r_global))
        q = np.zeros((r_global.shape[0], 1))
        joint = wald_test(betas, r_global, q, vm)
        self.joint = joint
        self.regi = regi


class Wald:

    '''
    Chi sq. Wald statistic to test for restriction of coefficients.
    Implementation following Greene [Greene2003]_ eq. (17-24), p. 488

    ...

    Parameters
    ==========
    reg     : regression object
              Regression object from PySAL.spreg
    r       : array
              Array of dimension Rxk (R being number of restrictions) with constrain setup.
    q       : array
              Rx1 array with constants in the constraint setup. See Greene
              [1]_ for reference.

    Attributes
    ==========
    w       : float
              Wald statistic
    pvalue  : float
              P value for Wald statistic calculated as a Chi sq. distribution
              with R degrees of freedom

    '''

    def __init__(self, reg, r, q=None):
        if not q:
            q = np.zeros((r.shape[0], 1))
        self.w, self.pvalue = wald_test(reg.betas, r, q, reg.vm)


class Regimes_Frame:

    '''
    Setup framework to work with regimes. Basically it involves:
        * Dealing with the constant in a regimes world
        * Creating a sparse representation of X 
        * Generating a list of names of X taking into account regimes

    ...

    Parameters
    ==========
    x            : array
                   Two dimensional array with n rows and one column for each
                   independent (exogenous) variable, excluding the constant
    regimes      : list
                   List of n values with the mapping of each
                   observation to a regime. Assumed to be aligned with 'x'.
    constant_regi: [False, 'one', 'many']
                   Switcher controlling the constant term setup. It may take
                   the following values:

                     *  False: no constant term is appended in any way
                     *  'one': a vector of ones is appended to x and held
                               constant across regimes
                     * 'many': a vector of ones is appended to x and considered
                               different per regime (default)
    cols2regi    : list, 'all'
                   Argument indicating whether each
                   column of x should be considered as different per regime
                   or held constant across regimes (False).
                   If a list, k booleans indicating for each variable the
                   option (True if one per regime, False to be held constant).
                   If 'all' (default), all the variables vary by regime.
    names         : None, list of strings
                   Names of independent variables for use in output

    Returns
    =======
    x            : csr sparse matrix
                   Sparse matrix containing X variables properly aligned for
                   regimes regression. 'xsp' is of dimension (n, k*r) where 'r'
                   is the number of different regimes
                   The structure of the alignent is X1r1 X2r1 ... X1r2 X2r2 ...
    names        : None, list of strings
                   Names of independent variables for use in output
                   conveniently arranged by regimes. The structure of the name
                   is "regimeName_-_varName"
    kr           : int
                   Number of variables/columns to be "regimized" or subject
                   to change by regime. These will result in one parameter
                   estimate by regime for each variable (i.e. nr parameters per
                   variable)
    kf           : int
                   Number of variables/columns to be considered fixed or
                   global across regimes and hence only obtain one parameter
                   estimate
    nr           : int
                   Number of different regimes in the 'regimes' list

    '''

    def __init__(self, x, regimes, constant_regi, cols2regi, names=None, yend=False):
        if cols2regi == 'all':
            cols2regi = [True] * x.shape[1]
        else:
            if yend:
                cols2regi = cols2regi[-x.shape[1]:]
            else:
                cols2regi = cols2regi[0:x.shape[1]]
        if constant_regi:
            x = np.hstack((np.ones((x.shape[0], 1)), x))
            if constant_regi == 'one':
                cols2regi.insert(0, False)
            elif constant_regi == 'many':
                cols2regi.insert(0, True)
            else:
                raise Exception, "Invalid argument (%s) passed for 'constant_regi'. Please secify a valid term." % str(
                    constant)
        try:
            x = regimeX_setup(x, regimes, cols2regi,
                              self.regimes_set, constant=constant_regi)
        except AttributeError:
            self.regimes_set = _get_regimes_set(regimes)
            x = regimeX_setup(x, regimes, cols2regi,
                              self.regimes_set, constant=constant_regi)

        kr = len(np.where(np.array(cols2regi) == True)[0])
        if yend:
            self.kr += kr
            self.kf += len(cols2regi) - kr
            self.kryd = kr
        else:
            self.kr = kr
            self.kf = len(cols2regi) - self.kr
            self.kryd = 0
        self.nr = len(set(regimes))

        if names:
            names = set_name_x_regimes(
                names, regimes, constant_regi, cols2regi, self.regimes_set)

        return (x, names)


def wald_test(betas, r, q, vm):
    '''
    Chi sq. Wald statistic to test for restriction of coefficients.
    Implementation following Greene [Greene2003]_ eq. (17-24), p. 488

    ...

    Parameters
    ==========
    betas   : array
              kx1 array with coefficient estimates
    r       : array
              Array of dimension Rxk (R being number of restrictions) with constrain setup.
    q       : array
              Rx1 array with constants in the constraint setup. See Greene
              [1]_ for reference.
    vm      : array
              kxk variance-covariance matrix of coefficient estimates

    Returns
    =======
    w       : float
              Wald statistic
    pvalue  : float
              P value for Wald statistic calculated as a Chi sq. distribution
              with R degrees of freedom

    '''
    rbq = np.dot(r, betas) - q
    rvri = la.inv(np.dot(r, np.dot(vm, r.T)))