#!/usr/bin/env python

from gimpfu import *
import glob
import inspect
import os
import re
import sys

class Tronify:
    def __init__(self):
        # Method names of the form "f_42_55..." are used as image filters and
        # act upon a range of frame indices e.g. 42-55. Construct a list of
        # these filters in the form of [("f_42_55...", 42, 55), ...]
        p = re.compile(r"^f_(\d+)_(\d+).*$")
        self.filters = [(r.group(0), int(r.group(1)), int(r.group(2))) \
            for r in [p.search(f) for f in dir(Tronify)] if r]

        self.scale = 1
        self.testMode = False
        self.blue = (70, 172, 255)
        self.darkBlue = (47, 115, 237)
        self.red = (255, 0, 0)
        self.white = (255, 255, 255)
        self.ringsColor = self.blue
        self.ringsFeather = 0
        self.ringsThreshold = 70
        # self.chromaOut = False
        self.chromaOut = (0, 0, 0)

    def animate(self, initial, final, f1 = 0, f2 = 0):
        # linear
        if f1 == 0: f1, f2 = self.frameRange()
        def _animate(y1, y2):
            m = (y2 - y1) / float(f2 - f1)
            b = y1 - m * f1
            return int(m * self.i + b)
        if type(initial) is tuple:
            return tuple(map(_animate, initial, final))
        return _animate(initial, final)

    def animatePeak(self, initial, final, f1 = 0, f2 = 0):
        # linear peak at midpoint /\
        if f1 == 0: f1, f2 = self.frameRange()
        if f2 - f1 < 2: return final
        fMid = (f1 + f2) / 2.0
        def _animatePeak(y1, y2):
            if f1 <= self.i <= fMid:
                m = (y2 - y1) / float(fMid - f1)
                b = y1 - m * f1
            else:
                m = (y1 - y2) / float(f2 - fMid)
                b = y2 - m * fMid
            return int(m * self.i + b)
        if type(initial) is tuple:
            return tuple(map(_animatePeak, initial, final))
        return _animatePeak(initial, final)

    def frameRange(self):
        p = re.compile(r"^f_(\d+)_(\d+).*$")
        r = p.search(inspect.stack()[2][3])
        if not r: r = p.search(inspect.stack()[3][3])
        if not r: raise ValueError("Couldn't infer frame range.")
        return (int(r.group(1)), int(r.group(2)))

    def tronifyDir(self, threadId, numThreads, framesDir, testMode):
        self.testMode = testMode
        if testMode: print(filter(lambda f: f[0], self.filters))
        # Determine frames (and file extensions) to be processed.
        framesList = [re.match(r"frame-(\d+)\.(\w{3})", os.path.basename(f)).groups()
            for f in sorted(glob.glob(framesDir + "/frame-*[0-9].???"))
            if not os.path.isfile(f[0:-4] + ".done." + f[-3:])]
        segmentSize = len(framesList) // numThreads # e.g. 100 // 6 = 16
        remainder = len(framesList) % numThreads # e.g 100 % 6 = 4
        # We want e.g. 17, 17, 17, 17, 16, 16
        segmentSize += 1
        segmentStart = (threadId - 1) * segmentSize + 1
        segmentEnd = segmentStart + segmentSize - 1
        if threadId > remainder:
            segmentStart -= (threadId - 1) - remainder
            segmentEnd -= threadId - remainder
        # print("Thread " + str(threadId) + " has indices " + str(segmentStart) + "-" + str(segmentEnd))
        for i in range(segmentStart, segmentEnd + 1):
            n = framesList[i - 1][0]
            extension = framesList[i - 1][1]
            infile = os.path.join(framesDir, "frame-" + n + "." + extension)
            outfile = os.path.join(framesDir, "frame-" + n + ".done." + extension)
            image = pdb.gimp_file_load(infile, infile)
            drawable = pdb.gimp_image_get_active_layer(image)
            self.tronifyImage(image, drawable, int(n))
            pdb.gimp_file_save(image, drawable, outfile, outfile)
            pdb.gimp_image_delete(image)
            if testMode: break

    def tronifyImage(self, image, drawable, i):
        self.i = i
        # save instance variables
        state = vars(self).copy()
        del state["filters"]
        # apply filters and tronify
        filters = [r[0] for r in self.filters if r[1] <= i <= r[2]]
        for f in filters: getattr(self, f)()
        self.__tronifyImage(image, drawable)
        # restore instance variables
        for k,v in state.iteritems(): setattr(self, k, v)

    def __tronifyImage(self, image, drawable):
        s = self.scale

        # create channel: disc
        pdb.gimp_image_select_item(image, CHANNEL_OP_REPLACE, drawable) # all but alpha
        channelDisc = pdb.gimp_selection_save(image)

        # create channel: white
        pdb.gimp_context_set_antialias(True)
        pdb.gimp_context_set_sample_criterion(SELECT_CRITERION_V)
        pdb.gimp_context_set_sample_threshold_int(self.ringsThreshold)
        pdb.gimp_image_select_color(image, CHANNEL_OP_SUBTRACT, drawable, (0, 0, 0))
        channelWhite = pdb.gimp_selection_save(image)

        # create channel: rings
        pdb.gimp_image_select_item(image, CHANNEL_OP_REPLACE, channelDisc)
        pdb.gimp_image_select_item(image, CHANNEL_OP_SUBTRACT, channelWhite)
        channelRings = pdb.gimp_selection_save(image)

        # white: add grainy noise and gray-ify
        pdb.gimp_image_select_item(image, CHANNEL_OP_REPLACE, channelWhite)
        pdb.plug_in_hsv_noise(image, drawable, 4, 0, s*64, s*64)
        pdb.gimp_colorize(drawable, 199, 11, -56)

        # rings: colorarization
        pdb.gimp_image_select_item(image, CHANNEL_OP_REPLACE, channelRings)
        pdb.gimp_context_set_foreground(self.ringsColor)
        pdb.gimp_edit_fill(drawable, 0)
        pdb.script_fu_drop_shadow(image, drawable, 0, s*1, s*3.0, (0, 0, 0), 100, 0)
        pdb.gimp_selection_feather(image, s*self.ringsFeather)
        pdb.gimp_edit_fill(drawable, 0)
        pdb.gimp_edit_fill(drawable, 0)

        # apply chroma key mask
        if self.chromaOut is not False:
            pdb.gimp_context_set_background(self.chromaOut)
            pdb.gimp_layer_flatten(drawable)

    def f_001_138(self):
        self.ringsColor = self.darkBlue
        self.ringsFeather = 0

    def f_036_048(self):
        self.ringsColor = self.animate(self.darkBlue, self.white)
        self.ringsFeather = self.animate(20, 80)
        self.ringsThreshold = self.animate(self.ringsThreshold, 90)

    def f_049_063(self):
        self.ringsColor = self.animate(self.white, self.darkBlue)
        self.ringsFeather = self.animate(80, 20)
        self.ringsThreshold = self.animate(90, self.ringsThreshold)

    def f_082_092(self):
        self.ringsColor = self.animate(self.darkBlue, self.white)
        self.ringsFeather = self.animate(20, 80)
        self.ringsThreshold = self.animate(self.ringsThreshold, 90)

    def f_093_109(self):
        self.ringsColor = self.animate(self.white, self.darkBlue)
        self.ringsFeather = self.animate(80, 20)
        self.ringsThreshold = self.animate(90, self.ringsThreshold)

from gimpfu import *

def tronify_dir(outPipe, threadId, numThreads, framesDir, testMode):
    sys.stdout = open(outPipe, "w", 0)
    pngCompression = 1 # 0-9
    pdb.file_png_set_defaults(0, pngCompression, 1, 0, 0, 1, 1, 0, 0)
    t = Tronify()
    t.tronifyDir(threadId, numThreads, framesDir, testMode)
    sys.stdout.close()

register(
    "python_fu_tronify_dir",
    "TRONify",
    "Create TRON-like effects on series of images.",
    "David Fleming",
    "David Fleming",
    "2024",
    "TRONify Images...",
    "",
    [
        # (PF_IMAGE, "image", "Input image", None),
        # (PF_DRAWABLE, "drawable", "Input layer", None),
        (PF_STRING, "outPipe", "Output Pipe", ""),
        (PF_INT, "threadId", "Thread ID", 1),
        (PF_INT, "numThreads", "Number of Threads", 1),
        (PF_STRING, "framesDir", "Frames Directory", ""),
        (PF_BOOL, "testMode", "Test Mode", False),
    ],
    [],
    tronify_dir, menu="<Image>/Filters")

main()