# Copyright 2023 DARWIN EU®
#
# This file is part of CohortSurvival
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#' Add time and event status to a cohort table
#'
#' Add the columns needed by standard survival modelling functions, such as
#' `survival::Surv()`, to an OMOP cohort table. This is a lower-level helper:
#' it creates `time` and `status` but does not fit a Kaplan-Meier curve or
#' return a `summarised_result`.
#'
#' `time` is the number of days from target cohort entry to the first applicable
#' event or censoring date. Censoring can occur at the end of observation, at
#' target cohort exit when `censorOnCohortExit = TRUE`, at `censorOnDate`, or at
#' `followUpDays`. `status` is `1` for people with the outcome event and `0`
#' for censored records. Records with an outcome in the washout window are kept
#' in the table with `time` and `status` set to `NA`, so they can be removed by
#' downstream analyses.
#'
#' @param x Cohort table to add survival information to.
#' @param cdm CDM reference created by CDMConnector.
#' @param outcomeCohortTable Name of the cohort table containing the outcome of
#' interest.
#' @param outcomeCohortId ID of event cohorts to include. Only one outcome
#' (and so one ID) can be considered. It can either be a
#' cohort_definition_id value or a cohort_name.
#' @param outcomeDateVariable Variable containing date of outcome event. This is
#' usually `"cohort_start_date"`.
#' @param outcomeWashout Washout time in days for the outcome. If an individual
#' has an outcome during the washout period before target cohort entry, `status`
#' and `time` will be set to `NA`. Use `Inf` for any prior outcome and `0` for
#' no pre-index washout.
#' @param censorOnCohortExit If TRUE, an individual's follow up will be
#' censored at their target cohort exit.
#' @param censorOnDate If not NULL, an individual's follow up will be censored
#' at the given date. This can be a scalar Date or the name of a date column in
#' `x`.
#' @param followUpDays Number of days to follow up individuals (lower bound 1,
#' upper bound Inf). Follow-up is censored at this value.
#' @param name Name of the new table, if NULL a temporary table is returned.
#'
#' @return A cohort table with two additional columns. The `time` column
#' contains the number of days to event or censoring. The `status` column
#' indicates whether the patient had the event (`1`) or was censored (`0`).
#' @export
#'
#' @examples
#' \donttest{
#'
#' cdm <- mockMGUS2cdm()
#' cdm$mgus_diagnosis <- cdm$mgus_diagnosis |>
#'   addCohortSurvival(
#'     cdm = cdm,
#'     outcomeCohortTable = "death_cohort",
#'     outcomeCohortId = 1
#'   )
#'
#' cdm$mgus_diagnosis |>
#'   dplyr::select(subject_id, cohort_start_date, time, status) |>
#'   dplyr::collect()
#' }
#'
addCohortSurvival <- function(x,
                              cdm,
                              outcomeCohortTable,
                              outcomeCohortId = 1,
                              outcomeDateVariable = "cohort_start_date",
                              outcomeWashout = Inf,
                              censorOnCohortExit = FALSE,
                              censorOnDate = NULL,
                              followUpDays = Inf,
                              name = NULL) {

  validateExtractSurvivalInputs(
    cdm = cdm,
    cohortTable = x,
    outcomeCohortTable = outcomeCohortTable,
    outcomeCohortId = outcomeCohortId,
    outcomeWashout = outcomeWashout,
    censorOnCohortExit = censorOnCohortExit,
    censorOnDate = censorOnDate,
    followUpDays = followUpDays
  )

  comp <- newTable(name)

  # drop columns if they already exist
  x <- x |>
    dplyr::select(!dplyr::any_of(c("days_to_exit",
                                   "time",
                                   "status")))

  # get time to end of observation period
  x <- x |>
    PatientProfiles::addFutureObservation(
      indexDate = "cohort_start_date",
      futureObservationName = "days_to_exit"
    )

  if(outcomeWashout == 0) {
    # get any events before or after index date
    x <- x |>
      dplyr::mutate(event_in_washout = 0L) |>
      PatientProfiles::addCohortIntersectDays(
        indexDate = "cohort_start_date",
        targetCohortTable = outcomeCohortTable,
        targetCohortId = outcomeCohortId,
        targetDate = outcomeDateVariable,
        window = c(0, Inf),
        order = "first",
        nameStyle = "days_to_event"
      ) |>
      dplyr::mutate(
        days_to_event = as.numeric(.data$days_to_event)
      )
  } else {
    # get any events before or after index date
    x <- x |>
      PatientProfiles::addCohortIntersectFlag(
        indexDate = "cohort_start_date",
        targetCohortTable = outcomeCohortTable,
        targetCohortId = outcomeCohortId,
        window = c(-outcomeWashout,-1),
        nameStyle = "event_in_washout"
      ) |>
      PatientProfiles::addCohortIntersectDays(
        indexDate = "cohort_start_date",
        targetCohortTable = outcomeCohortTable,
        targetCohortId = outcomeCohortId,
        targetDate = outcomeDateVariable,
        window = c(0, Inf),
        order = "first",
        nameStyle = "days_to_event"
      ) |>
      dplyr::mutate(
        days_to_event = as.numeric(.data$days_to_event)
      )
  }

  # whatever comes first

  # censor at first of
  # 1) outcome,
  # 2) end of observation period
  # 3) cohort exit (if censorOnCohortExit is TRUE)
  # 4) followUpDays (if followUpDays is not Inf)

  if (isTRUE(censorOnCohortExit)) {
    x <- x |>
      dplyr::mutate(days_end_cohort = clock::date_count_between(
        start = .data$cohort_start_date,
        end = .data$cohort_end_date,
        precision = "day"
      )) |>
      dplyr::mutate(days_to_event = dplyr::if_else(
        .data$days_to_event <= .data$days_end_cohort,
        .data$days_to_event, as.numeric(NA)
      )) |>
      dplyr::mutate(days_to_exit = dplyr::if_else(
        .data$days_to_exit < .data$days_end_cohort,
        .data$days_to_exit, .data$days_end_cohort
      )) |>
      dplyr::select(!"days_end_cohort") |>
      dplyr::compute(name = comp$name, temporary = comp$temporary,
                     logPrefix = "CohortSurvival_addCohortSurvival_censor_on_cohort_exit_")
  }

  if (!is.null(censorOnDate)) {
    # create a helper censor_date column either from a scalar Date or a cohort column
    if (is.character(censorOnDate)) {
      x <- x |>
        dplyr::mutate(censor_date = .data[[censorOnDate]])
    } else {
      x <- x |>
        dplyr::mutate(censor_date = .env$censorOnDate)
    }

    x <- x |>
      dplyr::mutate(days_to_censor = clock::date_count_between(
        start = .data$cohort_start_date,
        end = .data$censor_date,
        precision = "day"
      )) |>
      dplyr::mutate(days_to_event = dplyr::if_else(
        .data$days_to_event >= .data$days_to_censor,
        as.numeric(NA), .data$days_to_event
      )) |>
      dplyr::mutate(days_to_exit = dplyr::if_else(
        .data$days_to_exit < .data$days_to_censor,
        .data$days_to_exit, .data$days_to_censor
      )) |>
      dplyr::compute(name = comp$name, temporary = comp$temporary,
                     logPrefix = "CohortSurvival_addCohortSurvival_censor_on_date")
  }

  if (followUpDays != Inf) {
    x <- x |>
      dplyr::mutate(days_to_event = dplyr::if_else(
        .data$days_to_event <= .env$followUpDays,
        .data$days_to_event, as.numeric(NA)
      )) |>
      dplyr::mutate(days_to_exit = dplyr::if_else(
        .data$days_to_exit < .env$followUpDays,
        .data$days_to_exit, .env$followUpDays
      )) |>
      dplyr::compute(name = comp$name, temporary = comp$temporary,
                     logPrefix = "CohortSurvival_addCohortSurvival_follow_up_days_")
  }

  # now just using days_to_event and days_to_exit
  # add status variable (1 if event, 0 if not)
  # add time variable (days to event for those with event,
  # days to exit if no event)
  x <- x |>
    dplyr::mutate(status = dplyr::if_else(
      !is.na(.data$days_to_event), 1, 0
    )) |>
    dplyr::mutate(time = dplyr::if_else(.data$status == 1,
                                        .data$days_to_event, .data$days_to_exit
    ))

  # if censorOnDate was used, now that `status` and `time` exist
  # apply the censor_date effect (set to NA for cohort_start_date > censor_date)
  if (!is.null(censorOnDate)) {
    x <- x |>
      dplyr::compute(name = comp$name, temporary = comp$temporary,
                     logPrefix = "CohortSurvival_addCohortSurvival_censor_on_date_materialized") |>
      dplyr::mutate(
        status = dplyr::if_else(.data$cohort_start_date > .data$censor_date, NA, .data$status),
        time = dplyr::if_else(.data$cohort_start_date > .data$censor_date, NA, .data$time)
      ) |>
      dplyr::select(!c("days_to_censor", "censor_date"))
  }

  # for anyone with an outcome in the washout
  # we keep them, but their time and event will be set to NA
  # (ie they won't contribute to any analysis)
  x <- x |>
    dplyr::mutate(
      status = dplyr::if_else(!is.na(.data$event_in_washout) &
                                .data$event_in_washout == 1, NA,
                              .data$status
      ),
      time = dplyr::if_else(!is.na(.data$event_in_washout) &
                              .data$event_in_washout == 1, NA,
                            .data$time
      )
    )
  # likewise if we censor on a date prior to their cohort start date
  # earlier we applied cohort-start > censor_date adjustments when creating censor_date

  x <- x |>
    dplyr::select(!c("event_in_washout", "days_to_event")) |>
    dplyr::compute(name = comp$name, temporary = comp$temporary,
                   logPrefix = "CohortSurvival_addCohortSurvival_clean")

  return(x)
}

