from tokenizer import sent_tokenizer, word_tokenizer
import biviewer
import pdb
import re
import progressbar
import collections
import string
from unidecode import unidecode
import codecs

import yaml
from pprint import pprint

import numpy as np
import math

import difflib

import sklearn
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.grid_search import GridSearchCV

from sklearn.feature_extraction import DictVectorizer

from sklearn import cross_validation
from sklearn import metrics
from sklearn import svm
from sklearn.linear_model import SGDClassifier
from sklearn.externals import six

from collections import defaultdict

from sklearn.metrics import precision_recall_fscore_support
import random
import operator

from sklearn.cross_validation import KFold

from journalreaders import PdfReader

import cPickle as pickle

from sklearn.metrics import f1_score, make_scorer, fbeta_score, accuracy_score
import nltk
from nltk.corpus import stopwords

REGEX_QUOTE_PRESENT = re.compile("Quote\:")
REGEX_QUOTE = re.compile("\"(.*?)\"") # retrive blocks of text in quotes
REGEX_ELLIPSIS = re.compile("\s*[\[\(]?\s?\.\.+\s?[\]\)]?\s*") # to catch various permetations of "..." and "[...]"

SIMPLE_WORD_TOKENIZER = re.compile("[a-zA-Z]{2,}") # regex of the rule used by sklearn CountVectorizer

CORE_DOMAINS = ["Random sequence generation", "Allocation concealment", "Blinding of participants and personnel",
                "Blinding of outcome assessment", "Incomplete outcome data", "Selective reporting"]
                # "OTHER" is generated in code, not in the mapping file
                # see data/domain_names.txt for various other criteria
                # all of these are available via QualityQuoteReader

ALL_DOMAINS = CORE_DOMAINS[:] # will be added to later

RoB_CLASSES = ["YES", "NO", "UNKNOWN"]

STOP_WORDS = set(stopwords.words('english'))

# @TODO move me
domain_str = lambda d: d.lower().replace(" ", "_")

def show_most_informative_features(vectorizer, clf, n=1000):
    ###
    # note that in the multi-class case, clf.coef_ will
    # have k weight vectors, which I believe are one per
    # each class (i.e., each is a classifier discriminating
    # one class versus the rest).
    c_f = sorted(zip(clf.coef_[0], vectorizer.get_feature_names()))

    if n == 0:
        n = len(c_f)/2

    top = zip(c_f[:n], c_f[:-(n+1):-1])
    print
    print "%d most informative features:" % (n, )
    out_str = []
    for (c1, f1), (c2, f2) in top:
        out_str.append("\t%.4f\t%-15s\t\t%.4f\t%-15s" % (c1, f1, c2, f2))
    feature_str = "\n".join(out_str)
    return feature_str

def show_most_informative_features_ynu(vectorizer, clf, n=10):
    ###
    # note that in the multi-class case, clf.coef_ will
    # have k weight vectors, which I believe are one per
    # each class (i.e., each is a classifier discriminating
    # one class versus the rest).

    combinations =  ["NO vs (YES + UNKNOWN)", "UNKNOWN vs (YES + NO)", "YES vs (NO + UNKNOWN)"]

    out_str = []

    for i, combination in enumerate(combinations):

        out_str.append(combination)
        out_str.append("*" * 20)

        c_f = sorted(zip(clf.coef_[i], vectorizer.get_feature_names()))

        if n == 0:
            n = len(c_f)/2

        top = zip(c_f[:n], c_f[:-(n+1):-1])

        for (c1, f1), (c2, f2) in top:
            out_str.append("\t%.4f\t%-15s\t\t%.4f\t%-15s" % (c1, f1, c2, f2))
    feature_str = "\n".join(out_str)
    return feature_str

def load_domain_map(filename="data/domain_names.txt"):

    with codecs.open(filename, 'rb', 'utf-8') as f:
        raw_data = yaml.load(f)

    mapping = {}

    for key, value in raw_data.iteritems():
        for synonym in value:
            mapping[synonym] = key

    return mapping

class QualityQuoteReader2():
    """
    iterates through Cochrane Risk of Bias information
    v2 maintains unique ids for all source studies + does not filter by whether a quote is present
    returns list of quotes where they are available, else none
    """

    def __init__(self, sent=False, test_mode=False):
        self.BiviewerView = collections.namedtuple('BiViewer_View', ['uid', 'cochrane', 'studypdf'])
        self.pdfviewer = biviewer.PDFBiViewer()
        self.domain_map = load_domain_map()
        if test_mode:
            self.test_mode_break_point = 500
        else:
            self.test_mode_break_point = None

    def __iter__(self):
        """
        run through PDF/Cochrane data
        preprocesses PDF text
        and maps domain title to one of the core Risk of Bias domains if possible
        """

        progress_bar_limit = len(self.pdfviewer) if self.test_mode_break_point is None else self.test_mode_break_point
        p = progressbar.ProgressBar(progress_bar_limit, timer=True)

        for uid, study in enumerate(self.pdfviewer):

            if self.test_mode_break_point and (uid >= self.test_mode_break_point):
                break

            p.tap()
            quality_data = study.cochrane["QUALITY"]
            for domain in quality_data:

                domain["QUOTES"] = self.preprocess_cochrane(domain["DESCRIPTION"])
                try:
                    domain["DOMAIN"] = self.domain_map[domain["DOMAIN"]] # map domain titles to our core categories
                except:
                    domain["DOMAIN"] = "OTHER"

            yield self.BiviewerView(uid=uid, cochrane={"QUALITY": quality_data}, studypdf=self.preprocess_pdf(study.studypdf))

    def __len__(self):
        return len(self.pdfviewer) if self.test_mode_break_point is None else self.test_mode_break_point

    def preprocess_pdf(self, pdftext):
        pdftext = unidecode(pdftext)
        pdftext = re.sub("\n", " ", pdftext) # preprocessing rule 1
        return pdftext

    def preprocess_cochrane(self, rawtext):

        # regex clean up of cochrane strings
        processedtext = unidecode(rawtext)
        processedtext = re.sub(" +", " ", processedtext)

        # extract all parts in quotes
        quotes = REGEX_QUOTE.findall(processedtext)

        # then split at any ellipses
        quote_parts = []
        for quote in quotes:
            quote_parts.extend(REGEX_ELLIPSIS.split(quote))
        return quote_parts

class PDFMatcher():
    """
    matches and generates sent tokens from pdf text
    """
    def __init__(self, quotes=None, pdftext=None):
        # load a sequence matcher; turn autojunk off (since buggy for long strings)
        self.sequencematcher = difflib.SequenceMatcher(None, autojunk=False)

        if quotes:
            self.quotes = self.load_quotes(quotes)
        if pdftext:
            self.pdftext = self.load_pdftext(pdftext)

    def load_quotes(self, quotes):
        self.quotes = quotes

    def load_pdftext(self, pdftext):
        self.pdftext = pdftext
        self.lenpdf = len(pdftext)
        self.sequencematcher.set_seq2(self.pdftext)
        self.sent_indices =  sent_tokenizer.span_tokenize(self.pdftext)

    def _overlap(self, t1, t2):
        """
        finds out whether two tuples overlap
        """
        t1s, t1e = t1
        t2s, t2e = t2

        # true if either start of t1 is inside t2, or start of t2 is inside t1
        return (t2s <= t1s <= t2e) or (t1s <= t2s <= t1e)

    def generate_X(self):
        X = []
        # go through sentence indices
        # make X (list of sentences)
        for (start_i, end_i) in self.sent_indices:
            X.append(self.pdftext[start_i: end_i])
        return X

    def generate_y(self, min_char_match=20):
        """
        returns X: list of sentence strings
                y: numpy vector of 1, -1 (for positive/negative examples)
        """
        good_match = False # this will be set to True if sufficent matching characters in
                           # at least one of the parts of the quotes

        match_indices = []

        # go through quotes, match using difflib
        # and keep any matches which are long enough so likely true matches
        for quote in self.quotes:

            self.sequencematcher.set_seq1(quote)

            best_match = self.sequencematcher.find_longest_match(0, len(quote), 0, self.lenpdf)

            # only interested in good quality matches
            if best_match.size > min_char_match:
                good_match = True
                match_indices.append((best_match.b, best_match.b + best_match.size)) # add (start_i, end_i) tuples (of PDF indices)

        y = []

        if not good_match:
            # if quality criteria not met, leave here
            # (i.e. return empty lists [], [])
            return y

        # otherwise continue and generate feature and answer vectors

        # get indices of sentences (rather than split)
        sent_indices = sent_tokenizer.span_tokenize(self.pdftext)

        # go through sentence indices
        # make X (list of sentences)
        # and calculate y, if there is *any* overlap with matched quoted text then
        # y = True
        for (start_i, end_i) in sent_indices:

            # if any overlaps with quotes, then y = True, else False
            if any((self._overlap((start_i, end_i), match_tuple) for match_tuple in match_indices)):
                y.append(1)
            else:
                y.append(-1)
        return y

