#------------------------------------------------------------------------------
#------------------------------------------------------------------------------
#        2.3. Imputation Method for MER Historical Distributions
#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

# Objective: This file observes the ratios between the shares of the PPP and MER 
# distributions in 1980 for the residual regions and exploits these ratios to extend 
# this comparative behavior backward from 1980 to 1820, in order to generate full 
# historical MER distributions for the residual regions. The estimation of continental 
# and world MER distributions is carried out in the main.do code by aggregating 
# countries and residual regions.

library(haven)
library(dplyr)
library(purrr)
library (wid)
library(tidyr)
library(stringr)


project_dir <- "/Users/manuelestebanarias/Desktop/2025_replication_extension_nonmain_mer"
#-------------------------------------------------------------------------------
#          A.  Generate adjustment ratios PPP-MER 1980 
#-------------------------------------------------------------------------------

#--------- 1. Import wid distributions -----------------------------------------
current_80 <- read_dta(
  file.path(project_dir, "11_ptinc_hweal_1980.dta")
)

#--------- 2. Collaps distributions into standard percentiles ------------------
# Create fractile groups
breaks <- c(0,10000,25000,50000,75000, 90000, 99000, 100000)

data80 <- current_80 %>%
  mutate(
    ftile_index = cut(p, breaks = breaks, include.lowest = TRUE, right = FALSE, labels = FALSE),
    ftile = breaks[ftile_index]  # this gives you the left edge of the interval
  )

# Calculate population
data80 <- data80 %>%
  group_by(iso, name) %>%
  arrange(p) %>%
  mutate(
    lead_p = lead(p),
    pop = (lead_p - p) / 1e5,
    pop = ifelse(is.na(pop), 1 / 1e5, pop)
  ) %>%
  ungroup() %>%
  select(-lead_p)

# Collapse data by iso-year-ftile
data80 <- data80 %>%
  group_by(iso, name, ftile) %>%
  summarise(
    t_PPP = min(t_PPP, na.rm = TRUE),
    s_PPP = sum(s_PPP, na.rm = TRUE),       
    a_PPP = sum(a_PPP*pop, na.rm = TRUE),
    
    t_MER = min(t_MER, na.rm = TRUE),
    s_MER = sum(s_MER, na.rm = TRUE),       
    a_MER = sum(a_MER*pop, na.rm = TRUE),
    
    pop = sum(pop, na.rm = TRUE),
    .groups = "drop"
  )


data80 <- data80 %>%
  rename(p = ftile) %>%
  mutate(
    a_PPP = a_PPP/pop,
    a_MER = a_MER/pop,
    t_PPP = t_PPP/pop,
    t_MER = t_MER/pop,
  )     

#--------- 3. Estimate adjustment ratios ----------------------------------------

deltas <- data80 %>%
  mutate(
    d_t = t_MER - t_PPP,
    d_s = s_MER - s_PPP,
    r_t = t_MER / t_PPP
  ) %>%
  mutate (
    r_t = if_else(is.na(r_t), 1, r_t)
  ) %>%
  select(iso, name, p, d_s, d_t,r_t)


#-------------------------------------------------------------------------------
#          B.  Extend PPP series to MER 
#-------------------------------------------------------------------------------

aggregates_hist <- download_wid (
  indicators = c("ahweal","anninc"),
  areas      = c("OA-PPP","OB-PPP","OC-PPP","OD-PPP","OE-PPP","OH-PPP","OI-PPP","OJ-PPP","OK-PPP","OL-PPP","QM-PPP"),
  year       = 1820:1980,
  perc       = "p0p100",
  ages       = 992,
  pop        = "i")

aggregates_hist <- aggregates_hist %>%
  pivot_wider(
    names_from  = variable,
    values_from = value
  ) %>%
  select(-percentile)
  

#--------- 1. Import wid distributions -----------------------------------------
pretax_hist <- download_wid (
  indicators = c("sptinc","aptinc","tptinc"),
  areas      = c("OA-PPP","OB-PPP","OC-PPP","OD-PPP","OE-PPP","OH-PPP","OI-PPP","OJ-PPP","OK-PPP","OL-PPP","QM-PPP"),
  year       = 1820:1980,
  ages       = 992)

pretax_hist$variable <- substr(pretax_hist$variable, 1, 1)

pretax_hist <- pretax_hist %>%
  rename(
    p       = percentile
  )

