%pylab --no-import-all inline
from __future__ import division
from collections import Counter
import matplotlib.pyplot as plt
import urllib
import itertools
import random 
import re

def Keyboard(rows):
    "A keyboard is a {letter: location} map, e.g. {'Q':Point(0, 0), 'A': Point(0.5, 1)}."
    return {ch: Point(x/2, y) 
            for (y, row) in enumerate(rows)
            for (x, ch) in enumerate(row)
            if ch != ' '}

class Point(complex):
    "A point in the (x, y) plane."
    def __repr__(self): return 'Point({}, {})'.format(self.x, self.y)
    x = property(lambda p: p.real)
    y = property(lambda p: p.imag)
    
def distance(A, B):
    "The distance between two points."
    return abs(A - B)

qwerty = Keyboard(('Q W E R T Y U I O P',
                   ' A S D F G H J K L ',
                   '   Z X C V B N M   '))
qwerty

W, O, R, D = qwerty['W'], qwerty['O'], qwerty['R'], qwerty['D'], 
distance(W, O) + distance(O, R) + distance(R, D)

def path_length(word, kbd=qwerty):
    "The total path length for a word on this keyboard: the sum of the segment lengths."
    return sum(distance(kbd[word[i]], kbd[word[i+1]])
               for i in range(len(word)-1))

path_length('WORD')

path_length('TO')

path_length('TYPEWRITER') == 1 + 4 + 7 + 1 + 2 + 4 + 3 + 2 + 1 == 25

WORDS = set(urllib.urlopen('http://norvig.com/ngrams/TWL06.txt').read().split())

len(WORDS)

max(WORDS, key=path_length)

def print_top(n, sequence, key=None, fmt='{:.1f} {}'):
    "Find the top n elements of sequence as ranked by key function, and print them."
    for x in sorted(sequence, key=key, reverse=True)[:n]:
        print(fmt.format(key(x), x))
              
print_top(10, WORDS, path_length)

def path_length_ratio(word, kbd=qwerty): return path_length(word, kbd) / len(word)

print_top(10, WORDS, path_length_ratio)

def on_top_row(word): return all(L in 'QWERTYUIOP' for L in word)

print_top(10, filter(on_top_row, WORDS), len)

def Workload(text):
    """Create a Workload--a dict of the form {'AB': 1000, ...} 
    saying how often each letter pair occurs in text."""
    pairs = filter(str.isalpha, bigrams(text))
    return normalize(Counter(segment(A, B) for (A, B) in pairs))

def bigrams(sequence):
    "Every sequence of two adjacent elements (overlapping) in sequence."
    return (sequence[i:i+2] for i in range(len(sequence) - 1))
    
def segment(A, B):
    "Canonical two-character segment: segment('A', 'B') == segment('B', 'A') == 'AB'."
    return A+B if (A < B) else B+A

def normalize(dictionary):
    "Normalize a {key: val} dict so that the sum of the vals is 1.0."
    total = sum(dictionary.values())
    for k in dictionary:
        dictionary[k] /= total
    return dictionary

Workload('SHOT IS GOOD -- GOOOOOOOOOOOAL!')

TEXT     = urllib.urlopen('http://norvig.com/ngrams/smaller.txt').read().upper()
WORKLOAD = Workload(TEXT)

WORKLOAD.most_common(10)

len(WORKLOAD)

def workload_average(kbd, workload=WORKLOAD):
    "The average segment length over a workload of segments."
    return sum(distance(kbd[A], kbd[B]) * workload[A+B]
               for (A, B) in workload)

workload_average(qwerty)

def plot_kbd(kbd, K=20):
    "Plot the keyboard with square keys, K units on a side."
    H = K / 2       ## (K is Key width/height; H is half K)
    # Draw a square for each key
    for L in kbd:
        x, y = K * kbd[L].x, -K * kbd[L].y
        plot_square(x, y, H, label=L)
    # Show the plot
    plt.axis('equal'); plt.axis('off'); plt.show()

