# this line makes the code compatible with Python 2 and 3
from __future__ import print_function, division

# this line makes Jupyter show figures in the notebook
%matplotlib inline

"""
Explanation: Introduction to PMFs
Copyright 2015 Allen Downey
License: Creative Commons Attribution 4.0 International
End of explanation
"""


import random
import matplotlib.pyplot as plt
from collections import Counter
from itertools import izip

class Hist(Counter):
    
    def __add__(self, other):
        """Returns the Pmf of the sum of elements from self and other."""
        return Hist(x + y for x, y in product(self.elements(), other.elements()))
    
    def choice(self):
        """Chooses a random element."""
        return random.choice(list(self.elements()))
    
    def plot(self, **options):
        """Plots the Pmf."""
        plt.bar(*zip(*self.items()), **options)
        plt.xlabel('Values')
        plt.ylabel('Counts')
    
    def ranks(self):
        """Returns ranks and counts as lists."""
        return izip(*enumerate(sorted(self.values(), reverse=True)))

"""
Explanation: Hist objects
A histogram is a map from each possible value to the number of times it appears.  A map can be a mathematical function or, as in the examples below, a Python data structure that provides the ability to look up a value and get its probability.
Counter is a data structure provided by Python; I am defining a new data structure, called a Hist, that has all the features of a Counter, plus a few more that I define.
End of explanation
"""


hist = Hist('allen')
hist

"""
Explanation: As an example, I'll make a Hist of the letters in my name:
End of explanation
"""


hist['l']

"""
Explanation: We can look up a letter and get the corresponding count:
End of explanation
"""


for letter in hist:
    print(letter, hist[letter])

"""
Explanation: Or loop through all the letters and print their counts:
End of explanation
"""


hist.most_common()

"""
Explanation: Counter provides most_common, which makes a list of (element, count) pairs:
End of explanation
"""


for letter, count in hist.most_common():
    print(letter, count)

"""
Explanation: Here they are in a more readable form:
End of explanation
"""


for i in range(10):
    print(hist.choice())

"""
Explanation: I defined choice, which returns a random element from the Hist.  On average, 'l' should appear twice as often as the other letters.
End of explanation
"""


def is_anagram(word1, word2):
    return Hist(word1) == Hist(word2)

"""
Explanation: One (perhaps surprising) thing you can use Hists for: checking whether two words are anagrams of each  other.  If two words are anagrams, they have the same Hist:
End of explanation
"""


is_anagram('allen', 'nella')

"""
Explanation: Here's a simple test:
End of explanation
"""


is_anagram('tachymetric', 'mccarthyite')

"""
Explanation: And my favorite anagram pair:
End of explanation
"""


is_anagram('abcd', 'abccd')

"""
Explanation: And here's a false one, just to make sure:
End of explanation
"""


d6 = Hist([1,2,3,4,5,6])
d6

"""
Explanation: So far the elements in the Hists have been letters (actually strings), but in statistics it is more common to work with numerical elements.  Here's a Hist that represents the possible outcomes of a six-sided die:
End of explanation
"""


import seaborn as sns
COLORS = sns.color_palette()

d6.plot(color=COLORS[3])

"""
Explanation: Hist provides a plot function:
End of explanation
"""


d6.elements()

"""
Explanation: elements returns an iterator
End of explanation
"""


list(d6.elements())

"""
Explanation: Which is easier to see if you convert to a list:
End of explanation
"""


from itertools import product

product(d6.elements(), d6.elements())

"""
Explanation: The product of two iterators is an iterator that enumerates all pairs:
End of explanation
"""


list(product(d6.elements(), d6.elements()))

"""
Explanation: Here are the elements of the product:
End of explanation
"""


list(x + y for x, y in product(d6.elements(), d6.elements()))

"""
Explanation: Now we can compute the sum of all pairs:
End of explanation
"""


Hist(x + y for x, y in product(d6.elements(), d6.elements()))

"""
Explanation: And finally make a Hist of the sums:
End of explanation
"""


twice = d6 + d6
twice

"""
Explanation: But all of that is provided by __add__, which we can call using the + operator:
End of explanation
"""


twice.plot(color=COLORS[2])

"""
Explanation: Now we can plot the histogram of outcomes from rolling two dice:
End of explanation
"""


thrice = twice + d6
thrice

"""
Explanation: Or three dice:
End of explanation
"""


thrice.plot(color=COLORS[1])

