if (!require("pacman")) install.packages("pacman")


pacman::p_load(data.table)
pacman::p_load(here)
pacman::p_load(matrixStats)


destinations_selection_and_sampling <- function
(
base_filt, ## dataset with only firm and country dimension, zeros when missing transaction
filt_par, ## parameters that define the criteria for country selection and sample construction
sampling = T ## If TRUE the function selects sampling, if FALSE the function operates only the selection of countries and provide the complete database
)
{
  
  
  
## Define a function converting the NAs to zeros
f_zero = function(DT) 
{
  for (i in names(DT))
    DT[is.na(get(i)), (i):=0]
}
## Define a function converting the zeros to NAs
f_nan = function(DT) 
{
  for (i in names(DT))
    DT[get(i)==0, (i):=NA]
}

### Store the columns of the dataset defining the points of the time series
export_cols <- colnames(base_filt)[3:ncol(base_filt)]
f_zero(base_filt)

## Force numeric format for columns containing export values
base_filt[,(export_cols) := lapply(.SD, function(x) as.numeric(x)), .SDcols = export_cols]


################ Defining and applying the filters on the country dimension ############################

cat('\nStarting the filtering procedure with a flows number threshold of',filt_par$flow_filter, 'and an export value share threshold of', filt_par$share_of_exp_value_filter,'\n')


## Define the list of exclusion (to be updated while applying the filter)
countries_excl_by_active_flows <- unique(base_filt$PYOD)

country_N <- base_filt[,
                       lapply(.SD, function(x)sum(x>0, na.rm = T)),   ## compute the sum
                       keyby = c('PYOD'),                  ## for every 'year,month, etc..'
                       .SDcols = export_cols]

avg_flows <- cbind(country_N[,1],rowMeans(country_N[,2:ncol(country_N)]))


## Filter 1: compile the exclusion list based on the number of trade flows targeting a given country/destination
for(i in export_cols)
{
  cs_sel <- country_N[,c('PYOD',i),with = F]  
  colnames(cs_sel) <- c('pyod','value')
  
  cs <- cs_sel[value < filt_par$flow_filter]$pyod
  
  countries_excl_by_active_flows <- intersect(countries_excl_by_active_flows,cs)
}  

## Define the list of exclusion (to be updated while applying the filter)
countries_excl_by_export_value_share <- unique(base_filt$PYOD)

country_aggregates <- base_filt[,
                                  lapply(.SD,function(x) as.double(x)),   ## convert to double
                                  keyby = c('PYOD'),                      ## for every destination
                                  .SDcols = export_cols
                                  ][,                                            ## All the transactions
                                    lapply(.SD,function(x) sum(x, na.rm = T)),   ## compute the sum
                                    keyby = c('PYOD'),                           ## for every 'year,month, etc..'
                                    .SDcols = export_cols]

country_value_shares <- country_aggregates[,c(list(PYOD = PYOD),lapply(.SD,function(x) x/sum(x))),.SDcols = export_cols]



## Filter 2: Compile the selection list based on the share the countries in the total export
avg_shares <- as.data.frame(cbind(PYOD = country_value_shares[,PYOD],
                                  SHARE = round(rowMeans(country_value_shares[,2:ncol(country_value_shares)]),4)))
setDT(avg_shares)

for(col in export_cols)
{  
  country_excl <- country_value_shares$PYOD[country_value_shares[,col,with = F] < filt_par$share_of_exp_value_filter]
  countries_excl_by_export_value_share <- intersect(countries_excl_by_export_value_share,country_excl)
}

# define the excluded countries as the intersection of both exclusion lists
excluded <- c(intersect(countries_excl_by_active_flows,countries_excl_by_export_value_share),
                                 'XH','XI','XZ','XW','XP','TL','WS','XA','XF','XG','XR')




########## Filter out the excluded countries from the original dataset #################

base_exp <- base_filt[!(PYOD%in%excluded)][order(PYOD,SIREN)]


before_filtering <- base_filt[,lapply(.SD,function(x) sum(x, na.rm = T)),
                              .SDcols = export_cols]
  
after_filtering <- base_exp[, lapply(.SD, function(x)sum(x, na.rm = T)),              ## compute the sum
                          .SDcols = export_cols]

share_preserved <- sum(after_filtering)/sum(before_filtering)

cat('\nThe country filtering procedure excludes the',round((1-share_preserved)*100,4),'% of the total export value, thus preserving the',
    round(share_preserved*100,4),'%\n')



#################### Filtering out flows without any point in growth rates ######################

cat('\nRemoving the time series with less than three points in growth rates...\n')

f_nan(base_exp) ## When calculating the growth rates counting an NA when there are no traade transactions

mat_base <- as.matrix(base_exp[,3:ncol(base_exp)])
gr_mat <- log(1+mat_base[,5:ncol(mat_base)]) - log(1+mat_base[,1:(ncol(mat_base)-4)])
gr_base <- setDT(as.data.frame(gr_mat))
base_exp <- base_exp[rowSums(!is.na(gr_base)&!gr_base==0, na.rm = T)>2] ### excluding series with more than 2 NAs
gr_base <- gr_base[rowSums(!is.na(gr_base)&!gr_base==0, na.rm = T)>2]
base_exp <- base_exp[rowVars(as.matrix(gr_base), na.rm = T)!=0] ### Excluding series with null variance, degenerate series (around 1 in 60000)



after_the_filter_on_growth_rates <- base_exp[, lapply(.SD, function(x)sum(x, na.rm = T)),          ## compute the sum
                                              .SDcols = export_cols]

share_preserved_final <- sum(after_the_filter_on_growth_rates)/sum(before_filtering)


cat('\nOn top of the filter on the countries, excluding the times series with less than 3 points in growth rates removes',round((1-share_preserved_final)*100,4),
    '% of the total export value, thus preserving the',
    round(share_preserved_final*100,4),'% \n')


countries <- unique(base_exp$PYOD)

save(countries, file = here('data_raw','countries.RData'))
firms <- unique(base_exp$SIREN)
vec_firms <- unique(base_exp$SIREN)

cat('\nCountries/destinations = ',length(countries))
cat('\nFirms = ',length(vec_firms))
cat('\nExport flows/time series = ',nrow(base_exp),'\n')


################## Sampling the firms ###########################

sel_firms <- list()
for(i in 1:filt_par$num_samples)
{
  set.seed(i)  
  sel_firms[[i]] <- sample(vec_firms,filt_par$num_firms,replace = F)  
  
  base_sel <- base_exp[SIREN %in% sel_firms[[i]]]
  countries <- intersect(countries,unique(base_sel$PYOD))
}


## Construct the folder name using the parameter of the selections: number of destinations and number of firms in each subsample
folder <- paste0('results_samples_',length(countries),'_',filt_par$num_firms/1000,'k')
dir.create(here::here('estimation_scripts','Inputs',folder), showWarnings = FALSE)
dir.create(here::here('estimation_scripts','Inputs',folder,'samples'), showWarnings = FALSE)

dir.create(here::here('results_analysis','Inputs',folder), showWarnings = FALSE)
dir.create(here::here('results_analysis','Inputs',folder,'samples'), showWarnings = FALSE)



#### Saving the filtered complete dataset and the weights (both aggregate level and firm level) 
### associated to this
f_nan(base_exp)
mat_base <- as.matrix(base_exp[PYOD %in% countries,3:ncol(base_exp)])
gr_mat <- log(1+mat_base[,5:ncol(mat_base)]) - log(1+mat_base[,1:(ncol(mat_base)-4)])
gr_base <- setDT(as.data.frame(gr_mat))
base_exp_gr <- cbind(base_exp[PYOD %in% countries,1:2],gr_base)

weight_full <- copy(base_exp)
weight_full <-  weight_full[PYOD %in% countries, (export_cols) := lapply(.SD, function(x) x/sum(x, na.rm = T)), .SDcols = export_cols]

weight_siren <-  base_exp[PYOD %in% countries][,
                                               paste0(export_cols,'w') := lapply(.SD, function(x) x/sum(x, na.rm = T)),
                                               by = c('SIREN'), .SDcols = export_cols][,c('SIREN','PYOD',paste0(export_cols,'w')), with = F]


f_nan(weight_full)
f_nan(weight_siren)

### Save the full dataset of growth rates for possible estimation of the factor model decomposition
write.csv(base_exp_gr, here::here('estimation_scripts','Inputs',folder,'samples',paste0("data_sample_gr0.csv")))

### Save the full dataset of growth rates and the related weights for results analysis
write.csv(base_exp_gr,here::here('results_analysis','Inputs',folder,'samples',paste0("data_sample_gr0.csv")))
write.csv(weight_full,here::here('results_analysis','Inputs',folder,'samples',paste0("data_sample_w_sel0.csv")))
write.csv(weight_siren, here::here('results_analysis','Inputs',folder,'samples',paste0("data_sample_w_firm0.csv")))




################### Subsampling the dataset #########################

n_flows <- 0
cuts_per <- 0



### Before sampling: construct the aggregate growth rates coming from the filtered dataset and save it
 # agg_vec <- colSums(base_exp[,-c(1,2)], na.rm = T)
 # agg_gr <- log(agg_vec[5:length(agg_vec)])-log(agg_vec[1:(length(agg_vec)-4)])
 # agg_gr_df <- as.data.frame(cbind(Quarters = names(agg_vec[5:length(agg_vec)]), agg_gr))
 # 
 # write.csv(agg_gr_df,here::here(folder,paste0("agg_gr.csv")), row.names = F)


if(sampling)
{
  
n_flows <- NULL
cuts_per <- NULL

for(i in 1:filt_par$num_samples)
{
  base_sel <- base_exp[PYOD %in% countries & SIREN %in% sel_firms[[i]]]
  c <- base_sel[, lapply(.SD, function(x)sum(x, na.rm = T)),              ## compute the sum
                .SDcols = export_cols]
  
  ### Calculation of the growth rate matrix
  f_nan(base_sel)
  mat_base <- as.matrix(base_sel[,3:ncol(base_sel)])
  gr_mat <- log(1+mat_base[,5:ncol(mat_base)]) - log(1+mat_base[,1:(ncol(mat_base)-4)])
  base_sel_gr <- setDT(as.data.frame(gr_mat))
  base_sel_gr <- cbind(base_sel[,1:2],gr_mat)
  
  cuts_per <- c(cuts_per,sum(c))
  n_flows <- c(n_flows,nrow(base_sel))
  
  save_cols <- colnames(base_sel)[-c(1,2)]
  weight_sel <- copy(base_sel)
  weight_sel[ , (save_cols) := lapply(.SD, function(x) x/sum(x, na.rm = T)), .SDcols = save_cols]

  weight_siren <-  base_sel[ ,
                             paste0(export_cols,'w') :=lapply(.SD, function(x) x/sum(x, na.rm = T)),
                             by = c('SIREN'), .SDcols = export_cols][,c('PYOD','SIREN',paste0(export_cols,'w')), with = F]
  
  
  f_nan(base_sel)
  f_nan(weight_sel)
  f_nan(weight_siren)
  
  print(length(unique(base_sel$PYOD)))
  print(nrow(base_sel))
  print(nrow(weight_sel))

  ### Save the samples of growth rates in the Inputs folder for the factor model decomposition
  write.csv(base_sel_gr,here::here('estimation_scripts','Inputs',folder,'samples',paste0("data_sample_gr",i,".csv")))
 
  ## Save the sample of growth rate and all the related weights in the results analysis folder
  write.csv(base_sel_gr,here::here('results_analysis','Inputs', folder,'samples',paste0("data_sample_gr",i,".csv")))
  write.csv(weight_sel,here::here('results_analysis', 'Inputs',folder,'samples',paste0("data_sample_w_sel",i,".csv")))
  write.csv(weight_siren,here::here('results_analysis', 'Inputs',folder,'samples',paste0("data_sample_w_firm",i,".csv")))
}
}



fileConn <- file(here::here('estimation_scripts','Inputs',folder,"filtering_sampling_info.txt"))
writeLines(c("Selection parameters:"
             ,do.call(paste0,Map(paste0,names(filt_par),' = ',filt_par,'; '))   
             ,"Number of countries/destinations:"
             ,length(countries)
             ,"Number of firms after selection"
             ,length(unique(base_exp$SIREN))
             ,"Number of firms per sample"
             ,filt_par$num_firms
             ,"Percentage of total value preserved after destination exclusion:"
             , share_preserved
             ,"Percentage of total value preserved after excluding the series with few points in growth rates:"
             , share_preserved_final
             ,"Avg. number of flows taken in samples"
             , mean(n_flows)
             ,"Avg. percentage value of the total export taken in the samples" 
             , mean(cuts_per)/sum(before_filtering)
             ,"Avg. percentage value of total export of the filtered dataset taken in the samples"
             , mean(cuts_per)/sum(after_the_filter_on_growth_rates)), fileConn)
close(fileConn)

return(base_exp[PYOD %in% countries])
}
