############################################################################################################
# Clean the ancestral distance data from Spolaore and Wacziarg (2009) and Spolaore and Wacziarg (2017).
# This is used to calculate the ancestral distance measure (genetic and linguistic difference).
############################################################################################################
import country_converter as coco
import pandas as pd
import os

PROJ_DIR = os.getenv("PROJ_DIR")

cc = coco.CountryConverter()

def classify(name):
    """Map country names to ISO-3 codes."""
    if "Kampuchea, Democratic" in name:
        return "KHM"
    elif "Cape Verde" in name:
        return "CPV"
    elif "Channel Islands" in name:
        return "GB-CHA"
    elif "Czechoslovakia" in name:
        return "CSK"
    elif "Faeroe Islands" in name:
        return "FRO"
    elif "Germany, Federal Republic of" in name:
        return "DEU"
    elif "Guinea Bissau" in name:
        return "GNB"
    elif "Korea,Dem.Rep." in name:
        return "PRK"
    elif "Netherlands Antilles" in name:
        return "ANT"
    elif "St Christopher and Nevis" in name:
        return "KNA"
    elif "St Lucia" in name:
        return "LCA"
    elif "St. Vincent" in name:
        return "VCT"
    elif "Swaziland" in name:
        return "SWZ"
    elif "U.S.A" in name:
        return "USA"
    elif "Virgin Islands (U.S.)" in name:
        return "VIR"
    elif "Western Samoa" in name:
        return "WSM"
    elif "Yemen, Arab Republic of" in name:
        return "YEM"
    elif "Yugoslavia" in name:
        return "YUG"
    elif "Myanmar(Burma)" in name:
        return "MMR"
    elif "German Democratic Republic" in name:
        return "DE"
    elif "Yemen, People's Democratic Republic of" in name:
        return "YEM"
    elif "Laos" in name:
        return "LAO"
    elif "U.S.S.R." in name:
        return "SUHH"
    else:
        return cc.convert(name, to='ISO3')

################################################
#### Spolaore and Wacziarg 2009. ###############
################################################

cols = [
    "wacziarg_1", "wacziarg_2", "country_1","country_2",
    "fst_distance_dominant", "nei_distance_dominant", "fst_distance_weighted", "nei_distance_weighted",
    "fst_stderr_dominant", "nei_stderr_dominant", "fst_stderr_weighted", "nei_stderr_weighted",
    "lingprox_dominant", "lingprox_weighted"
]

df = pd.read_stata(os.path.join(PROJ_DIR, "data", "ancestral_distance", "Spolaore and Wacziarg 2009", "bilateral.dta"), columns=cols)
iso3_1 = {x: classify(x) for x in df.country_1.unique().tolist()}
iso3_1 = pd.DataFrame.from_dict(iso3_1, orient="index", columns=["iso3_1"]).reset_index(names="country_1")

iso3_2 = {x: classify(x) for x in df.country_2.unique().tolist()}
iso3_2 = pd.DataFrame.from_dict(iso3_2, orient="index", columns=["iso3_2"]).reset_index(names="country_2")

data2009 = df.merge(iso3_1, how="left", on="country_1").merge(iso3_2, how="left", on="country_2")
data2009 = data2009.sort_values(by=["iso3_1", "iso3_2"])

# 2009 data: Assert that duplicate records have identical genetic and linguistic distance measures
vars2009 = [
    "fst_distance_dominant", "nei_distance_dominant", "fst_distance_weighted", "nei_distance_weighted",
    "fst_stderr_dominant", "nei_stderr_dominant", "fst_stderr_weighted", "nei_stderr_weighted",
    "lingprox_dominant", "lingprox_weighted"]

assert data2009[data2009.duplicated(subset=["iso3_1", "iso3_2"], keep=False)].groupby(["iso3_1", "iso3_2"])[vars2009].std().max().max() == 0

# Drop values with missing weighted linguistic distance data. Those are the ones with duplicate entries.
data2009 = data2009.dropna(subset=["fst_distance_weighted", "nei_distance_weighted"], how="all")

#### Spolaore and Wacziarg 2017.
df = pd.read_stata(os.path.join(PROJ_DIR, "data", "ancestral_distance", "Spolaore and Wacziarg 2017", "newgendist.dta"))

iso3_1 = {x: classify(x) for x in df.country_1.unique().tolist()}
iso3_1 = pd.DataFrame.from_dict(iso3_1, orient="index", columns=["iso3_1"]).reset_index(names="country_1")

iso3_2 = {x: classify(x) for x in df.country_2.unique().tolist()}
iso3_2 = pd.DataFrame.from_dict(iso3_2, orient="index", columns=["iso3_2"]).reset_index(names="country_2")

data2017 = df.merge(iso3_1, how="left", on="country_1").merge(iso3_2, how="left", on="country_2")
data2017 = data2017.sort_values(by=["iso3_1", "iso3_2"])

# 2017 data: Assert that there are no duplicate records once we drop rows with completely missing data
vars2017 = [
    'new_gendist_weighted', 'new_relgendist_weighted_usa',
    'new_relgendist_weighted_uk', 'new_gendist_plurality',
    'new_relgendist_plur_usa', 'new_relgendist_plur_uk', 'new_gendist_1500',
    'new_relgendist_1500_usa', 'new_relgendist_1500_uk',]

data2017 = data2017.dropna(subset=vars2017, how="all")
assert data2017[data2017.duplicated(subset=["iso3_1", "iso3_2"], keep=False)].shape[0] == 0

# Merge 2009 and 2017 data
data = data2009[["iso3_1", "iso3_2"] + vars2009].rename(columns={x: x + "_2009" for x in vars2009})
data = data.merge(data2017[["iso3_1", "iso3_2"] + vars2017].rename(columns={x: x + "_2017" for x in vars2017}), how="outer", on=["iso3_1", "iso3_2"])

assert data[data.duplicated(subset=["iso3_1", "iso3_2"], keep=False)].sum().sum() == 0

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