# This file produces the final dataset for reproducing Tables S3 and S4.

import os
import pandas as pd
import numpy as np
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from warnings import simplefilter
simplefilter(action="ignore", category=pd.errors.PerformanceWarning)

# Setup paths
PROJ_DIR = os.getenv("PROJ_DIR")

####################################
###     LOAD WORLD DATA          ###
####################################

world = pd.read_stata(os.path.join(PROJ_DIR, "out", "data_out", "world.dta"))
gcd_countries = world["iso3"].unique().tolist()

########################################################################
### ADD NEW VARIABLES BEFORE CREATING DISTANCE MATRIX, e.g. INDICES! ###
########################################################################

# Add treatment effect estimate on cooperation and beliefs (by country)
gcd = pd.read_stata(os.path.join(PROJ_DIR, "out", "data_out", "gcd.dta"), columns=["iso3", "treatment", "wgt_expfull", "coop", "coop_belief"], convert_categoricals=False)
gcd["coop_wgt"] = gcd["coop"] * gcd["wgt_expfull"]
gcd["coop_belief_wgt"] = gcd["coop_belief"] * gcd["wgt_expfull"]
means = gcd.groupby(["iso3", "treatment"], as_index=False)[["coop_wgt", "coop_belief_wgt"]].mean()
means = means.pivot_table(index='iso3', columns='treatment', values=['coop_wgt', "coop_belief_wgt"], aggfunc='first').reset_index()
means["te_effect_coop"] = means["coop_wgt"][1.0] - means["coop_wgt"][0.0]
means["te_effect_belief"] = means["coop_belief_wgt"][1.0] - means["coop_belief_wgt"][0.0]
world = world.merge(means.droplevel(1,axis=1)[["iso3", "te_effect_coop", "te_effect_belief"]], how="left", on="iso3")


##########################################
### DISTANCE MATRIX FOR GCD VARIABLES ###
##########################################

# Cross-join
df1 = world.assign(key=1)
df2 = world.assign(key=1)
df_cross = pd.merge(df1, df2, on='key').drop('key', axis=1)

# Filter out pairs where country is paired with itself
df_cross = df_cross[df_cross['iso3_x'] != df_cross['iso3_y']]

diff_frames = []
for col in world.select_dtypes(include=['int64', 'float64', 'float', 'float32', 'int', 'int32']).columns:
    if col != 'iso3':
        diff_frames.append(pd.DataFrame({f'd_{col}': abs(df_cross[f'{col}_x'] - df_cross[f'{col}_y'])}))

# Concatenate all difference dataframes along with the original df_cross
df_cross = pd.concat([df_cross] + diff_frames, axis=1)
df_cross = df_cross.rename(columns={col: col.rstrip('_x') + '_1' for col in df_cross.columns if col.endswith('_x')})
df_cross = df_cross.rename(columns={col: col.rstrip('_y') + '_2' for col in df_cross.columns if col.endswith('_x')})

# Rename columns for clarity
df_cross = df_cross.rename(columns={'iso3_x': 'iso3_1', 'iso3_y': 'iso3_2'})

# Keep only relevant columns
df_result = df_cross[['iso3_1', 'iso3_2'] + [col for col in df_cross.columns if col.startswith('d_')]]

# Drop columns with too many missings
thresh = len(df_result) * 0.8
df_result = df_result.dropna(axis=1, thresh=thresh)

# Generate country fixed effects
dummies = []
for iso3 in gcd_countries:
    dummies.append(pd.DataFrame({f"dummy_{iso3}": ((df_result['iso3_1'] == iso3) | (df_result['iso3_2'] == iso3)).astype(int)}))

# Combine data and defragmentate the dataframe with a copy.
df_result = pd.concat([df_result] + dummies, axis=1)
df_final = df_result.copy()

# Get sample subset where only one member of the symmetric country pairs is used
df_final["country_pair_singleton"] = (df_final["iso3_1"] < df_final["iso3_2"]).astype(int)

