text: # 4. import_survival.py
  # %matplotlib inline
  import pandas as pd
  import numpy as np
  import matplotlib.pyplot as plt
  import seaborn as sns;
  
  sns.set()
  
  from lifelines import KaplanMeierFitter, CoxTimeVaryingFitter
  
  from scipy.optimize import curve_fit
  
  
  def exponential_model(z, a, b):
      return a * np.exp(-b * z)
  
  
  def add_remaining_useful_life(df):
      # Get the total number of cycles for each unit
      grouped_by_unit = df.groupby(by="unit_nr")
      max_cycle = grouped_by_unit["time_cycles"].max()
  
      # Merge the max cycle back into the original frame
      result_frame = df.merge(max_cycle.to_frame(name='max_cycle'), left_on='unit_nr', right_index=True)
  
      # Calculate remaining useful life for each row
      remaining_useful_life = result_frame["max_cycle"] - result_frame["time_cycles"]
      result_frame["RUL"] = remaining_useful_life
  
      # drop max_cycle as it's no longer needed
      result_frame = result_frame.drop("max_cycle", axis=1)
      return result_frame
  
  
  def evaluate(y_true, y_hat, label='test'):
      mse = mean_squared_error(y_true, y_hat)
      rmse = np.sqrt(mse)
      variance = r2_score(y_true, y_hat)
      print('{} set RMSE:{}, R2:{}'.format(label, rmse, variance))
  
  
  # load data FD001.py
  # define filepath to read data
  dir_path = './CMAPSSData/'
  
  # define column names for easy indexing
  index_names = ['unit_nr', 'time_cycles']
  setting_names = ['setting_1', 'setting_2', 'setting_3']
  sensor_names = ['s_{}'.format(i) for i in range(1, 22)]
  col_names = index_names + setting_names + sensor_names
  
  # read data
  train = pd.read_csv((dir_path + 'train_FD001.txt'), sep='\s+', header=None, names=col_names)
  test = pd.read_csv((dir_path + 'test_FD001.txt'), sep='\s+', header=None, names=col_names)
  y_test = pd.read_csv((dir_path + 'RUL_FD001.txt'), sep='\s+', header=None, names=['RUL'])
  
  # inspect first few rows
  train.head()
  
  train = add_remaining_useful_life(train)
  train[index_names + ['RUL']].head()
  
  # clip RUL as discussed in SVR and problem framing analysis
  train['RUL'].clip(upper=125, inplace=True)
  
  # drop non-informative features, derived from EDA
  drop_sensors = ['s_1', 's_5', 's_6', 's_10', 's_16', 's_18', 's_19']
  drop_labels = setting_names + drop_sensors
  train.drop(labels=drop_labels, axis=1, inplace=True)
  
  remaining_sensors = ['s_2', 's_3', 's_4', 's_7', 's_8', 's_9',
                       's_11', 's_12', 's_13', 's_14', 's_15', 's_17', 's_20', 's_21']
  
  train['breakdown'] = 0
  idx_last_record = train.reset_index().groupby(by='unit_nr')['index'].last()  # engines breakdown at the last cycle
  train.at[idx_last_record, 'breakdown'] = 1
  
  train['start'] = train['time_cycles'] - 1
  train.tail()  # check results
  
  cut_off = 200
  train_censored = train[train['time_cycles'] <= cut_off].copy()
  
  data = train_censored[index_names + ['breakdown']].groupby('unit_nr').last()
  
  plt.figure(figsize=(15, 7))
  survival = KaplanMeierFitter()
  survival.fit(data['time_cycles'], data['breakdown'])
  survival.plot()
  plt.ylabel("Probability of survival")
  plt.show()
  plt.close()
  
  # Cox Proportional Hazards models
  
  train_cols = index_names + remaining_sensors + ['start', 'breakdown']
  predict_cols = ['time_cycles'] + remaining_sensors + ['start', 'breakdown']  # breakdown value will be 0
  
  ctv = CoxTimeVaryingFitter()
  ctv.fit(train_censored[train_cols], id_col="unit_nr", event_col='breakdown',
          start_col='start', stop_col='time_cycles', show_progress=True, step_size=1)
  
  ctv.print_summary()
  
  plt.figure(figsize=(10, 5))
  ctv.plot()
  plt.show()
  plt.close()
  
  df = train_censored.groupby("unit_nr").last()
  df = df[df['breakdown'] == 0]  # get engines from dataset which are still functioning so we can predict their RUL
  df_to_predict = df[df['breakdown'] == 0].copy()
  
  predictions = ctv.predict_log_partial_hazard(df_to_predict[predict_cols])
  predictions.rename(columns={0: "predictions"}, inplace=True)
  
  df_last = train.groupby('unit_nr').last()
  predictions['RUL'] = df_to_predict['RUL']
  predictions.head(10)
  
  plt.figure(figsize=(15, 5))
  plt.plot(predictions['RUL'], predictions['predictions'], '.b')
  xlim = plt.gca().get_xlim()
  plt.xlim(xlim[1], xlim[0])
  plt.xlabel('RUL')
  plt.ylabel('log_partial_hazard')
  
  plt.show()
  
  X = train.loc[train['unit_nr'].isin(df_to_predict.index)]
  X_unique = len(X['unit_nr'].unique())
  
  plt.figure(figsize=(15, 5))
  
  for i in range(1, X_unique, 2):
      X_sub = X.loc[X['unit_nr'] == i]
      predictions = ctv.predict_partial_hazard(X_sub)[0].values
      plt.plot(X_sub['time_cycles'].values, np.log(predictions))
  
  plt.xlabel('time_cycles')
  plt.ylabel('log_partial_hazard')
  plt.show()
  
  # Regressing log-partial hazard to RUL
  
  df_hazard = train_censored.copy()
  df_hazard['hazard'] = ctv.predict_log_partial_hazard(df_hazard)
  
  df_hazard.plot('hazard', 'RUL', 'scatter', figsize=(15, 5))
  plt.xlabel('hazard')
  plt.ylabel('RUL')
  plt.show()
  
  popt, pcov = curve_fit(exponential_model, df_hazard['hazard'], df_hazard['RUL'])
  print(popt)
  
  # prep test set
  test = X_test.drop(labels=drop_labels, axis=1)
  test['breakdown'] = 0
  test['start'] = test['time_cycles'] - 1
  
  # predict and evaluate
  y_hat = exponential_model(df_hazard['hazard'], *popt)
  evaluate(df_hazard['RUL'], y_hat, 'train')
  
  y_pred = ctv.predict_log_partial_hazard(test.groupby('unit_nr').last())
  y_hat = exponential_model(y_pred, *popt)
  evaluate(y_test, y_hat)
  
  # Repeat on complete dataset
  ctv2 = CoxTimeVaryingFitter()
  ctv2.fit(train[train_cols], id_col="unit_nr", event_col='breakdown',
           start_col='start', stop_col='time_cycles', show_progress=True)
  
  train['hazard'] = ctv2.predict_log_partial_hazard(train)
  popt2, pcov2 = curve_fit(exponential_model, train['hazard'], train['RUL'])
  
  y_hat = exponential_model(train['hazard'], *popt2)
  evaluate(train['RUL'], y_hat, 'train')
  
  y_pred = ctv2.predict_log_partial_hazard(test.groupby('unit_nr').last())
  y_hat = exponential_model(y_pred, *popt2)
  evaluate(y_test, y_hat)
  