def plot_square(x, y, H, label='', style='k-'):
    "Plot a square with center (x, y), half-width H, and optional label."
    plt.plot([x-H, x+H, x+H, x-H, x-H], 
             [y-H, y-H, y+H, y+H, y-H], style)  
    plt.annotate(label, (x-H/4, y-H/4)) # H/4 seems to place label well.
    
def show_kbd(kbd, name='keyboard'):
    "Plot a keyboard and print statistics about it."
    plot_kbd(kbd)    
    print('workload average = {:.1f} for {}'
          .format(workload_average(kbd), name))

show_kbd(qwerty)

def improve(kbd, swaps=1000, scorer=workload_average):
    "Minimize scorer(kbd) by swapping keys and keeping improvements."
    score = scorer(kbd)
    letters = list(kbd)
    for _ in range(swaps):
        A, B = random.sample(letters, 2)   # Step 1: pick two keys
        swap(kbd, A, B)                    # Step 2: swap them
        score2 = scorer(kbd)
        if score2 < score:                 # Step 3: If better, keep them
            score = score2                 # (and record the new best total)
        else:
            swap(kbd, B, A)                # Step 4: swap back if not better
    return kbd

def swap(kbd, A, B): kbd[A], kbd[B] = kbd[B], kbd[A]

def improved(kbd, swaps=1000, scorer=workload_average): 
    "Minimize scorer(kbd) by swapping keys and keeping improvements; don't modify kbd."
    return improve(kbd.copy(), swaps, scorer)

show_kbd(improved(qwerty, 1000))

def improve_scores(kbd, swaps=1000, scorer=workload_average):
    "A version of 'improve' that yields the workload average after each swap."
    total = scorer(kbd)
    letters = list(kbd)
    T = sum(WORKLOAD.values())             # <<< NEW
    for _ in range(swaps):
        A, B = random.sample(letters, 2)   # Step 1: pick two keys
        swap(kbd, A, B)                         # Step 2: swap them
        total2 = scorer(kbd)
        if total2 < total:                 # Step 3: If better, keep them
            total = total2                 # (and record the new best total)
        else:
            swap(kbd, B, A)                     # Step 4: swap back if not better
        yield total / T                    # <<< NEW
    
def workload_plot(kbd, swaps):
    plt.ylabel('Workload average segment length') 
    plt.xlabel('Number of swaps');  
    plt.plot(list(improve_scores(kbd.copy(), swaps)))                         

workload_plot(qwerty, 3000)

for run in range(10):
    workload_plot(qwerty, 3000)

def repeated_improved(kbd, repeats=20, swaps=1000, scorer=workload_average):
    "Try improved(kbd, swaps) multiple times and take the best result."
    return min([improved(kbd, swaps, scorer) for run in range(repeats)],
               key=scorer)

show_kbd(repeated_improved(qwerty, 1, 3000), "3000 swaps repeated once")
show_kbd(repeated_improved(qwerty, 3, 1000), "1000 swaps repeated three times")
show_kbd(repeated_improved(qwerty, 6, 500),  "500 swaps repeated six times")

%%time
show_kbd(repeated_improved(qwerty, 1, 3000), "3000 swaps repeated once")
show_kbd(repeated_improved(qwerty, 3, 1000), "1000 swaps repeated three times")
show_kbd(repeated_improved(qwerty, 6, 500),  "500 swaps repeated six times")

%%time 
show_kbd(repeated_improved(qwerty, 1, 30000), "30000 swaps repeated once")
show_kbd(repeated_improved(qwerty, 15, 2000), "2000 swaps repeated fifteen times")
show_kbd(repeated_improved(qwerty, 60, 500),  "500 swaps repeated sixty times")

