#!/usr/bin/env python

import argparse
from collections import deque, namedtuple
import sys

import numpy as np
from matplotlib import animation
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation
from skimage.io import imread
from skimage.util import img_as_float


class Image(object):
    def __init__(self, filepath, normalize):
        if filepath.endswith(".npy"):
            data = np.load(filepath)
        else:
            data = img_as_float(imread(filepath))

        if normalize:
            data = data - data.min()
            data = data / (data.max() + 1e-7)

        self.filepath = filepath
        self.data = data


class ImageList(object):
    def __init__(self, filenames, normalize=False):
        self._filenames = filenames
        self._cache = deque([], 5)
        self._normalize = normalize

    def __len__(self):
        return len(self._filenames)

    def __getitem__(self, i):
        # Return it if it is in the cache
        for item in self._cache:
            if item[0] == i:
                return item[1]

        # Read it and add it in the cache
        item = (
            i,
            Image(self._filenames[i], self._normalize)
        )
        self._cache.append(item)

        return item[1]


def update(img, img_plot, title, args):
    if args.show_files:
        title.set_text(img.filepath)
    img_plot.set_array(img.data)

    return img_plot, title


def main(argv=None):
    parser = argparse.ArgumentParser(
        description="Animate a list of images read from stdin"
    )
    parser.add_argument(
        "--size",
        type=lambda x: tuple(map(int, x.split("x"))),
        default="6x4",
        help="Set the size of the figure in inches"
    )
    parser.add_argument(
        "--tight",
        action="store_true",
        help="Set the tight layout"
    )
    parser.add_argument(
        "--margins",
        type=lambda x: dict(
            (a, float(b)) for a,b in
            [xi.split("=") for xi in x.split(",") if len(xi) > 0]
        ),
        default="left=0,right=1",
        help="Adjust the margins of the plot"
    )
    parser.add_argument(
        "--delay", "-d",
        type=int,
        default=200,
        help="Delay between frames in milliseconds"
    )
    parser.add_argument(
        "--no_axis",
        dest="axis",
        action="store_false",
        help="Hide the axis"
    )
    parser.add_argument(
        "--show_files",
        action="store_true",
        help="Show the filenames as title"
    )
    parser.add_argument(
        "--vmin",
        type=float,
        default=0,
        help="Set the minimum of the heatmap range"
    )
    parser.add_argument(
        "--vmax",
        type=float,
        default=1,
        help="Set the maximum of the heatmap range"
    )
    parser.add_argument(
        "--cmap",
        default="viridis",
        help="Set the matplotlib colormap"
    )
    parser.add_argument(
        "--normalize",
        action="store_true",
        help="Normalize the images to [0, 1]"
    )
    parser.add_argument(
        "--save", "-s",
        help="Save the animation in a video file"
    )
    parser.add_argument(
        "--codec",
        help="Define a codec to save the video with"
    )
    parser.add_argument(
        "--extra_args",
        type=lambda x: x.split(",") if x else None,
        help="Set a comma separated list of extra arguments for the video"
    )

    args = parser.parse_args()

    images = ImageList([l.strip() for l in sys.stdin], args.normalize)
    figure, axis = plt.subplots(figsize=args.size)
    img_plot = axis.imshow(images[0].data, vmin=args.vmin, vmax=args.vmax,
                           animated=True, cmap=args.cmap)
    title = axis.text(0, 0, "", ha="left", va="top", color="white")
    if not args.axis:
        axis.axis("off")
    if args.tight:
        plt.tight_layout()
    plt.subplots_adjust(**args.margins)
    img_animation = FuncAnimation(
        figure,
        update,
        fargs=(img_plot, title, args),
        frames=images,
        interval=args.delay,
        repeat=False,
        save_count=2,
        blit=False
    )

    if args.save:
        writer = animation.writers["ffmpeg"](
            fps=int(round(1000./args.delay)),
            bitrate=1800,
            metadata=dict(artist="Home-tools 'movie'"),
            codec=args.codec,
            extra_args=args.extra_args
        )
        img_animation.save(args.save, writer=writer)
    plt.show()


if __name__ == "__main__":
    main()
