#------------------------------------------------------------------------------
#------------------------------------------------------------------------------
# 3.1. Imputation Method for Distributional data for Additional Countries-Years
#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

# Objective: This file imports the pretax data, uses residual regions to complete 
# missing countries (i.e. non-main countries), and relies on completed fiscal data 
# to emulate the path of the top 1% of each country (and therefore of the residual 
# region). The remainder of the series is then adjusted to keep the distributions 
# internally consistent. Finally, the G-P interpolation is used to complete the 
# series.

#----------------------  Index  ------------------------------------------------
# -------- 0. Packages
#--------- 1. Import Data 
# -------------- 1.1. Import current data (already gpintereised and corrected) 
# -------------- 1.2. Import fiscal data 
#--------- 2.  Complete Percentiles Top 1% if not available
#--------- 3. Collapse data into g-pinter input format
#--------- 4. Calculate the correction ratios
#--------- 5. Transform the series
# -------------- 5.1 Find an optimal penalization lambda
# -------------- 5.2 Generate Pretax series series based on the fiscal data
# -------------- 5.3 Use the best_lambda for penalizing transformation
#--------- 6. Run the Gpinter
# ------------------------------------------------------------------------------

# -------- 0. Packages ---------------------------------------------------------
library(haven)
library(dplyr)
library(tidyr)
library(purrr)
library(zoo)
library(wid)

project_dir <- "/Users/manuelestebanarias/Desktop/2025_replication_extension_nonmain_mer"

#--------- 1. Import Data      -------------------------------------------------

aggregates_hist <- download_wid (
  indicators = c("anninc"),
  year       = 1850:1979,
  perc       = "p0p100",
  ages       = 992,
  pop        = "i")

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


# -------------- 1.1. Import current data (already gpintereised and corrected) -
current_wid <- current_wid  %>%
  filter(name == "historical_sptinc992j") %>%
  filter(year>=1850 & year <1980)

current_wid <- download_wid (
  indicators = c("sptinc","aptinc","tptinc"),
  year       = 1850:1979,
  ages       = 992) %>%
  mutate(name = "historical_sptinc992j")

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

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

current_wid  <- current_wid  %>%
  # 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)

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

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

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


# Manual corrections
  # Assign Eastern Europe to BG, PL, HR, HU
new_isos <- c("BG","HU", "HR", "PL")

current_wid <- current_wid %>%
  filter(iso != "QM") %>%
  bind_rows(
    current_wid %>%
      filter(iso == "QM") %>%
      crossing(new_iso = new_isos) %>%   # use a temporary column
      mutate(iso = new_iso) %>%          # replace old iso
      select(-new_iso)                   # remove the temp column
  )


# Assign Western Europe to IE, FI, GR, PT, CH
new_isos <- c("FI","IE", "GR", "PT","CH")

current_wid <- current_wid %>%
  filter(iso != "OC") %>%
  bind_rows(
    current_wid %>%
      filter(iso == "OC") %>%
      crossing(new_iso = new_isos) %>%   # use a temporary column
      mutate(iso = new_iso) %>%          # replace old iso
      select(-new_iso)                   # remove the temp column
  )

# Assign Mena to TN
current_wid <- current_wid %>%
  mutate(
    iso = if_else(iso=="OE", "TN", iso)
  )
# Assign Other sub Saharan Africa to GH, MU, MW, TZ, UG, ZM, ZW, SC, CM
new_isos <- c("GH", "MU", "MW", "TZ", "UG", "ZM","ZW", "SC","CM")

current_wid <- current_wid %>%
  filter(iso != "OJ") %>%
  bind_rows(
    current_wid %>%
      filter(iso == "OJ") %>%
      crossing(new_iso = new_isos) %>%   # use a temporary column
      mutate(iso = new_iso) %>%          # replace old iso
      select(-new_iso)                   # remove the temp column
  )

  # asing OJ to the SG and MY
new_isos <- c("SG","MY")

current_wid <- current_wid %>%
  filter(iso != "OI") %>%
  bind_rows(
    current_wid %>%
      filter(iso == "OI") %>%
      crossing(new_iso = new_isos) %>%   # use a temporary column
      mutate(iso = new_iso) %>%          # replace old iso
      select(-new_iso)                   # remove the temp column
  ) %>%
  filter(iso!="SG" & year!=1970)


