import pandas as pd

# Read the tables from the wikipedia page, and return the second one (index 1), which corresponds to the
# gun death rate per state.
gun_death_rate = pd.read_html("https://en.wikipedia.org/wiki/Firearm_death_rates_in_the_United_States_by_state",
             header=0,
             attrs={"class": ["sortable","wikitable", "plainrowheaders"]})[1]

# Read the tables from the gun laws wikipedia page, and return the 51 tables corresponding to the gun laws
# in each state, plus the district of columbia
per_state_gun_laws = pd.read_html("https://en.wikipedia.org/wiki/Gun_laws_in_the_United_States_by_state",header=0)[1:-7]

# Let's extract the number of restrictions (ie, count the number of laws) for each state
number_of_restrictions_per_state = []
for state_laws in per_state_gun_laws:
    # Everywhere there is a "Yes", we say that there is an additional restriction,
    # except for the "open carry" row, which is reversed, disallowing open carry is
    # a restriction
    open_carry_row = (state_laws[state_laws["Subject/Law"].str.contains("Open carry|Open Carry")].index[0])

    handgun_restrictions = (state_laws["Handguns"].str.contains("Yes") 
                            + 0.5*state_laws["Handguns"].str.contains("Partial"))
    handgun_restrictions[open_carry_row] = 1 - handgun_restrictions[open_carry_row]

    # Annoyingly, wikipedia uses two different table headings
    try:
        longgun_restrictions = (state_laws["Long Guns"].str.contains("Yes")
                                + 0.5*state_laws["Long Guns"].str.contains("Partial"))
    except:
        longgun_restrictions = (state_laws["Long guns"].str.contains("Yes")
                                + 0.5*state_laws["Long guns"].str.contains("Partial"))

    longgun_restrictions[open_carry_row] = 1 - longgun_restrictions[open_carry_row]

    number_of_restrictions_per_state.append(handgun_restrictions.sum() + longgun_restrictions.sum())

# Thankfully the states are alphabetical in both wikipedia pages, so they match up without any more work.
# We just need to remove the district of columbia (8th element, 0-indexed)
del number_of_restrictions_per_state[8]
full_dataset = gun_death_rate.assign(Laws=number_of_restrictions_per_state)

# Now a bit of cleanup:
# Winsconsin has a weird table, which we need to fix manually
full_dataset.loc[full_dataset["State"]=="Wisconsin","Laws"] += 3

full_dataset

%matplotlib inline
import seaborn as sb

sb.regplot(x="Laws",y="Rate",data=full_dataset)