
def build_prompt_generic(model_entry):
    # template_code = model_entry["template_code"]  # from registry
    # prompt = f""" # Complete the missing part of the following code template using PyMC such that you obtain the posterior samples via pm.sample(). Do not include any additional text, explanations, or code beyond what is necessary. The missing model block should appear inside `with pm.Model() as m:`.
    # {template_code}
    # """

    ## NEW
    template_code = model_entry["template_code"]  # from registry
    description = model_entry["description"]
    # prompt = f""" # Complete the missing part of the following code template using PyMC such that you obtain the posterior samples through `pm.sample(1000, tune = 1000, chains = 4, return_inferencedata=True, idata_kwargs={{"log_likelihood": True}})`. Do not include any additional text, explanations, or code beyond what is necessary. The missing model block should appear inside `with pm.Model() as m:`.
    # {template_code}
    # """

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

mesquite_template = """
import pymc as pm
import numpy as np
import arviz as az
import matplotlib.pyplot as plt
N = 46

canopy_height = np.array([
    1, 1.33, 0.6, 1.2, 1.05, 1, 0.9, 1.3, 0.6, 0.8, 0.6, 0.95, 1.2, 0.7, 1.2, 1.1,
    0.64, 1.2, 1.2, 1.3, 1, 0.9, 1, 0.5, 0.9, 0.6, 1.4, 2.5, 1.4, 1.3, 1.1, 1.5,
    0.8, 1.1, 1.3, 0.85, 1.5, 1.2, 1.15, 1.5, 0.7, 1.9, 1.5, 1.5, 1.5, 0.62
])

density = np.array([
    1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 1,
    5, 9, 1, 1, 1, 3, 1, 3, 7, 1, 2, 2, 2, 3, 1, 1, 2, 2, 1, 1
])

diam1 = np.array([
    1.8, 1.7, 2.8, 1.3, 3.3, 1.4, 1.5, 3.9, 1.8, 2.1, 0.8, 1.3, 1.2, 1.5, 2.8, 1.4,
    1.5, 2.4, 1.9, 2.3, 2.1, 2.4, 1, 1.3, 1.1, 1.3, 2.5, 5.2, 2, 1.6, 1.4, 3.2, 1.9,
    2.4, 2.5, 2.1, 2.4, 2.4, 1.9, 2.7, 1.3, 2.9, 2.1, 4.1, 2.8, 1.27
])

diam2 = np.array([
    1.15, 1.35, 2.55, 0.85, 1.9, 1.4, 0.5, 2.3, 1.35, 1.6, 0.63, 0.95, 0.9, 0.7, 1.7,
    0.85, 0.6, 2.4, 1.55, 1.6, 1.7, 1.3, 0.4, 0.6, 0.7, 1.2, 2.3, 4, 1.6, 1.6, 1, 1.9,
    1.8, 2.4, 1.8, 1.5, 2.2, 1.7, 1.2, 2.5, 1.1, 2.7, 1, 3.8, 2.5, 1
])

group = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

total_height = np.array([
    1.3, 1.35, 2.16, 1.8, 1.55, 1.2, 1, 1.7, 0.8, 1.2, 0.9, 1.35, 1.4, 1, 1.7, 1.5,
    0.65, 1.5, 1.7, 1.7, 1.5, 1.5, 1.2, 0.7, 1.2, 0.8, 1.7, 3, 1.7, 1.6, 1.5, 1.9, 1.1,
    1.6, 2, 1.25, 2, 1.3, 1.45, 2.2, 0.7, 1.9, 1.8, 2, 2.2, 0.92
])

weight =  np.array([
    401.3, 513.7, 1179.2, 308, 855.2, 268.7, 155.5, 1253.2, 328, 614.6, 60.2, 269.6,
    448.4, 120.4, 378.7, 266.4, 138.9, 1020.8, 635.7, 621.8, 579.8, 326.8, 66.7, 68,
    153.1, 256.4, 723, 4052, 345, 330.9, 163.5, 1160, 386.6, 693.5, 674.4, 217.5, 771.3,
    341.7, 125.7, 462.5, 64.5, 850.6, 226, 1745.1, 908, 213.5
])

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

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


# eight_school, dugongs, ark, gp, glm, surgical

models_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"
    },
    # {
    #     "name": "mesquite",
    #     "template_code": mesquite_template,
    #     "reference_samples_path": "reference_samples/mesquite.nc",
    #     "description": "Total production (biomass) of mesquite leaves as function of easily measured parameters of the plant, before actual harvesting takes place. Two groups: 26 mesquite bushes and 20 mesquite bushes measured at a different time of year.",
    #     "reference": "https://github.com/stan-dev/posteriordb/blob/master/posterior_database/models/stan/surgical_model.stan"
    # }
]
