import numpy as np
import pymc as pm
import pandas as pd

def model(data):
    """
    Logistic growth model for disease spread with Negative Binomial likelihood.
    
    Assumptions:
    - Disease spread follows a logistic growth mean trend.
    - Total population is 50 (carrying capacity).
    - Observed counts exhibit overdispersion relative to Binomial, modeled via Negative Binomial.
    - Priors are weakly informative.
    """
    # Extract data
    time = data['Time']
    infected = data['Infected_Count']
    N = 50  # Total population
    
    # Create PyMC model
    with pm.Model() as model:
        # Define feature container for time
        time_data = pm.Data('Time', time, dims='obs_id')
        
        # Priors for logistic growth parameters
        # r: growth rate (positive, steep rise observed)
        r = pm.HalfNormal('r', sigma=10)
        # t0: inflection point (centered around 1.0, sigma=1.0 to be flexible)
        t0 = pm.Normal('t0', mu=1.0, sigma=1.0)
        
        # Logistic function: I(t) = N / (1 + exp(-r*(t - t0)))
        # Compute the expected number of infected
        expected_infected = N / (1 + pm.math.exp(-r * (time_data - t0)))
        
        # Clip expected infected to be within [1, N-1] to avoid log issues or zero/negative means
        # We want mean > 0 for NegativeBinomial. 
        # Also, NegativeBinomial parametrization in PyMC uses mu and alpha (or n, p).
        # Using mu and alpha parametrization: mu is mean, alpha is dispersion.
        # Var = mu + mu^2 / alpha
        
        # Clip mean to be strictly positive and <= N
        mu = pm.math.minimum(pm.math.maximum(expected_infected, 0.5), N - 0.5)
        
        # Dispersion parameter for Negative Binomial
        # High alpha means low dispersion (closer to Poisson/Binomial)
        # Low alpha means high dispersion
        alpha = pm.HalfNormal('alpha', sigma=5)
        
        # Observed likelihood - Negative Binomial
        # mu is the mean, alpha is the dispersion parameter
        y_obs = pm.NegativeBinomial('y_obs', mu=mu, alpha=alpha, observed=infected, dims='obs_id')
    
    return model

def gen_model(observed_data):
    built_model = model({column: observed_data[column].to_numpy() for column in observed_data.columns})
    with built_model:
        trace = pm.sample(200, tune=200, target_accept=0.90, chains=1, cores=1, random_seed=42, idata_kwargs={"log_likelihood": True})
        posterior_predictive = pm.sample_posterior_predictive(trace, random_seed=314, return_inferencedata=False)
    return built_model, posterior_predictive, trace