import argparse
import os
import os.path as osp

import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import mmcv
import numpy as np

try:
    import imageio
except ImportError:
    imageio = None


def parse_args():
    parser = argparse.ArgumentParser(description='Create GIF for demo')
    parser.add_argument(
        'image_dir',
        help='directory where result '
        'images save path generated by ‘analyze_results.py’')
    parser.add_argument(
        '--out',
        type=str,
        default='result.gif',
        help='gif path where will be saved')
    args = parser.parse_args()
    return args


def _generate_batch_data(sampler, batch_size):
    batch = []
    for idx in sampler:
        batch.append(idx)
        if len(batch) == batch_size:
            yield batch
            batch = []
    if len(batch) > 0:
        yield batch


def create_gif(frames, gif_name, duration=2):
    """Create gif through imageio.

    Args:
        frames (list[ndarray]): Image frames
        gif_name (str): Saved gif name
        duration (int): Display interval (s),
            Default: 2
    """
    if imageio is None:
        raise RuntimeError('imageio is not installed,'
                           'Please use “pip install imageio” to install')
    imageio.mimsave(gif_name, frames, 'GIF', duration=duration)


def create_frame_by_matplotlib(image_dir,
                               nrows=1,
                               fig_size=(300, 300),
                               font_size=15):
    """Create gif frame image through matplotlib.

    Args:
        image_dir (str): Root directory of result images
        nrows (int): Number of rows displayed, Default: 1
        fig_size (tuple): Figure size of the pyplot figure.
           Default: (300, 300)
        font_size (int): Font size of texts. Default: 15

    Returns:
        list[ndarray]: image frames
    """

    result_dir_names = os.listdir(image_dir)
    assert len(result_dir_names) == 2
    # Longer length has higher priority
    result_dir_names.reverse()

    images_list = []
    for dir_names in result_dir_names:
        images_list.append(mmcv.scandir(osp.join(image_dir, dir_names)))

    frames = []
    for paths in _generate_batch_data(zip(*images_list), nrows):

        fig, axes = plt.subplots(nrows=nrows, ncols=2)
        fig.suptitle('Good/bad case selected according '
                     'to the COCO mAP of the single image')

        det_patch = mpatches.Patch(color='salmon', label='prediction')
        gt_patch = mpatches.Patch(color='royalblue', label='ground truth')
        # bbox_to_anchor may need to be finetuned
        plt.legend(
            handles=[det_patch, gt_patch],
            bbox_to_anchor=(1, -0.18),
            loc='lower right',
            borderaxespad=0.)

        if nrows == 1:
            axes = [axes]

        dpi = fig.get_dpi()
        # set fig size and margin
        fig.set_size_inches(
            (fig_size[0] * 2 + fig_size[0] // 20) / dpi,
            (fig_size[1] * nrows + fig_size[1] // 3) / dpi,
        )

        fig.tight_layout()
        # set subplot margin
        plt.subplots_adjust(
            hspace=.05,
            wspace=0.05,
            left=0.02,
            right=0.98,
            bottom=0.02,
            top=0.98)

        for i, (path_tuple, ax_tuple) in enumerate(zip(paths, axes)):
            image_path_left = osp.join(
                osp.join(image_dir, result_dir_names[0], path_tuple[0]))
            image_path_right = osp.join(
                osp.join(image_dir, result_dir_names[1], path_tuple[1]))
            image_left = mmcv.imread(image_path_left)
            image_left = mmcv.rgb2bgr(image_left)
            image_right = mmcv.imread(image_path_right)
            image_right = mmcv.rgb2bgr(image_right)

            if i == 0:
                ax_tuple[0].set_title(
                    result_dir_names[0], fontdict={'size': font_size})
                ax_tuple[1].set_title(
                    result_dir_names[1], fontdict={'size': font_size})
            ax_tuple[0].imshow(
                image_left, extent=(0, *fig_size, 0), interpolation='bilinear')
            ax_tuple[0].axis('off')
            ax_tuple[1].imshow(
                image_right,
                extent=(0, *fig_size, 0),
                interpolation='bilinear')
            ax_tuple[1].axis('off')

        canvas = fig.canvas
        s, (width, height) = canvas.print_to_buffer()
        buffer = np.frombuffer(s, dtype='uint8')
        img_rgba = buffer.reshape(height, width, 4)
        rgb, alpha = np.split(img_rgba, [3], axis=2)
        img = rgb.astype('uint8')

        frames.append(img)

    return frames


def main():
    args = parse_args()
    frames = create_frame_by_matplotlib(args.image_dir)
    create_gif(frames, args.out)


if __name__ == '__main__':
    main()