<commit_before>"""
HOW THIS WORKS:
1. Message objects are created via /dashboard/messages. Message objects have a string indicating their audience (e.g. all registered hackers, all checked-in, mailing list people, teammates listed on people's applications)
2. Message.enqueue_tasks() enqueues a bunch of Tas objects in the `messages` queue, by:
 2.1 calling Message.get_query() to get a datastore query
 2.2 batching the entities returned from the query into groups of 40 (batch_size)
 2.3 creating Task Queue tasks for each batch, containing message ID and entity ID for all entities
3. MessagesTaskQueueWork.post is called by the task queue for each task.
4. MessagesTaskQueueWork.post loads the Message object and all entities, then concurrently calls Message.send_to_entity_async(entity) on each entity. These are responsible for looking the audience, the current entity (which may be a Hacker or an EmailListEntry, or maybe something else in the future), and sending the appropriate emails or SMSs
"""

import webapp2
from template import template, template_string
from google.appengine.ext import ndb
from send_email import send_email
import re
from google.appengine.api import taskqueue
from registration import Hacker
from email_list import EmailListEntry
import sms
from config import isAdmin

class Message(ndb.Model):
	added = ndb.DateTimeProperty(auto_now_add=True)

	audience = ndb.StringProperty(choices=[None, 'registered', 'invited-friends', 'mailing-list-unregistered', 'waitlisted'], default=None)

	email_from_template = ndb.BooleanProperty(default=False)
	email_subject = ndb.TextProperty()
	email_html = ndb.TextProperty()

	sms_text = ndb.TextProperty()

	def kick_off(self):
		q = taskqueue.Queue("messages")
		q.add(taskqueue.Task(url='/dashboard/messages/message_task_queue_work', params={"kickoff_message_key": self.key.urlsafe()}))

	def send_to_email(self, email, template_args={}):
		# does the actual work of sending
		emails = [email]
		assert self.email_subject, "No email subject provided. Is email unchecked?"
		# if self.audience == 'invited-friends':
		# 	hacker = Hacker.query(Hacker.email==email).fetch()[0]
		# 	if hacker.teammates:
		# 		emails_found = [email.lower() for email in hacker.teammates.split(',')]
		# 		matching_hackers =  Hacker.query(Hacker.email.IN(emails_found)).fetch()
		# 		emails_already_registered = [h.email for h in matching_hackers]
		# 		emails = []
		# 		for email in emails_found:
		# 			if email not in emails_already_registered:
		# 				emails.append(email)
		# 				template_args["invited_by"] = hacker
		if self.email_from_template:
			html = template("emails/" + self.email_html + ".html", template_args)
		else:
			html = template_string(self.email_html, template_args)
			html = template("emails/generic.html", dict({"content": html}.items() + template_args.items()))
		subject = template_string(self.email_subject, template_args)
		send_email(html, subject, emails)

	def send_to_phone(self, phone):
		# actual work of sms'ing
		print "Sending SMS '{0}' to {1}".format(self.sms_text, phone)
		sms.send(phone, self.sms_text)

	def get_query(self):
		if self.audience == 'registered':
			return Hacker.query()
		elif self.audience == 'invited-friends':
			return Hacker.query(Hacker.teammates != None)
		elif self.audience == 'mailing-list-unregistered':
			return EmailListEntry.query()
		elif self.audience == 'waitlisted':
			return Hacker.query(Hacker.admitted_email_sent_date == None and acker.waitlist_email_sent_date == None)
		elif self.audience == None:
			return None
		else:
			assert 0, "Unknown audience"

	def enqueue_tasks(self):
		if self.get_query() == None:
			print "No audience selected. No emails sent through messages system."
			return
		q = taskqueue.Queue("messages")
		# max task size is 100kb; max # of tasks added per batch is 100
		task_futures = []

		batch_size = 40 # 2k users / 40 = 50 tasks. within our limits.
		batch_of_entity_keys = []
		def send_batch():
			params = {"message_key": self.key.urlsafe(), "entity_keys": ','.join(batch_of_entity_keys[:])}
			task = taskqueue.Task(params=params, url='/dashboard/messages/message_task_queue_work')
			task_future = q.add_async(task)
			task_futures.append(task_future)
			batch_of_entity_keys[:] = []

		for key in self.get_query().iter(keys_only=True):
			batch_of_entity_keys.append(key.urlsafe())
			if len(batch_of_entity_keys) >= batch_size: send_batch()
		if len(batch_of_entity_keys): send_batch()

		for f in task_futures:
			f.get_result()

	@ndb.tasklet
	def send_to_entity_async(self, entity):
		try:
			if self.audience == 'invited-friends':
				# don't actually send to the hacker -- send to their friends
				hacker = entity
				if hacker.teammates:
					emails = [email.lower() for email in hacker.teammates.split(',')]
					matching_hackers = yield Hacker.query(Hacker.email.IN(emails)).fetch_async()
					emails_already_registered = [h.email for h in matching_hackers]
					for email in emails:
						if email not in emails_already_registered:
							self.send_to_email(email, {"invited_by": hacker})
			elif self.audience == 'registered': # send emails directly to hackers
				hacker = entity
				if hacker.email and self.email_subject:
					self.send_to_email(hacker.email, {"hacker": hacker})
				if hacker.phone_number and self.sms_text:
					self.send_to_phone(hacker.phone_number)
			elif self.audience == 'mailing-list-unregistered':
				email = entity.email
				is_registered = (yield Hacker.query(Hacker.email == email).count_async()) > 0
				if not is_registered:
					self.send_to_email(email, {})
		except Exception as e:
			print "Failed to send email '{0}' to '{1} because {2}'".format(self.email_subject, entity.email, e)

class MessagesDashboardHandler(webapp2.RequestHandler):
	def get(self):
		if not isAdmin(): return self.redirect('/')
		self.response.write(template("messages_dashboard.html", {}))

	def post(self):
		if not isAdmin(): return self.redirect('/')
		audience = None if self.request.get('audience') == '' else self.request.get('audience')
		message = Message(audience=audience)
		if self.request.get('email'):
			message.email_subject = self.request.get('email-subject')
			if self.request.get('email-name'):
				message.email_from_template = True
				message.email_html = self.request.get('email-name')
			else:
				message.email_html = self.request.get('email-html')
		if self.request.get('sms'):
			message.sms_text = self.request.get('sms-text')
		if self.request.get('test'):
			recip = self.request.get('test-recipient')
			if '@' in recip:
				message.send_to_email(recip)
				self.response.write("Sent email")
			elif len(recip) > 0:
				message.send_to_phone(recip)
				self.response.write("Sent SMS")
		else:
			message.put()
			message.kick_off()
			self.response.write(template("messages_dashboard.html", {"status": "Sent!"}))

class MessagesTaskQueueWork(webapp2.RequestHandler):
	def post(self):
		if not isAdmin(): return self.redirect('/')
		if self.request.get('kickoff_message_key'):
			ndb.Key(urlsafe=self.request.get('kickoff_message_key')).get().enqueue_tasks()
		elif self.request.get('entity_keys'):
			keys = [ndb.Key(urlsafe=self.request.get('message_key'))] + map(lambda k: ndb.Key(urlsafe=k), self.request.get('entity_keys').split(','))
			msg_and_entities = ndb.get_multi(keys)
			message = msg_and_entities[0]
			entities = msg_and_entities[1:]
			futures = [message.send_to_entity_async(entity) for entity in entities]
			for f in futures: f.wait()
<commit_msg>changes to admissions<commit_after>"""
HOW THIS WORKS:
1. Message objects are created via /dashboard/messages. Message objects have a string indicating their audience (e.g. all registered hackers, all checked-in, mailing list people, teammates listed on people's applications)
2. Message.enqueue_tasks() enqueues a bunch of Tas objects in the `messages` queue, by:
 2.1 calling Message.get_query() to get a datastore query
 2.2 batching the entities returned from the query into groups of 40 (batch_size)
 2.3 creating Task Queue tasks for each batch, containing message ID and entity ID for all entities
3. MessagesTaskQueueWork.post is called by the task queue for each task.
4. MessagesTaskQueueWork.post loads the Message object and all entities, then concurrently calls Message.send_to_entity_async(entity) on each entity. These are responsible for looking the audience, the current entity (which may be a Hacker or an EmailListEntry, or maybe something else in the future), and sending the appropriate emails or SMSs
"""

import webapp2
from template import template, template_string
from google.appengine.ext import ndb
from send_email import send_email
import re
from google.appengine.api import taskqueue
from registration import Hacker
from email_list import EmailListEntry
import sms
from config import isAdmin

