from pylab import *

from shogun.Mathematics import LogDetEstimator
from shogun.Mathematics import ProbingSampler
from shogun.Mathematics import NormalSampler
from shogun.Library import SerialComputationEngine
from shogun.Mathematics import LogRationalApproximationCGM
from shogun.Mathematics import RealSparseMatrixOperator
from shogun.Mathematics import LanczosEigenSolver
from shogun.Mathematics import CCGMShiftedFamilySolver
from shogun.Mathematics import Statistics

from scipy.sparse import csc_matrix
from scipy.sparse import spdiags

# create diagonal matrix with challenging eigenspectrum
n=100
difficulty_level=2 # should be in [1,infty]
min_eig=1e-10
data=abs(randn(n)**difficulty_level)+min_eig
plot(sort(data))
title("Eigenspectrum")
matrix=csc_matrix(spdiags(data, 0, n, n))
true_log_det=sum(log(data))

# create shogun representation and prepare things for log-det sampler
op=RealSparseMatrixOperator(csc_matrix(matrix))
engine=SerialComputationEngine()
linear_solver=CCGMShiftedFamilySolver()
accuracy=1e-5
eigen_solver=LanczosEigenSolver(op)
eigen_solver.compute()
op_func=LogRationalApproximationCGM(op, engine, eigen_solver, linear_solver, accuracy)

print "computed eigenvalues:", eigen_solver.get_min_eigenvalue(), eigen_solver.get_max_eigenvalue()
print "true eigenvalues:", min(data), max(data)

def plot_estimate_convergence(estimates, true_value):
    plot(cumsum(estimates)/(arange(n_estimates)+1))
    plot([0,n_estimates], [true_value, true_value])
    legend(["Estimates", "True"])

# log det with probing sampler
trace_sampler=ProbingSampler(op)
log_det_estimator=LogDetEstimator(trace_sampler, op_func, engine)
n_estimates=10
estimates=log_det_estimator.sample(n_estimates)
plot_estimate_convergence(estimates, true_log_det)
title("Probe sampler")
print "Probe sampler:", mean(estimates)
print "True Value", true_log_det
print estimates

# log det with normal sampler
trace_sampler=NormalSampler(n)
log_det_estimator=LogDetEstimator(trace_sampler, op_func, engine)
n_estimates=100
estimates=log_det_estimator.sample(n_estimates)
plot_estimate_convergence(estimates, true_log_det)
title("Normal sampler")
print "Normal sampler:", mean(estimates)
print "True Value", true_log_det

normal_sampler=NormalSampler(n)
normal_sampler.precompute()
print normal_sampler.sample(0)

probing_sampler=ProbingSampler(op)
probing_sampler.precompute()
print probing_sampler.sample(0)