pretax_hist <- pretax_hist %>%
  # Parsing the percentiles
  mutate(
    p_min = NA_real_,
    p_max = NA_real_
  ) %>%
  mutate(
    p_min = if_else(
      str_detect(p, "^p([0-9\\.]+)p([0-9\\.]+)$"),
      round(1000 * as.numeric(str_match(p, "^p([0-9\\.]+)p([0-9\\.]+)$")[,2])),
      p_min
    ),
    p_max = if_else(
      str_detect(p, "^p([0-9\\.]+)p([0-9\\.]+)$"),
      round(1000 * as.numeric(str_match(p, "^p([0-9\\.]+)p([0-9\\.]+)$")[,3])),
      p_max
    )
  ) %>%
  mutate(
    p_min = if_else(
      str_detect(p, "^p([0-9\\.]+)$"),
      round(1000 * as.numeric(str_match(p, "^p([0-9\\.]+)$")[,2])),
      p_min
    ),
    p_max = if_else(is.na(p_max), 1000*100, p_max)
  ) %>%
  mutate(
    p_max = case_when(
      is.na(p_max) & between(p_min, 0, 98000)   ~ p_min + 1000,
      is.na(p_max) & between(p_min, 99000, 99800) ~ p_min + 100,
      is.na(p_max) & between(p_min, 99900, 99980) ~ p_min + 10,
      is.na(p_max) & between(p_min, 99990, 99999) ~ p_min + 1,
      TRUE ~ p_max
    )
  ) %>%
  # Reconstruct p 
  mutate(
    p = if_else(
      !is.na(p_max),
      paste0("p", round(p_min / 1e3, 3), "p", round(p_max / 1e3, 3)),
      p
    )
  ) %>%
  # Keeping the g-percentiles
  mutate(
    n = round(p_max - p_min, 1)
  )%>%
  filter(n %in% c(1, 10, 100, 1000)) %>%
  filter(!(n == 1000 & p_min >= 99000)) %>%
  filter(!(n == 100  & p_min >= 99900)) %>%
  filter(!(n == 10   & p_min >= 99990)) %>%
  # Clean columns
  select(-p_max, -p) %>%
  rename(p = p_min)

pretax_hist <- pretax_hist %>%
  pivot_wider(
    names_from  = variable,
    values_from = value
  )%>%
  mutate(name = "historical_sptinc992j")

pretax_hist <- pretax_hist %>%
  left_join(aggregates_hist, by = c("country", "year")) 

pretax_hist <- pretax_hist %>%
  mutate(
    a = a/anninc992i,
    t = t/anninc992i
  ) %>%
  select(- ahweal992i,- anninc992i) %>%
  rename(iso=country) 


wealth_hist <- download_wid (
    indicators = c("shweal","ahweal","thweal"),
    areas      = c("OA-PPP","OB-PPP","OC-PPP","OD-PPP","OE-PPP","OH-PPP","OI-PPP","OJ-PPP","OK-PPP","OL-PPP","QM-PPP"),
    year       = 1820:1980,
    ages       = 992)

wealth_hist$variable <- substr(wealth_hist$variable, 1, 1)

wealth_hist <- wealth_hist %>%
  rename(
    p       = percentile
  )

wealth_hist <- wealth_hist %>%
  # Parsing the percentiles
  mutate(
    p_min = NA_real_,
    p_max = NA_real_
  ) %>%
  mutate(
    p_min = if_else(
      str_detect(p, "^p([0-9\\.]+)p([0-9\\.]+)$"),
      round(1000 * as.numeric(str_match(p, "^p([0-9\\.]+)p([0-9\\.]+)$")[,2])),
      p_min
    ),
    p_max = if_else(
      str_detect(p, "^p([0-9\\.]+)p([0-9\\.]+)$"),
      round(1000 * as.numeric(str_match(p, "^p([0-9\\.]+)p([0-9\\.]+)$")[,3])),
      p_max
    )
  ) %>%
  mutate(
    p_min = if_else(
      str_detect(p, "^p([0-9\\.]+)$"),
      round(1000 * as.numeric(str_match(p, "^p([0-9\\.]+)$")[,2])),
      p_min
    ),
    p_max = if_else(is.na(p_max), 1000*100, p_max)
  ) %>%
  mutate(
    p_max = case_when(
      is.na(p_max) & between(p_min, 0, 98000)   ~ p_min + 1000,
      is.na(p_max) & between(p_min, 99000, 99800) ~ p_min + 100,
      is.na(p_max) & between(p_min, 99900, 99980) ~ p_min + 10,
      is.na(p_max) & between(p_min, 99990, 99999) ~ p_min + 1,
      TRUE ~ p_max
    )
  ) %>%
  # Reconstruct p 
  mutate(
    p = if_else(
      !is.na(p_max),
      paste0("p", round(p_min / 1e3, 3), "p", round(p_max / 1e3, 3)),
      p
    )
  ) %>%
  # Keeping the g-percentiles
  mutate(
    n = round(p_max - p_min, 1)
  )%>%
  filter(n %in% c(1, 10, 100, 1000)) %>%
  filter(!(n == 1000 & p_min >= 99000)) %>%
  filter(!(n == 100  & p_min >= 99900)) %>%
  filter(!(n == 10   & p_min >= 99990)) %>%
  # Clean columns
  select(-p_max, -p) %>%
  rename(p = p_min)

