import os
#os.chdir('/Users/q6600sl/IPython_NB')
from lifemodels import s_models
%matplotlib inline

"""
Explanation: UNKOWN genotype demography experiment
Summary

to be filled
to install run pip install git+https://github.com/ctzhu/lifemodels, assuming a running python installation is inplace
I have to make WIN32 dynamic libraries, before that the following notebook can only run in MacOX
End of explanation
"""


#Read the data
original_df = pd.read_csv('/Users/q6600sl/Downloads/SP_12-22-15.txt', sep=' ')
#1st step: create a survival object
surv_df    = s_models.survt_df(original_df)

original_df.head()

"""
Explanation: · Read the data and create a survival object using .survt_df() method
The only thing I add is the column name for the genotypes ('Genotype')
End of explanation
"""


gp3_fit    = s_models.distfit_df(surv_df, 'gp3')

gp3_fit.mdl_all_free

"""
Explanation: · Specify the model and create the model fit object using .distfit_df() method
End of explanation
"""


gp3_fit.var

"""
Explanation: · We can get the variance of the parameter estimates
It might not always be there, due to float overflow and underflow problems
End of explanation
"""


(gp3_fit.logq('Mut-1'), gp3_fit.logq_nc('Mut-1'))

medium_df = pd.DataFrame({'With_C'   : map(gp3_fit.logq, gp3_fit.mdl_all_free.index),
                          'Without_C': map(gp3_fit.logq_nc, gp3_fit.mdl_all_free.index)},
                         index=gp3_fit.mdl_all_free.index)
np.exp(medium_df)

ax = medium_df.plot(kind='bar', alpha=0.5)
ax.set_ylim(3, 4.5)

right_y = ax.twinx()
right_y.set_ylim(3, 4.5)
_ = right_y.set_yticklabels(map('{:.2f}'.format,  np.exp(right_y.get_yticks())))

ax.set_ylabel('Medium survival time in log(days)')
right_y.set_ylabel('Medium survival time in days')
plt.title('Medium lifespan estimated based on Gompertz Makeham Model', y=1.05)

"""
Explanation: · Get the estimated log time-to-reach-%-of-death using .logq() method
.logq(key, q), key is genotype here. q is that target % of death. q is default to 50%, therefore by default it returns the estimate for medium lifespan.
.logq() method and .logq_nc() method are similar: the latter ignores the Makehamm term, if there is one.
End of explanation
"""


(gp3_fit.logm('Mut-1'), gp3_fit.logm_nc('Mut-1'))

"""
Explanation: · Get the estimated log mean lifespan using .logm() method
End of explanation
"""


gp3_fit.logt_var('Mut-1', 3.9)

"""
Explanation: · We can also calculate the variance of time-to-percentage-death estimiates
End of explanation
"""


f, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 5))

plot_par = {'logy' : True,
            'alpha': 0.5,
            'color': list('BBBBBR')} 

np.exp(gp3_fit.mdl_all_free).ix[:, 'alpha'].plot(kind='bar', ax=ax1, **plot_par)
np.exp(gp3_fit.mdl_all_free).ix[:, 'beta'].plot(kind='bar', ax=ax2, **plot_par)
np.exp(gp3_fit.mdl_all_free).ix[:, 'c'].plot(kind='bar', ax=ax3, **plot_par)

ax1.set_title(r'Gompertz $\alpha$')
ax2.set_title(r'Gompertz $\beta$')
ax3.set_title(r'Gompertz $C$')

plt.suptitle('Gompertz Makeham Model parameter estimates', fontsize=16, y=0.99)

"""
Explanation: · Plot for parameter estimates
End of explanation
"""


f, axes = plt.subplots(2, 3, figsize=(14, 8), sharex=True, sharey=True)
axes = axes.ravel()
for ax, key in zip(axes, gp3_fit.mdl_all_free.index.tolist()):
    gp3_fit.plot(key, np.linspace(0,80,100),ax,'hzd')
    ax.set_ylim((-8,1))
    ax.set_title(key)
plt.suptitle('log-harzard plot', fontsize=20)

f, axes = plt.subplots(2, 3, figsize=(14, 8), sharex=True, sharey=True)
axes = axes.ravel()
for ax, key in zip(axes, gp3_fit.mdl_all_free.index.tolist()):
    gp3_fit.plot(key, np.linspace(0,80,100),ax,'sf')
    ax.set_ylim((-0.01,1.01))
    ax.set_title(key)
plt.suptitle('survival plot', fontsize=20)

"""
Explanation: · Plot harzard plot, survival plot and other plots using .plot() method
End of explanation
"""


