#!/usr/bin/env python
# coding=utf-8
import argparse

from cla.util import util


def main():
    parser = argparse.ArgumentParser(description='chinese language tools command line.')
    parser.add_argument("action", help="action to perform, including "
                                       "extract, cut, learn, doc2vec, "
                                       "similar, classify, topic")
    parser.add_argument("input", nargs="?", help="input for the action")
    parser.add_argument("-o", "--output-file", help="file to save output")
    parser.add_argument("-p", "--port", default=8888, type=int, help="server port to listen")
    parser.add_argument("-m", "--stored-model", help="saved word2vec model")
    parser.add_argument("-t", "--training-data", help="labeled data to train classifier")
    parser.add_argument("-s", "--skip-lines", nargs="+", help="skip lines with the specified prefixes when cutting")
    parser.add_argument("-e", "--encoding", default="utf-8", help="specify file encoding")
    parser.add_argument("-D", "--daemon", action="store_true", help="run as a background server, supports classify")
    parser.add_argument("-C", "--cleanup", action="store_true", help="cleaup meaningless words")
    parser.add_argument("--source-type", default="qq", help="source type of the input, supported types are qq")
    parser.add_argument("--min-length", default=20, type=int, help="min sentence length for topic learning")
    args = parser.parse_args()

    if args.action == "extract":
        if not args.input:
            print "Source file is required!"
            return

        print "Prepare extracting sentences from " + args.input
        if args.source_type == "qq":
            result_path = util.process_qq_history(args.input, encoding=args.encoding, output_path=args.output_file)
            print "Done. Extracted sentences saved to " + result_path
        else:
            print "Unsupported type: " + args.source_type

    elif args.action == "cut":
        if not args.input:
            print "File contains sentences is required!"
            return

        print "Cutting sentences in " + args.input
        result_path = util.cut_words_in(args.input,
                                        encoding=args.encoding,
                                        output_path=args.output_file,
                                        skip_prefixes=args.skip_lines,
                                        cleanup=args.cleanup)
        print "Done. Words saved ot " + result_path

    elif args.action == "learn":
        from cla.learn.word2vec import VectorModel
        if not args.input:
            print "File contains words is required!"
            return

        if args.stored_model:
            model = VectorModel(source_file_path=args.stored_model)
            result_file = args.stored_model
            update = True
            print "Loaded previous model: " + args.stored_model
        else:
            result_file = "word2vec.model"
            model = VectorModel()
            update = False

        print "Learning vector presentation of " + args.input
        model.train(source_corpus_path=args.input, update=update)

        if args.output_file:
            result_file = args.output_file
        model.save(result_file)
        print "Done. Model saved to " + result_file

    elif args.action == "doc2vec":
        from cla.learn.word2vec import VectorModel
        if not args.input:
            print "File contains words is required!"
            return
        if not args.stored_model:
            print "A stored word2vec model is required! (--stored-model)"
            return

        model = VectorModel(source_file_path=args.stored_model)
        if not args.output_file:
            import os
            _, filename = os.path.split(args.input)
            result_path = "vector_of_" + filename
        else:
            result_path = args.output_file

        with open(args.input, 'r') as input_file, open(result_path, 'w') as output_file:
            for line in input_file:
                line = line.decode(encoding=args.encoding).strip()
                if not line:
                    continue

                vector = model.to_vector(line.split(" "))
                output_file.write(" ".join(map(str, vector)) + "\n")
        print "Done. Vectors saved to " + result_path

    elif args.action == "similar":
        from cla.learn.word2vec import VectorModel
        if not args.input:
            print "A input word is required!"
            return
        if not args.stored_model:
            print "A stored word2vec model is required! (--stored-model)"
            return

        model = VectorModel(source_file_path=args.stored_model)
        for word, similarity in model.model.similar_by_word(args.input.decode(args.encoding)):
            print word.encode(args.encoding) + "\t%.5f" % similarity

    elif args.action == "classify":
        from cla.learn.classifier import TraditionalClassifier

        if not args.stored_model:
            print "A stored word2vec model is required! (--stored-model)"
            return
        if not args.training_data:
            print "Training labeled data is required! (--training-data)"
            return

        if not args.input and not args.daemon:
            print "File contains words is required!"
            return

        classifier = TraditionalClassifier(vector_model_path=args.stored_model, training_data_path=args.training_data)
        if args.daemon:
            def func(uri_data):
                data = {
                    "label": 0,
                    "words": []
                }
                words = util.cut_words(uri_data.decode(encoding=args.encoding))
                data["label"] = int(classifier.classify([words])[0])
                for a_word in words:
                    data["words"].append(a_word.encode(encoding=args.encoding))
                return data

            server(args.port, func)
        else:
            if not args.output_file:
                import os
                _, filename = os.path.split(args.input)
                result_path = "classified_" + filename
            else:
                result_path = args.output_file

            with open(args.input, 'r') as input_file, open(result_path, 'w') as output_file:
                for line in input_file:
                    line = line.decode(encoding=args.encoding)
                    classified = classifier.classify([line])[0]
                    output_file.write(classified + " " + line.encode(encoding=args.encoding) + "\n")

    elif args.action == "topic":
        if not args.input:
            print "File contains sentences is required!"
            return
        if not args.output_file:
            import os
            _, filename = os.path.split(args.input)
            result_path = "topic_of_" + filename
        else:
            result_path = args.output_file

        # =====
        # KMeans method.
        # -----
        #
        # from cla.learn.topic import TopicClustering
        #
        # if not args.stored_model:
        #     print "A stored word2vec model is required! (--stored-model)"
        #     return
        #
        # model = TopicClustering(document_path=args.input, vector_model_path=args.stored_model,
        #                         cut=False, n_clusters=20)
        # r = {}
        # with open(args.input, 'r') as input_file:
        #     for line in input_file:
        #         candidates = model.cluster(util.cut_words(line.decode(encoding=args.encoding)))
        #         if len(candidates) < 1:
        #             continue
        #         document_topic = candidates[0]
        #         if document_topic not in r:
        #             r[document_topic] = []
        #         r[document_topic].append(line)
        #
        # with open(result_path, 'w') as output:
        #     for topic_id in r:
        #         topic = "Cluster " + str(topic_id)
        #         output.write("=====\n" + topic + "\n-----\n")
        #         for sentence in r[topic_id]:
        #             output.write(sentence)
        #         output.write("\n")
        # =====
        # LDA method.
        # -----
        #
        from cla.learn.topic import TopicModel
        model = TopicModel(args.input, cut=False, num_topics=20, min_length=args.min_length)
        r = {}
        with open(args.input, 'r') as input_file:
            for line in input_file:
                decoded_line = line.decode(encoding=args.encoding)
                if len(decoded_line) < args.min_length:
                    continue

                candidates = model.identify_topic(util.cut_words(decoded_line))
                if len(candidates) < 1:
                    continue

                document_topic, _ = candidates[0]
                if document_topic not in r:
                    r[document_topic] = []
                r[document_topic].append(line)

        with open(result_path, 'w') as output:
            for topic_id in r:
                topic_words = model.topic_words(topic_id, 10)
                topic = str(topic_id) + ": " + " ".join([x for x, _ in topic_words]).encode(encoding=args.encoding)
                output.write("=====\n" + topic + "\n-----\n")
                for sentence in r[topic_id]:
                    output.write(sentence)
                output.write("\n")

    else:
        print "Unsupported action: " + args.action


