# -*- coding: utf-8 -*-
from __future__ import print_function
import numpy as np
import pandas as pd

from lifelines.fitters import UnivariateFitter
from lifelines.utils import _preprocess_inputs, _additive_estimate, StatError, inv_normal_cdf,\
    median_survival_times
from lifelines.plotting import plot_loglogs


class KaplanMeierFitter(UnivariateFitter):

    """
    Class for fitting the Kaplan-Meier estimate for the survival function.

    KaplanMeierFitter( alpha=0.95)

    alpha: The alpha value associated with the confidence intervals.

    """

    def fit(self, durations, event_observed=None, timeline=None, entry=None, label='KM_estimate',
            alpha=None, left_censorship=False, ci_labels=None):
        """
        Parameters:
          duration: an array, or pd.Series, of length n -- duration subject was observed for
          timeline: return the best estimate at the values in timelines (postively increasing)
          event_observed: an array, or pd.Series, of length n -- True if the the death was observed, False if the event
             was lost (right-censored). Defaults all True if event_observed==None
          entry: an array, or pd.Series, of length n -- relative time when a subject entered the study. This is
             useful for left-truncated (not left-censored) observations. If None, all members of the population
             were born at time 0.
          label: a string to name the column of the estimate.
          alpha: the alpha value in the confidence intervals. Overrides the initializing
             alpha for this call to fit only.
          left_censorship: True if durations and event_observed refer to left censorship events. Default False
          ci_labels: add custom column names to the generated confidence intervals
                as a length-2 list: [<lower-bound name>, <upper-bound name>]. Default: <label>_lower_<alpha>


        Returns:
          self, with new properties like 'survival_function_'.

        """
        # if the user is interested in left-censorship, we return the cumulative_density_, no survival_function_,
        estimate_name = 'survival_function_' if not left_censorship else 'cumulative_density_'
        v = _preprocess_inputs(durations, event_observed, timeline, entry)
        self.durations, self.event_observed, self.timeline, self.entry, self.event_table = v
        self._label = label
        alpha = alpha if alpha else self.alpha
        log_survival_function, cumulative_sq_ = _additive_estimate(self.event_table, self.timeline,
                                                                   self._additive_f, self._additive_var,
                                                                   left_censorship)

        if entry is not None:
            # a serious problem with KM is that when the sample size is small and there are too few early
            # truncation times, it may happen that is the number of patients at risk and the number of deaths is the same.
            # we adjust for this using the Breslow-Fleming-Harrington estimator
            n = self.event_table.shape[0]
            net_population = (self.event_table['entrance'] - self.event_table['removed']).cumsum()
            if net_population.iloc[:int(n / 2)].min() == 0:
                ix = net_population.iloc[:int(n / 2)].argmin()
                raise StatError("""There are too few early truncation times and too many events. S(t)==0 for all t>%.1f. Recommend BreslowFlemingHarringtonFitter.""" % ix)

        # estimation
        setattr(self, estimate_name, pd.DataFrame(np.exp(log_survival_function), columns=[self._label]))
        self.__estimate = getattr(self, estimate_name)
        self.confidence_interval_ = self._bounds(cumulative_sq_[:, None], alpha, ci_labels)
        self.median_ = median_survival_times(self.__estimate, left_censorship=left_censorship)

        # estimation methods
        self.predict = self._predict(estimate_name, label)
        self.subtract = self._subtract(estimate_name)
        self.divide = self._divide(estimate_name)

        # plotting functions
        self.plot = self._plot_estimate(estimate_name)
        setattr(self, "plot_" + estimate_name, self.plot)
        self.plot_loglogs = plot_loglogs(self)
        return self

    def _bounds(self, cumulative_sq_, alpha, ci_labels):
        # See http://courses.nus.edu.sg/course/stacar/internet/st3242/handouts/notes2.pdf
        alpha2 = inv_normal_cdf((1. + alpha) / 2.)
        df = pd.DataFrame(index=self.timeline)
        v = np.log(self.__estimate.values)

        if ci_labels is None:
            ci_labels = ["%s_upper_%.2f" % (self._label, alpha), "%s_lower_%.2f" % (self._label, alpha)]
        assert len(ci_labels) == 2, "ci_labels should be a length 2 array."

        df[ci_labels[0]] = np.exp(-np.exp(np.log(-v) + alpha2 * np.sqrt(cumulative_sq_) / v))
        df[ci_labels[1]] = np.exp(-np.exp(np.log(-v) - alpha2 * np.sqrt(cumulative_sq_) / v))
        return df

    def _additive_f(self, population, deaths):
        np.seterr(invalid='ignore', divide='ignore')
        return (np.log(population - deaths) - np.log(population))

    def _additive_var(self, population, deaths):
        np.seterr(divide='ignore')
        return (1. * deaths / (population * (population - deaths))).replace([np.inf], 0)

