def test_beta():
    for matrix in [real_matrix(), complex_matrix()]:
        for beta in [0, 1, 3.4]:
            matrix_plus_beta = matrix + beta * sparse.eye(*matrix.shape)
            for mode in ["auto", "supernodal", "simplicial"]:
                L = cholesky(matrix, beta=beta).L()
                assert factor_of(cholesky(matrix, beta=beta),
                                 matrix_plus_beta)

def sparse_build_np_models(kernel, trans_samples, ter_samples, ter_rew_samples, lamb):
    Xa, Ra, Xpa = zip(*trans_samples)
    Xa_term, Ra_term = zip(*ter_samples)
    Xa = [ np.vstack((xa, xa_term)) if xa_term.size > 0 else xa for xa, xa_term in izip(Xa, Xa_term) ]
    Ra = [ np.hstack((ra, ra_term)) if ra_term.size > 0 else ra for ra, ra_term in izip(Ra, Ra_term) ] 
    
    k = len(trans_samples) 
    
    # build the K_a,b matrices
    Kab = dict()
    KabT = dict()
    for a,b in product(xrange(k), xrange(k)):
        if Xa_term[b].size > 0:
            Kab[(a,b)] = np.hstack((kernel(Xa[a], Xpa[b]), 
                                    np.zeros((Xa[a].shape[0], Xa_term[b].shape[0]))))
        else:
            Kab[(a,b)] = kernel(Xa[a], Xpa[b])
        Kab[(a,b)] = csr_matrix(Kab[(a,b)] * (np.abs(Kab[(a,b)]) > 1e-3))
        KabT[(a,b)] = Kab[(a,b)].T.tocsr()
    
    # build the K_a, D_a matrices
    Ka = [kernel(Xa[i], Xa[i])  for i in xrange(k)]
    Dainv = [csc_matrix(Ka[i] * (np.abs(Ka[i]) > 1e-3)) + lamb*scipy.sparse.eye(*Ka[i].shape) for i in xrange(k)]
#    print np.linalg.matrix_rank(Dainv[2].toarray()), Dainv[2].shape, np.linalg.cond(Dainv[2].toarray())
#    print np.linalg.eig(Dainv[2].toarray())[0]
#    print [np.linalg.eig(Dainv[i].toarray())[0].min() for i in xrange(3)]
#    plt.spy(Dainv[2].toarray())
#    plt.show()  
#    print Dainv[2].shape

    index = (squareform(pdist(Xa[2])) == 0.0).nonzero()    
    
#    print (squareform(pdist(Xa[2])) == 0.0).nonzero()
#    print Xa[2][index[0][:5],:]
#    print Xa[2][index[1][:5],:]
#    
#    splu(Dainv[0])    
#    cholesky(Dainv[0])
#    splu(Dainv[1])
#    cholesky(Dainv[1])
#    splu(Dainv[2])
    cholesky(Dainv[2])
        
    
    Da= [cholesky(Dainv[i]) for i in xrange(k)]
#    Da = [splu(Dainv[i]) for i in xrange(k)]
        
    # build K_ter matrix
    Kterma = [ np.hstack((kernel(ter_rew_samples[0], Xpa[i]),
                          np.zeros((ter_rew_samples[0].shape[0], Xa_term[i].shape[0])))) if Xa_term[i].size > 0
                else kernel(ter_rew_samples[0], Xpa[i]) for i in xrange(k)]
    K_ter = kernel(ter_rew_samples[0], ter_rew_samples[0])
    D_ter = cholesky(csc_matrix(K_ter*(np.abs(K_ter) > 1e-3)) + lamb*scipy.sparse.eye(*K_ter.shape))
#    D_ter = splu(csc_matrix(K_ter*(np.abs(K_ter) > 1e-3)) + lamb*scipy.sparse.eye(*K_ter.shape))
    R_ter = ter_rew_samples[1]
    
    return kernel, Kab, KabT, Da, Dainv, Ra, Kterma, D_ter, R_ter, Xa

