import readers.pos_corpus as pcc
import sequences.extended_feature as exfc
import sequences.structured_perceptron as spc

MAX_SENT_SIZE = 1000
MAX_NR_SENTENCES = 100000
MODEL_DIR = "../models/wsj_postag/"


def build_corpus_features():
    corpus = pcc.PostagCorpus()
    train_seq = corpus.read_sequence_list_conll(
        "../data/train-02-21.conll",
        max_sent_len=MAX_SENT_SIZE,
        max_nr_sent=MAX_NR_SENTENCES)
    corpus.add_sequence_list(train_seq)
    features = exfc.ExtendedFeatures(corpus)
    features.build_features()
    corpus.save_corpus(MODEL_DIR)
    features.save_features(MODEL_DIR + "features.txt")
    return corpus, features


def train_pos(corpus, features):
    model = spc.StructuredPercetron(corpus, features)
    model.nr_rounds = 10
    model.train_supervised(corpus.sequence_list.seq_list)
    model.save_model(MODEL_DIR)
    return model


def eval_model(corpus, features, model):
    dev_seq = corpus.read_sequence_list_conll("../data/dev-22.conll")
    pred_dev = model.viterbi_decode_corpus_log(dev_seq.seq_list)
    eval_dev = model.evaluate_corpus(dev_seq.seq_list, pred_dev)
    print "Accuracy on wsj development %f" % eval_dev
    test_seq = corpus.read_sequence_list_conll("../data/test-23.conll")
    pred_test = model.viterbi_decode_corpus_log(test_seq.seq_list)
    eval_test = model.evaluate_corpus(test_seq.seq_list, pred_test)
    print "Accuracy on wsj test %f" % eval_test


def eval_brown(corpus, features, model):
    categories = [
        'adventure',
        'belles_lettres',
        'editorial',
        'fiction',
        'government',
        'hobbies',
        'humor',
        'learned',
        'lore',
        'mystery',
        'news',
        'religion',
        'reviews',
        'romance',
        'science_fiction']
    for cat in categories:
        brown_seq = corpus.read_sequence_list_brown(categories=cat)
        brown_pred = model.viterbi_decode_corpus_log(brown_seq.seq_list)
        brown_eval = model.evaluate_corpus(brown_seq.seq_list, brown_pred)
        print "Accuracy on Brown cat %s: %f" % (cat, brown_eval)


def load_model():
    corpus = pcc.PostagCorpus()
    corpus.load_corpus(MODEL_DIR)
    features = exfc.ExtendedFeatures(corpus)
    features.load_features(MODEL_DIR + "features.txt", corpus)
    model = spc.StructuredPercetron(corpus, features)
    model.load_model(MODEL_DIR)
    return corpus, features, model


def main():
    print "Building corpus"
    corpus, features = build_corpus_features()
    print "Training model"
    model = train_pos(corpus, features)
    print "Testing on wsj"
    eval_model(corpus, features, model)
    print "Testing on brown"
    eval_brown(corpus, features, model)
    # print "Loading models"
    # corpus,features,model = load_model()
    # print "Testing on wsj"
    # eval_model(corpus,features,model)
    # print "Testing on brown"
    # eval_brown(corpus,features,model)
