secret_santa.py

from __future__ import print_function
from __future__ import division

import io
import random
import numpy as np
from PIL import Image

from skimage.transform import resize
import quantumrandom  # https://qrng.anu.edu.au/
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.image import MIMEImage
from pulp import LpProblem, LpMinimize, LpVariable, LpBinary, lpSum

from keras.models import Sequential
from keras.layers import Dense, LSTM, Conv2D, Activation, BatchNormalization
from keras.layers import UpSampling2D, Reshape
from keras.utils import get_file


people = [
{\
    'name' : 'Name1',\
    'email' : 'name1@email.com'},
{\
    'name' : 'Name2',\
    'email' : 'name2@email.com'},
{\
    'name' : 'Name3',\
    'email' : 'name3@email.com'}]


def build_generator(img_shape, noise_dim):

    # Generator
    generator = Sequential()

    s = img_shape[0] // 4
    nb_channels = img_shape[-1]

    generator.add(Dense(128 * s * s, input_dim=noise_dim,
                  kernel_initializer='glorot_uniform'))
    generator.add(BatchNormalization(momentum=0.9))
    generator.add(Activation('relu'))

    generator.add(Reshape((s, s, 128)))

    generator.add(UpSampling2D(size=(2, 2)))

    generator.add(Conv2D(16, kernel_size=(3, 3), padding='same',
                  kernel_initializer='glorot_uniform'))
    generator.add(BatchNormalization(momentum=0.9))
    generator.add(Activation('relu'))

    generator.add(Conv2D(16, kernel_size=(3, 3), padding='same',
                  kernel_initializer='glorot_uniform'))
    generator.add(BatchNormalization(momentum=0.9))
    generator.add(Activation('relu'))

    generator.add(UpSampling2D(size=(2, 2)))

    generator.add(Conv2D(32, kernel_size=(3, 3), padding='same',
                  kernel_initializer='glorot_uniform'))
    generator.add(BatchNormalization(momentum=0.9))
    generator.add(Activation('relu'))

    generator.add(Conv2D(nb_channels, kernel_size=(5, 5), padding='same',
                  activation='tanh', kernel_initializer='glorot_uniform'))

    path = get_file('generator.h5',
        origin='https://jcboyd.github.io/assets/secret-santa/generator.h5')
    generator.load_weights(path)

    return generator


def get_image(model):

    gen_noise = np.random.randn(1, 100)
    sample = model.predict(gen_noise)

    sample = resize(sample.squeeze(), output_shape=(256, 256, 3), order=0)
    sample = np.clip(127.5 * sample + 127.5, 0, 255).astype('uint8')

    img = Image.fromarray(sample, 'RGB')

    return img


# character ordering for RNN
chars = [u'\n', u' ', u'!', u'"', u'&', u"'", u'(', u')', u',',  u'-', u'.',
         u'0', u'1', u'2', u'3', u'4', u'5', u'7', u'8', u'9', u':', u';',
         u'>', u'?', u'a', u'b', u'c', u'd', u'e', u'f', u'g', u'h', u'i',
         u'j', u'k', u'l', u'm', u'n', u'o', u'p', u'q', u'r', u's', u't',
         u'u', u'v', u'w', u'x', u'y', u'z', u'\x92', u'\xe8', u'\xeb']


def build_rnn(max_length=40):

    model = Sequential()
    model.add(LSTM(128, input_shape=(max_length, len(chars))))
    model.add(Dense(len(chars), activation='softmax'))

    path = get_file('rnn_weights.h5',
        origin='https://jcboyd.github.io/assets/secret-santa/rnn_weights.h5')
    model.load_weights(path)
    return model


def get_verse(model, verse='wishing you a very merry cbio christmas\n'):

    max_lines = 4
    max_length = 40

    while verse.count('\n') < max_lines:

        x_pred = np.zeros((1, max_length, len(chars)))

        for t, char in enumerate(verse[-max_length:]):
            x_pred[0, t, chars.index(char)] = 1.

        # anneal probabilties - no alarms/surpises please!
        preds = model.predict(x_pred, verbose=0)[0].astype('float64')
        logits = np.exp(np.log(preds) / 0.02)
        probs = logits / np.sum(logits)

        idx = np.random.choice(probs.shape[0], p=probs)
        next_char = chars[idx]
        verse += next_char

    return verse