def sample_mult_normal_given_precision(mu, precision, num_of_samples):
	from scikits.sparse import cholmod
	factor = cholmod.cholesky(precision)
	D = factor.D(); D = numpy.reshape(D, (len(D),1))
	samples = factor.solve_Lt(numpy.random.normal(size=(len(mu),num_of_samples))/numpy.sqrt(D))
	samples = numpy.repeat(mu, num_of_samples, axis=1)+ factor.apply_Pt(samples)
	return samples

def test_complex():
    c = complex_matrix()
    fc = cholesky(c)
    r = real_matrix()
    fr = cholesky(r)
    
    assert factor_of(fc, c)

    assert np.allclose(fc(np.arange(4)),
                       (c.todense().I * np.arange(4)[:, np.newaxis]).ravel())
    assert np.allclose(fc(np.arange(4) * 1j),
                       (c.todense().I * (np.arange(4) * 1j)[:, np.newaxis]).ravel())
    assert np.allclose(fr(np.arange(4)),
                       (r.todense().I * np.arange(4)[:, np.newaxis]).ravel())
    # If we did a real factorization, we can't do solves on complex arrays:
    assert_raises(CholmodError, fr, np.arange(4) * 1j)

def getModel(img,var,lmat,pmat,cmat,rmat,reg,niter=10,onlyRes=False):
    from scikits.sparse.cholmod import cholesky
    import numpy

    omat = pmat*lmat
    rhs = omat.T*(img/var)

    B = omat.T*cmat*omat

    res = 0.
    regs = [reg]
    lhs = B+regs[-1]*rmat

    F = cholesky(lhs)
    fit = F(rhs)
    for i in range(niter):
        res = fit.dot(rmat*fit)

        delta = reg*1e3
        lhs2 = B+(reg+delta)*rmat

        T = (2./delta)*(numpy.log(F.cholesky(lhs2).L().diagonal()).sum()-numpy.log(F.L().diagonal()).sum())
        reg = (omat.shape[0]-T*reg)/res
        if abs(reg-regs[-1])/reg<0.005:
            break
        regs.append(reg)
        lhs = B+regs[-1]*rmat
        F = F.cholesky(lhs)
        fit = F(rhs)
    print reg,regs
    res = -0.5*res*regs[-1] + -0.5*((omat*fit-img)**2/var).sum()
    if onlyRes:
        return res,reg
    model = (omat*fit)
    return res,reg,fit,model

def test_spinv():
    X = sparse.rand(500,500,density=0.03)
    K = (X.T * X).tocsc() + sparse.identity(500)
    for mode in ("simplicial", "supernodal"):
        L = cholesky(K, mode=mode)
        # Full inverse
        inv_K = L.inv().todense()
        for form in ("lower", "upper", "full"):
            # Sparse inverse
            spinv_K = L.spinv(form=form)
            # 1) Check that the non-zero elements of spinv_K are
            # correct
            spinv_K = spinv_K.tocoo()
            i = spinv_K.row
            j = spinv_K.col
            x = spinv_K.data
            assert(np.allclose(x, inv_K[i,j]))
            # 2) Check that the elements that are non-zero in K
            # are correct in spinv_K
            K = K.tocoo()
            i = K.row
            j = K.col
            spinv_K = spinv_K.todense()
            if form == "lower":
                Z = np.tril(inv_K)
            elif form == "upper":
                Z = np.triu(inv_K)
            else:
                Z = inv_K
            assert(np.allclose(spinv_K[i,j], Z[i,j]))

