# Rust-free versions of the Modular Stairs Kit maps (Quixel Megascans wlombbzdw; asked for by the user on
# 2026-10-12: "please remove the rust on stairs"). The kit's own files in Assets/ are only read.
#
#   T_wlombbzdw_4k_B.png, _ORM.png, _N.png  ->  *_norust.png in this folder
#
# In the kit's maps rust is the part that is warm in colour and not metal (ORM blue near 0, roughness ~0.75),
# while the clean galvanised steel is neutral grey, metal ~0.96, roughness ~0.34. So:
#   mask      = how much a pixel is rust: from (1 - metallic) and from red-minus-blue of the base colour
#   base      = rust pixels are filled from the clean steel around them (weighted blur), a little of the
#               original grain is put back, then all colour cast is taken out
#   ORM       = roughness and metallic of rust pixels are filled from the clean steel the same way; AO is kept
#   normal    = softened inside the rust mask only, so the scabs do not stay as bumps
#
# Run:  python make_norust.py [kit texture folder]
import os, sys
import numpy as np
from PIL import Image
Image.MAX_IMAGE_PIXELS = None

SRC = (r"D:\HADAL - The Descent\Assets\VaultCache\FabLibrary\Modular_Stairs_Kit-da827c1d\gltf\high"
       r"\modular_stairs_kit_wlomb_extracted\Textures")
HERE = os.path.dirname(os.path.abspath(__file__))
GRAIN = 0.55            # how much of the original fine grain is kept on the filled areas
NORMAL_SOFTEN = 0.85    # how far the normal map is flattened inside the rust mask

def box(a, r):
    """Box blur of a 2-D float array, radius r pixels, edges repeated."""
    if r <= 0:
        return a
    for ax in (0, 1):
        p = np.concatenate([np.repeat(a.take([0], ax), r + 1, ax), a, np.repeat(a.take([-1], ax), r, ax)], ax)
        c = np.cumsum(p, ax, dtype=np.float64)
        n = a.shape[ax]
        hi = c.take(range(2 * r + 1, 2 * r + 1 + n), ax); lo = c.take(range(0, n), ax)
        a = ((hi - lo) / (2 * r + 1)).astype(np.float32)
    return a

def blur(a, r):
    """Three box passes: close to a Gaussian."""
    return box(box(box(a, r), r), r)

def fill(ch, w, radii=(6, 20, 64, 220)):
    """Value of the clean pixels around each pixel (w = weight of 'clean'); wider search where little is clean."""
    out = np.zeros_like(ch); got = np.zeros_like(ch)
    for r in radii:
        bw = blur(w, r); bc = blur(ch * w, r)
        conf = np.clip(bw / 0.25, 0, 1) * (1 - got)
        out += conf * bc / np.maximum(bw, 1e-4)
        got += conf
    mean = float((ch * w).sum() / max(float(w.sum()), 1.0))
    return out + (1 - got) * mean

def load(name):
    return np.asarray(Image.open(os.path.join(SRC, name)).convert('RGB'), dtype=np.float32) / 255.0

def save(a, name):
    Image.fromarray(np.clip(a * 255.0 + 0.5, 0, 255).astype(np.uint8), 'RGB').save(os.path.join(HERE, name), compress_level=6)

if __name__ == '__main__':
    if len(sys.argv) > 1:
        SRC = sys.argv[1]
    b = load('T_wlombbzdw_4k_B.png'); o = load('T_wlombbzdw_4k_ORM.png')
    warm = b[..., 0] - b[..., 2]
    metal = o[..., 2]
    m = np.maximum(np.clip((0.90 - metal) / 0.45, 0, 1), np.clip((warm - 0.025) / 0.045, 0, 1))
    m = np.clip(blur(m, 2) * 1.6, 0, 1)                   # close small gaps, soft edge
    w = (1.0 - m) ** 2
    print('rust share of the atlas: %.1f %%' % (100 * float((m > 0.5).mean())))

    lum = 0.2126 * b[..., 0] + 0.7152 * b[..., 1] + 0.0722 * b[..., 2]
    grain = lum - blur(lum, 3)
    lum_fill = fill(lum, w)
    y = lum * (1 - m) + (lum_fill + GRAIN * grain) * m
    save(np.repeat(y[..., None], 3, axis=2), 'T_wlombbzdw_4k_B_norust.png')

    r_fill = fill(o[..., 1], w); r_grain = o[..., 1] - blur(o[..., 1], 3)
    o2 = o.copy()
    o2[..., 1] = o[..., 1] * (1 - m) + (r_fill + 0.3 * r_grain) * m
    o2[..., 2] = metal * (1 - m) + fill(metal, w) * m
    save(o2, 'T_wlombbzdw_4k_ORM_norust.png')

    n = load('T_wlombbzdw_4k_N.png') * 2.0 - 1.0
    k = (m * NORMAL_SOFTEN)[..., None]
    soft = np.stack([blur(n[..., i], 5) for i in range(3)], -1)
    n2 = n * (1 - k) + soft * k
    n2[..., 2] = np.maximum(n2[..., 2], 1e-3)
    n2 /= np.linalg.norm(n2, axis=2, keepdims=True)
    save(n2 * 0.5 + 0.5, 'T_wlombbzdw_4k_N_norust.png')
    print('written to', HERE)