#!/usr/bin/env python

"""
pyltx is a script with some useful tools to working with LaTeX.
Run "pyltx -h" to see what it can do.
Some basic functionality include:

1. Compiling latex with multiple repetitions to ensure correct references,
   with options to clean previous output files and build a file with an
   integrated bibliography.
2. Consolidate all active files in a separate directory. Useful for generating
   a clean working directory for final upload for publication.
"""


import argparse
import subprocess
import os
import sys
import re
import os
import shutil
import glob


def compile(args):
    name = args.filename.rsplit(".")[0]
    if not os.path.exists("%s.tex" % name):
        print("%s.tex does not exist!" % name)
        sys.exit(-1)
    if args.clean:
        torm = [name + "." + ext for ext in ["aux", "blg", "log", "bbl"]]
        print("Removing %s." % ", ".join(torm))
        try:
            subprocess.check_output(["rm"] + torm)
        except subprocess.CalledProcessError:
            print("No files to delete! Continuing...")

    print("Running pdflatex...")
    output = subprocess.check_output(["pdflatex", name])
    if args.verbose:
        print(output)
    count = 1
    while ("Rerun" in output or "natbib Warning" in output) and count < 5:
        # This repetitive compilation process is necessary to ensure that
        # bibliography and references are properly updated.
        count += 1
        print("Rerunning %d" % count)
        print("Running bibtex...")
        output = subprocess.check_output(["bibtex", name])
        if args.verbose:
            print(output)
        print("Running pdflatex...")
        output = subprocess.check_output(["pdflatex", name])
        if args.verbose:
            print(output)

    if args.bbl:
        print("Building full bbl file...")
        with open("%s.bbl" % name) as f:
            bbl = f.read()

        with open("%s.tex" % name) as f:
            lines = []
            for l in f:
                if l.startswith("\\bibliography"):
                    lines.append(bbl.strip())
                else:
                    lines.append(l.strip())

        with open("%s_bbl.tex" % name, "wt") as f:
            f.write("\n".join(lines))

        print("%s_bbl.tex written!" % name)



def finalize(args):
    name = args.filename.rsplit(".")[0]
    if not os.path.exists("%s.tex" % name):
        print("%s.tex does not exist!" % name)
        sys.exit(-1)

    tex = "%s.tex" % name

    srcdir = os.path.dirname(os.path.abspath(tex))

    with open(tex) as f:
        contents = f.read()

    dest = args.output_dir

    os.mkdir(dest)

    def process_pattern(p):
        for m in re.finditer(p, contents, re.MULTILINE):
            fname = m.group(2)
            found = False
            for ext in ["", ".pdf", ".eps"]:
                if os.path.exists(fname + ext):
                    fname += ext
                    found = True
                    break
            if not found:
                fnames = glob.glob(fname + "*")
                if not fnames:
                    print("%s not found!")
                    sys.exit(-1)
                else:
                    fname = fnames[0]

            print("Copying %s..." % fname)
            shutil.copy(os.path.join(srcdir, fname),
                os.path.join(dest, fname))

    process_pattern("\\includegraphics(\[[^\]]+\])*\{([^\}]+)\}")
    process_pattern("bibliography(\[[^\]]+\])*\{([^\}]+)\}")

    shutil.copy(tex, os.path.join(dest, args.filename))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="""
    pyltx is a convenient script that provides useful tools for LaTeX. "
    "Type \"pyltx -h\" for help.""",
        epilog="Author: Shyue Ping Ong")

    subparsers = parser.add_subparsers()

    parser_compile = subparsers.add_parser("compile", help="Compile a file")

    parser_compile.add_argument("filename", metavar="filename", default="main",
                                type=str, nargs="?",
                                help="Filename to compile.")
    parser_compile.add_argument(
        "-v", "--verbose", dest="verbose",
        action="store_true",
        help="Verbose mode. Output from all commands are printed.")

    parser_compile.add_argument(
        "-c", "--clean", dest="clean",
        action="store_true",
        help="Clean the output files before compilation. Latex usually has the"
             " usual aux, blg, bbl files. This will delete all those files to"
             " force a clean compile.")
    parser_compile.add_argument(
        "-b", "--bbl", dest="bbl",
        action="store_true",
        help="Build a file with a full bbl integrated rather than in a "
             "separate bibtex. Some journals require this. The integrated file"
             " is named <filename>__bbl.tex.")
    parser_compile.set_defaults(func=compile)

    parser_finalize = subparsers.add_parser(
        "finalize",
        help="Consolidate all actually used files in a new directory. Useful"
             " for preparing final submission. Files which do not appear in"
             " the actual .tex file will not be included in the new dir.")

    parser_finalize.add_argument(
        "filename", metavar="filename", default="main",
        type=str, nargs="?", help="Filename to compile.")
    parser_finalize.add_argument("-o", "--output_dir", dest="output_dir",
        type=str, default="final",
        help="Directory to output files to.")
    parser_finalize.set_defaults(func=finalize)

    args = parser.parse_args()

    try:
        a = getattr(args, "func")
    except AttributeError:
        parser.print_help()
        sys.exit(0)
    args.func(args)