%pylab inline

import scipy.stats as stats
figsize(12.5,3)
colors = ["#348ABD", "#A60628", "#7A68A6", "#467821"]

x = np.linspace(0,1)
y1, y2 = stats.beta.pdf(x, 1,1), stats.beta.pdf(x, 10,10)

p = plt.plot( x, y1, 
    label='An objective prior \n(uninformative, \n"Principle of Indifference" )'  )
plt.fill_between( x, 0, y1, color = p[0].get_color(), alpha = 0.3 )

p = plt.plot( x,y2 ,
     label = "A subjective prior \n(informative)" )
plt.fill_between( x, 0, y2, color = p[0].get_color(), alpha = 0.3 )

p = plt.plot( x[25:], 2*np.ones(25), label = "another subjective prior" )
plt.fill_between( x[25:], 0, 2, color = p[0].get_color(), alpha = 0.3 )

plt.ylim(0,4)

plt.ylim(0, 4)
leg = plt.legend(loc = "upper left")
leg.get_frame().set_alpha(0.4)
plt.title("Comparing objective vs. subjective priors for an unknown probability" );

figsize( 12.5, 5 )
gamma = stats.gamma

parameters = [ (1, 0.5), (9, 2), (3, 0.5 ), (7, 0.5) ]


x = linspace( 0.001 ,20, 150 )
for alpha, beta in parameters:
    y = gamma.pdf(x, alpha, scale=1./beta)
    lines = plt.plot( x, y, label = "(%.1f,%.1f)"%(alpha,beta), lw = 3 )
    plt.fill_between( x, 0, y, alpha = 0.2, color = lines[0].get_color() )
    plt.autoscale(tight=True)
    
plt.legend(title=r"$\alpha, \beta$ - parameters")

import pymc as mc

n = 4
for i in range( 10 ):
    ax = plt.subplot( 2, 5, i+1)
    if i >= 5:
        n = 15
    plt.imshow( mc.rwishart( n+1, np.eye(n) ), interpolation="none", 
                cmap = plt.cm.hot ) 
    
    ax.axis("off")
plt.suptitle("Random matrices from a Wishart Distribution" );

figsize( 12.5, 5 )
import scipy.stats as stats

params = [ (2,5), (1,1), (0.5, 0.5), ( 5, 5), (20, 4), (5, 1)]

x = np.linspace( 0.01, .99, 100 )
beta = stats.beta
for a, b in params:
    y = beta.pdf( x, a, b )
    lines = plt.plot( x, y, label = "(%.1f,%.1f)"%(a,b), lw = 3 )
    plt.fill_between( x, 0, y, alpha = 0.2, color = lines[0].get_color() )
    plt.autoscale(tight=True)
plt.ylim(0)
plt.legend(loc = 'upper left', title="(a,b)-parameters");

from pymc import rbeta

rand = np.random.rand

class Bandits(object):
    """
    This class represents N bandits machines.

    parameters:
        p_array: a (n,) Numpy array of probabilities >0, <1.

    methods:
        pull( i ): return the results, 0 or 1, of pulling 
                   the ith bandit.
    """
    def __init__(self, p_array):
        self.p = p_array
        self.optimal = np.argmax(p_array)
        
    def pull( self, i ):
        #i is which arm to pull
        return rand() < self.p[i]
    
    def __len__(self):
        return len(self.p)

    
class BayesianStrategy( object ):
    """
    Implements a online, learning strategy to solve
    the Multi-Armed Bandit problem.
    
    parameters:
        bandits: a Bandit class with .pull method
    
    methods:
        sample_bandits(n): sample and train on n pulls.

    attributes:
        N: the cumulative number of samples
        choices: the historical choices as a (N,) array
        bb_score: the historical score as a (N,) array
    """
    
    def __init__(self, bandits):
        
        self.bandits = bandits
        n_bandits = len( self.bandits )
        self.wins = np.zeros( n_bandits )
        self.trials = np.zeros(n_bandits )
        self.N = 0
        self.choices = []
        self.bb_score = []

    
    def sample_bandits( self, n=1 ):
        
        bb_score = np.zeros( n )
        choices = np.zeros( n )
        
        for k in range(n):
            #sample from the bandits's priors, and select the largest sample
            choice = np.argmax( rbeta( 1 + self.wins, 1 + self.trials - self.wins) )
            
            #sample the chosen bandit
            result = self.bandits.pull( choice )
            
            #update priors and score
            self.wins[ choice ] += result
            self.trials[ choice ] += 1
            bb_score[ k ] = result 
            self.N += 1
            choices[ k ] = choice
            
        self.bb_score = np.r_[ self.bb_score, bb_score ]
        self.choices = np.r_[ self.choices, choices ]
        return 