class Message(ndb.Model):
	added = ndb.DateTimeProperty(auto_now_add=True)

	audience = ndb.StringProperty(choices=[None, 'registered', 'invited-friends', 'mailing-list-unregistered', 'waitlisted'], default=None)

	email_from_template = ndb.BooleanProperty(default=False)
	email_subject = ndb.TextProperty()
	email_html = ndb.TextProperty()

	sms_text = ndb.TextProperty()

	def kick_off(self):
		q = taskqueue.Queue("messages")
		q.add(taskqueue.Task(url='/dashboard/messages/message_task_queue_work', params={"kickoff_message_key": self.key.urlsafe()}))

	def send_to_email(self, email, template_args={}):
		# does the actual work of sending
		emails = [email]
		assert self.email_subject, "No email subject provided. Is email unchecked?"
		# if self.audience == 'invited-friends':
		# 	hacker = Hacker.query(Hacker.email==email).fetch()[0]
		# 	if hacker.teammates:
		# 		emails_found = [email.lower() for email in hacker.teammates.split(',')]
		# 		matching_hackers =  Hacker.query(Hacker.email.IN(emails_found)).fetch()
		# 		emails_already_registered = [h.email for h in matching_hackers]
		# 		emails = []
		# 		for email in emails_found:
		# 			if email not in emails_already_registered:
		# 				emails.append(email)
		# 				template_args["invited_by"] = hacker
		if self.email_from_template:
			html = template("emails/" + self.email_html + ".html", template_args)
		else:
			html = template_string(self.email_html, template_args)
			html = template("emails/generic.html", dict({"content": html}.items() + template_args.items()))
		subject = template_string(self.email_subject, template_args)
		send_email(html, subject, emails)

	def send_to_phone(self, phone):
		# actual work of sms'ing
		print "Sending SMS '{0}' to {1}".format(self.sms_text, phone)
		sms.send(phone, self.sms_text)

	def get_query(self):
		if self.audience == 'registered':
			return Hacker.query()
		elif self.audience == 'invited-friends':
			return Hacker.query(Hacker.teammates != None)
		elif self.audience == 'mailing-list-unregistered':
			return EmailListEntry.query()
		elif self.audience == 'waitlisted':
			return Hacker.query(Hacker.admitted_email_sent_date == None and acker.waitlist_email_sent_date != None)
		elif self.audience == None:
			return None
		else:
			assert 0, "Unknown audience"

	def enqueue_tasks(self):
		if self.get_query() == None:
			print "No audience selected. No emails sent through messages system."
			return
		q = taskqueue.Queue("messages")
		# max task size is 100kb; max # of tasks added per batch is 100
		task_futures = []

		batch_size = 40 # 2k users / 40 = 50 tasks. within our limits.
		batch_of_entity_keys = []
		def send_batch():
			params = {"message_key": self.key.urlsafe(), "entity_keys": ','.join(batch_of_entity_keys[:])}
			task = taskqueue.Task(params=params, url='/dashboard/messages/message_task_queue_work')
			task_future = q.add_async(task)
			task_futures.append(task_future)
			batch_of_entity_keys[:] = []

		for key in self.get_query().iter(keys_only=True):
			batch_of_entity_keys.append(key.urlsafe())
			if len(batch_of_entity_keys) >= batch_size: send_batch()
		if len(batch_of_entity_keys): send_batch()

		for f in task_futures:
			f.get_result()

	@ndb.tasklet
	def send_to_entity_async(self, entity):
		try:
			if self.audience == 'invited-friends':
				# don't actually send to the hacker -- send to their friends
				hacker = entity
				if hacker.teammates:
					emails = [email.lower() for email in hacker.teammates.split(',')]
					matching_hackers = yield Hacker.query(Hacker.email.IN(emails)).fetch_async()
					emails_already_registered = [h.email for h in matching_hackers]
					for email in emails:
						if email not in emails_already_registered:
							self.send_to_email(email, {"invited_by": hacker})
			elif self.audience == 'registered': # send emails directly to hackers
				hacker = entity
				if hacker.email and self.email_subject:
					self.send_to_email(hacker.email, {"hacker": hacker})
				if hacker.phone_number and self.sms_text:
					self.send_to_phone(hacker.phone_number)
			elif self.audience == 'mailing-list-unregistered':
				email = entity.email
				is_registered = (yield Hacker.query(Hacker.email == email).count_async()) > 0
				if not is_registered:
					self.send_to_email(email, {})
		except Exception as e:
			print "Failed to send email '{0}' to '{1} because {2}'".format(self.email_subject, entity.email, e)

class MessagesDashboardHandler(webapp2.RequestHandler):
	def get(self):
		if not isAdmin(): return self.redirect('/')
		self.response.write(template("messages_dashboard.html", {}))

	def post(self):
		if not isAdmin(): return self.redirect('/')
		audience = None if self.request.get('audience') == '' else self.request.get('audience')
		message = Message(audience=audience)
		if self.request.get('email'):
			message.email_subject = self.request.get('email-subject')
			if self.request.get('email-name'):
				message.email_from_template = True
				message.email_html = self.request.get('email-name')
			else:
				message.email_html = self.request.get('email-html')
		if self.request.get('sms'):
			message.sms_text = self.request.get('sms-text')
		if self.request.get('test'):
			recip = self.request.get('test-recipient')
			if '@' in recip:
				message.send_to_email(recip)
				self.response.write("Sent email")
			elif len(recip) > 0:
				message.send_to_phone(recip)
				self.response.write("Sent SMS")
		else:
			message.put()
			message.kick_off()
			self.response.write(template("messages_dashboard.html", {"status": "Sent!"}))

class MessagesTaskQueueWork(webapp2.RequestHandler):
	def post(self):
		if not isAdmin(): return self.redirect('/')
		if self.request.get('kickoff_message_key'):
			ndb.Key(urlsafe=self.request.get('kickoff_message_key')).get().enqueue_tasks()
		elif self.request.get('entity_keys'):
			keys = [ndb.Key(urlsafe=self.request.get('message_key'))] + map(lambda k: ndb.Key(urlsafe=k), self.request.get('entity_keys').split(','))
			msg_and_entities = ndb.get_multi(keys)
			message = msg_and_entities[0]
			entities = msg_and_entities[1:]
			futures = [message.send_to_entity_async(entity) for entity in entities]
			for f in futures: f.wait()
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""Machine learning model which uses Wikipedia data to predicts ILI levels.

Usage:
  model.py <year_start> <year_end> <dataset_path> <incidence_path> <keywords_file> [--p] [--f] [--v] [--no-future]

  <year_start>      The first influenza season we want to predict.
  <year_end>        The last influenza season we want to predict.
  <dataset_path>    The path to dataset files (the ones which contain Wikipedia pageviews).
  <incidence_path>  The path to the files which specify the ILI levels for each week.
  <keywords_file>   The path to the file which contains the Wikipedia pages used.
  -p, --poisson     Use the Poisson model + LASSO instead of the linear one.
  -f, --file        Write informations to file
  -v, --verbose     Output more informations
  -n, --no-future   Use a different method to train the model (avoid using seasonal influnza which are
                    later than the one we want to predict)
  -h, --help        Print this help message
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import csv
from sklearn.linear_model import LassoCV
from sklearn.metrics import mean_squared_error
from tabulate import tabulate
from docopt import docopt

from models_utils import generate, generate_labels, generate_labels_one_year, generate_one_year, standardize_data, generate_keywords, stz, get_important_pages

from cvglmnetCoef import cvglmnetCoef
from cvglmnetPredict import cvglmnetPredict
from cvglmnet import cvglmnet

#########################
#### INITIALIZATION #####
#########################

# Parse the command line
arguments = docopt(__doc__)

# Year selected
year_start = int(arguments["<year_start>"])
year_end = int(arguments["<year_end>"])+1

# Model type
model_type = "Linear Model"
using_poisson = False
if arguments["--p"]:
    using_poisson = True
    model_type="Poisson Model"

# Feature path and labels
path_features = arguments["<dataset_path>"]
path_labels = arguments["<incidence_path>"]
keywords = arguments["<keywords_file>"]

# Selected columns that will be extracted from the dataframes
selected_columns = generate_keywords(keywords)

all_features_values = pd.DataFrame()
all_predicted_values = []
all_true_labels = pd.DataFrame()
total_weeks = []
all_weighted_feature=dict()

# The pair start-end seasons
year_sel = (year_start, year_end)

######################
##### ALGORITHM ######
######################

# For each season, train a model with the complete dataset (without the season)
# and predict the ILI incidence.
for year_selected in range(year_sel[0], year_sel[1]):

    print("------------")
    print("[*] ", year_selected )

    # Generate the list of excluded years
    excluded=[]
    if arguments["--no-future"]:
        for i in range(2007, 2018):
            if i>year_selected:
                excluded.append(i)

    # Data generation from data files
    dataset = generate(year_selected, excluded, path_features)[selected_columns]
    labels = generate_labels(year_selected, excluded, path_labels)
    data = generate_one_year(year_selected, path_features)[selected_columns]
    labels_test = generate_labels_one_year(year_selected, path_labels)
    weeks = labels_test["week"]

    # Create a copy of the dataset where we fills NaN values
    # with 0.
    dataset_zero = dataset.fillna(0)
    data_zero = data.fillna(0)

    # Standardize data
    train, test = standardize_data(dataset_zero, data_zero)

    # Create a Lasso Cross-Validation instance which will be
    # trained on the dataset in which NaN values are replaced
    # with the column mean.
    if not using_poisson:
        model = LassoCV(max_iter=100000, n_jobs=-1, n_alphas=1000, random_state=1)
        model.fit(train, labels["incidence"].fillna(0))
        result = model.predict(test)
    else:
        model = cvglmnet(x=train.as_matrix().copy(), y=labels["incidence"].fillna(0).copy().as_matrix(), family='poisson', alpha=1.0,
                       ptype="mse", parallel=True, nfolds=10)
        result = cvglmnetPredict(model, test.as_matrix(), ptype='response', s="lambda_min")
        coeff = cvglmnetCoef(model, s="lambda_min")

    # Get the feature coefficients
    if not using_poisson:
        coeff = model.coef_
    else:
        coeff = cvglmnetCoef(model, s="lambda_min")

    # Add the pair (value, coeff) to a dictionary
    for i in list(zip(coeff, selected_columns)):
        if (i[0] != 0):
            if all_weighted_feature.get(i[1], None) is None:
                all_weighted_feature[i[1]] = [i[0]]
            else:
                all_weighted_feature[i[1]].append(i[0])

    # Add the value to global variables
    all_features_values= all_features_values.append(data.fillna(0))
    all_predicted_values.extend(result)
    all_true_labels = all_true_labels.append(labels_test)
    total_weeks.extend(weeks)

    # If we are verbose print more informations
    if arguments["--v"]:
        index_peak = labels_test['incidence'].fillna(0).idxmax(axis=1)
        index_peak_m = np.argmax(result)
        print("[*] Influenza Season Peak (Week):", labels_test.iloc[index_peak]["week"])
        print("[*] Influenza Seasons Predicted Peak (Week): ", labels_test.iloc[index_peak_m]["week"])
        print("[*] Influenza Season Peak Value: ", labels_test.iloc[index_peak]["incidence"])
        print("[*] Influenza Season Predicted Peak Value: ", result[index_peak_m])
        print("------------")

# Get important pages to generate the plot we need
important_pages = get_important_pages(all_weighted_feature, 5)

