import os
import json
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
import logging
import argparse
import numpy as np
from gluoncv.utils import viz
import boto3
import cv2

from urlparse import urlparse  # python 2

# from urllib.parse import urlparse  # python 3


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
s3 = boto3.resource('s3')

# construct the argument parser and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image-folder", required=True, help="folder containing images to do inference on")
ap.add_argument("-l", "--label-fname", required=True, help="label file name")
ap.add_argument("-f", "--output-fname", required=True, help="generated pdf file name")
ap.add_argument("-c", "--confidence-filter", required=False, default=0.2, help="confidence filter to use")

COL = 3
MAX_ROW = 5
SHAPE = 512
CLASSES = ['blue box', 'yellow box']


def read_manifest_file(file_path):
    with open(file_path, 'r') as f:
        output = [json.loads(line.strip()) for line in f.readlines()]
        return output


def download_file(uri, working_dir):
    fname = uri.split('/')[-1]
    image_local_path = os.path.join(working_dir, fname)
    s3_paths = urlparse(uri)
    s3_bucket = s3_paths.netloc
    s3_key = s3_paths.path.lstrip('/')
    s3.Bucket(s3_bucket).download_file(s3_key, image_local_path)
    return image_local_path


def title_page(title, pdf):
    '''Create a title page with only text.'''
    plt.figure(figsize=(10, 10), facecolor='white', dpi=100)
    plt.text(0.1, 0.5, s=title, fontsize=20)
    plt.axis('off')
    pdf.savefig()
    plt.close()


def page_loop(page_number, beginning_index, predictions, image_folder, pdf, confidence_filter):
    '''Loop over a single image page of the output pdf.'''

    row = min(MAX_ROW, int((len(predictions) - beginning_index) / COL + 1))
    last_index = min(len(predictions), beginning_index + row * COL)
    logger.info(
        "page number: {}, row number: {}, beginning index: {}, last index: {}".format(page_number, row, beginning_index,
                                                                                      last_index))

    fig, axes = plt.subplots(nrows=row, ncols=COL,
                             figsize=(20, 20),
                             facecolor='white', dpi=100)
    fig.suptitle('Page {}'.format(page_number), fontsize=24)

    for i in range(beginning_index, last_index):
        entry = predictions[i]
        try:

            image_path = os.path.join(image_folder, entry['image'])

            img = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
            # Resize image to fit network input
            # origs = cv2.resize(img, (SHAPE,SHAPE))
            origs = img
            # tensors, origs = yolo_transform.load_test(image_path, short=SHAPE)
            height, width, _c = origs.shape
            labels = np.array(entry['prediction'])

            if (len(entry['prediction']) > 0):
                cid = labels[:, 0]
                score = labels[:, 1]
                bbox = labels[:, range(2, 6)]

                bbox[:, (0, 2)] *= width
                bbox[:, (1, 3)] *= height
            else:
                cid = []
                score = []
                bbox = []
            img_name = entry['image'].split(".")[0]
            if row == 1:
                ax = axes[(i - beginning_index) % COL]
            else:
                ax = axes[int((i - beginning_index) / COL), (i - beginning_index) % COL]
            ax.set_title("{}".format(img_name))
            ax.get_xaxis().set_visible(False)
            ax.get_yaxis().set_visible(False)
            axis = viz.plot_bbox(origs, bbox, score, cid, thresh=confidence_filter,
                                 class_names=CLASSES, ax=ax)
        except:
            logger.error(json.dumps(entry, indent=2))
            raise
    next_page_number = page_number + 1
    next_beginning_index = last_index
    pdf.savefig()
    plt.close()

    return next_page_number, next_beginning_index


def get_ground_truth_labels(s3_uri):
    ground_truth_file = download_file(s3_uri, ".")
    validation_lines = read_manifest_file(ground_truth_file)
    ground_truth_labels = {}
    for line in validation_lines:
        img_name = line["source-ref"].split("/")[-1].split(".")[0]
        annotations = line["bb"]["annotations"]
        number_of_boxes = len(annotations)
        ground_truth_labels[img_name] = {"bb": annotations, "n": number_of_boxes}
    logger.info("read {} lines from {}".format(len(ground_truth_labels), s3_uri))
    return ground_truth_labels


def main():
    args = vars(ap.parse_args())
    image_folder = args["image_folder"]
    output_fname = args["output_fname"]
    label_fname = args['label_fname']
    confidence_filter_score = float(args["confidence_filter"])
    logger.info("doing inference for images in directory: {}".format(image_folder))
    logger.info("confidence filter: {}".format(confidence_filter_score))
    logger.info("output file name: {}".format(output_fname))
    predictions = read_manifest_file(label_fname)
    logger.info("label file name: {}. number of labels: {}".format(label_fname, len(predictions)))

    with PdfPages(output_fname) as pdf:
        title_page('Pig head ground truth labels\n confidence filter: {} \n'.format(confidence_filter_score), pdf)
        next_page_number, next_beginning_index = page_loop(0, 0, predictions, image_folder,
                                                           pdf=pdf,
                                                           confidence_filter=confidence_filter_score
                                                           )

        while next_beginning_index < len(predictions):
            next_page_number, next_beginning_index = page_loop(next_page_number, next_beginning_index, predictions,
                                                               image_folder,
                                                               pdf=pdf,
                                                               confidence_filter=confidence_filter_score
                                                               )


if __name__ == "__main__":
    main()