keyboards = {
    'qwerty': Keyboard(('Q W E R T Y U I O P',
                        ' A S D F G H J K L ',  
                        '   Z X C V B N M   ')),

    '4-by-7': Keyboard((' A B C D E F ',
                        'G H I J K L M',
                        ' N O P Q R S ',
                        'T U V W X Y Z')),

    '5-by-6': Keyboard((' A B C D E ',
                        ' F G H I J ',
                        'K L M N O P',
                        ' Q R S T U ',
                        ' V W X Y Z '))
    }

# Compare keyboards
for name in keyboards:
    show_kbd(keyboards[name], name)

def report(keyboards, scorer=workload_average):
    "Iterate through a dict of {name: kbd} pairs, showing kbd before and after repeated_improved(kbd)."
    for (name, kbd) in keyboards.items():
        show_kbd(repeated_improved(kbd, scorer=scorer), 
                 'repeated improved ' + name)

%%time
report(keyboards)

cat = ''.join    ## Function to join strings together into one big string

# Order letters by frequency in TEXT
ordered_letters = cat(sorted(list(qwerty), reverse=True, key=lambda L: TEXT.count(L)))

def recentered(kbd):
    "Put the most frequent letters in the center of this keyboard."
    center = mean(kbd.values())
    ordered_locations = sorted(kbd.values(), key=lambda point: abs(point-center))
    return dict(zip(ordered_letters, ordered_locations))

def mean(numbers): return sum(numbers) / len(numbers)

ordered_letters

show_kbd(recentered(qwerty), 'recentered qwerty')

for _ in range(10):
    workload_plot(recentered(qwerty), 3000)

show_kbd(repeated_improved(recentered(qwerty)), 'repeated improved recentered qwerty')

def workload_lower_bound(kbd=recentered(qwerty), workload=WORKLOAD):
    "Given a keyboard, find a lower bound on the workload average."
    letters = list(kbd)
    def distances(L):
        "Distances to L key, shortest distance first."
        return sorted(abs(kbd[A] - kbd[L]) for A in letters if A != L)
    def counts(L):
        "Workload counts for segments involving L, biggest count first."
        return sorted([workload[segment(A, L)] for A in letters if A != L],
                      reverse=True)
    return (sum(dot_product(distances(L), counts(L)) for L in kbd)
            / sum(workload.values()) / 2) # Divide by 2 because we counted each segment twice

def dot_product(A, B):
    "Σ_i{A_i * B_i}"
    return sum(A[i] * B[i] for i in range(len(A)))

for name in keyboards:
    print name, workload_lower_bound(keyboards[name])

def neighboring_keys(kbd, radius=1.5):
    "Build a dict of {Letter:NeighboringLetters}, e.g. {'Q':'QAW', ...}."
    def neighboring_letters(A): 
        return cat(B for B in kbd if distance(kbd[A], kbd[B]) <= radius)
    return {A: neighboring_letters(A) for A in kbd}

def neighboring_keys(kbd, radius=1.5):
    "Build a dict of {Letter:NeighboringLetters}, e.g. {'Q':'QAW', ...}."
    def neighboring_letters(A): 
        return 
    return {A: cat(B for B in kbd if distance(kbd[A], kbd[B]) <= radius)
            for A in kbd}

cat = ''.join  ## Function to join letters (or strings) into one string

qwerty_neighbors = neighboring_keys(qwerty) 
qwerty_neighbors

choices = [qwerty_neighbors[L] for L in 'HELLO']
choices

paths = {cat(letters) for letters in itertools.product(*choices)}

len(paths)

random.sample(paths, 8)

WORDS & paths

def confusions(word, neighbors=neighboring_keys(qwerty), words=WORDS):
    "All valid words whose paths could be confused with the path for the given word."
    choices = [neighbors[L] for L in word]
    return {cat(letters) for letters in itertools.product(*choices)} & words

confusions('HELLO')

confusions('WORLD')

confusions('TESTING')

%time confusions('SOMETHING')

[len(neighboring_keys(qwerty)[L]) for L in 'SOMETHING']

