import pymc as pm
import numpy as np
import math
from matplotlib import pyplot as plt
from IPython.core.pylabtools import figsize

figsize(11, 9)
%matplotlib inline

time = 10
sampling_freq = 30
pi2 = math.pi*2
x = np.linspace(0, time, time*sampling_freq)

# Simulate some data
signal = 5*np.sin(x*pi2*2 + pi2/4)

plt.title("Original signal")
plt.plot(x, signal)

# Corrupt the signal
y = signal + np.random.randn(len(x))

plt.title("Corrupted signal")
plt.plot(x, y)

# Probabilistic Model
A = pm.Uniform("A", 0, 100)
f = pm.Uniform("f", 0, sampling_freq/2) # Nyquist
p = pm.Uniform("p", 0, pi2)
var = pm.Uniform("tau", 0, 2)

@pm.deterministic
def mu_(x=x, A=A, f=f, p=p):
    return A*np.sin(x*f*pi2 + p)

observation = pm.Normal("observation", mu=mu_, tau=1/(var*var), value=y, observed=True)

model = pm.Model([A, f, p, var, observation])
mcmc = pm.MCMC(model)
mcmc.sample(40000, 10000, 1)

# Reconstructed signal
A_m = mcmc.trace('A')[:].mean()
f_m = mcmc.trace('f')[:].mean()
p_m = mcmc.trace('p')[:].mean()

print "A ~ " + str(A_m) + ", f ~ " + str(f_m) + ", p ~ " + str(p_m)

r = A_m*np.sin(x*f_m*pi2 + p_m)
plt.title("Reconstructed signal")
plt.plot(x, r)