class SentenceModel():
    """
    predicts whether sentences contain risk of bias information
    - uses data from Cochrane quotes only
    """

    def __init__(self, test_mode=False):
        self.quotereader = QualityQuoteReader2(test_mode=test_mode) # this now runs through all studies

    def generate_data(self, uid_filter=None):
        """
        tokenizes and processes the raw text from pdfs and cochrane
        saves in self.X_list and self.y_list (both dicts)

        """

        test_domains = CORE_DOMAINS # may change later to access various "other" domains

        # one feature matrix X across all domains
        self.X_list = [] # key will be unique id, value will be text
        # but multiple y vectors; one for each test domain
        self.y_list = {domain: [] for domain in test_domains}

        self.y_judgements = {domain: [] for domain in test_domains} # used in subclasses to make hybrid models

        self.X_uids = []
        self.y_uids = {domain: [] for domain in test_domains}

        for uid, cochrane_data, pdf_data in self.quotereader:

            if uid_filter is not None and uid not in uid_filter:
                continue

            matcher = PDFMatcher()
            matcher.load_pdftext(pdf_data)

            X_study = matcher.generate_X()

            self.X_list.extend(X_study)
            self.X_uids.extend([uid] * len(X_study))

            domains_done_already = [] # for this particular study
            # (we're ignoring multiple quotes per domain at the moment and getting the first...)

            for domain in cochrane_data["QUALITY"]:

                if domain["DOMAIN"] not in test_domains or domain["DOMAIN"] in domains_done_already:

                    continue # skip if a domain is repeated in a study (though note that this is likely due to different RoB per *outcome* which is ignored here)

                if domain["QUOTES"]:
                    matcher.load_quotes(domain["QUOTES"])
                    y_study = matcher.generate_y()

                    self.y_list[domain["DOMAIN"]].extend(y_study)
                    self.y_uids[domain["DOMAIN"]].extend([uid] * len(y_study))
                    self.y_judgements[domain["DOMAIN"]].extend([domain["RATING"]] * len(y_study))

                    domains_done_already.append(domain["DOMAIN"])

        self.y = {domain: np.array(self.y_list[domain]) for domain in test_domains}

        self.X_uids = np.array(self.X_uids)
        self.y_uids = {domain: np.array(self.y_uids[domain]) for domain in test_domains}
        self.y_judgements = {domain: np.array(self.y_judgements[domain]) for domain in test_domains}

        # self.vectorize()

    def vectorize(self):

        self.vectorizer = ModularCountVectorizer()
        # self.X = self.vectorizer.fit_transform(self.X_list, max_features=50000)
        self.X = self.vectorizer.fit_transform(self.X_list)

    def load_text(self, filename):
        """
        loads the original text of all PDFs, to debugging and looking at predicted text from corpus
        NB this is a large file
        """
        with open(filename, 'rb') as f:
            self.X_list = pickle.load(f)

    def __len__(self):
        """
        returns the total number of studies (not features)
        """
        return len(self.quotereader)

    def len_domain(self, domain):
        return len(np.unique(self.y_uids[domain]))

    def domain_X_filter(self, domain):
        """
        returns X_filter for a domain
        """
        y_study_uids = np.unique(self.y_uids[domain])
        X_filter = np.nonzero([(X_uid in y_study_uids) for X_uid in self.X_uids])[0]
        return X_filter

    def domain_uids(self, domain):
        unique_study_ids = np.unique(self.y_uids[domain])
        return unique_study_ids

    def X_y_uid_filtered(self, uids, domain):
        X_all = self.X_domain_all(domain=domain)
        y_all = self.y_domain_all(domain=domain)

        filter_ids = np.nonzero([(y_study_id in uids) for y_study_id in self.y_uids[domain]])[0]
        X_filtered = X_all[filter_ids]
        y_filtered = y_all[filter_ids]

        return X_filtered, y_filtered

    def get_all_domains(self):
        return self.y.keys()

    def X_get_sentence(self, select_sent_id, domain):

        y_study_ids = np.unique(self.y[domain].study_ids)
        X_filter = np.nonzero([X_study_id in y_study_ids for X_study_id in self.X.study_ids])[0]
        return self.X_list.data[X_filter[select_sent_id]]

    def X_domain_all(self, domain):
        """
        retrieve X data for a domain
        """

        X_filter = self.domain_X_filter(domain)
        return self.X[X_filter]

    def y_domain_all(self, domain):
        return self.y[domain]

    # def X_y_filtered(self, filter_ids, domain):

    #     X_all = self.X_domain_all(domain=domain)
    #     y_all = self.y_domain_all(domain=domain)

    #     # np.unique always returns ordered ids
    #     unique_study_ids = np.unique(self.y_uids[domain])

    #     mapped_ids = [unique_study_ids[filter_id] for filter_id in filter_ids]

    #     filter_ids = np.nonzero([(y_study_id in mapped_ids) for y_study_id in self.y_uids[domain]])[0]

    #     X_filtered = X_all[filter_ids]
    #     y_filtered = y_all[filter_ids]

    #     return X_filtered, y_filtered

class DocumentLevelModel(SentenceModel):
    """
    for predicting the risk of bias
    as "HIGH", "LOW", or "UNKNOWN" for a document
    using a binary bag-of-words as features for each document
    """

    def generate_data(self, uid_filter=None, binarize=False):
        """
        tokenizes and processes the raw text from pdfs and cochrane
        saves in self.X_list and self.y_list (both dicts)

        """

        test_domains = CORE_DOMAINS # may change later to access various "other" domains

        # one feature matrix X across all domains
        self.X_list = [] # key will be unique id, value will be text
        # but multiple y vectors; one for each test domain
        self.y_list = {domain: [] for domain in test_domains}

        self.X_uids = []
        self.y_uids = {domain: [] for domain in test_domains}

        for uid, cochrane_data, pdf_data in self.quotereader:

            if uid_filter is not None and uid not in uid_filter:
                continue

            X_study = [pdf_data] # this time the X is the whole pdf data

            self.X_list.extend(X_study)
            self.X_uids.extend([uid] * len(X_study))

            domains_done_already = [] # for this particular study
            # (we're ignoring multiple quotes per domain at the moment and getting the first...)

            for domain in cochrane_data["QUALITY"]:

                if domain["DOMAIN"] not in test_domains or domain["DOMAIN"] in domains_done_already:
                    continue # skip if a domain is repeated in a study (though note that this is likely due to different RoB per *outcome* which is ignored here)

                if binarize:
                    y_study = 1 if domain["RATING"]=="YES" else -1 # binarize
                else:
                    y_study = domain["RATING"]

                self.y_list[domain["DOMAIN"]].append(y_study)
                self.y_uids[domain["DOMAIN"]].append(uid)

                domains_done_already.append(domain["DOMAIN"])

        self.y = {domain: np.array(self.y_list[domain]) for domain in test_domains}

        self.X_uids = np.array(self.X_uids)
        self.y_uids = {domain: np.array(self.y_uids[domain]) for domain in test_domains}

class MultiTaskDocumentModel(DocumentLevelModel):
    '''
    The idea here is to train a single model across all domains. Basically:

        y_ij = sign{(w0 + w_j) * x_i}

    for document x_i, where a w_j is learned for each domain and w0 is a shared
    weight vector (across domains).
    '''

    def vectorize(self):

        self.vectorizer = ModularCountVectorizer()

        self.vectorizer.builder_clear()

        self.X_mt_labels = [] # which rows correspond to which doc/interactions?
        self.y_mt = []
        self.uids_to_row_indices = {}
        self.row_indices_to_uids, self.row_indices_to_domains = [], []
        # number of rows in the 'multi-task' matrix, which
        # will vary depending on how many docs have labels
        # for how many domains
        n_rows = 0 # (equal to len(self.X_mt_labels)

        '''
        the vectorizer wants all the documents at once,
        so here we are going to build them up. we're only
        going to add interaction copies for a given document
        for those domains that we have an associated label.
        '''
        docs = []
        # which indices in docs correspond to copies for
        # which domains?
        domains_to_interaction_doc_indices = defaultdict(list)
        for i, doc in enumerate(self.X_list):
            # `intercept' document
            uid = self.X_uids[i]

            # add |CORE_DOMAINS| copies for each instance.
            for domain in CORE_DOMAINS:
                d_str = domain_str(domain)
                if uid in self.domain_uids(domain):
                    # get the label (note that we match up the
                    # uid to do this)
                    uids_to_lbls = dict(zip(self.y_uids[domain],
                                        self.y_domain_all(domain=domain)))
                    #y_index = self.y_uids(domain).index(uid)
                    #domain_ys = self.y_domain_all(domain=domain)
                    #self.y_mt.append(domain_ys[y_index])
                    self.y_mt.append(uids_to_lbls[uid])

                    # append interaction copy of doc
                    docs.append(doc)
                    self.row_indices_to_uids.append(uid)
                    self.row_indices_to_domains.append(domain)
                    self.X_mt_labels.append("%s-%s" % (i, d_str))
                    domains_to_interaction_doc_indices[d_str].append(n_rows)

                    n_rows += 1

        '''
        now actually ad documents and interaction copies to
        the vectorizer.
        '''
        #for i, doc in enumerate(self.X_list):
        # `intercept' document
        self.vectorizer.builder_add_docs(docs)

        for domain in CORE_DOMAINS:
            d_str = domain_str(domain)
            interaction_list = []
            for i in xrange(len(docs)):
                if i in domains_to_interaction_doc_indices[d_str]:
                    interaction_list.append(docs[i])
                else:
                    interaction_list.append("")
            self.vectorizer.builder_add_docs(interaction_list, prefix=d_str+"-")

        # BCW -- attempting to upper bound features!
        # note that this will keep the <max_features> most common
        # features, regardless of whether or not they are 'interaction'
        # features
        # self.X = self.vectorizer.builder_fit_transform(max_features=50000)
        self.X = self.vectorizer.builder_fit_transform(low=2)

    def X_y_uid_filtered(self, uids, domain=None):
        X_indices, y = [], []
        for i in xrange(self.X.shape[0]):
            if domain is None and self.row_indices_to_uids[i] in uids:
                # if domain is None, return big multi-task design matrix
                # -- you'll want to do this, e.g., for training
                X_indices.append(i)
                y.append(self.y_mt[i])
            elif domain == self.row_indices_to_domains[i] and self.row_indices_to_uids[i] in uids:
                # otherwise (when domain is not None), return
                # instances for only the target domain
                # (e.g., for testing)
                X_indices.append(i)
                y.append(self.y_mt[i])

        return self.X[X_indices], y

class MultiTaskHybridDocumentModel(MultiTaskDocumentModel):
    '''
    same as the MultiTaskDocumentModel, except takes in sentence
    level modelling too into the mix
    '''

    def vectorize(self):

        self.vectorizer = ModularCountVectorizer()
        self.vectorizer.builder_clear()

        self.X_mt_labels = [] # which rows correspond to which doc/interactions?
        self.y_mt = []
        self.uids_to_row_indices = {}
        self.row_indices_to_uids, self.row_indices_to_domains = [], []
        # number of rows in the 'multi-task' matrix, which
        # will vary depending on how many docs have labels
        # for how many domains
        n_rows = 0 # (equal to len(self.X_mt_labels)

        '''
        the vectorizer wants all the documents at once,
        so here we are going to build them up. we're only
        going to add interaction copies for a given document
        for those domains that we have an associated label.
        '''
        docs = []
        high_prob_sents = defaultdict(list)
        # which indices in docs correspond to copies for
        # which domains?
        domains_to_interaction_doc_indices = defaultdict(list)
        for i, doc in enumerate(self.X_list):
            # `intercept' document
            uid = self.X_uids[i]

            # add |CORE_DOMAINS| copies for each instance.
            for domain in CORE_DOMAINS:
                d_str = domain_str(domain)
                if uid in self.domain_uids(domain):
                    # get the label (note that we match up the
                    # uid to do this)
                    uids_to_lbls = dict(zip(self.y_uids[domain],
                                        self.y_domain_all(domain=domain)))
                    #y_index = self.y_uids(domain).index(uid)
                    #domain_ys = self.y_domain_all(domain=domain)
                    #self.y_mt.append(domain_ys[y_index])
                    self.y_mt.append(uids_to_lbls[uid])

                    # append interaction copy of doc
                    docs.append(doc)

                    high_prob_sents[domain].append(self.get_sent_predictions_for_doc(doc, domain))

                    for high_prob_domain in CORE_DOMAINS:
                        if high_prob_domain != domain:
                            high_prob_sents[high_prob_domain].append("")

                    self.row_indices_to_uids.append(uid)
                    self.row_indices_to_domains.append(domain)
                    self.X_mt_labels.append("%s-%s" % (i, d_str))
                    domains_to_interaction_doc_indices[d_str].append(n_rows)

                    n_rows += 1

        '''
        now actually add documents and interaction copies to
        the vectorizer.
        '''
        #for i, doc in enumerate(self.X_list):
        # `intercept' document
        self.vectorizer.builder_add_docs(docs)

        for domain in CORE_DOMAINS:
            d_str = domain_str(domain)
            interaction_list, sent_interaction_list = [], []
            for i in xrange(len(docs)):
                if i in domains_to_interaction_doc_indices[d_str]:
                    interaction_list.append(docs[i])
                    sent_interaction_list.append(high_prob_sents[domain][i])
                else:
                    interaction_list.append("")
                    sent_interaction_list.append("")

            self.vectorizer.builder_add_docs(interaction_list, prefix=d_str+"-doc-")
            self.vectorizer.builder_add_docs(sent_interaction_list, prefix=d_str+"-sent-")

        self.X = self.vectorizer.builder_fit_transform(max_features=200000, low=3)
        # self.X = self.vectorizer.builder_fit_transform(max_features=50000)

        ####
        # maybe record the feature indices here that are to receive
        # different 'amounts' of regularization
        ####

    def get_sent_predictions_for_doc(self, doc, domain):

        # tokenize into sentences
        sents = sent_tokenizer.tokenize(doc)

        # vectorize the sentences
        X_sents = self.sent_vectorizer.transform(sents)

        # get predicted 1 / -1 for the sentences
        pred_class = self.sent_clfs[domain].predict(X_sents)

        # get the sentences which are predicted 1
        positive_sents = [sent for sent, pred in zip(sents, pred_class) if pred==1]

        # make a single string per doc
        rob_sents = " ".join(positive_sents)

        return rob_sents

    def set_sent_model(self, sent_clfs, sent_vectorizer):
        """
        set a model which takes in a list of sentences;
        and returns -1 or 1
        """
        self.sent_clfs = sent_clfs
        self.sent_vectorizer = sent_vectorizer