def prefixes(words):
    "Return a set of prefixes (1 to N characters) of this collection of words."
    return {word[:i] for word in words for i in range(1, len(word)+1)}

prefixes(['THESE', 'THEY', 'THOSE'])

PREFIXES = prefixes(WORDS)

len(PREFIXES), len(WORDS)

def confusions(word, neighbors=neighboring_keys(qwerty), words=WORDS, prefixes=PREFIXES):
    "All valid words whose paths could be confused with the path for this word."
    results = set() # A set of words that are confusions of 'word'
    Q = ['']        # A queue of partial paths 
    while Q:
        path = Q.pop()
        if len(path) < len(word):
            for L in neighbors[word[len(path)]]:
                if path + L in prefixes:
                    Q.append(path + L)
        elif path in words:
            results.add(path)
    return results

confusions('HELLO')

%time confusions('SOMETHING')

def words(text): 
    "Extract a list of all words from a text. Make everything uppercase."
    return re.findall('[A-Z]+', text.upper())

test_words = words("""Hello world! Testing something: confusion paths on 
multiple distinct inputs of various lengths. See if everything works 
perfect, or poorly, or somewhere between the extremes.""")

def test_confusions(words=test_words):
    "Run through some test words, printing and summarizing the confusions."
    total = 0
    for word in sorted(set(words)):
        others = sorted(set(confusions(word)) - {word})
        total += len(others)
        print('{}({}): {}'.format(word, len(others), ' '.join(others)))
    print 'Total of {} confusions for {} words'.format(total, len(words))

%time test_confusions()

def plot_kbd(kbd, K=20, words=()):
    "Plot the keyboard with square keys, K units on a side."
    H = K / 2       ## (K is Key width/height; H is half K)
    # Draw a square for each key, and plot paths for words
    for L in kbd:
        x, y = K * kbd[L].x, -K * kbd[L].y
        plot_square(x, y, H, label=L)
    plot_paths(kbd, K, words)
    # Show the plot and print the workload average
    plt.axis('equal'); plt.axis('off'); plt.show()
    
def plot_paths(kbd, K, words):
    "Plot paths for each word."
    for (i, word) in enumerate(words):
        Xs = [+K * kbd[L].x for L in word]
        Ys = [-K * kbd[L].y for L in word]
        plt.plot(Xs, Ys, '-o')

plot_kbd(qwerty, words=['HELLO', 'JELLO'])

def plot_paths(kbd, K, words):
    "Plot paths for each word, each with a different offset (and color)."
    Q = K / 4  ## Q is a quarter of a key width
    offsets = [Point(-Q, -Q), Point(-Q, +Q),  Point(Q, +Q), Point(Q, -Q)]
    for (i, word) in enumerate(words):
        Xs = [+K * kbd[L].x + offsets[i % 4].x for L in word]
        Ys = [-K * kbd[L].y + offsets[i % 4].y for L in word]
        plt.plot(Xs, Ys, '-o')

plot_kbd(qwerty, words=['HELLO', 'JELLO'])

plot_kbd(qwerty, words=confusions('HELLO'))

confusions('JELLO')

plot_kbd(qwerty, words=['JELLO', 'BROOK'])

plot_kbd(qwerty, words=['JE', 'BR'])

JE = qwerty['E'] - qwerty['J']
BR = qwerty['R'] - qwerty['B']
distance(JE, BR)

def similar_segments(kbd, path, word):
    "Do the last two letters of path form a similar segment to corresponding letters of word?"
    n = len(path)
    if n < 2: return True
    # Define endpoints (P1, P2), (W1, W2); then vectors P and W; then difference between.
    (P1, P2), (W1, W2) = (path[n-2:n], word[n-2:n])
    P = kbd[P2] - kbd[P1]
    W = kbd[W2] - kbd[W1]
    return distance(P, W) < 2

similar_segments(qwerty, 'BR', 'JELLO')

similar_segments(qwerty, 'HE', 'JELLO')

similar_segments(qwerty, 'HEP', 'JELLO')

