import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

data = pd.read_csv("president_county_candidate.csv")

def GetFirstDigit(num) -> int:
    first_digit = num
    while (first_digit >= 10):
        first_digit = first_digit // 10
    return (1 if first_digit == 0 else first_digit)

biden = [0] * 10
trump = [0] * 10

# get occurance data
for index, row in data.iterrows():
    digit = GetFirstDigit(data['votes'][index])
    if (data['candidate'][index] == 'Joe Biden'):
        biden[digit] += 1
    elif (data['candidate'][index] == 'Donald Trump'):
        trump[digit] += 1

num_counties = len(data.county.unique())    # should be 3007, but we get the value from data anyway
if (num_counties != 3007): # bit of a sanity check
    print('Incomplete statistics')

# get occurance data as measure of probability
for index in range(10):
    biden[index] = biden[index] / num_counties
    trump[index] = trump[index] / num_counties


fig, ax = plt.subplots()
plt.xlim((1, 9))

benford = [0, .301, .176, .125, .097, .079, .067, .058, .051, 0.046]

ax.plot(biden, label='Biden', color='b')
ax.plot(trump, label='Trump', color='r')
ax.plot(benford, label='True Benford', color='y')

ax.set_ylabel('Probability of Occurance')
ax.set_xlabel('Most Significant Integer')
ax.set_title('2020 Votes for Major Candidates vs Benford\'s Law')
ax.legend()

plt.show()