wealth_hist <- wealth_hist %>%
  pivot_wider(
    names_from  = variable,
    values_from = value
  )%>%
  mutate(name = "historical_sptinc992j")


wealth_hist <- wealth_hist %>%
  left_join(aggregates_hist, by = c("country", "year")) %>%
  mutate(
    a = a/ahweal992i,
    t = t/ahweal992i
  ) 

wealth_hist <- wealth_hist %>%
  select(-anninc992i, -ahweal992i)%>%
  mutate(
    name = "historical_shweal992j"
  )%>%
  rename(iso=country) 

# complete the subregions within OH.
wealth_hist <- wealth_hist %>%
  bind_rows(
    wealth_hist %>%
      filter(iso == "OH") %>%
      mutate(iso = "OK")
  )
wealth_hist <- wealth_hist %>%
  bind_rows(
    wealth_hist %>%
      filter(iso == "OH") %>%
      mutate(iso = "OL")
  )

current_wid<-rbind(pretax_hist, wealth_hist)


#--------- 2. Collapse data into g-pinter input format ------------------------- 
# Create frac-tile groups
breaks <- c(0,10000,25000,50000,75000, 90000, 99000, 100000)

data_org <- current_wid %>%
  mutate(
    ftile_index = cut(p, breaks = breaks, include.lowest = TRUE, right = FALSE, labels = FALSE),
    ftile = breaks[ftile_index]  # this gives you the left edge of the interval
  )

# Calculate population
data_org <- data_org %>%
  group_by(iso, year, name) %>%
  arrange(p) %>%
  mutate(
    lead_p = lead(p),
    pop = (lead_p - p) / 1e5,
    pop = ifelse(is.na(pop), 1 / 1e5, pop)
  ) %>%
  ungroup() %>%
  select(-lead_p)

# Collapse data by iso-year-ftile
data_org <- data_org %>%
  group_by(iso, year, name, ftile) %>%
  summarise(
    t = min(t, na.rm = TRUE),
    s = sum(s, na.rm = TRUE),       
    a = sum(a*pop, na.rm = TRUE),
    pop = sum(pop, na.rm = TRUE),
    .groups = "drop"
  )


data_org <- data_org %>%
  rename(p = ftile) %>%
  mutate(
    a = a/pop,
    t = t/pop,
  )   

# Verify consistency
data_org <- data_org %>%
  group_by(iso, year, name) %>%
  mutate(
    shr = sum(s, na.rm = TRUE),
    apop = a * pop,
    avg = sum(apop, na.rm = TRUE)
  ) %>%
  ungroup()

# Check: average and share should be close to 1
stopifnot(all(abs(data_org$shr - 1) < 1e-5))
stopifnot(all(abs(data_org$avg - 1) < 1e-5))

#  Cleanup
data_org <- data_org %>%
  select(-shr, -avg, -apop) 

#--------- 3. Transform series from PPP to MER  -------------------------------- 
# Transform
data<- data_org %>%
  left_join(deltas, by = c("iso","p","name")) %>%
  mutate(
    s_ppp = s,
    t_ppp = t,
    a_ppp = a,
    #t = t *r_t,
    s = s + d_s
  )

# Ensure shares consistency
data <- data %>%
  group_by(iso, year, name) %>%
  mutate(
    shr  = sum(s, na.rm = TRUE),
    apop = a * pop,
    avg  = sum(apop, na.rm = TRUE)
  ) %>%
  ungroup()

data<- data %>%
  mutate(
    s=s/shr
  )

# Correct the rest of the distribution
data <- data %>%
  mutate(
    a = (s*1)/pop
  )

# Drop observations that violate the average monotonicity
data<- data%>%
  group_by(name, iso, year) %>%
  mutate(
    a=if_else(lag(a)>a & p>0 & p<90000, NA, a)
  )

# Replase dropped averages
data <- data %>%
  group_by(iso, year, name) %>%
  mutate(
    a75 = a[p == 75000][1],
    a90 = a[p == 90000][1],
    a99 = a[p == 99000][1],
    
    m = mean(c(a75, a90, a99), na.rm = TRUE),
    alpha = pmax(1, a99 / m - 1)
 )


