from git import Repo
repo = Repo('./CUQIpy')
branch = repo.active_branch
# Sample with NUTS + time it
import time
import numpy as np
from cuqi.testproblem import Deconvolution1D
from cuqi.distribution import Gaussian, Gamma, JointDistribution, GMRF
from cuqi.experimental.mcmc import NUTS
import matplotlib.pyplot as plt

# Inverse problem
np.random.seed(0) 
testproblem_obj = Deconvolution1D(dim=128, phantom='sinc', noise_std=0.001)
posterior = testproblem_obj.posterior

# NUTS
Ns=15
Nb=15
print(branch.name)
print('NUTS')
start = time.time()
np.random.seed(0)
NUTS_sampler = NUTS(posterior, step_size=None)
NUTS_sampler.warmup(Nb, tune_freq=1/Nb).sample(Ns)

end = time.time()
print('Time:', end - start)
print('samples norm:', np.linalg.norm(NUTS_sampler.get_samples().samples))
plt.plot([np.linalg.norm(NUTS_sampler.get_samples().samples[:,i]) for i in range(Ns+Nb)])
plt.ylim([0, 14])
# add branch name in title
plt.title(branch.name)
plt.figure()
plt.plot(NUTS_sampler.num_tree_node_list)
plt.title(branch.name)

import cuqi
import numpy as np
from cuqi.distribution import Gamma, Gaussian, GMRF, JointDistribution, LMRF
from cuqi.experimental.mcmc import NUTS, HybridGibbs, Conjugate, LinearRTO, ConjugateApprox, UGLA
from cuqi.testproblem import Deconvolution1D
from git import Repo
import matplotlib.pyplot as plt
import time
repo = Repo('../')
branch = repo.active_branch

# Forward problem
np.random.seed(0)
A, y_data, info = Deconvolution1D(dim=128, phantom='sinc', noise_std=0.001).get_components()

# Bayesian Inverse Problem
s = Gamma(1, 1e-4)
x = GMRF(np.zeros(A.domain_dim), 50)
y = Gaussian(A@x, lambda s: 1/s)

# Posterior
target = JointDistribution(y, x, s)(y=y_data)

Nb=40
sampling_strategy = {
    "x" : NUTS(max_depth=7),
    "s" : Conjugate()
}

# Here we do 10 internal steps with NUTS for each Gibbs step
num_sampling_steps = {
    "x" : 1,
    "s" : 1
}

sampler = HybridGibbs(target, sampling_strategy, num_sampling_steps)
# start time
start_time = time.time()
sampler.warmup(Nb)
sampler.sample(40)
samples = sampler.get_samples()
# end time
end_time = time.time()
print(end_time - start_time)


samples["x"].plot_ci(exact=info.exactSolution)
plt.title(branch.name +" time: "+str(end_time - start_time))