magpie.py

#!/usr/bin/env python3
#
# Copyright (C) 2024
#
"""
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

import os
import re
import subprocess
import sys
import warnings

import userhook

from common import FloatFormat, Profile
from string import Template


LICENSE_HEADER = "".join(map("// {}\n".format, __doc__.splitlines()))

# format: (texture format, swizzle, cast to go to texture format, values to pad cast)
format_map = {
    "float": ("R16_FLOAT", ".x", "", []),
    "vec2":  ("R16G16_FLOAT", ".xy", "", []),
    "float2":  ("R16G16_FLOAT", ".xy", "", []),
    "vec3":  ("R16G16B16A16_FLOAT", ".xyz", "vec4", ["0.0"]),
    "float3":  ("R16G16B16A16_FLOAT", ".xyz", "float4", ["0.0"]),
    "vec4":  ("R16G16B16A16_FLOAT", ".xyzw", "", []),
    "float4":  ("R16G16B16A16_FLOAT", ".xyzw", "", [])
}


class MagpieBase(userhook.UserHook):
    HEADERS = ["PASS", "DESC", "STYLE", "IN", "OUT", "BLOCK_SIZE", "NUM_THREADS"]
    output_width = None
    output_height = None
    compute_pass = 0

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.registered_textures = {}
        self.out_textures = {}

    def magpie_header(self, sortname=None):
        scriptpath = os.path.dirname(os.path.realpath(__file__))

        normalized_argv = [(os.path.relpath(x, scriptpath) if ('\\' in x or '/' in x) else x) for x in sys.argv]

        giturl = subprocess.run(["git", "-C", scriptpath, "remote", "get-url", "origin"], capture_output=True).stdout.decode().strip()
        giturl = re.sub(r"^.*@", "https://", giturl)
        giturl = re.sub(".git$", "", giturl)

        headers = []
        headers += ["// This file is generated by the scripts available at %s" % giturl]
        headers += ["// Please don't edit this file directly."]
        headers += ["// Generated by: %s" % " ".join(normalized_argv)]
        headers += [LICENSE_HEADER]
        headers += [
            "//!MAGPIE EFFECT",
            "//!VERSION 4"
        ]
        headers += ["//!SORT_NAME %s" % sortname] if sortname else []
        headers += [""]

        return "\n".join(headers) + "\n"

    def set_transform(self, mul_x, mul_y, offset_x, offset_y):
        hook = self.hook[0]
        if hook == "INPUT":
            base_width = "INPUT_WIDTH"
            base_height = "INPUT_HEIGHT"
        else:
            base_width = self.out_textures[hook]["width"]
            base_height = self.out_textures[hook]["height"]

        self.output_width = "%s * %s" % (base_width, mul_x)
        self.output_height = "%s * %s" % (base_height, mul_y)

        if offset_x != 0.0 or offset_y != 0.0:
            warnings.warn("Offsets aren't implemented. Defined offset: %s %s" % (offset_x, offset_y))

    def tex_headers(self, name, filename=None, format=None, filter="POINT", width=None, height=None):
        headers  = ["//!TEXTURE"]
        headers += ["//!SOURCE %s"  % filename] if filename else []
        headers += ["//!FORMAT %s"  % format  ] if format   else []

        # OUTPUT is a special texture on MagpieFX v4.
        # Its size should be set to the size of the last pass
        # If it won't be set then effect supports arbitrary resizing
        if name == "OUTPUT":
            headers += ["//$output_width"]
            headers += ["//$output_height"]
        else:
            headers += ["//!WIDTH  %s"  % width   ] if width    else []
            headers += ["//!HEIGHT %s"  % height  ] if height   else []

        headers += ["Texture2D %s;" % name    ]
        headers += [""]

        if filter:
            headers += [self.sampler_headers(name, filter).rstrip(), ""]

        self.registered_textures[name] = {
            'filename': filename,
            'format': format,
            'width': width,
            'height': height
        }
        return "\n".join(headers) + "\n"

    def sampler_headers(self, name, filter="POINT"):
        headers = [
            "//!SAMPLER",
            "//!FILTER %s" % filter,
            "SamplerState sam_%s;" % name,
            ""
        ]
        return "\n".join(headers) + "\n"

    def generate_tex_magpie(self, lut_name, weights, lut_width, lut_height, float_format=FloatFormat.float32, overwrite=False):
        from dds import DDSTexture

        assert len(weights) == lut_width * lut_height * 4

        tex_format, short_format_name = {
            FloatFormat.float16dx: ("R16G16B16A16_FLOAT", "f16"),
            FloatFormat.float32dx: ("R32G32B32A32_FLOAT", "f32")
        }[float_format]

        ddst = DDSTexture()
        ddst.format = tex_format
        ddst.width = lut_width
        ddst.height = lut_height
        ddst.data = weights

        mode = 'xb'
        if overwrite:
            mode = 'wb'

        filename = "%s_%s.dds" % (lut_name, short_format_name)
        try:
            with open(filename, mode=mode) as f:
                f.write(bytes(ddst))
        except FileExistsError:
            if not overwrite:
                pass
            else:
                raise

        return self.tex_headers(lut_name, filename=filename, format=tex_format, filter="LINEAR")

    def hlsl_defines(self):
        return """$temp_textures
