import argparse
import json
import os

from image_prediction.pipeline import load_pipeline
from image_prediction.transformer.transformers.coordinate.pdfnet import PDFNetCoordinateTransformer
from image_prediction.utils.pdf_annotation import annotate_pdf


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("pdf")

    args = parser.parse_args()

    return args


def main(args):
    pipeline = load_pipeline(verbose=True)

    pdf_path = args.pdf

    with open(pdf_path, "rb") as f:
        predictions = list(pipeline(f.read()))

    print(json.dumps(predictions, indent=2))

    annotate_pdf(
        pdf_path, predictions, os.path.join("/tmp", os.path.basename(pdf_path.replace(".pdf", "_annotated.pdf")))
    )


if __name__ == "__main__":
    args = parse_args()
    main(args)