library(tidyverse) # Add the tidyverse package to my current library.
library(haven) # Handle labelled data.
library(broom)
library(splitstackshape) #transform wide data (with stacked variables) to long data

## Import data
for (i in 1:2) {
  assign(paste0("wave", i), #assign is similar to <-; paste0 is to combine wave and i into a name, i ranges from 1 to 6. 
         read_dta(paste0("anchor", i, "_50percent_Eng.dta"))
  )
} 



## Clean data
clean_fun <- function(df) {  df %>% 
    transmute(
      id=zap_label(id), #remove label of id
      age=zap_label(age), #remove label of age
      wave=as.numeric(wave),
      sex=as_factor(sex_gen), #make sex_gen as a factor
      relstat=as_factor(relstat), #make relstat as a factor
      relstat=case_when(relstat== "-7 Incomplete data" ~ as.character(NA), #specify when is missing for relstat
                        TRUE ~ as.character(relstat))%>% as_factor(), #make relstat as a factor again
      hlt=case_when(hlt1<0 ~ as.numeric(NA),  #specify when hlt1 is missing 
                    TRUE ~ as.numeric(hlt1)),
      sat=case_when(sat6<0 ~ as.numeric(NA), #specify when sat6 is missing
                    TRUE ~ as.numeric(sat6)),
      ptner=case_when(
        relstat %in% c("1 Never married single","6 Divorced/separated single","9 Widowed single") ~ "No",
        # when relstat has any of the three situations, I assign "Nevermarried" to new variable "marital1"
        relstat %in% c("2 Never married LAT","3 Never married COHAB",
                       "4 Married COHAB","5 Married noncohabiting",
                       "7 Divorced/separated LAT","8 Divorced/separated COHAB",
                       "10 Widowed LAT", "11 Widowed COHAB") ~ 'Yes') %>% as_factor()
    )%>% drop_na()  
  }
wave1a <- clean_fun(wave1)
wave2a <- clean_fun(wave2)
#wave3a <- clean_fun(wave3)
#wave4a <- clean_fun(wave4)
#wave5a <- clean_fun(wave5)
#wave6a <- clean_fun(wave6)



## Sample selection
wave1b <- wave1a %>% filter(ptner=="No")%>%
  rename(wave.1=wave, age.1=age, sex.1=sex, relstat.1=relstat, hlt.1=hlt, sat.1=sat, ptner.1=ptner) #rename variables

wave2b <- wave2a %>% 
  rename(wave.2=wave, age.2=age, sex.2=sex, relstat.2=relstat, hlt.2=hlt, sat.2=sat, ptner.2=ptner)

#wave3b <- wave3a %>% 
#  rename(wave.3=wave, age.3=age, sex.3=sex, relstat.3=relstat, hlt.3=hlt, sat.3=sat, ptner.3=ptner)

#wave4b <- wave4a %>% 
#  rename(wave.4=wave, age.4=age, sex.4=sex, relstat.4=relstat, hlt.4=hlt, sat.4=sat, ptner.4=ptner)

#wave5b <- wave5a %>% 
#  rename(wave.5=wave, age.5=age, sex.5=sex, relstat.5=relstat, hlt.5=hlt, sat.5=sat, ptner.5=ptner)

#wave6b <- wave6a %>% 
#  rename(wave.6=wave, age.6=age, sex.6=sex, relstat.6=relstat, hlt.6=hlt, sat.6=sat, ptner.6=ptner)

wide_data <- left_join(wave1b, wave2b, by = "id") 
#%>%  # left join wave1b and wave2b
#  left_join(wave3b, by = "id") %>% # left join with wave3b
#  left_join(wave4b, by = "id") %>% # left join with wave4b
#  left_join(wave5b, by = "id") %>% # left join with wave5b
#  left_join(wave6b, by = "id") # left join with wave6b
#by using left_join I keep those single in wave1 and follow them



## Check the participation over time

wide_data$check <- paste(wide_data$wave.1, wide_data$wave.2, sep='-')
#wide_data$check <- paste(wide_data$wave.1, wide_data$wave.2, wide_data$wave.3,
#                         wide_data$wave.4, wide_data$wave.5, wide_data$wave.6, sep='-')

table(wide_data$check)
## 
##  1-2 1-NA 
## 1971  620
## Generate a panel dataset 
#Step 1: make a panel dataset (person-period dataset)
long_data<- merged.stack(wide_data, #dataset for transfrom
                          var.stubs = c("age", "wave", "sex","relstat", "hlt", "sat", "ptner"), 
                          #var.stubs is to specify the prefixes of the variable groups
                          sep = ".") %>%  
  #sep is to specify the character that separates the "variable name" from the "times" in the source
  drop_na(wave)
#drop the observations which did not join the wave

#Step 2: generate transition variables
long_data <- long_data %>% 
  group_by(id) %>% 
  mutate(
    havepartner=case_when( ptner!=dplyr::lag(ptner, 1) & dplyr::lag(ptner, 1)=="No" & ptner=="Yes" ~ 1,
                          TRUE ~ 0)
  )

#step3: remove observation upon and after first union dissolution
long_data <- long_data %>% 
  group_by(id) %>% mutate(
     wave=as.numeric(wave),
     partnerwave=case_when(havepartner==1 ~ wave), 
     partnerwave1=min(partnerwave, na.rm = TRUE)
  ) #identify the timing of first union formation and first union dissolution
#There will be many warnings, whenever you attempt to find the minimum or maximum value of a vector that has a length of zero. It wonโ€™t actually prevent your code from running.
long_data$partnerwave1[is.infinite(long_data$partnerwave1)] <- 0 #those who are untreated the value is 0.