#!/usr/bin/env python

import matplotlib.pyplot as plt
from database import Database as db
from lifelines.estimation import KaplanMeierFitter
from matplotlib.pyplot import subplot
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.serif'] = ['Cambria Math'] + plt.rcParams['font.serif']

kmf = KaplanMeierFitter()
sx = subplot()

def survival_lifeline():
    f = db()
    df = f.survival_read2().head(n=1000000)
    for r in df['IP_SOURCE'].unique():
        ix = df['IP_SOURCE'] == r
        kmf.fit(df['T'].ix[ix], df['C'].ix[ix], label = r)
        print(kmf.survival_function_)
        kmf.plot(ax=sx, ci_show=False, drawstyle='steps-mid', )

survival_lifeline()
plt.show()




