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 logistic growth dynamics.
    - Total population is fixed at 50.
    - Observed counts are subject to overdispersion relative to the logistic mean.
    - Negative Binomial likelihood is used to model the count data, allowing for 
      extra variance beyond the Binomial constraint.
    """
    # 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). HalfNormal with sigma=5 allows for rapid growth.
        r = pm.HalfNormal('r', sigma=5)
        # t0: inflection point. Centered at 1.0 with reasonable uncertainty.
        t0 = pm.Normal('t0', mu=1.0, sigma=0.5)
        
        # 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 ensure it stays within valid range for probability calculation
        # though NegativeBinomial takes mu directly, we might want to ensure positivity.
        # The logistic function naturally stays between 0 and N.
        
        # Overdispersion parameter for Negative Binomial
        # alpha is the dispersion parameter. Smaller alpha means more dispersion.
        # Using a HalfNormal prior for alpha.
        alpha = pm.HalfNormal('alpha', sigma=5)
        
        # Observed likelihood: Negative Binomial
        # pm.NegativeBinomial('y_obs', mu=expected_infected, alpha=alpha, observed=infected, dims='obs_id')
        # Note: In some PyMC versions, the parameterization might differ. 
        # Standard parameterization is mu (mean) and alpha (dispersion).
        # Variance = mu + mu^2 / alpha
        y_obs = pm.NegativeBinomial('y_obs', mu=expected_infected, 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