fill_block <- function(p, pL, pU, aL, aU, alpha) {
  aL + (aU - aL) * (((p - pL) / (pU - pL))^(1 + alpha))
}


data <- data %>%
  group_by(iso, year, name) %>%
  arrange(p) %>%
  mutate(
    observed = !is.na(a)
  ) %>%
  mutate(
    # upper bound
    pU = map_dbl(row_number(), ~ {
      idx <- which(observed & p > p[.x])
      if (length(idx) == 0) NA_real_ else p[min(idx)]
    }),
    aU = map_dbl(row_number(), ~ {
      idx <- which(observed & p > p[.x])
      if (length(idx) == 0) NA_real_ else a[min(idx)]
    }),
    # lower bound
    pL = map_dbl(row_number(), ~ { 
      idx <- which(observed & p < p[.x])
      if(length(idx)==0) p[which.min(p[observed])] else p[max(idx)]
    }),
    aL = map_dbl(row_number(), ~ { 
      idx <- which(observed & p < p[.x])
      if(length(idx)==0) a[which.min(p[observed])] else a[max(idx)]
    })
  ) %>%
  mutate(
    a2 = if_else(
      is.na(a) & !is.na(pU) & !is.na(aL),
      fill_block(p, pL, pU, aL, aU, alpha),
      NA
    ),
    a = if_else(
      is.na(a),
      a2,
      a
    )
  ) %>%
  ungroup()



data <- data %>%
  group_by(iso, year,name) %>%
  mutate(
    total = sum(a * pop, na.rm = TRUE),
    a = a / total
  ) %>%
  ungroup()

# correct shares again if necessary
data <- data %>%
  mutate(
     s = a*pop
  )

# generate tresholds 
data <- data %>%  
  group_by(name, iso, year) %>%
  arrange(p, .by_group = TRUE) %>%
  mutate(
    t = ((a - lag(a)) / 2) + lag(a)
  ) %>%
  mutate(
    #t=if_else(p==0, t_ppp*(a_ppp/a)*pop, t)
    t=if_else(p==0, t_ppp*pop + (a-a_ppp), t)
  )


# Verify consistency
data <- data %>%
  group_by(iso, year, name) %>%
  mutate(
    shr  = sum(s, na.rm = TRUE),
    apop = a * pop,
    avg  = sum(apop, na.rm = TRUE)
  ) %>%
  ungroup()

# Check: average and share should be close to 1
stopifnot(all(abs(data$shr - 1) < 1e-4))
stopifnot(all(abs(data$avg - 1) < 1e-4))

#  Cleanup
data <- data %>%
  select(iso, year, p, name, pop, t, s, a)  # %>%
  #mutate(
  #  iso = paste0(iso,"-MER")
  #)

#------------------------------------------------------------------------------
#------------------------------------------------------------------------------
#            6. Run the Gpinter
#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

library(tidyverse)
library(haven)
library(gpinter)
library(here)
library(progressr)

handlers(handler_progress(
  format = ":spin :current/:total [:bar] :percent in :elapsedfull ETA: :eta"
))
options(dplyr.summarise.inform = FALSE)

gperc <- c(
  seq(0, 99000, 1000), seq(99100, 99900, 100),
  seq(99910, 99990, 10), seq(99991, 99999, 1)
)

data <- bind_rows(data) %>%
  arrange(iso, year, p)


groups <- data %>%
  select(iso, year, pop ,p, name, a, t) %>%
  drop_na() %>%
  group_by(iso, year, name) %>%
  mutate(average = sum(a * pop, na.rm = TRUE)) %>%
  arrange(iso, year, name, p) %>%
  group_by(iso, year, name) %>%
  filter(n() > 2)%>%
  group_split()


with_progress({
  p <- progressor(along = groups)
  gpinterized <- groups %>% map_dfr(~ {
    dist <- tabulation_fit(p = .x$p/1e5, threshold = .x$t,  bracketavg = .x$a, average = .x$average[1], fast = TRUE)
    tab <- generate_tabulation(dist, gperc/1e5) %>% as_tibble() %>% transmute(
      iso = first(.x$iso),
      year = first(.x$year),
      name = first(.x$name),
      p = round(fractile*1e5),
      t = threshold,
      ts = top_share,
      bs = bottom_share,
      s = bracket_share,
      a = bracket_average
    )
    p()
    return(tab)
  })
})

write_dta(gpinterized, here(
  file.path(project_dir,"output3_hist_pretax_hweal_mer_gpinterized.dta")))
