import numpy as np
import pymc as pm
import pytensor.tensor as pt
import numpy.testing as npt
from tqdm import tqdm

def run_reproduction(seed):
    rng = np.random.RandomState(seed)
    n_outputs = 2
    X = rng.randn(20, 3)
    y = rng.randn(n_outputs, 20)
    Xnew = rng.randn(30, 3)
    pnew = rng.randn(n_outputs, 30)

    with pm.Model() as latent_model:
        cov_func = pm.gp.cov.ExpQuad(3, [0.1, 0.2, 0.3])
        mean_func = pm.gp.mean.Constant(0.5)
        latent_gp = pm.gp.Latent(mean_func=mean_func, cov_func=cov_func)
        latent_f = latent_gp.prior("f", X, n_outputs=n_outputs, reparameterize=True)
        latent_p = latent_gp.conditional("p", Xnew)

    with pm.Model() as marginal_model:
        cov_func = pm.gp.cov.ExpQuad(3, [0.1, 0.2, 0.3])
        mean_func = pm.gp.mean.Constant(0.5)
        marginal_gp = pm.gp.Marginal(mean_func=mean_func, cov_func=cov_func)

        marginal_f = marginal_gp.marginal_likelihood("f", X, y, sigma=0.0)
        marginal_p = marginal_gp.conditional("p", Xnew)

    lp_latent = latent_p.eval({latent_f: y})
    lp_marginal = marginal_p.eval()

    npt.assert_allclose(lp_latent, lp_marginal, rtol=1e-4)

if __name__ == "__main__":
    iterations = 500    
    for i in tqdm(range(iterations)):
        try:
            run_reproduction(seed=i)
        except AssertionError as e:
            print(f"\nCaught failure at seed: {i}")
            print(e)
            break