import pymc as pm
import numpy as np

def model(data):
    """
    Bayesian von Bertalanffy growth model with Log-Normal likelihood.
    
    This model addresses the critique that a simple Normal likelihood with constant
    variance may not capture the heteroscedasticity and skewness in dugong length data.
    By using a Log-Normal likelihood, we assume the errors are multiplicative,
    which is often more appropriate for biological growth data.
    """
    age = data['age']
    length = data['length']
    
    with pm.Model() as dugong_model:
        # Data container for age
        age_data = pm.Data('age', age, dims="obs_id")
        
        # Priors for von Bertalanffy growth parameters
        # L_inf: Asymptotic length. HalfNormal with sigma=3.0 allows for lengths up to ~6m (3 sigma), 
        # which is biologically plausible for dugongs.
        L_inf = pm.HalfNormal('L_inf', sigma=3.0)
        
        # k: Growth rate coefficient. HalfNormal with sigma=2.0 is weakly informative.
        k = pm.HalfNormal('k', sigma=2.0)
        
        # t0: Theoretical age at zero length. Normal centered at 0 with sigma=1.0.
        # This allows for negative t0 (early growth) or positive t0 (delayed start).
        t0 = pm.Normal('t0', mu=0.0, sigma=1.0)
        
        # Expected length according to von Bertalanffy growth function
        vb_mean = L_inf * (1 - pm.math.exp(-k * (age_data - t0)))
        
        # Ensure vb_mean is positive to take log. 
        # In practice, with reasonable priors and data, this should be positive.
        # We add a small epsilon for numerical stability if needed, but 
        # the model should constrain it naturally.
        mu_log = pm.math.log(vb_mean)
        
        # Observation model: Log-Normal likelihood
        # sigma is the standard deviation on the log scale.
        sigma = pm.HalfNormal('sigma', sigma=0.5)
        
        # Observed likelihood
        y_obs = pm.LogNormal('y_obs', mu=mu_log, sigma=sigma, observed=length, dims="obs_id")
    
    return dugong_model

def gen_model(observed_data):
    # Convert pandas DataFrame to dictionary of numpy arrays
    data_dict = {column: observed_data[column].to_numpy() for column in observed_data.columns}
    
    built_model = model(data_dict)
    with built_model:
        # Sample from posterior
        trace = pm.sample(
            200, 
            tune=200, 
            target_accept=0.90, 
            chains=1, 
            cores=1, 
            random_seed=42, 
            idata_kwargs={"log_likelihood": True}
        )
        
        # Generate posterior predictive samples
        posterior_predictive = pm.sample_posterior_predictive(
            trace, 
            random_seed=314, 
            return_inferencedata=False
        )
        
    return built_model, posterior_predictive, trace