class HybridDocModel(DocumentLevelModel):
    """
    for predicting the risk of bias
    as "HIGH", "LOW", or "UNKNOWN" for a document
    using a binary bag-of-words as features for each document
    """

    def vectorize(self, domain=None):

        if domain is None:
            raise TypeError("this class requires domain specific vectorization")

        self.vectorizer = ModularCountVectorizer()
        self.vectorizer.builder_clear()

        X_filter = self.domain_X_filter(domain)
        predictions = self.get_sent_predictions_for_domain(domain)

        self.vectorizer.builder_add_docs([self.X_list[i] for i in X_filter])
        self.vectorizer.builder_add_docs(predictions, prefix="high-prob-sent-", weighting=10)

        self.X = self.vectorizer.builder_fit_transform()

    def get_sent_predictions_for_domain(self, domain):

        uids = self.domain_uids(domain)

        predictions = []

        for uid in uids:

            # get the index of the study with specified uid
            study_index = np.nonzero(self.X_uids==uid)[0][0]

            # tokenize into sentences
            sents = sent_tokenizer.tokenize(self.X_list[study_index])

            # vectorize the sentences
            X_sents = self.sent_vectorizer.transform(sents)

            # get predicted 1 / -1 for the sentences
            pred_class = self.sent_clf.predict(X_sents)

            # get the sentences which are predicted 1
            positive_sents = [sent for sent, pred in zip(sents, pred_class) if pred==1]

            # make a single string per doc
            doc = " ".join(positive_sents)

            predictions.append(doc)

        return predictions

    def set_sent_model(self, doc_clf, doc_vectorizer):
        """
        set a model which takes in a list of sentences;
        and returns -1 or 1
        """
        self.sent_clf = doc_clf
        self.sent_vectorizer = doc_vectorizer

    def X_y_uid_filtered(self, uids, domain):
        X_all = self.X
        y_all = self.y_domain_all(domain=domain)

        filter_ids = np.nonzero([(y_study_id in uids) for y_study_id in self.y_uids[domain]])[0]
        X_filtered = X_all[filter_ids]
        y_filtered = y_all[filter_ids]

        return X_filtered, y_filtered

# class HybridDocModel2(HybridDocModel):
#     """
#     for predicting the risk of bias
#     as "HIGH", "LOW", or "UNKNOWN" for a document
#     using a binary bag-of-words as features for each document
#     """

#     def vectorize(self, domain=None):

#         if domain is None:
#             raise TypeError("this class requires domain specific vectorization")

#         self.vectorizer = ModularCountVectorizer()
#         self.vectorizer.builder_clear()
#         X_filter = self.domain_X_filter(domain)

#         predictions = self.get_sent_predictions_for_domain(domain)

#         # self.vectorizer.builder_add_docs([self.X_list[i] for i in X_filter])

#         self.vectorizer.builder_add_docs(predictions, prefix="high-prob-sent-")

#         self.X = self.vectorizer.builder_fit_transform()

class HybridModel(SentenceModel):
    """
    predicts whether sentences contain risk of bias information
    - uses real RoB judgements
    """

    def vectorize(self, domain=None, interaction_classes=["YES", "NO"]):

        if domain is None:
            raise TypeError("this class requires domain specific vectorization")

        self.vectorizer = ModularCountVectorizer()

        self.vectorizer.builder_clear()

        X_filter = self.domain_X_filter(domain)

        sents = [self.X_list[i] for i in X_filter]

        self.vectorizer.builder_add_docs(sents)

        for interaction_class in interaction_classes:

            self.vectorizer.builder_add_interaction_features(
                sents, self.y_judgements[domain]==interaction_class, prefix="rob-" + interaction_class + "-")

        self.X = self.vectorizer.builder_fit_transform()

    def X_y_uid_filtered(self, uids, domain):
        X_all = self.X
        y_all = self.y_domain_all(domain=domain)

        filter_ids = np.nonzero([(y_study_id in uids) for y_study_id in self.y_uids[domain]])[0]
        X_filtered = X_all[filter_ids]
        y_filtered = y_all[filter_ids]

        return X_filtered, y_filtered

class HybridModelProbablistic(HybridModel):
    """
    predicts whether sentences contain risk of bias information
    - requires a model to be passed in which can predice RoB judgements from
    full text document
    """

    def vectorize(self, domain=None, interaction_classes=["YES", "NO"], use_vectorizer=None):

        if domain is None:
            raise TypeError("this class requires domain specific vectorization")

        if use_vectorizer is None:
            self.vectorizer = ModularCountVectorizer()
        else:
            self.vectorizer = use_vectorizer

        self.vectorizer.builder_clear()
        X_filter = self.domain_X_filter(domain)

        predictions = self.get_doc_predictions_for_domain(domain)

        sents = [self.X_list[i] for i in X_filter]
        self.vectorizer.builder_add_docs(sents)

        for interaction_class in interaction_classes:
            self.vectorizer.builder_add_interaction_features(sents, predictions==interaction_class, prefix="rob-" + interaction_class + "-")

        if use_vectorizer is None:
            self.X = self.vectorizer.builder_fit_transform()
        else:
            self.X = self.vectorizer.builder_transform()

    def get_doc_predictions_for_domain(self, domain):

        uids = self.domain_uids(domain)

        predictions = []

        for uid in uids:

            # get the indices of all sentences in the study with specified uid
            X_filter = np.nonzero(self.X_uids==uid)[0]

            # make a single string per doc
            doc = " ".join([self.X_list[i] for i in X_filter])
            # vectorize the docs, then predict using the model
            X_doc = self.doc_vectorizer.transform(doc)
            prediction = self.doc_clf.predict(X_doc)

            # add the same prediction for each sentence
            predictions.extend([prediction[0]] * len(X_filter))

        return np.array(predictions)

    def set_doc_model(self, doc_clf, doc_vectorizer):
        """
        set a model which takes in a full text doc;
        outputs a doc class "YES", "NO", or "UNKNOWN"
        """
        self.doc_clf = doc_clf
        self.doc_vectorizer = doc_vectorizer

def _document_frequency(X):
    """Count the number of non-zero values for each feature in csc_matrix X."""
    return np.diff(X.indptr)

class ModularCountVectorizer():
    """
    Similar to CountVectorizer from sklearn, but allows building up
    of feature matrix gradually, and adding prefixes to feature names
    (to identify interaction terms)
    """

    def __init__(self, *args, **kwargs):
        self.data = []
        self.vectorizer = DictVectorizer(*args, **kwargs)

    def _transform_X_to_dict(self, X, prefix=None, weighting=1):
        """
        makes a list of dicts from a document
        1. word tokenizes
        2. creates {word1:1, word2:1...} dicts
        (note all set to '1' since the DictVectorizer we use assumes all missing are 0)
        """
        return [self._dict_from_word_list(
            self._word_tokenize(document, prefix=prefix), weighting=1) for document in X]

    def _word_tokenize(self, text, prefix=None, stopword=True):
        """
        simple word tokenizer using the same rule as sklearn
        punctuation is ignored, all 2 or more letter characters are a word
        """

        stop_word_list = STOP_WORDS if stopword else set()

        if prefix:
            return [prefix + word.lower() for word in SIMPLE_WORD_TOKENIZER.findall(text)
                        if not word.lower() in stop_word_list]
        else:
            return [word.lower() for word in SIMPLE_WORD_TOKENIZER.findall(text)
                        if not word.lower() in stop_word_list]

    def _dict_from_word_list(self, word_list, weighting=1):
        return {word: weighting for word in word_list}

    def _dictzip(self, dictlist1, dictlist2):
        """
        zips together two lists of dicts of the same length
        """
        # checks lists must be the same length
        if len(dictlist1) != len(dictlist2):
            raise IndexError("Unable to combine featuresets with different number of examples")

        output = []

        for dict1, dict2 in zip(dictlist1, dictlist2):
            output.append(dict(dict1.items() + dict2.items()))
            # note that this *overwrites* any duplicate keys with the key/value from dictlist2!!

        return output

    def transform(self, X, prefix=None):
        # X is a list of document strings
        # word tokenizes each one, then passes to a dict vectorizer
        dict_list = self._transform_X_to_dict(X, prefix=prefix)
        return self.vectorizer.transform(dict_list)

    def fit_transform(self, X, prefix=None, max_features=None, low=None):
        # X is a list of document strings
        # word tokenizes each one, then passes to a dict vectorizer
        dict_list = self._transform_X_to_dict(X, prefix=prefix)
        X = self.vectorizer.fit_transform(dict_list)

        if max_features is not None or low is not None:
            X, removed = self._limit_features(X.tocsc(),
                        self.vectorizer.vocabulary_, low=low, limit=max_features)
            print "pruned %s features!" % len(removed)
            X = X.tocsc()

        return self.vectorizer.fit_transform(dict_list)

    def get_feature_names(self):
        return self.vectorizer.get_feature_names()

    def builder_clear(self):
        self.builder = []
        self.builder_len = 0

    def builder_add_docs(self, X, prefix = None, weighting=1):
        #pdb.set_trace()
        if not self.builder:
            self.builder_len = len(X)
            self.builder = self._transform_X_to_dict(X, prefix=prefix, weighting=weighting)
        else:
            X_dicts = self._transform_X_to_dict(X, prefix=prefix, weighting=weighting)
            self.builder = self._dictzip(self.builder, X_dicts)

    def builder_add_interaction_features(self, X, interactions, prefix=None):
        if prefix is None:
            raise TypeError('Prefix is required when adding interaction features')

        doc_list = [(sent if interacting else "") for sent, interacting in zip(X, interactions)]
        self.builder_add_docs(doc_list, prefix)

    def builder_fit_transform(self, max_features=None, low=2):
        X = self.vectorizer.fit_transform(self.builder)
        if max_features is not None or low is not None:
            X, removed = self._limit_features(X.tocsc(),
                        self.vectorizer.vocabulary_, low=low, limit=max_features)
            print "pruned %s features!" % len(removed)
            X = X.tocsc()

        return X #self.vectorizer.fit_transform(self.builder)

    def builder_transform(self):
        return self.vectorizer.transform(self.builder)

    def _limit_features(self, cscmatrix, vocabulary, high=None, low=None,
                        limit=None):
        """Remove too rare or too common features.

        Prune features that are non zero in more samples than high or less
        documents than low, modifying the vocabulary, and restricting it to
        at most the limit most frequent.

        This does not prune samples with zero features.
        """
        if high is None and low is None and limit is None:
            return cscmatrix, set()

        # Calculate a mask based on document frequencies
        dfs = _document_frequency(cscmatrix)
        mask = np.ones(len(dfs), dtype=bool)
        if high is not None:
            mask &= dfs <= high
        if low is not None:
            mask &= dfs >= low
        if limit is not None and mask.sum() > limit:
            # backward compatibility requires us to keep lower indices in ties!
            # (and hence to reverse the sort by negating dfs)
            mask_inds = (-dfs[mask]).argsort()[:limit]
            new_mask = np.zeros(len(dfs), dtype=bool)
            new_mask[np.where(mask)[0][mask_inds]] = True
            mask = new_mask

        new_indices = np.cumsum(mask) - 1  # maps old indices to new
        removed_terms = set()
        for term, old_index in list(six.iteritems(vocabulary)):
            if mask[old_index]:
                vocabulary[term] = new_indices[old_index]
            else:
                del vocabulary[term]
                removed_terms.add(term)
        kept_indices = np.where(mask)[0]

        return cscmatrix[:, kept_indices], removed_terms

