import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline
"""
# Improving Matplotlib Aesthetics to emulate FiveThirtyEight.
"""
"""
### Set style as default 538 style
We will edit it later.
"""
import matplotlib.style as style
style.available
"""
Note the style "fivethirtyeight" that we will begin with.
"""
style.use("fivethirtyeight")
"""
### Import data
"""
world = pd.read_csv("../input/world-happiness-report-2021/world-happiness-report.csv")
world.head()
"""
### Get top 5 happiest 2021 to compare to the US over time
"""
# get top 5 happiest countries, 2020
top5_2020 = world[world["year"]==2020].sort_values(by="Life Ladder", ascending = False).head()["Country name"]

# subset by top 5 happiest and least happy countries, plus the US
subset = world[
    (world["Country name"].isin(top5_2020)) 
    | (world["Country name"]=="United States")
]
subset
"""
### Initial plot | Happiness over time
"""
fig, ax = plt.subplots(figsize = (14,8))

for country, group in subset.groupby("Country name"):
    group.plot(
        ax = ax, 
        x = "year", 
        y = "Life Ladder", 
        label = country
    )

plt.show()
"""
### Not bad for a first plot. Let's add some things.
"""
fig, ax = plt.subplots(figsize = (14,8))

for country, group in subset.groupby("Country name"):
    group.plot(
        ax = ax, 
        x = "year", 
        y = "Life Ladder", 
        label = country
    )
    
### add the following: ###

# y-tick label size
ax.tick_params(axis="both", which="major", labelsize=18)

# capitalize x-label
ax.set_xlabel("Year")

# bold x-axis
plt.axhline(6.5, color="black", linewidth=1.5, alpha=0.6)

plt.show()
"""
### Getting better...
"""
fig, ax = plt.subplots(figsize = (14,8))

for country, group in subset.groupby("Country name"):
    group.plot(
        ax = ax, 
        x = "year", 
        y = "Life Ladder", 
        label = country
    )
    
ax.tick_params(axis="both", which="major", labelsize=18)
ax.set_xlabel("Year")
plt.axhline(6.5, color="black", linewidth=1.5, alpha=0.6)

### add the following: ###

# add v-line next to y-tick labels
ax.set_xlim(2004.5,2021)
plt.axvline(2004.75, color="black", linewidth=1.5, alpha=0.5)

# add signature bar
ax.text(
    x = 2003.7,
    y = 6.17,
    s = "   twitch.tv/MitchsWorkshop                                                                                        Source: Gallup World Poll   ",
    fontsize = 17,
    color = "#F0F0F0",
    backgroundcolor = "grey"
)

# add title
ax.text(
    x = 2004,
    y = 8.3,
    s = "The World's Top 5 Happiest Countries, and the United States",
    fontsize = 21,
    fontweight = "bold",
    fontfamily = "poppins"
)

# add subtitle
ax.text(
    x = 2004,
    y = 8.141,
    s = "Note: the y-axis is not base-0. This is purely a demonstration of improved plot aesthetics.\nAll scores are out of 10.",
    fontsize = 18
)

# custom colors
colors = [
    "#3c9dbd",
    "#c2762b",
    "#cc274d",
    "#29802c",
    "#565e5e",
    "#623aa6"
]
leg = ax.get_legend()
for i,c in enumerate(colors):
    plt.gca().get_lines()[i].set_color(c)

plt.show()
"""
### Now let's remove the legend
"""
fig, ax = plt.subplots(figsize = (14,8))

for country, group in subset.groupby("Country name"):
    group.plot(
        ax = ax, 
        x = "year", 
        y = "Life Ladder", 
        label = country
    )
    
ax.tick_params(axis="both", which="major", labelsize=18)
ax.set_xlabel("Year")
plt.axhline(6.5, color="black", linewidth=1.5, alpha=0.6)
ax.set_xlim(2004.5,2021)
plt.axvline(2004.75, color="black", linewidth=1.5, alpha=0.5)

# signature bar
ax.text(
    x = 2003.7,
    y = 6.17,
    s = "   twitch.tv/MitchsWorkshop                                                                                        Source: Gallup World Poll   ",
    fontsize = 17,
    color = "#F0F0F0",
    backgroundcolor = "grey"
)

# title
ax.text(
    x = 2004,
    y = 8.3,
    s = "The World's Top 5 Happiest Countries, and the United States",
    fontsize = 21,
    fontweight = "bold",
    fontfamily = "poppins"
)

# add subtitle
ax.text(
    x = 2004,
    y = 8.141,
    s = "Note: the y-axis is not base-0. This is purely a demonstration of improved plot aesthetics.\nAll scores are out of 10.",
    fontsize = 18
)

# custom colors
colors = [
    "#3c9dbd",
    "#c2762b",
    "#cc274d",
    "#29802c",
    "#565e5e",
    "#623aa6"
]

# change line colors
for i,c in enumerate(colors):
    plt.gca().get_lines()[i].set_color(c)
    
### add the following: ###

# Denmark
plt.text(
    x = 2005.5,
    y = 7.88,
    s = "Denmark",
    rotation = -29,
    color = colors[0],
    fontsize = 14,
    fontweight = "bold"
)

# Finland
plt.text(
    x = 2006.2,
    y = 7.7,
    s = "Finland",
    rotation = 0,
    color = colors[1],
    fontsize = 14,
    fontweight = "bold"
)

# Iceland
plt.text(
    x = 2008,
    y = 6.95,
    s = "Iceland",
    rotation = 42,
    color = colors[2],
    fontweight = "bold"
)

# Switzerland
plt.text(
    x = 2012,
    y = 7.6,
    s = "Switzerland",
    rotation = -37,
    color = colors[4],
    fontweight = "bold"
)

# United States
plt.text(
    x = 2010,
    y = 6.985,
    s = "United States",
    rotation = -20,
    color = colors[5],
    fontweight = "bold"
)

# Iceland (whom I originally forgot. Sorry, Iceland.)
plt.text(
    x = 2013,
    y = 7.25,
    s = "Iceland",
    rotation = -25,
    color = colors[3],
    fontweight = "bold"
)

# disable legend
ax.get_legend().remove()

plt.savefig("after.png", bbox_inches = "tight") # for twitter @MitchsWorkshop
plt.show()
"""
### I build things like this live on [Twitch](http://twitch.tv/MitchsWorkshop)! Stop by and ask questions or just hang out and talk data!
"""