"""
COX-PH Baseline
Leon Zheng
"""

import pandas as pd
from sksurv.linear_model import CoxPHSurvivalAnalysis
from sklearn.base import BaseEstimator
from sksurv.util import Surv
import numpy as np
from metrics import cindex


class CoxPhRegression(BaseEstimator):

    def __init__(self, alpha=1e-8, threshold=0.9):
        self.alpha = alpha
        self.threshold = threshold
        self.model = CoxPHSurvivalAnalysis(alpha=alpha)

    def set_data(self, input_train, output_train, input_test):
        self.input_train = input_train
        self.output_train = output_train
        self.input_test = input_test

    def fit(self, X, y):
        structured_y = Surv.from_dataframe('Event', 'SurvivalTime', y)
        self.model.fit(X, structured_y)
        return self

    def predict(self, X):
        prediction = self.model.predict_survival_function(X)
        y_pred = []
        for pred in prediction:
            time = pred.x
            survival_prob = pred.y
            i_pred = 0
            while i_pred < len(survival_prob) - 1 and survival_prob[i_pred] > self.threshold:
                i_pred += 1
            y_pred.append(time[i_pred])
        return pd.DataFrame(np.array([[y, np.nan] for y in y_pred]), index=X.index, columns=['SurvivalTime', 'Event'])

    def score(self, X, y):
        y_pred = self.predict(X)
        return cindex(y, y_pred)