# If we're running on Colab, install empiricaldist
# https://pypi.org/project/empiricaldist/

import sys
IN_COLAB = 'google.colab' in sys.modules

if IN_COLAB:
    !pip install empiricaldist

# Get utils.py

from os.path import basename, exists

def download(url):
    filename = basename(url)
    if not exists(filename):
        from urllib.request import urlretrieve
        local, _ = urlretrieve(url, filename)
        print('Downloaded ' + local)
    
download('https://github.com/AllenDowney/ThinkBayes2/raw/master/soln/utils.py')

from utils import set_pyplot_params
set_pyplot_params()

"""
Explanation: Survival Analysis
End of explanation
"""


download('https://github.com/AllenDowney/ThinkBayes2/raw/master/examples/Osogbo2021.csv')

"""
Explanation: The following cell downloads the data.
End of explanation
"""


import pandas as pd

df = pd.read_csv('Osogbo2021.csv')
df.head()

import numpy as np

log_duration = np.log10(duration)

from empiricaldist import Cdf
from utils import decorate

cdf = Cdf.from_seq(log_duration)
cdf.plot()
    
decorate(xlabel='Survival time (log10 days)', 
         ylabel='CDF',
         title='Distribution raw survival times')

mu = log_duration.mean()
sigma = log_duration.std()
mu, sigma

from empiricaldist import Pmf

def make_uniform(qs, name=None, **options):
    """Make a Pmf that represents a uniform distribution."""
    pmf = Pmf(1.0, qs, **options)
    pmf.normalize()
    if name:
        pmf.index.name = name
    return pmf

import numpy as np

qs = np.linspace(2, 4, num=101)
prior_mu = make_uniform(qs, name='mean')

"""
Explanation: I'll use Pandas to load the data into a DataFrame.
End of explanation
"""


qs = np.linspace(0.01, 1, num=90)
prior_sigma = make_uniform(qs, name='std')

"""
Explanation: I chose the lower and upper bounds by trial and error.
I'll explain how when we look at the posterior distribution.
Here's the prior distribution for sigma:
End of explanation
"""


from utils import make_joint

prior = make_joint(prior_mu, prior_sigma)

"""
Explanation: Now we can use make_joint to make the joint prior distribution.
End of explanation
"""


complete = (df['CENSOR'] == 0)
data_complete = log_duration[complete]
data_ongoing = log_duration[~complete]

len(data_complete), len(data_ongoing)

"""
Explanation: And we'll start by working with the data from the control group.
End of explanation
"""


mu_mesh, sigma_mesh, data_mesh = np.meshgrid(
    prior.columns, prior.index, data_complete)

mu_mesh.shape

from scipy.stats import norm

densities = norm(mu_mesh, sigma_mesh).pdf(data_mesh)
densities.shape

likelihood = densities.prod(axis=2)
likelihood.shape

from utils import normalize

posterior = prior * likelihood
normalize(posterior)
posterior.shape

"""
Explanation: Likelihood
First update
End of explanation
"""


mu_mesh, sigma_mesh, data_mesh = np.meshgrid(
    prior.columns, prior.index, data_ongoing)

mu_mesh.shape

from scipy.stats import norm

densities = norm(mu_mesh, sigma_mesh).sf(data_mesh)
densities.shape

likelihood = densities.prod(axis=2)
likelihood.shape

from utils import normalize

posterior2 = posterior * likelihood
normalize(posterior2)
posterior2.shape

import matplotlib.pyplot as plt
from utils import plot_contour

plot_contour(posterior2, cmap='Blues')

decorate(xlabel='Mean (mu)', 
         ylabel='Standard deviation (sigma)',
         title='Joint posterior distributions of mu and sigma')

"""
Explanation: Second update
End of explanation
"""


from utils import marginal



"""
Explanation: Posterior Marginal Distributions
I'll use marginal, which we saw in <<_MarginalDistributions>>, to extract the posterior marginal distributions for the population means.
End of explanation
"""


pmf_mean = marginal(posterior, 0)
pmf_mean.plot()

pmf_mean2 = marginal(posterior2, 0)
pmf_mean2.plot()

decorate(xlabel='Population mean (mu)', 
         ylabel='PDF', 
         title='Posterior distributions of mu')

pmf_std = marginal(posterior, 1)
pmf_std.plot()

pmf_std2 = marginal(posterior2, 1)
pmf_std2.plot()

decorate(xlabel='Population std (sigma)', 
         ylabel='PDF')

"""
Explanation: Here's what they look like:
End of explanation
"""


