# AUTOGENERATED! DO NOT EDIT! File to edit: 59_hazard.Cox.ipynb (unless otherwise specified).

__all__ = ['ProportionalHazard']

# Cell
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.preprocessing import MaxAbsScaler, StandardScaler

from ..losses import hazard_loss
from .ph import PieceWiseHazard

# torch.Tensor.ndim = property(lambda x: x.dim())

# Cell
class ProportionalHazard(nn.Module):
    """
    Hazard proportional to time and feature component as shown above.
    parameters:
    - breakpoints: time points where hazard would change
    - max_t: maximum point of time to plot to.
    - dim: number of input dimensions of x
    - h: (optional) number of hidden units (for x only).
    """
    def __init__(self, breakpoints:np.array, t_scaler:MaxAbsScaler, x_scaler:StandardScaler,
                 dim:int, h:tuple=(), **kwargs):
        super().__init__()
        self.baseλ = PieceWiseHazard(breakpoints, t_scaler)
        self.x_scaler = x_scaler
        nodes = (dim,) + h + (1,)
        self.layers = nn.ModuleList([nn.Linear(a,b, bias=False)
                                   for a,b in zip(nodes[:-1], nodes[1:])])

    def forward(self, t, t_section, x):
        logλ, Λ = self.baseλ(t, t_section)

        for layer in self.layers[:-1]:
            x = F.relu(layer(x))
        log_hx = self.layers[-1](x)

        logλ += log_hx
        Λ = torch.exp(log_hx + torch.log(Λ))
        return logλ, Λ

    def survival_function(self, t:np.array, x:np.array) -> torch.Tensor:
        if len(t.shape) == 1:
            t = t[:,None]
        t = self.baseλ.t_scaler.transform(t)
        if len(x.shape) == 1:
            x = x[None, :]
        if len(x) == 1:
            x = np.repeat(x, len(t), axis=0)
        x = self.x_scaler.transform(x)


        with torch.no_grad():
            x = torch.Tensor(x)
            # get the times and time sections for survival function
            breakpoints = self.baseλ.breakpoints[1:].cpu().numpy()
            t_sec_query = np.searchsorted(breakpoints.squeeze(), t.squeeze())
            # convert to pytorch tensors
            t_query = torch.Tensor(t)
            t_sec_query = torch.LongTensor(t_sec_query)

            # calculate cumulative hazard according to above
            _, Λ = self.forward(t_query, t_sec_query, x)
            return torch.exp(-Λ)


    def plot_survival_function(self, t:np.array, x:np.array) -> None:
        s = self.survival_function(t, x)

        # plot
        plt.figure(figsize=(12,5))
        plt.plot(t, s)
        plt.xlabel('Time')
        plt.ylabel('Survival Probability')
        plt.show()