//!COMMON
#include "prescalers.hlsli"

#define LAST_PASS $last_pass

"""

    def reset(self):
        super().reset()

        for name in self.HEADERS:
            self.header[name] = None

        self.header["IN"] = list(self.hook)
        self.header["OUT"] = "OUTPUT"

    def bind_tex(self, tex):
        self.header["IN"].append(tex)

    def save_tex(self, tex):
        self.header["OUT"] = tex
        assert tex not in self.out_textures

        self.out_textures[tex] = {
            'width': self.output_width or "INPUT_WIDTH",
            'height': self.output_height or "INPUT_HEIGHT",
            'format': "R16G16B16A16_FLOAT",
            'sample_format': "float4"
        }

    def save_format(self, value):
        assert value in format_map
        format = format_map[value][0]
        self.out_textures[self.header["OUT"]]["sample_format"] = value
        self.out_textures[self.header["OUT"]]["format"] = format

    def set_compute(self, bw, bh, tw=None, th=None):
        self.header["BLOCK_SIZE"] = "%d, %d" % (bw, bh)
        if tw and th:
            self.header["NUM_THREADS"] = "%d, %d" % (tw, th)
        else:
            self.header["NUM_THREADS"] = "%d, %d" % (bw, bh)

    def setup_image_conversion(self):
        GLSL = self.add_glsl

        try:
            profile = self.profile
        except AttributeError:
            # assume LUMA if profile is not defined
            profile = Profile.luma

        if profile == Profile.luma:
            GLSL("#define GET_SAMPLE(x) dot(x.rgb, rgb2y)")
        else:
            GLSL("#define GET_SAMPLE(x) x")

        if self.header["OUT"] != "OUTPUT":
            tex_name = self.header["OUT"]
            tex = self.out_textures[tex_name]
            sample_format = tex["sample_format"]
            _, swizzle, cast, pack_values = format_map[sample_format]

            GLSL("#define imageStore(out_image, pos, val) imageStoreOverride(pos, val%s)" % swizzle)
            GLSL("void imageStoreOverride(uint2 pos, %s value) {" % sample_format)
            GLSL("    %s[pos] = %s(value %s);" % (tex_name, cast, ", ".join([""] + pack_values) ))
            GLSL("}")
        elif profile == Profile.luma:
            # out_image is not used on magpie
            GLSL("#define imageStore(out_image, pos, val) imageStoreOverride(pos, val.x)")
            GLSL("""
void imageStoreOverride(uint2 pos, float value) {
    float2 UV = mul(rgb2uv, INPUT.SampleLevel(sam_INPUT_LINEAR, HOOKED_map(pos), 0).rgb);
    OUTPUT[pos] = float4(mul(yuv2rgb, float3(value.x, UV)), 1.0);
}""")
        else:
            GLSL("#define imageStore(out_image, pos, val) imageStoreOverride(pos, val)")
            GLSL("""
