
import pymc as pm
import numpy as np
import arviz as az
import matplotlib.pyplot as plt

# Given Data
N = 11  # Number of observations
x = np.array([-10, -8, -6, -4, -2, 0, 2, 4, 6, 8, 10])
y = np.array([4.75906, 1.59423, 2.99548, 5.27501, 1.66472, 2.24347, 2.8914, 4.08681, 4.60588, 0.802364, 3.92136])
k = np.array([40, 37, 29, 12, 4, 3, 9, 19, 77, 82, 33])

with pm.Model() as m:
	μ = pm.Normal("μ", mu=0, sigma=10)
	σ = pm.HalfNormal("σ", sigma=10)
	λ = pm.Deterministic("λ", pm.math.exp(μ))

	# Define the Poisson likelihood
	obs = pm.Poisson("obs", mu=λ, observed=y)

	# Define the Gaussian Process
	kernel = pm.gp.cov.ExpQuad(1, ls=1)
	gp = pm.gp.Latent(cov_func=kernel)
	f = gp.prior("f", X=x)

	# Define the GP likelihood
	gp_like = pm.Poisson("gp_like", mu=f, observed=k)

	# Sample the posterior
	trace = pm.sample(1000, tune=1000, chains=4, return_inferencedata=True, idata_kwargs={"log_likelihood": True})

	# Posterior diagnostics
	summary = az.summary(trace)