"""
Explanation: Notice that this is looking more and more like a bell curve:
End of explanation
"""


import thinkstats2

dct_file = '2002FemPreg.dct'
dat_file = '2002FemPreg.dat.gz'

dct = thinkstats2.ReadStataDct(dct_file)
preg = dct.ReadFixedWidth(dat_file, compression='gzip')

preg.head()

"""
Explanation: As the number of dice increases, the result converges to a normal distribution, also known as a Gaussian distribution.
Are first babies more likely to be late?
This is one of the first topics I wrote about in my blog, and still the most popular, with more than 100,000 page views:
http://allendowney.blogspot.com/2011/02/are-first-babies-more-likely-to-be-late.html
I used data from the National Survey of Family Growth (NSFG):
End of explanation
"""


preg.outcome.value_counts().sort_index()

"""
Explanation: The variable outcome encodes the outcome of the pregnancy.  Outcome 1 is a live birth.
End of explanation
"""


preg.pregordr.value_counts().sort_index()

"""
Explanation: pregorder is 1 for first pregnancies, 2 for others.
End of explanation
"""


live = preg[preg.outcome == 1]
firsts = live[live.birthord == 1]
others = live[live.birthord != 1]
len(firsts), len(others)

"""
Explanation: I selected live births, then split into first babies and others.
End of explanation
"""


firsts.prglngth.mean(), others.prglngth.mean()

"""
Explanation: The mean pregnancy lengths are slightly different:
End of explanation
"""


diff = firsts.prglngth.mean() - others.prglngth.mean()
diff

"""
Explanation: The difference is 0.078 weeks:
End of explanation
"""


diff * 7 * 24

"""
Explanation: Which is 13 hours.  Note: the best units to report are often not the units you computed.
End of explanation
"""


first_hist = Hist(firsts.prglngth)
other_hist = Hist(others.prglngth)

"""
Explanation: Let's see if we can visualize the difference in the histograms:
End of explanation
"""


def plot_distributions(dist1, dist2):
    dist1.plot(width=-0.45, align='edge', color=COLORS[3], label='firsts')
    dist2.plot(width=0.45, align='edge', color=COLORS[4], label='others')
    plt.xlim(33.5, 43.5)
    plt.legend()

"""
Explanation: I used some plotting options to put two bar charts side-by-side:
End of explanation
"""


plot_distributions(first_hist, other_hist)

"""
Explanation: Here are the two histograms:
End of explanation
"""


import numpy as np

class Pmf(Hist):
    
    def normalize(self):
        total = sum(self.values())
        for element in self:
            self[element] /= total
        return self
    
    def plot_cumulative(self, **options):
        xs, ps = zip(*sorted(self.iteritems()))
        cs = np.cumsum(ps, dtype=np.float)
        cs /= cs[-1]
        plt.plot(xs, cs, **options)

"""
Explanation: Remember that the vertical axis is counts.  In this case, we are comparing counts with different totals, which might be misleading.
An alternative is to compute a probability mass function (PMF), which divides the counts by the totals, yielding a map from each element to its probability.
The probabilities are "normalized" to add up to 1.
End of explanation
"""


first_pmf = Pmf(firsts.prglngth).normalize()
other_pmf = Pmf(others.prglngth).normalize()

"""
Explanation: Now we can compare PMFs fairly.
End of explanation
"""


plot_distributions(first_pmf, other_pmf)

"""
Explanation: Using PMFs, we see that some of the difference at 39 weeks was an artifact of the different samples sizes:
End of explanation
"""


first_pmf.plot_cumulative(linewidth=4, color=COLORS[3], label='firsts')
other_pmf.plot_cumulative(linewidth=4, color=COLORS[4], label='others')
plt.xlim(23.5, 44.5)
plt.legend(loc='upper left')

"""
Explanation: Even so, it is not easy to compare PMFs.  One more alternative is the cumulative mass function (CMF), which shows, for each $t$, the total probability up to and including $t$.
End of explanation
"""


def iterate_words(filename):
    """Read lines from a file and split them into words."""
    for line in open(filename):
        for word in line.split():
            yield word.strip()

"""
Explanation: The CDFs are similar up to week 38.  After that, first babies are more likely to be born late.
Note: don't be afraid of thick lines.  Differences that are only visible with thin lines are unlikely to be real.
Word Frequencies
Next topic: let's look at histograms of words, bigrams and trigrams.
The following function reads lines from a file or URL and splits them into words:
End of explanation
"""