# Print MSE of the two models
print("------------")
print("MSE: ", mean_squared_error(all_true_labels["incidence"].fillna(0), all_predicted_values))
print("------------")

# Print Pearson Coefficient
print("------------")
print("Pearson Correlation Coeff: ", np.corrcoef(all_predicted_values, all_true_labels["incidence"].fillna(0), rowvar=False)[0][1])
print("------------")

# Print important pages
print(tabulate(important_pages, headers=["Page name", "Weight"]))

###########################
#### GRAPHS GENERATION ####
###########################

# Plot some informations
fig = plt.figure(3, figsize=(15, 6))
ax = fig.add_subplot(111)

# Set up the graph title, x and y axis titles
plt.title("Influenza Seasons "+str(year_sel[0]-1)+" - "+str(year_sel[1]-1), fontsize=18)
plt.ylabel("Incidence on 100000 people", fontsize=17)
plt.xlabel("Year-Week", fontsize=17)

# Generate the x axis labels in such way to
# have the year-week pair only every two ticks.
weeks_used = []
for k, v in enumerate(total_weeks):
    if k%2==0:
        weeks_used.append(v)
    else:
        weeks_used.append(" ")

# Set up the axes ticks
plt.xticks(range(0, len(weeks_used)), weeks_used, rotation="vertical", fontsize=15)
plt.yticks(fontsize=15)

# Plot the model result and the incidence
plt.plot(range(0, len(all_predicted_values)), all_predicted_values, 'r-', label=model_type, linewidth=3)
plt.plot(range(0, len(all_true_labels["incidence"])), all_true_labels["incidence"].fillna(0), 'k-', label="Incidence", linewidth=3)

# Add dotted line into the graph to delimit the influenza seasons.
for i in range(1, (year_sel[1]-year_sel[0])):
    plt.axvline(26*i, linestyle="dashed", color="k", linewidth=3)

# Update some graphical options
plt.grid()
plt.legend(fontsize=16, loc="upper right")
plt.tight_layout()

# Save the graph on file
plt.savefig("season_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+".png", dpi=150)

# Feature plot
figf = plt.figure(4, figsize=(15, 6))
axf = fig.add_subplot(111)

# Generate the graph showing the pageview variation of the top 5 Wikipedia's pages
# choosen by the model.
plt.title("Pageview variation "+str(year_sel[0]-1)+" - "+str(year_sel[1]-1), fontsize=18)
plt.ylabel("Pageview value", fontsize=17)
plt.xlabel("Year-Week", fontsize=17)

# Plot the axes labels
plt.xticks(range(0, len(weeks_used)), weeks_used, rotation="vertical", fontsize=15)
plt.yticks(fontsize=15)

# Plot all the pageview data and the ILI incidence
std_all_features_values = stz(all_features_values)
for key, value in important_pages:
 plt.plot(range(0, len(std_all_features_values[key])), std_all_features_values[key], '-', label=key, linewidth=3)
plt.plot(range(0, len(all_true_labels["incidence"])), stz(all_true_labels["incidence"]), 'k-', label="Incidence", linewidth=3)

# Add dotted line into the graph to delimit the influenza seasons.
for i in range(1, (year_sel[1]-year_sel[0])):
    plt.axvline(26*i, linestyle="dashed", color="k", linewidth=3)

# Update some graphical options
plt.grid()
plt.legend(fontsize=16, loc="best")
plt.tight_layout()

# Save again the graph on file
plt.savefig("season_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_features.png", dpi=150)

# If --f then we write some information to file
if arguments["--f"]:
    with open(str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_information.csv", 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile, delimiter=',',
                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
        spamwriter.writerow(['mse', mean_squared_error(all_true_labels["incidence"].fillna(0), all_predicted_values)])
        for p in important_pages:
            spamwriter.writerow([p[0], float(p[1])])<commit_msg>[model] Small refactoring of evaluation results.<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""Machine learning model which uses Wikipedia data to predicts ILI levels.

Usage:
  model.py <year_start> <year_end> <dataset_path> <incidence_path> <keywords_file> [--p] [--f] [--v] [--no-future]

  <year_start>      The first influenza season we want to predict.
  <year_end>        The last influenza season we want to predict.
  <dataset_path>    The path to dataset files (the ones which contain Wikipedia pageviews).
  <incidence_path>  The path to the files which specify the ILI levels for each week.
  <keywords_file>   The path to the file which contains the Wikipedia pages used.
  -p, --poisson     Use the Poisson model + LASSO instead of the linear one.
  -f, --file        Write informations to file
  -v, --verbose     Output more informations
  -n, --no-future   Use a different method to train the model (avoid using seasonal influnza which are
                    later than the one we want to predict)
  -h, --help        Print this help message
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import csv
from sklearn.linear_model import LassoCV
from sklearn.metrics import mean_squared_error
from tabulate import tabulate
from docopt import docopt

from models_utils import generate, generate_labels, generate_labels_one_year, generate_one_year, standardize_data, generate_keywords, stz, get_important_pages

from cvglmnetCoef import cvglmnetCoef
from cvglmnetPredict import cvglmnetPredict
from cvglmnet import cvglmnet

#########################
#### INITIALIZATION #####
#########################

# Parse the command line
arguments = docopt(__doc__)

# Year selected
year_start = int(arguments["<year_start>"])
year_end = int(arguments["<year_end>"])+1

# Model type
model_type = "Linear Model"
using_poisson = False
if arguments["--p"]:
    using_poisson = True
    model_type="Poisson Model"

# Feature path and labels
path_features = arguments["<dataset_path>"]
path_labels = arguments["<incidence_path>"]
keywords = arguments["<keywords_file>"]

# Selected columns that will be extracted from the dataframes
selected_columns = generate_keywords(keywords)

all_features_values = pd.DataFrame()
all_predicted_values = []
all_true_labels = pd.DataFrame()
total_weeks = []
all_weighted_feature=dict()

# The pair start-end seasons
year_sel = (year_start, year_end)

######################
##### ALGORITHM ######
######################

# For each season, train a model with the complete dataset (without the season)
# and predict the ILI incidence.
for year_selected in range(year_sel[0], year_sel[1]):

    print("------------")
    print("[*] ", year_selected )

    # Generate the list of excluded years
    excluded=[]
    if arguments["--no-future"]:
        for i in range(2007, 2018):
            if i>year_selected:
                excluded.append(i)

    # Data generation from data files
    dataset = generate(year_selected, excluded, path_features)[selected_columns]
    labels = generate_labels(year_selected, excluded, path_labels)
    data = generate_one_year(year_selected, path_features)[selected_columns]
    labels_test = generate_labels_one_year(year_selected, path_labels)
    weeks = labels_test["week"]

    # Create a copy of the dataset where we fills NaN values
    # with 0.
    dataset_zero = dataset.fillna(0)
    data_zero = data.fillna(0)

    # Standardize data
    train, test = standardize_data(dataset_zero, data_zero)

    # Create a Lasso Cross-Validation instance which will be
    # trained on the dataset in which NaN values are replaced
    # with the column mean.
    if not using_poisson:
        model = LassoCV(max_iter=100000, n_jobs=-1, n_alphas=1000, random_state=1)
        model.fit(train, labels["incidence"].fillna(0))
        result = model.predict(test)
    else:
        model = cvglmnet(x=train.as_matrix().copy(), y=labels["incidence"].fillna(0).copy().as_matrix(), family='poisson', alpha=1.0,
                       ptype="mse", parallel=True, nfolds=10)
        result = cvglmnetPredict(model, test.as_matrix(), ptype='response', s="lambda_min")
        coeff = cvglmnetCoef(model, s="lambda_min")

    # Get the feature coefficients
    if not using_poisson:
        coeff = model.coef_
    else:
        coeff = cvglmnetCoef(model, s="lambda_min")

    # Add the pair (value, coeff) to a dictionary
    for i in list(zip(coeff, selected_columns)):
        if (i[0] != 0):
            if all_weighted_feature.get(i[1], None) is None:
                all_weighted_feature[i[1]] = [i[0]]
            else:
                all_weighted_feature[i[1]].append(i[0])

    # Add the value to global variables
    all_features_values= all_features_values.append(data.fillna(0))
    all_predicted_values.extend(result)
    all_true_labels = all_true_labels.append(labels_test)
    total_weeks.extend(weeks)

    # If we are verbose print more informations
    if arguments["--v"]:
        index_peak = labels_test['incidence'].fillna(0).idxmax(axis=1)
        index_peak_m = np.argmax(result)
        print("[*] Influenza Season Peak (Week):", labels_test.iloc[index_peak]["week"])
        print("[*] Influenza Seasons Predicted Peak (Week): ", labels_test.iloc[index_peak_m]["week"])
        print("[*] Influenza Season Peak Value: ", labels_test.iloc[index_peak]["incidence"])
        print("[*] Influenza Season Predicted Peak Value: ", result[index_peak_m])
        print("------------")

# Get important pages to generate the plot we need
important_pages = get_important_pages(all_weighted_feature, 5)

# Print MSE of the two models
mse = mean_squared_error(all_true_labels["incidence"].fillna(0), all_predicted_values)
print("------------")
print("MSE: ", mse)
print("------------")

# Print Pearson Coefficient
pcc = np.corrcoef(all_predicted_values, all_true_labels["incidence"].fillna(0), rowvar=False)[0][1]
print("------------")
print("Pearson Correlation Coeff: ", pcc)
print("------------")

# Print important pages
print(tabulate(important_pages, headers=["Page name", "Weight"]))

###########################
#### GRAPHS GENERATION ####
###########################

# Plot some informations
fig = plt.figure(3, figsize=(15, 6))
ax = fig.add_subplot(111)

# Set up the graph title, x and y axis titles
plt.title("Influenza Seasons "+str(year_sel[0]-1)+" - "+str(year_sel[1]-1), fontsize=18)
plt.ylabel("Incidence on 100000 people", fontsize=17)
plt.xlabel("Year-Week", fontsize=17)

# Generate the x axis labels in such way to
# have the year-week pair only every two ticks.
weeks_used = []
for k, v in enumerate(total_weeks):
    if k%2==0:
        weeks_used.append(v)
    else:
        weeks_used.append(" ")

# Set up the axes ticks
plt.xticks(range(0, len(weeks_used)), weeks_used, rotation="vertical", fontsize=15)
plt.yticks(fontsize=15)

# Plot the model result and the incidence
plt.plot(range(0, len(all_predicted_values)), all_predicted_values, 'r-', label=model_type, linewidth=3)
plt.plot(range(0, len(all_true_labels["incidence"])), all_true_labels["incidence"].fillna(0), 'k-', label="Incidence", linewidth=3)

# Add dotted line into the graph to delimit the influenza seasons.
for i in range(1, (year_sel[1]-year_sel[0])):
    plt.axvline(26*i, linestyle="dashed", color="k", linewidth=3)

# Update some graphical options
plt.grid()
plt.legend(fontsize=16, loc="upper right")
plt.tight_layout()

# Save the graph on file
plt.savefig("season_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+".png", dpi=150)

# Feature plot
figf = plt.figure(4, figsize=(15, 6))
axf = fig.add_subplot(111)

# Generate the graph showing the pageview variation of the top 5 Wikipedia's pages
# choosen by the model.
plt.title("Pageview variation "+str(year_sel[0]-1)+" - "+str(year_sel[1]-1), fontsize=18)
plt.ylabel("Pageview value", fontsize=17)
plt.xlabel("Year-Week", fontsize=17)

# Plot the axes labels
plt.xticks(range(0, len(weeks_used)), weeks_used, rotation="vertical", fontsize=15)
plt.yticks(fontsize=15)

# Plot all the pageview data and the ILI incidence
std_all_features_values = stz(all_features_values)
for key, value in important_pages:
 plt.plot(range(0, len(std_all_features_values[key])), std_all_features_values[key], '-', label=key, linewidth=3)
plt.plot(range(0, len(all_true_labels["incidence"])), stz(all_true_labels["incidence"]), 'k-', label="Incidence", linewidth=3)

# Add dotted line into the graph to delimit the influenza seasons.
for i in range(1, (year_sel[1]-year_sel[0])):
    plt.axvline(26*i, linestyle="dashed", color="k", linewidth=3)

# Update some graphical options
plt.grid()
plt.legend(fontsize=16, loc="best")
plt.tight_layout()

# Save again the graph on file
plt.savefig("season_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_features.png", dpi=150)

# If --f then we write some information to file
if arguments["--f"]:
    with open(str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_information.csv", 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile, delimiter=',',
                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
        spamwriter.writerow(['mse', mean_squared_error(all_true_labels["incidence"].fillna(0), all_predicted_values)])
        for p in important_pages:
            spamwriter.writerow([p[0], float(p[1])])
<|endoftext|>


<commit_before>import tempfile
from subprocess import check_output, CalledProcessError, call

from . import var as v

def _isDeviceAvailable():
    """
    Private Function to check if device is available;
    To be used by only functions inside module
    :return: True or False
    """
    result = getserialno()
    if result[1].strip() == "unknown":
        return False
    else:
        return True

def version():
    """
    Display the version of adb
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_VERSION]
    return _exec_command(adb_full_cmd)

def bugreport(dest_file="default.log"):
    """
    Prints dumpsys, dumpstate, and logcat data to the screen, for the purposes of bug reporting
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_BUGREPORT]
    try:
        dest_file_handler = open(dest_file, "w")
    except IOError:
        print "IOError: Failed to create a log file"

    # We have to check if device is available or not before executing this command
    # as adb bugreport will wait-for-device infinitely and does not come out of
    # loop
    # Execute only if device is available only
    if _isDeviceAvailable():
        result = _exec_command_to_file(adb_full_cmd, dest_file_handler)
        return (result, "Success: Bug report saved to: " + dest_file)
    else:
        return (0, "Device Not Found")

def push(src, dest):
    """
    Push object from host to target
    :param src: string path to source object on host
    :param dest: string destination path on target
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_PUSH, src, dest]
    return _exec_command(adb_full_cmd)

def pull(src, dest):
    """
    Pull object from target to host
    :param src: string path of object on target
    :param dest: string destination path on host
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_PULL, src, dest]
    return _exec_command(adb_full_cmd)

def devices(opts=[]):
    """
    Get list of all available devices including emulators
    :param opts: list command options (e.g. ["-r", "-a"])
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_DEVICES, _convert_opts(opts)]
    return _exec_command(adb_full_cmd)

def shell(cmd):
    """
    Execute shell command on target
    :param cmd: string shell command to execute
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_SHELL, cmd]
    return _exec_command(adb_full_cmd)

def install(apk, opts=[]):
    """
    Install *.apk on target
    :param apk: string path to apk on host to install
    :param opts: list command options (e.g. ["-r", "-a"])
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_INSTALL, _convert_opts(opts), apk]
    return _exec_command(adb_full_cmd)

def uninstall(app, opts=[]):
    """
    Uninstall app from target
    :param app: app name to uninstall from target (e.g. "com.example.android.valid")
    :param opts: list command options (e.g. ["-r", "-a"])
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_UNINSTALL, _convert_opts(opts), app]
    return _exec_command(adb_full_cmd)