void imageStoreOverride(uint2 pos, float4 value) {
    OUTPUT[pos] = value;
}""")

    def lookup_texture(self, tex):
        try:
            return self.registered_textures[tex]
        except KeyError:
            return self.out_textures[tex]

    def texture_defines(self):
        GLSL = self.add_glsl

        GLSL("")
        self.setup_image_conversion()

        for tex in self.header["IN"]:

            GLSL("")

            width = self.lookup_texture(tex)["width"]
            height = self.lookup_texture(tex)["height"]
            sample_format = self.lookup_texture(tex).get("sample_format", "vec4")
            swizzle = format_map[sample_format][1]
            if tex == "INPUT":
                GLSL("#define {0}_tex(pos) GET_SAMPLE(vec4(texture({0}, pos)))".format(tex))
                GLSL("static const float2 {0}_size = float2(GetInputSize());".format(tex))
                GLSL("static const float2 {0}_pt = float2(GetInputPt());".format(tex))
            elif width and height:
                GLSL("#define {0}_tex(pos) ({1}(texture({0}, pos){2}))".format(tex, sample_format, swizzle))
                width = width.replace("INPUT_WIDTH", "GetInputSize().x").replace("INPUT_HEIGHT", "GetInputSize().y")
                height = height.replace("INPUT_WIDTH", "GetInputSize().x").replace("INPUT_HEIGHT", "GetInputSize().y")
                GLSL("static const float2 {0}_size = float2({1}, {2});".format(tex, width, height))
                GLSL("static const float2 {0}_pt = float2(1.0/({0}_size.x), 1.0/({0}_size.y));".format(tex))
            else:
                GLSL("#define {0}_tex(pos) (vec4(texture({0}, pos)))".format(tex))

        GLSL("")
        GLSL("#define HOOKED_tex(pos) {0}_tex(pos)".format(self.hook[0]))
        GLSL("#define HOOKED_size {0}_size".format(self.hook[0]))
        GLSL("#define HOOKED_pt {0}_pt".format(self.hook[0]))

    def function_header_compute(self):
        GLSL = self.add_glsl

        GLSL("")
        GLSL("#define CURRENT_PASS $current_pass")

        self.texture_defines()

        GLSL("")
        GLSL("void Pass${current_pass}(uint2 blockStart, uint3 threadId) {")

    def finish(self, text):
        mappings = {
            "output_width": "",
            "output_height": "",
            "last_pass": self.compute_pass
        }
        if self.output_width and self.output_height:
            mappings["output_width"]  = "!WIDTH  %s" % self.output_width
            mappings["output_height"] = "!HEIGHT %s" % self.output_height

        temp_textures = []
        for tex_name, tex in self.out_textures.items():
            if tex_name not in self.registered_textures:
                width = tex["width"]
                height = tex["height"]
                format = tex["format"]
                texture = self.tex_headers(tex_name, width=width, height=height, format=format)
                temp_textures.append(texture)
        mappings["temp_textures"] = "\n".join(temp_textures)
        return Template(text).substitute(mappings)


class MagpieHook(userhook.UserHook):
    def generate(self):
        self.compute_pass += 1
        self.header["PASS"] = "$current_pass"
        headers = []
        for name in self.HEADERS:
            if name in self.header:
                value = self.header[name]
                if isinstance(value, list):
                    if name in ["IN"]:
                        headers.extend(["//!%s %s" % (name, ", ".join(value))] if value else [])
                    else:
                        for arg in value:
                            headers.append("//!%s %s" % (name, arg.strip()))
                elif isinstance(value, str):
                    headers.append("//!%s %s" % (name, value.strip()))

        hook = "\n".join(headers + self.glsl + [""])

        self.add_mappings(current_pass=self.compute_pass)
        # safe_substitute because finish passes text through substitution again
        hook = Template(hook).safe_substitute(self.mappings)
        return hook