# FAIRY TALES
# By The Brothers Grimm
# http://www.gutenberg.org/cache/epub/2591/pg2591.txt'
wc = Hist(iterate_words('pg2591.txt'))

"""
Explanation: Here's an example using a book from Project Gutenberg.  wc is a histogram of word counts:
End of explanation
"""


wc.most_common(20)

"""
Explanation: Here are the 20 most common words:
End of explanation
"""


ranks, counts = wc.ranks()
plt.plot(ranks, counts, linewidth=10, color=COLORS[5])
plt.xlabel('rank')
plt.ylabel('count')

"""
Explanation: Word frequencies in natural languages follow a predictable pattern called Zipf's law (which is an instance of Stigler's law, which is also an instance of Stigler's law).
We can see the pattern by lining up the words in descending order of frequency and plotting their ranks (1st, 2nd, 3rd, ...) versus counts (6507, 5250, 2707):
End of explanation
"""


ranks, counts = wc.ranks()
plt.plot(ranks, counts, linewidth=4, color=COLORS[5])
plt.xlabel('rank')
plt.ylabel('count')
plt.xscale('log')
plt.yscale('log')

"""
Explanation: Huh.  Maybe that's not so clear after all.  The problem is that the counts drop off very quickly.  If we use the highest count to scale the figure, most of the other counts are indistinguishable from zero.
Also, there are more than 10,000 words, but most of them appear only a few times, so we are wasting most of the space in the figure in a regime where nothing is happening.
This kind of thing happens a lot.  A common way to deal with it is to compute the log of the quantities or to plot them on a log scale:
End of explanation
"""


from itertools import tee

def pairwise(iterator):
    """Iterates through a sequence in overlapping pairs.
    
    If the sequence is 1, 2, 3, the result is (1, 2), (2, 3), (3, 4), etc.
    """
    a, b = tee(iterator)
    next(b, None)
    return izip(a, b)

"""
Explanation: This (approximately) straight line is characteristic of Zipf's law.
n-grams
On to the next topic: bigrams and trigrams.
End of explanation
"""


bigrams = Hist(pairwise(iterate_words('pg2591.txt')))

"""
Explanation: bigrams is the histogram of word pairs:
End of explanation
"""


bigrams.most_common(20)

"""
Explanation: And here are the 20 most common:
End of explanation
"""


def triplewise(iterator):
    a, b, c = tee(iterator, 3)
    next(b)
    next(c)
    next(c)
    return izip(a, b, c)

"""
Explanation: Similarly, we can iterate the trigrams:
End of explanation
"""


trigrams = Hist(triplewise(iterate_words('pg2591.txt')))

# Uncomment this line to run the analysis with Elvis Presley lyrics
#trigrams = Hist(triplewise(iterate_words('lyrics-elvis-presley.txt')))

"""
Explanation: And make a histogram:
End of explanation
"""


trigrams.most_common(20)

"""
Explanation: Here are the 20 most common:
End of explanation
"""


from collections import defaultdict

d = defaultdict(Hist)
for a, b, c in trigrams:
    d[a, b][c] += trigrams[a, b, c]

"""
Explanation: And now for a little fun.  I'll make a dictionary that maps from each word pair to a Hist of the words that can follow.
End of explanation
"""


d['the', 'blood']

"""
Explanation: Now we can look up a pair and see what might come next:
End of explanation
"""


d['into', 'the'].most_common(10)

"""
Explanation: Here are the most common words that follow "into the":
End of explanation
"""


d['said', 'the'].most_common(10)

"""
Explanation: Here are the words that follow "said the":
End of explanation
"""


d['said', 'the'].choice()

"""
Explanation: Hist provides choice, which chooses a random word with probability proportional to count:
End of explanation
"""


prefix = 'said', 'the'
suffix = d[prefix].choice()
suffix

"""
Explanation: Given a prefix, we can choose a random suffix:
End of explanation
"""


prefix = prefix[1], suffix
prefix

"""
Explanation: Then we can shift the words and compute the next prefix:
End of explanation
"""


for i in range(100):
    suffix = d[prefix].choice()
    print(suffix, end=' ')
    prefix = prefix[1], suffix

"""
Explanation: Repeating this process, we can generate random new text that has the same correlation structure between words as the original:
End of explanation
"""