def IRLS(Y, W, nore, reData, estCov, curr_bbeta):
    N = W.shape[0]
    nofe = W.shape[1] - nore
    err = 1
    logistic_input = W*curr_bbeta;
    curr_pred = np.asarray(np.divide(np.exp(logistic_input),(1+np.exp(logistic_input))))
    inv_curr_covar = invrecovar(reData, estCov); # Construct the covariance matrix b given the estimate of covariance matrix of random effects
    hessian_from_logprior = block_diag((inv_curr_covar,np.zeros((nofe,nofe))))
    start_time = time.time()
    curr_b = curr_bbeta[np.arange(nore)] 
    while err>0.00001:
      old_b = curr_b;
      old_bbeta = curr_bbeta; 
      curr_gradient = -((W.T)*(Y-curr_pred) -sp.vstack(((inv_curr_covar*curr_b), np.zeros((nofe,1)))));
      hessain_from_loglik = (W.T)*(sp.csr_matrix(((curr_pred*(1-curr_pred)).ravel(), (np.arange(N), np.arange(N)))))*W;                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                
      curr_hessian = hessain_from_loglik+hessian_from_logprior;
      factor = cholesky(curr_hessian);
      delta_bbeta = factor(curr_gradient);
      delta_bbeta = np.reshape(delta_bbeta, (len(delta_bbeta),1));
      curr_bbeta = curr_bbeta - delta_bbeta;
      logistic_input = W*curr_bbeta;
      curr_pred = np.asarray(np.divide(np.exp(logistic_input),(1+np.exp(logistic_input))));
      curr_b = curr_bbeta[0:nore];
      curr_beta = curr_bbeta[nore:];
      err = np.sqrt(sum(np.square(curr_bbeta - old_bbeta))/len(old_bbeta));
      #print repr(err)
    pred_error = np.sqrt(sum(np.square(curr_pred - Y))/len(curr_pred)); # current prediction mean square error
    print 'Current training RMSE = '+repr(pred_error);
    hessian_from_loglik = (W.T)*(sp.csr_matrix(((curr_pred*(1-curr_pred)).ravel(), (np.arange(N), np.arange(N)))))*W;                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                
    curr_hessian = hessian_from_logprior+hessian_from_loglik;
    curr_b = sp.csr_matrix(curr_b);
    log_c1 = sp.csr_matrix((W*curr_bbeta).T)*Y - sum(np.log(1+np.exp(W*curr_bbeta))) - 0.5*curr_b.T*(inv_curr_covar)*(curr_b); # refer to sampling paper
    # log_lik_part = - 0.5*inv_curr_covar.shape[1]*np.log(2*np.pi);

    # for i in range(len(reData)):
    #   p = len(np.unique(reData[i][1]));
    #   (sign, logdet) = np.linalg.slogdet(estCov[i]);
    #   log_lik_part = log_lik_part - p*logdet;

    # log_lik = log_c1 + log_lik_part;
    factor1 = cholesky(curr_hessian);
    factor2 = cholesky(inv_curr_covar);
    log_lik = -factor1.logdet()+factor2.logdet();

    return ([curr_bbeta, curr_pred, curr_hessian, log_c1, log_lik])

def logdet_cholmod(R,U,rows,cols, psd_tolerance=1e-6, factor=None):
  from scikits.sparse.cholmod import cholesky
  if not is_psd_cholmod(R, U, rows, cols, psd_tolerance, factor):
    return -np.Inf
  X = build_sparse(R, U, rows, cols)
#  print factor
  filled_factor = cholesky(X) if factor is None else factor.cholesky(X)
  D = filled_factor.D()
  return np.sum(np.log(D))

 def cholsolve(A, b, fac=None):
     if fac is None:
         fac = cholmod.cholesky(A)
     else:
         # store a representation of the factorised matrix and update it in
         # place
         fac.cholesky_inplace(A)
     x = fac.solve_A(b)
     return x, fac

def is_psd_cholmod(R,U,rows,cols,tolerance=1e-6,factor=None):
  X = build_sparse(R, U, rows, cols)
  from scikits.sparse.cholmod import cholesky
  try:
    full_factor = cholesky(X) if factor is None else factor.cholesky(X)
  except:
    return False
  D = full_factor.D()
  return np.all(D>tolerance)

def test_deprecation():
    f = cholesky(sparse.eye(5, 5))
    b = np.ones(5)
    for dep_method in "solve_P", "solve_Pt":
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            getattr(f, dep_method)(b)
            assert len(w) == 1
            assert issubclass(w[-1].category, DeprecationWarning)
            assert "deprecated" in str(w[-1].message)