def getserialno():
    """
    Get serial number for all available target devices
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_GETSERIALNO]
    return _exec_command(adb_full_cmd)

def wait_for_device():
    """
    Block execution until the device is online
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_WAITFORDEVICE]
    return _exec_command(adb_full_cmd)

def sync():
    """
    Copy host->device only if changed
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_SHELL ,v.ADB_COMMAND_SYNC]
    return _exec_command(adb_full_cmd)

def start_server():
    """
    Startd adb server daemon on host
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_START_SERVER]
    return _exec_command(adb_full_cmd)

def kill_server():
    """
    Kill adb server daemon on host
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_KILL_SERVER]
    return _exec_command(adb_full_cmd)

def get_state():
    """
    Get state of device connected per adb
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_GET_STATE]
    return _exec_command(adb_full_cmd)

def _convert_opts(opts):
    """
    Convert list with command options to single string value
    with 'space' delimeter
    :param opts: list with space-delimeted values
    :return: string with space-delimeted values
    """
    return ' '.join(opts)

def _exec_command(adb_cmd):
    """
    Format adb command and execute it in shell
    :param adb_cmd: list adb command to execute
    :return: string '0' and shell command output if successful, otherwise
    raise CalledProcessError exception and return error code
    """
    t = tempfile.TemporaryFile()
    final_adb_cmd = []
    for e in adb_cmd:
        if e != '':  # avoid items with empty string...
            final_adb_cmd.append(e)  # ... so that final command doesn't
            # contain extra spaces
    print('\n*** Executing ' + ' '.join(adb_cmd) + ' ' + 'command')

    try:
        output = check_output(final_adb_cmd, stderr=t)
    except CalledProcessError as e:
        t.seek(0)
        result = e.returncode, t.read()
    else:
        result = 0, output
        print('\n' + result[1])

    return result

def _exec_command_to_file(adb_cmd, dest_file_handler):
    """
    Format adb command and execute it in shell and redirects to a file
    :param adb_cmd: list adb command to execute
    :param dest_file_handler: file handler to which output will be redirected
    :return: string '0' and writes shell command output to file if successful, otherwise
    raise CalledProcessError exception and return error code
    """
    t = tempfile.TemporaryFile()
    final_adb_cmd = []
    for e in adb_cmd:
        if e != '':  # avoid items with empty string...
            final_adb_cmd.append(e)  # ... so that final command doesn't
            # contain extra spaces
    print('\n*** Executing ' + ' '.join(adb_cmd) + ' ' + 'command')

    try:
        output = call(final_adb_cmd, stdout=dest_file_handler, stderr=t)
    except CalledProcessError as e:
        t.seek(0)
        result = e.returncode, t.read()
    else:
        result = output
        dest_file_handler.close()

    return result<commit_msg>fixed print statement for python3<commit_after>import tempfile
from subprocess import check_output, CalledProcessError, call

from . import var as v

def _isDeviceAvailable():
    """
    Private Function to check if device is available;
    To be used by only functions inside module
    :return: True or False
    """
    result = getserialno()
    if result[1].strip() == "unknown":
        return False
    else:
        return True

def version():
    """
    Display the version of adb
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_VERSION]
    return _exec_command(adb_full_cmd)

def bugreport(dest_file="default.log"):
    """
    Prints dumpsys, dumpstate, and logcat data to the screen, for the purposes of bug reporting
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_BUGREPORT]
    try:
        dest_file_handler = open(dest_file, "w")
    except IOError:
        print("IOError: Failed to create a log file")

    # We have to check if device is available or not before executing this command
    # as adb bugreport will wait-for-device infinitely and does not come out of
    # loop
    # Execute only if device is available only
    if _isDeviceAvailable():
        result = _exec_command_to_file(adb_full_cmd, dest_file_handler)
        return (result, "Success: Bug report saved to: " + dest_file)
    else:
        return (0, "Device Not Found")

def push(src, dest):
    """
    Push object from host to target
    :param src: string path to source object on host
    :param dest: string destination path on target
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_PUSH, src, dest]
    return _exec_command(adb_full_cmd)

def pull(src, dest):
    """
    Pull object from target to host
    :param src: string path of object on target
    :param dest: string destination path on host
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_PULL, src, dest]
    return _exec_command(adb_full_cmd)

def devices(opts=[]):
    """
    Get list of all available devices including emulators
    :param opts: list command options (e.g. ["-r", "-a"])
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_DEVICES, _convert_opts(opts)]
    return _exec_command(adb_full_cmd)

def shell(cmd):
    """
    Execute shell command on target
    :param cmd: string shell command to execute
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_SHELL, cmd]
    return _exec_command(adb_full_cmd)

def install(apk, opts=[]):
    """
    Install *.apk on target
    :param apk: string path to apk on host to install
    :param opts: list command options (e.g. ["-r", "-a"])
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_INSTALL, _convert_opts(opts), apk]
    return _exec_command(adb_full_cmd)

def uninstall(app, opts=[]):
    """
    Uninstall app from target
    :param app: app name to uninstall from target (e.g. "com.example.android.valid")
    :param opts: list command options (e.g. ["-r", "-a"])
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_UNINSTALL, _convert_opts(opts), app]
    return _exec_command(adb_full_cmd)

