#!/usr/bin/env python3

import os
import re
import subprocess as subp
import tempfile
from typing import List

import numpy as np
from PIL import Image, ImageOps, ImageFilter, ImageMorph, ImageChops, ImageDraw, ImageFont
try:
    import scipy.ndimage
    import scipy.signal

    HAVE_SCIPY = True
except ModuleNotFoundError:
    HAVE_SCIPY = False

## BEGIN Thresholding

def threshold(image: Image, level: int) -> Image:
    """Threshold an image at the given level."""
    return image.convert("L").point(lambda p: p >= level and 255, mode="1")

def compute_otsu_level(image: Image) -> int:
    """Compute the Otsu level for thresholding."""
    hist = image.histogram()
    total = sum(hist)

    sum_b = 0
    w_b = 0
    mx = 0
    level = 0
    sum_1 = sum((a*b for a, b in zip(range(256), hist)))
    for i in range(1, 256):
        w_f = total - w_b
        if w_b > 0 and w_f > 0:
            m_f = (sum_1 - sum_b) / w_f
            val = w_b * w_f * ((sum_b / w_b) - m_f) * ((sum_b / w_b) - m_f)
            if val >= mx:
                level = i
                mx = val
        w_b += hist[i]
        sum_b += (i-1) * hist[i]

    return level

def otsu(image: Image) -> Image:
    """Threshold an image using the Otsu level.

    Otsu thresholding typically does very well when the image is first upsampled, then    thresholded, e.g. via

      otsu(img.resize((img.width*2, img.height*2), resample=Image.Resampling.LANCZOS))

    This is useful when e.g. the original image is low-resolution; upsampling then
    thresholding will typically give a better result than thresholding and then
    upsampling.
    """
    image = image.convert("L")
    o = compute_otsu_level(image)
    return image.point(lambda i: i > o, mode="1")

def local_otsu(image, window_size=100):
    """Threshold the image via the local Otsu method."""
    output = Image.new("1", image.size)
    for x_window in range(math.ceil(image.width / window_size)):
        for y_window in range(math.ceil(image.height / window_size)):
            a = x_window*window_size
            b = min(x_window*window_size + window_size, image.width)
            c = y_window*window_size
            d = min(y_window*window_size + window_size, image.height)

            output.paste(otsu(image.crop((a, c, b, d))), (a, c))

    return output

## END Thresholding

## BEGIN Image processing

def difference(a: Image, b: Image, level: int = 50):
    """Compute a difference mask of two images.

    We do:

    1. Compute the difference between `a` and `b`
    2. Threshold this to remove noise using `level`
    3. Use grayscale dilation to expand this slightly
    4. Threshold
    """
    ## TODO: Is step 2 necessary? At worst we should just change some artifacting
    return ImageMorph.MorphOp(op_name="dilation8").apply(threshold(ImageChops.difference(a, b), level).convert("L"))[1].convert("1")

def remove_line(a: Image, b: Image, c: Image, level: int = 50) -> Image:
    """Remove a moving line from a sequence of frames.

    The three images `a`, `b`, and `c` should each have the line in a different spot.
    We do:

    1. Compute the difference mask between `a` and `b`
    2. Paste these parts of image `c` into `a`.

    If the images are of varying quality, `a` should be the best.
    """
    a = a.copy()
    a.paste(c, None, mask=difference(a, b, level=level))
    
    return a

def do_savgol(scene: Image, cutoff: float) -> Image:
    """Run a 2D Savitzky-Golay filter on the image, then threshold at `cutoff`.

    This is very good at smoothing out thresholded images and should be used as the
    final step; it is computationally expensive.
    """
    if not HAVE_SCIPY:
        raise ModuleNotFoundError("Need scipy for savgol filtering")
    a = np.array(scene).astype(float)
    b = scipy.signal.savgol_filter(
        scipy.signal.savgol_filter(a, 11, 4, axis=1), 
        11, 
        4, 
        axis=0
    )
    return Image.fromarray(b > cutoff)

def split_lines(image: Image, min_dist: int = 15) -> List[Image]:
    """Split an image into lines, each separated by at least `min_dist`."""
    ## TODO: Morphological opening here to remove noise?
    image = otsu(image)
    arr = np.array(image)

    extents = []
    last = None
    for n, row in enumerate(arr):
        if False not in row:
            if last is None:
                continue
            else:
                extents.append((last, n))
                last = None
        else:
            if last is None:
                last = n

    if last is not None:
        extents.append((last, len(arr)))

    found = True
    while found:
        found = False
        for i in range(len(extents)-1):
            a, b = extents[i]
            a2, b2 = extents[i+1]
            if a2 - b < min_dist:
                extents.pop(i)
                extents.pop(i)
                extents.insert(i, (a, b2))
                found = True
                break

    return extents

## END Image processing

## BEGIN Crop/pad