def _random_correlated_image(mean, sigma, image_shape, alpha=0.3, rng=None):
    """
    Creates a random image with correlated neighbors.
    pixel covariance is sigma^2, direct neighors pixel covariance is alpha * sigma^2.

    Parameters
    ----------
    mean : the mean value of the image pixel values.
    sigma : the std dev of image pixel values.
    image_shape : tuple, shape = (3, )
    alpha : the neighbors correlation factor.
    rng : random number generator (a numpy.random.RandomState instance).
    """
    dim_x, dim_y, dim_z = image_shape
    dim_image = dim_x * dim_y * dim_z

    correlated_image = 0
    for neighbor in [(1, 0, 0), (0, 1, 0), (0, 0, 1)]:
        corr_data = []
        corr_i = []
        corr_j = []
        for i, j, k in [(0, 0, 0), neighbor]:
            d2 = 1.0 * (i*i + j*j + k*k)
            ind = np.asarray(np.mgrid[0:dim_x-i, 0:dim_y-j, 0:dim_z-k], dtype=np.int)
            ind = ind.reshape((3, (dim_x - i) * (dim_y - j) * (dim_z - k)))
            corr_i.extend(np.ravel_multi_index(ind, (dim_x, dim_y, dim_z)).tolist())
            corr_j.extend(np.ravel_multi_index(ind + np.asarray([i, j, k])[:, None],
                                          (dim_x, dim_y, dim_z)).tolist())
            if i>0 or j>0 or k>0:
                corr_i.extend(np.ravel_multi_index(ind + np.asarray([i, j, k])[:, None],
                                              (dim_x, dim_y, dim_z)).tolist())
                corr_j.extend(np.ravel_multi_index(ind, (dim_x, dim_y, dim_z)).tolist())
            if i==0 and j==0 and k==0:
                corr_data.extend([3.0] * ind.shape[1])
            else:
                corr_data.extend([alpha * 3.0] * 2 * ind.shape[1])

        correlation = scisp.csc_matrix((corr_data, (corr_i, corr_j)), shape=(dim_image, dim_image))

        factor = cholesky(correlation)
        L = factor.L()
        P = factor.P()[None, :]
        P = scisp.csc_matrix((np.ones(dim_image),
                              np.vstack((P, np.asarray(range(dim_image))[None, :]))),
                             shape=(dim_image, dim_image))

        sq_correlation = P.dot(L)
        X = rng.normal(0, 1, dim_image)
        Y = sq_correlation.dot(X)
        Y = Y.reshape((dim_x, dim_y, dim_z))
        X = X.reshape((dim_x, dim_y, dim_z))
        correlated_image += Y
    correlated_image /= 3

    return correlated_image * sigma + mean

def test_spinv_complex():
    # Complex matrices do not work yet, but this test checks that some
    # reasonable error is produced
    C = complex_matrix()
    L = cholesky(C)
    try:
        # This should result CholmodError because spinv is not
        # implemented for complex matrices - yet
        invC = L.spinv()
        assert(False)
    except CholmodError:
        assert(True)

def get_posterior_params(W, V, M, K_inv):
	"""
	INPUT: 	W, sparse matrix format, same dimension as K_inv
			V, a numpy vector of size K_inv.shape[0]
			M: the mean vector in numpy format
			K_inv: precision matrix in sparse form
	"""
	from scikits.sparse import cholmod
	posterior_precision = K_inv + W
	posterior_precision_factor = cholmod.cholesky(posterior_precision)
	posterior_mean = posterior_precision_factor(V + K_inv*M)
	return (posterior_mean, posterior_precision) 

def test_update_downdate():
    m = real_matrix()
    f = cholesky(m)
    L = f.L()[f.P(), :]
    assert factor_of(f, m)
    f.update_inplace(L)
    assert factor_of(f, 2 * m)
    f.update_inplace(L)
    assert factor_of(f, 3 * m)
    f.update_inplace(L, subtract=True)
    assert factor_of(f, 2 * m)
    f.update_inplace(L, subtract=True)
    assert factor_of(f, m)