current_wid_1980 <- read_dta(file.path(project_dir,"/22_1980_ptinc.dta"))
current_wid_hist <- read_dta(file.path(project_dir,"/24_ptinc_historical.dta"))

current_wid <- bind_rows(current_wid, current_wid_CL, current_wid_1980, current_wid_hist)

#Drop regional data
current_wid <- current_wid %>%
  filter(!iso %in% c("OA", "OB", "OC", "OD", "OE", "OH", "OI", "OJ", "OK", "OL", "QM")) %>%
  filter(year>=1850 & year <=1980)

# -------------- 1.2. Import fiscal data 
fiscal_data<-read_dta(file.path(project_dir,"/23_fiinc_historical_top.dta")) 

#fiscal_data <- fiscal_data%>%
#  mutate(
#    n_available = rowSums(!is.na(select(., si, sj, st)))
#  )

fiscal_data <- fiscal_data %>%
  mutate(
    sj = if_else(iso=="US", NA, sj), # T is more complete
    st = if_else(iso=="FR", NA, st), #Both available,
  
  )
  
fiscal_data_2 <- fiscal_data%>%
  group_by(iso, year) %>%
  mutate(has_p99 = any(p == 99000)) %>%
  ungroup()%>%
  filter(has_p99==FALSE) %>%
  select(-has_p99)


#--------- 2.  Complete Percentiles Top 1% if not available -------------------
# Create fractile groups
breaks99 <- c(0,10000,25000,50000,75000, 90000, 99000, 99900, 100000)

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

# Calculate population
data_org99 <- data_org99 %>%
  group_by(iso, year) %>%
  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_org99 <- data_org99 %>%
  group_by(iso, year, 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_org99 <- data_org99 %>%
  rename(p = ftile) %>%
  mutate(
    a = a/pop
  )   

# Verify consistency
data_org99 <- data_org99 %>%
  group_by(iso, year) %>%
  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_org99$shr - 1) < 1e-6))
stopifnot(all(abs(data_org99$avg - 1) < 1e-6))

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


# Make sure data is sorted
data_org99 <- data_org99 %>%
  mutate(
    benchmark=1
  )%>%
  arrange(iso, p, year)

# Generate full year observations and Interpolate missing years between 1870 and 1980
data99 <- data_org99 %>%
  group_by(iso, p) %>%
  complete(year = full_seq(1850:1980, 1)) %>%
  arrange(iso, p, year) %>%
  mutate(
    t = if (sum(!is.na(t)) > 1) {
      approx(year[!is.na(t)], t[!is.na(t)], xout = year, rule = 2)$y
    } else {
      # If only one observation, just fill all with that single value (or NA if none)
      rep(t[!is.na(t)][1], length(year))
    },
    s = if (sum(!is.na(s)) > 1) {
      approx(year[!is.na(s)], s[!is.na(s)], xout = year, rule = 2)$y
    } else {
      rep(s[!is.na(s)][1], length(year))
    },
    a = if (sum(!is.na(a)) > 1) {
      approx(year[!is.na(a)], a[!is.na(a)], xout = year, rule = 2)$y
    } else {
      rep(a[!is.na(a)][1], length(year))
    }
  ) %>%
  ungroup() %>%
  mutate(
    pop = case_when(
      p == 0     ~ 0.1,
      p == 10000 ~ 0.15,
      p == 25000 ~ 0.25,
      p == 50000 ~ 0.25,
      p == 75000 ~ 0.15,
      p == 90000 ~ 0.09,
      p == 99000 ~ 0.009,
      p == 99900 ~ 0.001,
      TRUE       ~ pop
    )
  ) %>%
  arrange(iso, year, p)

# Keep countries with fiscal top 0,1
data_ratio99 <- data99 %>%
  left_join(fiscal_data_2, by = c("iso", "year","p")) %>%
  group_by(iso) %>%
  filter(any(!is.na(si) | !is.na(sj) )) %>%
  ungroup() 

# Calcualte ratio 
data_ratio99 <- data_ratio99 %>%
  group_by(iso, year) %>%
  mutate(s_99000 = s[p == 99000],
         s_99900 = s[p == 99900]
  ) %>%
  ungroup() %>%
  mutate( 
    ratio99_999=s_99000/s_99900
  ) %>%
  select(-s_99000, -s_99900)

