from __future__ import print_function, division

import numpy as np
import thinkstats2
import thinkplot

import redline

%matplotlib inline
formats = ['png', 'pdf']

gap_times = np.array(redline.OBSERVED_GAP_TIMES)
len(gap_times), min(gap_times), max(gap_times)

pmf_z = thinkstats2.Pmf(gap_times, label='actual')
cdf_z = pmf_z.MakeCdf()

pmf_zb = redline.BiasPmf(pmf_z, label='biased')
cdf_zb = pmf_zb.MakeCdf()

pmf_y = redline.PmfOfWaitTime(pmf_z)
cdf_y = pmf_y.MakeCdf()

pmf_yb = redline.PmfOfWaitTime(pmf_zb)
cdf_yb = pmf_yb.MakeCdf()

pmf_z.Mean() / 60, pmf_y.Mean() / 60

pmf_zb.Mean() / 60, pmf_yb.Mean() / 60

thinkplot.PrePlot(2)
thinkplot.Plot(cdf_z.xs/60, cdf_z.ps, label=cdf_z.label)
thinkplot.Config(xlabel='time between trains (minutes)', ylabel='CDF', loc='lower right')
thinkplot.Save('redline1', formats=formats)

thinkplot.PrePlot(2)
thinkplot.Plot(cdf_z.xs/60, cdf_z.ps, label=cdf_z.label)
thinkplot.Plot(cdf_zb.xs/60, cdf_zb.ps, label=cdf_zb.label)
thinkplot.Config(xlabel='time between trains (minutes)', ylabel='CDF', 
                 ylim=[0, 1], loc='lower right')
thinkplot.Save('redline2', formats=formats)