import numpy as np
import pymc as pm
import theano.tensor as tt


def model(data):
    """
    Build a PyMC model for the death_process / infection spread problem.
    
    The data consists of Time and Infected_Count observations from a population of 50.
    We need to model the number of infected individuals as a function of time.
    
    A common model for disease spread is the logistic growth model (SIR-like), where:
    I(t) = N / (1 + ((N - I0) / I0) * exp(-r * t))
    
    However, since we're observing noisy counts at specific times, we can use a simpler
    parametric model or a stochastic process model.
    
    Given the summary statistics:
    - Time ranges from ~0 to ~2
    - Infected_Count ranges from 1 to 47
    - Mean infected count is ~30
    
    Let's use a logistic growth model with parameters:
    - r: growth rate
    - I0: initial infected count
    
    The logistic function: I(t) = N / (1 + ((N - I0) / I0) * exp(-r * t))
    
    We'll place priors on r and I0, and use a likelihood that accounts for the discrete
    nature of the counts (e.g., Normal or Poisson approximation, or just Normal for simplicity
    since counts are relatively large).
    """
    
    N = 50  # Total population
    
    times = data['Time']
    infected_counts = data['Infected_Count']
    
    with pm.Model() as pm_model:
        # Priors on parameters
        # Growth rate r: positive, reasonable range
        r = pm.HalfNormal('r', sigma=5)
        
        # Initial infected count I0: between 1 and N
        # Using a Beta distribution scaled to [1, N]
        # Or directly use a discrete uniform or normal truncated
        I0 = pm.Beta('I0', alpha=2, beta=5, transform=pm.distributions.transforms.ordered)
        # Scale I0 to be between 1 and N-1
        I0_scaled = I0 * (N - 2) + 1
        
        # Deterministic: predicted infected count at each time point
        # Logistic growth: I(t) = N / (1 + ((N - I0) / I0) * exp(-r * t))
        # Handle edge cases where I0 is very small
        # Using a safer formulation
        
        # Let's use a different parameterization to avoid numerical issues
        # I(t) = N * sigmoid(r * (t - t0))
        # where t0 is the inflection point
        # sigmoid(x) = 1 / (1 + exp(-x))
        
        # Alternative: use r and K (carrying capacity, which is N)
        # I(t) = N / (1 + exp(-r * (t - t_mid)))
        # where t_mid is the time at which I(t) = N/2
        
        t_mid = pm.Normal('t_mid', mu=1.0, sigma=0.5)
        
        # Predicted infected count using logistic function
        I_pred = N / (1 + tt.exp(-r * (times - t_mid)))
        
        # Likelihood: observed counts are approximately normal around the predicted values
        # with some noise. Since counts are integers, we could use a Poisson or Negative Binomial,
        # but for simplicity and given the moderate counts, a Normal likelihood with a small
        # standard deviation works well.
        
        sigma = pm.HalfNormal('sigma', sigma=5)
        
        # Observed likelihood
        infected_obs = pm.Normal('infected_obs', mu=I_pred, sigma=sigma, observed=infected_counts)
    
    return pm_model