import argparse
from itertools import starmap
from pathlib import Path

import numpy as np
import pdf2image
from PIL import Image
from funcy import lmap

from cv_analysis.figure_detection.figure_detection_pipeline import make_figure_detection_pipeline
from cv_analysis.utils.draw import draw_rectangles


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--pdf_path", "-p", required=True)
    parser.add_argument("--output_folder", "-o", required=True)
    return parser.parse_args()


def annotate_figures(images):
    pipeline = make_figure_detection_pipeline()
    result = map(pipeline, images)
    annotated_images = starmap(draw_rectangles, zip(images, result))
    return annotated_images


def save_as_pdf(images, output_folder, file_name):
    Path(output_folder).mkdir(parents=True, exist_ok=True)
    images = lmap(Image.fromarray, images)
    images[0].save(f"{output_folder}/{file_name}_annotated_figures.pdf", save_all=True, append_images=images)


if __name__ == "__main__":
    args = parse_args()
    pages = lmap(np.array, pdf2image.convert_from_path(args.pdf_path))
    annotated_pages = annotate_figures(images=pages)
    save_as_pdf(annotated_pages, args.output_folder, Path(args.pdf_path).stem)