import pandas as pd

# Define our header
col_names = [
    "year",
    "num_males_with_income",
    "male_median_income_curr_dollars",
    "male_median_income_2019_dollars",
    "num_females_with_income",
    "female_median_income_curr_dollars",
    "female_median_income_2019_dollars",
]

# Load Asian census data XLS, skipping all headers
dfa = pd.read_excel(
    r'p08a.xlsx',
    skiprows=8,
    # Make sure PD doesn't use header row for our DF
    header=None,
    # Define col names
    names=col_names,
)
# Load White census data XLS, skipping all headers
dfw = pd.read_excel(
    r'p08w.xlsx',
    skiprows=8,
    # Make sure PD doesn't use header row for our DF
    header=None,
    # Define cold names
    names=col_names
)
# Splinter off rows into age group DFs for both sets of data
dfa1524 = dfa.iloc[:20]
dfa2534 = dfa.iloc[25:45]
dfa3544 = dfa.iloc[50:70]
dfa4554 = dfa.iloc[75:95]
dfa5564 = dfa.iloc[100:120]
dfa6574 = dfa.iloc[125:145]
dfa75 = dfa.iloc[150:170]

dfw1524 = dfw.iloc[:20]
dfw2534 = dfw.iloc[25:45]
dfw3544 = dfw.iloc[50:70]
dfw4554 = dfw.iloc[75:95]
dfw5564 = dfw.iloc[100:120]
dfw6574 = dfw.iloc[125:145]
dfw75 = dfw.iloc[150:170]

# Add Age Range col to each DF
dfa1524.insert(0, 'age_range', '15-24')
dfa2534.insert(0, 'age_range', '25-34')
dfa3544.insert(0, 'age_range', '35-44')
dfa4554.insert(0, 'age_range', '45-54')
dfa5564.insert(0, 'age_range', '55-64')
dfa6574.insert(0, 'age_range', '65-74')
dfa75.insert(0, 'age_range', 'Over 75')

dfw1524.insert(0, 'age_range', '15-24')
dfw2534.insert(0, 'age_range', '25-34')
dfw3544.insert(0, 'age_range', '35-44')
dfw4554.insert(0, 'age_range', '45-54')
dfw5564.insert(0, 'age_range', '55-64')
dfw6574.insert(0, 'age_range', '65-74')
dfw75.insert(0, 'age_range', 'Over 75')

# Stack cleaned DF's vertically
dfa = pd.concat([
    dfa1524,
    dfa2534,
    dfa3544,
    dfa4554,
    dfa5564,
    dfa6574,
    dfa75
], axis=0)

dfw = pd.concat([
    dfw1524,
    dfw2534,
    dfw3544,
    dfw4554,
    dfw5564,
    dfw6574,
    dfw75
], axis=0)

# Add Race col
dfa.insert(0, 'race', 'asian')
dfw.insert(0, 'race', 'white')

# Clean garbage chars in Year col using regex
dfa['year'] = dfa['year'].replace(to_replace=r'(\s\(\d+\))', value='', regex=True)
dfw['year'] = dfw['year'].replace(to_replace=r'(\s\(\d+\))', value='', regex=True)

# Stack our cleaned + normalized data into a single DF
df = pd.concat([
    dfa,
    dfw
], axis=0)

# Convert the DF col types to conform to our CensusRecord model
df = df.astype({
    "race": str,
    "age_range": str,
    "year": int,
    "num_males_with_income": int,
    "male_median_income_curr_dollars": float,
    "male_median_income_2019_dollars": float,
    "num_females_with_income": int,
    "female_median_income_curr_dollars": float,
    "female_median_income_2019_dollars": float,
})

# Pickle the DF
df.to_pickle("./res.pkl")

[((15, 6, 22, 1), 'pandas.read_excel', 'pd.read_excel', (), '', True, 'import pandas as pd\n'), ((24, 6, 31, 1), 'pandas.read_excel', 'pd.read_excel', (), '', True, 'import pandas as pd\n'), ((67, 6, 75, 10), 'pandas.concat', 'pd.concat', (), '', True, 'import pandas as pd\n'), ((77, 6, 85, 10), 'pandas.concat', 'pd.concat', (), '', True, 'import pandas as pd\n'), ((96, 5, 99, 10), 'pandas.concat', 'pd.concat', (), '', True, 'import pandas as pd\n')]