# Clean the data on cultural distance from the World Values Survey.
# Transform the square matrix of distances to a long-format dataset.
# This is being used to construct the input data for Table S.3 and S.4

import os
import pandas as pd
import pycountry

PROJ_DIR = os.getenv("PROJ_DIR")

categories = ["", "beliefs", "group_memberships", "law", "misc", "political", "sexuality", "social_relations"]
data_list = []
renaming = {"fst_cultural_group_memberships": "fst_cultural_groups", "fst_cultural_social_relations": "fst_cultural_social"}

for suffix in categories:
    if suffix != "":
        suffix = "_" + suffix
    df = pd.read_csv(os.path.join(PROJ_DIR, "data", "cultural_distance", "matrix_all" + suffix + ".csv"))
    df.columns = [x.replace("1994-2014", "").replace("Rep.", "Republic").replace("Serbia and Montenegro", "SCG") for x in df.columns ]
    df.index = df.columns
    countries = df.columns

    exceptions = ["SCG", "Turkey"]
    def exceptions_handler(x):
        if x == "SCG":
            return "SCG"
        if x == "Turkey":
            return "TUR"

    iso3 = [pycountry.countries.search_fuzzy(x)[0].alpha_3  if not x in exceptions else exceptions_handler(x) for x in df.columns]
    df.index = iso3
    df.columns = iso3
    
    mapping = pd.DataFrame(zip(countries, iso3), columns=["cname", "iso3"])
    
    # Reshape data to long and remove the diagonal entries.
    df_long = df.stack().reset_index().drop_duplicates(keep="first")
    df_long.columns = ['iso3_1', 'iso3_2', 'fst_cultural'+suffix]
    df_long = df_long[df_long['iso3_1'] != df_long['iso3_2']]
    
    # Add country names
    df_long = df_long.merge(mapping, how="left", left_on=["iso3_1"], right_on=["iso3"]).rename(columns={"cname": "cname_1"}).drop("iso3", axis=1)
    df_long = df_long.merge(mapping, how="left", left_on=["iso3_2"], right_on=["iso3"]).rename(columns={"cname": "cname_2"}).drop("iso3", axis=1)

    # Add identifier for similar pairs
    df_long["cpair_order"] = (df_long["iso3_1"] < df_long["iso3_2"]).astype(int)
    data_list.append(df_long)


df = pd.concat([x.set_index(["iso3_1", "iso3_2", "cname_1", "cname_2", "cpair_order"], verify_integrity=True) for x in data_list], axis=1)
df = df.reset_index()
df = df.rename(columns=renaming)

# Quality checks
assert df.duplicated().sum() == 0
assert df.shape[0] < 10000

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