def sentence_prediction_test(class_weight={1: 5, -1:1}, model=SentenceModel(test_mode=True)):
    print
    print
    print

    print "Sentence level prediction"
    print "=" * 40
    print

    s = model

    print "Model name:\t" + s.__class__.__name__
    print s.__doc__

    print "class_weight=%s" % (str(class_weight),)

    s.generate_data()
    s.vectorize()

    for test_domain in CORE_DOMAINS:
        print ("*"*40) + "\n\n" + test_domain + "\n\n" + ("*" * 40)

        domain_uids = s.domain_uids(test_domain)
        no_studies = len(domain_uids)

        kf = KFold(no_studies, n_folds=5, shuffle=False, indices=True)

        # # tuned_parameters = {"alpha": np.logspace(-4, -1, 10)}
        # tuned_parameters = [{"alpha": np.logspace(-4, -1, 5)}, {"class_weight": [{1: i, -1: 1} for i in np.logspace(0, 1, 5)]}]
        # clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring='recall')

        print "making scorer"
        ftwo_scorer = make_scorer(fbeta_score, beta=2)

        tuned_parameters = [{"alpha": np.logspace(-4, -1, 10)}, {"class_weight": [{1: i, -1: 1} for i in np.logspace(0, 1, 10)]}]
        clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring=ftwo_scorer)

        metrics = []

        for fold_i, (train, test) in enumerate(kf):

            X_train, y_train = s.X_y_uid_filtered(domain_uids[train], test_domain)
            X_test, y_test = s.X_y_uid_filtered(domain_uids[test], test_domain)

            clf.fit(X_train, y_train)

            y_preds = clf.predict(X_test)

            fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds))[:,1]

            metrics.append(fold_metric) # get the scores for positive instances

            print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])

            # if not sample and list_features:
            #     # not an obvious way to get best features for ensemble
            #     print show_most_informative_features(s.vectorizer, clf)

        # summary score

        summary_metrics = np.mean(metrics, axis=0)
        print "=" * 40
        print "mean score:\tprecision %.2f, recall %.2f, f-score %.2f" % (summary_metrics[0], summary_metrics[1], summary_metrics[2])

        # then train all for most informative features
        clf = SGDClassifier(loss="hinge", penalty="L2", alpha=0.01, class_weight={1: 5, -1: 1})
        X_all = s.X_domain_all(test_domain)
        y_all = s.y_domain_all(test_domain)

        clf.fit(X_all, y_all)

        print show_most_informative_features(s.vectorizer, clf)

def stupid_sentence_prediction_test(model=SentenceModel(test_mode=False)):
    print
    print
    print

    print "Sentence level prediction"
    print "=" * 40
    print

    s = model

    print "Model name:\t" + s.__class__.__name__
    print s.__doc__

    # print "class_weight=%s" % (str(class_weight),)

    s.generate_data()
    s.vectorize()

    for test_domain in CORE_DOMAINS:
        print ("*"*40) + "\n\n" + test_domain + "\n\n" + ("*" * 40)

        domain_uids = s.domain_uids(test_domain)
        no_studies = len(domain_uids)

        kf = KFold(no_studies, n_folds=5, shuffle=False, indices=True)

        print "making scorer"

        metrics = []

        for fold_i, (train, test) in enumerate(kf):

            X_test, y_test = s.X_y_uid_filtered(domain_uids[test], test_domain)

            y_preds = np.array([1] * len(y_test))

            fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds))[:,1]

            fold_metric = np.append(fold_metric, accuracy_score(y_test, y_preds))

            metrics.append(fold_metric) # get the scores for positive instances

            print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f, precision %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2], fold_metric[3])

            # if not sample and list_features:
            #     # not an obvious way to get best features for ensemble
            #     print show_most_informative_features(s.vectorizer, clf)

        # summary score

        summary_metrics = np.mean(metrics, axis=0)
        print "=" * 40
        print "mean score:\tprecision %.2f, recall %.2f, f-score %.2f, precision %.2f" % (summary_metrics[0], summary_metrics[1], summary_metrics[2], summary_metrics[3])

def binary_doc_prediction_test(model=DocumentLevelModel(test_mode=False)):
    print
    print
    print

    print "Binary doc prediction"
    print "=" * 40
    print

    s = model

    s.generate_data(binarize=True)
    s.vectorize()

    for test_domain in CORE_DOMAINS:
        print ("*"*40) + "\n\n" + test_domain + "\n\n" + ("*" * 40)

        domain_uids = s.domain_uids(test_domain)
        no_studies = len(domain_uids)

        kf = KFold(no_studies, n_folds=5, shuffle=False, indices=True)

        # # tuned_parameters = {"alpha": np.logspace(-4, -1, 10)}
        # tuned_parameters = [{"alpha": np.logspace(-4, -1, 5)}, {"class_weight": [{1: i, -1: 1} for i in np.logspace(0, 1, 5)]}]
        # clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring='recall')

        # print "making scorer"
        # ftwo_scorer = make_scorer(fbeta_score, beta=2)
        tuned_parameters = {"alpha": np.logspace(-4, -1, 10), "class_weight": [{1: i, -1: 1} for i in np.logspace(-1, 1, 10)]}
        clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring="f1")

        metrics = []

        for fold_i, (train, test) in enumerate(kf):

            X_train, y_train = s.X_y_uid_filtered(domain_uids[train], test_domain)
            X_test, y_test = s.X_y_uid_filtered(domain_uids[test], test_domain)

            clf.fit(X_train, y_train)

            y_preds = clf.predict(X_test)

            fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds))[:,1]

            metrics.append(fold_metric) # get the scores for positive instances

            print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])

            # if not sample and list_features:
            #     # not an obvious way to get best features for ensemble
            #     print show_most_informative_features(s.vectorizer, clf)

        # summary score

        summary_metrics = np.mean(metrics, axis=0)
        print "=" * 40
        print "mean score:\tprecision %.2f, recall %.2f, f-score %.2f" % (summary_metrics[0], summary_metrics[1], summary_metrics[2])

        # then train all for most informative features
        clf = SGDClassifier(loss="hinge", penalty="L2", alpha=0.01, class_weight={1: 5, -1: 1})
        X_all = s.X_domain_all(test_domain)
        y_all = s.y_domain_all(test_domain)

        clf.fit(X_all, y_all)

        print show_most_informative_features(s.vectorizer, clf)

def multitask_document_prediction_test(model=MultiTaskDocumentModel(test_mode=False),
                                        test_domain=CORE_DOMAINS[0]):
    print "multitask!"
    d = model
    d.generate_data(binarize=True) # some variations use the quote data internally
                      # for sentence prediction (for additional features)

    # d.X_uids contains all the UIds.
    # d.y_uids contains a dictionary mapping domains to the UIds for
    # which we have labels (in said domain)
    #pdb.set_trace()
    all_uids = d.X_uids
    d.vectorize()

    ####
    # the major change here is we don't loop through the domains!
    tuned_parameters = {"alpha": np.logspace(-4, -1, 10)}
    clf = GridSearchCV(SGDClassifier(loss="log", penalty="L2"),
                                tuned_parameters, scoring='f1')

    kf = KFold(len(all_uids), n_folds=5, shuffle=False) ### TODO 250 is totally random
    metrics = defaultdict(list)

    for fold_i, (train, test) in enumerate(kf):

        print "Training on fold", fold_i,
        # note that we do *not* pass in a domain here, because
        # we use *all* domain training data
        X_train, y_train = d.X_y_uid_filtered(all_uids[train])
        print "done!"

        clf.fit(X_train, y_train)

        print "Testing on fold", fold_i,
        for domain in CORE_DOMAINS:
            # multitask uses same trained model for all domains, but test on separate test data
            X_test, y_test = d.X_y_uid_filtered(all_uids[test], domain)

            y_preds = clf.predict(X_test)

            fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds))[:,1]
            metrics[domain].append(fold_metric) # get the scores for positive instances (save them up since all in the wrong order here!)
        print "done!"

    # then recreate in the right order for printout
    for domain in CORE_DOMAINS:

        print
        print domain
        print "*" * 60
        print

        for fold_i, fold_metric in enumerate(metrics[domain]):
            print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])
        # summary score

        summary_metrics = np.mean(metrics[domain], axis=0)
        print "=" * 40
        print "mean score:\tprecision %.2f, recall %.2f, f-score %.2f" % (summary_metrics[0], summary_metrics[1], summary_metrics[2])

def dummy_document_prediction():

    print "dummy!"

    d = DocumentLevelModel(test_mode=False)
    d.generate_data(binarize=True) # some variations use the quote data internally
                      # for sentence prediction (for additional features)

    d.vectorize()
    all_uids = d.X_uids

    kf = KFold(len(all_uids), n_folds=5, shuffle=False) ### TODO 250 is totally random
    metrics = defaultdict(list)

    for fold_i, (train, test) in enumerate(kf):

        print "Testing on fold", fold_i,
        for domain in CORE_DOMAINS:
            # multitask uses same trained model for all domains, but test on separate test data

            X_test, y_test = d.X_y_uid_filtered(all_uids[test], domain)

            y_preds = np.array(([1] * len(y_test))) # everything gets marked low risk of bias

            fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds))[:,1]
            metrics[domain].append(fold_metric) # get the scores for positive instances (save them up since all in the wrong order here!)
        print "done!"

    # then recreate in the right order for printout
    for domain in CORE_DOMAINS:

        print
        print domain
        print "*" * 60
        print

        for fold_i, fold_metric in enumerate(metrics[domain]):
            print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])
        # summary score

        summary_metrics = np.mean(metrics[domain], axis=0)
        print "=" * 40
        print "mean score:\tprecision %.2f, recall %.2f, f-score %.2f" % (summary_metrics[0], summary_metrics[1], summary_metrics[2])

