import pandas as pd

primary_data = 'data/data_dle_psč.csv'
psc_obce_data = 'data/obce_psč_předpočítané.csv'
hierarchy_data = 'data/vazby-cr.csv'


def preprocess_primary():
    df = pd.read_csv(primary_data, sep=';', encoding='utf-8')
    df = df.loc[df["Complete?"] == "Complete"]
    df = df.drop(columns=['Complete?'])
    df = df.replace({"muž": "M", "žena": "F"})

    return df

def preprocess_psc_obce():
    df = pd.read_csv(psc_obce_data, sep=',', encoding='utf-8')

    return df

def preprocess_hierarchy():
    df = pd.read_csv(hierarchy_data, sep=';', encoding='cp1250')
    df = df.drop(columns=['COBCE_KOD'])
    df = df.drop_duplicates()

    df['OKRES_KOD'] = df['OKRES_KOD'].fillna(9999).astype(int) # Praha

    # validate uniqueness of OBEC_KOD
    duplicate_obec_kod = df[df.duplicated(subset='OBEC_KOD', keep=False)]
    assert duplicate_obec_kod.empty, "Duplicate OBEC_KOD found"

    # validate correctly asssigning Praha to OKRES_KOD
    assert set(df[df['OKRES_KOD'] == 9999]['OBEC_KOD']) == {554782}, "Invalid OKRES_KOD for OBEC_KOD"

    return df


if __name__ == '__main__':
    df_primary = preprocess_primary()
    df_obce = preprocess_psc_obce()
    df_hierarchy = preprocess_hierarchy()

    df_merged = df_primary.merge(df_obce, left_on='PSČ', right_on='PSČ', how='inner')
    df_merged = df_merged.merge(df_hierarchy, left_on='Kód obce', right_on='OBEC_KOD', how='inner')
    print(df_merged)

    df_grouped = df_merged.groupby(['OKRES_KOD']).size().reset_index(name='Count')

    print(df_grouped)