# Merge the treatment effects variables
df_final = df_final.merge(
    world[["iso3", "te_effect_coop", "te_effect_belief"]],
    how="left",
    left_on="iso3_1",
    right_on="iso3"
).rename(columns={
    "te_effect_coop": "te_effect_coop_1",
    "te_effect_belief": "te_effect_belief_1"
}).drop(columns=["iso3"])
df_final = df_final.merge(
    world[["iso3", "te_effect_coop", "te_effect_belief"]],
    how="left",
    left_on="iso3_2",
    right_on="iso3"
).rename(columns={
    "te_effect_coop": "te_effect_coop_2",
    "te_effect_belief": "te_effect_belief_2"
}).drop(columns=["iso3"])

###################################
####### ADD OTHER DATA SETS #######
###################################

# ADD CULTURAL DISTANCE
culture_dist = pd.read_csv(os.path.join(PROJ_DIR, "out", "data_out", "temp", "cultural_distance.csv"))

# ADD LINGUISTIC AND GENETIC DISTANCE DATA FOR THE ANCESTRAL DISTANCE MEASURE (SEE BELOW)
genes = pd.read_csv(os.path.join(PROJ_DIR, "out", "data_out", "temp", "ancestral_distance.csv"))
genes_swapped = genes.copy()
genes_swapped['iso3_1'], genes_swapped['iso3_2'] = genes_swapped['iso3_2'], genes_swapped['iso3_1']

genes_dist = pd.concat([genes, genes_swapped], ignore_index=True)
genes_dist = genes_dist.drop_duplicates(subset=['iso3_1', 'iso3_2'])

# ADD ETHNOGRAPHIC ATLAS DISTANCE
ea = pd.read_stata(os.path.join(PROJ_DIR, "out", "data_out", "temp", "ethno_atlas_distance.dta"), convert_categoricals=False)

final = df_final.merge(culture_dist, how="left", on=["iso3_1", "iso3_2"]).merge(
    genes_dist, how="left", on=["iso3_1", "iso3_2"]).merge(
            ea, how="left", left_on=["iso3_1", "iso3_2"], right_on=["iso3_1", "iso3_2"]
        )

###################################
####### POST-MERGE CLEANING #######
###################################

# Keep only a single member of each country pair, i.e., the lower triangle
final = final[final["country_pair_singleton"] == 1]

# Constant indicator for some regressions.
final["all"] = 1

# Standardize cultural distance variables
for var in final.columns:
    if var.startswith("fst_cultural"):
        mean = final[var].mean(skipna=True)
        sd = final[var].std(skipna=True)
        final[var+"_std"] = (final[var] - mean) / sd

# Standardize distances
mean = final["d_ea"].mean(skipna=True)
sd = final["d_ea"].std(skipna=True)
final["d_ea_std"] = (final["d_ea"] - mean) / sd

# Ancestral distance: Create index variable
final["lingdist_dominant_2009"] = 1.0 - (final["lingprox_dominant_2009"] ** 0.5 / 15.0)
final["lingdist_weighted_2009"] = 1.0 - (final["lingprox_weighted_2009"] ** 0.5 / 15.0)

for var in ["fst_distance_dominant_2009", "fst_distance_weighted_2009", "lingdist_dominant_2009", "lingdist_weighted_2009", "new_gendist_weighted_2017", "new_gendist_plurality_2017"]:
    mean = final[var].mean(skipna=True)
    sd = final[var].std(skipna=True)
    final[var+"_std"] = (final[var] - mean) / sd


final["ancestral_index"] = (final["fst_distance_weighted_2009_std"] + final["new_gendist_weighted_2017_std"] + final["lingdist_weighted_2009_std"]) / 3.0
final["ancestral_index"] = (final["ancestral_index"] - final["ancestral_index"].mean(skipna=True)) / final["ancestral_index"].std(skipna=True)

# Standardize GCDS variables
cols = ["d_gcd_altruism_std_wgtexpctrl", "d_gcd_patience_std_wgtexpctrl", "d_gcd_risktaking_std_wgtexpctrl", "d_gcd_coop_belief_wgtexpctrl", "d_gcd_norm_wgtexpctrl", "d_gcd_pg_preference_wgtexpctrl"]
for col in cols:
    mean = final[col].mean(skipna=True)
    sd = final[col].std(skipna=True)
    final[col+"_std"] = (final[col] - mean) / sd

###################################
#######        EXPORT       #######
###################################

final.to_stata(os.path.join(PROJ_DIR, "out", "data_out", "world_distance.dta"), write_index=False)
