"""

Coxnet: CoxPH with Lasso penalty. On Owkin dataset.

Leon Zheng

"""

import preprocessing

from sksurv.linear_model import CoxnetSurvivalAnalysis

from sklearn.model_selection import cross_validate, RandomizedSearchCV

from sksurv.util import Surv

import numpy as np

import pandas as pd

# Features

# features = ['original_shape_Compactness2',

#  'original_shape_SphericalDisproportion',

#  'original_shape_SurfaceVolumeRatio',

#  'original_firstorder_Kurtosis',

#  'original_firstorder_MeanAbsoluteDeviation',

#  'original_firstorder_Minimum',

#  'original_glcm_ClusterProminence',

#  'original_glcm_Contrast',

#  'original_glcm_DifferenceEntropy',

#  'original_glcm_DifferenceAverage',

#  'original_glcm_JointEnergy',

#  'original_glcm_Id',

#  'original_glcm_Idm',

#  'original_glcm_Imc1',

#  'original_glcm_Imc2',

#  'original_glcm_Idmn',

#  'original_glcm_Idn',

#  'original_glrlm_ShortRunEmphasis',

#  'original_glrlm_LongRunEmphasis',

#  'original_glrlm_GrayLevelNonUniformity',

#  'original_glrlm_RunPercentage',

#  'original_glrlm_ShortRunLowGrayLevelEmphasis',

#  'original_glrlm_LongRunLowGrayLevelEmphasis',

#  'original_glrlm_LongRunHighGrayLevelEmphasis',

#  'Nstage',

#  'age',

#  'SourceDataset']

# radiomics_features = ['original_shape_Sphericity', 'original_shape_SurfaceVolumeRatio',

#                       'original_shape_Maximum3DDiameter', 'original_glcm_JointEntropy', 'original_glcm_Id',

#                       'original_glcm_Idm']

# clinical_features = ['SourceDataset', 'Nstage']

# features = radiomics_features + clinical_features

features = ['Mstage',

            'Nstage',

            'SourceDataset',

            'age',

            'original_shape_VoxelVolume',

            'original_firstorder_Maximum',

            'original_firstorder_Mean',

            'original_glcm_ClusterProminence',

            'original_glcm_Idm',

            'original_glcm_Idn',

            'original_glrlm_RunPercentage']

# Read data

input_train, output_train, input_test = preprocessing.load_owkin_data()

input_train = input_train[features]

input_test = input_test[features]

input_train, input_test = preprocessing.normalizing_input(input_train, input_test)

structured_y = Surv.from_dataframe('Event', 'SurvivalTime', output_train)

# Coxnet

# coxnet = CoxnetSurvivalAnalysis()

# print(cross_validate(coxnet, input_train, structured_y, cv=5))

# Grid search

tuned_params = {"l1_ratio": np.linspace(0.01, 0.02, 100),

                "n_alphas": range(140, 160, 1),

                }

grid_search = RandomizedSearchCV(CoxnetSurvivalAnalysis(), tuned_params, cv=5, n_jobs=4, n_iter=1000)

grid_search.fit(input_train, structured_y)

print(grid_search.best_score_)

best_params = grid_search.best_params_

print(best_params)

# Prediction

def predict(model, X, threshold=0.9):

    prediction = 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] > 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'])

coxph = CoxnetSurvivalAnalysis(**best_params, fit_baseline_model=True)

coxph.fit(input_train, structured_y)

y_pred = predict(coxph, input_test)

y_pred.to_csv('submission.csv')