import os
import sys
import argparse
import time
import traceback
import subprocess
import re

import utils
import infer
import config
from utils import logger


def load_existing_results(eval_result_file):
    evals = {}
    with utils.open_file(eval_result_file) as f:
        for line in f:
            line = line.strip()
            if not line: continue
            pos = line.find(" ")
            pass_id, ret = int(line[len("Pass="):pos]), line[pos + 1:]
            evals[pass_id] = ret
    return evals


__PATTERN_CHUNK_F1 = re.compile("chunk_f1=(\d+(\.\d+)?)")


def find_best_pass(evals):
    results = []
    for pass_id, eval_ret in evals.iteritems():
        chunk_f1 = float(__PATTERN_CHUNK_F1.search(eval_ret).group(1))
        results.append((pass_id, chunk_f1))

    results.sort(key=lambda item: (-item[1], item[0]))
    return results[0][0]


def eval_one_pass(infer_obj, conf, model_path, data_path, eval_script):
    if not os.path.exists("tmp"): os.makedirs("tmp")
    # model file is not ready
    if not os.path.exists(model_path): return False

    output_path = os.path.join("tmp", "%s_%s.txt.gz" % (
        os.path.basename(model_path), os.path.basename(data_path)))
    with utils.open_file(output_path, "w") as output:
        try:
            infer_obj.infer(model_path, data_path, output)
        except Exception as ex:
            traceback.print_exc()
            return None

    cmd = [
        "python", eval_script, output_path, data_path, "--fuzzy", "--schema",
        conf.label_schema
    ]
    logger.info("cmd: %s" % " ".join(cmd))
    eval_ret = subprocess.check_output(cmd)
    if "chunk_f1" not in eval_ret:
        raise ValueError("Unknown error in cmd \"%s\"" % " ".join(cmd))

    return eval_ret


def run_eval(infer_obj,
             conf,
             model_dir,
             input_path,
             eval_script,
             log_file,
             start_pass_id,
             end_pass_id,
             force_rerun=False):
    if not force_rerun and os.path.exists(log_file):
        evals = load_existing_results(log_file)
    else:
        evals = {}
    with utils.open_file(log_file, "w") as log:
        for i in xrange(start_pass_id, end_pass_id + 1):
            if i in evals:
                eval_ret = evals[i]
            else:
                pass_id = "%05d" % i
                model_path = os.path.join(model_dir,
                                          "params_pass_%s.tar.gz" % pass_id)
                logger.info("Waiting for model %s ..." % model_path)
                while True:
                    eval_ret = eval_one_pass(infer_obj, conf, model_path,
                                             input_path, eval_script)
                    if eval_ret:
                        evals[i] = eval_ret
                        break

                    # wait for one minute and rerun
                    time.sleep(60)
            print >> log, "Pass=%d %s" % (i, eval_ret.rstrip())
            log.flush()
    return evals


def parse_cmd():
    parser = argparse.ArgumentParser()
    parser.add_argument("model_dir")
    parser.add_argument("data_type", choices=["ann", "ir"], default="ann")
    parser.add_argument(
        "--val_eval_output", help="validation set evaluation result file")
    parser.add_argument(
        "--tst_eval_output", help="test set evaluation result file")
    parser.add_argument("--start_pass_id", type=int, default=0)
    parser.add_argument(
        "--end_pass_id", type=int, default=24, help="this pass is included")
    parser.add_argument("--force_rerun", action="store_true")
    return parser.parse_args()


__eval_scripts = {
    "ann": "data/evaluation/evaluate-tagging-result.py",
    "ir": "data/evaluation/evaluate-voting-result.py",
}

__val_data = {
    "ann": "./data/data/validation.ann.json.gz",
    "ir": "./data/data/validation.ir.json.gz",
}

__tst_data = {
    "ann": "./data/data/test.ann.json.gz",
    "ir": "./data/data/test.ir.json.gz",
}


def main(args):
    conf = config.InferConfig()
    conf.vocab = utils.load_dict(conf.word_dict_path)
    logger.info("length of word dictionary is : %d." % len(conf.vocab))

    if args.val_eval_output:
        val_eval_output = args.val_eval_output
    else:
        val_eval_output = "eval.val.%s.txt" % args.data_type

    if args.tst_eval_output:
        tst_eval_output = args.tst_eval_output
    else:
        tst_eval_output = "eval.tst.%s.txt" % args.data_type

    eval_script = __eval_scripts[args.data_type]
    val_data_file = __val_data[args.data_type]
    tst_data_file = __tst_data[args.data_type]

    infer_obj = infer.Infer(conf)
    val_evals = run_eval(
        infer_obj,
        conf,
        args.model_dir,
        val_data_file,
        eval_script,
        val_eval_output,
        args.start_pass_id,
        args.end_pass_id,
        force_rerun=args.force_rerun)

    best_pass_id = find_best_pass(val_evals)

    tst_evals = run_eval(
        infer_obj,
        conf,
        args.model_dir,
        tst_data_file,
        eval_script,
        tst_eval_output,
        start_pass_id=best_pass_id,
        end_pass_id=best_pass_id,
        force_rerun=args.force_rerun)

    logger.info("Best Pass=%d" % best_pass_id)
    logger.info("Validation: %s" % val_evals[best_pass_id])
    logger.info("Test      : %s" % tst_evals[best_pass_id])


if __name__ == "__main__":
    main(parse_cmd())