import matplotlib.image as mpimg
img = mpimg.imread('waitt.gif')
imshow(img)  # Image from http://mahalanobis.twoday.net/stories/3486587/

import numpy as np
import pymc as mc
import statsmodels.api as sm

a_interval = 3.0
tau1 = 2.0

unbiased_a = mc.Normal('unbiased_a', mu=a_interval, tau=tau1)

@mc.stochastic(name='unbiased_abs_a', trace=True)
def unbiased_abs_a(value=5.0, mu=a_interval, tau=tau1):
    if value < 0:
        return np.NINF
    return mc.normal_like(value, mu, tau)

@mc.stochastic(name='biased_a', trace=True)
def biased_a(value=5.0, mu=a_interval, tau=tau1):
    if value < 0:
        return np.NINF
    return np.sum(np.log(value) + mc.normal_like(value, mu, tau))

model = mc.MCMC([biased_a, unbiased_a, unbiased_abs_a])
model.sample(iter=100000, burn=1000)

x = model.trace('biased_a')[:]
y = model.trace('unbiased_a')[:]
z = model.trace('unbiased_abs_a')[:]

figure(figsize=(10, 5))
hist(x, alpha=0.5, label="Biased A", bins=50, normed=True)
hist(z, alpha=0.5, label="Unbiased A", bins=50, normed=True)
legend()
title("Comparison of Bus Arrival Time")
xlabel("Time (min)")
ylabel("Frequency")
savefig("dist_comp_bus_a.png")

figure(figsize=(10, 5))

ecdf_x = sm.distributions.ECDF(x)
ecdf_z = sm.distributions.ECDF(z)

x_x = np.linspace(min(x), max(x))
y_x = ecdf_x(x_x)

x_z = np.linspace(min(z), max(z))
y_z = ecdf_z(x_z)

plot(x_x, y_x, label="Biased A")
plot(x_z, y_z, label="Unbiased A")

legend()
xlabel("Time")
ylabel("CDF")
title("Comparison of Bus Arrival Time (CDF)")
savefig("dist_comp_bus_a_cdf.png")

b_interval = 15.0
tau = .5

unbiased_b = mc.Normal('unbiased_b', mu=b_interval, tau=tau)

@mc.stochastic(name='unbiased_abs_b', trace=True)
def unbiased_abs_b(value=5.0, mu=b_interval, tau=tau):
    if value < 0:
        return np.NINF
    return mc.normal_like(value, mu, tau)

@mc.stochastic(name='biased_b', trace=True)
def biased_b(value=5.0, mu=b_interval, tau=tau):
    if value < 0:
        return np.NINF
    return np.sum(np.log(value) + mc.normal_like(value, mu, tau))

model2 = mc.MCMC([biased_b, unbiased_b, unbiased_abs_b])
model2.sample(iter=100000, burn=1000)

x2 = model2.trace('biased_b')[:]
y2 = model2.trace('unbiased_b')[:]
z2 = model2.trace('unbiased_abs_b')[:]

figure(figsize=(10, 5))
hist(x2, alpha=0.5, label="Biased B", bins=50, normed=True)
hist(z2, alpha=0.5, label="Unbiased B", bins=50, normed=True)
legend()
xlabel("Time (min)")
ylabel("Frequency")
title("Comparison of Bus Arrival Time")
savefig("dist_comp_bus_b.png")

figure(figsize=(10, 5))

ecdf_x2 = sm.distributions.ECDF(x2)
ecdf_z2 = sm.distributions.ECDF(z2)

x_x2 = np.linspace(min(x2), max(x2))
y_x2 = ecdf_x2(x_x2)

x_z2 = np.linspace(min(z2), max(z2))
y_z2 = ecdf_z2(x_z2)

plot(x_x2, y_x2, label="Biased B")
plot(x_z2, y_z2, label="Unbiased B")

legend()
xlabel("Time")
ylabel("CDF")
title("Comparison of Bus Arrival Time (CDF)")
savefig("dist_comp_bus_b_cdf.png")

@mc.deterministic
def waittime_a(sample_a=biased_a):
    return sample_a * np.random.random()

@mc.deterministic
def waittime_b(sample_b=biased_b):
    return sample_b * np.random.random()