import scipy.stats as stats
figsize( 11.0, 10)

beta = stats.beta
x = np.linspace(0.001,.999,200)

def plot_priors(bayesian_strategy, prob, lw = 3, alpha = 0.2, plt_vlines = True):
    ## plotting function
    wins = bayesian_strategy.wins
    trials = bayesian_strategy.trials
    for i in range( prob.shape[0] ):
        y = beta( 1+wins[i], 1 + trials[i] - wins[i] )
        p = plt.plot( x, y.pdf(x), lw = lw )
        c = p[0].get_markeredgecolor()
        plt.fill_between(x,y.pdf(x),0, color = c, alpha = alpha, 
                         label="underlying probability: %.2f"%prob[i])
        if plt_vlines:
            plt.vlines( prob[i], 0, y.pdf(prob[i]) ,
                       colors = c, linestyles = "--", lw = 2 )
        plt.autoscale(tight = "True")
        plt.title("Posteriors After %d pull"%bayesian_strategy.N +\
                    "s"*(bayesian_strategy.N>1) )
        plt.autoscale(tight=True)
    return

hidden_prob = np.array([0.85, 0.60, 0.75] )
bandits = Bandits( hidden_prob )
bayesian_strat = BayesianStrategy( bandits )

draw_samples = [1, 1, 3, 10, 10, 25, 50, 100, 200, 600 ]

for j,i in enumerate(draw_samples):
    subplot( 5, 2, j+1) 
    bayesian_strat.sample_bandits(i)
    plot_priors( bayesian_strat, hidden_prob )
    #plt.legend()
    plt.autoscale( tight = True )
plt.tight_layout()

from IPython.core.display import HTML

#try executing the below command twice if the first time doesn't work
HTML(filename = "BanditsD3.html" )

figsize( 12.5, 5 )

from other_strats import *

#define a harder problem
hidden_prob = np.array([0.15, 0.2, 0.1, 0.05] )
bandits = Bandits( hidden_prob )


#define regret
def regret( probabilities, choices ):
    w_opt = probabilities.max()
    return ( w_opt - probabilities[choices.astype(int)] ).cumsum()


#create new strategies

strategies= [upper_credible_choice, 
            bayesian_bandit_choice, 
            ucb_bayes , 
            max_mean,
            random_choice ]
algos = []
for strat in strategies:
    algos.append( GeneralBanditStrat( bandits, strat ))

#train 10000 times
for strat in algos:
    strat.sample_bandits( 10000)
    
#test and plot
for i,strat in enumerate(algos):
    _regret = regret( hidden_prob, strat.choices )
    plt.plot( _regret, label = strategies[i].__name__, lw = 3)

plt.title("Total Regret of Bayesian Bandits Strategy vs. Random guessing" )
plt.xlabel("Number of pulls")
plt.ylabel("Regret after $n$ pulls");
plt.legend(loc = "upper left");

#this can be slow, so I recommend NOT running it. 

trials = 500
expected_total_regret = np.zeros( ( 10000, 3 ) )

for i_strat, strat in enumerate( strategies[:-2] ):
    for i in range(trials):
        general_strat = GeneralBanditStrat( bandits, strat )
        general_strat.sample_bandits(10000)
        _regret =  regret( hidden_prob, general_strat.choices ) 
        expected_total_regret[:,i_strat] += _regret
    
    plot(expected_total_regret[:,i_strat]/trials, lw =3, label = strat.__name__  )
        
plt.title("Expected Total Regret of Multi-armed Bandit strategies" )
plt.xlabel("Number of pulls")
plt.ylabel("Exepected Total Regret \n after $n$ pulls");
plt.legend(loc = "upper left");

