library(dplyr, warn.conflicts = FALSE)
library(survival)

test_that("tsesimp: weibull aft", {
  # the eventual survival time
  shilong1 <- shilong %>%
    arrange(bras.f, id, tstop) %>%
    group_by(bras.f, id) %>%
    filter(row_number() == n()) %>%
    select(-c("ps", "ttc", "tran"))
  
  # the last value of time-dependent covariates before pd
  shilong2 <- shilong %>%
    filter(pd == 0 | tstart <= dpd) %>%
    arrange(bras.f, id, tstop) %>%
    group_by(bras.f, id) %>%
    filter(row_number() == n()) %>%
    select(bras.f, id, ps, ttc, tran)
  
  # combine baseline and time-dependent covariates
  shilong3 <- shilong1 %>%
    left_join(shilong2, by = c("bras.f", "id"))
  
  # apply the two-stage method
  fit1 <- tsesimp(
    data = shilong3, time = "tstop", event = "event",
    treat = "bras.f", censor_time = "dcut", pd = "pd",
    pd_time = "dpd", swtrt = "co", swtrt_time = "dco",
    base_cov = "",
    base2_cov = c("agerand", "sex.f", "tt_Lnum", "rmh_alea.c",
                  "pathway.f", "ps", "ttc", "tran"),
    aft_dist = "weibull", alpha = 0.05,
    recensor = TRUE, swtrt_control_only = FALSE, offset = 1,
    boot = FALSE)

  # numeric code of treatment and apply administrative censoring
  data1 <- shilong3 %>% 
    mutate(treat = ifelse(bras.f == "CT", 0, 1),
           swtrt = 1*co,
           censor_time = ifelse(event == 1, dcut, tstop))
  
  tablist <- lapply(0:1, function(h) {
    df1 <- data1 %>% 
      filter(treat == h & pd == 1) %>%
      mutate(time = tstop - dpd + 1)
    
    fit_aft <- survreg(Surv(time, event) ~ swtrt + agerand + sex.f 
                       + tt_Lnum + rmh_alea.c + pathway.f
                       + ps + ttc + tran, 
                       data = df1, dist = "weibull")
    
    psi = -fit_aft$coefficients[2]

    data1 %>% 
      filter(treat == h) %>%
      mutate(u = ifelse(swtrt == 1, 
                        ifelse(pd == 1, dpd-1 + (tstop-dpd+1)*exp(psi), 
                               dco-1 + (tstop-dco+1)*exp(psi)), tstop),
             c = pmin(censor_time, censor_time*exp(psi)),
             time_ts = pmin(u, c),
             event_ts = event*(u <= c))
  })
  data2 <- do.call(rbind, tablist)
  
  fit <- coxph(Surv(time_ts, event_ts) ~ treat, data = data2)
  hr1 <- exp(as.numeric(c(fit$coefficients[1], confint(fit)[1,])))
  
  expect_equal(hr1, c(fit1$hr, fit1$hr_CI))
})