import pandas as pd
import numpy as np
df = pd.read_csv("data/pandas/titanic.csv")
print(df.shape)
df[["Survived", "Pclass", "Sex", "Age", "Fare"]].head(3)

df["child"] = np.where(df["Age"] < 16, "child", "adult")
df.groupby("child")["Survived"].mean().round(2)

step1 = df[df["Embarked"] == "S"]
step2 = step1[step1["Age"] >= 18]
step3 = step2.groupby("Pclass")["Age"].mean()
step4 = step3.sort_values(ascending=False)
step4.round(1)

(
    df
    .query("Embarked == 'S' and Age >= 18")
    .groupby("Pclass")["Age"]
    .mean()
    .sort_values(ascending=False)
    .round(1)
)

result = df
    .query("Age > 60")

result = (df.query("Age > 60")
          .sort_values("Age", ascending=False))
result[["Name", "Age", "Survived"]].head(3)

fare = (df
    .query("Sex == 'female' and Pclass == 1")
    ["Fare"]
    .median())
print(round(fare, 2))

min_fare = 100
(df
    .query("Fare > @min_fare and Pclass == 1")
    .sort_values("Fare", ascending=False)
    [["Name", "Fare", "Survived"]]
    .head(3))

(df
    .loc[lambda d: d["Cabin"].notna()]
    .groupby("Pclass")
    .size())

(df
    .query("Embarked in ['C', 'Q']")
    .groupby("Embarked")["Survived"]
    .mean()
    .round(2))

share = (df
    .query("Pclass == 3")
    ["Cabin"]
    .notna()
    .mean())
print(round(share, 3))

(df
    .sort_values(["Pclass", "Age"])
    .groupby("Pclass")
    .head(1)
    .reset_index(drop=True)
    [["Pclass", "Name", "Age", "Survived"]])

titles = (df
    .assign(Title=lambda d: d["Name"].str.extract(
        r",\s*([^.]+)\.", expand=False))
    ["Title"]
    .value_counts()
    .head(5))
titles

(df
    .assign(family=lambda d: d["SibSp"] + d["Parch"] + 1,
            alone=lambda d: d["family"] == 1)
    .groupby("alone")["Survived"]
    .mean()
    .round(2))

(df
    .assign(Age=lambda d: d["Age"].fillna(d["Age"].median()))
    ["Age"]
    .describe()
    .round(1))

print(df["Age"].isna().sum())

(df.assign(family=lambda d: d["SibSp"] + d["Parch"] + 1)
   .assign(alone=df["family"] == 1))

n_alone = (df.assign(family=lambda d: d["SibSp"] + d["Parch"] + 1)
             .assign(alone=lambda d: d["family"] == 1)
             ["alone"].sum())
print(n_alone)

(df
    .assign(family=lambda d: d["SibSp"] + d["Parch"] + 1,
            per_person=lambda d: d["Fare"] / d["family"])
    .nlargest(3, "per_person")
    [["Name", "Fare", "family", "per_person"]]
    .round(2))

summary = (df
    .dropna(subset=["Embarked"])
    .groupby(["Embarked", "Pclass"])
    .agg(rate=("Survived", "mean"),
         people=("Survived", "size"))
    .reset_index()
    .sort_values("rate", ascending=False))
summary.head().round(2)

(df
    .merge(summary, on=["Embarked", "Pclass"], how="left")
    [["Embarked", "Pclass", "Survived", "rate"]]
    .head(3)
    .round(2))

(df
    .dropna(subset=["Embarked"])
    .pivot_table(index="Pclass", columns="Embarked",
                 values="Fare", aggfunc="median")
    .round(2))

import matplotlib.pyplot as plt
ax = (df
    .groupby(["Pclass", "Sex"])["Survived"]
    .mean()
    .unstack()
    .plot(kind="bar", rot=0, figsize=(7, 4)))
ax.set_title("Titanic survival rate by class and sex")
ax.set_xlabel("Passenger class")
ax.set_ylabel("Share who survived")
plt.show()

(df
    .dropna(subset=["Embarked"])
    .groupby("Embarked")["Age"]
    .median()
    .sort_values(ascending=False))

def add_age_band(d):
    bins = [0, 12, 18, 35, 60, 100]
    names = ["child", "teen", "young", "adult", "senior"]
    band = pd.cut(d["Age"], bins=bins, labels=names)
    return d.assign(age_band=band)

(df.pipe(add_age_band)["age_band"]
   .value_counts(sort=False))

def survival_by(d, col):
    return d.groupby(col, observed=True)["Survived"].mean()

(df
    .pipe(add_age_band)
    .pipe(survival_by, col="age_band")
    .round(2))

def peek(d, label):
    print(label, d.shape)
    return d

rate = (df
    .pipe(peek, "start")
    .dropna(subset=["Age"])
    .pipe(peek, "with age")
    .query("Sex == 'female'")
    .pipe(peek, "women")["Survived"].mean())
print(round(rate, 2))

(df
    .pipe(lambda d: d[d["Fare"] > d["Fare"].median()])
    ["Pclass"]
    .value_counts())

(df.copy()
    .dropna(subset=["Embarked"], inplace=True)
    .groupby("Embarked").size())

(df
    .dropna(subset=["Embarked"])
    .groupby("Embarked").size())

print(df["Age"].isna().sum(), (df["Sex"] == "male").sum())

bands = ["Low", "Mid", "High", "Top"]
fares = (df
    .assign(band=lambda d: pd.qcut(d["Fare"], 4, labels=bands))
    .groupby(["band", "Sex"], observed=True)["Survived"]
    .mean()
    .unstack()
    .round(2))
fares

(df
    .assign(band=lambda d: pd.qcut(d["Fare"], 4, labels=bands))
    ["band"]
    .value_counts(sort=False))

(df
    .assign(band=lambda d: pd.qcut(d["Fare"], 4, labels=bands))
    .groupby("band", observed=True)["Pclass"]
    .value_counts()
    .unstack(fill_value=0))

(df
    .query("Pclass == 3")
    .assign(band=lambda d: pd.qcut(d["Fare"], 2,
                                   labels=["cheaper", "dearer"]))
    .groupby(["band", "Sex"], observed=True)["Survived"]
    .mean().unstack().round(2))

ax = fares.plot(kind="bar", rot=0, figsize=(7, 4))
ax.set_title("Survival rate by fare band and sex")
ax.set_xlabel("Fare band (quarters of passengers)")
ax.set_ylabel("Share who survived")
plt.show()

low, top = fares.loc["Low", "male"], fares.loc["Top", "male"]
print(f"Men: {low:.0%} in the Low band, {top:.0%} in the Top band")