[((48, 12, 48, 74), 'lifelines.utils._preprocess_inputs', '_preprocess_inputs', ({(48, 31, 48, 40): 'durations', (48, 42, 48, 56): 'event_observed', (48, 58, 48, 66): 'timeline', (48, 68, 48, 73): 'entry'}, {}), '(durations, event_observed, timeline, entry)', False, 'from lifelines.utils import _preprocess_inputs, _additive_estimate, StatError, inv_normal_cdf, median_survival_times\n'), ((52, 48, 54, 83), 'lifelines.utils._additive_estimate', '_additive_estimate', ({(52, 67, 52, 83): 'self.event_table', (52, 85, 52, 98): 'self.timeline', (53, 67, 53, 83): 'self._additive_f', (53, 85, 53, 103): 'self._additive_var', (54, 67, 54, 82): 'left_censorship'}, {}), '(self.event_table, self.timeline, self._additive_f, self.\n    _additive_var, left_censorship)', False, 'from lifelines.utils import _preprocess_inputs, _additive_estimate, StatError, inv_normal_cdf, median_survival_times\n'), ((70, 23, 70, 94), 'lifelines.utils.median_survival_times', 'median_survival_times', (), '', False, 'from lifelines.utils import _preprocess_inputs, _additive_estimate, StatError, inv_normal_cdf, median_survival_times\n'), ((80, 28, 80, 46), 'lifelines.plotting.plot_loglogs', 'plot_loglogs', ({(80, 41, 80, 45): 'self'}, {}), '(self)', False, 'from lifelines.plotting import plot_loglogs\n'), ((85, 17, 85, 50), 'lifelines.utils.inv_normal_cdf', 'inv_normal_cdf', ({(85, 32, 85, 49): '(1.0 + alpha) / 2.0'}, {}), '((1.0 + alpha) / 2.0)', False, 'from lifelines.utils import _preprocess_inputs, _additive_estimate, StatError, inv_normal_cdf, median_survival_times\n'), ((86, 13, 86, 46), 'pandas.DataFrame', 'pd.DataFrame', (), '', True, 'import pandas as pd\n'), ((87, 12, 87, 42), 'numpy.log', 'np.log', ({(87, 19, 87, 41): 'self.__estimate.values'}, {}), '(self.__estimate.values)', True, 'import numpy as np\n'), ((98, 8, 98, 52), 'numpy.seterr', 'np.seterr', (), '', True, 'import numpy as np\n'), ((102, 8, 102, 34), 'numpy.seterr', 'np.seterr', (), '', True, 'import numpy as np\n'), ((99, 16, 99, 43), 'numpy.log', 'np.log', ({(99, 23, 99, 42): '(population - deaths)'}, {}), '(population - deaths)', True, 'import numpy as np\n'), ((99, 46, 99, 64), 'numpy.log', 'np.log', ({(99, 53, 99, 63): 'population'}, {}), '(population)', True, 'import numpy as np\n'), ((64, 22, 64, 171), 'lifelines.utils.StatError', 'StatError', ({(64, 32, 64, 170): "('There are too few early truncation times and too many events. S(t)==0 for all t>%.1f. Recommend BreslowFlemingHarringtonFitter.'\n     % ix)"}, {}), "(\n    'There are too few early truncation times and too many events. S(t)==0 for all t>%.1f. Recommend BreslowFlemingHarringtonFitter.'\n     % ix)", False, 'from lifelines.utils import _preprocess_inputs, _additive_estimate, StatError, inv_normal_cdf, median_survival_times\n'), ((67, 50, 67, 79), 'numpy.exp', 'np.exp', ({(67, 57, 67, 78): 'log_survival_function'}, {}), '(log_survival_function)', True, 'import numpy as np\n'), ((93, 42, 93, 52), 'numpy.log', 'np.log', ({(93, 49, 93, 51): '-v'}, {}), '(-v)', True, 'import numpy as np\n'), ((94, 42, 94, 52), 'numpy.log', 'np.log', ({(94, 49, 94, 51): '-v'}, {}), '(-v)', True, 'import numpy as np\n'), ((93, 64, 93, 87), 'numpy.sqrt', 'np.sqrt', ({(93, 72, 93, 86): 'cumulative_sq_'}, {}), '(cumulative_sq_)', True, 'import numpy as np\n'), ((94, 64, 94, 87), 'numpy.sqrt', 'np.sqrt', ({(94, 72, 94, 86): 'cumulative_sq_'}, {}), '(cumulative_sq_)', True, 'import numpy as np\n')]