# complete 99 fiinc
data_ratio99 <- data_ratio99 %>%
  group_by(iso, year) %>%
  mutate(
    si_999 = si[p == 99900],
    sj_999 = sj[p == 99900],
    st_999 = st[p == 99900],
  ) %>%
  ungroup() %>%
  mutate( 
    si = if_else(p == 99000 & !is.na(si_999), si_999*ratio99_999, si),
    sj = if_else(p == 99000 & !is.na(sj_999), sj_999*ratio99_999, sj),
    st = if_else(p == 99000 & !is.na(st_999), st_999*ratio99_999, st)
  ) %>%
  select(-si_999, -sj_999, -st_999)

# generate a new fiscal_data_2
fiscal_data_2_99<- data_ratio99 %>%
  filter(p==99000) %>%
  filter(!(is.na(si) & is.na(sj) & is.na(st))) %>%
  select(iso, year, p, si, sj, st)
  
# Append it to fiscal_data
fiscal_data <- bind_rows(fiscal_data, fiscal_data_2_99) 
fiscal_data <- fiscal_data %>% arrange(iso, year, p)

rm(data_org99, data_ratio99, fiscal_data_2, fiscal_data_2_99, breaks99)

#--------- 3. Collapse data into g-pinter input format ------------------------- 
# Create fractile 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) %>%
  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, 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
  )   

# Verify consistency
data_org <- data_org %>%
  group_by(iso, year) %>%
  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-6))
stopifnot(all(abs(data_org$avg - 1) < 1e-6))

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


# Step 1. Make sure data is sorted
data_org <- data_org %>%
  mutate(
    benchmark=1
  )%>%
  arrange(iso, p, year)

#-----------------------------------------------------------------------------------------------------------------------------------
#-----------------------------------------------------------------------------------------------------------------------------------
data_org_backup<-data_org  %>%
  select(iso, year, p, s) %>%
  rename(s_real=s)

data_org3 <- data_org %>%
  filter(year %in% c(1820, 1850, 1880, 1900, 1910, 1920, 1930, 1940, 1950, 1960, 1970, 1980))
#-----------------------------------------------------------------------------------------------------------------------------------
#------------------------------------------------------------------------------------------------------------------------------------

# Generate full year observations and Interpolate missing years between 1870 and 1980
data <- data_org %>%
  group_by(iso, p) %>%
  complete(year = full_seq(1850:1980, 1)) %>%
  arrange(iso, p, year) %>%
  mutate(
    t = if (sum(!is.na(t)) > 1) {
      approx(year[!is.na(t)], t[!is.na(t)], xout = year, rule = 2)$y
    } else {
      # If only one observation, just fill all with that single value (or NA if none)
      rep(t[!is.na(t)][1], length(year))
    },
    s = if (sum(!is.na(s)) > 1) {
      approx(year[!is.na(s)], s[!is.na(s)], xout = year, rule = 2)$y
    } else {
      rep(s[!is.na(s)][1], length(year))
    },
    a = if (sum(!is.na(a)) > 1) {
      approx(year[!is.na(a)], a[!is.na(a)], xout = year, rule = 2)$y
    } else {
      rep(a[!is.na(a)][1], length(year))
    }
  ) %>%
  ungroup() %>%
  mutate(
    pop = case_when(
      p == 0     ~ 0.1,
      p == 10000 ~ 0.15,
      p == 25000 ~ 0.25,
      p == 50000 ~ 0.25,
      p == 75000 ~ 0.15,
      p == 90000 ~ 0.09,
      p == 99000 ~ 0.01,
      TRUE       ~ pop
    )
  ) %>%
  arrange(iso, year, p)

#--------- 4. Calculate the correction ratios --------------------------------
# Keep only countries with at least one coyered fiinc to be extended
data_ratio <- data %>%
  left_join(fiscal_data, by = c("iso", "year","p")) %>%
  group_by(iso) %>%
  filter(any(!is.na(si) | !is.na(sj) | !is.na(st) )) %>%
  ungroup() 

# Interoolate the si and sj
library(dplyr)
library(zoo)