def test_convenience():
    A_dense_seed = np.array([[10, 0, 3, 0],
                             [0, 5, 0, -2],
                             [3, 0, 5, 0],
                             [0, -2, 0, 2]])
    for mode in ("simplicial", "supernodal"):
        for dtype in (float, complex):
            A_dense = np.array(A_dense_seed, dtype=dtype)
            A_sp = sparse.csc_matrix(A_dense)
            f = cholesky(A_sp, mode=mode)
            assert np.allclose(f.det(), np.linalg.det(A_dense))
            assert np.allclose(f.logdet(), np.log(np.linalg.det(A_dense)))
            assert np.allclose(f.slogdet(), [1, np.log(np.linalg.det(A_dense))])
            assert np.allclose((f.inv() * A_sp).todense(), np.eye(4))

 def __init__(self, verts, tris, m=1.0):
     self._verts = verts
     self._tris = tris
     # precompute some stuff needed later on
     e01 = verts[tris[:,1]] - verts[tris[:,0]]
     e12 = verts[tris[:,2]] - verts[tris[:,1]]
     e20 = verts[tris[:,0]] - verts[tris[:,2]]
     self._triangle_area = .5 * veclen(np.cross(e01, e12))
     unit_normal = normalized(np.cross(normalized(e01), normalized(e12)))
     self._un = unit_normal
     self._unit_normal_cross_e01 = np.cross(unit_normal, -e01)
     self._unit_normal_cross_e12 = np.cross(unit_normal, -e12)
     self._unit_normal_cross_e20 = np.cross(unit_normal, -e20)
     # parameters for heat method
     h = np.mean(map(veclen, [e01, e12, e20]))
     t = m * h ** 2
     # pre-factorize poisson systems
     Lc, vertex_area = compute_mesh_laplacian(verts, tris, area_type='lumped_mass')
     A = sparse.spdiags(vertex_area, 0, len(verts), len(verts))
     #self._factored_AtLc = splu((A - t * Lc).tocsc()).solve
     self._factored_AtLc = cholesky((A - t * Lc).tocsc(), mode='simplicial')
     #self._factored_L = splu(Lc.tocsc()).solve
     self._factored_L = cholesky(Lc.tocsc(), mode='simplicial')

def random_projection_cholmod_csc(X,k,factor=None):
  from scikits.sparse.cholmod import cholesky
  # TODO: add factor update
  full_factor = cholesky(X) if factor is None else factor.cholesky(X)
  (n,_) = X.shape
  D = full_factor.D()
  D12 = np.sqrt(D)
  Q = np.random.rand(k,n)
  Q[Q<0.5] = -1.0/np.sqrt(k)
  Q[Q>=0.5] = 1.0/np.sqrt(k)
  zA = (Q / D12).T
  zB = full_factor.solve_Lt(zA)
  A = full_factor.apply_Pt(zB).T
  return A

    def solve (self, recover_cov=False):
        n = len (self.nodes)
        L = sparse.coo_matrix ((self._data, (self._rows,self._cols)), 
                shape=(self._Ns*n,self._Ns*n))
        e = np.asarray (self._vector)
        factor = cholmod.cholesky (L.tocsc ())

        x = factor (e)
        self.x = x.reshape (n,self._Ns)

        if recover_cov:
            S = factor (np.eye (self._Ns*n))
            Smarginal = [S[ii:ii+2,ii:ii+2].ravel ().tolist () for ii in range (0,self._Ns*n,self._Ns)]
            self.cov = np.asarray (Smarginal)
        else: self.cov = np.zeros ((n,4))

    def __init__(self, A, **kwargs):
        """CholeskyOperator corresponding to the sparse matrix `A`.

        CholeskyOperator(A)

        The sparse matrix `A` must be in one of the SciPy sparse
        matrix formats.
        """
        m, n = A.shape
        if m != n:
            raise ValueError("Input matrix must be square")

        self.__factor = cholesky(A)

        super(CholeskyOperator, self).__init__(n, n, matvec=self.cholesky_matvec, symmetric=True, **kwargs)