def getserialno():
    """
    Get serial number for all available target devices
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_GETSERIALNO]
    return _exec_command(adb_full_cmd)

def wait_for_device():
    """
    Block execution until the device is online
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_WAITFORDEVICE]
    return _exec_command(adb_full_cmd)

def sync():
    """
    Copy host->device only if changed
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_SHELL ,v.ADB_COMMAND_SYNC]
    return _exec_command(adb_full_cmd)

def start_server():
    """
    Startd adb server daemon on host
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_START_SERVER]
    return _exec_command(adb_full_cmd)

def kill_server():
    """
    Kill adb server daemon on host
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_KILL_SERVER]
    return _exec_command(adb_full_cmd)

def get_state():
    """
    Get state of device connected per adb
    :return: result of _exec_command() execution
    """
    adb_full_cmd = [v.ADB_COMMAND_PREFIX, v.ADB_COMMAND_GET_STATE]
    return _exec_command(adb_full_cmd)

def _convert_opts(opts):
    """
    Convert list with command options to single string value
    with 'space' delimeter
    :param opts: list with space-delimeted values
    :return: string with space-delimeted values
    """
    return ' '.join(opts)

def _exec_command(adb_cmd):
    """
    Format adb command and execute it in shell
    :param adb_cmd: list adb command to execute
    :return: string '0' and shell command output if successful, otherwise
    raise CalledProcessError exception and return error code
    """
    t = tempfile.TemporaryFile()
    final_adb_cmd = []
    for e in adb_cmd:
        if e != '':  # avoid items with empty string...
            final_adb_cmd.append(e)  # ... so that final command doesn't
            # contain extra spaces
    print('\n*** Executing ' + ' '.join(adb_cmd) + ' ' + 'command')

    try:
        output = check_output(final_adb_cmd, stderr=t)
    except CalledProcessError as e:
        t.seek(0)
        result = e.returncode, t.read()
    else:
        result = 0, output
        print('\n' + result[1])

    return result

def _exec_command_to_file(adb_cmd, dest_file_handler):
    """
    Format adb command and execute it in shell and redirects to a file
    :param adb_cmd: list adb command to execute
    :param dest_file_handler: file handler to which output will be redirected
    :return: string '0' and writes shell command output to file if successful, otherwise
    raise CalledProcessError exception and return error code
    """
    t = tempfile.TemporaryFile()
    final_adb_cmd = []
    for e in adb_cmd:
        if e != '':  # avoid items with empty string...
            final_adb_cmd.append(e)  # ... so that final command doesn't
            # contain extra spaces
    print('\n*** Executing ' + ' '.join(adb_cmd) + ' ' + 'command')

    try:
        output = call(final_adb_cmd, stdout=dest_file_handler, stderr=t)
    except CalledProcessError as e:
        t.seek(0)
        result = e.returncode, t.read()
    else:
        result = output
        dest_file_handler.close()

    return result<|endoftext|>


<commit_before># Third-party
import astropy.time as atime
import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np

# Package
from .log import log as logger
from .celestialmechanics import SimulatedRVOrbit

__all__ = ['plot_rv_curves']

_truth_color = '#006837'
_prev_result_color = '#2166AC'

def plot_rv_curves(samples, t_grid, n_plot=None, rv_unit=None, data=None,
                   ax=None, plot_kwargs=dict(), data_plot_kwargs=dict(),
                   add_labels=True, trend_t_offset=0.):
    """
    Plot radial velocity curves for the input set of orbital parameter
    samples over the input grid of times.

    Parameters
    ----------
    samples : dict
        TODO describe
    t_grid : array_like, `~astropy.time.Time`
        Array of times. Either in BMJD or as an Astropy time object.
    n_plot : int (optional)
        The maximum number of samples to plot. Defaults to 128.
    rv_unit : `~astropy.units.UnitBase` (optional)
        The units to use when plotting RV's.
    data : `~thejoker.data.RVData` (optional)
        Over-plot the data as well.
    ax : `~matplotlib.Axes` (optional)
        A matplotlib axes object to plot on to. If not specified, will
        create a new figure and plot on that.
    plot_kwargs : dict
        Passed to `matplotlib.pyplot.plot()`.
    data_plot_kwargs : dict
        Passed to `thejoker.data.RVData.plot()`.
    add_labels : bool (optional)
        Add labels to the axes or not.
    TODO: trend_t_offset

    Returns
    -------
    fig : `~matplotlib.Figure`

    """

    if ax is None:
        fig,ax = plt.subplots(1,1)
    else:
        fig = ax.figure

    if isinstance(t_grid, atime.Time):
        t_grid = t_grid.tcb.mjd

    if n_plot is None:
        n_plot = len(samples['P'])

    # scale the transparency of the lines
    Q = 4. # HACK
    line_alpha = 0.05 + Q / (n_plot + Q)

    if rv_unit is None:
        rv_unit = u.km/u.s

    # default plotting style
    style = plot_kwargs.copy()
    style.setdefault('linestyle', '-')
    style.setdefault('alpha', line_alpha)
    style.setdefault('marker', '')
    style.setdefault('color', '#555555')

    # plot orbits over the data
    model_rv = np.zeros((n_plot, len(t_grid)))
    for i in range(n_plot):
        this_samples = dict()
        for k in samples.keys():
            this_samples[k] = samples[k][i]
        this_samples.pop('jitter')

        orbit = SimulatedRVOrbit(**this_samples)
        model_rv[i] = orbit.generate_rv_curve(t_grid, trend_t_offset).to(rv_unit).value
    ax.plot(t_grid, model_rv.T, **style)

    if data is not None:
        data_style = data_plot_kwargs.copy()
        data_style.setdefault('rv_unit', rv_unit)
        data_style.setdefault('markersize', 4.)

        if data_style['rv_unit'] != rv_unit:
            raise u.UnitsError("Data plot units don't match rv_unit!")

        data.plot(ax=ax, **data_style)

        _rv = data.rv.to(rv_unit).value
        drv = _rv.max()-_rv.min()
        ax.set_ylim(_rv.min() - 0.2*drv, _rv.max() + 0.2*drv)

    ax.set_xlim(t_grid.min(), t_grid.max())
    if add_labels:
        ax.set_xlabel('BMJD')
        ax.set_ylabel('RV [{}]'.format(rv_unit.to_string(format='latex_inline')))

    return fig
<commit_msg>remove colors [ci skip]<commit_after># Third-party
import astropy.time as atime
import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np

# Package
from .log import log as logger
from .celestialmechanics import SimulatedRVOrbit

__all__ = ['plot_rv_curves']

def plot_rv_curves(samples, t_grid, n_plot=None, rv_unit=None, data=None,
                   ax=None, plot_kwargs=dict(), data_plot_kwargs=dict(),
                   add_labels=True, trend_t_offset=0.):
    """
    Plot radial velocity curves for the input set of orbital parameter
    samples over the input grid of times.

    Parameters
    ----------
    samples : dict
        TODO describe
    t_grid : array_like, `~astropy.time.Time`
        Array of times. Either in BMJD or as an Astropy time object.
    n_plot : int (optional)
        The maximum number of samples to plot. Defaults to 128.
    rv_unit : `~astropy.units.UnitBase` (optional)
        The units to use when plotting RV's.
    data : `~thejoker.data.RVData` (optional)
        Over-plot the data as well.
    ax : `~matplotlib.Axes` (optional)
        A matplotlib axes object to plot on to. If not specified, will
        create a new figure and plot on that.
    plot_kwargs : dict
        Passed to `matplotlib.pyplot.plot()`.
    data_plot_kwargs : dict
        Passed to `thejoker.data.RVData.plot()`.
    add_labels : bool (optional)
        Add labels to the axes or not.
    TODO: trend_t_offset

    Returns
    -------
    fig : `~matplotlib.Figure`

    """

    if ax is None:
        fig,ax = plt.subplots(1,1)
    else:
        fig = ax.figure

    if isinstance(t_grid, atime.Time):
        t_grid = t_grid.tcb.mjd

    if n_plot is None:
        n_plot = len(samples['P'])

    # scale the transparency of the lines
    Q = 4. # HACK
    line_alpha = 0.05 + Q / (n_plot + Q)

    if rv_unit is None:
        rv_unit = u.km/u.s

    # default plotting style
    style = plot_kwargs.copy()
    style.setdefault('linestyle', '-')
    style.setdefault('alpha', line_alpha)
    style.setdefault('marker', '')
    style.setdefault('color', '#555555')

    # plot orbits over the data
    model_rv = np.zeros((n_plot, len(t_grid)))
    for i in range(n_plot):
        this_samples = dict()
        for k in samples.keys():
            this_samples[k] = samples[k][i]
        this_samples.pop('jitter')

        orbit = SimulatedRVOrbit(**this_samples)
        model_rv[i] = orbit.generate_rv_curve(t_grid, trend_t_offset).to(rv_unit).value
    ax.plot(t_grid, model_rv.T, **style)

    if data is not None:
        data_style = data_plot_kwargs.copy()
        data_style.setdefault('rv_unit', rv_unit)
        data_style.setdefault('markersize', 4.)

        if data_style['rv_unit'] != rv_unit:
            raise u.UnitsError("Data plot units don't match rv_unit!")

        data.plot(ax=ax, **data_style)

        _rv = data.rv.to(rv_unit).value
        drv = _rv.max()-_rv.min()
        ax.set_ylim(_rv.min() - 0.2*drv, _rv.max() + 0.2*drv)

    ax.set_xlim(t_grid.min(), t_grid.max())
    if add_labels:
        ax.set_xlabel('BMJD')
        ax.set_ylabel('RV [{}]'.format(rv_unit.to_string(format='latex_inline')))

    return fig
<|endoftext|>


<commit_before>#!/usr/bin/env python3
import logging
import re
import os

from gi.repository import Gst
from configparser import SafeConfigParser
from lib.args import Args
from vocto.transitions import Composites, Transitions

testPatternCount = -1

GST_TYPE_VIDEO_TEST_SRC_PATTERN = [
    "smpte",
    "snow",
    "black",
    "white",
    "red",
    "green",
    "blue",
    "checkers-1",
    "checkers-2",
    "checkers-4",
    "checkers-8",
    "circular",
    "blink",
    "smpte75",
    "zone-plate",
    "gamut",
    "chroma-zone-plate",
    "solid",
    "ball",
    "smpte100",
    "bar"
]

class VocConfigParser(SafeConfigParser):

    log = logging.getLogger('VocConfigParser')

    def getList(self, section, option):
        option = self.get(section, option).strip()
        if len(option) == 0:
            return []

        unfiltered = [x.strip() for x in option.split(',')]
        return list(filter(None, unfiltered))

    def getSources(self):
        return self.getList('mix', 'sources')

    def getAudioSource(self):
        return self.get('mix', 'audiosource', fallback=None)

    def getSlidesSource(self):
        return self.get('mix', 'slides_source_name', fallback=None)

    def getSourceKind(self, source):
        return self.get('source.{}'.format(source), 'kind', fallback='test')

    def getDeckLinkDeviceNumber(self, source):
        return self.get('source.{}'.format(source), 'devicenumber', fallback=0)

    def getDeckLinkAudioConnection(self, source):
        return self.get('source.{}'.format(source), 'audio_connection', fallback='auto')

    def getDeckLinkVideoConnection(self, source):
        return self.get('source.{}'.format(source), 'video_connection', fallback='auto')

    def getDeckLinkVideoMode(self, source):
        return self.get('source.{}'.format(source), 'video_mode', fallback='1080i50')

    def getDeckLinkVideoFormat(self, source):
        return self.get('source.{}'.format(source), 'video_format', fallback='auto')

    def getImageURI(self,source):
        if self.has_option('source.{}'.format(source), 'imguri'):
            return self.get('source.{}'.format(source), 'imguri')
        else:
            path = os.path.abspath(self.get('source.{}'.format(source), 'file'))
            if not os.path.isfile(path):
                self.log.error("image file '%s' could not be found" % path)
            return "file://{}".format(path)

    def getLocation(self,source):
        return self.get('source.{}'.format(source), 'location')

    def getTestPattern(self, source):
        pattern = self.get('source.{}'.format(
            source), 'pattern', fallback=None)
        if not pattern:
            global testPatternCount
            testPatternCount += 1
            pattern = GST_TYPE_VIDEO_TEST_SRC_PATTERN[testPatternCount]
            self.log.info("Pattern unspecified, picking pattern '{} ({})'"
                          .format(pattern, testPatternCount))
        return pattern

    def getSourceDeinterlace(self, source):
        return self.get('source.{}'.format(source), 'deinterlace', fallback='no')

    def getVolume(self, source):
        return self.getfloat("source.{}".format(source), 'volume', fallback=0.0)

    def setShowVolume(self, show=True):
        self.add_section_if_missing('audio')
        self.set('audio', 'volumecontrol', "true" if show else "false")

    def getVideoCaps(self, section='mix'):
        return self.get(section, 'videocaps')

    def getAudioCaps(self, section='mix'):
        return self.get(section, 'audiocaps')

    def getNumAudioStreams(self):
        return self.getint('mix', 'audiostreams', fallback=1)

    def getVideoSize(self, section='mix'):
        caps = Gst.Caps.from_string(
            self.getVideoCaps(section)).get_structure(0)
        _, width = caps.get_int('width')
        _, height = caps.get_int('height')
        return (width, height)

    def getVideoRatio(self, section='mix'):
        width, height = self.getVideoSize(section)
        return float(width)/float(height)

    def getFramerate(self, section='mix'):
        caps = Gst.Caps.from_string(
            self.getVideoCaps(section)).get_structure(0)
        (_, numerator, denominator) = caps.get_fraction('framerate')
        return (numerator, denominator)

    def getFramesPerSecond(self, section='mix'):
        num, denom = self.getFramerate(section)
        return float(num) / float(denom)

    def getVideoSystem(self):
        return self.get('videodisplay', 'system', fallback='gl')

    def getPlayAudio(self):
        return self.getboolean('audio', 'play', fallback=True)

    def getVolumeControl(self):
        # Check if there is a fixed audio source configured.
        # If so, we will remove the volume sliders entirely
        # instead of setting them up.
        return (self.getboolean('audio', 'volumecontrol', fallback=True)
                or self.getboolean('audio', 'forcevolumecontrol', fallback=False))

    def getStreamBlankerEnabled(self):
        return self.getboolean('stream-blanker', 'enabled', fallback=False)

    def getStreamBlankerSources(self):
        if self.getStreamBlankerEnabled():
            return self.getList('stream-blanker', 'sources')
        else:
            return []

    def getStreamBlankerVolume(self):
        return self.getfloat('stream-blanker', 'volume', fallback=0.0)

    def getMirrorsEnabled(self):
        return self.getboolean('mirrors', 'enabled', fallback=False)

    def getMirrorsSources(self):
        if self.getMirrorsEnabled():
            if self.has_option('mirrors', 'sources'):
                return self.getList('mirrors', 'sources')
            else:
                return self.getSources()
        else:
            return []

    def getOutputBuffers(self, channel):
        return self.getint('output-buffers', channel, fallback=500)

    def getPreviewVaapi(self):
        if self.has_option('previews', 'vaapi'):
            return self.get('previews', 'vaapi')
        return None

    def getPreviewCaps(self):
        if self.has_option('previews', 'videocaps'):
            return self.getVideoCaps('previews')
        else:
            return self.getVideoCaps()

    def getPreviewSize(self):
        width = self.getint('previews', 'width') if self.has_option(
            'previews', 'width') else 320
        height = self.getint('previews', 'height') if self.has_option(
            'previews', 'height') else int(width * 9 / 16)
        return(width, height)

    def getDeinterlacePreviews(self):
        return self.getboolean('previews', 'deinterlace')

    def getPreviewsEnabled(self):
        return self.getboolean('previews', 'enabled')

    def getLivePreviewEnabled(self):
        return self.getboolean('previews', 'live', fallback=False)

    def getPreviewDecoder(self):
        if self.has_option('previews', 'vaapi'):
            return self.get('previews', 'vaapi')
        else:
            return 'jpeg'

    def getComposites(self):
        return Composites.configure(self.items('composites'), self.getVideoSize())

    def getTargetComposites(self):
        return Composites.targets(self.getComposites())

    def getTransitions(self, composites):
        return Transitions.configure(self.items('transitions'),
                                     composites,
                                     fps=self.getFramesPerSecond())

    def getPreviewNameOverlay(self):
        return self.getboolean('previews', 'nameoverlay',fallback=True)

    def hasSource(self, source):
        return self.has_section('source.{}'.format(source))

    def hasOverlay(self):
        return self.has_section('overlay')

    def getOverlayAutoOff(self):
        return self.getboolean('overlay', 'auto-off', fallback=True)

    def getOverlayUserAutoOff(self):
        return self.getboolean('overlay', 'user-auto-off', fallback=False)
<commit_msg>Added some default values for configuration items which had none before<commit_after>#!/usr/bin/env python3
import logging
import re
import os

from gi.repository import Gst
from configparser import SafeConfigParser
from lib.args import Args
from vocto.transitions import Composites, Transitions

testPatternCount = -1

GST_TYPE_VIDEO_TEST_SRC_PATTERN = [
    "smpte",
    "snow",
    "black",
    "white",
    "red",
    "green",
    "blue",
    "checkers-1",
    "checkers-2",
    "checkers-4",
    "checkers-8",
    "circular",
    "blink",
    "smpte75",
    "zone-plate",
    "gamut",
    "chroma-zone-plate",
    "solid",
    "ball",
    "smpte100",
    "bar"
]

class VocConfigParser(SafeConfigParser):

    log = logging.getLogger('VocConfigParser')

    def getList(self, section, option):
        option = self.get(section, option).strip()
        if len(option) == 0:
            return []

        unfiltered = [x.strip() for x in option.split(',')]
        return list(filter(None, unfiltered))

    def getSources(self):
        return self.getList('mix', 'sources')

    def getAudioSource(self):
        return self.get('mix', 'audiosource', fallback=None)

    def getSlidesSource(self):
        return self.get('mix', 'slides_source_name', fallback=None)

    def getSourceKind(self, source):
        return self.get('source.{}'.format(source), 'kind', fallback='test')

    def getDeckLinkDeviceNumber(self, source):
        return self.get('source.{}'.format(source), 'devicenumber', fallback=0)

    def getDeckLinkAudioConnection(self, source):
        return self.get('source.{}'.format(source), 'audio_connection', fallback='auto')

    def getDeckLinkVideoConnection(self, source):
        return self.get('source.{}'.format(source), 'video_connection', fallback='auto')

    def getDeckLinkVideoMode(self, source):
        return self.get('source.{}'.format(source), 'video_mode', fallback='1080i50')

    def getDeckLinkVideoFormat(self, source):
        return self.get('source.{}'.format(source), 'video_format', fallback='auto')

    def getImageURI(self,source):
        if self.has_option('source.{}'.format(source), 'imguri'):
            return self.get('source.{}'.format(source), 'imguri')
        else:
            path = os.path.abspath(self.get('source.{}'.format(source), 'file'))
            if not os.path.isfile(path):
                self.log.error("image file '%s' could not be found" % path)
            return "file://{}".format(path)

    def getLocation(self,source):
        return self.get('source.{}'.format(source), 'location')

    def getTestPattern(self, source):
        pattern = self.get('source.{}'.format(
            source), 'pattern', fallback=None)
        if not pattern:
            global testPatternCount
            testPatternCount += 1
            pattern = GST_TYPE_VIDEO_TEST_SRC_PATTERN[testPatternCount]
            self.log.info("Pattern unspecified, picking pattern '{} ({})'"
                          .format(pattern, testPatternCount))
        return pattern

    def getSourceDeinterlace(self, source):
        return self.get('source.{}'.format(source), 'deinterlace', fallback='no')

    def getVolume(self, source):
        return self.getfloat("source.{}".format(source), 'volume', fallback=0.0)

    def setShowVolume(self, show=True):
        self.add_section_if_missing('audio')
        self.set('audio', 'volumecontrol', "true" if show else "false")

    def getVideoCaps(self):
        return self.get('mix', 'videocaps', fallback="video/x-raw,format=I420,width=1920,height=1080,framerate=25/1,pixel-aspect-ratio=1/1")

    def getAudioCaps(self, section='mix'):
        return self.get(section, 'audiocaps', fallback="audio/x-raw,format=S16LE,channels=2,layout=interleaved,rate=48000")

    def getNumAudioStreams(self):
        return self.getint('mix', 'audiostreams', fallback=1)

    def getVideoSize(self):
        caps = Gst.Caps.from_string(
            self.getVideoCaps()).get_structure(0)
        _, width = caps.get_int('width')
        _, height = caps.get_int('height')
        return (width, height)

    def getVideoRatio(self):
        width, height = self.getVideoSize()
        return float(width)/float(height)

    def getFramerate(self):
        caps = Gst.Caps.from_string(
            self.getVideoCaps()).get_structure(0)
        (_, numerator, denominator) = caps.get_fraction('framerate')
        return (numerator, denominator)

    def getFramesPerSecond(self):
        num, denom = self.getFramerate()
        return float(num) / float(denom)

    def getVideoSystem(self):
        return self.get('videodisplay', 'system', fallback='gl')

    def getPlayAudio(self):
        return self.getboolean('audio', 'play', fallback=True)

    def getVolumeControl(self):
        # Check if there is a fixed audio source configured.
        # If so, we will remove the volume sliders entirely
        # instead of setting them up.
        return (self.getboolean('audio', 'volumecontrol', fallback=True)
                or self.getboolean('audio', 'forcevolumecontrol', fallback=False))

    def getStreamBlankerEnabled(self):
        return self.getboolean('stream-blanker', 'enabled', fallback=False)

    def getStreamBlankerSources(self):
        if self.getStreamBlankerEnabled():
            return self.getList('stream-blanker', 'sources')
        else:
            return []

    def getStreamBlankerVolume(self):
        return self.getfloat('stream-blanker', 'volume', fallback=0.0)

    def getMirrorsEnabled(self):
        return self.getboolean('mirrors', 'enabled', fallback=False)

    def getMirrorsSources(self):
        if self.getMirrorsEnabled():
            if self.has_option('mirrors', 'sources'):
                return self.getList('mirrors', 'sources')
            else:
                return self.getSources()
        else:
            return []

    def getOutputBuffers(self, channel):
        return self.getint('output-buffers', channel, fallback=500)

    def getPreviewVaapi(self):
        if self.has_option('previews', 'vaapi'):
            return self.get('previews', 'vaapi')
        return None

    def getPreviewCaps(self):
        if self.has_option('previews', 'videocaps'):
            return self.get('previews', 'videocaps', fallback='video/x-raw,width=1024,height=576,framerate=25/1')
        else:
            return self.getVideoCaps()

    def getPreviewSize(self):
        width = self.getint('previews', 'width') if self.has_option(
            'previews', 'width') else 320
        height = self.getint('previews', 'height') if self.has_option(
            'previews', 'height') else int(width * 9 / 16)
        return(width, height)

    def getDeinterlacePreviews(self):
        return self.getboolean('previews', 'deinterlace', fallback=False)

    def getPreviewsEnabled(self):
        return self.getboolean('previews', 'enabled', fallback=False)

    def getLivePreviewEnabled(self):
        return self.getboolean('previews', 'live', fallback=False)

    def getPreviewDecoder(self):
        if self.has_option('previews', 'vaapi'):
            return self.get('previews', 'vaapi')
        else:
            return 'jpeg'

    def getComposites(self):
        return Composites.configure(self.items('composites'), self.getVideoSize())

    def getTargetComposites(self):
        return Composites.targets(self.getComposites())

    def getTransitions(self, composites):
        return Transitions.configure(self.items('transitions'),
                                     composites,
                                     fps=self.getFramesPerSecond())

    def getPreviewNameOverlay(self):
        return self.getboolean('previews', 'nameoverlay',fallback=True)

    def hasSource(self, source):
        return self.has_section('source.{}'.format(source))

    def hasOverlay(self):
        return self.has_section('overlay')

    def getOverlayAutoOff(self):
        return self.getboolean('overlay', 'auto-off', fallback=True)

    def getOverlayUserAutoOff(self):
        return self.getboolean('overlay', 'user-auto-off', fallback=False)
<|endoftext|>


<commit_before>"""
HOW THIS WORKS:
1. Message objects are created via /dashboard/messages. Message objects have a string indicating their audience (e.g. all registered hackers, all checked-in, mailing list people, teammates listed on people's applications)
2. Message.enqueue_tasks() enqueues a bunch of Tas objects in the `messages` queue, by:
 2.1 calling Message.get_query() to get a datastore query
 2.2 batching the entities returned from the query into groups of 40 (batch_size)
 2.3 creating Task Queue tasks for each batch, containing message ID and entity ID for all entities
3. MessagesTaskQueueWork.post is called by the task queue for each task.
4. MessagesTaskQueueWork.post loads the Message object and all entities, then concurrently calls Message.send_to_entity_async(entity) on each entity. These are responsible for looking the audience, the current entity (which may be a Hacker or an EmailListEntry, or maybe something else in the future), and sending the appropriate emails or SMSs
"""

import webapp2
from template import template, template_string
from google.appengine.ext import ndb
from send_email import send_email
import re
from google.appengine.api import taskqueue
from registration import Hacker
from email_list import EmailListEntry
import sms
from config import isAdmin

class Message(ndb.Model):
	added = ndb.DateTimeProperty(auto_now_add=True)

	audience = ndb.StringProperty(choices=[None, 'registered', 'invited-friends', 'mailing-list-unregistered', 'waitlisted'], default=None)

	email_from_template = ndb.BooleanProperty(default=False)
	email_subject = ndb.TextProperty()
	email_html = ndb.TextProperty()

	sms_text = ndb.TextProperty()

	def kick_off(self):
		q = taskqueue.Queue("messages")
		q.add(taskqueue.Task(url='/dashboard/messages/message_task_queue_work', params={"kickoff_message_key": self.key.urlsafe()}))

	def send_to_email(self, email, template_args={}):
		# does the actual work of sending
		emails = [email]
		assert self.email_subject, "No email subject provided. Is email unchecked?"
		# if self.audience == 'invited-friends':
		# 	hacker = Hacker.query(Hacker.email==email).fetch()[0]
		# 	if hacker.teammates:
		# 		emails_found = [email.lower() for email in hacker.teammates.split(',')]
		# 		matching_hackers =  Hacker.query(Hacker.email.IN(emails_found)).fetch()
		# 		emails_already_registered = [h.email for h in matching_hackers]
		# 		emails = []
		# 		for email in emails_found:
		# 			if email not in emails_already_registered:
		# 				emails.append(email)
		# 				template_args["invited_by"] = hacker
		if self.email_from_template:
			html = template("emails/" + self.email_html + ".html", template_args)
		else:
			html = template_string(self.email_html, template_args)
			html = template("emails/generic.html", dict({"content": html}.items() + template_args.items()))
		subject = template_string(self.email_subject, template_args)
		send_email(html, subject, emails)

	def send_to_phone(self, phone):
		# actual work of sms'ing
		print "Sending SMS '{0}' to {1}".format(self.sms_text, phone)
		sms.send(phone, self.sms_text)

	def get_query(self):
		if self.audience == 'registered':
			return Hacker.query()
		elif self.audience == 'invited-friends':
			return Hacker.query(Hacker.teammates != None)
		elif self.audience == 'mailing-list-unregistered':
			return EmailListEntry.query()
		elif self.audience == 'waitlisted':
			return Hacker.query(Hacker.admitted_email_sent_date == None and Hacker.waitlist_email_sent_date != None)
		elif self.audience == None:
			return None
		else:
			assert 0, "Unknown audience"

	def enqueue_tasks(self):
		if self.get_query() == None:
			print "No audience selected. No emails sent through messages system."
			return
		q = taskqueue.Queue("messages")
		# max task size is 100kb; max # of tasks added per batch is 100
		task_futures = []

		batch_size = 40 # 2k users / 40 = 50 tasks. within our limits.
		batch_of_entity_keys = []
		def send_batch():
			params = {"message_key": self.key.urlsafe(), "entity_keys": ','.join(batch_of_entity_keys[:])}
			task = taskqueue.Task(params=params, url='/dashboard/messages/message_task_queue_work')
			task_future = q.add_async(task)
			task_futures.append(task_future)
			batch_of_entity_keys[:] = []

		for key in self.get_query().iter(keys_only=True):
			batch_of_entity_keys.append(key.urlsafe())
			if len(batch_of_entity_keys) >= batch_size: send_batch()
		if len(batch_of_entity_keys): send_batch()

		for f in task_futures:
			f.get_result()

	@ndb.tasklet
	def send_to_entity_async(self, entity):
		try:
			if self.audience == 'invited-friends':
				# don't actually send to the hacker -- send to their friends
				hacker = entity
				if hacker.teammates:
					emails = [email.lower() for email in hacker.teammates.split(',')]
					matching_hackers = yield Hacker.query(Hacker.email.IN(emails)).fetch_async()
					emails_already_registered = [h.email for h in matching_hackers]
					for email in emails:
						if email not in emails_already_registered:
							self.send_to_email(email, {"invited_by": hacker})
			elif self.audience == 'registered': # send emails directly to hackers
				hacker = entity
				if hacker.email and self.email_subject:
					self.send_to_email(hacker.email, {"hacker": hacker})
				if hacker.phone_number and self.sms_text:
					self.send_to_phone(hacker.phone_number)
			elif self.audience == 'mailing-list-unregistered':
				email = entity.email
				is_registered = (yield Hacker.query(Hacker.email == email).count_async()) > 0
				if not is_registered:
					self.send_to_email(email, {})
			elif self.audience == 'waitlisted':
				hacker = entity
				hacker.waitlist_email_sent_date = datetime.datetime.now()
				self.send_to_email(hacker.email, {"hacker": hacker})
				hacker.put()

		except Exception as e:
			print "Failed to send email '{0}' to '{1} because {2}'".format(self.email_subject, entity.email, e)

class MessagesDashboardHandler(webapp2.RequestHandler):
	def get(self):
		if not isAdmin(): return self.redirect('/')
		self.response.write(template("messages_dashboard.html", {}))

	def post(self):
		if not isAdmin(): return self.redirect('/')
		audience = None if self.request.get('audience') == '' else self.request.get('audience')
		message = Message(audience=audience)
		if self.request.get('email'):
			message.email_subject = self.request.get('email-subject')
			if self.request.get('email-name'):
				message.email_from_template = True
				message.email_html = self.request.get('email-name')
			else:
				message.email_html = self.request.get('email-html')
		if self.request.get('sms'):
			message.sms_text = self.request.get('sms-text')
		if self.request.get('test'):
			recip = self.request.get('test-recipient')
			if '@' in recip:
				message.send_to_email(recip)
				self.response.write("Sent email")
			elif len(recip) > 0:
				message.send_to_phone(recip)
				self.response.write("Sent SMS")
		else:
			message.put()
			message.kick_off()
			self.response.write(template("messages_dashboard.html", {"status": "Sent!"}))

class MessagesTaskQueueWork(webapp2.RequestHandler):
	def post(self):
		#if not isAdmin(): return self.redirect('/')
		if self.request.get('kickoff_message_key'):
			ndb.Key(urlsafe=self.request.get('kickoff_message_key')).get().enqueue_tasks()
		elif self.request.get('entity_keys'):
			keys = [ndb.Key(urlsafe=self.request.get('message_key'))] + map(lambda k: ndb.Key(urlsafe=k), self.request.get('entity_keys').split(','))
			msg_and_entities = ndb.get_multi(keys)
			message = msg_and_entities[0]
			entities = msg_and_entities[1:]
			futures = [message.send_to_entity_async(entity) for entity in entities]
			for f in futures: f.wait()
<commit_msg>fixes to ALL THE QUERIES OH GOD<commit_after>"""
HOW THIS WORKS:
1. Message objects are created via /dashboard/messages. Message objects have a string indicating their audience (e.g. all registered hackers, all checked-in, mailing list people, teammates listed on people's applications)
2. Message.enqueue_tasks() enqueues a bunch of Tas objects in the `messages` queue, by:
 2.1 calling Message.get_query() to get a datastore query
 2.2 batching the entities returned from the query into groups of 40 (batch_size)
 2.3 creating Task Queue tasks for each batch, containing message ID and entity ID for all entities
3. MessagesTaskQueueWork.post is called by the task queue for each task.
4. MessagesTaskQueueWork.post loads the Message object and all entities, then concurrently calls Message.send_to_entity_async(entity) on each entity. These are responsible for looking the audience, the current entity (which may be a Hacker or an EmailListEntry, or maybe something else in the future), and sending the appropriate emails or SMSs
"""

import webapp2
from template import template, template_string
from google.appengine.ext import ndb
from send_email import send_email
import re
from google.appengine.api import taskqueue
from registration import Hacker
from email_list import EmailListEntry
import sms
from config import isAdmin

class Message(ndb.Model):
	added = ndb.DateTimeProperty(auto_now_add=True)

	audience = ndb.StringProperty(choices=[None, 'registered', 'invited-friends', 'mailing-list-unregistered', 'waitlisted'], default=None)

	email_from_template = ndb.BooleanProperty(default=False)
	email_subject = ndb.TextProperty()
	email_html = ndb.TextProperty()

	sms_text = ndb.TextProperty()

	def kick_off(self):
		q = taskqueue.Queue("messages")
		q.add(taskqueue.Task(url='/dashboard/messages/message_task_queue_work', params={"kickoff_message_key": self.key.urlsafe()}))

	def send_to_email(self, email, template_args={}):
		# does the actual work of sending
		emails = [email]
		assert self.email_subject, "No email subject provided. Is email unchecked?"
		# if self.audience == 'invited-friends':
		# 	hacker = Hacker.query(Hacker.email==email).fetch()[0]
		# 	if hacker.teammates:
		# 		emails_found = [email.lower() for email in hacker.teammates.split(',')]
		# 		matching_hackers =  Hacker.query(Hacker.email.IN(emails_found)).fetch()
		# 		emails_already_registered = [h.email for h in matching_hackers]
		# 		emails = []
		# 		for email in emails_found:
		# 			if email not in emails_already_registered:
		# 				emails.append(email)
		# 				template_args["invited_by"] = hacker
		if self.email_from_template:
			html = template("emails/" + self.email_html + ".html", template_args)
		else:
			html = template_string(self.email_html, template_args)
			html = template("emails/generic.html", dict({"content": html}.items() + template_args.items()))
		subject = template_string(self.email_subject, template_args)
		send_email(html, subject, emails)

	def send_to_phone(self, phone):
		# actual work of sms'ing
		print "Sending SMS '{0}' to {1}".format(self.sms_text, phone)
		sms.send(phone, self.sms_text)

	def get_query(self):
		if self.audience == 'registered':
			return Hacker.query()
		elif self.audience == 'invited-friends':
			return Hacker.query(Hacker.teammates != None)
		elif self.audience == 'mailing-list-unregistered':
			return EmailListEntry.query()
		elif self.audience == 'waitlisted':
			return Hacker.query(Hacker.admitted_email_sent_date == None, Hacker.waitlist_email_sent_date != None)
		elif self.audience == None:
			return None
		else:
			assert 0, "Unknown audience"

	def enqueue_tasks(self):
		if self.get_query() == None:
			print "No audience selected. No emails sent through messages system."
			return
		q = taskqueue.Queue("messages")
		# max task size is 100kb; max # of tasks added per batch is 100
		task_futures = []

		batch_size = 40 # 2k users / 40 = 50 tasks. within our limits.
		batch_of_entity_keys = []
		def send_batch():
			params = {"message_key": self.key.urlsafe(), "entity_keys": ','.join(batch_of_entity_keys[:])}
			task = taskqueue.Task(params=params, url='/dashboard/messages/message_task_queue_work')
			task_future = q.add_async(task)
			task_futures.append(task_future)
			batch_of_entity_keys[:] = []

		for key in self.get_query().iter(keys_only=True):
			batch_of_entity_keys.append(key.urlsafe())
			if len(batch_of_entity_keys) >= batch_size: send_batch()
		if len(batch_of_entity_keys): send_batch()

		for f in task_futures:
			f.get_result()

	@ndb.tasklet
	def send_to_entity_async(self, entity):
		try:
			if self.audience == 'invited-friends':
				# don't actually send to the hacker -- send to their friends
				hacker = entity
				if hacker.teammates:
					emails = [email.lower() for email in hacker.teammates.split(',')]
					matching_hackers = yield Hacker.query(Hacker.email.IN(emails)).fetch_async()
					emails_already_registered = [h.email for h in matching_hackers]
					for email in emails:
						if email not in emails_already_registered:
							self.send_to_email(email, {"invited_by": hacker})
			elif self.audience == 'registered': # send emails directly to hackers
				hacker = entity
				if hacker.email and self.email_subject:
					self.send_to_email(hacker.email, {"hacker": hacker})
				if hacker.phone_number and self.sms_text:
					self.send_to_phone(hacker.phone_number)
			elif self.audience == 'mailing-list-unregistered':
				email = entity.email
				is_registered = (yield Hacker.query(Hacker.email == email).count_async()) > 0
				if not is_registered:
					self.send_to_email(email, {})
			elif self.audience == 'waitlisted':
				hacker = entity
				hacker.waitlist_email_sent_date = datetime.datetime.now()
				self.send_to_email(hacker.email, {"hacker": hacker})
				hacker.put()

		except Exception as e:
			print "Failed to send email '{0}' to '{1} because {2}'".format(self.email_subject, entity.email, e)

class MessagesDashboardHandler(webapp2.RequestHandler):
	def get(self):
		if not isAdmin(): return self.redirect('/')
		self.response.write(template("messages_dashboard.html", {}))

	def post(self):
		if not isAdmin(): return self.redirect('/')
		audience = None if self.request.get('audience') == '' else self.request.get('audience')
		message = Message(audience=audience)
		if self.request.get('email'):
			message.email_subject = self.request.get('email-subject')
			if self.request.get('email-name'):
				message.email_from_template = True
				message.email_html = self.request.get('email-name')
			else:
				message.email_html = self.request.get('email-html')
		if self.request.get('sms'):
			message.sms_text = self.request.get('sms-text')
		if self.request.get('test'):
			recip = self.request.get('test-recipient')
			if '@' in recip:
				message.send_to_email(recip)
				self.response.write("Sent email")
			elif len(recip) > 0:
				message.send_to_phone(recip)
				self.response.write("Sent SMS")
		else:
			message.put()
			message.kick_off()
			self.response.write(template("messages_dashboard.html", {"status": "Sent!"}))

class MessagesTaskQueueWork(webapp2.RequestHandler):
	def post(self):
		#if not isAdmin(): return self.redirect('/')
		if self.request.get('kickoff_message_key'):
			ndb.Key(urlsafe=self.request.get('kickoff_message_key')).get().enqueue_tasks()
		elif self.request.get('entity_keys'):
			keys = [ndb.Key(urlsafe=self.request.get('message_key'))] + map(lambda k: ndb.Key(urlsafe=k), self.request.get('entity_keys').split(','))
			msg_and_entities = ndb.get_multi(keys)
			message = msg_and_entities[0]
			entities = msg_and_entities[1:]
			futures = [message.send_to_entity_async(entity) for entity in entities]
			for f in futures: f.wait()
<|endoftext|>