# This file generates the Ethnographic Atlas distance variable for Tables S3 and S4.
# Note: Run this file only after the cleaning_R files have been run.
import os
import pandas as pd
import warnings
warnings.simplefilter(action='ignore', category=pd.errors.PerformanceWarning)

# 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()

### Load cleaned Ethnographic Atlas ###
path = os.path.join(PROJ_DIR, "data", "ethnographic_atlas_distance", "Ancestral_Characteristics_Database_Extended_Eastern_Europe_Siberia_and_WES_Version_1_1.dta")
df = pd.read_stata(path)
cols = ["isocode"] + [x for x in df.columns if x.startswith("v") and "_" in x]
df = df[cols]

# Get variable labels
with pd.read_stata(path, iterator=True) as iterator:
    varlabels = iterator.variable_labels()

# Remove columns with missing data and countries that are not part of the gcd data
df = df[df["isocode"].isin(gcd_countries)]
df = df.dropna(how="any", axis=1)

grp1_vars = [col for col in df.columns if col.endswith("_grp1")]
keywords = ["missing data", "no data", "insufficient information", "not coded"]
no_data_vars = [var for var in grp1_vars if any(keyword in varlabels[var].lower() for keyword in keywords)]

# Rescale by share of non-missing data
for col in df.columns:
    if not col in no_data_vars + ["isocode"]:
        matching_no_data_var = col.split("_")[0] + "_grp1"
        if matching_no_data_var in no_data_vars:
            df[col] = df[col] / (1.0 - df[matching_no_data_var])

# The above process creates missings because if all data is missing (e.g. _grp1 = 1), the other variables have value "0" rather than missing: So creating missings is not a bug

# Drop columns with any missings
columns_without_missing_values = df.columns[~df.isnull().any()]
df = df[columns_without_missing_values]

# Remove "no data" columns -- not needed anymore
df = df.drop(columns=no_data_vars)

# Create variance estimates
for col in df.columns:
    if col not in "isocode":
        df[col+"_std"] = df[col].std()
        df[col] = df[col] / df[col+"_std"]
        del df[col+"_std"]

df = df.set_index("isocode").transpose()

# Create distance matrix
dist_tmp = []
ea_countries = df.columns.tolist()
for c1 in ea_countries:
    for c2 in ea_countries:
        if c1 != c2:
            dist_tmp.append((c1, c2, ((df[c1] - df[c2]) ** 2).mean()))

dist = pd.DataFrame(dist_tmp)
dist.columns = ["iso3_1", "iso3_2", "d_ea"]

# Export data
dist.to_stata(os.path.join(PROJ_DIR, "out", "data_out", "temp", "ethno_atlas_distance.dta"), write_index=False)
dist.to_csv(os.path.join(PROJ_DIR, "out", "data_out", "temp", "ethno_atlas_distance.csv"), index=False)