def send_mail(sender, receiver, verse, img):
    """
    Composes email for sender-receiver pair
    """
    msgRoot = MIMEMultipart('related')
    msgRoot['Subject'] = 'Christmas Party -- Secret Santa'
    msgRoot['From'] = 'Ninja Santa <ninja.santa@email.com>'
    msgRoot['To'] = '%s <%s>' % (sender['name'], sender['email'])
    msgRoot.preamble = 'This is a multi-part message in MIME format.'

    msgText = MIMEText("""<h1>Hi %s!</h1>
    <p>You are the Secret Santa of <b>%s!</b></p>
    <img src="cid:image1" style="width:128px;height:128px;">
    <p><i>%s</i></p>
    <p>This email was auto-generated from the cluster.
    Random shuffle seeded with quantum random data
    (<a href="https://qrng.anu.edu.au/">https://qrng.anu.edu.au/</a>).
    Secret santas assigned with integer programming.
    Verse composed by LSTM.
    Complete script available 
    <a href=https://gist.github.com/jcboyd/6716981917396b6ee7af28be2329a929>
    here.</a>
    And <a href=https://github.com/jcboyd/ganta-claus>this</a>
    is how GANta was trained.</p>
    """ % (sender['name'].split(" ")[0], receiver['name'], verse), 'html')

    msgRoot.attach(msgText)

    # attach image
    outbuf = io.BytesIO()
    img.save(outbuf, format='PNG')
    msgImage = MIMEImage(outbuf.getvalue())

    # Define the image's ID as referenced above
    msgImage.add_header('Content-ID', '<image1>')
    msgRoot.attach(msgImage)

    smtpObj = smtplib.SMTP(host='localhost', port=25)
    smtpObj.sendmail('ninja.santa@email.com', \
        sender['email'], \
        msgRoot.as_string())


def secret_santa_assignment(num_people):
    """
    Solves an integer programming problem to make sender-receiver assignments
    """
    p = LpProblem('p', LpMinimize)
    # objective is dummy variable
    p += 0

    # square matrix of sender-receiver variables
    vars = [[LpVariable('x_%d_%d' % (i, j), 0, 1, LpBinary)\
        for j in range(num_people)]\
        for i in range(num_people)]

    for i, sender in enumerate(vars):
        for j, receiever in enumerate(sender):
            # no one sends to themselves
            if i == j:
                p += receiever >= 0
                p += -receiever >= 0 

        # everyone sends to exactly one person
        p += lpSum(sender) == 1

        # everyone receives from exactly one person
        col = [vars[k][i] for k in range(num_people)]
        p += lpSum(col) == 1
    p.solve()

    assignment = np.array(
        [[vars[i][j].value()\
        for i in range(num_people)]\
        for j in range(num_people)])

    return assignment


if __name__ == '__main__':
    # Seed with quantum randomness--can't be too careful...
    # random.seed(int(quantumrandom.hex(), 16))
    # random.shuffle(people)

    # Create assignment
    num_people = len(people)
    assignment = secret_santa_assignment(num_people)

    # Assert that assignment is valid
    assert all(np.diag(assignment) == 0)
    assert all(assignment.sum(axis=0) == 1)
    assert all(assignment.sum(axis=1) == 1)

    rnn = build_rnn()
    generator = build_generator((32, 32, 3), noise_dim=100)
    
    for i, sender in enumerate(people):
        receiver = people[np.argmax(assignment[i])]
        try:
            verse = get_verse(rnn).replace('\n', '<br/>')
            img = get_image(generator)
            send_mail(sender, receiver, verse, img)
            print('Sent to %s...' % (sender['email']))
        except:
            print('Message failed to send')
            continue