def multitask_hybrid_document_prediction_test(model=MultiTaskHybridDocumentModel(test_mode=True)):

    print "multitask! and hybrid! :)"
    d = model
    d.generate_data(binarize=True) # some variations use the quote data internally
                      # for sentence prediction (for additional features)

    # d.X_uids contains all the UIds.
    # d.y_uids contains a dictionary mapping domains to the UIds for
    # which we have labels (in said domain)
    #pdb.set_trace()
    all_uids = d.X_uids
    # d.vectorize()

    ####
    # the major change here is we don't loop through the domains!
    tuned_parameters = {"alpha": np.logspace(-4, -1, 10)}
    clf = GridSearchCV(SGDClassifier(loss="log", penalty="L2"),
                                tuned_parameters, scoring='f1')

    kf = KFold(len(all_uids), n_folds=5, shuffle=False)
    metrics = defaultdict(list)

    print "...generating sentence data...",
    s = SentenceModel(test_mode=False)
    s.generate_data()
    s.vectorize()
    print "done!"

    sent_tuned_parameters = [{"alpha": np.logspace(-4, -1, 5)}, {"class_weight": [{1: i, -1: 1} for i in np.logspace(0, 2, 10)]}]

    for fold_i, (train, test) in enumerate(kf):

        sent_clfs = defaultdict(list)

        for domain in CORE_DOMAINS:
            sents_X, sents_y = s.X_domain_all(domain=domain), s.y_domain_all(domain=domain)

            sent_clfs[domain] = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"),
                                                tuned_parameters, scoring='recall')

        print "Training on fold", fold_i,

        d.set_sent_model(sent_clfs, s.vectorizer)
        d.vectorize()

        # note that we do *not* pass in a domain here, because
        # we use *all* domain training data
        X_train, y_train = d.X_y_uid_filtered(all_uids[train])
        sent_clfs[domain].fit(sents_X, sents_y)

        clf.fit(X_train, y_train)
        print "done!"

        print "Testing on fold", fold_i,
        for domain in CORE_DOMAINS:
            # multitask uses same trained model for all domains, but test on separate test data
            X_test, y_test = d.X_y_uid_filtered(all_uids[test], domain)

            y_preds = clf.predict(X_test)

            fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds))[:,1]
            metrics[domain].append(fold_metric) # get the scores for positive instances (save them up since all in the wrong order here!)
        print "done!"

    # then recreate in the right order for printout
    for domain in CORE_DOMAINS:

        print
        print domain
        print "*" * 60
        print

        for fold_i, fold_metric in enumerate(metrics[domain]):
            print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])
        # summary score

        summary_metrics = np.mean(metrics[domain], axis=0)
        print "=" * 40
        print "mean score:\tprecision %.2f, recall %.2f, f-score %.2f" % (summary_metrics[0], summary_metrics[1], summary_metrics[2])

def document_prediction_test(model=DocumentLevelModel(test_mode=False)):

    print "Document level prediction"
    print "=" * 40
    print

    d = model
    d.generate_data() # some variations use the quote data internally
                                                # for sentence prediction (for additional features)

    d.vectorize()

    for test_domain in CORE_DOMAINS:
        print ("*"*40) + "\n\n" + test_domain + "\n\n" + ("*" * 40)

        # f1_prefer_nos = make_scorer(f1_score, pos_label="NO")

        tuned_parameters = {"alpha": np.logspace(-4, -1, 10)}
        clf = GridSearchCV(SGDClassifier(loss="log", penalty="L2"), tuned_parameters, scoring='f1')

        # clf = SGDClassifier(loss="hinge", penalty="L2")

        domain_uids = d.domain_uids(test_domain)
        no_studies = len(domain_uids)

        kf = KFold(no_studies, n_folds=5, shuffle=False)

        metrics = []

        for fold_i, (train, test) in enumerate(kf):

            X_train, y_train = d.X_y_uid_filtered(domain_uids[train], test_domain)
            X_test, y_test = d.X_y_uid_filtered(domain_uids[test], test_domain)

            clf.fit(X_train, y_train)

            y_preds = clf.predict(X_test)

            fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds, labels=RoB_CLASSES))[:3]

            print ('fold %d\t' % (fold_i)) + '\t'.join(RoB_CLASSES)

            # for metric_type, scores in zip(["prec.", "recall", "f1"], fold_metric):
            #     print "%s\t%.2f\t%.2f\t%.2f" % (metric_type, scores[0], scores[1], scores[2])

            # print

            # print clf.best_params_

            #### START CONFUSION

            real_no_indices = (y_test=="NO")
            print "The actual NOs were predicted as..."
            print collections.Counter(y_preds[real_no_indices])

            #### END CONFUSION

            metrics.append(fold_metric) # get the scores for positive instances

            # print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])

        mean_scores = np.mean(metrics, axis=0)

        print "=" * 40
        print 'means \t' + '\t'.join(RoB_CLASSES)

        for metric_type, scores in zip(["prec.", "recall", "f1"], mean_scores):
            print "%s\t%.2f\t%.2f\t%.2f" % (metric_type, scores[0], scores[1], scores[2])
        print

        # then train all for most informative features
        clf = SGDClassifier(loss="hinge", penalty="L2", alpha=0.01)
        X_all = d.X_domain_all(test_domain)
        y_all = d.y_domain_all(test_domain)

        clf.fit(X_all, y_all)

        print show_most_informative_features_ynu(d.vectorizer, clf)

def simple_hybrid_prediction_test(model=HybridModel(test_mode=True)):

    print "Hybrid prediction"
    print "=" * 40
    print

    s = model
    s.generate_data() # some variations use the quote data internally
                                                # for sentence prediction (for additional features)

    for test_domain in CORE_DOMAINS:

        s.vectorize(test_domain)

        print ("*"*40) + "\n\n" + test_domain + "\n\n" + ("*" * 40)

        domain_uids = s.domain_uids(test_domain)
        no_studies = len(domain_uids)

        kf = KFold(no_studies, n_folds=5, shuffle=False)

        # tuned_parameters = [{"alpha": np.logspace(-4, -1, 5)}, {"class_weight":  [{1: i, -1: 1} for i in np.logspace(0, 1, 5)]}]
        # clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring='f1')

        print "making scorer"
        ftwo_scorer = make_scorer(fbeta_score, beta=2)

        tuned_parameters = [{"alpha": np.logspace(-4, -1, 10)}, {"class_weight": [{1: i, -1: 1} for i in np.logspace(0, 1, 10)]}]
        clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring=ftwo_scorer)

        metrics = []

        for fold_i, (train, test) in enumerate(kf):

            X_train, y_train = s.X_y_uid_filtered(domain_uids[train], test_domain)
            X_test, y_test = s.X_y_uid_filtered(domain_uids[test], test_domain)

            clf.fit(X_train, y_train)

            y_preds = clf.predict(X_test)

            fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds))[:,1]

            metrics.append(fold_metric) # get the scores for positive instances

            print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])

            metrics.append(fold_metric) # get the scores for positive instances

        # summary score

        summary_metrics = np.mean(metrics, axis=0)
        print "=" * 40
        print "mean score:\tprecision %.2f, recall %.2f, f-score %.2f" % (summary_metrics[0], summary_metrics[1], summary_metrics[2])

# def simple_hybrid_prediction_test(model=HybridModel(test_mode=True)):

#     print "Hybrid prediction"
#     print "=" * 40
#     print

#     s = model
#     s.generate_data() # some variations use the quote data internally
#                                                 # for sentence prediction (for additional features)

#     for test_domain in CORE_DOMAINS:

#         s.vectorize(test_domain)

#         print ("*"*40) + "\n\n" + test_domain + "\n\n" + ("*" * 40)

#         domain_uids = s.domain_uids(test_domain)
#         no_studies = len(domain_uids)

#         kf = KFold(no_studies, n_folds=5, shuffle=False)

#         tuned_parameters = [{"alpha": np.logspace(-4, -1, 5)}, {"class_weight":  [{1: i, -1: 1} for i in np.logspace(0, 1, 5)]}]
#         clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring='f1')

#         metrics = []

#         for fold_i, (train, test) in enumerate(kf):

#             X_train, y_train = s.X_y_uid_filtered(domain_uids[train], test_domain)
#             X_test, y_test = s.X_y_uid_filtered(domain_uids[test], test_domain)

#             clf.fit(X_train, y_train)

#             y_preds = clf.predict(X_test)

#             fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds))[:,1]

#             metrics.append(fold_metric) # get the scores for positive instances

#             print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])

#             metrics.append(fold_metric) # get the scores for positive instances

#         # summary score

#         summary_metrics = np.mean(metrics, axis=0)
#         print "=" * 40
#         print "mean score:\tprecision %.2f, recall %.2f, f-score %.2f" % (summary_metrics[0], summary_metrics[1], summary_metrics[2])

def true_hybrid_prediction_test(model, test_mode=False):

    print "True Hybrid prediction"
    print "=" * 40
    print

    s = model
    s.generate_data() # some variations use the quote data internally
                                                # for sentence prediction (for additional features)

    s_cheat = HybridModel(test_mode=False)
    s_cheat.generate_data()

    for test_domain in CORE_DOMAINS:

        print ("*"*40) + "\n\n" + test_domain + "\n\n" + ("*" * 40)

        domain_uids = s.domain_uids(test_domain)
        no_studies = len(domain_uids)
        kf = KFold(no_studies, n_folds=5, shuffle=False)

        print "making scorer"
        ftwo_scorer = make_scorer(fbeta_score, beta=2)

        tuned_parameters = [{"alpha": np.logspace(-4, -1, 10)}, {"class_weight": [{1: i, -1: 1} for i in np.logspace(0, 1, 10)]}]
        clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring=ftwo_scorer)

        metrics = []

        for fold_i, (train, test) in enumerate(kf):

            print "training doc level model with test data, please wait..."

            d = DocumentLevelModel(test_mode=False)
            d.generate_data(uid_filter=domain_uids[train])
            d.vectorize()
            doc_X, doc_y = d.X_domain_all(domain=test_domain), d.y_domain_all(domain=test_domain)

            doc_tuned_parameters = {"alpha": np.logspace(-4, -1, 10)}
            doc_clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), doc_tuned_parameters, scoring='f1')

            doc_clf.fit(doc_X, doc_y)

            s.set_doc_model(doc_clf, d.vectorizer)

            s_cheat.vectorize(test_domain)
            s.vectorize(test_domain, use_vectorizer=s_cheat.vectorizer)

            X_train, y_train = s_cheat.X_y_uid_filtered(domain_uids[train], test_domain)
            # train on the *true* labels

            X_test, y_test = s.X_y_uid_filtered(domain_uids[test], test_domain)

            clf.fit(X_train, y_train)

            y_preds = clf.predict(X_test)

            fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds))[:,1]

            metrics.append(fold_metric) # get the scores for positive instances

            print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])

            metrics.append(fold_metric) # get the scores for positive instances

        # summary score

        summary_metrics = np.mean(metrics, axis=0)
        print "=" * 40
        print "mean score:\tprecision %.2f, recall %.2f, f-score %.2f" % (summary_metrics[0], summary_metrics[1], summary_metrics[2])

