import pandas as pd

def chances_to_meet(person_a_wait_minutes, person_b_wait_minutes, time_window_minutes):
    a = person_a_wait_minutes / time_window_minutes
    b = person_b_wait_minutes / time_window_minutes
    return a - 1/2 * a ** 2 + b - 1/2 * b ** 2

#%% Show probability of meeting given same times

person_a = [t*60/8 for t in range(9)]
person_b = [t*60/8 for t in range(9)]

meeting_times_1 = list(zip(person_a,person_b))
probabilites_1 = []

for pa, pb in meeting_times_1:
    probabilites_1.append(chances_to_meet(pa,pb,60))

df1 = pd.DataFrame(probabilites_1, index = meeting_times_1)
ax1 = df1.plot(title="Probability that Romeo and Juliet will meet within an hour window given equal wait times",
               legend=None)

ax1.set_xlabel("Minutes that (Romeo, Juliet) will wait before leaving")
ax1.set_ylabel("Probability of meeting each other")


#%% Show probability of meeting given one will wait indefinitely

person_a = [60 for t in range(9)]
person_b = [t*60/8 for t in range(9)]

meeting_times_1 = list(zip(person_a,person_b))
probabilites_1 = []

for pa, pb in meeting_times_1:
    probabilites_1.append(chances_to_meet(pa,pb,60))

df2 = pd.DataFrame(probabilites_1, index = meeting_times_1)
ax2 = df2.plot(title="Probability that Romeo and Juliet will meet within an hour window if Romeo waits the whole hour",
               legend=None)

ax2.set_xlabel("Minutes that (Romeo, Juliet) will wait before leaving")
ax2.set_ylabel("Probability of meeting each other")