import pandas as pd
import numpy as np

raw = pd.read_csv("data/pandas/salary_survey.csv")
q1, q3 = raw["annual_salary"].quantile([0.25, 0.75])
upper = q3 + 1.5 * (q3 - q1)
print("Upper fence:", upper)
print("Above it:", (raw["annual_salary"] > upper).sum())
print("Largest salary:", raw["annual_salary"].max())

print(raw.shape)
raw.columns.tolist()

raw.isna().mean().mul(100).round(1).sort_values(ascending=False).head(6)

raw["currency"].value_counts().head(6)

counts = raw["industry"].value_counts()
print(len(counts), "different industries")
print((counts == 1).sum(), "typed by only one person")

names = {"how_old_are_you": "age", "annual_salary": "salary",
         "other_monetary_comp": "bonus",
         "overall_years_of_professional_experience": "years_total",
         "years_of_experience_in_field": "years_field",
         "highest_level_of_education_completed": "education"}
keep = ["timestamp", "age", "industry", "job_title", "salary",
        "bonus", "currency", "country", "years_total",
        "years_field", "education", "gender"]
df = raw.rename(columns=names)[keep].copy()
df.head(3).T

print(df["country"].nunique(), "distinct country values")
df["country"].value_counts().head(8)

country = df["country"].str.strip().str.lower()
country = country.str.replace(".", "", regex=False)
print(country.nunique(), "after strip, lower and dots removed")
country.value_counts().head(5)

us = ["united states", "usa", "us", "united states of america",
      "america", "united state", "unites states", "united stated",
      "u s"]
uk = ["uk", "united kingdom", "england", "scotland", "wales",
      "northern ireland"]
fix = {**dict.fromkeys(us, "united states"),
       **dict.fromkeys(uk, "united kingdom")}
df["country"] = country.replace(fix).str.title()
df["country"].value_counts().head(5)

done = ["United States", "United Kingdom"]
has_unit = df["country"].str.contains("unit", case=False)
left = df.loc[has_unit & ~df["country"].isin(done), "country"]
print(len(left), "rows still unmapped")
left.value_counts().head(6)

us_like = df["country"].str.contains(r"^(?:the )?unit\w* ?sta",
                                    case=False)
uk_like = df["country"].str.contains("kingdom|kindom", case=False)
df.loc[us_like, "country"] = "United States"
df.loc[uk_like, "country"] = "United Kingdom"
left = df.loc[df["country"].str.contains("unit", case=False)
              & ~df["country"].isin(done), "country"]
print(len(left), "rows still unmapped")

df["gender"] = df["gender"].replace(
    {"Prefer not to answer": "Other or prefer not to answer"})
df["gender"].value_counts(dropna=False)

df["timestamp"] = pd.to_datetime(df["timestamp"],
                                 format="%m/%d/%Y %H:%M:%S")
df["timestamp"].agg(["min", "max"])

df["salary"].str.replace(",", "")

print("salary is stored as", df["salary"].dtype)
df["bonus"] = df["bonus"].fillna(0)
df["total_pay"] = df["salary"] + df["bonus"]
df[["salary", "bonus", "total_pay"]].head(3)

big_bonus = df["bonus"] > df["salary"]
print(big_bonus.sum(), "people report a bonus bigger than salary")
df.loc[big_bonus, ["job_title", "salary", "bonus"]].head(4)

df["years_field"].unique()

bands = ["1 year or less", "2-4 years", "5-7 years", "8-10 years",
         "11-20 years", "21-30 years", "31-40 years",
         "41 years or more"]
matches = df["years_total"].isin(bands)
print((~matches).sum(), "labels do not match a band")

for col in ["years_total", "years_field"]:
    tidy = df[col].str.replace(" - ", "-")
    df[col] = pd.Categorical(tidy, categories=bands, ordered=True)
    print(col, "missing:", df[col].isna().sum())
df["years_field"].value_counts(sort=False)

ages = ["under 18", "18-24", "25-34", "35-44", "45-54", "55-64",
        "65 or over"]
df["age"] = pd.Categorical(df["age"], categories=ages, ordered=True)
print(df["age"].isna().sum(), "unmatched")
df.dtypes

print(df.duplicated().sum(), "exact duplicate rows")

answers = df.columns.drop("timestamp")
dup = df.duplicated(subset=answers, keep=False)
print(df.duplicated(subset=answers).sum(), "repeat submissions")
cols = ["timestamp", "job_title", "salary"]
df[dup].sort_values(["job_title", "salary"])[cols].head(4)

before = len(df)
df = df.drop_duplicates(subset=answers, keep="first")
print(before, "->", len(df), "rows")

in_us = df["country"] == "United States"
other = (in_us & (df["currency"] != "USD")).sum()
print(other, "US answers in another currency")
us = df[in_us & (df["currency"] == "USD")].copy()
print(len(us), "US answers paid in dollars")

us["salary"].describe().round(0)

q1, q3 = us["salary"].quantile([0.25, 0.75])
fence = q3 + 1.5 * (q3 - q1)
high = us[us["salary"] > fence]
print(len(high), "above the fence of", fence)
high["industry"].value_counts().head(4)

us.nlargest(5, "salary")[["job_title", "industry", "salary"]]

too_low = us["salary"] < 15_000
print(too_low.sum(), "salaries below 15,000")
us.loc[too_low, ["job_title", "salary"]].head(6)

low = us.loc[too_low, "salary"]
print((low < 1_000).sum(), "of them are under 1,000")
low.describe().round(0)

clean = us[~too_low]
print(len(clean), "rows kept out of", len(raw))
clean.isna().sum()[lambda s: s > 0]

import os
os.makedirs("output/pandas", exist_ok=True)
clean.to_csv("output/pandas/l21_salary_clean.csv", index=False)
back = pd.read_csv("output/pandas/l21_salary_clean.csv")
print(back.shape)

print("years_field:", back["years_field"].dtype)
print("timestamp:  ", back["timestamp"].dtype)

clean.to_parquet("output/pandas/l21_salary_clean.parquet")
back = pd.read_parquet("output/pandas/l21_salary_clean.parquet")
print("years_field:", back["years_field"].dtype)
print("timestamp:  ", back["timestamp"].dtype)
print(back["years_field"].cat.ordered, "ordered")

log = pd.Series({"raw answers": len(raw),
                 "after duplicates": len(df),
                 "US, paid in USD": len(us),
                 "believable salary": len(clean)})
pd.DataFrame({"rows": log, "percent": (log / len(raw) * 100).round(1)})

pay = clean.groupby("years_field", observed=True)["salary"]
by_years = pay.agg(["median", "count"])
by_years

import matplotlib.pyplot as plt

ax = by_years["median"].plot.bar(figsize=(9, 4), rot=30)
ax.set_title("Median US salary by years in field (2021 survey)")
ax.set_xlabel("Years of experience in the field")
ax.set_ylabel("Median salary (USD)")
plt.show()

top = clean["industry"].value_counts().head(3).index
subset = clean[clean["industry"].isin(top)]
table = subset.pivot_table(index="years_field", columns="industry",
                           values="salary", aggfunc="median",
                           observed=True)
table.round(-3)

gap = table["Computing or Tech"] / table["Nonprofits"]
gap.round(2)

pay_cols = ["salary", "total_pay"]
clean.groupby("industry")[pay_cols].median().loc[top].round(-3)