import logging
from logging.handlers import TimedRotatingFileHandler
import pandas as pd
from flask import Flask, abort, redirect, request, jsonify
from flask_restful import Resource, Api, reqparse
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.linear_model import LogisticRegression
import scipy
import pickle

def people_tokenizer(s):
    people = s.split(';')
    return filter(None, people)

app = Flask(__name__)
api = Api(app)
parser = reqparse.RequestParser()
parser.add_argument('Body', type=str, help='The body of the email to be sorted.')
parser.add_argument('Subject', type=str, help='The subject of the email to be sorted.')
parser.add_argument('ToRecipients', type=str, help='The to line of the email to be sorted, semi-colon separated')
parser.add_argument('Sender', type=str, help='The sender of the email to be sorted')
parser.add_argument('CcRecipients', type=str, help='The cc line of the email to be sorted, semi-colon separated')
parser.add_argument('SentDateTime', type=str, help='The datetime that the email was sent.')

class HelloWorld(Resource):
    def get(self):
        return {}

# TODO: Serve N predictions in one request to improve performance
class Prediction(Resource):
    def post(self):
        args = parser.parse_args()
        app.logger.debug(args)

        # userId = '8lWZYw-u-yNbGBkC4B--ip77K1oVwwyZTHKLeD7rm7k'
        userId = request.args['userId']

        # Load message as pandas dataframe
        query_message = pd.DataFrame(args, index=[0])
        app.logger.debug(query_message.head(5))

        # TODO: Load model from local file, since we only have one user right now.
        # TODO: Could hold in memory to limit i/o on every request
        try:
            app.logger.debug('Reading model for user {0} from disk'.format(userId))
            pkl_name = '{0}.pkl'.format(userId)
            pipeline = pickle.load(open(pkl_name, 'rb'))
        
        except Exception as ex:
            app.logger.error(str(ex))
            raise ex

        app.logger.debug(pipeline)
        sub_vect = pipeline.named_steps['sub_vocab']
        sub_tfidf = pipeline.named_steps['sub_tfidf']
        people_vect = pipeline.named_steps['people_vocab']
        trained_model = pipeline.named_steps['classifier']

        # Perform bag of words vectorization on Subject
        subject_counts = sub_vect.transform(query_message['Subject'])
        app.logger.debug('Performed bag of words feature construction on subject')

        # Convert Subject from bools to tf-idf values
        subject_tfidf = sub_tfidf.transform(subject_counts)

        # Create a new column that merges CC, To, From
        query_message['CcRecipients'].fillna('', inplace=True)
        query_message['ToRecipients'].fillna('', inplace=True)
        query_message['Sender'].fillna('', inplace=True)
        query_message['People'] = query_message['Sender'] + ';' + query_message['CcRecipients'] + ';' + query_message['ToRecipients']
        app.logger.debug(query_message.head(5))

        # Convert people column into email address counts
        people_counts = people_vect.transform(query_message['People'])
        app.logger.debug('Performed one hot encodig of people column')

        # Merge people counts with subject tf-idf
        feature_matrix = scipy.sparse.hstack([people_counts, subject_tfidf])

        # Make prediction on resulting query message
        class_predictions = trained_model.predict_proba(feature_matrix)
        app.logger.debug(class_predictions)
        prediction_idx = -1
        best_prediction = 0
        for idx, prediction in enumerate(class_predictions[0]):
            if prediction > best_prediction:
                prediction_idx = idx
                best_prediction = prediction

        return { 'class_label': trained_model.classes_[prediction_idx], 'probability': best_prediction }

api.add_resource(Prediction, '/predict/')
api.add_resource(HelloWorld, '/hello/')

if __name__ == '__main__':

    logHandler = TimedRotatingFileHandler('./logs/predictions.log', when="D", interval=1, backupCount=0)
    logHandler.setLevel(logging.DEBUG)
    app.logger.setLevel(logging.DEBUG)
    app.logger.addHandler(logHandler)

    app.run(port=5000, host='0.0.0.0')