m = mc.MCMC([biased_a, biased_b, waittime_a, waittime_b])
m.sample(iter=100000, burn=1000)

wt_a = m.trace('waittime_a')[:]
wt_b = m.trace('waittime_b')[:]

figure(figsize=(10, 5))
hist(wt_a, alpha=0.5, bins=50, range=(0, 20), label="Bus A", normed=True)
hist(wt_b, alpha=0.5, bins=50, range=(0, 20), label="Bus B", normed=True)
title("Comparison of Bus Wait Time")
xlabel("Time (min)")
ylabel("Frequency")
legend()
savefig("dist_comp_wait_time.png")

def hop_on_fast_arrival(a_wait, b_wait):    
    a_arrival = a_wait + 15
    b_arrival = b_wait + 9
    
    result = np.empty(len(a_wait))
    
    result[(b_wait <= a_wait) & (b_arrival <= a_arrival)] = 1
    result[(b_wait <= a_wait) & (b_arrival > a_arrival)] = 0
    
    result[(b_wait > a_wait) & (b_arrival >= a_arrival)] = 1
    result[(b_wait > a_wait) & (b_arrival < a_arrival)] = 0
    return result

result = hop_on_fast_arrival(wt_a, wt_b)

np.count_nonzero(result)

np.count_nonzero(result) / float(len(result))

def wait_for_b(a_wait, b_wait, k=3.0):    
    a_arrival = a_wait + 15
    b_arrival = b_wait + 9
    
    result = np.empty(len(a_wait))
    
    result[(b_wait <= a_wait) & (b_arrival <= a_arrival)] = 1
    result[(b_wait <= a_wait) & (b_arrival > a_arrival)] = 0
    
    # a가 먼저 오고 k 분이 넘게 지난 상황. b를 기다렸다 탄다.
    result[(b_wait > a_wait) & (a_wait >= k) & (b_arrival <= a_arrival)] = 1
    result[(b_wait > a_wait) & (a_wait >= k) & (b_arrival > a_arrival)] = 0

    # a가 먼저 오고 아직 k 분 못되게 기다린 상황. 바로 a를 탄다.
    result[(b_wait > a_wait) & (a_wait < k) & (b_arrival <= a_arrival)] = 0
    result[(b_wait > a_wait) & (a_wait < k) & (b_arrival > a_arrival)] = 1
    return result

grid = np.linspace(0, 10, 50)
v = []
for k in grid:
    d = wait_for_b(wt_a, wt_b, k)
    v.append(np.count_nonzero(d)/float(len(d)))

figure(figsize=(10, 5))
plot(grid, v)
xlabel("Wait Time (min)")
ylabel("Winning Rate")
title("Winning Rate for Various Wait Time")
savefig("wait_time_success_rate.png")

print grid[np.argmax(v)], np.max(v)

np.max(v)/float(len(result))

def hop_on_fast_arrival_time(a_wait, b_wait):    
    a_arrival = a_wait + 15
    b_arrival = b_wait + 9
    
    result = np.empty(len(a_wait))
    
    result[b_wait <= a_wait] = b_arrival[b_wait <= a_wait]
    result[b_wait > a_wait] = a_arrival[b_wait > a_wait]
    return result

strategy1 = hop_on_fast_arrival_time(wt_a, wt_b)

figure(figsize=(10, 5))
hist(strategy1, bins=50)
xlabel("Time (min)")
ylabel("Frequency")
title("Home Arrival Time - FCFS")
savefig("home_arrival_time_fcfs.png")

print "Mean arrival time: {}".format(strategy1.mean())
print "STD of arrival time: {}".format(strategy1.std())
print "Median arrival time: {}".format(np.median(strategy1))

def wait_for_b_time(a_wait, b_wait, k=3.0):    
    a_arrival = a_wait + 15
    b_arrival = b_wait + 9
    
    result = np.empty(len(a_wait))
    
    result[b_wait <= a_wait] = b_arrival[b_wait <= a_wait]    
    # a가 먼저 오고 k 분이 넘게 지난 상황. b를 기다렸다 탄다.
    result[(b_wait > a_wait) & (a_wait >= k)] = b_arrival[(b_wait > a_wait) & (a_wait >= k)]
    # a가 먼저 오고 아직 k 분 못되게 기다린 상황. 바로 a를 탄다.
    result[(b_wait > a_wait) & (a_wait < k)] = a_arrival[(b_wait > a_wait) & (a_wait < k)]
    return result