plt.figure()
[pl1, pl2, pl3] = plt.plot( expected_total_regret[:, [0,1,2]], lw = 3  )
plt.xscale("log")
plt.legend([ pl1, pl2, pl3], 
           ["Upper Credible Bound", "Bayesian Bandit", "UCB-Bayes"],
            loc="upper left")
plt.ylabel("Exepected Total Regret \n after $\log{n}$ pulls");
plt.title( "log-scale of above" );
plt.ylabel("Exepected Total Regret \n after $\log{n}$ pulls");

figsize( 12.0, 8)
beta = stats.beta
hidden_prob = beta.rvs(1,13, size = 35 )
print hidden_prob
bandits = Bandits( hidden_prob )
bayesian_strat = BayesianStrategy( bandits )

for j,i in enumerate([100, 200, 500, 1300 ]):
    subplot( 2, 2, j+1) 
    bayesian_strat.sample_bandits(i)
    plot_priors( bayesian_strat, hidden_prob, lw = 2, alpha = 0.0, plt_vlines=False )
    #plt.legend()
    plt.xlim(0, 0.5 )


import scipy.stats as stats
figsize(11., 5 )
colors = ["#348ABD", "#A60628", "#7A68A6", "#467821"]

normal = stats.norm
x = np.linspace( -0.15, 0.15, 100 )

expert_prior_params = {"AAPL":(0.05, 0.03),
                 "GOOG":(-0.03, 0.04 ), 
                 "TSLA": (-0.02, 0.01), 
                 "AMZN": (0.03, 0.02 ), 
                 }

for i, (name, params) in enumerate(expert_prior_params.iteritems() ):
    subplot(2,2,i)
    y = normal.pdf( x, params[0], scale = params[1] )
    #plt.plot( x, y, c = colors[i] )
    plt.fill_between(x, 0, y, color = colors[i], linewidth=2,
                     edgecolor = colors[i], alpha = 0.6)
    plt.title(name + " prior" )
    plt.vlines(0, 0, y.max(), "k","--", linewidth = 0.5 )
    plt.xlim(-0.15, 0.15 )
plt.tight_layout()

import pymc as mc

n_observations = 100 #we will truncate the the most recent 100 days.

prior_mu = np.array( [ x[0] for x in expert_prior_params.values() ] )
prior_std = np.array( [ x[1] for x in expert_prior_params.values() ] )

inv_cov_matrix = mc.Wishart( "inv_cov_matrix", n_observations, diag(prior_std**2) )
mu = mc.Normal( "returns", prior_mu, 1, size = 4 )

# I wish I could have used Pandas as a prereq for this book, but oh well.

import datetime
import ystockquote as ysq

stocks = ["AAPL", "GOOG", "TSLA", "AMZN" ]

enddate =  datetime.datetime.now().strftime("%Y-%m-%d") #today's date.
startdate = "2012-09-01"

stock_closes = {}
stock_returns = {}
CLOSE = 6

for stock in stocks:
    x = np.array( ysq.get_historical_prices( stock, startdate, enddate ) )
    stock_closes[stock] = x[1:,CLOSE].astype(float)

#create returns:

for stock in stocks:
    _previous_day = np.roll(stock_closes[stock],-1)
    stock_returns[stock] = ((stock_closes[stock] - _previous_day)/_previous_day)[:n_observations]

dates = map( lambda x: datetime.datetime.strptime(x, "%Y-%m-%d" ), x[1:n_observations+1,0] ) 

figsize(12.5, 4)

for _stock, _returns in stock_returns.iteritems():
    p = plot( (1+_returns)[::-1].cumprod()-1, '-o', label = "%s"%_stock, 
                markersize=4, markeredgecolor="none" )

plt.xticks( arange(100)[::-8], 
            map(lambda x: datetime.datetime.strftime(x, "%Y-%m-%d"), dates[::8] ),
             rotation=60);

plt.legend(loc = "upper left")
plt.title("Return space")
plt.ylabel("Return of $1 on first date, x100%");

figsize(11., 5 )
returns = np.zeros( (n_observations,4 ))

