<gh_stars>1-10
# -*- coding: utf-8 -*-
"""csa_api.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1OfXVnydZA5UonDed1gTcBbFmpMLW4YTY
"""


from surv_ci.train_utils import train_survci
from surv_ci.train_utils import _get_padded_features, _get_padded_targets
from surv_ci.train_utils import _reshape_tensor_with_nans
from surv_ci.utilities import get_parameters, softmax_out, sample_weibull, sample_lognormal
from surv_ci.model import survci
from lifelines.utils import concordance_index
from surv_ci.losses import conditional_loss,unconditional_loss,mse_loss,imb_loss,mse_total,factual_loss

import torch
import numpy as np

np.random.seed(1234)
torch.manual_seed(seed=1234)

if torch.cuda.is_available():
  device = torch.device("cuda")
else:
  device = torch.device("cpu")
  
class survciBase():

  def __init__(self, k=3, layers=None, distribution="Weibull",
               temp=1000., discount=1.0,imb_func='lin_disc',p_ipm=0.5,p_alpha=1e-2,p_beta =1e-4 ):
    self.k = k
    self.layers = layers
    self.dist = distribution
    self.temp = temp
    self.discount = discount
    self.fitted = False
    self.imb_func = imb_func
    self.p_ipm = p_ipm
    self.p_alpha = p_alpha
    self.p_beta = p_beta

  def _gen_torch_model(self, inputdim, optimizer, num_treatments):
    """Helper function to return a torch model."""
    return survci(inputdim,k=self.k, layers=self.layers,
                                     dist=self.dist,
                                     temp=self.temp,
                                     discount=self.discount,
                                     optimizer=optimizer,
                                     num_treatments=2,imb_func=self.imb_func,p_ipm=self.p_ipm,p_alpha=self.p_alpha, p_beta = self.p_beta)
  

  def fit(self, x, t, e,w, vsize=0.15, val_data=None,
          iters=1, learning_rate=1e-3, batch_size=100,
          elbo=True, optimizer="Adam", random_state=1234):



    processed_data = self._prepocess_training_data(x, t, e,w,
                                                   vsize, val_data,
                                                   random_state)
    x_train, t_train, e_train,w_train, x_val, t_val, e_val,w_val = processed_data

    inputdim = x_train.shape[-1]

   
    model = self._gen_torch_model(inputdim, optimizer, num_treatments=2)
    model, _ = train_survci(model,
                         x_train, t_train, e_train,w_train,
                         x_val, t_val, e_val,w_val,
                         n_iter=iters,
                         lr=learning_rate,
                         elbo=elbo,
                         bs=batch_size)

    self.torch_model = model.eval()
    self.fitted = True

    return self    

  def compute_loss(self, x, t, e,w):

    if not self.fitted:
      raise Exception("The model has not been fitted yet. Please fit the " +
                      "model using the `fit` method on some training data " +
                      "before calling `_eval_nll`.")
      
    processed_data = self._prepocess_training_data(x, t, e,w, 0, None, 0)
    _, _, _,_, x_val, t_val, e_val,w_val = processed_data

    x_val, t_val, e_val,w_val = x_val,\
        _reshape_tensor_with_nans(t_val),\
        _reshape_tensor_with_nans(e_val),\
        _reshape_tensor_with_nans(w_val)
    loss = 0
    ll_loss = float(conditional_loss(self.torch_model,x_val, t_val, e_val,w_val, elbo=False))
    imb = imb_loss(self.torch_model,x_val,w_val)
    mse = mse_loss(self.torch_model,x_val,t_val,w_val)
    l_f = factual_loss(self.torch_model,x_val,t_val,e_val,w_val)
    # mse = mse_total(self.torch_model,x_val,t_val,e_val,w_val)
    loss = ll_loss + self.torch_model.p_alpha*imb + self.torch_model.p_beta*mse
    #loss = ll_loss+self.torch_model.p_alpha*imb + self.torch_model.p_beta*l_f
    
    return loss.detach().numpy() 

  def compute_mse_factual(self, x, t, e,w):

    if not self.fitted:
      raise Exception("The model has not been fitted yet. Please fit the " +
                      "model using the `fit` method on some training data " +
                      "before calling `_eval_nll`.")
      
    processed_data = self._prepocess_training_data(x, t, e,w, 0, None, 0)
    _, _, _,_, x_val, t_val, e_val,w_val = processed_data

    x_val, t_val, e_val,w_val = x_val,\
        _reshape_tensor_with_nans(t_val),\
        _reshape_tensor_with_nans(e_val),\
        _reshape_tensor_with_nans(w_val)
    #loss = 0
    #ll_loss = float(conditional_loss(self.torch_model,x_val, t_val, e_val,w_val, elbo=False))
    #imb = imb_loss(self.torch_model,x_val,w_val)
    mse = mse_loss(self.torch_model,x_val,t_val,w_val)
    # mse = mse_total(self.torch_model,x_val,t_val,e_val,w_val)
    #loss = ll_loss + self.torch_model.p_alpha*imb + self.torch_model.p_beta*mse

    return mse.detach().numpy() 

  def _prepocess_test_data(self, x):
    return torch.from_numpy(x)

  def _prepocess_training_data(self, x, t, e, w,vsize, val_data, random_state):

    idx = list(range(x.shape[0]))
    np.random.seed(random_state)
    np.random.shuffle(idx)
    x_train, t_train, e_train,w_train = x[idx], t[idx], e[idx], w[idx]

    x_train = torch.from_numpy(x_train).double()
    t_train = torch.from_numpy(t_train).double()
    e_train = torch.from_numpy(e_train).double()
    w_train = torch.from_numpy(w_train).double()

    if val_data is None:

      vsize = int(vsize*x_train.shape[0])
      x_val, t_val, e_val,w_val = x_train[-vsize:], t_train[-vsize:], e_train[-vsize:], w_train[-vsize:]

      x_train = x_train[:-vsize]
      t_train = t_train[:-vsize]
      e_train = e_train[:-vsize]
      w_train = w_train[:-vsize]

    else:

      x_val, t_val, e_val,w_val = val_data

      x_val = torch.from_numpy(x_val).double()
      t_val = torch.from_numpy(t_val).double()
      e_val = torch.from_numpy(e_val).double()
      w_val = torch.from_numpy(w_val).double()

    return (x_train, t_train, e_train,w_train,
            x_val, t_val, e_val,w_val)
    
  def predict_dist_parameters(self,x,w):
    x = self._prepocess_test_data(x)
    if self.fitted:
      return get_parameters(self.torch_model,x,w)
    else:
      raise Exception("The model has not been fitted yet. Please fit the " +
                      "model using the `fit` method on some training data " +
                      "before calling `predict_time`.")
      
  def compute_ci(self,x,t,e,w):
    processed_data = self._prepocess_training_data(x, t, e,w, 0, None, 0)
    _, _, _,_, x_val, t_val, e_val,w_val = processed_data

    x_val, t_val, e_val,w_val = x_val,\
        _reshape_tensor_with_nans(t_val),\
        _reshape_tensor_with_nans(e_val),\
        _reshape_tensor_with_nans(w_val)
    if self.fitted:
      treated_idx = torch.where(w_val>0)[0]                        
      control_idx = torch.where(w_val<1)[0]
      shape_co,scale_co,logits_co, shape_tr,scale_tr, logits_tr = get_parameters(self.torch_model,x_val,w_val)  #Predicted Factual Parameters 
      shape_co ,scale_co = softmax_out(shape_co, scale_co, logits_co)
      shape_tr, scale_tr = softmax_out(shape_tr,scale_tr,logits_tr)
      shape_co = shape_co.detach().numpy()
      scale_co  = scale_co.detach().numpy()
      shape_tr = shape_tr.detach().numpy()
      scale_tr = scale_tr.detach().numpy()
      if self.torch_model.dist == 'LogNormal':
        t_co_samples = sample_lognormal(mu=shape_co, sigma=np.exp(scale_co))
        t_tr_samples = sample_lognormal(mu=shape_tr, sigma=np.exp(scale_tr))
      elif self.torch_model.dist == 'Weibull':
        t_co_samples = sample_weibull(shape=shape_co, scale=np.exp(scale_co))
        t_tr_samples = sample_weibull(shape=shape_tr, scale=np.exp(scale_tr))
      else:
        print('Sampling Distribution function not defined')
      t_pred_co = np.median(t_co_samples,axis=1)
      t_pred_tr = np.median(t_tr_samples,axis=1)
      c_index_co = concordance_index(event_times=t_val[control_idx],predicted_scores=t_pred_co,event_observed=e_val[control_idx])
      c_index_tr = concordance_index(event_times=t_val[treated_idx],predicted_scores=t_pred_tr,event_observed=e_val[treated_idx])
      ci_index = (c_index_co + c_index_tr) * 0.5
      return ci_index
    else:
      raise Exception("The model has not been fitted yet. Please fit the " +
                      "model using the `fit` method on some training data " +
                      "before calling `predict_time`.")