def hybrid_doc_prediction_test(model=HybridDocModel(test_mode=True)):

    print "Hybrid doc level prediction"
    print "=" * 40
    print

    d = model
    d.generate_data() # some variations use the quote data internally
                                                # for sentence prediction (for additional features)

    for test_domain in CORE_DOMAINS:
        print ("*"*40) + "\n\n" + test_domain + "\n\n" + ("*" * 40)

        domain_uids = d.domain_uids(test_domain)
        no_studies = len(domain_uids)
        kf = KFold(no_studies, n_folds=5, shuffle=False)
        tuned_parameters = {"alpha": np.logspace(-4, -1, 5)}
        clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring='f1')

        metrics = []

        for fold_i, (train, test) in enumerate(kf):

            s = SentenceModel(test_mode=False)
            s.generate_data(uid_filter=domain_uids[train])
            s.vectorize()
            sents_X, sents_y = s.X_domain_all(domain=test_domain), s.y_domain_all(domain=test_domain)

            sent_tuned_parameters = [{"alpha": np.logspace(-4, -1, 5)}, {"class_weight": [{1: i, -1: 1} for i in np.logspace(0, 2, 10)]}]
            sent_clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring='recall')
            sent_clf.fit(sents_X, sents_y)
            d.set_sent_model(sent_clf, s.vectorizer)
            d.vectorize(test_domain)

            X_train, y_train = d.X_y_uid_filtered(domain_uids[train], test_domain)
            X_test, y_test = d.X_y_uid_filtered(domain_uids[test], test_domain)

            clf.fit(X_train, y_train)

            y_preds = clf.predict(X_test)

            fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds, labels=RoB_CLASSES))[:3]

            print ('fold %d\t' % (fold_i)) + '\t'.join(RoB_CLASSES)

            for metric_type, scores in zip(["prec.", "recall", "f1"], fold_metric):
                print "%s\t%.2f\t%.2f\t%.2f" % (metric_type, scores[0], scores[1], scores[2])

            print

            metrics.append(fold_metric) # get the scores for positive instances

            # print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])

        mean_scores = np.mean(metrics, axis=0)

        print "=" * 40
        print 'means \t' + '\t'.join(RoB_CLASSES)

        for metric_type, scores in zip(["prec.", "recall", "f1"], mean_scores):
            print "%s\t%.2f\t%.2f\t%.2f" % (metric_type, scores[0], scores[1], scores[2])
        print

def binary_hybrid_doc_prediction_test(model=HybridDocModel(test_mode=True)):

    print "Binary hybrid doc level prediction"
    print "=" * 40
    print

    d = model
    d.generate_data(binarize=True) # some variations use the quote data internally
                                                # for sentence prediction (for additional features)

    for test_domain in CORE_DOMAINS:

        print ("*"*40) + "\n\n" + test_domain + "\n\n" + ("*" * 40)

        domain_uids = d.domain_uids(test_domain)
        no_studies = len(domain_uids)
        kf = KFold(no_studies, n_folds=5, shuffle=False)
        tuned_parameters = {"alpha": np.logspace(-4, -1, 10), "class_weight": [{1: i, -1: 1} for i in np.logspace(-1, 1, 10)]}
        clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring='f1')

        metrics = []

        for fold_i, (train, test) in enumerate(kf):

            s = SentenceModel(test_mode=True)
            s.generate_data(uid_filter=domain_uids[train])
            s.vectorize()
            sents_X, sents_y = s.X_domain_all(domain=test_domain), s.y_domain_all(domain=test_domain)

            sent_tuned_parameters = [{"alpha": np.logspace(-4, -1, 5)}, {"class_weight": [{1: i, -1: 1} for i in np.logspace(0, 2, 10)]}]
            sent_clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring='recall')
            sent_clf.fit(sents_X, sents_y)
            d.set_sent_model(sent_clf, s.vectorizer)
            d.vectorize(test_domain)

            X_train, y_train = d.X_y_uid_filtered(domain_uids[train], test_domain)
            X_test, y_test = d.X_y_uid_filtered(domain_uids[test], test_domain)

            clf.fit(X_train, y_train)

            y_preds = clf.predict(X_test)

            fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds))[:,1]

            metrics.append(fold_metric) # get the scores for positive instances

            print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])

            metrics.append(fold_metric) # get the scores for positive instances

            # print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])

        summary_metrics = np.mean(metrics, axis=0)
        print "=" * 40
        print "mean score:\tprecision %.2f, recall %.2f, f-score %.2f" % (summary_metrics[0], summary_metrics[1], summary_metrics[2])

        # then train all for most informative features

        s = SentenceModel(test_mode=True)
        s.generate_data(uid_filter=domain_uids)
        s.vectorize()
        sents_X, sents_y = s.X_domain_all(domain=test_domain), s.y_domain_all(domain=test_domain)

        sent_tuned_parameters = [{"alpha": np.logspace(-4, -1, 5)}, {"class_weight": [{1: i, -1: 1} for i in np.logspace(0, 2, 10)]}]
        sent_clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring='recall')
        sent_clf.fit(sents_X, sents_y)
        d.set_sent_model(sent_clf, s.vectorizer)
        d.vectorize(test_domain)

        clf = SGDClassifier(loss="hinge", penalty="L2", alpha=0.1, class_weight={1: 1, -1: 1})
        X_all, y_all = d.X_y_uid_filtered(domain_uids, test_domain)

        clf.fit(X_all, y_all)

        print show_most_informative_features(s.vectorizer, clf)

def binary_hybrid_doc_prediction_test2(model=HybridDocModel(test_mode=True)):

    print "Binary hybrid doc level prediction version 2 (maybe quicker!!)"
    print "=" * 40
    print

    d = model
    d.generate_data(binarize=True) # some variations use the quote data internally
                                                # for sentence prediction (for additional features)

    for test_domain in CORE_DOMAINS:

        print ("*"*40) + "\n\n" + test_domain + "\n\n" + ("*" * 40)

        domain_uids = d.domain_uids(test_domain)
        no_studies = len(domain_uids)
        kf = KFold(no_studies, n_folds=5, shuffle=False)
        tuned_parameters = {"alpha": np.logspace(-4, -1, 10), "class_weight": [{1: i, -1: 1} for i in np.logspace(-1, 1, 10)]}
        clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring='f1')

        metrics = []

        s = SentenceModel(test_mode=True)
        s.generate_data(uid_filter=domain_uids)
        s.vectorize()

        for fold_i, (train, test) in enumerate(kf):

            sents_X, sents_y = s.X_y_uid_filtered(domain_uids[test], test_domain)
            # sent_tuned_parameters = [{"alpha": np.logspace(-4, -1, 5)}, {"class_weight": [{1: i, -1: 1} for i in np.logspace(0, 2, 10)]}]
            # sent_clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring='recall')

            sent_tuned_parameters = [{"alpha": np.logspace(-4, -1, 5)}]
            sent_clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2", class_weight={1:5, -1:1}), sent_tuned_parameters, scoring='recall')

            sent_clf.fit(sents_X, sents_y)
            d.set_sent_model(sent_clf, s.vectorizer)
            d.vectorize(test_domain)

            X_train, y_train = d.X_y_uid_filtered(domain_uids[train], test_domain)
            X_test, y_test = d.X_y_uid_filtered(domain_uids[test], test_domain)

            clf.fit(X_train, y_train)

            # print show_most_informative_features(s.vectorizer, clf.best_estimator_)
            print show_most_informative_features(s.vectorizer, clf)

        #     y_preds = clf.predict(X_test)

        #     fold_metric = np.array(sklearn.metrics.precision_recall_fscore_support(y_test, y_preds))[:,1]

        #     metrics.append(fold_metric) # get the scores for positive instances

        #     print "fold %d:\tprecision %.2f, recall %.2f, f-score %.2f" % (fold_i, fold_metric[0], fold_metric[1], fold_metric[2])

        #     metrics.append(fold_metric) # get the scores for positive instances

        # summary_metrics = np.mean(metrics, axis=0)
        # print "=" * 40
        # print "mean score:\tprecision %.2f, recall %.2f, f-score %.2f" % (summary_metrics[0], summary_metrics[1], summary_metrics[2])

        # # then train all for most informative features

        # sents_X, sents_y = s.X_domain_all(domain=test_domain), s.y_domain_all(domain=test_domain)

        # sent_tuned_parameters = [{"alpha": np.logspace(-4, -1, 5)}, {"class_weight": [{1: i, -1: 1} for i in np.logspace(0, 2, 10)]}]
        # sent_clf = GridSearchCV(SGDClassifier(loss="hinge", penalty="L2"), tuned_parameters, scoring='recall')
        # sent_clf.fit(sents_X, sents_y)

        # d.set_sent_model(sent_clf, s.vectorizer)
        # d.vectorize(test_domain)

        # clf = SGDClassifier(loss="hinge", penalty="L2", alpha=0.5, class_weight={1: 1, -1: 1})
        # X_all, y_all = d.X_y_uid_filtered(domain_uids, test_domain)
        # clf.fit(X_all, y_all)

def main():
    # dummy_document_prediction()
    stupid_sentence_prediction_test()
    # true_hybrid_prediction_test(model=HybridModelProbablistic(test_mode=False))
    # sentence_prediction_test(model=SentenceModel(test_mode=False))
    # simple_hybrid_prediction_test(model=HybridModel(test_mode=False))
    # binary_doc_prediction_test()
    #print "Try weighting sentences better"
    #binary_hybrid_doc_prediction_test2()
    # binary_hybrid_doc_prediction_test()

    # hybrid_doc_prediction_test(model=HybridDocModel2(test_mode=False))
    # document_prediction_test(model=DocumentLevelModel(test_mode=False))

    # multitask_document_prediction_test(model=MultiTaskDocumentModel(test_mode=False))
    # multitask_hybrid_document_prediction_test(model=MultiTaskHybridDocumentModel(test_mode=False))

if __name__ == '__main__':
    main()


#!/usr/bin/env python
"""
=====================
Lasso path using LARS
=====================

Computes Lasso Path along the regularization parameter using the LARS
algorithm on the diabetes dataset. Each color represents a different
feature of the coefficient vector, and this is displayed as a function
of the regularization parameter.

"""
print(__doc__)

# Author: Fabian Pedregosa <fabian.pedregosa@inria.fr>
#         Alexandre Gramfort <alexandre.gramfort@inria.fr>
# License: BSD 3 clause

import numpy as np
import matplotlib.pyplot as plt

from sklearn import linear_model
from sklearn import datasets

diabetes = datasets.load_diabetes()
X = diabetes.data
y = diabetes.target

print("Computing regularization path using the LARS ...")
alphas, _, coefs = linear_model.lars_path(X, y, method='lasso', verbose=True)

xx = np.sum(np.abs(coefs.T), axis=1)
xx /= xx[-1]

plt.plot(xx, coefs.T)
ymin, ymax = plt.ylim()
plt.vlines(xx, ymin, ymax, linestyle='dashed')
plt.xlabel('|coef| / max|coef|')
plt.ylabel('Coefficients')
plt.title('LASSO Path')
plt.axis('tight')
plt.show()


"""Bayesian Data Analysis, 3rd ed
Chapter 5, demo 2

Hierarchical model for SAT-example data (BDA3, p. 102)

"""

