#!/usr/bin/env python

import argparse
from functools import lru_cache
from pathlib import Path
import re

import cv2
import numpy as np

from tqdm import tqdm


def speed_type(s):
    try:
        s = float(s)
        def constant_speed(t):
            return s
        return constant_speed
    except ValueError:
        def u(x):
            return 1.0 if x > 0 else 0
        def pulse(x, t_on, t_off):
            remainder = x % (t_on + t_off)
            return 1 if remainder < t_on else 0
        def fspeed(t):
            return eval(s, dict(u=u, t=t, np=np, pulse=pulse))
        return fspeed


class InfiniteList:
    def __init__(self, lst):
        self.lst = lst

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

    def __getitem__(self, idx):
        return self.lst[idx % len(self)]


def collect_files(folder, pattern):
    files = []
    folder = Path(folder)
    for f in folder.iterdir():
        if re.match(pattern, f.name):
            files.append(f)
    return files


def image_reader(resize, cache_size=32):
    @lru_cache(maxsize=cache_size)
    def _reader(f):
        if isinstance(f, Path):
            f = str(f)
        im = cv2.imread(f)
        if im.shape != resize:
            im = cv2.resize(im, resize)
        return im
    return _reader


class InfiniteAnimations:
    def __init__(self, *, frames, rows, columns, image_reader):
        self.frames = frames
        self.rows = rows
        self.columns = columns or len(frames)
        self.reader = image_reader

        if self.rows * self.columns != len(self.frames):
            raise ValueError(("rows*columns should equal the "
                              "number of animations"))

        frame_size = image_reader(self.frames[0][0]).shape[:2]
        self.frame_width = frame_size[1]
        self.frame_height = frame_size[0]
        self.width = self.columns * self.frame_width
        self.height = self.rows * self.frame_height

    def copy_to(self, canvas, x, y, frame_id):
        size = canvas.shape[:2]
        xmin = x
        xmax = x + size[1]
        ymin = y
        ymax = y + size[0]
        cx = cy = 0

        while cy < size[0]:
            x = xmin
            cx = 0
            while cx < size[1]:
                idx_x = (x % self.width) // self.frame_width
                idx_y = (y % self.height) // self.frame_height
                idx = idx_y * self.columns + idx_x

                frame = self.reader(self.frames[idx][frame_id])

                sx = (x % self.width) % self.frame_width
                sy = (y % self.height) % self.frame_height
                src = frame[sy:, sx:]
                tgt = canvas[cy:, cx:]

                if src.shape[0] > tgt.shape[0]:
                    src = src[:tgt.shape[0]]
                else:
                    tgt = tgt[:src.shape[0]]

                if src.shape[1] > tgt.shape[1]:
                    src = src[:, :tgt.shape[1]]
                else:
                    tgt = tgt[:, :src.shape[1]]

                tgt[...] = src

                cx += tgt.shape[1]
                x += tgt.shape[1]
            y += tgt.shape[0]
            cy += tgt.shape[0]


def main(argv=None):
    parser = argparse.ArgumentParser(
        description="Animate images on a moving 2d plane"
    )

    parser.add_argument(
        "animations",
        nargs="+",
        help="The folders containing the animation frames"
    )

    parser.add_argument(
        "--resize",
        type=lambda x: tuple(int(xi) for xi in x.split("x")),
        default=(512, 512),
        help="Resize the images as they are being read"
    )
    parser.add_argument(
        "--pattern",
        default=".*?[0-9]{3}.png$",
        help=("Regex pattern to match with the filename for the frames of each "
              "animation (default: '.*?[0-9]{3}.png$')")
    )
    parser.add_argument(
        "--canvas",
        type=lambda x: tuple(int(xi) for xi in x.split("x")),
        default=(512, 512),
        help="The size of the final canvas (default: 512x512)"
    )
    parser.add_argument(
        "--columns",
        type=int,
        default=0,
        help=("Define the number of columns with separate animations "
              "(default: 0 which means all)")
    )
    parser.add_argument(
        "--rows",
        type=int,
        default=1,
        help="Define the number of rows with separate animations (default: 1)"
    )
    parser.add_argument(
        "--xspeed",
        type=speed_type,
        default=speed_type(1),
        help=("How many horizontal pixels to add to the initial position "
              "every frame (default: 1)")
    )
    parser.add_argument(
        "--yspeed",
        type=speed_type,
        default=speed_type(0),
        help=("How many vertical pixels to add to the initial position "
              "every frame (default: 1)")
    )
    parser.add_argument(
        "--fps",
        type=int,
        default=25,
        help="How many frames per second to use for the video (default: 25)"
    )
    parser.add_argument(
        "--n_frames",
        type=int,
        default=100,
        help="How many frames to write (default: 100)"
    )
    parser.add_argument(
        "--output_file",
        default="out.mp4",
        help="Define the output file (default: out.mp4)"
    )

    args = parser.parse_args(argv)

    animation = InfiniteAnimations(
        frames=[
            InfiniteList(sorted(collect_files(animation, args.pattern)))
            for animation in args.animations
        ],
        rows=args.rows,
        columns=args.columns,
        image_reader=image_reader(args.resize)
    )
    canvas = np.zeros(args.canvas + (3,), dtype=np.uint8)

    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    writer = cv2.VideoWriter(
        args.output_file,
        fourcc,
        args.fps,
        (canvas.shape[1], canvas.shape[0])
    )
    x = 0
    y = 0
    for i in tqdm(range(args.n_frames)):
        animation.copy_to(canvas, int(round(x)), int(round(y)), i)
        x += args.xspeed(i)
        y += args.yspeed(i)
        writer.write(canvas)
    writer.release()


if __name__ == "__main__":
    main()