validateExtractSurvivalInputs <- function(cdm,
                                          cohortTable,
                                          outcomeCohortTable,
                                          outcomeCohortId,
                                          outcomeWashout,
                                          censorOnCohortExit,
                                          censorOnDate,
                                          followUpDays) {
  omopgenerics::validateCdmArgument(cdm)
  omopgenerics::validateCohortArgument(cdm[[outcomeCohortTable]])
  checkExposureCohortId(cohortTable)
  if (is.character(censorOnDate)) {
    omopgenerics::assertCharacter(censorOnDate, length = 1)
  } else {
    omopgenerics::assertDate(censorOnDate, null = TRUE)
  }
  checkCensorOnDate(cohortTable, censorOnDate)
  outcomeCohortId <- omopgenerics::validateCohortIdArgument(outcomeCohortId, cdm[[outcomeCohortTable]],
                                                            null = FALSE)
  omopgenerics::assertLogical(censorOnCohortExit, length = 1)
  omopgenerics::assertNumeric(followUpDays, length = 1, min = 1, integerish = TRUE)
  omopgenerics::assertNumeric(outcomeWashout, length = 1, min = 0, integerish = TRUE)
}

newTable <- function(name, call = parent.frame()) {
  omopgenerics::assertCharacter(name, length = 1, null = TRUE, na = TRUE, call = call)
  if (is.null(name) || is.na(name)) {
    x <- list(name = NULL, temporary = TRUE)
  } else {
    x <- list(name = name, temporary = FALSE)
  }
  return(x)
}