data_ratio <- data_ratio %>%
  filter(p == 99000) %>%
  group_by(iso) %>%
  arrange(year) %>%
  mutate(
    si_interp = na.approx(si, x = year, na.rm = FALSE),
    sj_interp = na.approx(sj, x = year, na.rm = FALSE),
    st_interp = na.approx(st, x = year, na.rm = FALSE)
  ) %>%
  ungroup() 

data_ratio <- data_ratio %>%
  filter(p == 99000) %>%
  group_by(iso) %>%
  arrange(year) %>%
  mutate(
    si_filled = na.locf(na.locf(si_interp, na.rm = FALSE), fromLast = TRUE, na.rm = FALSE),
    sj_filled = na.locf(na.locf(sj_interp, na.rm = FALSE), fromLast = TRUE, na.rm = FALSE),
    st_filled = na.locf(na.locf(st_interp, na.rm = FALSE), fromLast = TRUE, na.rm = FALSE)
  ) %>%
  ungroup() %>%
  mutate(
    si_interp = if_else(is.na(si_interp), si_filled, si_interp),
    sj_interp = if_else(is.na(sj_interp), sj_filled, sj_interp),
    st_interp = if_else(is.na(st_interp), st_filled, st_interp)
  )


data_ratio <- data_ratio%>%
  mutate(
    si = if_else(is.na(si), si_interp, si),
    sj = if_else(is.na(sj), sj_interp, sj),
    st = if_else(is.na(st), st_interp, st)
  ) %>%
  select(-si_interp,-sj_interp,-st_interp) %>%
  arrange(iso, year, p)

#Generate a ratio of the top1
data_ratio <-data_ratio %>%
  arrange(iso, year, p) %>%
  mutate(
    ratio99_0  = if_else(benchmark==1 & !is.na(si) & p==99000, s/si, NA),
    ratio99_0  = if_else(benchmark==1 & !is.na(sj) & p==99000, s/sj, ratio99_0),
    ratio99_0  = if_else(benchmark==1 & !is.na(st) & p==99000, s/st, ratio99_0)
  ) 

# Interpolate ratios if possible
data_ratio <- data_ratio %>%
  group_by(iso) %>%
  arrange(year, .by_group = TRUE) %>%
  mutate(ratio99 = na.approx(ratio99_0, x = year, na.rm = FALSE)) %>%
  ungroup()

# Extrapolate ratios if necessary
data_ratio <- data_ratio %>%
  arrange(iso, year) %>%
  group_by(iso,p) %>%
    fill(ratio99, .direction = "up") %>%  # fill backward from next available value
  ungroup() %>%
  group_by(iso,p) %>%
    fill(ratio99, .direction = "down") %>%  # fill forward in time
  ungroup() %>%
  select(iso, year, p, si, sj, st, ratio99)

# Keep only available years
data_ratio <- data_ratio %>%
  filter(!(is.na(sj) & is.na(si) & is.na(st)))

data_ratio <- data_ratio %>%
  semi_join(fiscal_data, by = c("iso", "year"))

#--------- 5. Transform the series --------------------------------------------
# Estimate proportions between percentiles
data_annex_s <- data %>%
  pivot_wider(
    id_cols = c(iso, year),  # keep these as identifiers
    names_from = p,                    # each value of p becomes a new column
    values_from = s,                   # the corresponding s values
    names_prefix ="s_"              # optional prefix for clarity 
    ) %>%
  mutate(
    r_s_90_99= s_90000/s_99000,
    r_s_75_90= s_75000/s_90000
  ) 

data_annex_t <- data %>%
  pivot_wider(
    id_cols = c(iso, year),  # keep these as identifiers
    names_from = p,                    # each value of p becomes a new column
    values_from = t,                   # the corresponding s values
    names_prefix ="t_"              # optional prefix for clarity 
  )  %>%
  mutate(
    r_t_90_99 = t_90000/t_99000,
    r_t_75_90 = t_75000/t_90000
  ) 

data <- data %>%
  group_by(iso) %>%
  mutate(
    corr= cor(t, s, method = "pearson") 
  )%>%
  ungroup() 


