# Configure Jupyter so figures appear in the notebook
%matplotlib inline

# Configure Jupyter to display the assigned value after an assignment
%config InteractiveShell.ast_node_interactivity='last_expr_or_assign'

# import classes from thinkbayes2
from thinkbayes2 import Pmf, Suite

import thinkbayes2
import thinkplot

"""
Explanation: Think Bayes
This notebook presents example code and exercise solutions for Think Bayes.
Copyright 2018 Allen B. Downey
MIT License: https://opensource.org/licenses/MIT
End of explanation
"""


from random import random

def flip(p):
    return random() < p

"""
Explanation: The Dungeons and Dragons club
Suppose there are 10 people in my Dungeons and Dragons club; on any game day, each of them has a 70% chance of showing up.
Each player has one character and each character has 6 attributes, each of which is generated by rolling and adding up 3 6-sided dice.
At the beginning of the game, I ask whose character has the lowest attribute. The wizard says, "My constitution is 5; does anyone have a lower attribute?", and no one does.
The warrior says "My strength is 16; does anyone have a higher attribute?", and no one does.
How many characters are in the party?
The prior
There are three ways to compute the prior distribution:


Simulation


Convolution


Analytic distribution


First, simulation.  Here's a function that flips a coin with probability p:
End of explanation
"""


flips = [flip(0.7) for i in range(10)]

"""
Explanation: We can use it to flip a coin for each member of the club.
End of explanation
"""


sum(flips)

"""
Explanation: And count the number that show up on game day.
End of explanation
"""


def game_day(n, p):
    flips = [flip(p) for i in range(n)]
    return sum(flips)

game_day(10, 0.7)

"""
Explanation: Let's encapsulate that in a function that simulates a game day.
End of explanation
"""


sample = [game_day(10, 0.7) for i in range(1000)]
pmf_sample = Pmf(sample)
thinkplot.Hist(pmf_sample)

"""
Explanation: If we run that function many times, we get a sample from the distribution of the number of players.
End of explanation
"""


def coin(p):
    return Pmf({1:p, 0:1-p})

"""
Explanation: The second method is convolution.  Instead of flipping a coin, we can create a Pmf object that represents the distribution of outcomes from a single flip.
End of explanation
"""


player = coin(0.7)
player.Print()

"""
Explanation: Here's what it looks like.
End of explanation
"""


(player + player).Print()

"""
Explanation: If we have two players, there are three possible outcomes:
End of explanation
"""


prior = sum([player]*10)
prior.Print()

"""
Explanation: If we have 10 players, we can get the prior distribution like this:
End of explanation
"""


thinkplot.Hist(pmf_sample, color='C0')
thinkplot.Pmf(prior, color='C1')

thinkplot.decorate(xlabel='Number of players',
                   ylabel='PMF')

"""
Explanation: Now we can compare the results of simulation and convolution:
End of explanation
"""


help(thinkbayes2.MakeBinomialPmf)

"""
Explanation: Finally, we can use an analytic distribution.  The distribution we just computed is the binomial distribution, which has the following PMF:
$ PMF(k; n, p) = P(k ~|~ n, p) = {n \choose k}\,p^{k}(1-p)^{n-k}$
We could evalate the right hand side in Python, or use MakeBinomialPmf:
End of explanation
"""


binomial = thinkbayes2.MakeBinomialPmf(10, 0.7)
thinkplot.Pmf(prior, color='C1')
thinkplot.Pmf(binomial, color='C2', linestyle='dotted')

thinkplot.decorate(xlabel='Number of players',
                   ylabel='PMF')

"""
Explanation: And we can confirm that the analytic result matches what we computed by convolution.
End of explanation
"""


thinkplot.Pmf(prior, color='gray')
del prior[0]
del prior[1]
prior.Normalize()
thinkplot.Pmf(prior, color='C1')

thinkplot.decorate(xlabel='Number of players',
                   ylabel='PMF')

"""
Explanation: Since two players spoke, we can eliminate the possibility of 0 or 1 players:
End of explanation
"""


d6 = Pmf([1,2,3,4,5,6])
d6.Print()

"""
Explanation: Likelihood
There are three components of the likelihood function:


The probability that the highest attribute is 16.


The probability that the lowest attribute is 5.


The probability that the lowest and highest attributes are held by different players.


To compute the first component, we have to compute the distribution of the maximum of $6n$ attributes, where $n$ is the number of players.
Here is the distribution for a single die.
End of explanation
"""


thrice = sum([d6] * 3)
thinkplot.Pdf(thrice)
thinkplot.decorate(xlabel='Attribute',
                   ylabel='PMF')

"""
Explanation: And here's the distribution for the sum of three dice.
End of explanation
"""


cdf_thrice = thrice.MakeCdf()
thinkplot.Cdf(cdf_thrice)
thinkplot.decorate(xlabel='Attribute',
                   ylabel='CDF')

"""
Explanation: Here's the CDF for the sum of three dice.
End of explanation
"""