strategy2 = wait_for_b_time(wt_a, wt_b, k=3.15)

figure(figsize=(10, 5))
hist(strategy2, bins=50)
xlabel("Time (min)")
ylabel("Frequency")
title("Home Arrival Time - WFB")
savefig("home_arrival_time_wfb.png")

print "Mean arrival time: {}".format(strategy2.mean())
print "STD of arrival time: {}".format(strategy2.std())
print "Median arrival time: {}".format(np.median(strategy2))

def wait_twice(a_wait, b_wait, k=3.0, l=3.0, seed=42):  
    if k < 0:
        k = 0
    if l < 0:
        l = 0
    # 빠른 버스가 먼저 오면 탄다.
    # 느린 버스가 먼저 오면 지금 기다린 시간 k에 따라 다른 전략을 취한다
    # 그 후에도 한 번 더 느린 버스가 오면 기다린 시간 l에 따라 다른 전략을 취한다
    a_arrival = a_wait + 15
    b_arrival = b_wait + 9
    
    result = np.empty(len(a_wait))
    
    # b가 먼저 오면 b를 탄다.
    result[b_wait <= a_wait] = b_arrival[b_wait <= a_wait]
    # a가 먼저 오고 아직 k 분 못되게 기다린 상황. 바로 a를 탄다.
    result[(b_wait > a_wait) & (a_wait < k)] = a_arrival[(b_wait > a_wait) & (a_wait < k)]
    
    # a가 먼저 오고 k 분이 넘게 지난 상황. b를 일단 기다려 본다.
    cond = (b_wait > a_wait) & (a_wait >= k)
    size = np.count_nonzero(cond)
    # 두 번째 버스는 unbiased distribution에서 뽑아야 함
    next_buses_a = abs(np.random.normal(a_interval, np.sqrt(1/tau), size=size))
    # 두 번째 버스 a가 오기 전 총 대기 시간
    wait_total = next_buses_a + a_wait[cond]
    
    tmp = np.empty(size)
    b_slices = b_arrival[cond]
        
    # 다음 버스가 b이면 b를 탄다.
    tmp[wait_total >= b_wait[cond]] = b_slices[wait_total >= b_wait[cond]]
    # 다음 버스가 또 a이면 고민을 한다.
    # 추가 대기 시간이 l 이상이면 b를 기다렸다 탄다.
    tmp[(wait_total < b_wait[cond]) & (next_buses_a >= l)] = b_slices[(wait_total < b_wait[cond]) & (next_buses_a >= l)]
    # 추가 대기 시간이 l 미만이면 그냥 a를 탄다.
    tmp[(wait_total < b_wait[cond]) & (next_buses_a < l)] = wait_total[(wait_total < b_wait[cond]) & (next_buses_a < l)] + 9

    result[cond] = tmp
    return result

strategy3 = wait_twice(wt_a, wt_b, k=0.0, l=6.0)

print "Mean arrival time: {}".format(strategy3.mean())
print "STD of arrival time: {}".format(strategy3.std())
print "Median arrival time: {}".format(np.median(strategy3))

figure(figsize=(10, 5))
hist(strategy3, bins=50)
xlabel("Time (min)")
ylabel("Frequency")
title("Home Arrival Time - Optimal")
savefig("home_arrival_time_optimal.png")

grid = np.linspace(0, 20, 50)
v = []
for i in grid:
    for j in grid:
        v.append(wait_twice(wt_a, wt_b, k=i, l=j).mean())

figure(figsize=(10, 5))
plot(v)
xlabel("Patched Time")
ylabel("Arrival Time")
title("Comparison of Different Wait Time (2-D)")
savefig("home_arrival_time_comp_2d.png")

figure(figsize=(10, 5))
for i in range(0, 3):
    plot(grid, v[i*len(grid):(i+1)*len(grid)], label="{} min".format(i))
legend()
xlabel("Time (min)")
ylabel("Arrival Time")
title("Comparison of Different Wait Time (1-D)")
savefig("home_arrival_time_comp_1d.png")

wait_twice(wt_a, wt_b, k=0, l=10).mean()