# -------------- 5.1 Find an optimal penalization lambda
# Function for fining the optimal lambda
apply_correction <- function(data, lambda) {
  data %>%
    left_join(data_org_backup, by = c("iso", "year", "p")) %>%
    left_join(data_ratio, by = c("iso", "year", "p")) %>%
    mutate(
      s_org = s,
      s_trans = if_else(!is.na(ratio99) & !is.na(si) & p == 99000, si * ratio99, s),
      s_trans = if_else(!is.na(ratio99) & !is.na(sj) & p == 99000, sj * ratio99, s_trans),
      s_trans = if_else(!is.na(ratio99) & !is.na(st) & p == 99000, st * ratio99, s_trans),
      diff_s = s_trans - s_org
    ) %>%
    group_by(iso) %>%
    mutate(
      sd_diff = sd(diff_s, na.rm = TRUE),
      diff_s_penalized = diff_s / (1 + abs(diff_s / (lambda * sd_diff)))
    ) %>%
    ungroup() %>%
    mutate(s_corr = s + diff_s_penalized)
}

# Define a loss function
loss_function <- function(lambda, data) {
  data_corr <- apply_correction(data, lambda)
  
  # Compare only where s_real exists
  data_compare <- data_corr %>%
    filter(!is.na(s_real)) %>%
    mutate(error = (s_corr - s_real)^2)
  
  mean(data_compare$error, na.rm = TRUE)  # Mean Squared Error
}

# Grid search over candidate lambdas
lambda_candidates <- seq(0.1, 50, by = 0.5) # Not: This was run until 200 and the result was still 17.1, so it is not necesary

lambda_results <- tibble(
  lambda = lambda_candidates,
  mse = map_dbl(lambda_candidates, ~ loss_function(.x, data))
)

# Best lambda
best_lambda <- lambda_results %>% filter(mse == min(mse)) %>% pull(lambda)


# -------------- 5.2 Generate Pretax series series based on the fiscal data
data <- data %>%
  left_join(data_ratio, by = c("iso", "year","p")) %>%
  mutate(
    s_org=s,
    s_trans = if_else(!is.na(ratio99) & !is.na(si) & p==99000, si * ratio99, s),
    s_trans = if_else(!is.na(ratio99) & !is.na(sj) & p==99000, sj * ratio99, s_trans),
    s_trans = if_else(!is.na(ratio99) & !is.na(st) & p==99000, st * ratio99, s_trans)
  ) 
  
# -------------- 5.3 Use the best_lambda for penalizing trasnformation
data <- data %>%
  mutate(
    diff_s=s_trans-s_org
  ) %>%
  group_by(iso) %>%
  mutate(
    sd_diff = sd(diff_s, na.rm = TRUE),
    diff_s_penalized = diff_s / (1 + abs(diff_s / (best_lambda* sd_diff)))
  ) %>%
  ungroup() %>%
  mutate(
    s = if_else(!is.na(diff_s_penalized), s + diff_s_penalized, s)
  )

data <- data %>%
  left_join(data_annex_s, by = c("iso", "year")) %>%
  left_join(data_annex_t, by = c("iso", "year")) %>%
  group_by(iso) %>%
  filter(any(!is.na(ratio99))) %>%
  ungroup() 

# Correct the top threshold
data <- data %>%
  mutate(
    t_org = t,
    ratio_s99 = if_else(            p==99000, s/s_99000,   NA),
    ratio_t99 = if_else(            p==99000, t/s_org,     NA),
    t = if_else(!is.na(ratio_t99) & p==99000, s*ratio_t99, t)
  )

# re-weight the rest of the shares and correct under the same value the  thresholds
data <- data %>%
  group_by(iso, year) %>%      # <-- adjust grouping as needed
  mutate(
    rest_curr =    sum(s[p <= 99000], na.rm = TRUE),
    rest_expe = 1- sum(s[p >  99000], na.rm = TRUE),
    adj_ratio = rest_expe/rest_curr
  ) %>%
  ungroup() %>%
  mutate(
    s = if_else(p <= 99000,  s*adj_ratio, s),
    t = if_else(p <= 99000,  t*adj_ratio, t)
  )

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

# Verify consistency
data <- data %>%
  group_by(iso, year) %>%
  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-6))
stopifnot(all(abs(data$avg - 1) < 1e-6))


#  Cleanup
#data <- data %>%
#  filter(is.na(benchmark))

data <- data %>%
  group_by(iso, year) %>%
  filter(any(!is.na(ratio99))) %>%
  ungroup()

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


#------------------------------------------------------------------------------
#------------------------------------------------------------------------------
#            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("output5_pretax_nonbenchmark_gpinterized.dta"))
