from scipy.sparse import eye
from scipy.io import mmread
from matplotlib import pyplot as plt

matFile='apache2.mtx'
M = mmread(matFile)
rows = M.shape[0]
cols = M.shape[1]
A = M + eye(rows, cols) * 10000.0
plt.title("A")
plt.spy(A, precision = 1e-2, marker = '.', markersize = 0.01)
plt.show()

from modshogun import RealSparseMatrixOperator, LanczosEigenSolver

op = RealSparseMatrixOperator(A.tocsc())

# Lanczos iterative Eigensolver to compute the min/max Eigenvalues which is required to compute the shifts
eigen_solver = LanczosEigenSolver(op)
# we set the iteration limit high to compute the eigenvalues more accurately, default iteration limit is 1000
eigen_solver.set_max_iteration_limit(2000)

# computing the eigenvalues
eigen_solver.compute()
print 'Minimum Eigenvalue:', eigen_solver.get_min_eigenvalue()
print 'Maximum Eigenvalue:', eigen_solver.get_max_eigenvalue()

# We can specify the power of the sparse-matrix that is to be used for coloring, default values will apply a
# 2-distance greedy graph coloring algorithm on the sparse-matrix itself. Matrix-power, if specified, is computed in O(lg p)
from modshogun import ProbingSampler

trace_sampler = ProbingSampler(op)
# apply the graph coloring algorithm and generate the number of colors, i.e. number of trace samples
trace_sampler.precompute()
print 'Number of colors used:', trace_sampler.get_num_samples()

from modshogun import SerialComputationEngine, CGMShiftedFamilySolver, LogRationalApproximationCGM

engine = SerialComputationEngine()
cgm = CGMShiftedFamilySolver()
# setting the iteration limit (set this to higher value for higher condition number)
cgm.set_iteration_limit(100)

# accuracy determines the number of contour points in the rational approximation (i.e. number of shifts in the systems)
accuracy = 1E-15

# we create a operator-log-function using the sparse matrix operator that uses CG-M to solve the shifted systems
op_func = LogRationalApproximationCGM(op, engine, eigen_solver, cgm, accuracy)
op_func.precompute()
print 'Number of shifts:', op_func.get_num_shifts()

from modshogun import LogDetEstimator

# number of log-det samples (use a higher number to get better estimates)
# (this is 5 times number of colors estimate in practice, so usually 1 probing estimate is enough)
num_samples = 5

log_det_estimator = LogDetEstimator(trace_sampler, op_func, engine)
estimates = log_det_estimator.sample(num_samples)

estimated_logdet = mean(estimates)
print 'Estimated log(det(A)):', estimated_logdet

# the following method requires massive amount of memory, for demonstration purpose
# the following code is commented out and direct value obtained from running it once is used

# from modshogun import Statistics
# actual_logdet = Statistics.log_det(A)

actual_logdet = 7120357.73878
print 'Actual log(det(A)):', actual_logdet

plt.hist(estimates)
plt.plot([actual_logdet, actual_logdet], [0,len(estimates)], linewidth=3)
plt.show()

from scipy.sparse import csc_matrix

m = mmread('west0479.mtx')
# computing a spd with added ridge
B = csc_matrix(m.transpose() * m + identity(m.shape[0]) * 1000.0)

fig = plt.figure(figsize=(12, 4))
ax = fig.add_subplot(1,2,1)
ax.set_title('B')
ax.spy(B, precision = 1e-5, marker = '.', markersize = 2.0)

ax = fig.add_subplot(1,2,2)
ax.set_title('lower Cholesky factor')
dense_matrix = B.todense()
L = cholesky(dense_matrix)
ax.spy(csc_matrix(L), precision = 1e-5, marker = '.', markersize = 2.0)

plt.show()

op = RealSparseMatrixOperator(B)
eigen_solver = LanczosEigenSolver(op)

# computing log-det estimates using probing sampler
probing_sampler = ProbingSampler(op)
cgm.set_iteration_limit(500)

op_func = LogRationalApproximationCGM(op, engine, eigen_solver, cgm, 1E-5)
log_det_estimator = LogDetEstimator(probing_sampler, op_func, engine)
num_probing_estimates = 100
probing_estimates = log_det_estimator.sample(num_probing_estimates)

# computing log-det estimates using Gaussian sampler
from modshogun import NormalSampler, Statistics

num_colors = probing_sampler.get_num_samples()
normal_sampler = NormalSampler(op.get_dimension())
log_det_estimator = LogDetEstimator(normal_sampler, op_func, engine)
num_normal_estimates = num_probing_estimates * num_colors
normal_estimates = log_det_estimator.sample(num_normal_estimates)

