# dependencies
from lifelines import KaplanMeierFitter
from lifelines.statistics import pairwise_logrank_test
from bokeh.io import output_file, output_notebook
from bokeh.plotting import figure, show
from bokeh.models import ColumnDataSource, HoverTool,Band
from bokeh.layouts import row, column, gridplot
from bokeh.models.widgets import Tabs, Panel
from bokeh import palettes
import itertools
from itertools import cycle
import pandas as pd
import numpy as np

# helper functions   
def make_step(x,y):
    """
    Makes a step function using x and y coordinates
    ----------
    x: x coordinates in a list
    y: y coordinates in a list
    -------
    
    """
    add_x = x[1:]
    new_x = [None]*(len(x)+len(add_x))
    new_x[::2] = x
    new_x[1::2] = add_x
    add_y = y[:-1]
    new_y = [None]*(len(x)+len(add_y))
    new_y[::2] = y
    new_y[1::2] = add_y
    return new_x, new_y

def color_negative_red(val):
    color = 'red' if val > 0.05 else 'green'
    return 'color: %s' % color

class time_to_event(object):
    """
    Class for fitting the Kaplan-Meier estimate to record time to event
    
    """

    def fit_plot(self,dataframe,durations,group=None,event_observed=None):
        """
        Fit the model to right censored data and plot survival function
        Parameters
        ----------
          dataframe: Pandas dataframe being fit. Should have columns with durations, and optionally group and event observed.
          durations: Time to event for data point. For example, if the event is defined as "when supply for a clothing line is over" and say Women's size Small Tshirt supplies last for 45 days, then duration is 45.
          group: Column of dataframe that represents strata of dataframe (such as clusters or customer segments)
          event_observed: Whether event was observed (True) or not (False).
        Returns
        -------
        Interactive plot showing survival curves
        """
        # Determine where the visualization will be rendered
        output_notebook()  # Render inline in a Jupyter Notebook
        if group is None:
            kmf = KaplanMeierFitter()
            if event_observed is not None:
                kmf.fit(dataframe[durations], dataframe[event_observed])
            else:
                kmf.fit(dataframe[durations])
            
            # Extract kmf estimate values 
            x_temp = kmf.survival_function_.index
            y_temp = kmf.survival_function_['KM_estimate'].values
            
            # Make step transformations
            x,y = make_step(x_temp,y_temp)
            conf_upper = kmf.confidence_interval_survival_function_['KM_estimate_upper_0.95']
            conf_lower = kmf.confidence_interval_survival_function_['KM_estimate_lower_0.95']
            conf_upper = make_step(conf_upper,y_temp)[0]
            conf_lower = make_step(conf_lower,y_temp)[0]
            
            # Getting the kmf results into dataframe format
            df_for_plot = pd.DataFrame()
            df_for_plot['x'] = x
            df_for_plot['y'] = y
            df_for_plot['upper'] = conf_upper
            df_for_plot['lower'] = conf_lower

            # Instantiate a figure() object
            fig = figure(plot_height=400,
            plot_width=600,
            x_axis_label='Timeline',
            y_axis_label='Some metric',
            title='Graph',
            )
            fig.multi_line(xs=[df_for_plot['x']], ys=[df_for_plot['y']],
            line_width=4,
            line_color = 'firebrick')

            # confidence intervals
            fig.multi_line(xs=[df_for_plot['x'],df_for_plot['x']], ys=[df_for_plot['upper'],df_for_plot['lower']],
            line_width=2,alpha=0.5,line_dash='dashed',legend='confidence bound')

            #add hovering tooltips
            tooltips = [
            ('Time','$x'),
            ('Survival Probability', '$y')
            ]
            fig.legend.click_policy="hide"
            fig.add_tools(HoverTool(tooltips=tooltips))

            show(fig)
            
        else:
            list_segment_vals_x = []
            list_segment_vals_y = []
            for segment in dataframe[group].unique():
                df_to_fit = dataframe[dataframe[group]==segment]
                kmf = KaplanMeierFitter()
                if event_observed is not None:
                    kmf.fit(df_to_fit[durations], df_to_fit[event_observed])
                else:
                    kmf.fit(df_to_fit[durations])
                x_temp = kmf.survival_function_.index
                y_temp = kmf.survival_function_['KM_estimate'].values
            
                # Make step transformations
                x,y = make_step(x_temp,y_temp)
                list_segment_vals_x.append(x)
                list_segment_vals_y.append(y)
                
            # Instantiate a figure() object
            fig = figure(plot_height=400,
                         plot_width=600,
                         x_axis_label='Timeline',
                         y_axis_label='Some metric',
                         title='Graph',
                        )
            count = 0
            for val_1,val_2 in zip(list_segment_vals_x,list_segment_vals_y):
                fig.line(x=val_1, y=val_2,
                     line_width=2,legend=[str(x) for x in dataframe['segments'].unique()][count],
                     color=palettes.Category20_5[count])
                count+=1

            #add hovering tooltips
            tooltips = [
                        ('Time','$x'),
                        ('Survival Probability', '$y')
                       ]
            fig.add_tools(HoverTool(tooltips=tooltips))
            fig.legend.click_policy="hide"
            show(fig)
        return self
            
    def significance_results(self,dataframe,durations,group,event_observed=None):
        """
        Fit the model to right censored data and plot survival function
        Parameters
        ----------
          dataframe: Pandas dataframe being fit. Should have columns with durations, group and event observed.
          durations: Time to event for data point. For example, if the event is defined as "when supply for a clothing line is over" and say Women's size Small Tshirt supplies last for 45 days, then duration is 45.
          group: Column of dataframe that represents strata of dataframe (such as clusters or customer segments)
          event_observed: Whether event was observed (True) or not (False).
        Returns
        -------
        Table showing significance results at the 0.05 level.
        """
        df_pairwise = pd.DataFrame()
        if event_observed is not None:
            df_pairwise =pairwise_logrank_test(dataframe['durations'],dataframe['segments'],
                                           dataframe['event_observed']).summary[['p']]
        else:
            df_pairwise =pairwise_logrank_test(dataframe['durations'],dataframe['segments']).summary[['p']]
        df_pairwise.columns = ['p-value']
        df_pairwise = df_pairwise.style.applymap(color_negative_red)
        
        return df_pairwise
            
        