#!/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 = False
        self.testMode = False
        self.chromaIn1 = (0, 0, 255)
        self.chromaIn2 = (23, 19, 207)
        self.chromaIn3 = (24, 20, 169)
        self.blueLines = (70, 172, 255)
        self.darkBlue = (47, 115, 237)
        self.greenLines = (0, 255, 0)
        self.orangeLines = (228, 188, 0)
        self.redLines = (255, 0, 0)
        self.whiteLines = (255, 255, 255)
        self.linesColor = self.blueLines
        self.linesFeather = 0
        self.monitorLines = self.orangeLines
        self.monitorFeather = 0
        # 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 f_090_260_monitor_flash(self):
        if 0 <= self.i % 90 <= 5:
            f1 = self.i / 90 * 90
            self.monitorLines = self.animate((255, 255, 255), self.monitorLines, f1, f1 + 5)
            self.monitorFeather = self.animate(20, self.monitorFeather, f1, f1 + 5)

    def f_015_025_flash(self):
        self.linesColor = self.animate((255, 255, 255), self.linesColor)
        self.linesFeather = self.animate(5, self.linesFeather)

    def f_045_060_flash(self):
        self.linesColor = self.animate((255, 255, 255), self.linesColor)
        self.linesFeather = self.animate(10, self.linesFeather)

    def f_120_150_power_surge_up(self):
        self.linesColor = self.animate(self.linesColor, (255, 255, 255))
        self.linesFeather = self.animate(self.linesFeather, 50)
        self.monitorLines = self.animate(self.orangeLines, (119, 98, 0))
        self.monitorFeather = 0

    def f_151_230_full_power(self):
        self.linesColor = (255, 255, 255)
        self.linesFeather = 50
        self.monitorLines = (119, 98, 0)
        self.monitorFeather = 0

    def f_231_260_drain(self):
        self.linesColor = self.animate((255, 255, 255), self.linesColor)
        self.linesFeather = self.animate(50, self.linesFeather)
        self.monitorLines = self.animate((119, 98, 0), self.orangeLines)
        self.monitorFeather = 0

    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 __tronifyImage(self, image, drawable):
        if self.scale is False:
            w = pdb.gimp_image_width(image)
            h = pdb.gimp_image_height(image)
            if w == 1280: self.scale = 1 # 1280x720
            if w == 1920: self.scale = 1.5 # 1920x1080
            if w == 3840: self.scale = 3 # 3840x2160
            if self.scale is False: return
        s = self.scale
        c64Points = map(lambda x: s*x, (343, 553, 415, 570, 465, 608, 530, 720, 343, 720))

        # create channel: chroma key
        pdb.gimp_context_set_antialias(True)
        pdb.gimp_context_set_sample_criterion(SELECT_CRITERION_COMPOSITE)
        pdb.gimp_context_set_sample_threshold_int(25)
        pdb.gimp_image_select_color(image, CHANNEL_OP_REPLACE, drawable, self.chromaIn1)
        pdb.gimp_image_select_color(image, CHANNEL_OP_ADD, drawable, self.chromaIn2)
        pdb.gimp_image_select_color(image, CHANNEL_OP_ADD, drawable, self.chromaIn3)
        pdb.gimp_selection_grow(image, s*2) # encroach into the non-chroma some.
        # select-out bottom edge
        pdb.gimp_image_select_rectangle(image, CHANNEL_OP_SUBTRACT, 0, s*718, s*1280, s*2)
        channelChroma = pdb.gimp_selection_save(image)
        pdb.gimp_layer_add_alpha(drawable)
        pdb.gimp_drawable_edit_clear(drawable)

        # create channel: unitard
        pdb.gimp_context_set_antialias(False)
        pdb.gimp_context_set_sample_criterion(SELECT_CRITERION_S)
        pdb.gimp_context_set_sample_threshold_int(15)
        pdb.gimp_image_select_color(image, CHANNEL_OP_REPLACE, drawable, (255, 255, 255)) # full white
        pdb.gimp_context_set_sample_criterion(SELECT_CRITERION_H)
        pdb.gimp_context_set_sample_threshold_int(45)
        pdb.gimp_image_select_color(image, CHANNEL_OP_ADD, drawable, (206, 213, 248)) # cool white
        # shrink out face pinholes and grow past edge of unitard
        pdb.gimp_selection_shrink(image, s*3)
        pdb.gimp_selection_grow(image, s*4)
        # select out computer monitor area
        pdb.gimp_image_select_rectangle(image, CHANNEL_OP_SUBTRACT, 0, s*433, s*345, s*287)
        channelUnitard = pdb.gimp_selection_save(image)

        # # create channel: face
        # pdb.gimp_image_select_rectangle(image, CHANNEL_OP_REPLACE, 0, 0, s*1280, s*433)
        # pdb.gimp_image_select_item(image, CHANNEL_OP_SUBTRACT, channelChroma)
        # pdb.gimp_image_select_item(image, CHANNEL_OP_SUBTRACT, channelUnitard)
        # channelFace = pdb.gimp_selection_save(image)

        # create channel: lines
        # pdb.gimp_image_select_item(image, CHANNEL_OP_REPLACE, channelUnitard)
        pdb.gimp_context_set_antialias(True)
        pdb.gimp_context_set_sample_criterion(SELECT_CRITERION_V)
        pdb.gimp_context_set_sample_threshold_int(80)
        pdb.gimp_image_select_color(image, CHANNEL_OP_INTERSECT, drawable, (0, 0, 0))
        # select out the C64 keyboard
        pdb.gimp_image_select_polygon(image, CHANNEL_OP_SUBTRACT, len(c64Points), c64Points)
        # select-out bottom table edge (for some reason it gets "lined in")
        pdb.gimp_image_select_rectangle(image, CHANNEL_OP_SUBTRACT, s*503, s*718, s*777, s*2)
        # select-out right table end
        pdb.gimp_image_select_rectangle(image, CHANNEL_OP_SUBTRACT, s*1038, s*610, s*242, s*110)
        channelLines = pdb.gimp_selection_save(image)

        # create channel: remainder
        pdb.gimp_selection_all(image)
        pdb.gimp_image_select_item(image, CHANNEL_OP_SUBTRACT, channelChroma)
        # pdb.gimp_image_select_item(image, CHANNEL_OP_SUBTRACT, channelFace)
        pdb.gimp_image_select_item(image, CHANNEL_OP_SUBTRACT, channelUnitard)
        pdb.gimp_image_select_item(image, CHANNEL_OP_SUBTRACT, channelLines)
        # select-in bottom edge
        pdb.gimp_image_select_rectangle(image, CHANNEL_OP_ADD, 0, s*718, s*1280, s*2)
        channelRemainder = pdb.gimp_selection_save(image)

        # create channel: monitorLines
        pdb.gimp_context_set_antialias(True)
        pdb.gimp_context_set_sample_criterion(SELECT_CRITERION_V)
        pdb.gimp_context_set_sample_threshold_int(40)
        pdb.gimp_image_select_color(image, CHANNEL_OP_INTERSECT, drawable, (0, 0, 0))
        pdb.gimp_image_select_rectangle(image, CHANNEL_OP_SUBTRACT, s*345, 0, s*935, s*720)
        # select-out bottom edge
        pdb.gimp_image_select_rectangle(image, CHANNEL_OP_SUBTRACT, 0, s*718, s*1280, s*2)
        channelMonitorLines = pdb.gimp_selection_save(image)

        # apply chroma key mask
        if self.chromaOut is not False:
            pdb.gimp_image_select_item(image, CHANNEL_OP_REPLACE, channelChroma)
            pdb.gimp_context_set_foreground(self.chromaOut)
            pdb.gimp_edit_fill(drawable, 0)

        # apply remainder: add grainy noise and gray-ify
        pdb.gimp_image_select_item(image, CHANNEL_OP_REPLACE, channelRemainder)
        pdb.plug_in_hsv_noise(image, drawable, 4, 0, s*32, 0)
        pdb.gimp_hue_saturation(drawable, ALL_HUES, 0, -85, -100)

        # # apply face: add grainy noise and gray-ify
        # pdb.gimp_image_select_item(image, CHANNEL_OP_REPLACE, channelFace)
        # pdb.plug_in_hsv_noise(image, drawable, 4, 0, 32, 0)
        # pdb.gimp_hue_saturation(drawable, ALL_HUES, 0, -85, -100)

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

        # lines: colorarization
        pdb.gimp_image_select_item(image, CHANNEL_OP_REPLACE, channelLines)
        pdb.gimp_context_set_foreground(self.linesColor)
        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.linesFeather)
        pdb.gimp_edit_fill(drawable, 0)
        pdb.gimp_edit_fill(drawable, 0)

        # monitor lines: colorarization
        pdb.gimp_image_select_item(image, CHANNEL_OP_REPLACE, channelMonitorLines)
        pdb.gimp_context_set_foreground(self.monitorLines)
        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.monitorFeather)
        pdb.gimp_edit_fill(drawable, 0)
        pdb.gimp_edit_fill(drawable, 0)

    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 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()
