import numpy as np
import pymc as pm


def model(data):
    """
    Logistic growth model for disease spread.
    
    Assumptions:
    - Disease spread follows logistic growth
    - Total population is 50 (carrying capacity)
    - Observed counts follow a Binomial distribution
    - Priors are weakly informative to allow flexibility
    """
    # 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)
        # t0: inflection point (centered around 1, since that's where we see ~25 infected)
        r = pm.HalfNormal('r', sigma=5)
        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)))
        
        # Convert to probability for Binomial distribution
        p = expected_infected / N
        
        # Clip to avoid numerical issues
        p = pm.math.minimum(pm.math.maximum(p, 1e-10), 1 - 1e-10)
        
        # Observed likelihood - Binomial
        y_obs = pm.Binomial('y_obs', n=N, p=p, 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