def server(port, func):
    import SocketServer
    import json
    from BaseHTTPServer import BaseHTTPRequestHandler

    # noinspection PyPep8Naming,PyClassHasNoInit
    class Handler(BaseHTTPRequestHandler):
        def do_OPTIONS(self):
            self.send_response(200)
            self.send_header("Access-Control-Allow-Origin", "*")
            self.send_header("Access-Control-Allow-Headers", "origin, content-type, accept")
            self.send_header("Access-Control-Allow-Credentials", "true")
            self.send_header("Access-Control-Allow-Methods", "GET, OPTIONS, HEAD")
            self.send_header("Access-Control-Max-Age", "1209600")
            self.end_headers()

        def do_GET(self):
            import urllib
            uri_data = urllib.unquote(self.path[1:])
            self.send_response(200)
            self.send_header("Content-type", "application/json")
            self.send_header("Access-Control-Allow-Origin", "*")
            self.end_headers()

            r = {
                "status": "error",
                "data": "unknown error."
            }

            try:
                r["data"] = func(uri_data)
                r["status"] = "ok"
            except Exception as e:
                r["data"] = str(e)

            self.wfile.write(json.dumps(r, ensure_ascii=False))

    httpd = SocketServer.TCPServer(("", port), Handler)

    print "Serving at port: ", port
    httpd.serve_forever()


if __name__ == '__main__':
    main()