similar_segments(qwerty, 'J', 'JELLO')

def confusions(word, neighbors=neighboring_keys(qwerty), 
               words=WORDS, prefixes=PREFIXES, kbd=qwerty):
    "All valid words whose paths could be confused with the path for this word."
    results = set() # A set of words that are confusions of 'word'
    Q = ['']        # A queue of partial paths 
    while Q:
        path = Q.pop()
        if len(path) < len(word):
            for L in neighbors[word[len(path)]]:
                newpath = path + L
                if (newpath in prefixes) and similar_segments(kbd, newpath, word):
                    Q.append(newpath)
        elif path in words:
            results.add(path)
    return results

test_confusions()

plot_kbd(qwerty, words=['SEE', 'DEE'])

similar_segments(qwerty, 'DE', 'SEE')

def similar_segments(kbd, path, word):
    "Do the last two letters of path form a similar segment to corresponding letters of word?"
    n = len(path)
    if n < 2: return True
    # Define endpoints (P1, P2), (W1, W2); then vectors P and W; then difference between.
    (P1, P2), (W1, W2) = (path[n-2:n], word[n-2:n])
    P = Point(kbd[P2] - kbd[P1])
    W = Point(kbd[W2] - kbd[W1])
    return (distance(P, W) < 2) and (P.x * W.x >= 0) and (P.y * W.y >= 0)

test_confusions()

WORDLOAD = normalize(Counter(words(TEXT)))

len(WORDLOAD)

WORDLOAD.most_common(10)

WORDLOAD_PREFIXES = prefixes(WORDLOAD)

CONFUSING = [w for w in WORDLOAD 
             if len(confusions(w, words=WORDLOAD, prefixes=WORDLOAD_PREFIXES)) > 1]

print(len(CONFUSING))
print(len(WORDLOAD))
print(len(CONFUSING) / len(WORDLOAD))

sum(WORDLOAD[w] for w in WORDLOAD 
    if len(confusions(w, words=WORDLOAD, prefixes=WORDLOAD_PREFIXES)) > 1)

def confusingness(kbd, wordload=WORDLOAD, prefixes=WORDLOAD_PREFIXES):
    "The proportion of words in wordload that are confused with other words."
    neighbors = neighboring_keys(kbd)
    return sum(WORDLOAD[w] for w in WORDLOAD 
               if len(confusions(w, neighbors, wordload, prefixes)) > 1)

def show_kbd(kbd, name='keyboard'):
    "Plot a keyboard and print statistics about it."
    plot_kbd(kbd)   
    print('workload average = {:.1f}; confusingness = {:.0%} for {}'
          .format(workload_average(kbd), confusingness(kbd), name))

for name in keyboards:
    show_kbd(keyboards[name], name)

report(keyboards)

%time confusingness(qwerty)

%time workload_average(qwerty)

%time kbd = improved(qwerty, swaps=20, scorer=confusingness)
show_kbd(kbd)

def combined_scorer(kbd):
    "The average of workload average and 6 * confusingness."
    return mean([workload_average(kbd), 6 * confusingness(kbd)])

def show_kbd(kbd, name='keyboard'):
    "Plot a keyboard and print statistics about it."
    plot_kbd(kbd)   
    W = workload_average(kbd)
    C = confusingness(kbd)
    print('workload average = {:.1f}; confusingness = {:.0%}; combined = {:.1f} for {}'
          .format(W, C, mean([W, 6 * C]), name))

for name in keyboards:
    show_kbd(keyboards[name], name)

%time kbd = improved(keyboards['4-by-7'], swaps=30, scorer=combined_scorer)
show_kbd(kbd)

show_kbd(Keyboard((' A B E D Q F ',
                   'G U M J I L H',
                   ' N O P W C S ',
                   'T X R V Y K Z')))

plot_kbd(qwerty, words=['PUTS', 'POTS', 'PITS', 'POUTS', 'PUTTS'])