def double_bbox(image: Image) -> [int, int, int, int]:
    """Get the bounding box of an image possibly surrounded by a black then white 
    outline.
    """
    image = threshold(image, 95)
    left, upper, right, lower = image.getbbox()

    other = image.crop((left, upper, right, lower))
    left2, upper2, right2, lower2 = ImageOps.invert(other).getbbox()

    return left2+left, upper2+upper, right2+left, lower2+upper

def crop_scene(i: Image, new_width: int) -> Image:
    """Remove the bounding box from a scene and resize to match `new_width`."""
    width, height = i.width, i.height
    left, upper, right, lower = double_bbox(i)
    i = i.crop((left, upper, right, lower))
    
    new_height = int((lower - upper)*new_width/(right-left)+1)
    ri = i.resize((new_width, new_height))

    return ri

def pad_scene(scene: Image, padding: int = 10) -> Image:
    """Add padding to a scene."""
    i = Image.new("L", size=(scene.width, scene.height+2*padding), color=255)
    i.paste(scene, (0, padding))

    return i

## END Crop/pad

## BEGIN FFmpeg

def get_scenes(file: str, cutoff: float = 0.05) -> List[float]:
    """Use FFmpeg to get a list of scenes in a video file."""
    res = subp.check_output(["ffmpeg", "-i", file, "-filter:v", "select='gt(scene,%s)',showinfo" % str(cutoff), "-f", "null", "-"], stderr=subp.STDOUT)

    ## FFmpeg's output is annoying to parse...
    times = []
    for line in res.decode("utf-8").split('\n'):
        if not "showinfo" in line:
            continue
        val = re.search(r'pts_time:([0-9.]+)', line)
        if not val:
            continue
        times.append(float(val.group(1)))

    return [0] + times

def do_screenshots(file: str, td: str, fps: int = 1) -> List[str]:
    """Use FFmpeg to take `fps` screenshots per second of `file`, storing them in 
    the folder `td`.

    Space-wise this is wasteful but is faster/simpler than taking only the screenshots
    we need.
    """
    subp.check_call(["ffmpeg", "-i", file, "-q:v", "3", "-vf", "fps=%d" % fps, os.path.join(td, "%09d.jpg")], stderr=subp.DEVNULL)
    snapshots = []
    for f in os.listdir(td):
        if not f.endswith(".jpg"):
            continue
        time = int(f.split('.')[0])
        snapshots.append((time, os.path.join(td, f)))

    return sorted(snapshots)

## END FFmpeg

## BEGIN Scenes

def filter_scenes(scenes: List[float], cutoff: float = 1) -> List[float]:
    """Filter out scenes that are too close together (false positives)."""
    new = []
    last = None
    for s in scenes[::-1]:
        if last is None or last - s > cutoff:
            new.append(s)
            last = s

    return list(reversed(new))