# average in groups of n_effective_samples
effective_estimates_normal = zeros(num_probing_estimates)
for i in range(num_probing_estimates):
    idx = i * num_colors
    effective_estimates_normal[i] = mean(normal_estimates[idx:(idx + num_colors)])

actual_logdet = Statistics.log_det(B)
print 'Actual log(det(B)):', actual_logdet
print 'Estimated log(det(B)) using probing sampler:', mean(probing_estimates)
print 'Estimated log(det(B)) using Gaussian sampler:', mean(effective_estimates_normal)
print 'Variance using probing sampler:',var(probing_estimates)
print 'Variance using Gaussian sampler:',var(effective_estimates_normal)

fig = plt.figure(figsize=(15, 4))
ax = fig.add_subplot(1,3,1)
ax.set_title('Probing sampler')
ax.plot(cumsum(probing_estimates)/(arange(len(probing_estimates))+1))
ax.plot([0,len(probing_estimates)], [actual_logdet, actual_logdet])
ax.legend(["Probing", "True"])

ax = fig.add_subplot(1,3,2)
ax.set_title('Gaussian sampler')
ax.plot(cumsum(effective_estimates_normal)/(arange(len(effective_estimates_normal))+1))
ax.plot([0,len(probing_estimates)], [actual_logdet, actual_logdet])
ax.legend(["Gaussian", "True"])

ax = fig.add_subplot(1,3,3)
ax.hist(probing_estimates)
ax.hist(effective_estimates_normal)
ax.plot([actual_logdet, actual_logdet], [0,len(probing_estimates)], linewidth=3)
plt.show()

from scipy.io import loadmat

def get_Q_y_A(kappa):
    # read the ozone data and create the matrix Q
    ozone = loadmat('ozone_data.mat')
    GiCG = ozone["GiCG"]
    G = ozone["G"]
    C0 = ozone["C0"]
    kappa = 13.1
    Q = GiCG + 2 * (kappa ** 2) * G + (kappa ** 4) * C0
    # also, added a ridge here
    Q = Q + eye(Q.shape[0], Q.shape[1]) * 10000.0
    
    plt.spy(Q, precision = 1e-5, marker = '.', markersize = 1.0)
    plt.show()
    
    # read y and A
    y = ozone["y_ozone"]
    A = ozone["A"]
    return Q, y, A
    
def log_det(A):
    op = RealSparseMatrixOperator(A)
    engine = SerialComputationEngine()
    eigen_solver = LanczosEigenSolver(op)
    probing_sampler = ProbingSampler(op)
    cgm = CGMShiftedFamilySolver()
    cgm.set_iteration_limit(100)
    op_func = LogRationalApproximationCGM(op, engine, eigen_solver, cgm, 1E-5)
    log_det_estimator = LogDetEstimator(probing_sampler, op_func, engine)
    num_estimates = 1
    return mean(log_det_estimator.sample(num_estimates))

def log_likelihood(tau, kappa):
    Q, y, A = get_Q_y_A(kappa)
    n = len(y);
    AtA = A.T.dot(A)
    M = Q + tau * AtA;

    # Computing log-determinants")
    logdet1 = log_det(Q)
    logdet2 = log_det(M)
        
    first = 0.5 * logdet1 + 0.5 * n * log(tau) - 0.5 * logdet2
    
    # computing the rest of the likelihood
    second_a = -0.5 * tau * (y.T.dot(y))
    
    second_b = array(A.T.dot(y))
    from scipy.sparse.linalg import spsolve
    second_b = spsolve(M, second_b)
    second_b = A.dot(second_b)
    second_b = y.T.dot(second_b)
    second_b = 0.5 * (tau ** 2) * second_b
        
    log_det_part = first
    quadratic_part = second_a + second_b
    const_part = -0.5 * n * log(2 * pi)
      
    log_marignal_lik = const_part + log_det_part + quadratic_part

    return log_marignal_lik

L = log_likelihood(1.0, 15.0)
print 'Log-likelihood estimate:', L

from modshogun import RealSparseMatrixOperator, ComplexDenseMatrixOperator

dim = 5
random.seed(10)
# create a random valued sparse matrix linear operator
A = csc_matrix(random.randn(dim, dim))
op = RealSparseMatrixOperator(A)

# creating a random vector
random.seed(1)
b = array(random.randn(dim))
v = op.apply(b)
print 'A.apply(b)=',v

