
def build_prompt_generic(model_entry):
    ## NEW
    template_code = model_entry["template_code"]  # from registry
    description = model_entry["description"]

    prompt = f""" # Complete the PyMC model definition within the 'with pm.Model() as m:' block below. Your output must define a complete Bayesian model with appropriate priors, likelihood, and then sample the posterior using, `pm.sample(1000, tune=1000, chains=4, return_inferencedata=True, idata_kwargs={{"log_likelihood": True}})`. Do not include any extra commentary or text outside the code. Follow best practices for expert-level Bayesian modeling.
    
# Description: {description}

    {template_code}    
    """

    return prompt

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

# Data
y = np.array([28, 8, -3, 7, -1, 1, 18, 12])
sigma = np.array([15, 10, 16, 11, 9, 11, 10, 18])

with pm.Model() as m:
"""
# Reference: https://www.pymc.io/projects/docs/en/stable/learn/core_notebooks/model_comparison.html

# ---------------------------------------------------------------------------
# Example: Dugongs Growth Model
# ---------------------------------------------------------------------------


gp_template = """
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:
"""

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

# Data
X = np.array([1, 1.5, 1.5, 1.5, 2.5, 4, 5, 5, 7, 8, 8.5, 9, 9.5, 9.5, 10, 12, 12, 13, 13, 14.5, 15.5, 15.5, 16.5, 17, 22.5, 29, 31.5])
y = np.array([1.8, 1.85, 1.87, 1.77, 2.02, 2.27, 2.15, 2.26, 2.47, 2.19, 2.26, 2.4, 2.39, 2.41, 2.5, 2.32, 2.32, 2.43, 2.47, 2.56, 2.65, 2.47, 2.64, 2.56, 2.7, 2.72, 2.57])

with pm.Model() as m:
"""

# Reference: https://areding.github.io/6420-pymc/unit6/Unit6-missingdata.html

# ---------------------------------------------------------------------------
# Example: Generalized Linear Model (GLM)
# ---------------------------------------------------------------------------
glm_template = """
import pymc as pm
import numpy as np
import arviz as az
import matplotlib.pyplot as plt

# Given Data
nyears = 40  # Number of years
year = np.array([-0.95, -0.9, -0.85, -0.8, -0.75, -0.7, -0.65, -0.6, -0.55, -0.5, -0.45, -0.4, -0.35, -0.3, -0.25, -0.2, -0.15, -0.1, -0.05, 0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1])
C = np.array([27, 42, 35, 55, 61, 19, 41, 74, 43, 42, 73, 37, 48, 49, 19, 72, 30, 18, 31, 71, 63, 51, 48, 73, 49, 54, 43, 59, 30, 24, 62, 55, 51, 47, 14, 27, 45, 20, 26, 19])
N = np.array([43, 83, 53, 91, 95, 24, 62, 91, 64, 57, 97, 56, 74, 66, 28, 92, 40, 23, 46, 96, 91, 75, 71, 100, 72, 77, 64, 68, 43, 32, 97, 92, 75, 84, 22, 58, 81, 37, 45, 39])

with pm.Model() as m:
"""


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

# Given Data
N = 12  # Number of observations
n = np.array([47, 148, 119, 810, 211, 196, 148, 215, 207, 97, 256, 360])
r = np.array([0, 18, 8, 46, 8, 13, 9, 31, 14, 8, 29, 24])

with pm.Model() as m:
"""

# ---------------------------------------------------------------------------
# Model Registry
# ---------------------------------------------------------------------------

datas_info = [
    {
        "name": "eight_schools",
        "template_code": eight_schools_template,
        "reference_samples_path": "reference_samples/eight_schools.nc",
        "description": "A hierarchical model for the 8-schools data.",
        "reference": "https://www.pymc.io/projects/docs/en/stable/learn/core_notebooks/model_comparison.html"
    },
    {
        "name": "dugongs",
        "template_code": dugongs_template,
        "reference_samples_path": "reference_samples/dugongs.nc",
        "description": "A growth model for dugongs with missing data.",
        "reference": "https://areding.github.io/6420-pymc/unit6/Unit6-missingdata.html"
    },
    {
        "name": "gp",
        "template_code": gp_template,
        "reference_samples_path": "reference_samples/rat_tumor.nc",
        "description": "Simulated data from a Poisson GP model.",
        "reference": "https://github.com/stan-dev/posteriordb/blob/master/posterior_database/models/stan/gp_regr.stan"
    },
    {
        "name": "glm",
        "template_code": glm_template,
        "reference_samples_path": "reference_samples/rat_tumor.nc",
        "description": "Simulated population counts of peregrines in the French Jura over 9 years",
        "reference": "https://github.com/stan-dev/posteriordb/blob/master/posterior_database/models/stan/gp_regr.stan"
    },
    {
        "name": "surgical",
        "template_code": surgical_template,
        "reference_samples_path": "reference_samples/surgical.nc",
        "description": "The mortality rates in 12 hospitals performing cardiac surgery on babies.",
        "reference": "https://github.com/stan-dev/posteriordb/blob/master/posterior_database/models/stan/surgical_model.stan"
    },
]