for i, (_stock,_returns) in enumerate(stock_returns.iteritems() ):
    
    returns[:,i] = _returns
    
    subplot(2,2,i)
    plt.hist( _returns, bins = 20, 
              normed = True, histtype="stepfilled",
              color = colors[i], alpha = 0.7 )
    plt.title(_stock + " returns" )
    plt.xlim(-0.15, 0.15 )

plt.tight_layout()
plt.suptitle("Histogram of daily returns", size =14);

obs = mc.MvNormal( "observed returns", mu, inv_cov_matrix, observed = True, value = returns )

model = mc.Model( [obs, mu, inv_cov_matrix] )
mcmc = mc.MCMC()

mcmc.sample( 150000, 100000, 3 )

figsize(12.5,4)

#examine the mean return first.
mu_samples = mcmc.trace("returns")[:]

for i in range(4):
    plt.hist(mu_samples[:,i], alpha = 0.8 - 0.05*i, bins = 30,
             histtype="stepfilled", normed=True, 
             label = "%s"%stock_returns.keys()[i])

plt.vlines( mu_samples.mean(axis=0), 0, 500, linestyle="--", linewidth = .5 )

plt.title("Posterior distribution of $\mu$, daily stock returns" )
plt.legend();

figsize(11.0,3)
for i in range(4):
    subplot(2,2,i+1)
    plt.hist(mu_samples[:,i], alpha = 0.8 - 0.05*i, bins = 30,
             histtype="stepfilled", normed=True, color = colors[i],
             label = "%s"%stock_returns.keys()[i])
    plt.title( "%s"%stock_returns.keys()[i] )
    plt.xlim(-0.15, 0.15 )
    
plt.suptitle("Posterior distribution of daily stock returns" )
plt.tight_layout()

inv_cov_samples = mcmc.trace("inv_cov_matrix")[:]
mean_covariance_matrix = np.linalg.inv( inv_cov_samples.mean(axis=0) )


def cov2corr( A ):
    """
    covariance matrix to correlation matrix.
    """
    d = np.sqrt(A.diagonal())
    A = ((A.T/d).T)/d
    #A[ np.diag_indices(A.shape[0]) ] = np.ones( A.shape[0] )
    return A


subplot(1,2,1)
plt.imshow( cov2corr(mean_covariance_matrix) , interpolation="none", 
                cmap = plt.cm.hot ) 
plt.xticks( arange(4), stock_returns.keys() )
plt.yticks( arange(4), stock_returns.keys() )
plt.colorbar(orientation="vertical")
plt.title("(mean posterior) Correlation Matrix" )

subplot(1,2,2)
plt.bar(np.arange(4),  np.sqrt(np.diag(mean_covariance_matrix)),
        color = "#348ABD", alpha = 0.7 )
plt.xticks( np.arange(4) + 0.5, stock_returns.keys() );
plt.title("(mean posterior) variances of daily stock returns" )

plt.tight_layout();


figsize( 12.5, 15)

p = 0.6
beta1_params = np.array( [1.,1.] )
beta2_params = np.array( [2,10] )
beta = stats.beta

x = np.linspace(0.00, 1, 125)
data = mc.rbernoulli(p, size=500)

figure()
for i,N in enumerate([0,4,8, 32,64, 128, 500]):
    s = data[:N].sum() 
    subplot(8,1,i+1)
    params1 = beta1_params + np.array( [s, N-s] )
    params2 = beta2_params + np.array( [s, N-s] )
    y1,y2 = beta.pdf( x, *params1), beta.pdf( x, *params2)
    plt.plot( x,y1, label = r"flat prior", lw =3 )
    plt.plot( x, y2, label = "biased prior", lw= 3 )
    plt.fill_between( x, 0, y1, color ="#348ABD", alpha = 0.15) 
    plt.fill_between( x, 0, y2, color ="#A60628", alpha = 0.15) 
    plt.legend(title = "N=%d"%N)
    plt.vlines( p, 0.0, 7.5, linestyles = "--", linewidth=1)
    #plt.ylim( 0, 10)#


from IPython.core.display import HTML
def css_styling():
    styles = open("../styles/custom.css", "r").read()
    return HTML(styles)
css_styling()