from __future__ import division
import numpy as np
from scipy.stats import norm
import scipy.io # For importing a matlab file
import matplotlib.pyplot as plt

# Edit default plot settings (colours from colorbrewer2.org)
plt.rc('font', size=14)
plt.rc('lines', color='#377eb8', linewidth=2)
plt.rc('axes', color_cycle=(plt.rcParams['lines.color'],)) # Disable color cycle

# SAT-example data (BDA3 p. 120)
# y is the estimated treatment effect
# s is the standard error of effect estimate
y = np.array([28,  8, -3,  7, -1,  1, 18, 12])
s = np.array([15, 10, 16, 11,  9, 11, 10, 18])
M = len(y)

# load the pre-computed results for the hierarchical model
# replace this with your own code in Ex 5.1*
hres_path = '../utilities_and_data/demo5_2.mat'
hres = scipy.io.loadmat(hres_path)
''' Content information of the precalculated results:
>>> scipy.io.whosmat('demo5_2.mat')
[('pxm', (8, 500),  'double'),
 ('t',   (1, 1000), 'double'),
 ('tp',  (1, 1000), 'double'),
 ('tsd', (8, 1000), 'double'),
 ('tm',  (8, 1000), 'double')]
'''
pxm = hres['pxm']
t   = hres['t'][0]
tp  = hres['tp'][0]
tsd = hres['tsd']
tm  = hres['tm']

# plot the separate, pooled and hierarchical models
fig, axes = plt.subplots(3, 1, sharex=True, figsize=(8,10))
x = np.linspace(-40, 60, 500)

# separate
lines = axes[0].plot(x, norm.pdf(x[:,None], y[1:], s[1:]), linewidth=1)
line, = axes[0].plot(x, norm.pdf(x, y[0], s[0]), 'r')
axes[0].legend((line, lines[1]), ('school A', 'other schools'),
               loc='upper left')
axes[0].set_yticks(())
axes[0].set_title('separate model')

# pooled
axes[1].plot(
    x,
    norm.pdf(
        x,
        np.sum(y/s**2)/np.sum(1/s**2),
        np.sqrt(1/np.sum(1/s**2))
    ),
    label='All schools'
)
axes[1].legend(loc='upper left')
axes[1].set_yticks(())
axes[1].set_title('pooled model')

# hierarchical
lines = axes[2].plot(x, pxm[1:].T, linewidth=1)
line, = axes[2].plot(x, pxm[0], 'r')
axes[2].legend((line, lines[1]), ('school A', 'other schools'),
               loc='upper left')
axes[2].set_yticks(())
axes[2].set_title('hierarchical model')
axes[2].set_xlabel('Treatment effect')

# plot various marginal and conditional posterior summaries
fig, axes = plt.subplots(3, 1, sharex=True, figsize=(8,10))

axes[0].plot(t, tp)
axes[0].set_yticks(())
axes[0].set_title(r'marginal posterior density $p(\tau|y)$')
axes[0].set_ylabel(r'$p(\tau|y)$', fontsize=20)
axes[0].set_xlim([0,35])

lines = axes[1].plot(t, tm[1:].T, linewidth=1)
line, = axes[1].plot(t, tm[0].T, 'r')
axes[1].legend((line, lines[1]), ('school A', 'other schools'),
               loc='upper left')
axes[1].set_title(r'conditional posterior means of effects '
                  r'$\operatorname{E}(\theta_j|\tau,y)$')
axes[1].set_ylabel(r'$\operatorname{E}(\theta_j|\tau,y)$', fontsize=20)

lines = axes[2].plot(t, tsd[1:].T, linewidth=1)
line, = axes[2].plot(t, tsd[0].T, 'r')
axes[2].legend((line, lines[1]), ('school A', 'other schools'),
               loc='upper left')
axes[2].set_title(r'standard deviations of effects '
                  r'$\operatorname{sd}(\theta_j|\tau,y)$')
axes[2].set_ylabel(r'$\operatorname{sd}(\theta_j|\tau,y)$', fontsize=20)
axes[2].set_xlabel(r'$\tau$', fontsize=20)

plt.show()


#
# Copyright 2015 Quantopian, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Tools to generate data sources.
"""
import numpy as np
import pandas as pd

from zipline.gens.utils import hash_args

from zipline.sources.data_source import DataSource

class DataFrameSource(DataSource):
    """
    Data source that yields from a pandas DataFrame.

    :Axis layout:
        * columns : sids
        * index : datetime

    :Note:
        Bars where the price is nan are filtered out.
    """

    def __init__(self, data, **kwargs):
        assert isinstance(data.index, pd.tseries.index.DatetimeIndex)
        # Only accept integer SIDs as the items of the DataFrame
        assert isinstance(data.columns, pd.Int64Index)
        # TODO is ffilling correct/necessary?
        # Forward fill prices
        self.data = data.fillna(method='ffill')
        # Unpack config dictionary with default values.
        self.start = kwargs.get('start', self.data.index[0])
        self.end = kwargs.get('end', self.data.index[-1])
        self.sids = self.data.columns

        # Hash_value for downstream sorting.
        self.arg_string = hash_args(data, **kwargs)

        self._raw_data = None

        self.started_sids = set()

    @property
    def mapping(self):
        return {
            'dt': (lambda x: x, 'dt'),
            'sid': (lambda x: x, 'sid'),
            'price': (float, 'price'),
            'volume': (int, 'volume'),
        }

    @property
    def instance_hash(self):
        return self.arg_string

    def raw_data_gen(self):
        for dt, series in self.data.iterrows():
            for sid, price in series.iteritems():
                # Skip SIDs that can not be forward filled
                if np.isnan(price) and \
                   sid not in self.started_sids:
                    continue
                self.started_sids.add(sid)

                event = {
                    'dt': dt,
                    'sid': sid,
                    'price': price,
                    # Just chose something large
                    # if no volume available.
                    'volume': 1e9,
                }
                yield event

    @property
    def raw_data(self):
        if not self._raw_data:
            self._raw_data = self.raw_data_gen()
        return self._raw_data

class DataPanelSource(DataSource):
    """
    Data source that yields from a pandas Panel.

    :Axis layout:
        * items : sids
        * major_axis : datetime
        * minor_axis : price, volume, ...

    :Note:
        Bars where the price is nan are filtered out.
    """

    def __init__(self, data, **kwargs):
        assert isinstance(data.major_axis, pd.tseries.index.DatetimeIndex)
        # Only accept integer SIDs as the items of the Panel
        assert isinstance(data.items, pd.Int64Index)
        # TODO is ffilling correct/necessary?
        # forward fill with volumes of 0
        self.data = data.fillna(value={'volume': 0})
        self.data = self.data.fillna(method='ffill')
        # Unpack config dictionary with default values.
        self.start = kwargs.get('start', self.data.major_axis[0])
        self.end = kwargs.get('end', self.data.major_axis[-1])
        self.sids = self.data.items

        # Hash_value for downstream sorting.
        self.arg_string = hash_args(data, **kwargs)

        self._raw_data = None

        self.started_sids = set()

    @property
    def mapping(self):
        mapping = {
            'dt': (lambda x: x, 'dt'),
            'sid': (lambda x: x, 'sid'),
            'price': (float, 'price'),
            'volume': (int, 'volume'),
        }

        # Add additional fields.
        for field_name in self.data.minor_axis:
            if field_name in ['price', 'volume', 'dt', 'sid']:
                continue
            mapping[field_name] = (lambda x: x, field_name)

        return mapping

    @property
    def instance_hash(self):
        return self.arg_string

    def raw_data_gen(self):
        for dt in self.data.major_axis:
            df = self.data.major_xs(dt)
            for sid, series in df.iteritems():
                # Skip SIDs that can not be forward filled
                if np.isnan(series['price']) and \
                   sid not in self.started_sids:
                    continue
                self.started_sids.add(sid)

                event = {
                    'dt': dt,
                    'sid': sid,
                }
                for field_name, value in series.iteritems():
                    event[field_name] = value

                yield event

    @property
    def raw_data(self):
        if not self._raw_data:
            self._raw_data = self.raw_data_gen()
        return self._raw_data


import os
import re
import pandas as pd
import string
import itertools
import numpy as np
import sys
import argparse
from collections import OrderedDict

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Create GO annotation and enrichment file')
    parser.add_argument('-i',type=str,dest='infile',required=True,help="Input file")
    parser.add_argument('-o',type=str,dest='out',required=True,help="Ouput file")
    parser.add_argument('-db',type=str,dest='db',required=True,help="GO Database file")
    args = parser.parse_args()
    print (args)

def sort_uniq(sequence):
    return (x[0] for x in itertools.groupby(sorted(sequence)))

path = "/home/zluna/Work/GO"
fout = open(args.out+"_anno.xls", 'w')
print("Gene_id", "GO_annotation", sep = '\t', file = fout)
go_db = pd.read_table(os.path.join(path, args.db), header = None)
eggout = pd.read_table(os.path.join(path, args.infile), header = None)
#pd.DataFrame.head(eggout)
#eggout.head(100)
dict = OrderedDict()
first_flag = 1
a = list(go_db[0])
for i in range(len(eggout)):
    gene_id = eggout[0][i]
    go_id = eggout[5][i]
    if pd.isnull(eggout[5][i]):
        go_id = ''
    #print(gene_id, kegg_id, type(kegg_id), sep ='\t')
    go_id = go_id.split(',')
    if len(go_id) == 0:
        continue
    go_term = '; '.join(list(go_db[go_db[2].isin(go_id)][0]))
    #print(gene_id, go_id, go_term, sep ='\t')
    go_sum = []
    sel_go_table = go_db[go_db[2].isin(go_id)]
    for j in range(len(sel_go_table)):
        go_sum.append(''.join(( list(sel_go_table[2])[j], "~", list(sel_go_table[0])[j])))
    print(gene_id, str(go_sum).strip('[]').replace(']','').replace("'","").replace(", ","; "), sep = '\t', file = fout)

    a = list(go_db[2])
### Use dictionary
    for k in range(len(a)):
        if str(go_sum).find(a[k]) != -1 :
            if a[k] not in dict.keys():
### The value must be list type, if just give the 'gene_id' as the value of key, it can not use 'append' method to add the new 'gene_id' to the existing key.
                dict[a[k]] = []
                dict[a[k]].append(gene_id)
            else:
                dict[a[k]].append(gene_id)

                #dict[a[j]] = [dict[a[j]], gene_id]
fout.close()

fout2 = open(args.out+"_enrich.xls", 'w')
print('GOID', 'Term', 'Genes', 'Gene_count', sep = '\t', file = fout2)
for key,values in  dict.items():
    print(key, list(go_db[go_db[2] == key][0]), str(values).strip('[]').replace(']','').replace("'",""), len(values), sep ='\t', file = fout2)
fout2.cloes()


# The MIT License (MIT)
# Copyright (c) 2016, 2017 by the ESA CCI Toolbox development team and contributors
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is furnished to do
# so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Description
===========

Index calculation operations

Functions
=========
"""
import xarray as xr
import pandas as pd

from cate.core.op import op, op_input
from cate.ops.select import select_var
from cate.ops.subset import subset_spatial
from cate.ops.anomaly import anomaly_external
from cate.core.types import PolygonLike, VarName, ValidationError
from cate.util.monitor import Monitor