def find_scenes(f: str, ffmpeg_cutoff: float, line: bool = False, scene_cutoff: float = 1, fps: int = 1) -> List[Image]:
    """Return a list of scenes in the file `f`."""
    ## Find the times we need
    times = filter_scenes(get_scenes(f, ffmpeg_cutoff), scene_cutoff)
    with tempfile.TemporaryDirectory() as td:
        ## Take screenshots
        snapshots = do_screenshots(f, td, fps=fps)

        ## Open the screenshots we need
        use = []
        for n, (t, nt) in enumerate(zip(times, times[1:]+[None]), 1):
            ## Figure out which screenshots work for this scene
            available = [b for a, b in snapshots if t < a and (nt is None or a < nt)]

            if line:
                ## Grab three shots to remove the line
                p1 = Image.open(available[len(available)//3])
                p2 = Image.open(available[len(available)//2])
                p3 = Image.open(available[2*(len(available)//3)])

                use.append(remove_line(p2, p1, p3))
            else:
                ## Grab the middle, usually the compression is better
                pick =  available[len(available)//2]
                use.append(Image.open(pick))

        return use

def split_scene(scene: Image, min_dist: int = 10) -> List[Image]:
    """Split a scene into lines."""
    todo = []
    extents = split_lines(scene, min_dist)
    for a, b in extents:
        todo.append(scene.crop((0, a, scene.width, b)))
    return todo

def merge_lines(images: List[Image], aspect: float = 8.5/11, center: bool = True, breaks: List[int] = None, center_breaks: bool = False) -> List[Images]:
    """Merge lines into images of maximum aspect ratio `aspect`.

    The resulting images will have the same width as the widest input line.
    
    The result may have aspect ratio larger than `aspect` if an input image already
    does; otherwise, they will all have height `max_width / aspect`.

    :param breaks: A list of forced page-breaks (e.g. page-turns).
    :param center: Center images horizontally.
    :param center_breaks: Space out lines on below-height pages.
    """
    width = max(i.width for i in images)

    todo = []
    output = []
    broken = []
    ## Split input images into groups of appropriate size
    for n, i in enumerate(images):
        height = sum((j.height for j in todo))
        if todo and (width/(height+i.height)) < aspect or n-1 in breaks:
            output.append(todo)
            broken.append(len(output))
            todo = []

        todo.append(i)

    output.append(todo)

    ## Merge them together
    for n, todo in enumerate(output, 1):
        height = sum((j.height for j in todo))
        outimage = Image.new("L", (width, max(height, int(width/aspect + 1))), color=255)

        if center and (center_breaks or n not in broken):
            padding = int((width/aspect - height)/(len(todo)+1))
        else:
            padding = 0

        cury = padding
        for t in todo:
            outimage.paste(t, (0, cury))
            cury += t.height + padding

        yield outimage

DEFAULT_FONT = ImageFont.load_default(64)

def number_images(images: List[Image], start: int = 0, in_place: bool = False, font = DEFAULT_FONT) -> List[Image]:
    res = []
    for n, i in enumerate(images, start):
        if not in_place:
            i = i.copy()
        d = ImageDraw.Draw(i)
        d.text((10, 10), str(n), font=font, fill=0)

        res.append(i)

    return res

## END Scenes

if __name__ == "__main__":
    import argparse as ap

    parser = ap.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)

    find = subparsers.add_parser("find", help="Find scenes in video")
    find.add_argument("input", help="Input video to process")
    find.add_argument("-c", "--cutoff", type=float, default=0.03, help="Default scene cutoff for FFmpeg")
    find.add_argument("-l", "--line", action="store_true", help="Remove scrolling line")
    find.add_argument("-t", "--scene-time", type=float, default=2, help="Length of shortest scene")
    find.add_argument("-f", "--format", help="Output format", default="png")
    find.add_argument("outdir", help="Output directory")

    merge = subparsers.add_parser("merge", help="Merge scenes into a PDF")
    merge.add_argument("input", help="Input directory")
    merge.add_argument("-l", "--line-dist", default=10, type=int, help="Minimum distance between lines")
    merge.add_argument("-a", "--aspect", help="Desired aspect ratio (width/height)", type=float, default=8.5/11)
    merge.add_argument("-p", "--padding", help="Padding between cropped lines", type=int, default=50)
    merge.add_argument("-e", "--center", action="store_true", help="Horizontally center lines")
    merge.add_argument("-b", "--center-breaks", action="store_true", help="Vertically center breaks on a page")
    merge.add_argument("-s", "--savgol", type=float, default=0, help="Savgol threshold (0 to disable)")
    merge.add_argument("-o", "--otsu", action="store_true", help="Run Otsu thresholding")
    merge.add_argument("-f", "--format", help="Output format", default="png")
    merge.add_argument("-d", "--dpi", type=int, default=300, help="PDF DPI")
    merge.add_argument("-n", "--number", action="store_true", help="Number lines in output")
    merge.add_argument("--no-merge", action="store_true", help="Don't merge the output, save them to the output folder instead")
    merge.add_argument("output", help="Output file")
    merge.add_argument("forced_breaks", nargs="*", type=int, help="Forced page breaks")

    args = parser.parse_args()

    if args.command == "find":
        os.makedirs(args.outdir, exist_ok=True)

        scenes = find_scenes(args.input, args.cutoff, line=args.line, scene_cutoff=args.scene_time)

        fmt = "out%07d."+args.format
        for n, s in enumerate(scenes):
            s.save(os.path.join(args.outdir, fmt % n))
    elif args.command == "merge":
        if args.no_merge:
            os.makedirs(args.output, exist_ok=True)

        scenes = []
        for f in sorted(os.listdir(args.input)):
            if not f.endswith('.'+args.format):
                continue
            scenes.append(Image.open(os.path.join(args.input, f)))

        new = []
        width = max(s.width for s in scenes)
        for s in scenes:
            s = crop_scene(s, width)
            for r in split_scene(s, args.line_dist):
                new.append(pad_scene(r, args.padding))

        if args.number:
            number_images(new, in_place=True)

        with tempfile.TemporaryDirectory() as outdir:
            outfiles = []
            fmt = "out%07d."+args.format

            for n, done in enumerate(merge_lines(new, aspect=args.aspect, center=args.center, center_breaks=args.center_breaks, breaks=args.forced_breaks), 1):
                name = os.path.join(outdir if not args.no_merge else args.output, fmt % n)
                outfiles.append(name)
                if args.otsu:
                    done = otsu(done.resize((done.width*2, done.height*2), resample=Image.Resampling.LANCZOS))
                if args.savgol > 0:
                    done = do_savgol(done, args.savgol)
                done.save(name)

            if not args.no_merge:
                subp.check_call(["magick", "-density", str(args.dpi), *outfiles, args.output])