# create a dense matrix linear operator
B = array(random.randn(dim, dim)).astype(complex)
op = ComplexDenseMatrixOperator(B)
print 'Dimension:', op.get_dimension()

from scipy.sparse import csc_matrix
from scipy.sparse import identity
from modshogun import ConjugateGradientSolver

# creating a random spd matrix
dim = 5
random.seed(10)
m = csc_matrix(random.randn(dim, dim))
a = m.transpose() * m + identity(dim)
Q = RealSparseMatrixOperator(a)

# creating a random vector
y = array(random.randn(dim))

# solve the system Qx=y
# the argument is set as True to gather convergence statistics (default is False)
cg = ConjugateGradientSolver(True)
cg.set_iteration_limit(20)
x = cg.solve(Q,y)

print 'x:',x

# verifying the result
print 'y:', y
print 'Qx:', Q.apply(x)

residuals = cg.get_residuals()
plt.plot(residuals)
plt.show()

from modshogun import ComplexSparseMatrixOperator
from modshogun import ConjugateOrthogonalCGSolver

# creating a random spd matrix
dim = 5
random.seed(10)
m = csc_matrix(random.randn(dim, dim))
a = m.transpose() * m + identity(dim)
a = a.astype(complex)

# adding a complex entry along the diagonal
for i in range(0, dim):
    a[i,i] += complex(random.randn(), random.randn())
Q = ComplexSparseMatrixOperator(a)

z = array(random.randn(dim))
# solve for the system Qx=z
cocg = ConjugateOrthogonalCGSolver(True)
cocg.set_iteration_limit(20)
x = cocg.solve(Q, z)

print 'x:',x

# verifying the result
print 'z:',z
print 'Qx:',real(Q.apply(x))

residuals = cocg.get_residuals()
plt.plot(residuals)
plt.show()

from modshogun import CGMShiftedFamilySolver

cgm = CGMShiftedFamilySolver()

# creating a random spd matrix
dim = 5
random.seed(10)
m = csc_matrix(random.randn(dim, dim))
a = m.transpose() * m + identity(dim)
Q = RealSparseMatrixOperator(a)

# creating a random vector
v = array(random.randn(dim))

# number of shifts (will be equal to the number of contour points)
num_shifts = 3;

# generating some random shifts
shifts = []
for i in range(0, num_shifts):
    shifts.append(complex(random.randn(), random.randn()))
sigma = array(shifts)
print 'Shifts:', sigma

# generating some random weights
weights = []
for i in range(0, num_shifts):
    weights.append(complex(random.randn(), random.randn()))
alpha = array(weights)
print 'Weights:',alpha

# solve for the systems
cgm = CGMShiftedFamilySolver(True)
cgm.set_iteration_limit(20)
x = cgm.solve_shifted_weighted(Q, v, sigma, alpha)

print 'x:',x

residuals = cgm.get_residuals()
plt.plot(residuals)
plt.show()

# verifying the result with cocg
x_s = array([0+0j] * dim)
for i in range(0, num_shifts):
    a_s = a.astype(complex)
    for j in range(0, dim):
        # moving the complex shift inside the operator
        a_s[j,j] += sigma[i]
    Q_s = ComplexSparseMatrixOperator(a_s)
    # multiplying the result with weight
    x_s += alpha[i] * cocg.solve(Q_s, v)
print 'x\':', x_s

from modshogun import DirectSparseLinearSolver

# creating a random spd matrix
dim = 5
random.seed(10)
m = csc_matrix(random.randn(dim, dim))
a = m.transpose() * m + identity(dim)
Q = RealSparseMatrixOperator(a)

# creating a random vector
y = array(random.randn(dim))

# solve the system Qx=y
chol = DirectSparseLinearSolver()
x = chol.solve(Q,y)

print 'x:',x

# verifying the result
print 'y:', y
print 'Qx:', Q.apply(x)

from modshogun import DirectLinearSolverComplex

# creating a random spd matrix
dim = 5
random.seed(10)
m = array(random.randn(dim, dim))
a = m.transpose() * m + identity(dim)
a = a.astype(complex)

# adding a complex entry along the diagonal
for i in range(0, dim):
    a[i,i] += complex(random.randn(), random.randn())
Q = ComplexDenseMatrixOperator(a)

z = array(random.randn(dim))
# solve for the system Qx=z
solver = DirectLinearSolverComplex()
x = solver.solve(Q, z)

print 'x:',x

# verifying the result
print 'z:',z
print 'Qx:',real(Q.apply(x))