_ALL_FILE_FILTER = dict(name='All Files', extensions=['*'])

@op(tags=['index'])
@op_input('file', file_open_mode='r', file_filters=[dict(name='NetCDF', extensions=['nc']), _ALL_FILE_FILTER])
@op_input('var', value_set_source='ds', data_type=VarName)
def enso_nino34(ds: xr.Dataset,
                var: VarName.TYPE,
                file: str,
                threshold: float = None,
                monitor: Monitor = Monitor.NONE) -> pd.DataFrame:
    """
    Calculate nino34 index, which is defined as a five month running mean of
    anomalies of monthly means of SST data in Nino3.4 region:: lon_min=-170
    lat_min=-5 lon_max=-120 lat_max=5.

    :param ds: A monthly SST dataset
    :param file: Path to the reference data file e.g. a climatology. A suitable reference dataset
    can be generated using the long_term_average operation
    :param var: Dataset variable (geophysial quantity) to use for index
    calculation.
    :param threshold: If given, boolean El Nino/La Nina timeseries will be
    calculated and added to the output dataset according to the given
    threshold. Where anomaly larger than the positive value of the threshold
    indicates El Nino and anomaly smaller than the negative of the given
    threshold indicates La Nina.
    :param monitor: a progress monitor.
    :return: A dataset that contains the index timeseries.
    """
    n34 = '-170, -5, -120, 5'
    name = 'ENSO N3.4 Index'
    return _generic_index_calculation(ds, var, n34, 5, file, name, threshold, monitor)

@op(tags=['index'])
@op_input('var', value_set_source='ds', data_type=VarName)
@op_input('file', file_open_mode='r', file_filters=[dict(name='NetCDF', extensions=['nc']), _ALL_FILE_FILTER])
@op_input('region', value_set=['N1+2', 'N3', 'N34', 'N4', 'custom'])
@op_input('custom_region', data_type=PolygonLike)
def enso(ds: xr.Dataset,
         var: VarName.TYPE,
         file: str,
         region: str = 'n34',
         custom_region: PolygonLike.TYPE = None,
         threshold: float = None,
         monitor: Monitor = Monitor.NONE) -> pd.DataFrame:
    """
    Calculate ENSO index, which is defined as a five month running mean of
    anomalies of monthly means of SST data in the given region.

    :param ds: A monthly SST dataset
    :param file: Path to the reference data file e.g. a climatology. A suitable reference dataset
    can be generated using the long_term_average operation
    :param var: Dataset variable to use for index calculation
    :param region: Region for index calculation, the default is Nino3.4
    :param custom_region: If 'custom' is chosen as the 'region', this parameter
    has to be provided to set the desired region.
    :param threshold: If given, boolean El Nino/La Nina timeseries will be
    calculated and added to the output dataset, according to the given
    threshold. Where anomaly larger than then positive value of the threshold
    indicates El Nino and anomaly smaller than the negative of the given
    threshold indicates La Nina.
    :param monitor: a progress monitor.
    :return: A dataset that contains the index timeseries.
    """
    regions = {'N1+2': '-90, -10, -80, 0',
               'N3': '-150, -5, -90, 5',
               'N3.4': '-170, -5, -120, 5',
               'N4': '160, -5, -150, 5',
               'custom': custom_region}
    converted_region = PolygonLike.convert(regions[region])
    if not converted_region:
        raise ValidationError('No region has been provided to ENSO index calculation')

    name = 'ENSO ' + region + ' Index'
    if 'custom' == region:
        name = 'ENSO Index over ' + PolygonLike.format(converted_region)

    return _generic_index_calculation(ds, var, converted_region, 5, file, name, threshold, monitor)

@op(tags=['index'])
@op_input('var', value_set_source='ds', data_type=VarName)
@op_input('file', file_open_mode='r', file_filters=[dict(name='NetCDF', extensions=['nc']), _ALL_FILE_FILTER])
def oni(ds: xr.Dataset,
        var: VarName.TYPE,
        file: str,
        threshold: float = None,
        monitor: Monitor = Monitor.NONE) -> pd.DataFrame:
    """
    Calculate ONI index, which is defined as a three month running mean of
    anomalies of monthly means of SST data in the Nino3.4 region.

    :param ds: A monthly SST dataset
    :param file: Path to the reference data file e.g. a climatology. A suitable reference dataset
    can be generated using the long_term_average operation
    :param var: Dataset variable to use for index calculation
    :param threshold: If given, boolean El Nino/La Nina timeseries will be
    calculated and added to the output dataset, according to the given
    threshold. Where anomaly larger than then positive value of the threshold
    indicates El Nino and anomaly smaller than the negative of the given
    threshold indicates La Nina.
    :param monitor: a progress monitor.
    :return: A dataset that containts the index timeseries
    """
    n34 = '-170, -5, -120, 5'
    name = 'ONI Index'
    return _generic_index_calculation(ds, var, n34, 3, file, name, threshold, monitor)

def _generic_index_calculation(ds: xr.Dataset,
                               var: VarName.TYPE,
                               region: PolygonLike.TYPE,
                               window: int,
                               file: str,
                               name: str,
                               threshold: float = None,
                               monitor: Monitor = Monitor.NONE) -> pd.DataFrame:
    """
    A generic index calculation. Where an index is defined as an anomaly
    against the given reference of a moving average of the given window size of
    the given given region of the given variable of the given dataset.

    :param ds: Dataset from which to calculate the index
    :param var: Variable from which to calculate index
    :param region: Spatial subset from which to calculate the index
    :param window: Window size for the moving average
    :param file: Path to the reference file
    :param threshold: Absolute threshold that indicates an ENSO event
    :param name: Name of the index
    :param monitor: a progress monitor.
    :return: A dataset that contains the index timeseries
    """
    var = VarName.convert(var)
    region = PolygonLike.convert(region)

    with monitor.starting("Calculate the index", total_work=2):
        ds = select_var(ds, var)
        ds_subset = subset_spatial(ds, region)
        anom = anomaly_external(ds_subset, file, monitor=monitor.child(1))
        with monitor.child(1).observing("Calculate mean"):
            ts = anom.mean(dim=['lat', 'lon'])
        df = pd.DataFrame(data=ts[var].values, columns=[name], index=ts.time.values)
        retval = df.rolling(window=window, center=True).mean().dropna()

    if threshold is None:
        return retval

    retval['El Nino'] = pd.Series((retval[name] > threshold),
                                  index=retval.index)
    retval['La Nina'] = pd.Series((retval[name] < -threshold),
                                  index=retval.index)
    return retval


import unittest
import sys

import numpy as np

from sklearn.mixture import DPGMM, VBGMM
from sklearn.mixture.dpgmm import log_normalize
from sklearn.datasets import make_blobs
from sklearn.utils.testing import assert_array_less, assert_equal
from sklearn.mixture.tests.test_gmm import GMMTester
from sklearn.externals.six.moves import cStringIO as StringIO

np.seterr(all='warn')

def test_class_weights():
    # check that the class weights are updated
    # simple 3 cluster dataset
    X, y = make_blobs(random_state=1)
    for Model in [DPGMM, VBGMM]:
        dpgmm = Model(n_components=10, random_state=1, alpha=20, n_iter=50)
        dpgmm.fit(X)
        # get indices of components that are used:
        indices = np.unique(dpgmm.predict(X))
        active = np.zeros(10, dtype=np.bool)
        active[indices] = True
        # used components are important
        assert_array_less(.1, dpgmm.weights_[active])
        # others are not
        assert_array_less(dpgmm.weights_[~active], .05)

def test_verbose_boolean():
    # checks that the output for the verbose output is the same
    # for the flag values '1' and 'True'
    # simple 3 cluster dataset
    X, y = make_blobs(random_state=1)
    for Model in [DPGMM, VBGMM]:
        dpgmm_bool = Model(n_components=10, random_state=1, alpha=20,
                           n_iter=50, verbose=True)
        dpgmm_int = Model(n_components=10, random_state=1, alpha=20,
                          n_iter=50, verbose=1)

        old_stdout = sys.stdout
        sys.stdout = StringIO()
        try:
            # generate output with the boolean flag
            dpgmm_bool.fit(X)
            verbose_output = sys.stdout
            verbose_output.seek(0)
            bool_output = verbose_output.readline()
            # generate output with the int flag
            dpgmm_int.fit(X)
            verbose_output = sys.stdout
            verbose_output.seek(0)
            int_output = verbose_output.readline()
            assert_equal(bool_output, int_output)
        finally:
            sys.stdout = old_stdout

def test_verbose_first_level():
    # simple 3 cluster dataset
    X, y = make_blobs(random_state=1)
    for Model in [DPGMM, VBGMM]:
        dpgmm = Model(n_components=10, random_state=1, alpha=20, n_iter=50,
                      verbose=1)

        old_stdout = sys.stdout
        sys.stdout = StringIO()
        try:
            dpgmm.fit(X)
        finally:
            sys.stdout = old_stdout

def test_verbose_second_level():
    # simple 3 cluster dataset
    X, y = make_blobs(random_state=1)
    for Model in [DPGMM, VBGMM]:
        dpgmm = Model(n_components=10, random_state=1, alpha=20, n_iter=50,
                      verbose=2)

        old_stdout = sys.stdout
        sys.stdout = StringIO()
        try:
            dpgmm.fit(X)
        finally:
            sys.stdout = old_stdout

def test_log_normalize():
    v = np.array([0.1, 0.8, 0.01, 0.09])
    a = np.log(2 * v)
    assert np.allclose(v, log_normalize(a), rtol=0.01)

def do_model(self, **kwds):
    return VBGMM(verbose=False, **kwds)

class DPGMMTester(GMMTester):
    model = DPGMM
    do_test_eval = False

    def score(self, g, train_obs):
        _, z = g.score_samples(train_obs)
        return g.lower_bound(train_obs, z)

class TestDPGMMWithSphericalCovars(unittest.TestCase, DPGMMTester):
    covariance_type = 'spherical'
    setUp = GMMTester._setUp

class TestDPGMMWithDiagCovars(unittest.TestCase, DPGMMTester):
    covariance_type = 'diag'
    setUp = GMMTester._setUp

class TestDPGMMWithTiedCovars(unittest.TestCase, DPGMMTester):
    covariance_type = 'tied'
    setUp = GMMTester._setUp

class TestDPGMMWithFullCovars(unittest.TestCase, DPGMMTester):
    covariance_type = 'full'
    setUp = GMMTester._setUp

class VBGMMTester(GMMTester):
    model = do_model
    do_test_eval = False

    def score(self, g, train_obs):
        _, z = g.score_samples(train_obs)
        return g.lower_bound(train_obs, z)

class TestVBGMMWithSphericalCovars(unittest.TestCase, VBGMMTester):
    covariance_type = 'spherical'
    setUp = GMMTester._setUp

class TestVBGMMWithDiagCovars(unittest.TestCase, VBGMMTester):
    covariance_type = 'diag'
    setUp = GMMTester._setUp

class TestVBGMMWithTiedCovars(unittest.TestCase, VBGMMTester):
    covariance_type = 'tied'
    setUp = GMMTester._setUp

class TestVBGMMWithFullCovars(unittest.TestCase, VBGMMTester):
    covariance_type = 'full'
    setUp = GMMTester._setUp