cdf_max_6 = cdf_thrice.Max(6)
thinkplot.Cdf(cdf_max_6)
thinkplot.decorate(xlabel='Attribute',
                   ylabel='CDF',
                   title='Maximum of 6 attributes')

"""
Explanation: The Max method raises the CDF to a power.  So here's the CDF for the maximum of six attributes.
End of explanation
"""


for n in range(2, 10, 2):
    cdf_max = cdf_thrice.Max(n*6)
    thinkplot.Cdf(cdf_max, label='n=%s'%n)

thinkplot.decorate(xlabel='Attribute',
                   ylabel='CDF',
                   title='Maximum of 6*n attributes')

"""
Explanation: If there are n players, there are 6*n attributes.  Here are the distributions for the maximum attribute of n players, for a few values of n.
End of explanation
"""


n = 7
cdf = cdf_thrice.Max(n*6)
thinkplot.Cdf(cdf, label='n=%s'%n)

sample_max = [max(cdf_thrice.Sample(42)) for i in range(1000)]
thinkplot.Cdf(thinkbayes2.Cdf(sample_max), label='sample')

thinkplot.decorate(xlabel='Attribute',
                   ylabel='CDF',
                   title='Maximum of 6*n attributes')

"""
Explanation: To check that, I'll compute the CDF for 7 players, and estimate it by simulation.
End of explanation
"""


def compute_cdf_min(cdf, k):
    """CDF of the min of k samples from cdf.
    
    cdf: Cdf object
    k: number of samples
    
    returns: new Cdf object
    """
    cdf_min = cdf.Copy()
    cdf_min.ps = 1 - (1 - cdf_min.ps)**k
    return cdf_min

"""
Explanation: Looks good.
Now, to compute the minimum, I have to write my own function, because Cdf doesn't provide a Min function.
End of explanation
"""


for n in range(2, 10, 2):
    cdf_min = compute_cdf_min(cdf_thrice, n*6)
    thinkplot.Cdf(cdf_min, label='n=%s'%n)

thinkplot.decorate(xlabel='Attribute',
                   ylabel='CDF',
                   title='Minimum of 6*n attributes')

"""
Explanation: Now we can compute the CDF of the minimum attribute for n players, for several values of n.
End of explanation
"""


n = 7
cdf = compute_cdf_min(cdf_thrice, n*6)
thinkplot.Cdf(cdf, label='n=%s'%n)

sample_min = [min(cdf_thrice.Sample(42)) for i in range(1000)]
thinkplot.Cdf(thinkbayes2.Cdf(sample_min), label='sample')

thinkplot.decorate(xlabel='Attribute',
                   ylabel='CDF',
                   title='Minimum of 6*n attributes')

"""
Explanation: And again we can check it by comparing to simulation results.
End of explanation
"""


like_min = {}
like_max = {}

for n in range(2, 11):
    cdf_min = compute_cdf_min(cdf_thrice, n*6)
    like_min[n] = cdf_min.MakePmf()
    cdf_max = cdf_thrice.Max(n*6)
    like_max[n] = cdf_max.MakePmf()
    print(like_min[n][5], like_max[n][16])

"""
Explanation: For efficiency and conciseness, it is helpful to precompute the distributions for the relevant values of n, and store them in dictionaries.
End of explanation
"""


def prob_same(n):
    return 5 / (6*n-1)

for n in range(2, 11):
    print(n, prob_same(n))

"""
Explanation: The output shows that the particular data we saw is symmetric: the chance that 16 is the maximum is the same as the chance that 5 is the minimum.
Finally, we need the probability that the minimum and maximum are held by the same person.  If there are n players, there are 6*n attributes.
Let's call the player with the highest attribute Max.  What is the chance that Max also has the lowest attribute?  Well Max has 5 more attributes, out of a total of 6*n-1 remaining attributes.
So here's prob_same as a function of n.
End of explanation
"""


class Dungeons(Suite):
    
    def Likelihood(self, data, hypo):
        """Probability of the data given the hypothesis.
        
        data: lowest attribute, highest attribute, boolean
              (whether the same person has both)
        hypo: number of players
        
        returns: probability
        """
        lowest, highest, same = data
        n = hypo
        
        p = prob_same(n)
        like = p if same else 1-p
        
        like *= like_min[n][lowest]
        like *= like_max[n][highest]

        return like

"""
Explanation: The update
Here's a class that implements this likelihood function.
End of explanation
"""


suite = Dungeons(prior)
thinkplot.Hist(suite)
thinkplot.decorate(xlabel='Number of players',
                   ylabel='PMF')
suite.Mean()

"""
Explanation: Here's the prior we computed above.
End of explanation
"""


suite.Update((5, 16, False))

"""
Explanation: And here's the update based on the data in the problem statement.
End of explanation
"""


thinkplot.Hist(suite)
thinkplot.decorate(xlabel='Number of players',
                   ylabel='PMF')
suite.Mean()

suite.Print()

"""
Explanation: Here's the posterior.
End of explanation
"""


suite.CredibleInterval()

sum(suite[n] for n in [5,6,7,8,9])

"""
Explanation: Based on the data, I am 94% sure there are between 5 and 9 players.
End of explanation
"""


